diff --git a/.circleci/config.yml b/.circleci/config.yml index b0a705966a2..2f01b6de4f3 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -2731,7 +2731,7 @@ jobs: - ~/.cache/uv - restore_cache: keys: - - ui-e2e-node-deps-v2-{{ checksum "ui/litellm-dashboard/package-lock.json" }} + - ui-e2e-node-deps-v3-{{ checksum "ui/litellm-dashboard/package-lock.json" }}-{{ checksum "tests/e2e/ui/package-lock.json" }} - run: name: Install Node dependencies and Playwright # The cimg/python:3.12-browsers image already ships the Chromium system @@ -2742,11 +2742,14 @@ jobs: command: | cd ui/litellm-dashboard npm ci + cd ../../tests/e2e/ui + npm ci npx playwright install chromium - save_cache: - key: ui-e2e-node-deps-v2-{{ checksum "ui/litellm-dashboard/package-lock.json" }} + key: ui-e2e-node-deps-v3-{{ checksum "ui/litellm-dashboard/package-lock.json" }}-{{ checksum "tests/e2e/ui/package-lock.json" }} paths: - ui/litellm-dashboard/node_modules + - tests/e2e/ui/node_modules - ~/.cache/ms-playwright - run: name: Build UI from source @@ -2777,10 +2780,10 @@ jobs: name: Seed database command: | PGPASSWORD=e2epassword psql -h localhost -p 5432 -U e2euser -d litellm_e2e \ - -f ui/litellm-dashboard/e2e_tests/fixtures/seed.sql + -f tests/e2e/ui/fixtures/seed.sql - run: name: Start mock LLM server - command: uv run --no-sync python ui/litellm-dashboard/e2e_tests/fixtures/mock_llm_server/server.py + command: uv run --no-sync python tests/e2e/ui/fixtures/mock_llm_server/server.py background: true - run: name: Start LiteLLM proxy @@ -2798,7 +2801,7 @@ jobs: command: | LITELLM_LICENSE="$LITELLM_LICENSE" \ uv run --no-sync python -m litellm.proxy.proxy_cli \ - --config ui/litellm-dashboard/e2e_tests/fixtures/config.yml \ + --config tests/e2e/ui/fixtures/config.yml \ --port 4000 background: true - run: @@ -2819,15 +2822,15 @@ jobs: # Forward LITELLM_LICENSE so license.spec.ts can detect that the # proxy was launched with a license and assert premium_user=true. command: | - cd ui/litellm-dashboard + cd tests/e2e/ui LITELLM_LICENSE="$LITELLM_LICENSE" \ - npx playwright test --config e2e_tests/playwright.config.ts + npx playwright test --config playwright.config.ts no_output_timeout: 10m - store_artifacts: - path: ui/litellm-dashboard/test-results + path: tests/e2e/ui/test-results destination: e2e-test-results - store_artifacts: - path: ui/litellm-dashboard/playwright-report + path: tests/e2e/ui/playwright-report destination: e2e-playwright-report e2e_ui_testing_server_root_path: @@ -2870,17 +2873,20 @@ jobs: - ~/.cache/uv - restore_cache: keys: - - ui-e2e-node-deps-v2-{{ checksum "ui/litellm-dashboard/package-lock.json" }} + - ui-e2e-node-deps-v3-{{ checksum "ui/litellm-dashboard/package-lock.json" }}-{{ checksum "tests/e2e/ui/package-lock.json" }} - run: name: Install Node dependencies and Playwright command: | cd ui/litellm-dashboard npm ci + cd ../../tests/e2e/ui + npm ci npx playwright install chromium - save_cache: - key: ui-e2e-node-deps-v2-{{ checksum "ui/litellm-dashboard/package-lock.json" }} + key: ui-e2e-node-deps-v3-{{ checksum "ui/litellm-dashboard/package-lock.json" }}-{{ checksum "tests/e2e/ui/package-lock.json" }} paths: - ui/litellm-dashboard/node_modules + - tests/e2e/ui/node_modules - ~/.cache/ms-playwright - run: name: Build UI from source @@ -2902,10 +2908,10 @@ jobs: name: Seed database command: | PGPASSWORD=e2epassword psql -h localhost -p 5432 -U e2euser -d litellm_e2e \ - -f ui/litellm-dashboard/e2e_tests/fixtures/seed.sql + -f tests/e2e/ui/fixtures/seed.sql - run: name: Start mock LLM server - command: uv run --no-sync python ui/litellm-dashboard/e2e_tests/fixtures/mock_llm_server/server.py + command: uv run --no-sync python tests/e2e/ui/fixtures/mock_llm_server/server.py background: true - run: name: Start LiteLLM proxy under a server root path @@ -2918,7 +2924,7 @@ jobs: command: | LITELLM_LICENSE="$LITELLM_LICENSE" \ uv run --no-sync python -m litellm.proxy.proxy_cli \ - --config ui/litellm-dashboard/e2e_tests/fixtures/config.yml \ + --config tests/e2e/ui/fixtures/config.yml \ --port 4000 background: true - run: @@ -2937,15 +2943,15 @@ jobs: - run: name: Run migration smoke under SERVER_ROOT_PATH command: | - cd ui/litellm-dashboard + cd tests/e2e/ui LITELLM_LICENSE="$LITELLM_LICENSE" \ - npx playwright test --config e2e_tests/migration.serverRootPath.config.ts + npx playwright test --config migration.serverRootPath.config.ts no_output_timeout: 10m - store_artifacts: - path: ui/litellm-dashboard/test-results + path: tests/e2e/ui/test-results destination: e2e-server-root-path-test-results - store_artifacts: - path: ui/litellm-dashboard/playwright-report + path: tests/e2e/ui/playwright-report destination: e2e-server-root-path-playwright-report build_docker_database_image: diff --git a/.circleci/scripts/classify_changes.sh b/.circleci/scripts/classify_changes.sh index 2c15428be6a..2ca2654a207 100755 --- a/.circleci/scripts/classify_changes.sh +++ b/.circleci/scripts/classify_changes.sh @@ -8,7 +8,7 @@ has_backend=false while IFS= read -r file || [ -n "$file" ]; do [ -n "$file" ] || continue case "$file" in - ui/*) has_client=true ;; + ui/* | tests/e2e/ui/*) has_client=true ;; docs/* | *.md | *.mdx) : ;; *) has_backend=true ;; esac diff --git a/.github/ISSUE_TEMPLATE/bug_report.yml b/.github/ISSUE_TEMPLATE/bug_report.yml index bbe4b76775d..665f8456f0b 100644 --- a/.github/ISSUE_TEMPLATE/bug_report.yml +++ b/.github/ISSUE_TEMPLATE/bug_report.yml @@ -30,7 +30,7 @@ body: id: steps-to-reproduce attributes: label: Steps to Reproduce - description: Please provide detailed steps to reproduce this bug(A curl/python code to reproduce the bug) + description: Please provide a numbered list of the exact steps to reproduce this bug (include a curl/python snippet to reproduce it). Number each step (1., 2., 3., ...) in the order you performed them. placeholder: | 1. config.yaml file/ .env file/ etc. 2. Run the following code... diff --git a/.github/pull_request_template.md b/.github/pull_request_template.md index d7e80b32749..1301bfb0e60 100644 --- a/.github/pull_request_template.md +++ b/.github/pull_request_template.md @@ -1,3 +1,18 @@ +## TLDR + + + +Problem this solves: + +- +- ... + +How it solves it: + +- +- ... + ## Relevant issues diff --git a/.github/workflows/_test-unit-base.yml b/.github/workflows/_test-unit-base.yml index 92230fc8892..7fd66e3325e 100644 --- a/.github/workflows/_test-unit-base.yml +++ b/.github/workflows/_test-unit-base.yml @@ -80,7 +80,7 @@ jobs: - name: Install dependencies if: steps.changes.outputs.decision != 'skip' run: | - .github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router + .github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router --extra saml - name: Generate Prisma client if: steps.changes.outputs.decision != 'skip' diff --git a/.github/workflows/codspeed.yml b/.github/workflows/codspeed.yml index 54a8e53d7a3..a69e50b5753 100644 --- a/.github/workflows/codspeed.yml +++ b/.github/workflows/codspeed.yml @@ -5,10 +5,24 @@ on: branches: - main - litellm_internal_staging + paths: + - "litellm/**" + - "tests/benchmarks/**" + - "pyproject.toml" + - "uv.lock" + - ".github/workflows/codspeed.yml" + - ".github/actions/setup-uv-with-retries/**" pull_request: branches: - main - litellm_internal_staging + paths: + - "litellm/**" + - "tests/benchmarks/**" + - "pyproject.toml" + - "uv.lock" + - ".github/workflows/codspeed.yml" + - ".github/actions/setup-uv-with-retries/**" # Allow CodSpeed to trigger backtest performance analysis # in order to generate initial data workflow_dispatch: diff --git a/.github/workflows/image-scan.yml b/.github/workflows/image-scan.yml index 8d791ca5bc7..4d4a3242399 100644 --- a/.github/workflows/image-scan.yml +++ b/.github/workflows/image-scan.yml @@ -9,6 +9,7 @@ on: - "litellm_**" paths: - docker/Dockerfile.non_root + - tests/proxy_migration_tests/test_offline_image_migration.py - uv.lock - ui/litellm-dashboard/package-lock.json - .github/workflows/image-scan.yml @@ -51,6 +52,23 @@ jobs: - name: Build runtime image run: docker build -f docker/Dockerfile.non_root -t litellm-image-scan:${{ github.sha }} . + # The prisma bake must migrate a fresh DB with no egress as an arbitrary + # non-root uid (OpenShift restricted-v2 / air-gapped / readOnlyRootFilesystem). + # `docker run` as the default uid with network hides a broken bake because + # the migration entrypoint exits 0 even when it applied nothing; asserting + # the schema was created is what catches it. + - name: Set up Python + uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0 + with: + python-version: "3.12" + + - name: Verify offline migration as a non-root uid + env: + LITELLM_IMAGE: litellm-image-scan:${{ github.sha }} + run: | + python -m pip install "pytest==9.0.3" + python -m pytest tests/proxy_migration_tests/test_offline_image_migration.py -v + # Scans the whole shipped artifact: OS/apk plus every language package # baked into the image, including ones no lockfile declares (e.g. prisma's # vendored node engine) that osv-scan cannot see. osv-scan stays the fast diff --git a/.github/workflows/mutation-test.yml b/.github/workflows/mutation-test.yml index 6684952b998..da4fe073a6a 100644 --- a/.github/workflows/mutation-test.yml +++ b/.github/workflows/mutation-test.yml @@ -55,7 +55,7 @@ jobs: - name: Install dependencies run: | - .github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router + .github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router --extra saml - name: Generate Prisma client env: diff --git a/.github/workflows/test-code-quality.yml b/.github/workflows/test-code-quality.yml index 9d28ca211cf..ae31395521a 100644 --- a/.github/workflows/test-code-quality.yml +++ b/.github/workflows/test-code-quality.yml @@ -115,6 +115,9 @@ jobs: - name: check_fastuuid_usage run: uv run --no-sync python ./tests/code_coverage_tests/check_fastuuid_usage.py + - name: check_e2e_no_raw_requests + run: uv run --no-sync python ./tests/code_coverage_tests/check_e2e_no_raw_requests.py + - name: memory_test run: uv run --no-sync python ./tests/code_coverage_tests/memory_test.py diff --git a/.github/workflows/test-litellm-ui-unit.yml b/.github/workflows/test-litellm-ui-unit.yml index f0f9f504752..5374a0059de 100644 --- a/.github/workflows/test-litellm-ui-unit.yml +++ b/.github/workflows/test-litellm-ui-unit.yml @@ -19,7 +19,7 @@ concurrency: jobs: ui-unit-tests: - runs-on: ubuntu-latest + runs-on: ubuntu-latest-16-cores timeout-minutes: 20 defaults: run: @@ -50,8 +50,8 @@ jobs: if [ -n "$BASE_SHA" ]; then echo "Pull request: running only tests related to changes since $BASE_SHA" npm run test -- --run --changed "$BASE_SHA" --passWithNoTests \ - --pool forks --poolOptions.forks.maxForks=4 + --pool forks --poolOptions.forks.maxForks=14 else echo "Push to $GITHUB_REF_NAME: running the full suite" - npm run test -- --run --pool forks --poolOptions.forks.maxForks=4 + npm run test -- --run --pool forks --poolOptions.forks.maxForks=14 fi diff --git a/.github/workflows/test-unit-proxy-endpoints.yml b/.github/workflows/test-unit-proxy-endpoints.yml index cbceabe91da..a8c853c0d1d 100644 --- a/.github/workflows/test-unit-proxy-endpoints.yml +++ b/.github/workflows/test-unit-proxy-endpoints.yml @@ -47,6 +47,7 @@ jobs: tests/test_litellm/proxy/rag_endpoints tests/test_litellm/proxy/realtime_endpoints tests/test_litellm/proxy/ui_crud_endpoints + tests/test_litellm/proxy/config_resolvers tests/test_litellm/proxy/utils workers: 2 reruns: 2 diff --git a/.github/workflows/test_server_root_path.yml b/.github/workflows/test_server_root_path.yml index f59cee29893..01f70511e79 100644 --- a/.github/workflows/test_server_root_path.yml +++ b/.github/workflows/test_server_root_path.yml @@ -106,8 +106,8 @@ jobs: with: node-version: "20" - - name: Install UI deps and Chromium - working-directory: ui/litellm-dashboard + - name: Install e2e deps and Chromium + working-directory: tests/e2e/ui run: | retry() { local attempt=1 @@ -131,17 +131,17 @@ jobs: retry npx playwright install --with-deps chromium - name: Run SERVER_ROOT_PATH redirect e2e - working-directory: ui/litellm-dashboard + working-directory: tests/e2e/ui env: SERVER_ROOT_PATH: ${{ matrix.root_path }} - run: npx playwright test --config=e2e_tests/serverRootPath.config.ts + run: npx playwright test --config=serverRootPath.config.ts - name: Upload Playwright artifacts on failure if: failure() uses: actions/upload-artifact@ea165f8d65b6e75b540449e92b4886f43607fa02 # v4.6.2 with: name: playwright-trace-${{ strategy.job-index }} - path: ui/litellm-dashboard/test-results/ + path: tests/e2e/ui/test-results/ retention-days: 7 - name: Cleanup diff --git a/.github/workflows/weekly_load_anomaly.yml b/.github/workflows/weekly_load_anomaly.yml new file mode 100644 index 00000000000..4c2103f026d --- /dev/null +++ b/.github/workflows/weekly_load_anomaly.yml @@ -0,0 +1,81 @@ +name: "Weekly Load Anomaly Check" + +on: + schedule: + - cron: "0 12 * * 6" + workflow_dispatch: + +permissions: + contents: read + +jobs: + weekly-load-anomaly: + if: github.event_name != 'schedule' || github.repository == 'BerriAI/litellm' + runs-on: ubuntu-latest + timeout-minutes: 45 + services: + postgres: + image: postgres:16.6 + env: + POSTGRES_USER: llmproxy + POSTGRES_PASSWORD: dbpassword9090 + POSTGRES_DB: litellm + ports: + - 5432:5432 + options: >- + --health-cmd "pg_isready -U llmproxy" + --health-interval 5s + --health-timeout 5s + --health-retries 10 + env: + DATABASE_URL: postgresql://llmproxy:dbpassword9090@localhost:5432/litellm + LITELLM_MASTER_KEY: sk-weekly-anomaly-check + ANTHROPIC_API_KEY: ${{ secrets.ANTHROPIC_API_KEY }} + AWS_BEARER_TOKEN_BEDROCK: ${{ secrets.AWS_BEARER_TOKEN_BEDROCK }} + steps: + - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 + with: + persist-credentials: false + + - name: Set up Python + uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0 + with: + python-version: "3.12" + + - name: Set up uv + uses: ./.github/actions/setup-uv-with-retries + with: + version: "0.10.9" + + - name: Install dependencies + run: | + .github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra proxy + + - name: Generate Prisma client + env: + PRISMA_BINARY_CACHE_DIR: ${{ runner.temp }}/prisma-cache + run: | + uv run --no-sync prisma generate --schema litellm/proxy/schema.prisma + + - name: Start the proxy + run: | + nohup uv run --no-sync litellm --config tests/e2e/load/weekly_anomaly_config.yml --port 4000 > proxy.log 2>&1 & + for _ in $(seq 1 90); do + if curl -fs http://localhost:4000/health/liveliness > /dev/null; then + exit 0 + fi + sleep 2 + done + echo "proxy never became live" + tail -n 100 proxy.log + exit 1 + + - name: Run the weekly session anomaly test + env: + E2E_WEEKLY_ANOMALY: "1" + run: | + uv run --no-sync pytest tests/e2e/load/test_weekly_session_anomaly_e2e.py -v --tb=short -rA + + - name: Show proxy log on failure + if: failure() + run: tail -n 300 proxy.log diff --git a/Dockerfile b/Dockerfile index 9977ebb82d7..a127cdabd59 100644 --- a/Dockerfile +++ b/Dockerfile @@ -64,6 +64,7 @@ RUN uv sync --frozen --no-install-project --no-install-workspace --no-default-gr --extra proxy-runtime \ --extra extra_proxy \ --extra semantic-router \ + --extra saml \ --python python3 # Copy full source tree @@ -84,6 +85,7 @@ RUN uv sync --frozen --no-default-groups --no-editable \ --extra proxy-runtime \ --extra extra_proxy \ --extra semantic-router \ + --extra saml \ --python python3 RUN HOME=/opt/prisma XDG_CACHE_HOME=/opt/prisma/.cache PRISMA_BINARY_CACHE_DIR=/opt/prisma/binaries \ diff --git a/backend/routes/allowlist.py b/backend/routes/allowlist.py index 02574ca505d..f3a028f5805 100644 --- a/backend/routes/allowlist.py +++ b/backend/routes/allowlist.py @@ -18,6 +18,7 @@ BACKEND_PATH_PREFIXES: tuple[str, ...] = ( "/team/", "/v2/team/", "/organization/", + "/v2/organization/", "/customer/", "/end_user/", "/sso/", diff --git a/docker/Dockerfile.database b/docker/Dockerfile.database index 34c9c606991..9ee076ce825 100644 --- a/docker/Dockerfile.database +++ b/docker/Dockerfile.database @@ -62,6 +62,7 @@ RUN uv sync --frozen --no-install-project --no-install-workspace --no-default-gr --extra proxy-runtime \ --extra extra_proxy \ --extra semantic-router \ + --extra saml \ --python python3 # Copy full source tree @@ -82,6 +83,7 @@ RUN uv sync --frozen --no-default-groups --no-editable \ --extra proxy-runtime \ --extra extra_proxy \ --extra semantic-router \ + --extra saml \ --python python3 RUN HOME=/opt/prisma XDG_CACHE_HOME=/opt/prisma/.cache PRISMA_BINARY_CACHE_DIR=/opt/prisma/binaries \ diff --git a/docker/Dockerfile.non_root b/docker/Dockerfile.non_root index 839f5da565c..946b4de6f5e 100644 --- a/docker/Dockerfile.non_root +++ b/docker/Dockerfile.non_root @@ -54,7 +54,6 @@ ENV UV_PROJECT_ENVIRONMENT=/app/.venv \ UV_LINK_MODE=copy \ PATH="/app/.venv/bin:${PATH}" \ LITELLM_NON_ROOT=true \ - PRISMA_BINARY_CACHE_DIR=/app/.cache/prisma-python/binaries \ XDG_CACHE_HOME=/app/.cache # Copy dependency metadata first for layer caching @@ -69,6 +68,7 @@ RUN --mount=type=cache,target=/app/.cache/uv,id=litellm-uv-cache \ --extra proxy-runtime \ --extra extra_proxy \ --extra semantic-router \ + --extra saml \ --python python3 # Copy full source tree @@ -95,6 +95,7 @@ RUN --mount=type=cache,target=/app/.cache/uv,id=litellm-uv-cache \ --extra proxy-runtime \ --extra extra_proxy \ --extra semantic-router \ + --extra saml \ --python python3 \ --no-sources-package litellm-proxy-extras; \ else \ @@ -103,10 +104,13 @@ RUN --mount=type=cache,target=/app/.cache/uv,id=litellm-uv-cache \ --extra proxy-runtime \ --extra extra_proxy \ --extra semantic-router \ + --extra saml \ --python python3; \ fi -RUN prisma generate --schema=./schema.prisma +RUN HOME=/opt/prisma XDG_CACHE_HOME=/opt/prisma/.cache PRISMA_BINARY_CACHE_DIR=/opt/prisma/binaries \ + npm_config_cache=/root/.npm \ + prisma generate --schema=./schema.prisma RUN sed -i 's/\r$//' docker/entrypoint.sh && chmod +x docker/entrypoint.sh && \ sed -i 's/\r$//' docker/prod_entrypoint.sh && chmod +x docker/prod_entrypoint.sh @@ -127,8 +131,6 @@ RUN for i in 1 2 3; do \ # the rest of the builder's /app is source and build metadata that must not # ship (manifest-scanning tools attribute everything in it to this image). # entrypoint.sh invokes litellm/proxy/prisma_migration.py by source path. -# Prisma caches live under /app/.cache here (XDG_CACHE_HOME / -# PRISMA_BINARY_CACHE_DIR) so the runtime prisma generate finds them. COPY --from=builder /app/.venv /app/.venv COPY --from=builder /app/docker /app/docker COPY --from=builder /app/schema.prisma /app/schema.prisma @@ -138,21 +140,35 @@ COPY --from=builder /app/litellm/proxy/prisma_migration.py /app/litellm/proxy/pr # enterprise.enterprise_hooks from it) COPY --from=builder /app/enterprise /app/enterprise COPY --from=builder /app/litellm-proxy-extras /app/litellm-proxy-extras -COPY --from=builder /app/.cache /app/.cache +# Prisma CLI + engines are baked under /opt/prisma, a fixed path every runtime +# uid can read and that no cache volume mount shadows (unlike /app/.cache or +# $HOME/.cache under readOnlyRootFilesystem + emptyDir or arbitrary-uid setups). +# PRISMA_CLI_QUERY_ENGINE_TYPE=binary makes the CLI use the baked binary query +# engine directly, so `prisma migrate deploy` on a fresh database needs no npm +# and no network access; without it the CLI looks for the library engine, which +# prisma stopped baking, and falls back to a download that fails offline or as a +# non-writable uid (#33650, #24554). +COPY --from=builder /opt/prisma /opt/prisma COPY --from=builder /var/lib/litellm/ui /var/lib/litellm/ui COPY --from=builder /var/lib/litellm/assets /var/lib/litellm/assets +# XDG_CACHE_HOME is intentionally left unset so it falls back to $HOME/.cache +# (/app/.cache, writable by the runtime uid). The prisma bake at the read-only +# /opt/prisma is anchored by PRISMA_BINARY_CACHE_DIR / PRISMA_CLI_PATH, so +# nothing needs XDG to point there; pointing it at the read-only bake would +# deny any XDG-aware library that writes a cache at runtime. ENV PATH="/app/.venv/bin:${PATH}" \ - PRISMA_BINARY_CACHE_DIR=/app/.cache/prisma-python/binaries \ + PRISMA_BINARY_CACHE_DIR=/opt/prisma/binaries \ + PRISMA_CLI_PATH=/opt/prisma/binaries/node_modules/.bin/prisma \ + PRISMA_CLI_QUERY_ENGINE_TYPE=binary \ HOME=/app \ LITELLM_NON_ROOT=true \ - XDG_CACHE_HOME=/app/.cache \ PRISMA_SKIP_POSTINSTALL_GENERATE=1 \ PRISMA_HIDE_UPDATE_MESSAGE=1 \ PRISMA_ENGINES_CHECKSUM_IGNORE_MISSING=1 \ PRISMA_OFFLINE_MODE=true -RUN mkdir -p /nonexistent /var/lib/litellm/assets /var/lib/litellm/ui && \ +RUN mkdir -p /nonexistent /app/.cache /var/lib/litellm/assets /var/lib/litellm/ui && \ chown -R nobody:nogroup /app /var/lib/litellm/ui /var/lib/litellm/assets /nonexistent && \ PRISMA_PATH=$(python -c "import os, prisma; print(os.path.dirname(prisma.__file__))") && \ chown -R nobody:nogroup "$PRISMA_PATH" && \ @@ -165,12 +181,14 @@ RUN mkdir -p /nonexistent /var/lib/litellm/assets /var/lib/litellm/ui && \ [ -n "$LITELLM_PROXY_EXTRAS_PATH" ] && chmod -R g=u "$LITELLM_PROXY_EXTRAS_PATH" || true && \ chmod -R g+w "$PRISMA_PATH" /var/lib/litellm/ui /var/lib/litellm/assets && \ [ -n "$LITELLM_PROXY_EXTRAS_PATH" ] && chmod -R g+w "$LITELLM_PROXY_EXTRAS_PATH" || true && \ - chmod -R g+rX "$PRISMA_PATH" /var/lib/litellm/ui /var/lib/litellm/assets /app/.cache + chmod -R g+rX "$PRISMA_PATH" /var/lib/litellm/ui /var/lib/litellm/assets && \ + chmod -R a+rX /opt/prisma && \ + test -x /opt/prisma/binaries/node_modules/.bin/prisma && \ + test -f /opt/prisma/binaries/node_modules/prisma/build/index.js && \ + ls /opt/prisma/binaries/node_modules/@prisma/engines/query-engine-* >/dev/null 2>&1 USER 65534 -RUN prisma generate --schema=./schema.prisma - EXPOSE 4000/tcp ENTRYPOINT ["/app/docker/prod_entrypoint.sh"] diff --git a/gateway/Dockerfile b/gateway/Dockerfile index da2f2c9c1e0..4b000912393 100644 --- a/gateway/Dockerfile +++ b/gateway/Dockerfile @@ -46,6 +46,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \ --extra proxy-runtime \ --extra extra_proxy \ --extra semantic-router \ + --extra bedrock-realtime \ --python python3 # Stage 2 — copy source and install the project + workspace members. @@ -57,6 +58,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \ --extra proxy-runtime \ --extra extra_proxy \ --extra semantic-router \ + --extra bedrock-realtime \ --python python3 RUN mkdir -p /home/nonroot && \ diff --git a/litellm-rust/crates/ai-gateway/src/messages/common_utils.rs b/litellm-rust/crates/ai-gateway/src/messages/common_utils.rs index 4b906155665..68ecc3f17c1 100644 --- a/litellm-rust/crates/ai-gateway/src/messages/common_utils.rs +++ b/litellm-rust/crates/ai-gateway/src/messages/common_utils.rs @@ -50,3 +50,15 @@ pub(super) fn has_header(headers: &[(String, String)], name: &str) -> bool { .iter() .any(|(key, _)| key.eq_ignore_ascii_case(name)) } + +pub(super) fn has_bearer_auth(headers: &[(String, String)]) -> bool { + headers.iter().any(|(name, value)| { + if !name.eq_ignore_ascii_case("authorization") { + return false; + } + let value = value.trim(); + value.len() > 7 + && value[..7].eq_ignore_ascii_case("bearer ") + && !value[7..].trim().is_empty() + }) +} diff --git a/litellm-rust/crates/ai-gateway/src/messages/prepare.rs b/litellm-rust/crates/ai-gateway/src/messages/prepare.rs index 624c3598fb0..9a027490eb6 100644 --- a/litellm-rust/crates/ai-gateway/src/messages/prepare.rs +++ b/litellm-rust/crates/ai-gateway/src/messages/prepare.rs @@ -3,7 +3,7 @@ use litellm_core::CoreResult; use litellm_core::messages::transformation::MessagesAuthStrategy; use litellm_core::routing_utils::provider::{CustomLlmProvider, get_custom_llm_provider}; -use super::common_utils::{has_header, messages_provider_config, string_headers}; +use super::common_utils::{has_bearer_auth, has_header, messages_provider_config, string_headers}; use super::types::{MessagesRequest, ProviderMessagesRequest}; pub(super) fn prepare_messages_call( @@ -33,7 +33,9 @@ pub(super) fn prepare_messages_call( let mut headers = string_headers(request.extra_headers)?; let auth_strategy = config.auth_strategy(); - if !has_header(&headers, auth_strategy.header_name()) { + let already_authorized = has_header(&headers, auth_strategy.header_name()) + || (config.accepts_bearer_auth() && has_bearer_auth(&headers)); + if !already_authorized { let api_key = config.resolve_api_key(request.api_key, &env_lookup)?; let auth_header = match auth_strategy { MessagesAuthStrategy::Bearer => { diff --git a/litellm-rust/crates/ai-gateway/src/messages/tests.rs b/litellm-rust/crates/ai-gateway/src/messages/tests.rs index a2d0f6fae23..23a53e98045 100644 --- a/litellm-rust/crates/ai-gateway/src/messages/tests.rs +++ b/litellm-rust/crates/ai-gateway/src/messages/tests.rs @@ -6,7 +6,7 @@ use tokio::io::{AsyncReadExt, AsyncWriteExt}; use tokio::net::{TcpListener, TcpStream}; use super::common_utils::{ - has_header, messages_provider_config, string_headers, truncate_error_body, + has_bearer_auth, has_header, messages_provider_config, string_headers, truncate_error_body, }; use super::{MessagesRequest, messages}; @@ -85,6 +85,34 @@ fn has_header_is_case_insensitive() { assert!(!has_header(&headers, "authorization")); } +#[test] +fn has_bearer_auth_requires_a_nonempty_bearer_token() { + assert!(has_bearer_auth(&[( + "Authorization".to_string(), + "Bearer tok".to_string() + )])); + assert!(has_bearer_auth(&[( + "authorization".to_string(), + "bearer tok".to_string() + )])); + assert!(!has_bearer_auth(&[( + "authorization".to_string(), + "Bearer ".to_string() + )])); + assert!(!has_bearer_auth(&[( + "authorization".to_string(), + String::new() + )])); + assert!(!has_bearer_auth(&[( + "authorization".to_string(), + "Basic abc".to_string() + )])); + assert!(!has_bearer_auth(&[( + "x-api-key".to_string(), + "sk".to_string() + )])); +} + #[tokio::test] async fn messages_round_trip_builds_azure_request_and_passes_response_through() { let listener = TcpListener::bind("127.0.0.1:0").await.expect("binds"); @@ -252,6 +280,112 @@ async fn messages_does_not_duplicate_auth_when_x_api_key_supplied() { assert!(!head.contains("rust-fallback-key"), "{head}"); } +#[tokio::test] +async fn messages_forwards_entra_id_bearer_without_requiring_api_key() { + let listener = TcpListener::bind("127.0.0.1:0").await.expect("binds"); + let addr = listener.local_addr().expect("addr"); + + let server = tokio::spawn(async move { + let (mut socket, _) = listener.accept().await.expect("accepts request"); + let request = read_http_request(&mut socket).await; + let response_body = + r#"{"id":"msg_3","type":"message","role":"assistant","content":[],"model":"m"}"#; + socket + .write_all(write_response(response_body).as_bytes()) + .await + .expect("writes response"); + request + }); + + let mut headers = Map::new(); + headers.insert( + "Authorization".to_string(), + Value::String("Bearer entra-token".to_string()), + ); + + messages(MessagesRequest { + model: "claude-sonnet-4-5", + body: json!({"model": "claude-sonnet-4-5", "max_tokens": 8, "messages": []}), + api_key: None, + api_base: Some(&format!("http://{addr}")), + custom_llm_provider: Some("azure_ai"), + extra_headers: Some(headers), + timeout: Some(Duration::from_secs(5)), + }) + .await + .expect("entra id request succeeds without api key"); + + let request = server.await.expect("server task completes"); + let head = request + .split_once("\r\n\r\n") + .expect("has body") + .0 + .to_ascii_lowercase(); + assert!(head.contains("authorization: bearer entra-token"), "{head}"); + assert!(!head.contains("x-api-key"), "{head}"); +} + +#[tokio::test] +async fn messages_requires_auth_when_no_key_and_no_header() { + let err = messages(MessagesRequest { + model: "claude-sonnet-4-5", + body: json!({"model": "claude-sonnet-4-5", "max_tokens": 8, "messages": []}), + api_key: None, + api_base: Some("http://127.0.0.1:1"), + custom_llm_provider: Some("azure_ai"), + extra_headers: None, + timeout: Some(Duration::from_millis(50)), + }) + .await + .expect_err("missing auth errors"); + + assert!(matches!(err, CoreError::Auth(_))); +} + +#[tokio::test] +async fn messages_ignores_malformed_authorization_and_uses_api_key() { + let listener = TcpListener::bind("127.0.0.1:0").await.expect("binds"); + let addr = listener.local_addr().expect("addr"); + + let server = tokio::spawn(async move { + let (mut socket, _) = listener.accept().await.expect("accepts request"); + let request = read_http_request(&mut socket).await; + let response_body = + r#"{"id":"msg_4","type":"message","role":"assistant","content":[],"model":"m"}"#; + socket + .write_all(write_response(response_body).as_bytes()) + .await + .expect("writes response"); + request + }); + + let mut headers = Map::new(); + headers.insert( + "Authorization".to_string(), + Value::String("Bearer ".to_string()), + ); + + messages(MessagesRequest { + model: "claude-sonnet-4-5", + body: json!({"model": "claude-sonnet-4-5", "max_tokens": 8, "messages": []}), + api_key: Some("sk-azure"), + api_base: Some(&format!("http://{addr}")), + custom_llm_provider: Some("azure_ai"), + extra_headers: Some(headers), + timeout: Some(Duration::from_secs(5)), + }) + .await + .expect("falls back to api key"); + + let request = server.await.expect("server task completes"); + let head = request + .split_once("\r\n\r\n") + .expect("has body") + .0 + .to_ascii_lowercase(); + assert!(head.contains("x-api-key: sk-azure"), "{head}"); +} + #[tokio::test] async fn messages_maps_provider_error_status_to_http_error() { let listener = TcpListener::bind("127.0.0.1:0").await.expect("binds"); diff --git a/litellm-rust/crates/core/src/messages/transformation.rs b/litellm-rust/crates/core/src/messages/transformation.rs index 3a34a58de6f..b478e20d24b 100644 --- a/litellm-rust/crates/core/src/messages/transformation.rs +++ b/litellm-rust/crates/core/src/messages/transformation.rs @@ -35,6 +35,10 @@ pub trait AnthropicMessagesProviderConfig: Sync { MessagesAuthStrategy::Header("x-api-key") } + fn accepts_bearer_auth(&self) -> bool { + false + } + fn default_headers(&self) -> &'static [(&'static str, &'static str)] { &[ ("anthropic-version", "2023-06-01"), diff --git a/litellm-rust/crates/core/src/providers/azure_ai/messages/transformation.rs b/litellm-rust/crates/core/src/providers/azure_ai/messages/transformation.rs index 6935bb4604b..7b958c77ba3 100644 --- a/litellm-rust/crates/core/src/providers/azure_ai/messages/transformation.rs +++ b/litellm-rust/crates/core/src/providers/azure_ai/messages/transformation.rs @@ -163,6 +163,10 @@ impl AnthropicMessagesProviderConfig for AzureAnthropicMessagesConfig { self.anthropic.auth_strategy() } + fn accepts_bearer_auth(&self) -> bool { + true + } + fn default_headers(&self) -> &'static [(&'static str, &'static str)] { self.anthropic.default_headers() } @@ -294,6 +298,11 @@ mod tests { ); } + #[test] + fn accepts_bearer_auth_for_entra_id() { + assert!(AZURE_ANTHROPIC_MESSAGES_CONFIG.accepts_bearer_auth()); + } + #[test] fn default_headers_match_python() { assert_eq!( diff --git a/litellm/__init__.py b/litellm/__init__.py index c9bdaea6af1..b3f49ad63ba 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -212,6 +212,9 @@ filter_invalid_headers: Optional[bool] = False add_user_information_to_llm_headers: Optional[bool] = ( None # adds user_id, team_id, token hash (params from StandardLoggingMetadata) to request headers ) +overwrite_user_with_key_hash: bool = ( + False # force the outgoing `user` param to the hashed api key, so providers see a stable, tamper-proof id +) store_audit_logs = False # Enterprise feature, allow users to see audit logs skip_system_message_in_guardrail: bool = False skip_tool_message_in_guardrail: bool = False @@ -428,7 +431,9 @@ default_team_settings: Optional[List] = None max_user_budget: Optional[float] = None default_max_internal_user_budget: Optional[float] = None max_internal_user_budget: Optional[float] = None -max_ui_session_budget: Optional[float] = 0.25 # $0.25 USD budgets for UI Chat sessions +max_ui_session_budget: Optional[float] = ( + 1.0 # USD budget for each dashboard login session (playground, test connection) +) internal_user_budget_duration: Optional[str] = None tag_budget_config: Optional[Dict[str, "BudgetConfig"]] = None max_end_user_budget: Optional[float] = None diff --git a/litellm/constants.py b/litellm/constants.py index 9dd80750f95..72da3641bcc 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -264,6 +264,9 @@ MAX_REDIS_BUFFER_DEQUEUE_COUNT = int(os.getenv("MAX_REDIS_BUFFER_DEQUEUE_COUNT", # Bounds asyncio.Queue() instances (log queues, spend update queues, etc.) to prevent unbounded memory growth LITELLM_ASYNCIO_QUEUE_MAXSIZE = int(os.getenv("LITELLM_ASYNCIO_QUEUE_MAXSIZE", 1000)) TOOL_POLICY_CACHE_TTL_SECONDS = int(os.getenv("TOOL_POLICY_CACHE_TTL_SECONDS", 60)) +GUARDRAIL_SCANNED_MESSAGES_CACHE_TTL_SECONDS = int( + os.getenv("GUARDRAIL_SCANNED_MESSAGES_CACHE_TTL_SECONDS", 24 * 60 * 60) +) # Aggregation threshold: default to 80% of the asyncio queue maxsize so the check can always trigger. # Must be < LITELLM_ASYNCIO_QUEUE_MAXSIZE; if set higher the aggregation logic will never fire. MAX_SIZE_IN_MEMORY_QUEUE = int(os.getenv("MAX_SIZE_IN_MEMORY_QUEUE", int(LITELLM_ASYNCIO_QUEUE_MAXSIZE * 0.8))) @@ -1147,6 +1150,7 @@ BEDROCK_CONVERSE_MODELS = [ "anthropic.claude-sonnet-4-5-20250929-v1:0", "anthropic.claude-fable-5", "anthropic.claude-sonnet-5", + "anthropic.claude-opus-5", "anthropic.claude-opus-4-8", "anthropic.claude-opus-4-7", "anthropic.claude-opus-4-6-v1:0", @@ -1530,6 +1534,7 @@ LITELLM_SETTINGS_SAFE_DB_OVERRIDES = [ # test_general_settings_ui_fields_are_db_overridable enforces that pairing. "enable_anthropic_prompt_caching", "anthropic_prompt_caching_ttl", + "max_ui_session_budget", ] SPECIAL_LITELLM_AUTH_TOKEN = ["ui-token"] DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL = int(os.getenv("DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL", 60)) diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index a40a8e1389c..96aed20529f 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -2191,6 +2191,13 @@ def batch_cost_calculator( return total_prompt_cost, total_completion_cost +def _summable_prompt_token_fields(prompt_tokens_details: BaseModel) -> List[str]: + field_names = list(type(prompt_tokens_details).model_fields) + if getattr(prompt_tokens_details, "cache_write_tokens", None) is None: + return field_names + return [attr for attr in field_names if attr != "cache_creation_tokens"] + + class BaseTokenUsageProcessor: @staticmethod def combine_usage_objects(usage_objects: List[Usage]) -> Usage: @@ -2225,7 +2232,7 @@ class BaseTokenUsageProcessor: # Check what keys exist in the model's prompt_tokens_details # Access model_fields on the class, not the instance, to avoid Pydantic 2.11+ deprecation warnings - for attr in type(usage.prompt_tokens_details).model_fields: + for attr in _summable_prompt_token_fields(usage.prompt_tokens_details): if ( hasattr(usage.prompt_tokens_details, attr) and not attr.startswith("_") diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index 856556f7c56..f639ad49d5e 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -1,3 +1,5 @@ +import contextvars +import hashlib import os import secrets from datetime import datetime @@ -46,7 +48,10 @@ if TYPE_CHECKING: dc = DualCache() -from litellm.constants import PRE_CALL_EXECUTED_GUARDRAILS_KEY +from litellm.constants import ( + GUARDRAIL_SCANNED_MESSAGES_CACHE_TTL_SECONDS, + PRE_CALL_EXECUTED_GUARDRAILS_KEY, +) from litellm.exceptions import ( BlockedPiiEntityError, GuardrailRaisedException, @@ -60,6 +65,10 @@ from litellm.exceptions import ( # proxy's metadata sanitizer. _PRE_CALL_EXECUTED_TOKEN = secrets.token_hex(16) +_guardrail_self_recorded: contextvars.ContextVar[bool] = contextvars.ContextVar( + "litellm_guardrail_self_recorded", default=False +) + def _strict_guardrail_modes_enabled() -> bool: """Whether guardrail-mode validation raises (default) or logs a warning. @@ -113,6 +122,8 @@ class CustomGuardrail(CustomLogger): on_sensitive_data: Optional[str] = None, sensitive_data_route_to_model: Optional[str] = None, sticky_session_routing: bool = True, + run_in_parallel: bool = False, + only_scan_new_messages: bool = False, **kwargs, ): """ @@ -131,6 +142,9 @@ class CustomGuardrail(CustomLogger): on_sensitive_data: Action when sensitive data is detected. 'block' (default) or 'route' sensitive_data_route_to_model: Model to route to when on_sensitive_data='route' sticky_session_routing: When True, all subsequent requests in the session use the same model + run_in_parallel: When True, this pre_call or post_call guardrail runs concurrently with + other opted-in guardrails of the same hook. Only safe for block-only guardrails that + do not mutate the request or response. """ self.guardrail_name = guardrail_name self.supported_event_hooks = supported_event_hooks @@ -145,6 +159,8 @@ class CustomGuardrail(CustomLogger): self.on_sensitive_data: Optional[str] = on_sensitive_data self.sensitive_data_route_to_model: Optional[str] = sensitive_data_route_to_model self.sticky_session_routing: bool = sticky_session_routing + self.run_in_parallel: bool = run_in_parallel + self.only_scan_new_messages: bool = only_scan_new_messages if supported_event_hooks: ## validate event_hook is in supported_event_hooks @@ -269,6 +285,100 @@ class CustomGuardrail(CustomLogger): """Extract session_id from request data.""" return get_session_id_from_request_data(request_data) + @staticmethod + def _scanned_text_hash(text: str) -> str: + """Stable content hash for a single scannable text segment. + + Hashing the exact text the provider would receive means an edited earlier + segment produces a different hash and gets re-scanned, while an unchanged + segment repeated on a later turn is skipped. + """ + return hashlib.sha256(text.encode("utf-8")).hexdigest() + + def _scanned_texts_cache_key(self, session_id: str) -> str: + return f"guardrail_scanned_texts:{self.guardrail_name}:{session_id}" + + async def filter_new_texts_for_session( + self, + texts: list[str] | None, + request_data: dict[str, object], + cache: DualCache, + ) -> list[str] | None: + """Return only the text segments not already scanned earlier in this session. + + Returns ``None`` when incremental scanning is inactive (feature off, no + session id, masking enabled, or the cache read failed). ``None`` signals + the caller to fall back to a full scan; a returned list (possibly empty) + signals the caller to scan only that subset and skip masking write-back. + """ + if not self.only_scan_new_messages or not texts: + return None + + if self.mask_request_content or self.mask_response_content: + verbose_logger.warning( + "Guardrail %s: only_scan_new_messages is not supported with masking; scanning full context.", + self.guardrail_name, + ) + return None + + session_id = get_session_id_from_request_data(request_data) + if not session_id: + verbose_logger.debug( + "Guardrail %s: only_scan_new_messages enabled but request has no session id; scanning full context.", + self.guardrail_name, + ) + return None + + try: + cached: object = await cache.async_get_cache(key=self._scanned_texts_cache_key(session_id)) + except Exception as e: # noqa: BLE001 # cache is best-effort; any failure must fall back to a full scan + verbose_logger.warning( + "Guardrail %s: failed to read scanned-message cache (%s); scanning full context.", + self.guardrail_name, + e, + ) + return None + + seen: set[str] = {str(h) for h in cached} if isinstance(cached, list) else set() + return [text for text in texts if self._scanned_text_hash(text) not in seen] + + async def mark_texts_scanned( + self, + texts: list[str] | None, + request_data: dict[str, object], + cache: DualCache, + ) -> None: + """Record the hashes of all text segments present on a successful (non-blocked) scan. + + Called only after the guardrail allows the request, so a blocked segment is + never marked scanned and will be re-checked if the client retries. + """ + if not self.only_scan_new_messages or not texts: + return + if self.mask_request_content or self.mask_response_content: + return + session_id = get_session_id_from_request_data(request_data) + if not session_id: + return + + cache_key = self._scanned_texts_cache_key(session_id) + current_hashes = [self._scanned_text_hash(text) for text in texts] + try: + existing: object = await cache.async_get_cache(key=cache_key) + existing_hashes: list[str] = [str(h) for h in existing] if isinstance(existing, list) else [] + merged: list[str] = list(dict.fromkeys(existing_hashes + current_hashes)) + await cache.async_set_cache( + key=cache_key, + value=merged, + ttl=GUARDRAIL_SCANNED_MESSAGES_CACHE_TTL_SECONDS, + ) + except Exception as e: # noqa: BLE001 # cache is best-effort; any failure must not block the request + verbose_logger.warning( + "Guardrail %s: failed to persist scanned-message cache (%s); next call will re-scan.", + self.guardrail_name, + e, + ) + def should_route_on_sensitive_data(self) -> bool: """ Returns True if this guardrail is configured to route requests @@ -856,6 +966,8 @@ class CustomGuardrail(CustomLogger): request_data["metadata"] = {} _append_guardrail_info(request_data["metadata"]) + _guardrail_self_recorded.set(True) + # Emit the otel guardrail span here, where every guardrail execution lands, # rather than relying on a post-call hook that does not fire on every path # (e.g. a pass-through request that passes its guardrails). @@ -1138,8 +1250,12 @@ def log_guardrail_information(func): (structured detections, tracing detail) than this decorator's "allow"/"mask"/raw-response default. To avoid double-recording in that case (which would emit two spans, two Datadog records, two spend-log - entries, etc.), snapshot the entry count before invocation: if the - wrapped function already appended its own entry, skip the auto-record. + entries, etc.), a context-local flag records whether the wrapped function + appended its own entry; if so, the auto-record is skipped. The flag is a + ``ContextVar`` rather than a count of entries in the shared ``request_data`` + so it stays correct when guardrails run concurrently (asyncio copies the + context into each gathered task): counting shared entries would let one + guardrail's append hide another guardrail's missing record. """ import functools import inspect @@ -1159,16 +1275,6 @@ def log_guardrail_information(func): return GuardrailEventHooks.post_call return None - def _count_recorded_guardrail_entries(request_data: dict) -> int: - total = 0 - for container_key in ("metadata", "litellm_metadata"): - container = request_data.get(container_key) - if isinstance(container, dict): - entries = container.get("standard_logging_guardrail_information") - if isinstance(entries, list): - total += len(entries) - return total - @functools.wraps(func) async def async_wrapper(*args, **kwargs): start_time = datetime.now() # Move start_time inside the wrapper @@ -1182,10 +1288,10 @@ def log_guardrail_information(func): original_inputs = kwargs.get("inputs") logging_obj = kwargs.get("logging_obj") - entries_before = _count_recorded_guardrail_entries(request_data) + self_recorded_token = _guardrail_self_recorded.set(False) try: response = await func(*args, **kwargs) - if _count_recorded_guardrail_entries(request_data) > entries_before: + if _guardrail_self_recorded.get(): return response return self._process_response( response=response, @@ -1197,7 +1303,7 @@ def log_guardrail_information(func): original_inputs=original_inputs, ) except Exception as e: - if _count_recorded_guardrail_entries(request_data) > entries_before: + if _guardrail_self_recorded.get(): raise return self._process_error( e=e, @@ -1208,6 +1314,7 @@ def log_guardrail_information(func): event_type=event_type, ) finally: + _guardrail_self_recorded.reset(self_recorded_token) _sync_guardrail_info_to_logging_obj(request_data, logging_obj) @functools.wraps(func) @@ -1223,10 +1330,10 @@ def log_guardrail_information(func): original_inputs = kwargs.get("inputs") logging_obj = kwargs.get("logging_obj") - entries_before = _count_recorded_guardrail_entries(request_data) + self_recorded_token = _guardrail_self_recorded.set(False) try: response = func(*args, **kwargs) - if _count_recorded_guardrail_entries(request_data) > entries_before: + if _guardrail_self_recorded.get(): return response return self._process_response( response=response, @@ -1236,7 +1343,7 @@ def log_guardrail_information(func): original_inputs=original_inputs, ) except Exception as e: - if _count_recorded_guardrail_entries(request_data) > entries_before: + if _guardrail_self_recorded.get(): raise return self._process_error( e=e, @@ -1245,6 +1352,7 @@ def log_guardrail_information(func): event_type=event_type, ) finally: + _guardrail_self_recorded.reset(self_recorded_token) _sync_guardrail_info_to_logging_obj(request_data, logging_obj) @functools.wraps(func) diff --git a/litellm/litellm_core_utils/duration_parser.py b/litellm/litellm_core_utils/duration_parser.py index 79036367652..b78a314dc45 100644 --- a/litellm/litellm_core_utils/duration_parser.py +++ b/litellm/litellm_core_utils/duration_parser.py @@ -9,9 +9,22 @@ duration_in_seconds is used in diff parts of the code base, example import re import time as time_module from datetime import datetime, time, timedelta, timezone, tzinfo -from typing import Optional, Tuple +from typing import Final, Optional, Tuple from zoneinfo import ZoneInfo +from litellm._logging import verbose_logger + +_BUDGET_DURATION_WORD_ALIASES: Final[dict[str, str]] = { + "hourly": "1h", + "daily": "24h", + "weekly": "7d", + "monthly": "30d", +} + + +def _normalize_duration(duration: str) -> str: + return _BUDGET_DURATION_WORD_ALIASES.get(duration.strip().lower(), duration) + def _extract_from_regex(duration: str) -> Tuple[int, str]: match = re.match(r"(\d+)(mo|[smhdw]?)", duration) @@ -48,7 +61,7 @@ def duration_in_seconds(duration: str) -> int: Returns time in seconds till when budget needs to be reset """ - value, unit = _extract_from_regex(duration=duration) + value, unit = _extract_from_regex(duration=_normalize_duration(duration)) if unit == "s": return value @@ -124,9 +137,13 @@ def get_next_standardized_reset_time( current_time, _ = _setup_timezone(current_time, timezone_str) # Parse duration - value, unit = _parse_duration(duration) + value, unit = _parse_duration(_normalize_duration(duration)) if value is None: - # Fall back to default if format is invalid + verbose_logger.warning( + "Unrecognized budget_duration %r; falling back to a next-midnight reset. " + "Use the format (e.g. '1h', '7d', '30d', '1mo').", + duration, + ) return current_time.replace(hour=0, minute=0, second=0, microsecond=0) + timedelta(days=1) # Midnight of the current day in the specified timezone diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 4462b01fe5c..e3c8c6e598f 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -1612,6 +1612,35 @@ class Logging(LiteLLMLoggingBaseClass): **kwargs, ) + async def dispatch_failure_handlers( + self, + exception: Exception, + traceback_exception: str, + prefer_async_handlers: bool = False, + ) -> None: + """Route failure logging to async and/or sync handlers for this request. + + Mirrors ``dispatch_success_handlers``: the sync ``failure_handler`` never runs + concurrently with ``async_failure_handler`` on the shared logging object, so the + two paths cannot mutate it at the same time. ``prefer_async_handlers`` only + bypasses the sync-SDK-only shortcut (e.g. ``async for`` on a stream from + ``completion()``); legacy string callbacks still run via + ``executor.submit(failure_handler)`` when configured. + """ + litellm_params = self.model_call_details.get("litellm_params", {}) or {} + sync_sdk = self._is_sync_litellm_request(litellm_params) + passthrough = self.call_type == CallTypes.pass_through.value + if sync_sdk and not prefer_async_handlers and not passthrough: + self.failure_handler(exception, traceback_exception) + return + + await self.async_failure_handler(exception, traceback_exception) + + if not self._should_run_sync_failure_callbacks_for_async_calls(): + return + + executor.submit(self.failure_handler, exception, traceback_exception) + def should_run_logging( self, event_type: Literal["async_success", "sync_success", "async_failure", "sync_failure"], @@ -3076,6 +3105,24 @@ class Logging(LiteLLMLoggingBaseClass): _filtered_success_callbacks = self._remove_internal_litellm_callbacks(_filtered_success_callbacks) return len(_filtered_success_callbacks) > 0 + def _should_run_sync_failure_callbacks_for_async_calls(self) -> bool: + """ + Returns: + - bool: True if sync failure callbacks should be run for async calls. eg. `langfuse`, `s3` + + Mirrors ``_should_run_sync_callbacks_for_async_calls`` but reads the failure + callback lists. Gating the legacy sync ``failure_handler`` on the success lists + would drop sync failure callbacks for any caller that configures only failure + callbacks, so streaming errors would be logged nowhere. + """ + _combined_sync_callbacks = self.get_combined_callback_list( + dynamic_success_callbacks=self.dynamic_failure_callbacks, + global_callbacks=litellm.failure_callback, + ) + _filtered_failure_callbacks = self._remove_internal_custom_logger_callbacks(_combined_sync_callbacks) + _filtered_failure_callbacks = self._remove_internal_litellm_callbacks(_filtered_failure_callbacks) + return len(_filtered_failure_callbacks) > 0 + def get_combined_callback_list(self, dynamic_success_callbacks: Optional[List], global_callbacks: List) -> List: if dynamic_success_callbacks is None: return list(global_callbacks) diff --git a/litellm/litellm_core_utils/llm_cost_calc/utils.py b/litellm/litellm_core_utils/llm_cost_calc/utils.py index 33bf546c239..85ed0665ebf 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/utils.py +++ b/litellm/litellm_core_utils/llm_cost_calc/utils.py @@ -457,7 +457,8 @@ def _parse_prompt_tokens_details(usage: Usage) -> PromptTokensDetailsResult: cache_creation_tokens = ( cast( Optional[int], - getattr(usage.prompt_tokens_details, "cache_creation_tokens", 0), + getattr(usage.prompt_tokens_details, "cache_write_tokens", 0) + or getattr(usage.prompt_tokens_details, "cache_creation_tokens", 0), ) or 0 ) @@ -906,10 +907,6 @@ def get_token_type_cost_breakdown( cache_read_tokens = prompt_tokens_details["cache_hit_tokens"] cache_creation_tokens = prompt_tokens_details["cache_creation_tokens"] cache_creation_token_details = prompt_tokens_details["cache_creation_token_details"] - # Some OpenAI-compatible providers (e.g. kimi-k2) report cache-write tokens - # under `cache_write_tokens`; mirror the total-cost normalization path. - if not cache_creation_tokens: - cache_creation_tokens = _coerce_token_count(getattr(usage.prompt_tokens_details, "cache_write_tokens", 0)) # Fall back to the private top-level counters the Usage constructor mirrors cache # tokens onto, so providers/callers that bypass prompt_tokens_details are covered. if not cache_read_tokens: diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index f518cbaadea..60dbf7c644a 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -2008,12 +2008,9 @@ class CustomStreamWrapper: if self.logging_obj is not None: self._record_partial_usage_for_failure() ## LOGGING - threading.Thread( - target=self.logging_obj.failure_handler, - args=(e, traceback_exception), - ).start() # log response - # Handle any exceptions that might occur during streaming - asyncio.create_task(self.logging_obj.async_failure_handler(e, traceback_exception)) + asyncio.create_task( + self.logging_obj.dispatch_failure_handlers(e, traceback_exception, prefer_async_handlers=True) + ) self._handle_stream_fallback_error(e) except (httpx.ReadError, httpx.RemoteProtocolError) as e: if self.received_finish_reason is None: @@ -2122,13 +2119,8 @@ class CustomStreamWrapper: if self.logging_obj is not None: self._record_partial_usage_for_failure() ## LOGGING - threading.Thread( - target=self.logging_obj.failure_handler, - args=(e, traceback_exception), - ).start() # log response - # Handle any exceptions that might occur during streaming asyncio.create_task( - self.logging_obj.async_failure_handler(e, traceback_exception) # type: ignore + self.logging_obj.dispatch_failure_handlers(e, traceback_exception, prefer_async_handlers=True) ) self._handle_stream_fallback_error(e) diff --git a/litellm/litellm_core_utils/url_utils.py b/litellm/litellm_core_utils/url_utils.py index 1cbb1ce973f..a83cb3bc69e 100644 --- a/litellm/litellm_core_utils/url_utils.py +++ b/litellm/litellm_core_utils/url_utils.py @@ -148,6 +148,15 @@ def _parse_url_destination_allowlist_entry( return _normalize_host(parsed.hostname), scheme, port +def provider_url_destination_candidates(value: str) -> Tuple[str, ...]: + return tuple( + candidate + for part in value.split(",") + for candidate in (part.strip(), part.strip().split("/", 1)[1] if "/" in part.strip() else "") + if candidate + ) + + def is_url_destination_allowed_by_host(url: str, allowed_hosts: List[str]) -> bool: """Return True when a credential-bearing provider URL is admin-allowlisted. diff --git a/litellm/llms/anthropic/chat/transformation.py b/litellm/llms/anthropic/chat/transformation.py index 5a0f274e3ca..e99f356f8f2 100644 --- a/litellm/llms/anthropic/chat/transformation.py +++ b/litellm/llms/anthropic/chat/transformation.py @@ -480,10 +480,21 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): """ Filter out unsupported fields from JSON schema for Anthropic's output_format API. - Anthropic's output_format doesn't support certain JSON schema properties: - - maxItems/minItems: Not supported for array types - - minimum/maximum: Not supported for numeric types - - minLength/maxLength: Not supported for string types + Anthropic's output_format doesn't support certain JSON schema properties. + These are constraints that cannot be enforced by the constrained-decoding + grammar Anthropic compiles the schema into, so the API rejects them with a + 400 ``invalid_request_error`` (e.g. "output_format.schema: For 'array' type, + property 'uniqueItems' is not supported"): + - maxItems/minItems/uniqueItems/contains/minContains/maxContains/prefixItems: array constraints + - minimum/maximum/exclusiveMinimum/exclusiveMaximum/multipleOf: numeric constraints + - minLength/maxLength: string constraints + - minProperties/maxProperties/patternProperties/propertyNames: object constraints + - dependentRequired/dependentSchemas/unevaluatedProperties: object constraints + - if/then/else/not: conditional and negation keywords + + ``oneOf`` is also rejected ("Schema type 'oneOf' is not supported") and is + rewritten to ``anyOf``, matching the Anthropic SDK. Unknown keywords are + ignored by the API, so anything not listed here passes through untouched. This mirrors the transformation done by the Anthropic Python SDK. See: https://platform.claude.com/docs/en/build-with-claude/structured-outputs#how-sdk-transformation-works @@ -504,33 +515,53 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): if not isinstance(schema, dict): return schema - # All numeric/string/array constraints not supported by Anthropic - unsupported_fields = { - "maxItems", - "minItems", # array constraints - "minimum", - "maximum", # numeric constraints - "exclusiveMinimum", - "exclusiveMaximum", # numeric constraints - "minLength", - "maxLength", # string constraints - } - - # Build description additions from removed constraints - constraint_descriptions: list = [] constraint_labels = { "minItems": "minimum number of items: {}", "maxItems": "maximum number of items: {}", + "uniqueItems": "all array items must be unique", + "contains": "array must contain an item matching: {}", + "minContains": "minimum number of matching items: {}", + "maxContains": "maximum number of matching items: {}", + "prefixItems": "leading items must match, in order: {}", "minimum": "minimum value: {}", "maximum": "maximum value: {}", "exclusiveMinimum": "exclusive minimum value: {}", "exclusiveMaximum": "exclusive maximum value: {}", + "multipleOf": "must be a multiple of {}", "minLength": "minimum length: {}", "maxLength": "maximum length: {}", + "minProperties": "minimum number of properties: {}", + "maxProperties": "maximum number of properties: {}", + "patternProperties": "properties whose names match each pattern must satisfy: {}", + "propertyNames": "property names must satisfy: {}", + "dependentRequired": "dependent required properties: {}", + "dependentSchemas": "dependent schemas: {}", + "unevaluatedProperties": "unevaluated properties must satisfy: {}", + "if": "conditional (if): {}", + "then": "conditional (then): {}", + "else": "conditional (else): {}", + "not": "must not match: {}", } - for field in unsupported_fields: - if field in schema: - constraint_descriptions.append(constraint_labels[field].format(schema[field])) + unsupported_fields = set(constraint_labels) + + # Build description additions from removed constraints. Iterating + # constraint_labels (not the set) keeps the note order deterministic across + # processes, so identical requests serialize identically regardless of + # PYTHONHASHSEED and stay cache-friendly. + constraint_descriptions: list = [] + for field, label in constraint_labels.items(): + if field not in schema: + continue + value = schema[field] + # A falsy boolean constraint (e.g. ``uniqueItems: false``) imposes no + # real requirement, so don't add a misleading advisory note for it. + if isinstance(value, bool) and not value: + continue + # Sub-schema constraints (e.g. ``contains``) are serialized as JSON so + # the advisory note preserves what the constraint actually required, + # instead of just noting that it existed. + note_value = json.dumps(value) if isinstance(value, (dict, list)) else value + constraint_descriptions.append(label.format(note_value)) result: Dict[str, Any] = {} @@ -557,11 +588,17 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): elif key == "$defs" and isinstance(value, dict): result[key] = {k: AnthropicConfig.filter_anthropic_output_schema(v) for k, v in value.items()} elif key == "anyOf" and isinstance(value, list): - result[key] = [AnthropicConfig.filter_anthropic_output_schema(item) for item in value] + result["anyOf"] = result.get("anyOf", []) + [ + AnthropicConfig.filter_anthropic_output_schema(item) for item in value + ] elif key == "allOf" and isinstance(value, list): result[key] = [AnthropicConfig.filter_anthropic_output_schema(item) for item in value] elif key == "oneOf" and isinstance(value, list): - result[key] = [AnthropicConfig.filter_anthropic_output_schema(item) for item in value] + # Anthropic rejects oneOf ("Schema type 'oneOf' is not supported"); + # the Anthropic SDK rewrites it to anyOf, so do the same. + result["anyOf"] = result.get("anyOf", []) + [ + AnthropicConfig.filter_anthropic_output_schema(item) for item in value + ] else: result[key] = value diff --git a/litellm/llms/bedrock/chat/converse_transformation.py b/litellm/llms/bedrock/chat/converse_transformation.py index c38b3593465..8ce2b982955 100644 --- a/litellm/llms/bedrock/chat/converse_transformation.py +++ b/litellm/llms/bedrock/chat/converse_transformation.py @@ -895,9 +895,8 @@ class AmazonConverseConfig(BaseConfig): if _tool_choice_value is not None: optional_params["tool_choice"] = _tool_choice_value if param == "parallel_tool_calls": - disable_parallel = not value optional_params["_parallel_tool_use_config"] = { - "tool_choice": {"disable_parallel_tool_use": disable_parallel} + "tool_choice": {"type": "auto", "disable_parallel_tool_use": not value} } if param == "thinking": if ( @@ -1208,6 +1207,22 @@ class AmazonConverseConfig(BaseConfig): return {} + @staticmethod + def _merge_parallel_tool_use_config(additional_request_params: dict, parallel_tool_use_config: dict) -> dict: + merged_entries = { + key: ( + { + **value, + **additional_request_params[key], + **{k: v for k, v in value.items() if k != "type"}, + } + if isinstance(additional_request_params.get(key), dict) and isinstance(value, dict) + else value + ) + for key, value in parallel_tool_use_config.items() + } + return {**additional_request_params, **merged_entries} + def _prepare_request_params( self, optional_params: dict, model: str, drop_params: bool = False ) -> Tuple[dict, dict, dict, Optional[OutputConfigBlock]]: @@ -1276,15 +1291,9 @@ class AmazonConverseConfig(BaseConfig): # Handle parallel_tool_calls configuration parallel_tool_use_config = additional_request_params.pop("_parallel_tool_use_config", None) if parallel_tool_use_config is not None and bedrock_converse_supports_parallel_tool_use_config(model): - for key, value in parallel_tool_use_config.items(): - if ( - key in additional_request_params - and isinstance(additional_request_params[key], dict) - and isinstance(value, dict) - ): - additional_request_params[key].update(value) - else: - additional_request_params[key] = value + additional_request_params = self._merge_parallel_tool_use_config( + additional_request_params, parallel_tool_use_config + ) additional_request_params.pop("parallel_tool_calls", None) diff --git a/litellm/llms/bedrock/realtime/handler.py b/litellm/llms/bedrock/realtime/handler.py index b48c37791c4..b7237d288ec 100644 --- a/litellm/llms/bedrock/realtime/handler.py +++ b/litellm/llms/bedrock/realtime/handler.py @@ -9,6 +9,8 @@ import contextlib import json from typing import Any, Optional +from pydantic import TypeAdapter + from litellm._logging import _redact_string, verbose_proxy_logger from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging @@ -16,6 +18,8 @@ from ..base_aws_llm import BaseAWSLLM from ..common_utils import BedrockError from .transformation import BedrockRealtimeConfig +_CLIENT_MODALITIES_ADAPTER: TypeAdapter["list[str] | None"] = TypeAdapter(list[str] | None) + class BedrockRealtime(BaseAWSLLM): """Handler for Bedrock Nova Sonic realtime speech-to-speech API.""" @@ -124,6 +128,9 @@ class BedrockRealtime(BaseAWSLLM): verbose_proxy_logger.debug("Bedrock Realtime: Bidirectional stream established") + await websocket.send_text(json.dumps(transformation_config.session_created_event(model, logging_obj))) + verbose_proxy_logger.debug("Bedrock Realtime: sent session.created to client on connect") + # Track state for transformation session_state = { "current_output_item_id": None, @@ -143,6 +150,7 @@ class BedrockRealtime(BaseAWSLLM): transformation_config, model, session_state, + logging_obj, ) ) @@ -179,6 +187,7 @@ class BedrockRealtime(BaseAWSLLM): transformation_config: BedrockRealtimeConfig, model: str, session_state: dict, + logging_obj: LiteLLMLogging | None = None, ): """Forward messages from client WebSocket to Bedrock stream.""" from aws_sdk_bedrock_runtime.models import ( @@ -210,6 +219,23 @@ class BedrockRealtime(BaseAWSLLM): for bedrock_message in transformed_messages: await send_to_bedrock(bedrock_message) + if logging_obj is not None: + client_message_type: str | None = None + requested_modalities: list[str] | None = None + with contextlib.suppress(Exception): + parsed_client_message = json.loads(message) + client_message_type = parsed_client_message.get("type") + if client_message_type == "session.update": + requested_modalities = _CLIENT_MODALITIES_ADAPTER.validate_python( + parsed_client_message.get("session", {}).get("modalities") + ) + if client_message_type == "session.update": + await client_ws.send_text( + json.dumps( + transformation_config.session_updated_event(model, logging_obj, requested_modalities) + ) + ) + except Exception as e: verbose_proxy_logger.debug(f"Client to Bedrock forwarding ended: {e}", exc_info=True) for close_message in transformation_config.session_close_messages(): diff --git a/litellm/llms/bedrock/realtime/transformation.py b/litellm/llms/bedrock/realtime/transformation.py index fe5f0584e03..24a40ebea1b 100644 --- a/litellm/llms/bedrock/realtime/transformation.py +++ b/litellm/llms/bedrock/realtime/transformation.py @@ -623,35 +623,42 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): verbose_logger.warning(f"Unknown message type: {message_type}") return [] - def transform_session_start_event( + def _session_object( self, - event: dict, model: str, logging_obj: LiteLLMLoggingObj, - ) -> OpenAIRealtimeStreamSessionEvents: - """ - Transform Bedrock sessionStart event to OpenAI session.created. - - Args: - event: Bedrock sessionStart event - model: Model ID - logging_obj: Logging object - - Returns: - OpenAI session.created event - """ - verbose_logger.debug("Handling sessionStart") - + modalities: list[str] | None = None, + ) -> OpenAIRealtimeStreamSession: session = OpenAIRealtimeStreamSession( id=logging_obj.litellm_trace_id, - modalities=["text", "audio"], + modalities=modalities if modalities is not None else ["text", "audio"], ) if model is not None and isinstance(model, str): session["model"] = model + return session + def session_created_event( + self, + model: str, + logging_obj: LiteLLMLoggingObj, + ) -> OpenAIRealtimeStreamSessionEvents: + """Build the OpenAI session.created event for this realtime session.""" return OpenAIRealtimeStreamSessionEvents( type="session.created", - session=session, + session=self._session_object(model, logging_obj), + event_id=str(uuid.uuid4()), + ) + + def session_updated_event( + self, + model: str, + logging_obj: LiteLLMLoggingObj, + modalities: list[str] | None = None, + ) -> OpenAIRealtimeStreamSessionEvents: + """Build the OpenAI session.updated ack reflecting the client's requested modalities.""" + return OpenAIRealtimeStreamSessionEvents( + type="session.updated", + session=self._session_object(model, logging_obj, modalities), event_id=str(uuid.uuid4()), ) @@ -1169,8 +1176,6 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): # Route to appropriate transformation method if "sessionStart" in event: - session_created = self.transform_session_start_event(event, model, logging_obj) - returned_messages.append(session_created) session_configuration_request = json.dumps({"configured": True}) elif "contentStart" in event: diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index c48d75439a7..ec1301e5923 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -2107,8 +2107,7 @@ class BaseLLMHTTPHandler: rust_messages_response = await self._maybe_rust_anthropic_messages( custom_llm_provider=custom_llm_provider, litellm_params=litellm_params, - stream=stream or False, - rust_stream_eligible=bool(stream) and not self._has_agentic_completion_hook(logging_obj), + has_agentic_hook=self._has_agentic_completion_hook(logging_obj), model=model, api_key=api_key, api_base=api_base, @@ -2266,8 +2265,7 @@ class BaseLLMHTTPHandler: *, custom_llm_provider: str, litellm_params: GenericLiteLLMParams, - stream: bool, - rust_stream_eligible: bool, + has_agentic_hook: bool, model: str, api_key: str | None, api_base: str | None, @@ -2279,7 +2277,7 @@ class BaseLLMHTTPHandler: return None if litellm_params.get("rust") is not True and not BaseLLMHTTPHandler._rust_env_enabled(): return None - if stream and not rust_stream_eligible: + if has_agentic_hook: return None from litellm.rust_bridge import messages as rust_messages_bridge diff --git a/litellm/llms/huggingface/embedding/handler.py b/litellm/llms/huggingface/embedding/handler.py index 39eb430db74..f72a79e084d 100644 --- a/litellm/llms/huggingface/embedding/handler.py +++ b/litellm/llms/huggingface/embedding/handler.py @@ -322,7 +322,7 @@ class HuggingFaceEmbedding(BaseLLM): task = get_hf_task_embedding_for_model(model=model, task_type=task_type, api_base=HF_HUB_URL) # print_verbose(f"{model}, {task}") embed_url = "" - if "https" in model: + if model.startswith(("http://", "https://")): embed_url = model elif api_base: embed_url = api_base diff --git a/litellm/llms/huggingface/embedding/transformation.py b/litellm/llms/huggingface/embedding/transformation.py index 13e38ab5560..6f27e3115eb 100644 --- a/litellm/llms/huggingface/embedding/transformation.py +++ b/litellm/llms/huggingface/embedding/transformation.py @@ -316,25 +316,6 @@ class HuggingFaceEmbeddingConfig(BaseConfig): return data - def get_api_base(self, api_base: Optional[str], model: str) -> str: - """ - Get the API base for the Huggingface API. - - Do not add the chat/embedding/rerank extension here. Let the handler do this. - """ - if "https" in model: - completion_url = model - elif api_base is not None: - completion_url = api_base - elif "HF_API_BASE" in os.environ: - completion_url = os.getenv("HF_API_BASE", "") - elif "HUGGINGFACE_API_BASE" in os.environ: - completion_url = os.getenv("HUGGINGFACE_API_BASE", "") - else: - completion_url = f"https://api-inference.huggingface.co/models/{model}" - - return completion_url - def validate_environment( self, headers: Dict, diff --git a/litellm/llms/oobabooga/chat/oobabooga.py b/litellm/llms/oobabooga/chat/oobabooga.py index fe2bb9dc6d1..40d88e8e125 100644 --- a/litellm/llms/oobabooga/chat/oobabooga.py +++ b/litellm/llms/oobabooga/chat/oobabooga.py @@ -34,7 +34,7 @@ def completion( optional_params=optional_params, litellm_params=litellm_params, ) - if "https" in model: + if model.startswith(("http://", "https://")): completion_url = model elif api_base: completion_url = api_base @@ -96,7 +96,7 @@ def embedding( encoding=None, ): # Create completion URL - if "https" in model: + if model.startswith(("http://", "https://")): embeddings_url = model elif api_base: embeddings_url = f"{api_base}/v1/embeddings" diff --git a/litellm/llms/sagemaker/chat/transformation.py b/litellm/llms/sagemaker/chat/transformation.py index 4e4e088f491..4447d63e5a1 100644 --- a/litellm/llms/sagemaker/chat/transformation.py +++ b/litellm/llms/sagemaker/chat/transformation.py @@ -143,7 +143,7 @@ class SagemakerChatConfig(OpenAIGPTConfig, BaseAWSLLM): raise SagemakerError(status_code=response.status_code, message=response.text) custom_stream_decoder = AWSEventStreamDecoder(model="", is_messages_api=True) - completion_stream = custom_stream_decoder.iter_bytes(response.iter_bytes(chunk_size=1024)) + completion_stream = custom_stream_decoder.iter_bytes(response.iter_bytes()) streaming_response = CustomStreamWrapper( completion_stream=completion_stream, @@ -189,7 +189,7 @@ class SagemakerChatConfig(OpenAIGPTConfig, BaseAWSLLM): raise SagemakerError(status_code=response.status_code, message=response.text) custom_stream_decoder = AWSEventStreamDecoder(model="", is_messages_api=True) - completion_stream = custom_stream_decoder.aiter_bytes(response.aiter_bytes(chunk_size=1024)) + completion_stream = custom_stream_decoder.aiter_bytes(response.aiter_bytes()) streaming_response = CustomStreamWrapper( completion_stream=completion_stream, diff --git a/litellm/llms/sagemaker/completion/handler.py b/litellm/llms/sagemaker/completion/handler.py index 4b87271fd44..c27a3c3528c 100644 --- a/litellm/llms/sagemaker/completion/handler.py +++ b/litellm/llms/sagemaker/completion/handler.py @@ -200,23 +200,12 @@ class SagemakerLLM(BaseAWSLLM): # Add model_id as InferenceComponentName header # boto3 doc: https://docs.aws.amazon.com/sagemaker/latest/APIReference/API_runtime_InvokeEndpoint.html prepared_request.headers.update({"X-Amzn-SageMaker-Inference-Component": model_id}) - sync_handler = _get_httpx_client() - sync_response = sync_handler.post( - url=prepared_request.url, + completion_stream = self.make_sync_call( + api_base=prepared_request.url, headers=prepared_request.headers, # type: ignore - data=prepared_request.body, - stream=stream, + data=cast(str, prepared_request.body), # cast-ok: signed body is a JSON str, mirrors async path + logging_obj=logging_obj, ) - - if sync_response.status_code != 200: - raise SagemakerError( - status_code=sync_response.status_code, - message=str(sync_response.read()), - ) - - decoder = AWSEventStreamDecoder(model="") - - completion_stream = decoder.iter_bytes(sync_response.iter_bytes(chunk_size=1024)) streaming_response = CustomStreamWrapper( completion_stream=completion_stream, model=model, @@ -334,6 +323,29 @@ class SagemakerLLM(BaseAWSLLM): litellm_params=litellm_params, ) + def make_sync_call( + self, + api_base: str, + headers: dict, + data: str, + logging_obj, + client=None, + ): + if client is None: + client = _get_httpx_client() + sync_response = client.post( + api_base, + headers=headers, + data=data, + stream=True, + ) + + if sync_response.status_code != 200: + raise SagemakerError(status_code=sync_response.status_code, message=str(sync_response.read())) + + decoder = AWSEventStreamDecoder(model="") + return decoder.iter_bytes(sync_response.iter_bytes()) + async def make_async_call( self, api_base: str, @@ -358,7 +370,7 @@ class SagemakerLLM(BaseAWSLLM): raise SagemakerError(status_code=response.status_code, message=response.text) decoder = AWSEventStreamDecoder(model="") - completion_stream = decoder.aiter_bytes(response.aiter_bytes(chunk_size=1024)) + completion_stream = decoder.aiter_bytes(response.aiter_bytes()) return completion_stream diff --git a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py index 624190a0b61..8dd0dc19b81 100644 --- a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py +++ b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py @@ -3333,27 +3333,37 @@ class ModelResponseIterator: return self.chunk_parser(chunk=json_chunk) - def handle_accumulated_json_chunk(self, chunk: str) -> Optional["ModelResponseStream"]: - chunk = litellm.CustomStreamWrapper._strip_sse_data_from_chunk(chunk) or "" - message = chunk.replace("\n\n", "") + def handle_accumulated_json_chunk(self, chunk: str, is_final: bool = False) -> Optional["ModelResponseStream"]: + message = litellm.CustomStreamWrapper._strip_sse_data_from_chunk(chunk) or "" + self.accumulated_json = (self.accumulated_json + message.replace("\n\n", "")).strip() - self.accumulated_json += message - - # json.loads on the whole buffer after every fragment is O(n^2) and - # holds the GIL, freezing the event loop for seconds on large responses - # (https://github.com/BerriAI/litellm/issues/26181). A complete Gemini - # chunk is a JSON object/array, so only attempt the parse once the - # buffer's last non-whitespace byte can close one. - stripped = self.accumulated_json.rstrip() - if not stripped or stripped[-1] not in "}]": + # Mid-stream, defer parsing until the buffer's last byte can close a value: + # attempting a parse after every fragment of one large object is O(n^2) and + # holds the GIL, freezing the event loop. At end of stream (is_final) no more + # data is coming, so drain whatever complete values remain regardless of the + # trailing byte, otherwise a complete leading value sitting behind a truncated + # trailing one would be silently dropped. + if not is_final and (not self.accumulated_json or self.accumulated_json[-1] not in "}]"): return None - try: - _data = json.loads(self.accumulated_json) - self.accumulated_json = "" # reset after successful parsing - return self.chunk_parser(chunk=_data) - except json.JSONDecodeError: - return None + # Peel one complete JSON value from the front of the buffer and keep the + # unconsumed tail. Running json.loads over the whole buffer would fail + # forever once it held more than one concatenated value ("Extra data") while + # never resetting the buffer, so the buffer grew without bound and pinned the + # core. raw_decode reports where the value ended, so concatenated values drain + # one call at a time. A leading non-dict value (never emitted by Gemini in + # practice) is consumed and skipped so it cannot block the dict values behind it. + decoder = json.JSONDecoder() + while self.accumulated_json: + try: + raw_value = decoder.raw_decode(self.accumulated_json) + except json.JSONDecodeError: + return None + decoded, end_index = cast("tuple[object, int]", raw_value) # cast-ok: raw_decode -> tuple[Any,int] + self.accumulated_json = self.accumulated_json[end_index:].strip() + if isinstance(decoded, dict): + return self.chunk_parser(chunk=decoded) + return None def _common_chunk_parsing_logic(self, chunk: str) -> Optional["ModelResponseStream"]: try: @@ -3378,7 +3388,9 @@ class ModelResponseIterator: chunk = self.response_iterator.__next__() except StopIteration: if self.chunk_type == "accumulated_json" and self.accumulated_json: - return self.handle_accumulated_json_chunk(chunk="") + result = self.handle_accumulated_json_chunk(chunk="", is_final=True) + if result is not None: + return result raise StopIteration except ValueError as e: raise RuntimeError(f"Error receiving chunk from stream: {e}") @@ -3400,7 +3412,9 @@ class ModelResponseIterator: chunk = await self.async_response_iterator.__anext__() except StopAsyncIteration: if self.chunk_type == "accumulated_json" and self.accumulated_json: - return self.handle_accumulated_json_chunk(chunk="") + result = self.handle_accumulated_json_chunk(chunk="", is_final=True) + if result is not None: + return result raise StopAsyncIteration except ValueError as e: raise RuntimeError(f"Error receiving chunk from stream: {e}") diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index d3917886060..d43eda39b1f 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -1502,6 +1502,222 @@ "supports_parallel_tool_use_config": true, "prompt_cache_min_tokens": 1024 }, + "anthropic.claude-opus-5": { + "bedrock_converse_supports_strict_tools": false, + "supports_adaptive_thinking": true, + "supports_mid_conversation_system": true, + "cache_creation_input_token_cost": 6.25e-06, + "cache_creation_input_token_cost_above_1hr": 1e-05, + "cache_read_input_token_cost": 5e-07, + "input_cost_per_token": 5e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 2.5e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_native_structured_output": true, + "supports_max_reasoning_effort": true, + "supports_output_config": true, + "supports_parallel_tool_use_config": true, + "prompt_cache_min_tokens": 512 + }, + "global.anthropic.claude-opus-5": { + "bedrock_converse_supports_strict_tools": false, + "supports_adaptive_thinking": true, + "supports_mid_conversation_system": true, + "cache_creation_input_token_cost": 6.25e-06, + "cache_creation_input_token_cost_above_1hr": 1e-05, + "cache_read_input_token_cost": 5e-07, + "input_cost_per_token": 5e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 2.5e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_native_structured_output": true, + "supports_max_reasoning_effort": true, + "supports_output_config": true, + "supports_parallel_tool_use_config": true, + "prompt_cache_min_tokens": 512 + }, + "us.anthropic.claude-opus-5": { + "bedrock_converse_supports_strict_tools": false, + "supports_adaptive_thinking": true, + "supports_mid_conversation_system": true, + "cache_creation_input_token_cost": 6.875e-06, + "cache_creation_input_token_cost_above_1hr": 1.1e-05, + "cache_read_input_token_cost": 5.5e-07, + "input_cost_per_token": 5.5e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 2.75e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_native_structured_output": true, + "supports_max_reasoning_effort": true, + "supports_output_config": true, + "supports_parallel_tool_use_config": true, + "prompt_cache_min_tokens": 512 + }, + "eu.anthropic.claude-opus-5": { + "bedrock_converse_supports_strict_tools": false, + "supports_adaptive_thinking": true, + "supports_mid_conversation_system": true, + "cache_creation_input_token_cost": 6.875e-06, + "cache_creation_input_token_cost_above_1hr": 1.1e-05, + "cache_read_input_token_cost": 5.5e-07, + "input_cost_per_token": 5.5e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 2.75e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_native_structured_output": true, + "supports_max_reasoning_effort": true, + "supports_output_config": true, + "supports_parallel_tool_use_config": true, + "prompt_cache_min_tokens": 512 + }, + "au.anthropic.claude-opus-5": { + "bedrock_converse_supports_strict_tools": false, + "supports_adaptive_thinking": true, + "supports_mid_conversation_system": true, + "cache_creation_input_token_cost": 6.875e-06, + "cache_creation_input_token_cost_above_1hr": 1.1e-05, + "cache_read_input_token_cost": 5.5e-07, + "input_cost_per_token": 5.5e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 2.75e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_native_structured_output": true, + "supports_max_reasoning_effort": true, + "supports_output_config": true, + "supports_parallel_tool_use_config": true, + "prompt_cache_min_tokens": 512 + }, + "jp.anthropic.claude-opus-5": { + "bedrock_converse_supports_strict_tools": false, + "supports_adaptive_thinking": true, + "supports_mid_conversation_system": true, + "cache_creation_input_token_cost": 6.875e-06, + "cache_creation_input_token_cost_above_1hr": 1.1e-05, + "cache_read_input_token_cost": 5.5e-07, + "input_cost_per_token": 5.5e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 2.75e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_native_structured_output": true, + "supports_max_reasoning_effort": true, + "supports_output_config": true, + "supports_parallel_tool_use_config": true, + "prompt_cache_min_tokens": 512 + }, "anthropic.claude-opus-4-8": { "bedrock_converse_supports_strict_tools": false, "supports_adaptive_thinking": true, @@ -2756,6 +2972,38 @@ "supports_xhigh_reasoning_effort": true, "supports_max_reasoning_effort": true }, + "azure_ai/claude-opus-5": { + "supports_mid_conversation_system": true, + "supports_adaptive_thinking": true, + "input_cost_per_token": 5e-06, + "output_cost_per_token": 2.5e-05, + "litellm_provider": "azure_ai", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "cache_creation_input_token_cost": 6.25e-06, + "cache_creation_input_token_cost_above_1hr": 1e-05, + "cache_read_input_token_cost": 5e-07, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_max_reasoning_effort": true, + "prompt_cache_min_tokens": 512 + }, "azure_ai/claude-opus-4-8": { "supports_mid_conversation_system": true, "supports_adaptive_thinking": true, @@ -11846,6 +12094,44 @@ "supports_output_config": true, "prompt_cache_min_tokens": 512 }, + "claude-opus-5": { + "cache_creation_input_token_cost": 6.25e-06, + "cache_creation_input_token_cost_above_1hr": 1e-05, + "cache_read_input_token_cost": 5e-07, + "input_cost_per_token": 5e-06, + "litellm_provider": "anthropic", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 2.5e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_native_structured_output": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_max_reasoning_effort": true, + "provider_specific_entry": { + "us": 1.1, + "fast": 2.0 + }, + "supports_output_config": true, + "supports_speed": true, + "prompt_cache_min_tokens": 512 + }, "claude-opus-4-8": { "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, @@ -36896,6 +37182,70 @@ "supports_xhigh_reasoning_effort": true, "supports_max_reasoning_effort": true }, + "vertex_ai/claude-opus-5": { + "supports_mid_conversation_system": true, + "supports_adaptive_thinking": true, + "cache_creation_input_token_cost": 6.25e-06, + "cache_creation_input_token_cost_above_1hr": 1e-05, + "cache_read_input_token_cost": 5e-07, + "input_cost_per_token": 5e-06, + "litellm_provider": "vertex_ai-anthropic_models", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 2.5e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_max_reasoning_effort": true, + "prompt_cache_min_tokens": 512 + }, + "vertex_ai/claude-opus-5@default": { + "supports_mid_conversation_system": true, + "supports_adaptive_thinking": true, + "cache_creation_input_token_cost": 6.25e-06, + "cache_creation_input_token_cost_above_1hr": 1e-05, + "cache_read_input_token_cost": 5e-07, + "input_cost_per_token": 5e-06, + "litellm_provider": "vertex_ai-anthropic_models", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 2.5e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_max_reasoning_effort": true, + "prompt_cache_min_tokens": 512 + }, "vertex_ai/claude-opus-4-8": { "supports_mid_conversation_system": true, "supports_adaptive_thinking": true, diff --git a/litellm/proxy/_experimental/mcp_server/auth/token_exchange.py b/litellm/proxy/_experimental/mcp_server/auth/token_exchange.py deleted file mode 100644 index cd41dd648ee..00000000000 --- a/litellm/proxy/_experimental/mcp_server/auth/token_exchange.py +++ /dev/null @@ -1,192 +0,0 @@ -""" -OAuth 2.0 Token Exchange (RFC 8693) handler for MCP servers. - -Exchanges a user's incoming JWT (subject_token) for a scoped access token -at an IDP's token exchange endpoint. The exchanged token is then used to -authenticate requests to the upstream MCP server. - -See: https://datatracker.ietf.org/doc/html/rfc8693 -""" - -import asyncio -import hashlib -import weakref -from typing import TYPE_CHECKING, Dict, Tuple - -import httpx - -from litellm._logging import verbose_logger -from litellm.caching.in_memory_cache import InMemoryCache -from litellm.constants import ( - MCP_OAUTH2_TOKEN_CACHE_DEFAULT_TTL, - MCP_OAUTH2_TOKEN_CACHE_MIN_TTL, - MCP_OAUTH2_TOKEN_EXPIRY_BUFFER_SECONDS, - MCP_TOKEN_EXCHANGE_CACHE_MAX_SIZE, -) -from litellm.llms.custom_httpx.http_handler import get_async_httpx_client -from litellm.proxy._experimental.mcp_server.auth.token_endpoint_auth import ( - build_token_endpoint_client_auth, -) -from litellm.types.llms.custom_http import httpxSpecialProvider -from litellm.types.mcp import DEFAULT_SUBJECT_TOKEN_TYPE - -if TYPE_CHECKING: - from litellm.types.mcp_server.mcp_server_manager import MCPServer - -# RFC 8693 grant type constant -TOKEN_EXCHANGE_GRANT_TYPE = "urn:ietf:params:oauth:grant-type:token-exchange" - - -class TokenExchangeHandler: - """Handles OAuth 2.0 Token Exchange (RFC 8693) for MCP servers. - - Caches exchanged tokens keyed by ``hash(subject_token + server_id)`` so - repeated calls with the same user token skip the IDP round-trip. - """ - - def __init__(self) -> None: - self._cache = InMemoryCache( - max_size_in_memory=MCP_TOKEN_EXCHANGE_CACHE_MAX_SIZE, - default_ttl=MCP_OAUTH2_TOKEN_CACHE_DEFAULT_TTL, - ) - # WeakValueDictionary so locks are GC'd once no coroutine holds a reference, - # preventing unbounded growth with many rotating user tokens. - self._locks: weakref.WeakValueDictionary[str, asyncio.Lock] = weakref.WeakValueDictionary() - - def _get_lock(self, cache_key: str) -> asyncio.Lock: - lock = self._locks.get(cache_key) - if lock is None: - lock = asyncio.Lock() - self._locks[cache_key] = lock - return lock - - @staticmethod - def _cache_key(subject_token: str, server_id: str) -> str: - raw = f"{subject_token}:{server_id}" - return hashlib.sha256(raw.encode()).hexdigest() - - async def exchange_token( - self, - subject_token: str, - server: "MCPServer", - ) -> str: - """Exchange *subject_token* for a scoped access token. - - Returns the exchanged ``access_token`` string (suitable for a - ``Bearer`` header). - - Raises ``ValueError`` on configuration or IDP errors. - """ - cache_key = self._cache_key(subject_token, server.server_id) - - # Fast path - cached = self._cache.get_cache(cache_key) - if cached is not None: - return cached - - # Slow path — one exchange at a time per (user, server) pair - async with self._get_lock(cache_key): - cached = self._cache.get_cache(cache_key) - if cached is not None: - return cached - - token, ttl = await self._do_exchange(subject_token, server) - self._cache.set_cache(cache_key, token, ttl=ttl) - return token - - async def _do_exchange( - self, - subject_token: str, - server: "MCPServer", - ) -> Tuple[str, int]: - """POST to the token exchange endpoint with RFC 8693 parameters. - - Returns ``(access_token, ttl_seconds)``. - """ - endpoint = server.token_exchange_endpoint or server.token_url - if not endpoint: - raise ValueError( - f"MCP server '{server.server_id}' has auth_type=oauth2_token_exchange " - f"but no token_exchange_endpoint or token_url configured" - ) - if not server.client_id or not server.client_secret: - raise ValueError( - f"MCP server '{server.server_id}' has auth_type=oauth2_token_exchange " - f"but missing client_id or client_secret" - ) - - client_auth = build_token_endpoint_client_auth( - auth_method=server.token_endpoint_auth_method, - client_id=server.client_id, - client_secret=server.client_secret, - ) - data: Dict[str, str] = { - "grant_type": TOKEN_EXCHANGE_GRANT_TYPE, - "subject_token": subject_token, - "subject_token_type": server.subject_token_type or DEFAULT_SUBJECT_TOKEN_TYPE, - **client_auth.body, - } - if server.audience: - data["audience"] = server.audience - if server.scopes: - data["scope"] = " ".join(server.scopes) - - verbose_logger.debug( - "Exchanging token for MCP server %s at %s (audience=%s)", - server.server_id, - endpoint, - server.audience, - ) - - client = get_async_httpx_client(llm_provider=httpxSpecialProvider.MCP) - post_kwargs = {"data": data, **({"headers": client_auth.headers} if client_auth.headers else {})} - try: - response = await client.post(endpoint, **post_kwargs) - response.raise_for_status() - except httpx.HTTPStatusError as exc: - verbose_logger.debug( - "Token exchange IDP error for MCP server %s (status %d)", - server.server_id, - exc.response.status_code, - ) - raise ValueError( - f"Token exchange for MCP server '{server.server_id}' failed with status {exc.response.status_code}" - ) from exc - - body = response.json() - if not isinstance(body, dict): - raise ValueError( - f"Token exchange response for MCP server '{server.server_id}' " - f"returned non-object JSON (got {type(body).__name__})" - ) - - access_token = body.get("access_token") - if not access_token: - raise ValueError(f"Token exchange response for MCP server '{server.server_id}' missing 'access_token'") - - raw_expires_in = body.get("expires_in") - try: - expires_in = int(raw_expires_in) if raw_expires_in is not None else MCP_OAUTH2_TOKEN_CACHE_DEFAULT_TTL - except (TypeError, ValueError): - expires_in = MCP_OAUTH2_TOKEN_CACHE_DEFAULT_TTL - - ttl = max( - expires_in - MCP_OAUTH2_TOKEN_EXPIRY_BUFFER_SECONDS, - MCP_OAUTH2_TOKEN_CACHE_MIN_TTL, - ) - - verbose_logger.info( - "Token exchange succeeded for MCP server %s (expires in %ds)", - server.server_id, - expires_in, - ) - return access_token, ttl - - def invalidate(self, subject_token: str, server_id: str) -> None: - """Remove a cached exchanged token (e.g. after a 401).""" - cache_key = self._cache_key(subject_token, server_id) - self._cache.delete_cache(cache_key) - - -# Module-level singleton -mcp_token_exchange_handler = TokenExchangeHandler() diff --git a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py index f1fcc95c532..a27d6b92843 100644 --- a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py +++ b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py @@ -25,6 +25,9 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.bridge_credenti from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import ( EnvelopeIdentity, ) +from litellm.proxy._experimental.mcp_server.outbound_credentials.session_credentials import ( + is_session_bearer_shaped, +) from litellm.proxy._types import ( UI_TEAM_ID, LiteLLM_TeamTable, @@ -124,6 +127,27 @@ def _has_client_supplied_mcp_auth( return bool(mcp_auth_header) or bool(mcp_server_auth_headers) +def _is_mcp_admitted_user_subject(user_api_key_auth: UserAPIKeyAuth | None) -> bool: + """True when this auth is a keyless subject admitted by the gateway session / bridge user + path, as opposed to a JWT or other keyless auth that merely lacks a ``team_id``. + + Reads the server-only ``mcp_admitted_user_subject`` field, set only by ``_reload_admitted_user``. It + is deliberately NOT a ``metadata`` key, which is caller-controlled at key creation and so forgeable + on a personal key to gain the team grant union or dodge the egress scrub; this field cannot be.""" + return user_api_key_auth is not None and user_api_key_auth.mcp_admitted_user_subject is True + + +def _is_aggregate_mcp_scope(route: str, mcp_servers: list[str] | None) -> bool: + """True when a request targets the aggregate ``/mcp`` endpoint rather than any named + server. Named targets arrive either through ``x-mcp-servers`` (``mcp_servers``) or a + path segment (``/mcp/{server}`` / ``/{server}/mcp``); the aggregate scope has neither. + The gateway-DCR session arm and challenge fire only here, so a per-server flow is never + affected.""" + if mcp_servers: + return False + return len(MCPRequestHandler._extract_target_server_names_from_path(route)) == 0 + + def _is_aggregate_gateway_dcr_challenge_scope( route: str, mcp_servers: list[str] | None, @@ -141,11 +165,9 @@ def _is_aggregate_gateway_dcr_challenge_scope( client. Fails closed to the original admission error otherwise.""" if not _is_litellm_auth_admission_error(exc): return False - if mcp_servers: - return False if _has_client_supplied_mcp_auth(mcp_auth_header, mcp_server_auth_headers): return False - return len(MCPRequestHandler._extract_target_server_names_from_path(route)) == 0 + return _is_aggregate_mcp_scope(route, mcp_servers) def _aggregate_gateway_dcr_challenge(request: Request, invalid_token: bool) -> HTTPException: @@ -362,6 +384,19 @@ class MCPRequestHandler: request=request, route=request_route, ) + elif ( + _is_aggregate_mcp_scope(request_route, mcp_servers) + and oauth2_headers + and is_session_bearer_shaped(oauth2_headers["Authorization"]) + ): + # A gateway DCR session bearer at the aggregate /mcp scope: open the identity-only session + # token and admit under the live litellm user. One that does not open fails closed with the + # aggregate invalid_token challenge; a non-session bearer falls through to the oauth2 arm. + validated_user_api_key_auth = await MCPRequestHandler._admit_gateway_session( + authorization_value=oauth2_headers["Authorization"], + request=request, + route=request_route, + ) elif oauth2_headers: # Authorization on a non-delegated server: the bearer must be a real # LiteLLM credential, so a failed validation is a genuine 401/403 and @@ -392,15 +427,80 @@ class MCPRequestHandler: bearer_presented=False, ) + # Leak-defense (single chokepoint): a gateway admission credential (session bearer or bridge + # envelope) is NEVER a valid upstream token. Scrub it from EVERY egress context so no + # client-forwarded, OBO, or passthrough path can send it upstream for replay. Anchored to the + # credential SHAPE, so a legitimate upstream/passthrough token is forwarded unchanged. + raw_headers = dict(headers) + ( + oauth2_headers, + raw_headers, + mcp_auth_header, + mcp_server_auth_headers, + ) = MCPRequestHandler._scrub_gateway_admission_credentials( + admitted=_is_mcp_admitted_user_subject(validated_user_api_key_auth), + oauth2_headers=oauth2_headers, + raw_headers=raw_headers, + mcp_auth_header=mcp_auth_header, + mcp_server_auth_headers=mcp_server_auth_headers, + ) + return ( validated_user_api_key_auth, mcp_auth_header, mcp_servers, mcp_server_auth_headers, oauth2_headers, - dict(headers), + raw_headers, ) + @staticmethod + def _is_gateway_admission_credential(value: str | None) -> bool: + """True when a header value is a gateway admission credential — a session bearer or bridge + envelope. It proves who signed in to the GATEWAY, never a valid UPSTREAM token, so it must never + be forwarded (a hostile upstream could capture and replay it against the aggregate ``/mcp`` scope).""" + return value is not None and (is_session_bearer_shaped(value) or is_bridge_envelope_shaped(value)) + + @staticmethod + def _scrub_gateway_admission_credentials( + admitted: bool, + oauth2_headers: dict[str, str] | None, + raw_headers: dict[str, str], + mcp_auth_header: str | None, + mcp_server_auth_headers: dict[str, dict[str, str]] | None, + ) -> tuple[dict[str, str] | None, dict[str, str], str | None, dict[str, dict[str, str]] | None]: + """Remove any gateway admission credential from EVERY egress header context, keyed on the credential + SHAPE: top-level ``Authorization`` (oauth2 + raw), the deprecated ``x-mcp-auth``, and per-server + ``x-mcp-{alias}-authorization``. A legitimate upstream/passthrough token is never gateway-shaped so + it survives (including the real upstream token the bridge arm injects per-server); an admitted + subject's top-level Authorization is dropped unconditionally as defense-in-depth.""" + cred = MCPRequestHandler._is_gateway_admission_credential + + # 1. Top-level Authorization → oauth2_headers. + authz = oauth2_headers.get("Authorization") if oauth2_headers else None + if admitted or cred(authz): + oauth2_headers = None + + # 2. raw_headers: drop the admitted subject's Authorization, and ANY header whose value is a + # gateway credential (covers x-mcp-auth and x-mcp-{alias}-authorization in their raw form). + raw_headers = { + k: v for k, v in raw_headers.items() if not ((admitted and k.lower() == "authorization") or cred(v)) + } + + # 3. Deprecated x-mcp-auth value. + if cred(mcp_auth_header): + mcp_auth_header = None + + # 4. Per-server x-mcp-{alias}-authorization values (drop the value, then any now-empty server dict). + if mcp_server_auth_headers: + stripped = { + alias: {h: val for h, val in hdrs.items() if not cred(val)} + for alias, hdrs in mcp_server_auth_headers.items() + } + mcp_server_auth_headers = {alias: hdrs for alias, hdrs in stripped.items() if hdrs} + + return oauth2_headers, raw_headers, mcp_auth_header, mcp_server_auth_headers + @staticmethod def _extract_target_server_names_from_path(path: str) -> List[str]: """ @@ -626,6 +726,62 @@ class MCPRequestHandler: case _: assert_never(result) + @staticmethod + async def _admit_gateway_session( + authorization_value: str, + request: Request, + route: str, + ) -> UserAPIKeyAuth: + """Open a gateway DCR session bearer and admit the live litellm user it references. + + Identity-only sibling of :meth:`_admit_dcr_bridge_delegate`: the session token seals no + upstream credential (those are vaulted per user, resolved at egress), so authorization is + resolved fresh via :meth:`_reload_admitted_user` + the centralized policy gate rather than a + mint-time snapshot. Pre-DB gates (size, IP, route allowlist) run first, mirroring the standard + pipeline. Fails closed with the aggregate ``invalid_token`` challenge on an expired, tampered, + foreign, or refresh token, or a missing/deactivated/policy-rejected user.""" + from litellm.proxy._experimental.mcp_server.outbound_credentials.session_credentials import ( + NotSessionBearer, + SessionBearerAdmitted, + SessionBearerInvalid, + resolve_session_bearer, + session_keys_from_master_key, + ) + from litellm.proxy.proxy_server import master_key + + if not master_key: + raise HTTPException(status_code=500, detail="Server misconfigured: master_key is not set") + + await MCPRequestHandler._run_pre_db_read_auth_checks(request=request, route=route) + + keys = session_keys_from_master_key(master_key) + result = resolve_session_bearer(authorization_value, keys, datetime.now(timezone.utc)) + match result: + case SessionBearerAdmitted(): + try: + admitted = await MCPRequestHandler._reload_admitted_user(result.principal.user_id) + await MCPRequestHandler._enforce_admitted_live_policy( + admitted=admitted, request=request, route=route + ) + except HTTPException as exc: + # A cryptographically valid bearer whose referenced user is now missing or + # SCIM-deactivated is an invalid_token at the aggregate scope: relay the RFC 9728 + # challenge so the DCR client re-authorizes, matching the SessionBearerInvalid + # arm, instead of a bare 401 with no WWW-Authenticate. A 503 (DB outage) is a + # transient availability failure, not an auth failure, so it passes through. + if exc.status_code == 401: + raise _aggregate_gateway_dcr_challenge(request, invalid_token=True) from exc + raise + return admitted + case SessionBearerInvalid(): + raise _aggregate_gateway_dcr_challenge(request, invalid_token=True) + case NotSessionBearer(): + # Unreachable: the arm is entered only for an is_session_bearer_shaped + # value. Kept for match exhaustiveness and fails closed regardless. + raise _aggregate_gateway_dcr_challenge(request, invalid_token=True) + case _: + assert_never(result) + @staticmethod async def _run_pre_db_read_auth_checks(request: Request, route: str) -> None: """Run the proxy-wide gates ``user_api_key_auth`` applies before any key lookup: the @@ -666,30 +822,18 @@ class MCPRequestHandler: @staticmethod async def _reload_admitted_user(user_id: str) -> UserAPIKeyAuth: - """Reload the live user an interactively-minted envelope references and admit them as - themselves. + """Reload the live user an interactively-minted envelope references and admit them as themselves. - The DCR client authenticates via SSO at the bridged authorize, which yields a user - subject rather than a virtual key, so the envelope admits under the user's own - identity: the reloaded ``user_id`` and the user's own MCP object permission ride on the - returned ``UserAPIKeyAuth``, and the SAME ``get_allowed_mcp_servers`` the key path uses then - computes which servers the user may reach, so the user's litellm MCP grants and access groups - gate the request exactly as a key's do. Only the user's OWN object permission is bound: a - ``UserAPIKeyAuth`` carries a single ``team_id`` while a user may belong to many teams, so - team-inherited MCP grants for a user are a follow-up (they need a many-teams union - ``get_allowed_mcp_servers`` does not do off one auth object). The caller's centralized policy - gate enforces the user's live budget and org state, and a SCIM-deactivated owner fails closed. + The user's own object permission and ``org_id`` ride on the returned ``UserAPIKeyAuth``, and the + SAME ``get_allowed_mcp_servers`` the key path uses gates the request. The ``mcp_admitted_user_subject`` + marker (set below) makes that resolver union the servers the user reaches through ANY of their teams + on top of these direct grants, each source bounded by ITS OWN org, so a user spanning organizations + cannot leak one org's servers past another's ceiling. - Error handling mirrors the key path's retryable-503 contract, but ``get_user_object`` defeats a - type-based check: where ``get_key_object`` raises a typed ``ProxyException`` for a missing key - and lets a DB outage propagate raw, ``get_user_object`` catches every DB failure and re-raises a - bare ``ValueError``, so a missing user and a real outage look identical and the original error - survives only as ``__context__``. ``_raise_503_if_db_unavailable`` therefore walks the cause - chain: a transient DB outage still surfaces as a retryable 503, while a missing user, or any - other non-outage resolution failure, fails closed as a 401 rather than an opaque 500. The - object-permission load shares this one boundary, so an outage there is classified the same - way (``get_object_permission`` itself swallows a failed load to ``None``, matching how - ``get_key_object`` best-effort-loads a key's object permission).""" + Error handling: ``get_user_object`` catches every DB failure and re-raises a bare ``ValueError``, so a + missing user and a real outage look identical (the cause survives only as ``__context__``). + ``_raise_503_if_db_unavailable`` walks the cause chain so an outage stays a retryable 503 while any + other failure fails closed as 401, not an opaque 500; the object-permission load shares that boundary.""" from litellm.proxy.auth.auth_checks import get_object_permission, get_user_object from litellm.proxy.proxy_server import prisma_client, user_api_key_cache @@ -721,12 +865,75 @@ class MCPRequestHandler: raise HTTPException(status_code=401, detail="Invalid or expired credential") if isinstance(user_object.metadata, dict) and user_object.metadata.get("scim_active") is False: raise HTTPException(status_code=401, detail="Invalid or expired credential") - return UserAPIKeyAuth( + admitted = UserAPIKeyAuth( user_id=user_object.user_id, user_role=user_object.user_role, + org_id=user_object.organization_id, object_permission=object_permission, object_permission_id=user_object.object_permission_id, + # Copy the live user's rate limits, as the standard user-subject path does: the parallel + # limiter reads these off the auth object and treats None as unlimited, so a keyless subject + # with them unset would outrun its user RPM/TPM. (Per-team mcp_rpm_limit is stamped below; + # per-KEY limits do not apply, there being no key.) + user_tpm_limit=user_object.tpm_limit, + user_rpm_limit=user_object.rpm_limit, ) + # Server-only marker, set AFTER construction: the before-validator strips it from any validated + # input, so caller-supplied data (key metadata, JWT claims) can never forge it. + admitted.mcp_admitted_user_subject = True + # Carry each granting team's per-server mcp_rpm_limit: this subject reaches servers through + # several teams under its own identity, so without this a cross-team user outruns every team's + # limit. Resolved from the same roster-checked sources as the grant union, so a team throttles + # only what it granted. + admitted.mcp_source_team_rpm_limits = await MCPRequestHandler._admitted_subject_team_rpm_limits(admitted) + return admitted + + @staticmethod + async def _admitted_subject_team_rpm_limits(auth: UserAPIKeyAuth) -> dict[str, dict[str, int]] | None: + """``team_id -> mcp_rpm_limit`` for every team this subject reaches servers through, each map + filtered to the servers THAT team's grant actually reaches. + + A limit rides the same scope as the access it bounds, so a roster team is charged only for a + server its OWN grant reaches (never one the user reaches through a different team, which would + drain a bucket shared by that team's keys for access it never provided). Grant scope comes from + the SAME ``get_allowed_mcp_servers(source)`` authorization uses; limit-map keys are names/aliases + so each is resolved to an id via ``expand_permission_list`` before the membership check. Returns + None (no descriptors) when nothing applies; a lookup failure narrows to None rather than raising, + since rate limiting must not deny a request authorization already allowed.""" + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + + try: + limits: dict[str, dict[str, int]] = {} + source_grants = await MCPRequestHandler.admitted_source_grants(auth) + for source, granted_ids in source_grants: + if not source.team_id: + continue + team_obj = await MCPRequestHandler._roster_team_object(source.team_id, auth) + team_limit = (team_obj.metadata or {}).get("mcp_rpm_limit") if team_obj is not None else None + if not isinstance(team_limit, dict) or not team_limit: + continue + applicable: dict[str, int] = {} + for server_name, rpm in team_limit.items(): + for server_id in global_mcp_server_manager.expand_permission_list([server_name]): + if server_id not in granted_ids: + continue + # Charge ONLY the source billing attributes the call to (same owner), so one + # cross-team user cannot drain several teams' shared buckets on a single call, + # and a server the user's OWN grant reaches charges no team bucket. + attributed = await MCPRequestHandler.attributing_source_for_server( + auth, server_id, source_grants=source_grants + ) + if attributed is not None and attributed.team_id == source.team_id: + applicable[server_name] = rpm + break + if applicable: + limits[source.team_id] = applicable + return limits or None + except Exception as e: # noqa: BLE001 # throttling metadata must never fail an allowed request + verbose_logger.warning(f"Failed to resolve per-team MCP rpm limits for admitted subject: {str(e)}") + return None @staticmethod async def _reload_admitted_key(key_hash: str) -> UserAPIKeyAuth: @@ -1075,6 +1282,8 @@ class MCPRequestHandler: @staticmethod async def get_allowed_mcp_servers( user_api_key_auth: Optional[UserAPIKeyAuth] = None, + *, + keyless_source: bool = False, ) -> List[str]: """ Get list of allowed MCP servers for the given user/key based on permissions. @@ -1096,6 +1305,13 @@ class MCPRequestHandler: from litellm.proxy.proxy_server import general_settings try: + # A keyless admitted subject resolves per source BEFORE any single-source rule here. Ordering + # matters: the no_mcp_servers opt-out below reads the caller's own object_permission, so above + # this branch a user's own opt-out would wrongly zero their TEAMS' grants too (each source is + # independent; an opt-out silences only its own source, inside the recursive call). + if _is_mcp_admitted_user_subject(user_api_key_auth) and user_api_key_auth is not None: + return await MCPRequestHandler._resolve_admitted_subject_servers(user_api_key_auth) + # Get allowed servers from key and team allowed_mcp_servers_for_key = await MCPRequestHandler._get_allowed_mcp_servers_for_key(user_api_key_auth) @@ -1125,8 +1341,18 @@ class MCPRequestHandler: # team's by default. With require_key_mcp_access_defined the # team is a ceiling rather than a default, so the key must # grant servers explicitly (or via an access group) to reach - # any — it inherits none. - base = set() if general_settings.get("require_key_mcp_access_defined", False) else team_set + # any — it inherits none. That ceiling is for VIRTUAL KEYS that + # can declare their own access; a keyless gateway/bridge-admitted + # user has no key to declare access on — team membership IS their + # only access path — so the flag must not zero their team grants. + # A keyless admitted subject returned above and never reaches this virtual-key ceiling, + # so require_key_mcp_access_defined can only ever zero a real key's inherited team grants. + # ``keyless_source`` marks one grant source of an admitted subject, which has no key + # to declare access on, so the flag must not zero its team grants. + require_key_access = ( + general_settings.get("require_key_mcp_access_defined", False) and not keyless_source + ) + base = team_set if not require_key_access else set() else: base = key_set & team_set # both restrict → intersect @@ -1185,24 +1411,290 @@ class MCPRequestHandler: ######################################################### # Apply org-level ceiling if org_id is set ######################################################### - if user_api_key_auth and user_api_key_auth.org_id: - allowed_mcp_servers_for_org = await MCPRequestHandler._get_allowed_mcp_servers_for_org( - user_api_key_auth - ) - if len(allowed_mcp_servers_for_org) > 0: - if has_lower_level_mcp_restrictions: - # Lower-level restrictions exist, so org can only cap them. - allowed_mcp_servers = [s for s in allowed_mcp_servers if s in allowed_mcp_servers_for_org] - else: - # No lower-level restrictions → org list becomes the ceiling - allowed_mcp_servers = allowed_mcp_servers_for_org - verbose_logger.debug(f"Applied org ceiling filter. Final allowed servers: {allowed_mcp_servers}") + allowed_mcp_servers = await MCPRequestHandler._apply_primary_org_ceiling( + allowed_mcp_servers, + user_api_key_auth, + has_lower_level_mcp_restrictions, + keyless_source=keyless_source, + ) return list(set(allowed_mcp_servers)) except Exception as e: verbose_logger.warning(f"Failed to get allowed MCP servers: {str(e)}") return [] + @staticmethod + async def _apply_primary_org_ceiling( + allowed_mcp_servers: list[str], + user_api_key_auth: UserAPIKeyAuth | None, + has_lower_level_mcp_restrictions: bool, + keyless_source: bool = False, + ) -> list[str]: + """Cap the resolved server list by this caller's org ceiling: an explicit org list intersects + lower-level restrictions (else becomes the ceiling); no org or an empty list leaves it unchanged. + + ``keyless_source`` governs both divergences for a keyless admitted source. An UNRESOLVABLE ceiling + fails CLOSED for it (its only org bound is this ceiling, so dropping it on a fault would escalate a + cross-org user) while a key stays fail-open. And an org list may only ever INTERSECT a source (the + admitted model unions grants, so a ceiling must not become one), whereas for a key it may + substitute, that being the key ceiling model.""" + if not (user_api_key_auth and user_api_key_auth.org_id): + return allowed_mcp_servers + allowed_mcp_servers_for_org = await MCPRequestHandler._get_allowed_mcp_servers_for_org(user_api_key_auth) + if allowed_mcp_servers_for_org is None: + verbose_logger.warning( + f"MCP org ceiling unresolved for org_id={user_api_key_auth.org_id!r}; " + f"{'denying (keyless admitted subject)' if keyless_source else 'leaving uncapped (key auth)'}" + ) + return [] if keyless_source else allowed_mcp_servers + if len(allowed_mcp_servers_for_org) == 0: + return allowed_mcp_servers + if has_lower_level_mcp_restrictions or keyless_source: + # Org can only cap lower-level restrictions. A keyless admitted source ALWAYS takes this + # arm: its model unions GRANTS, so an org list may only narrow a source, never become one. + capped = [s for s in allowed_mcp_servers if s in allowed_mcp_servers_for_org] + else: + # No lower-level restrictions → org list becomes the ceiling. + capped = allowed_mcp_servers_for_org + verbose_logger.debug(f"Applied org ceiling filter. Final allowed servers: {capped}") + return capped + + @staticmethod + def _scoped_source_auth( + auth: UserAPIKeyAuth, + *, + team_id: str | None, + org_id: str | None, + carry_user_grants: bool, + ) -> UserAPIKeyAuth: + """A plain, UNMARKED auth describing ONE grant source of an admitted subject. + + Only the fields the resolver consults are carried; everything else is left at its default on + purpose: no ``api_key``/``token`` (not a key), no budget/spend/rate-limit (the subject's own + user-level limits meter the request, and per-source copies would double descriptors), no + ``user_role`` (an admin role would grant every server at the server-manager wrapper). The + admission marker cannot be set via the constructor (a before-validator pops it), so each source + resolves as an ordinary caller and cannot re-enter the admitted path.""" + scoped = UserAPIKeyAuth( + user_id=auth.user_id, + team_id=team_id, + org_id=org_id, + parent_otel_span=auth.parent_otel_span, + ) + if carry_user_grants: + # The user's OWN grants. A team source carries none of these (the resolver loads the team's + # own object_permission from team_id); mixing them in would widen the team with grants it never made. + scoped.object_permission = auth.object_permission + scoped.object_permission_id = auth.object_permission_id + scoped.access_group_ids = auth.access_group_ids + return scoped + + @staticmethod + async def _admitted_subject_sources(auth: UserAPIKeyAuth) -> list[UserAPIKeyAuth]: + """The independent sources a keyless admitted subject reaches MCP servers through: their own + direct grants, plus every team they are a live roster member of. + + Each team source carries that TEAM's org (falling back to the user's), so the canonical resolver + applies the team's OWN owning-org ceiling — a cross-org user's teams are each bounded by their + own org, not the caller's home org. Roster membership is checked HERE (not per resolution) + because a user's cached ``teams`` array can name a team whose ``members_with_roles`` no longer + lists them; the roster is the source of truth for revocation.""" + from litellm.proxy.proxy_server import prisma_client + + sources = [ + MCPRequestHandler._scoped_source_auth(auth, team_id=None, org_id=auth.org_id, carry_user_grants=True) + ] + if not auth.user_id or prisma_client is None: + return sources + for team_id in await MCPRequestHandler._resolve_user_team_ids(auth.user_id, auth): + team_obj = await MCPRequestHandler._roster_team_object(team_id, auth) + if team_obj is None: + continue + sources.append( + MCPRequestHandler._scoped_source_auth( + auth, + team_id=team_id, + org_id=team_obj.organization_id or auth.org_id, + carry_user_grants=False, + ) + ) + return sources + + @staticmethod + async def _roster_team_object(team_id: str, auth: UserAPIKeyAuth) -> LiteLLM_TeamTable | None: + """The team row for ``team_id``, but ONLY when ``auth``'s user is a live roster member of it. + + The single owner of "is this team really one of this subject's sources", so the grant union + and the per-team rate limits cannot disagree about which teams count. A team lingering in the + user's cached ``teams`` array whose ``members_with_roles`` no longer lists them returns None + here, which is what revokes both its grants and its throttle in one place.""" + from litellm.proxy.auth.auth_checks import get_team_object + from litellm.proxy.proxy_server import ( + prisma_client, + proxy_logging_obj, + user_api_key_cache, + ) + + if prisma_client is None or not auth.user_id: + return None + try: + team_obj: LiteLLM_TeamTable | None = await get_team_object( + team_id=team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=auth.parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + ) + except Exception as e: # noqa: BLE001 # per-source isolation: one team's blip must not deny the others + # Fault isolation is per SOURCE: an unresolvable team contributes nothing (fail closed for + # it alone, access only narrows) while every other source stands. Raising would collapse the + # whole union to deny-all over one momentarily-unreadable row. + verbose_logger.warning(f"MCP admitted-subject source team {team_id!r} unresolvable, skipping: {str(e)}") + return None + if team_obj is None: + return None + member_user_ids = {getattr(m, "user_id", None) for m in (team_obj.members_with_roles or [])} - {None} + if auth.user_id not in member_user_ids: + return None + # A team (or its owning org) over budget is not a live grantor, exactly as it is not for a key + # pinned to it. Enforced via the SAME owners the key path uses (_team_max_budget_check / + # _organization_max_budget_check), targeted at the TEAM's org through the scoped source view, so + # no consumer of the source list ever sees an over-budget team. This is ENFORCEMENT of an + # already-exceeded state; ATTRIBUTION of new spend stays with the user (documented deferral). + from litellm.exceptions import BudgetExceededError + from litellm.proxy.auth.auth_checks import ( + _organization_max_budget_check, + _team_max_budget_check, + ) + + source_view = MCPRequestHandler._scoped_source_auth( + auth, team_id=team_id, org_id=team_obj.organization_id or auth.org_id, carry_user_grants=False + ) + try: + await _team_max_budget_check( + team_object=team_obj, valid_token=source_view, proxy_logging_obj=proxy_logging_obj + ) + await _organization_max_budget_check( + valid_token=source_view, + team_object=team_obj, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + except BudgetExceededError as e: + verbose_logger.info(f"MCP admitted-subject source team {team_id!r} over budget, not a grantor: {str(e)}") + return None + except Exception as e: # noqa: BLE001 # per-source isolation: a budget-check fault narrows, never raises + verbose_logger.warning(f"MCP budget check failed for source team {team_id!r}, skipping source: {str(e)}") + return None + return team_obj + + @staticmethod + async def admitted_source_grants(auth: UserAPIKeyAuth) -> list[tuple[UserAPIKeyAuth, set[str]]]: + """``(source, the servers that source grants)`` for every source of an admitted subject. + + THE owner of "which source reaches which server". The reachable union, the per-team throttle + scope, the tool union and billing attribution are all just different reads of this one + answer — computing it separately per consumer is how they drift (a throttle map scoped by + roster instead of by grant charged unrelated teams' buckets).""" + return [ + (source, set(await MCPRequestHandler.get_allowed_mcp_servers(source, keyless_source=True))) + for source in await MCPRequestHandler._admitted_subject_sources(auth) + ] + + @staticmethod + async def _resolve_admitted_subject_servers(auth: UserAPIKeyAuth) -> list[str]: + """Union of what each of the admitted subject's sources reaches, each answered by the + canonical resolver so no rule is reimplemented for this caller shape.""" + reachable: set[str] = set() + for _source, granted in await MCPRequestHandler.admitted_source_grants(auth): + reachable.update(granted) + return list(reachable) + + @staticmethod + async def billing_auth_for_tool_call(auth: UserAPIKeyAuth, tool_name: str) -> UserAPIKeyAuth: + """The auth object a tool call's SPEND should be recorded against. + + ``auth`` unchanged for any non-admitted caller (key/JWT billing byte-identical). For an admitted + subject whose call is reached through a team's grant, a copy carrying that team's ``team_id`` and + owning ``org_id`` so the team's budget accumulates and the right org is charged. Falls back to + user-level attribution (rather than guessing a team) when the tool name does not resolve to a + server, reusing the manager's own tool-name lookup.""" + if not _is_mcp_admitted_user_subject(auth): + return auth + try: + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + + server = global_mcp_server_manager._get_mcp_server_from_tool_name(tool_name) + if server is None: + return auth + source = await MCPRequestHandler.attributing_source_for_server(auth, server.server_id) + if source is None or not source.team_id: + return auth + billed = auth.model_copy() + billed.team_id = source.team_id + billed.org_id = source.org_id + return billed + except Exception as e: # noqa: BLE001 # attribution must never fail an authorized call + verbose_logger.warning(f"MCP billing attribution failed for {tool_name!r}, billing the user: {str(e)}") + return auth + + @staticmethod + async def attributing_source_for_server( + auth: UserAPIKeyAuth, + server_id: str, + source_grants: list[tuple[UserAPIKeyAuth, set[str]]] | None = None, + ) -> UserAPIKeyAuth | None: + """The source a billable call to ``server_id`` is attributed to, or None to bill the caller as + themselves (their own grant reaches it, or nothing does). + + The rule: a user's OWN grant is not "through a team", so it bills the user; otherwise the call + bills a granting team, deterministically the lowest ``team_id`` when several grant the server so + the pick is stable rather than dict-ordering-dependent. Reads the one grant owner, so the billed + team is always one that actually granted the server (restoring the team budget accrual and + owning-org charge that a keyless, team_id-less subject otherwise skipped).""" + source_grants = source_grants or await MCPRequestHandler.admitted_source_grants(auth) + granting = [(source, granted) for source, granted in source_grants if server_id in granted] + if not granting: + return None + for source, _granted in granting: + if source.team_id is None: + return None # the user's own grant reaches it: their spend, their org + return min((source for source, _ in granting), key=lambda s: s.team_id or "") + + @staticmethod + async def _resolve_admitted_subject_tools(server_id: str, auth: UserAPIKeyAuth) -> list[str] | None: + """Effective tool allowlist on ``server_id`` for an admitted subject, as the union over the + sources that actually grant that server. + + A source that does not grant the server contributes nothing, so its tool rules cannot leak + onto a server reached through a different source. A source that grants the server with no + tool restriction means the user can use every tool on it, so allow-all wins the union. When + no source grants the server the result is ``[]`` — deny all, fail closed.""" + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + + # An OPEN channel (allow_all_keys, the user's own BYOM) makes the server REACHABLE through the + # user, though no grant source names it — without this the union returns [], listable but + # uninvokable. Reachability is ALL it confers, NOT a ceiling waiver: the user's own + # mcp_tool_permissions and org tool ceiling still bind, exactly as a key's do on an allow_all server. + reachable_via_open_channel = server_id in await global_mcp_server_manager.operator_open_server_ids(auth) + + allowed: set[str] = set() + for source, granted in await MCPRequestHandler.admitted_source_grants(auth): + # The open channel is evaluated against the user's OWN source (team_id is None), so that + # source's restrictions apply to it; a team's rules never ride an open-channel server. + if server_id not in granted and not (reachable_via_open_channel and source.team_id is None): + continue + tools = await MCPRequestHandler.get_allowed_tools_for_server(server_id, source, keyless_source=True) + if tools is None: + return None + allowed.update(tools) + return sorted(allowed) + @staticmethod def _get_key_object_permission( user_api_key_auth: Optional[UserAPIKeyAuth] = None, @@ -1262,6 +1754,8 @@ class MCPRequestHandler: async def get_allowed_tools_for_server( server_id: str, user_api_key_auth: Optional[UserAPIKeyAuth] = None, + *, + keyless_source: bool = False, ) -> Optional[List[str]]: """ Get list of allowed tool names for a specific server based on key/team permissions. @@ -1278,6 +1772,12 @@ class MCPRequestHandler: return None try: + # FIRST statement, mirroring get_allowed_mcp_servers: a keyless admitted subject resolves per + # source and shares nothing with the single-credential prelude below. Ordering is the invariant: + # sat after the prelude, a fault in a lookup the subject never uses denied tools its teams grant. + if _is_mcp_admitted_user_subject(user_api_key_auth): + return await MCPRequestHandler._resolve_admitted_subject_tools(server_id, user_api_key_auth) + # Get key and team object permissions (already loaded in main auth flow) key_obj_perm = MCPRequestHandler._get_key_object_permission(user_api_key_auth) team_obj_perm = await MCPRequestHandler._get_team_object_permission(user_api_key_auth) @@ -1331,42 +1831,73 @@ class MCPRequestHandler: # No team restrictions → use key restrictions allowed_tools = cast(List[str], key_tools) - # Intersect with agent's tool permissions if agent_id is set - if user_api_key_auth.agent_id: - # Pre-fetch agent object_permission once to avoid duplicate DB query - agent_obj_perm = await MCPRequestHandler._get_agent_object_permission(user_api_key_auth) - agent_tools = await MCPRequestHandler._get_agent_tool_permissions_for_server( - server_id=server_id, - user_api_key_auth=user_api_key_auth, - agent_object_permission=agent_obj_perm, - ) - if agent_tools is not None: - if allowed_tools is not None: - allowed_tools = list(set(allowed_tools) & set(agent_tools)) - else: - allowed_tools = agent_tools - - # Apply org-level tool ceiling if org_id is set - if user_api_key_auth.org_id: - # _get_org_object_permission uses user_api_key_cache, so this is not a - # fresh DB round-trip when get_allowed_mcp_servers was already called. - org_obj_perm = await MCPRequestHandler._get_org_object_permission(user_api_key_auth) - org_tools = ( - global_mcp_server_manager.expand_tool_permissions(org_obj_perm.mcp_tool_permissions).get(server_id) - if org_obj_perm and org_obj_perm.mcp_tool_permissions - else None - ) - if org_tools is not None: - if allowed_tools is not None: - allowed_tools = list(set(allowed_tools) & set(org_tools)) - else: - allowed_tools = list(org_tools) - - return allowed_tools + return await MCPRequestHandler._apply_agent_and_org_tool_ceilings( + allowed_tools, server_id, user_api_key_auth, keyless_source=keyless_source + ) except Exception as e: verbose_logger.warning(f"Failed to get allowed tools for server: {str(e)}") - return None + # Fail CLOSED for a keyless admitted subject: ANY error must deny the server's tools ([]), + # not collapse to allow-all (None); key/JWT auth keeps its prior allow-all-on-error. Both + # keyless_source AND the marker are needed: each source resolves through an UNMARKED auth, so + # without keyless_source a fault under a source returns None and wins the union as allow-all. + return [] if (keyless_source or _is_mcp_admitted_user_subject(user_api_key_auth)) else None + + @staticmethod + async def _apply_agent_and_org_tool_ceilings( + allowed_tools: list[str] | None, + server_id: str, + user_api_key_auth: UserAPIKeyAuth, + keyless_source: bool = False, + ) -> list[str] | None: + """Narrow a key/team tool allowlist by the agent's tool permissions and the caller's org tool + ceiling. Each level only intersects; None at a level means no restriction from it. + + An UNRESOLVABLE org ceiling is decided per caller shape, mirroring the servers axis: a key stays + fail-open (skip the org step, keep the key/team/agent restrictions; letting the raise escape + would collapse them to allow-all, WIDER than before the fault), while a keyless source re-raises + so the outer handler denies that one source (its only org bound is this ceiling).""" + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + + if user_api_key_auth.agent_id: + # Pre-fetch agent object_permission once to avoid a duplicate DB query. + agent_obj_perm = await MCPRequestHandler._get_agent_object_permission(user_api_key_auth) + agent_tools = await MCPRequestHandler._get_agent_tool_permissions_for_server( + server_id=server_id, + user_api_key_auth=user_api_key_auth, + agent_object_permission=agent_obj_perm, + ) + if agent_tools is not None: + allowed_tools = ( + list(set(allowed_tools) & set(agent_tools)) if allowed_tools is not None else agent_tools + ) + + if user_api_key_auth.org_id: + # _get_org_object_permission uses user_api_key_cache, so this is not a fresh DB round-trip + # when get_allowed_mcp_servers was already called. + try: + org_obj_perm = await MCPRequestHandler._get_org_object_permission(user_api_key_auth) + except Exception as e: # noqa: BLE001 # unresolvable org ceiling, decided per caller shape + if keyless_source: + raise + verbose_logger.warning( + f"MCP org tool ceiling unresolvable for org_id={user_api_key_auth.org_id!r}; " + f"skipping org intersect, key/team/agent restrictions stand: {str(e)}" + ) + return allowed_tools + org_tools = ( + global_mcp_server_manager.expand_tool_permissions(org_obj_perm.mcp_tool_permissions).get(server_id) + if org_obj_perm and org_obj_perm.mcp_tool_permissions + else None + ) + if org_tools is not None: + allowed_tools = ( + list(set(allowed_tools) & set(org_tools)) if allowed_tools is not None else list(org_tools) + ) + + return allowed_tools @staticmethod async def is_tool_allowed_for_server( @@ -1537,10 +2068,96 @@ class MCPRequestHandler: @staticmethod async def _get_allowed_mcp_servers_for_team( - user_api_key_auth: Optional[UserAPIKeyAuth] = None, - ) -> List[str]: + user_api_key_auth: UserAPIKeyAuth | None = None, + ) -> list[str]: + """Get allowed MCP servers a caller inherits from the team it is pinned to. + + Exactly one team, or none. A subject that reaches servers through SEVERAL teams does not + fan out here: it is resolved one source per team in ``_resolve_admitted_subject_servers``, + and each of those sources pins a single ``team_id`` before reaching this point. Keeping the + fan-out here as well would be a second multi-team path to drift from that one. """ - Get allowed MCP servers for a team. + team_ids = await MCPRequestHandler._team_ids_for_mcp_grant(user_api_key_auth) + if not team_ids: + return [] + return await MCPRequestHandler._allowed_mcp_servers_for_single_team(team_ids[0], user_api_key_auth) + + @staticmethod + async def _team_ids_for_mcp_grant(user_api_key_auth: UserAPIKeyAuth | None) -> list[str]: + """The team ids whose MCP grants a caller inherits. + + A caller with an explicit ``team_id`` uses that single team; every other caller inherits no + team grants. That covers key auth and JWT auth (a keyless ``user_id`` auth with no team_id, + which must NOT silently gain the union across every team the user belongs to), and it covers + each single-source auth an admitted subject fans out into — those pin a team_id, so they land + on the first branch. The admitted subject itself never reaches here: it resolves per source + in ``_resolve_admitted_subject_servers`` before this point. The ``UI_TEAM_ID`` sentinel + resolves to no teams exactly as before.""" + if user_api_key_auth is None or not user_api_key_auth.team_id: + return [] + return [] if user_api_key_auth.team_id == UI_TEAM_ID else [user_api_key_auth.team_id] + + @staticmethod + async def _resolve_user_team_ids(user_id: str, user_api_key_auth: UserAPIKeyAuth) -> list[str]: + """The distinct team ids a user belongs to, from the live user record. Returns [] on + no DB, a missing user, or any resolution failure so a lookup blip narrows access + rather than raising; the caller's direct grants still apply.""" + from litellm.proxy.auth.auth_checks import get_user_object + from litellm.proxy.proxy_server import ( + prisma_client, + proxy_logging_obj, + user_api_key_cache, + ) + + if prisma_client is None: + return [] + try: + user_object = await get_user_object( + user_id=user_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + user_id_upsert=False, + parent_otel_span=user_api_key_auth.parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + ) + except Exception as e: # noqa: BLE001 # a team-resolution blip narrows access, never raises + verbose_logger.warning(f"Failed to resolve user teams for MCP grant: {str(e)}") + return [] + if user_object is None or not user_object.teams: + return [] + return list(dict.fromkeys(t for t in user_object.teams if t and t != UI_TEAM_ID)) + + @staticmethod + async def _team_granted_servers(team_obj: LiteLLM_TeamTable, team_access_group_servers: list[str]) -> set[str]: + """The raw MCP-server set a team grants (before any org ceiling): its object_permission (direct + ``mcp_servers``, the ``all_proxy_servers`` sentinel → the full registry, legacy access groups, + tool-perm-referenced servers) unioned with its unified ``access_group_ids`` servers.""" + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + + object_permissions = team_obj.object_permission + if object_permissions is None: + return set(team_access_group_servers) + if SpecialMCPServerName.all_proxy_servers.value in (object_permissions.mcp_servers or []): + return set(global_mcp_server_manager.get_registry().keys()) + legacy_access_group_servers = await MCPRequestHandler._get_mcp_servers_from_access_groups( + object_permissions.mcp_access_groups or [] + ) + return ( + set(global_mcp_server_manager.expand_permission_list(object_permissions.mcp_servers or [])) + | set(legacy_access_group_servers) + | set(global_mcp_server_manager.expand_tool_permissions(object_permissions.mcp_tool_permissions).keys()) + | set(team_access_group_servers) + ) + + @staticmethod + async def _allowed_mcp_servers_for_single_team( + team_id: str, + user_api_key_auth: UserAPIKeyAuth | None, + ) -> list[str]: + """Allowed MCP servers granted by ONE team (its raw grant, then capped by the team's own org + for a keyless admitted subject). Unions two sources: - Legacy team.object_permission (mcp_servers, mcp_access_groups, @@ -1551,9 +2168,6 @@ class MCPRequestHandler: the gate (no assigned_team_ids check needed here). """ try: - from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( - global_mcp_server_manager, - ) from litellm.proxy.auth.auth_checks import ( _get_mcp_server_ids_from_access_groups, get_team_object, @@ -1564,22 +2178,24 @@ class MCPRequestHandler: user_api_key_cache, ) - if user_api_key_auth is None or not user_api_key_auth.team_id or prisma_client is None: + if not team_id or team_id == UI_TEAM_ID or prisma_client is None: return [] - if user_api_key_auth.team_id == UI_TEAM_ID: - return [] - - team_obj: Optional[LiteLLM_TeamTable] = await get_team_object( - team_id=user_api_key_auth.team_id, + parent_otel_span = user_api_key_auth.parent_otel_span if user_api_key_auth is not None else None + team_obj: LiteLLM_TeamTable | None = await get_team_object( + team_id=team_id, prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, - parent_otel_span=user_api_key_auth.parent_otel_span, + parent_otel_span=parent_otel_span, proxy_logging_obj=proxy_logging_obj, ) if team_obj is None: return [] - + if team_obj.blocked: + # A blocked team grants nothing. The central policy gate enforces this for a key + # pinned to a single team_id, but a keyless admitted identity (no team_id) unions + # across all of its teams and would otherwise inherit a blocked team's MCP grants. + return [] team_access_group_servers = await _get_mcp_server_ids_from_access_groups( access_group_ids=team_obj.access_group_ids or [], prisma_client=prisma_client, @@ -1587,27 +2203,8 @@ class MCPRequestHandler: proxy_logging_obj=proxy_logging_obj, ) - object_permissions = team_obj.object_permission - if object_permissions is None: - return list(set(team_access_group_servers)) - - if SpecialMCPServerName.all_proxy_servers.value in (object_permissions.mcp_servers or []): - return list(global_mcp_server_manager.get_registry().keys()) - - direct_mcp_servers = global_mcp_server_manager.expand_permission_list(object_permissions.mcp_servers or []) - - legacy_access_group_servers = await MCPRequestHandler._get_mcp_servers_from_access_groups( - object_permissions.mcp_access_groups or [] - ) - - tool_perm_servers = list( - global_mcp_server_manager.expand_tool_permissions(object_permissions.mcp_tool_permissions).keys() - ) - - all_servers = ( - direct_mcp_servers + legacy_access_group_servers + tool_perm_servers + team_access_group_servers - ) - return list(set(all_servers)) + servers = await MCPRequestHandler._team_granted_servers(team_obj, team_access_group_servers) + return list(servers) except Exception as e: verbose_logger.warning(f"Failed to get allowed MCP servers for team: {str(e)}") return [] @@ -1621,7 +2218,11 @@ class MCPRequestHandler: ``get_object_permission`` helpers so MCP requests share the same ``user_api_key_cache`` entries as the rest of the proxy. """ - from litellm.proxy.auth.auth_checks import get_object_permission, get_org_object + from litellm.proxy.auth.auth_checks import ( + OrganizationNotFoundError, + get_object_permission, + get_org_object, + ) from litellm.proxy.proxy_server import ( prisma_client, proxy_logging_obj, @@ -1635,6 +2236,8 @@ class MCPRequestHandler: verbose_logger.debug("prisma_client is None") return None + # A team's organization_id can point at a deleted or not-yet-synced row; get_org_object raises + # OrganizationNotFoundError for that. That is a determinate ABSENCE (no ceiling), handled below. try: org_obj = await get_org_object( org_id=user_api_key_auth.org_id, @@ -1643,21 +2246,32 @@ class MCPRequestHandler: parent_otel_span=user_api_key_auth.parent_otel_span, proxy_logging_obj=proxy_logging_obj, ) - - if org_obj is None or not org_obj.object_permission_id: - return None - - return await get_object_permission( - object_permission_id=org_obj.object_permission_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - parent_otel_span=user_api_key_auth.parent_otel_span, - proxy_logging_obj=proxy_logging_obj, - ) - except Exception as e: - verbose_logger.warning(f"Failed to get org object permission: {str(e)}") + except OrganizationNotFoundError as e: + # CONFIRMED absent: places no ceiling. Every OTHER exception propagates as an unresolvable + # ceiling (denies for a keyless source, fail-open for a key); catching bare Exception here + # would treat a DB outage as "no org" and silently drop a real ceiling for its duration. + verbose_logger.debug(f"MCP org ceiling: org {user_api_key_auth.org_id!r} does not exist: {e}") return None + if org_obj is None or not org_obj.object_permission_id: + return None + + # The org NAMES a permission; failing to read it is INDETERMINATE and must not collapse into the + # None that means "no ceiling". Raise and let each caller pick fail-open or fail-closed. + object_permission = await get_object_permission( + object_permission_id=org_obj.object_permission_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=user_api_key_auth.parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + ) + if object_permission is None: + raise ValueError( + f"org {user_api_key_auth.org_id!r} names object_permission_id " + f"{org_obj.object_permission_id!r} which could not be loaded" + ) + return object_permission + @staticmethod async def _get_allowed_mcp_servers_for_org( user_api_key_auth: Optional[UserAPIKeyAuth] = None, @@ -1692,8 +2306,10 @@ class MCPRequestHandler: all_servers = direct_mcp_servers + access_group_servers + tool_perm_servers return list(set(all_servers)) except Exception as e: + # None = ceiling UNRESOLVED, distinct from [] = org places no restriction. Collapsing them + # let a DB fault silently drop a ceiling; the caller picks fail-open/closed from this signal. verbose_logger.warning(f"Failed to get allowed MCP servers for org: {str(e)}") - return [] + return None @staticmethod async def _get_allowed_mcp_servers_for_end_user( diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index 9a1b5cf4864..26241119dd8 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -3,6 +3,7 @@ import html as _html import json import secrets import time +from collections.abc import Mapping from datetime import datetime, timezone from typing import TYPE_CHECKING, Any, Dict, Literal, Optional, Tuple from urllib.parse import parse_qsl, urlencode, urlparse, urlunparse @@ -13,6 +14,7 @@ from fastapi.responses import HTMLResponse, JSONResponse, RedirectResponse, Resp from pydantic import BaseModel, ConfigDict, Field, SecretStr, ValidationError from litellm._logging import verbose_logger +from litellm.caching.in_memory_cache import InMemoryCache from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, httpxSpecialProvider, @@ -20,7 +22,9 @@ from litellm.llms.custom_httpx.http_handler import ( from litellm.proxy._experimental.mcp_server.auth.token_endpoint_auth import ( TokenEndpointAuthConfigError, build_token_endpoint_client_auth, + normalize_token_endpoint_auth_method, ) +from litellm.types.mcp_server.mcp_server_manager import MCPTokenEndpointAuthMethod from litellm.proxy._experimental.mcp_server.bridge_token_flow import ( _bridge_mint_error_response, _BridgeMintReady, @@ -29,6 +33,7 @@ from litellm.proxy._experimental.mcp_server.bridge_token_flow import ( _finish_bridge_mint, _prepare_bridge_mint, _prepare_bridge_refresh, + _reload_active_user_by_id, ) from litellm.proxy._experimental.mcp_server.faults import ( CallerRejected, @@ -39,6 +44,14 @@ from litellm.proxy._experimental.mcp_server.faults import ( dcr_fault_detail, render_token_fault, ) +from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import ( + aggregate_authorize, + aggregate_token, + complete_connect_flow, + is_gateway_dcr_client_id, + register_aggregate_client, + relative_request_url, +) from litellm.proxy._experimental.mcp_server.oauth_utils import ( TOKEN_NO_CACHE_HEADERS, get_request_base_url, @@ -111,6 +124,9 @@ def encode_state_with_base_url( client_redirect_uri: Optional[str] = None, litellm_user_id: str | None = None, mcp_server_id: str | None = None, + dcr_client_id: str | None = None, + dcr_client_secret: str | None = None, + dcr_token_endpoint_auth_method: MCPTokenEndpointAuthMethod | None = None, ) -> str: """ Encode the base_url, original state, and PKCE parameters using encryption. @@ -124,8 +140,18 @@ def encode_state_with_base_url( litellm_user_id: The SSO-authenticated litellm user captured at the bridge authorize (interactive dcr_bridge oauth_delegate only); the callback seals it into the gateway authorization code so the token mint can bind the envelope to this user - mcp_server_id: The bridge server the interactive flow targets, sealed alongside - litellm_user_id so the gateway code cannot be replayed against another server + mcp_server_id: The server the flow targets, sealed alongside litellm_user_id (bridge) or + dcr_client_id (ephemeral mint) so the gateway code cannot be replayed against another + server + dcr_client_id: The ephemeral DCR client the gateway minted at authorize for a + client-forwarded-token server with no caller-supplied client; the callback seals it + into the forwarded authorization code so the token exchange can authenticate with it + while the gateway stores nothing + dcr_client_secret: The minted client's secret, when the upstream issued one + dcr_token_endpoint_auth_method: The token-endpoint auth method the upstream's registration + response granted the minted client, sealed alongside the credentials so the exchange + authenticates the way the upstream expects instead of falling back to the server row's + configured method Returns: An encrypted string that encodes all values @@ -138,6 +164,9 @@ def encode_state_with_base_url( "client_redirect_uri": client_redirect_uri, "litellm_user_id": litellm_user_id, "mcp_server_id": mcp_server_id, + "dcr_client_id": dcr_client_id, + "dcr_client_secret": dcr_client_secret, + "dcr_token_endpoint_auth_method": dcr_token_endpoint_auth_method, } state_json = json.dumps(state_data, sort_keys=True) encrypted_state = encrypt_value_helper(state_json) @@ -217,14 +246,112 @@ def open_bridge_authorization_code(code: str) -> _BridgeAuthorizationCode | None return None +_PASSTHROUGH_AUTH_CODE_PREFIX = "llm_ptcode_" + + +class PassthroughAuthorizationCode(BaseModel): + """The ephemeral DCR client and upstream code the gateway seals into the authorization code it + forwards for a client-forwarded-token server (``true_passthrough`` / ``oauth_delegate``) whose + authorize fell through to gateway-side registration. These modes forbid the gateway from storing + an OAuth client identity, so the minted client survives only inside this sealed value: the + client echoes it back at the token endpoint, where the gateway recovers the client to + authenticate the upstream exchange. ``mcp_server_id`` binds the code to the server it was minted + for so it cannot be spent at another server's token endpoint.""" + + model_config = ConfigDict(frozen=True) + upstream_code: str = Field(min_length=1) + client_id: str = Field(min_length=1) + client_secret: str | None = None + token_endpoint_auth_method: MCPTokenEndpointAuthMethod | None = None + mcp_server_id: str = Field(min_length=1) + + +def seal_passthrough_authorization_code( + upstream_code: str, + client_id: str, + client_secret: str | None, + mcp_server_id: str, + token_endpoint_auth_method: MCPTokenEndpointAuthMethod | None = None, +) -> str: + """Seal the upstream authorization code together with the ephemeral DCR client that authorized + it. Encrypted with the same authenticated symmetric helper as the OAuth state and bridge codes, + so the client can neither read the (possibly confidential) client credentials nor forge a + code.""" + payload = json.dumps( + { + "upstream_code": upstream_code, + "client_id": client_id, + "client_secret": client_secret, + "token_endpoint_auth_method": token_endpoint_auth_method, + "mcp_server_id": mcp_server_id, + }, + sort_keys=True, + ) + return _PASSTHROUGH_AUTH_CODE_PREFIX + encrypt_value_helper(payload) + + +def open_passthrough_authorization_code(code: str) -> PassthroughAuthorizationCode | None: + """Recover the sealed ephemeral client and upstream code, or ``None`` when ``code`` is not a + gateway passthrough code or does not decrypt / validate, so a raw upstream code falls through to + the existing caller-supplied-client behavior.""" + if not code.startswith(_PASSTHROUGH_AUTH_CODE_PREFIX): + return None + decrypted = decrypt_value_helper( + code[len(_PASSTHROUGH_AUTH_CODE_PREFIX) :], "passthrough_authorization_code", return_original_value=False + ) + if not isinstance(decrypted, str): + return None + try: + return PassthroughAuthorizationCode.model_validate_json(decrypted) + except ValidationError: + return None + + +def redeem_passthrough_authorization_code( + code: str | None, mcp_server: MCPServer, code_verifier: str | None +) -> PassthroughAuthorizationCode | None: + """The single redemption gate for sealed passthrough codes: a raw or foreign code returns + ``None`` so the caller keeps its existing behavior, while a genuine sealed code must be spent + at the server it was minted for and must carry the PKCE verifier of the S256 flow that minted + it (the mint refuses downgraded flows, so a verifier-less redemption is an interception + attempt, not a legitimate client).""" + if not code: + return None + sealed = open_passthrough_authorization_code(code) + if sealed is None: + return None + if sealed.mcp_server_id != mcp_server.server_id: + raise HTTPException( + status_code=400, + detail="Authorization code was issued for a different MCP server", + ) + if not code_verifier: + raise HTTPException( + status_code=400, + detail="code_verifier is required to redeem this authorization code", + ) + return sealed + + +def _session_cookie_user_id(request: Request) -> str | None: + """The signed-in litellm user for a browser request, or ``None``. Thin wrapper so the + aggregate DCR flow's verbs receive the identity as a plain value instead of parsing + cookies themselves.""" + from litellm.proxy._experimental.mcp_server.byok_oauth_endpoints import ( # noqa: PLC0415 # circular import at module load + _user_id_from_session_cookie, + ) + + return _user_id_from_session_cookie(request) + + def _redirect_to_litellm_login(request: Request) -> RedirectResponse: """Send an unauthenticated browser through litellm login before the interactive bridge authorize can capture its identity. The bridge oauth_delegate flow seals the SSO user into the gateway code, - so a session is required; without one there is nothing to bind. After login the user re-initiates - the connection, which then finds the session cookie (the seamless return-to round-trip, which is - origin-validated against the control-plane URL, is a follow-up).""" + so a session is required; without one there is nothing to bind. A same-origin relative + ``return_to`` (honored by the SSO callback) brings the browser straight back to this authorize + request after login instead of stranding it on the dashboard.""" base_url = get_request_base_url(request) - return RedirectResponse(f"{base_url}/sso/key/generate") + return RedirectResponse(f"{base_url}/sso/key/generate?{urlencode({'return_to': relative_request_url(request)})}") # LIT-4197: some upstream authorization servers reject an over-long ``state`` @@ -499,6 +626,35 @@ def _raise_if_not_oauth2(mcp_server: MCPServer) -> None: ) +def _endpoint_not_configured_detail( + mcp_server: MCPServer, + endpoint_label: str, + manual_remedy: str, + issuer_remedy: str, +) -> str: + """The 400 detail for an unresolved OAuth endpoint, naming the likely cause for this server's + shape (LIT-4658): an anchored issuer whose metadata fell short, a configured (possibly + misconfigured) server url whose discovery failed, or no discovery source at all. Kept free of + URLs and issuer values because these endpoints are reachable pre-auth.""" + if mcp_server.issuer_is_anchored: + return ( + f"MCP server {endpoint_label} is not configured. Endpoint discovery anchored on the configured " + f"Issuer (RFC 8414) failed or its metadata did not include this endpoint; check the proxy logs " + f"for 'MCP OAuth' warnings from server load, verify the Issuer, or {manual_remedy}." + ) + if mcp_server.url: + return ( + f"MCP server {endpoint_label} is not configured. OAuth endpoint discovery against the configured " + f"server url did not resolve it; the url may be misconfigured. Check the proxy logs for " + f"'MCP OAuth' warnings from server load, verify the server url, or {manual_remedy}, or " + f"{issuer_remedy}." + ) + return ( + f"MCP server {endpoint_label} is not configured. Servers with no url (OpenAPI spec or stdio) run no " + f"resource discovery, so {manual_remedy}, or {issuer_remedy}." + ) + + def _raise_unless_oauth2_discovery_server( mcp_server: Optional[MCPServer], mcp_server_name: Optional[str], @@ -594,15 +750,17 @@ async def authorize_with_server( code_challenge_method: Optional[str] = None, response_type: Optional[str] = None, scope: Optional[str] = None, + ephemeral_dcr_client: "EphemeralDcrClient | None" = None, ): _raise_if_not_oauth2(mcp_server) if mcp_server.authorization_url is None: raise HTTPException( status_code=400, - detail=( - "MCP server authorization url is not configured. Servers with no url (OpenAPI " - "spec or stdio) run no resource discovery, so set Authorization URL and Token URL " - "manually, or set Issuer to discover them from the identity provider (RFC 8414)." + detail=_endpoint_not_configured_detail( + mcp_server, + "authorization url", + "set Authorization URL and Token URL manually", + "set Issuer to discover them from the identity provider (RFC 8414)", ), ) @@ -612,7 +770,10 @@ async def authorize_with_server( # calling this for its enforcement side effect, then falls through to the gateway # /callback flow below, which reads the original code_challenge names. bridge_challenge, bridge_method = _require_s256_pkce(code_challenge, code_challenge_method) - if _dcr_bridge_relays_client_registration(mcp_server): + # A gateway-minted ephemeral client is registered against {base}/callback, so its + # flow must run the short-circuit arm; the relay arm is only for clients that + # registered themselves through the front door and hold their own redirect binding. + if _dcr_bridge_relays_client_registration(mcp_server) and ephemeral_dcr_client is None: return _redirect_to_upstream_authorize( mcp_server=mcp_server, client_id=client_id, @@ -656,7 +817,12 @@ async def authorize_with_server( code_challenge_method=code_challenge_method, client_redirect_uri=redirect_uri, litellm_user_id=litellm_user_id, - mcp_server_id=mcp_server.server_id if litellm_user_id else None, + mcp_server_id=mcp_server.server_id if (litellm_user_id or ephemeral_dcr_client) else None, + dcr_client_id=ephemeral_dcr_client.client_id if ephemeral_dcr_client else None, + dcr_client_secret=ephemeral_dcr_client.client_secret if ephemeral_dcr_client else None, + dcr_token_endpoint_auth_method=ephemeral_dcr_client.token_endpoint_auth_method + if ephemeral_dcr_client + else None, ) relay_state = secrets.token_urlsafe(_OAUTH_STATE_HANDLE_BYTES) @@ -703,6 +869,7 @@ async def exchange_token_with_server( code_verifier: Optional[str], refresh_token: Optional[str] = None, scope: Optional[str] = None, + client_token_endpoint_auth_method: MCPTokenEndpointAuthMethod | None = None, ): _raise_if_not_oauth2(mcp_server) if grant_type not in ("authorization_code", "refresh_token"): @@ -711,22 +878,32 @@ async def exchange_token_with_server( if mcp_server.token_url is None: raise HTTPException( status_code=400, - detail=( - "MCP server token url is not configured. Servers with no url (OpenAPI spec or " - "stdio) run no resource discovery, so set Token URL manually, or set Issuer to " - "discover it from the identity provider (RFC 8414)." + detail=_endpoint_not_configured_detail( + mcp_server, + "token url", + "set Token URL manually", + "set Issuer to discover it from the identity provider (RFC 8414)", ), ) - # The id and secret must come from the same source. When the server-side client_id wins, - # falling back to the caller's secret pairs the persisted client with a foreign secret; the - # register short-circuit hands clients a placeholder secret ("dummy"), so a re-auth against a - # persisted public PKCE client (no stored secret) would send that placeholder and the IdP 401s. + # The id, secret, and token-endpoint auth method must come from the same source. When the + # server-side client_id wins, falling back to the caller's secret pairs the persisted client + # with a foreign secret; the register short-circuit hands clients a placeholder secret + # ("dummy"), so a re-auth against a persisted public PKCE client (no stored secret) would send + # that placeholder and the IdP 401s. Symmetrically, a caller-side client (an ephemeral mint + # recovered from a sealed code) must authenticate the way its own registration was granted, + # not the way the server row is configured; callers that carry no method keep the row's method + # as before. resolved_client_id = mcp_server.client_id if mcp_server.client_id else client_id resolved_client_secret = mcp_server.client_secret if mcp_server.client_id else client_secret + resolved_auth_method = ( + mcp_server.token_endpoint_auth_method + if mcp_server.client_id + else (client_token_endpoint_auth_method or mcp_server.token_endpoint_auth_method) + ) try: client_auth = build_token_endpoint_client_auth( - auth_method=mcp_server.token_endpoint_auth_method, + auth_method=resolved_auth_method, client_id=resolved_client_id, client_secret=resolved_client_secret, ) @@ -1229,7 +1406,7 @@ async def _persist_dcr_client_registration( return "failed" -def _client_supplied_redirect_uris(value: object) -> list[str] | None: +def client_supplied_redirect_uris(value: object) -> list[str] | None: """RFC 7591 redirect_uris must be a non-empty array of URI strings. Any other shape (not a list, an empty list, or a list holding a non-string or empty-string element) yields None so every register arm falls back to the gateway callback instead of echoing a malformed value back to the @@ -1241,6 +1418,142 @@ def _client_supplied_redirect_uris(value: object) -> list[str] | None: return uris if len(uris) == len(value) else None +async def _post_dcr_registration( + registration_url: str, + register_data: Mapping[str, object], + server_id: str, +) -> httpx.Response: + """POST an RFC 7591 registration to the upstream and return its response, relaying a classified + upstream rejection instead of a generic 500 and failing loud on an absent response.""" + headers = { + "Content-Type": "application/json", + "Accept": "application/json", + } + async_client = get_async_httpx_client(llm_provider=httpxSpecialProvider.Oauth2Register) + try: + response = await async_client.post( + registration_url, + headers=headers, + json=register_data, + ) + if response is not None: + response.raise_for_status() + except httpx.HTTPStatusError as exc: + status_code, detail = dcr_fault_detail(classify_upstream_dcr_rejection(exc.response, log_context=server_id)) + raise HTTPException(status_code=status_code, detail=detail) from exc + if response is None: + raise HTTPException( + status_code=502, + detail="MCP upstream registration endpoint returned no response", + ) + return response + + +class EphemeralDcrClient(BaseModel): + """A DCR client minted for a single authorize round trip and never stored by the gateway.""" + + model_config = ConfigDict(frozen=True) + client_id: str = Field(min_length=1) + client_secret: str | None = None + token_endpoint_auth_method: MCPTokenEndpointAuthMethod | None = None + + +_EPHEMERAL_DCR_CLIENT_CACHE = InMemoryCache(default_ttl=_OAUTH_STATE_COOKIE_TTL_SECONDS) +_EPHEMERAL_DCR_MINT_LOCKS: dict[str, asyncio.Lock] = {} + + +async def mint_ephemeral_dcr_client(request: Request, mcp_server: MCPServer) -> EphemeralDcrClient | None: + """Mint a throwaway OAuth client via the upstream's RFC 7591 registration endpoint for a + client-forwarded-token server whose authorize arrived with no client_id. Returns ``None`` when + the upstream exposes no registration endpoint, so the caller keeps its existing failure path. + The minted client is deliberately not persisted anywhere: ``true_passthrough`` / + ``oauth_delegate`` require the gateway to hold no OAuth client identity, so it survives only in + the encrypted OAuth state and the sealed authorization code the callback forwards. + + Reloading the authorize page or retrying a flow must not register a fresh upstream client every + time (an OAuth client identifies the application, not the user, so reuse is semantically + correct). A per-process TTL cache bounded to the OAuth state cookie's lifetime dedupes the mint + per (server, gateway origin), and a per-server lock single-flights concurrent mints (the + ``_OAUTH_METADATA_FETCH_LOCKS`` pattern; keyed by server_id alone so the lock registry stays + bounded by the server count even when the request origin varies) so parallel authorize requests + cannot each register an upstream client; the cache stamps nothing onto the server record and + correctness never depends on it because the sealed state carries the client through the flow.""" + if mcp_server.registration_url is None: + return None + request_base_url = get_request_base_url(request) + cache_key = f"mcp_ephemeral_dcr_client:{mcp_server.server_id}:{request_base_url}" + cached = _EPHEMERAL_DCR_CLIENT_CACHE.get_cache(cache_key) + if isinstance(cached, EphemeralDcrClient): + return cached + lock = _EPHEMERAL_DCR_MINT_LOCKS.setdefault(mcp_server.server_id, asyncio.Lock()) + async with lock: + cached_after_wait = _EPHEMERAL_DCR_CLIENT_CACHE.get_cache(cache_key) + if isinstance(cached_after_wait, EphemeralDcrClient): + return cached_after_wait + register_data: dict[str, object] = { + "client_name": mcp_server.server_name or mcp_server.server_id, + "redirect_uris": [f"{request_base_url}/callback"], + "grant_types": ["authorization_code", "refresh_token"], + "response_types": ["code"], + "token_endpoint_auth_method": "none", + } + response = await _post_dcr_registration( + registration_url=mcp_server.registration_url, + register_data=register_data, + server_id=mcp_server.server_id, + ) + try: + registration = _DcrClientRegistration.model_validate_json(response.text) + except ValidationError as exc: + raise HTTPException( + status_code=502, + detail="MCP upstream registration endpoint returned no usable client_id", + ) from exc + if not registration.client_id: + raise HTTPException( + status_code=502, + detail="MCP upstream registration endpoint returned no usable client_id", + ) + minted = EphemeralDcrClient( + client_id=registration.client_id, + client_secret=registration.client_secret, + token_endpoint_auth_method=normalize_token_endpoint_auth_method(registration.token_endpoint_auth_method), + ) + _EPHEMERAL_DCR_CLIENT_CACHE.set_cache(cache_key, minted) + return minted + + +async def resolve_ephemeral_dcr_client( + request: Request, + mcp_server: MCPServer, + code_challenge: str | None, + code_challenge_method: str | None, + redirect_uri: str, +) -> EphemeralDcrClient | None: + """The single owner of the gateway-side mint policy for a clientless authorize. Returns + ``None`` for servers whose mode does not permit gateway minting and for upstreams without a + registration endpoint, so those callers keep their existing failure paths: plain ``oauth2`` + keeps its persisted-client contract, and the interactive ``oauth_delegate`` dcr_bridge + sign-in has its own sealed-identity flow. ``true_passthrough`` mints regardless of the + ``dcr_bridge`` flag (the UI creates passthrough servers with the flag on by default): a + minted flow runs the bridge short-circuit arm, while the relay front door remains for + external clients that registered themselves. Flows that could never succeed fail loud + before any upstream registration: a missing ``authorization_url``, a downgraded PKCE pair + (without S256 the sealed code would be bearer-redeemable by any authenticated caller who + intercepts the redirect), or an untrusted ``redirect_uri`` (a rejected redirect must not be + usable to generate orphan IdP clients).""" + if not (mcp_server.is_true_passthrough or (mcp_server.is_oauth_delegate and not mcp_server.is_dcr_bridge)): + return None + if mcp_server.authorization_url is None: + raise HTTPException( + status_code=400, + detail="MCP server authorization url is not set", + ) + _require_s256_pkce(code_challenge, code_challenge_method) + validate_trusted_redirect_uri(request, redirect_uri) + return await mint_ephemeral_dcr_client(request, mcp_server) + + async def register_client_with_server( request: Request, mcp_server: MCPServer, @@ -1278,10 +1591,11 @@ async def register_client_with_server( if mcp_server.authorization_url is None: raise HTTPException( status_code=400, - detail=( - "MCP server authorization url is not configured. Servers with no url (OpenAPI " - "spec or stdio) run no resource discovery, so set Authorization URL and Token URL " - "manually, or set Issuer to discover them from the identity provider (RFC 8414)." + detail=_endpoint_not_configured_detail( + mcp_server, + "authorization url", + "set Authorization URL and Token URL manually", + "set Issuer to discover them from the identity provider (RFC 8414)", ), ) @@ -1302,30 +1616,11 @@ async def register_client_with_server( "response_types": response_types or (["code"] if bridge_relay else []), "token_endpoint_auth_method": token_endpoint_auth_method or ("none" if bridge_relay else ""), } - headers = { - "Content-Type": "application/json", - "Accept": "application/json", - } - - async_client = get_async_httpx_client(llm_provider=httpxSpecialProvider.Oauth2Register) - try: - response = await async_client.post( - mcp_server.registration_url, - headers=headers, - json=register_data, - ) - if response is not None: - response.raise_for_status() - except httpx.HTTPStatusError as exc: - status_code, detail = dcr_fault_detail( - classify_upstream_dcr_rejection(exc.response, log_context=mcp_server.server_id) - ) - raise HTTPException(status_code=status_code, detail=detail) from exc - if response is None: - raise HTTPException( - status_code=502, - detail="MCP upstream registration endpoint returned no response", - ) + response = await _post_dcr_registration( + registration_url=mcp_server.registration_url, + register_data=register_data, + server_id=mcp_server.server_id, + ) token_response = response.json() @@ -1358,6 +1653,18 @@ async def authorize( global_mcp_server_manager, ) + if mcp_server_name is None and client_id and is_gateway_dcr_client_id(client_id): + return aggregate_authorize( + request=request, + client_id=client_id, + redirect_uri=redirect_uri, + state=state, + code_challenge=code_challenge, + code_challenge_method=code_challenge_method, + response_type=response_type, + session_user_id=_session_cookie_user_id(request), + ) + lookup_name: Optional[str] = mcp_server_name or client_id client_ip = IPAddressUtils.get_mcp_client_ip(request) mcp_server = ( @@ -1421,6 +1728,25 @@ async def token_endpoint( global_mcp_server_manager, ) + if mcp_server_name is None and is_gateway_dcr_client_id(client_id): + from litellm.proxy.proxy_server import ( # noqa: PLC0415 # circular import at module load + master_key, + user_api_key_cache, + ) + + return await aggregate_token( + request=request, + grant_type=grant_type, + code=code, + redirect_uri=redirect_uri, + client_id=client_id, + code_verifier=code_verifier, + refresh_token=refresh_token, + master_key=master_key, + reload_user=_reload_active_user_by_id, + cache=user_api_key_cache, + ) + lookup_name = mcp_server_name or client_id client_ip = IPAddressUtils.get_mcp_client_ip(request) mcp_server = global_mcp_server_manager.get_mcp_server_by_name(lookup_name, client_ip=client_ip) @@ -1442,6 +1768,21 @@ async def token_endpoint( ) +@router.post("/authorize/complete") +async def authorize_complete(request: Request, flow: str = Form(...)): + """Finish an aggregate connect flow: mint the gateway authorization code for the + signed-in user and redirect back to the DCR client. POST plus the per-flow HttpOnly + cookie set at /authorize; an anonymous or bad-flow request just 400s.""" + from litellm.proxy.proxy_server import user_api_key_cache # noqa: PLC0415 # circular import at module load + + return await complete_connect_flow( + request=request, + flow_handle=flow, + session_user_id=_session_cookie_user_id(request), + cache=user_api_key_cache, + ) + + # Per RFC 6749 §4.1.2.1, an IdP that rejects an OAuth authorization request # redirects back to the configured redirect URI with ``error`` / # ``error_description`` / ``error_uri`` query params and no ``code``. The MCP @@ -1563,11 +1904,23 @@ async def callback( # envelope to this user. Every other flow forwards the raw code unchanged. litellm_user_id = state_data.get("litellm_user_id") mcp_server_id = state_data.get("mcp_server_id") + dcr_client_id = state_data.get("dcr_client_id") + dcr_client_secret = state_data.get("dcr_client_secret") forwarded_code = code if isinstance(litellm_user_id, str) and litellm_user_id and isinstance(mcp_server_id, str) and mcp_server_id: forwarded_code = seal_bridge_authorization_code( upstream_code=code, litellm_user_id=litellm_user_id, mcp_server_id=mcp_server_id ) + elif isinstance(dcr_client_id, str) and dcr_client_id and isinstance(mcp_server_id, str) and mcp_server_id: + forwarded_code = seal_passthrough_authorization_code( + upstream_code=code, + client_id=dcr_client_id, + client_secret=dcr_client_secret if isinstance(dcr_client_secret, str) and dcr_client_secret else None, + mcp_server_id=mcp_server_id, + token_endpoint_auth_method=normalize_token_endpoint_auth_method( + state_data.get("dcr_token_endpoint_auth_method") + ), + ) params = {"code": forwarded_code, "state": original_state} complete_returned_url = _append_query_params(redirect_uri, params) @@ -2158,7 +2511,7 @@ async def register_client(request: Request, mcp_server_name: Optional[str] = Non request_data = await _read_request_body(request=request) data: dict = {**request_data} - client_redirect_uris = _client_supplied_redirect_uris(data.get("redirect_uris")) + client_redirect_uris = client_supplied_redirect_uris(data.get("redirect_uris")) dummy_return = { "client_id": mcp_server_name or "dummy_client", @@ -2167,6 +2520,13 @@ async def register_client(request: Request, mcp_server_name: Optional[str] = Non } client_ip = IPAddressUtils.get_mcp_client_ip(request) if not mcp_server_name: + # A real DCR request carries redirect_uris (RFC 7591): route it to the aggregate DCR + # endpoint the aggregate authorization-server metadata advertises. A single-server + # deployment registers at /{server}/register instead (its bare-origin discovery + # advertises that), so this does not affect it. A request without redirect_uris is not + # a DCR request, so the legacy single-server-or-dummy fallback is kept for it. + if data.get("redirect_uris"): + return await register_aggregate_client(request=request, request_body=data) resolved = _resolve_oauth2_server_for_root_endpoints(client_ip=client_ip) if resolved: return await register_client_with_server( diff --git a/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py b/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py new file mode 100644 index 00000000000..58233c4c9e5 --- /dev/null +++ b/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py @@ -0,0 +1,637 @@ +"""The gateway-level DCR flow for the aggregate ``/mcp`` endpoint (``mcp_gateway_dcr``). + +An OAuth-only DCR client (Claude Desktop, Claude Code, MCP Inspector) pointed at the +aggregate ``/mcp`` endpoint discovers the gateway as its authorization server (PR 1 of +this track) and then walks the flow implemented here: + +1. ``POST /register``: stateless dynamic client registration. The ``client_id`` IS the + registration: the client's redirect URIs are sealed into it with the repo's + authenticated symmetric helper, so nothing is persisted and a forged or tampered + client_id simply fails to open. Clients are always public (``token_endpoint_auth_method + "none"``); PKCE S256 is what protects the code. +2. ``GET /authorize``: validates the client and redirect URI, requires S256 PKCE, and + interposes LiteLLM sign-in. Without a session cookie the browser is sent through + ``/sso/key/generate`` with a same-origin ``return_to`` so it lands back here after + login. With a session, the flow parameters and the SSO user are sealed into a per-flow + HttpOnly cookie (the same pattern as the upstream OAuth state relay) and the browser is + sent to the connect page, where the user authorizes individual servers (vaulting those + tokens server-side) before finishing. +3. ``POST /authorize/complete``: the deliberate finish step. A POST (not GET) bound to the + SameSite=Lax flow cookie, so a cross-site link cannot silently mint a code with the + victim's session, and the signed-in user must match the user sealed into the flow. + Mints a short-lived, single-use, gateway-sealed authorization code and redirects to the + client's registered redirect URI. +4. ``POST /token``: exchanges the code (PKCE-verified, client- and redirect-bound, + single-use) for the identity-only session tokens of + :mod:`.outbound_credentials.session_token`, re-validating that the litellm user is + still active first; the ``refresh_token`` grant rotates the pair the same way. + +Nothing here stores state server-side except the single-use code guard (a TTL cache +entry). Every sealed value is authenticated encryption over the proxy salt/master key +family, opened totally (bad input maps to an OAuth error, never a raise), and every +identity is a stable reference re-validated live at mint, refresh, and (in the admission +PR) tool-call time. Upstream server credentials never appear anywhere in this flow; they +are vaulted per user by the existing ``/v1/mcp`` authorize endpoints and resolved at +egress by user id. +""" + +from __future__ import annotations + +import hashlib +import hmac +import secrets +from base64 import urlsafe_b64encode +from collections.abc import Mapping +from datetime import datetime, timezone +from typing import Awaitable, Callable, Literal, TypeVar +from urllib.parse import parse_qsl, urlencode, urlparse, urlunparse + +from fastapi import HTTPException, Request +from fastapi.responses import JSONResponse, RedirectResponse, Response +from pydantic import BaseModel, ConfigDict, Field, ValidationError +from typing_extensions import assert_never + +from litellm._logging import verbose_logger +from litellm.caching.caching import DualCache +from litellm.proxy._experimental.mcp_server.oauth_utils import ( + TOKEN_NO_CACHE_HEADERS, + get_request_base_url, + is_loopback_redirect_host, + validate_redirect_uri_shape, +) +from litellm.proxy._experimental.mcp_server.outbound_credentials.session_credentials import ( + SessionRefreshOpened, + open_session_refresh_bearer, + session_keys_from_master_key, +) +from litellm.proxy._experimental.mcp_server.outbound_credentials.session_token import ( + SESSION_REFRESH_TTL_SECONDS, + MintedSessionToken, + SessionKeys, + SessionPrincipal, + mint_session_refresh_token, + mint_session_token, +) +from litellm.proxy.common_utils.encrypt_decrypt_utils import ( + decrypt_value_helper, + encrypt_value_helper, +) + +GATEWAY_DCR_CLIENT_ID_PREFIX = "llm_dcrc_" +"""Marker prefix on every gateway-issued DCR client_id so the root authorize/token +endpoints can route an aggregate-flow request without decrypting, and existing per-server +flows (whose client_ids are upstream-issued) are never captured by the aggregate arm.""" + +GATEWAY_AUTH_CODE_PREFIX = "llm_gcode_" +"""Marker prefix on the gateway-sealed authorization code, distinct from the bridge +``llm_bcode_`` so neither flow can consume the other's codes.""" + +CONNECT_FLOW_COOKIE_PREFIX = "mcp_connect_flow_" +"""Per-flow HttpOnly cookie holding the sealed connect flow, keyed by a short random +handle carried in the connect-page URL (the same handle-plus-cookie pattern as the +``mcp_oauth_state_`` upstream relay, for the same reasons: replica-safe with no +server-side session store, and the sealed value never appears in a URL).""" + +CONNECT_FLOW_TTL_SECONDS = 600 +GATEWAY_AUTH_CODE_TTL_SECONDS = 120 +_CLAIM_TTL_BUFFER_SECONDS = 60 +_USED_CODE_CACHE_PREFIX = "mcp_gateway_dcr_code_used:" +_USED_FLOW_CACHE_PREFIX = "mcp_gateway_dcr_flow_used:" +_USED_REFRESH_CACHE_PREFIX = "mcp_gateway_dcr_refresh_used:" + +MAX_REDIRECT_URIS = 3 +MAX_REDIRECT_URI_LENGTH = 256 +MAX_CLIENT_ID_LENGTH = 2048 +"""Registration bounds. They exist to bound the sealed client_id, which rides inside +every session-token claim set: 3 URIs of 256 bytes seal to roughly 1.2KB, comfortably +under this cap and under the session token's own 4KB ceiling. Claude Desktop and MCP +Inspector register one or two redirect URIs.""" + +MAX_STATE_LENGTH = 1024 +"""Bound on the client ``state`` sealed into the flow cookie and echoed on the auth-code +redirect. An unbounded ``state`` can push the sealed cookie past the browser's ~4KB cap +(silently dropped, breaking the flow); spec clients send a short opaque value.""" + +MIN_CODE_VERIFIER_LENGTH = 43 +MAX_CODE_VERIFIER_LENGTH = 128 +"""RFC 7636 section 4.1 bounds for the PKCE ``code_verifier``. Enforced so an out-of-range +verifier gets a clean ``invalid_request`` instead of an opaque PKCE-mismatch.""" + +_UNPREFIXED = "" +"""Prefix for a sealed value that carries no wire marker because it is never routed by +prefix (the connect flow lives only in its own per-handle cookie, opened by that one +handle). Named so the empty-string argument to ``_seal`` / ``_open_sealed`` reads as +deliberate rather than a typo.""" + +_CLIENT_RECORD_DEBUG_KEY = "gateway_dcr_client" +_CONNECT_FLOW_DEBUG_KEY = "gateway_connect_flow" +_AUTH_CODE_DEBUG_KEY = "gateway_authorization_code" + +ReloadUserFailure = Literal["unresolvable", "unavailable", "no_active_key"] +ReloadUser = Callable[[str], Awaitable[ReloadUserFailure | None]] +"""Injected live-user revalidation (the token endpoint's mirror of admission): +``None`` means the user is active; ``unavailable`` is a retryable DB outage; anything +else fails the grant closed.""" + + +class GatewayDcrClient(BaseModel): + """The registration record sealed into a gateway DCR ``client_id``. + + ``extra="forbid"`` so a sealed value of another type (an auth code, a connect flow) + that happened to decrypt under the shared key can never validate as a client record: + cross-type confusion is rejected at the model boundary, not left to differing required + fields.""" + + model_config = ConfigDict(frozen=True, extra="forbid") + redirect_uris: tuple[str, ...] = Field(min_length=1, max_length=MAX_REDIRECT_URIS) + iat: int + + +class _ConnectFlow(BaseModel): + """One in-flight authorize: the SSO user it belongs to and the client parameters + needed to mint the code at the finish step. Sealed into the per-flow cookie. ``jti`` + makes the flow single-use at complete; ``extra="forbid"`` rejects cross-type + confusion.""" + + model_config = ConfigDict(frozen=True, extra="forbid") + user_id: str = Field(min_length=1) + client_id: str = Field(min_length=1) + redirect_uri: str = Field(min_length=1) + state: str + code_challenge: str = Field(min_length=1) + jti: str = Field(min_length=1) + exp: int + + +class _GatewayAuthCode(BaseModel): + """The gateway-sealed authorization code: the user consent it represents and the + bindings the token endpoint must verify (client, redirect URI, PKCE challenge), + plus a ``jti`` for the single-use guard. ``extra="forbid"`` rejects cross-type + confusion.""" + + model_config = ConfigDict(frozen=True, extra="forbid") + user_id: str = Field(min_length=1) + client_id: str = Field(min_length=1) + redirect_uri: str = Field(min_length=1) + code_challenge: str = Field(min_length=1) + jti: str = Field(min_length=1) + iat: int + exp: int + + +def is_gateway_dcr_client_id(client_id: str | None) -> bool: + """Cheap prefix routing test so the root endpoints only enter the aggregate arm for + clients this flow registered; every other client_id keeps today's behavior.""" + return client_id is not None and client_id.startswith(GATEWAY_DCR_CLIENT_ID_PREFIX) + + +def _oauth_error(status_code: int, error: str, description: str) -> JSONResponse: + """RFC 6749 section 5.2 / RFC 7591 section 3.2.2 error body. Descriptions carry no + token, code, or URL material so they are safe to relay to any client.""" + return JSONResponse( + status_code=status_code, + content={"error": error, "error_description": description}, + headers=TOKEN_NO_CACHE_HEADERS, + ) + + +def _seal(prefix: str, payload: BaseModel) -> str: + return prefix + encrypt_value_helper(payload.model_dump_json()) + + +_SealedModelT = TypeVar("_SealedModelT", bound=BaseModel) + + +def _open_sealed(value: str, prefix: str, model: type[_SealedModelT], debug_key: str) -> _SealedModelT | None: + """Open a sealed value totally: anything that is not prefix-shaped, does not decrypt, + or does not validate returns ``None`` for the caller to map onto an OAuth error.""" + if not value.startswith(prefix): + return None + decrypted = decrypt_value_helper(value[len(prefix) :], debug_key, return_original_value=False) + if not isinstance(decrypted, str): + return None + try: + return model.model_validate_json(decrypted) + except ValidationError: + return None + + +def open_gateway_dcr_client(client_id: str) -> GatewayDcrClient | None: + return _open_sealed(client_id, GATEWAY_DCR_CLIENT_ID_PREFIX, GatewayDcrClient, _CLIENT_RECORD_DEBUG_KEY) + + +async def register_aggregate_client(request: Request, request_body: Mapping[str, object]) -> Response: + """RFC 7591 dynamic registration against the gateway itself, statelessly. + + Only ``redirect_uris`` is authoritative; every client is registered as a public + ``token_endpoint_auth_method "none"`` client regardless of what it asked for (RFC + 7591 lets the server override metadata), because the gateway never issues client + secrets: possession of a secret would add nothing over the mandatory S256 PKCE, and a + stateless registration has nowhere to keep one. Nothing is persisted, so open + registration cannot be used to fill storage. + + Redirect-URI *hygiene* is not decided here: :func:`validate_redirect_uri_shape` is + the single owner of that rule across the MCP OAuth surface, so allowlisted native + callbacks (``cursor://``) are accepted and fragments, missing hosts, userinfo + (``https://claude.ai@attacker.example/cb``) and backslash hosts are rejected exactly + as they are on /authorize and /callback. + + What this endpoint does decide is its own trust policy, which is deliberately wider + than :func:`validate_trusted_redirect_uri`'s: registration is *public*, so any https + client may register (that is what lets a hosted MCP client register at all), and the + controls are mandatory S256 PKCE plus the consent screen showing the client origin. + http is confined to loopback per RFC 8252 section 7.3. + """ + raw_uris = request_body.get("redirect_uris") + if not isinstance(raw_uris, list) or not raw_uris or len(raw_uris) > MAX_REDIRECT_URIS: + return _oauth_error( + 400, + "invalid_redirect_uri", + f"redirect_uris must be a list of 1 to {MAX_REDIRECT_URIS} URIs", + ) + if not all(isinstance(uri, str) and len(uri) <= MAX_REDIRECT_URI_LENGTH for uri in raw_uris): + return _oauth_error( + 400, + "invalid_redirect_uri", + f"each redirect URI must be a string of at most {MAX_REDIRECT_URI_LENGTH} characters", + ) + for uri in raw_uris: + parsed = urlparse(uri) + try: + if validate_redirect_uri_shape(parsed): + continue # allowlisted native callback, e.g. cursor:// + except HTTPException as exc: + # The shared validator speaks HTTP; RFC 7591 registration answers with an OAuth + # error object, so translate the shape without re-deciding the rule. + return _oauth_error(400, "invalid_redirect_uri", str(exc.detail)) + if parsed.scheme == "https" or (parsed.scheme == "http" and is_loopback_redirect_host(parsed)): + continue + return _oauth_error( + 400, + "invalid_redirect_uri", + "each redirect URI must be https, http on a loopback host, or a registered native callback", + ) + now = datetime.now(timezone.utc) + client_id = _seal( + GATEWAY_DCR_CLIENT_ID_PREFIX, GatewayDcrClient(redirect_uris=tuple(raw_uris), iat=int(now.timestamp())) + ) + if len(client_id) > MAX_CLIENT_ID_LENGTH: + return _oauth_error(400, "invalid_client_metadata", "registered metadata is too large") + return JSONResponse( + status_code=201, + content={ + "client_id": client_id, + "client_id_issued_at": int(now.timestamp()), + "redirect_uris": list(raw_uris), + "token_endpoint_auth_method": "none", + "grant_types": ["authorization_code", "refresh_token"], + "response_types": ["code"], + }, + ) + + +def _flow_cookie_name(handle: str) -> str: + return f"{CONNECT_FLOW_COOKIE_PREFIX}{handle}" + + +def _cookie_path_and_secure(request: Request) -> tuple[str, bool]: + parsed = urlparse(get_request_base_url(request)) + return parsed.path or "/", parsed.scheme == "https" + + +def _append_query_params(url: str, params: dict[str, str]) -> str: + parsed = urlparse(url) + query = parse_qsl(parsed.query, keep_blank_values=True) + list(params.items()) + return urlunparse(parsed._replace(query=urlencode(query))) + + +def relative_request_url(request: Request) -> str: + """The request's own path and query as a same-origin ``return_to`` target for the + login round-trip; relative by construction, so it can never leave the gateway.""" + path = request.url.path + return f"{path}?{request.url.query}" if request.url.query else path + + +def aggregate_authorize( + request: Request, + client_id: str, + redirect_uri: str, + state: str, + code_challenge: str | None, + code_challenge_method: str | None, + response_type: str | None, + session_user_id: str | None, +) -> Response: + """The aggregate authorize verb: validate the client, require S256 PKCE, interpose + LiteLLM sign-in, and hand the browser to the connect page with the flow sealed into a + per-flow cookie. + + Validation failures respond directly with 400 and never redirect: per RFC 6749 + section 4.1.2.1 an unvalidated redirect URI must not receive an error redirect, and + once the client is at fault there is no trusted place to send the browser. + """ + client = open_gateway_dcr_client(client_id) + if client is None: + return _oauth_error(400, "invalid_client", "unknown or malformed client_id") + if redirect_uri not in client.redirect_uris: + return _oauth_error(400, "invalid_request", "redirect_uri is not registered for this client") + if response_type != "code": + return _oauth_error(400, "unsupported_response_type", "response_type must be 'code'") + if not code_challenge or code_challenge_method != "S256": + return _oauth_error( + 400, + "invalid_request", + "PKCE is required: send code_challenge with code_challenge_method=S256", + ) + if len(state) > MAX_STATE_LENGTH: + return _oauth_error(400, "invalid_request", f"state must be at most {MAX_STATE_LENGTH} characters") + base_url = get_request_base_url(request) + if session_user_id is None: + login_url = f"{base_url}/sso/key/generate?{urlencode({'return_to': relative_request_url(request)})}" + return RedirectResponse(login_url, status_code=303) + now = datetime.now(timezone.utc) + handle = secrets.token_urlsafe(24) + flow = _ConnectFlow( + user_id=session_user_id, + client_id=client_id, + redirect_uri=redirect_uri, + state=state, + code_challenge=code_challenge, + jti=secrets.token_urlsafe(24), + exp=int(now.timestamp()) + CONNECT_FLOW_TTL_SECONDS, + ) + connect_url = _append_query_params( + f"{base_url}/ui/chat/integrations", + {"connect_flow": handle, "connect_client": _origin_only(redirect_uri)}, + ) + response = RedirectResponse(connect_url, status_code=303) + path, secure = _cookie_path_and_secure(request) + response.set_cookie( + key=_flow_cookie_name(handle), + value=_seal(_UNPREFIXED, flow), + max_age=CONNECT_FLOW_TTL_SECONDS, + path=path, + secure=secure, + httponly=True, + samesite="lax", + ) + return response + + +def _origin_only(url: str) -> str: + """Scheme+host for display on the connect page; never the full redirect URI, whose + path or query could carry values that do not belong in a page URL or logs.""" + parsed = urlparse(url) + return f"{parsed.scheme}://{parsed.netloc}" if parsed.netloc else "" + + +async def complete_connect_flow( + request: Request, + flow_handle: str, + session_user_id: str | None, + cache: DualCache, +) -> Response: + """The deliberate finish step of the connect flow: mint the gateway authorization + code and send the browser back to the client. + + Reached by POST so a cross-site GET cannot trigger it, and bound to the HttpOnly + per-flow cookie plus an exact match between the signed-in user and the user sealed + into the flow: a link crafted by another party dies here with ``access_denied`` + instead of minting a code for the victim's identity. The flow is single-use (an atomic + claim on its ``jti``), so a double-submit cannot mint two codes from one sign-in. + """ + sealed_flow = request.cookies.get(_flow_cookie_name(flow_handle)) + if sealed_flow is None: + return _oauth_error(400, "invalid_request", "unknown or expired connect flow") + flow = _open_sealed(sealed_flow, _UNPREFIXED, _ConnectFlow, _CONNECT_FLOW_DEBUG_KEY) + if flow is None: + return _oauth_error(400, "invalid_request", "unknown or expired connect flow") + now = datetime.now(timezone.utc) + if now.timestamp() >= flow.exp: + return _oauth_error(400, "invalid_request", "the connect flow has expired; restart the connection") + if session_user_id is None: + return _oauth_error(401, "login_required", "sign in to LiteLLM to finish connecting") + if session_user_id != flow.user_id: + return _oauth_error(403, "access_denied", "the signed-in user does not match this connect flow") + if not await _SingleUseGuard(cache).claim( + f"{_USED_FLOW_CACHE_PREFIX}{flow.jti}", CONNECT_FLOW_TTL_SECONDS + _CLAIM_TTL_BUFFER_SECONDS + ): + return _oauth_error(400, "invalid_request", "this connect flow was already completed; restart the connection") + code = _seal( + GATEWAY_AUTH_CODE_PREFIX, + _GatewayAuthCode( + user_id=flow.user_id, + client_id=flow.client_id, + redirect_uri=flow.redirect_uri, + code_challenge=flow.code_challenge, + jti=secrets.token_urlsafe(24), + iat=int(now.timestamp()), + exp=int(now.timestamp()) + GATEWAY_AUTH_CODE_TTL_SECONDS, + ), + ) + params = {"code": code, **({"state": flow.state} if flow.state else {})} + response = RedirectResponse(_append_query_params(flow.redirect_uri, params), status_code=303) + path, secure = _cookie_path_and_secure(request) + response.delete_cookie(key=_flow_cookie_name(flow_handle), path=path, secure=secure, httponly=True, samesite="lax") + return response + + +def _pkce_verifier_matches(code_verifier: str, code_challenge: str) -> bool: + """RFC 7636 S256 verification, total over hostile input. The comparison is over bytes + so a non-ASCII ``code_challenge`` (which reaches here unvalidated from the client's + authorize request) simply fails to match instead of raising ``TypeError`` the way + ``hmac.compare_digest`` does on two ``str`` with non-ASCII content. The verifier is + ASCII per spec; a compliant client's challenge is base64url and matches.""" + digest = hashlib.sha256(code_verifier.encode("ascii", "replace")).digest() + computed = urlsafe_b64encode(digest).rstrip(b"=") + return hmac.compare_digest(computed, code_challenge.encode("utf-8")) + + +class _SingleUseGuard: + """Atomic single-use claim for a one-time id (an auth-code, connect-flow ``jti``, or refresh-token + ``jti``) over the injected proxy cache. + + Uses an atomic increment rather than a get-then-set: two concurrent redemptions of the same id + cannot both observe "unused", because exactly one increment returns 1. The claim IS the gate, so it + fails closed. Crucially, the increment must be recorded in a backend SHARED across replicas, or the + single-use property is per-worker only (each replica's in-memory counter returns 1, so a captured + id replays through a different worker): + + - When a Redis backend is configured it is the SOLE authority: the claim goes straight to Redis + (``INCR`` is atomic across replicas), and any Redis fault fails the claim CLOSED — it never falls + back to the per-worker in-memory count (``DualCache.async_increment_cache`` does fall back, which + is exactly the replay window this avoids). + - With no Redis configured (single-replica) the in-memory increment is authoritative within the one + process. A multi-worker deployment must run Redis for the guarantee to hold across workers. + + The id's own TTL is the outer bound. For the auth code, PKCE binding is the primary defense against + interception; this makes the RFC 6749 4.1.2 single-use property reliable on top of it.""" + + def __init__(self, cache: DualCache) -> None: + self._cache = cache + + async def claim(self, key: str, ttl_seconds: int) -> bool: + """Atomically claim ``key``. ``True`` iff this caller is the first (increment to 1); ``False`` + on a replay (>1) or when the claim could not be recorded in the shared backend (fail closed).""" + from litellm.proxy.proxy_server import redis_usage_cache # noqa: PLC0415 # circular import at module load + + # Resolve the shared authority HERE rather than trusting the injected cache: callers pass + # user_api_key_cache, which only carries a redis_cache when enable_redis_auth_cache is set + # (off by default), so a guard that read its injected cache silently degraded every claim to + # a per-worker count on a stock multi-worker deployment. redis_usage_cache is the store the + # proxy already treats as cross-worker, so no call site can wire the guarantee away. + redis_cache = redis_usage_cache or getattr(self._cache, "redis_cache", None) + if redis_cache is not None: + # Shared, atomic authority for multi-replica deployments. Claim ONLY against Redis and fail + # CLOSED on any Redis fault (async_increment re-raises) rather than fall back to the + # per-worker in-memory count, which would let each replica observe count==1 and replay the id. + try: + count = await redis_cache.async_increment(key, 1, ttl=ttl_seconds) + except Exception as e: # noqa: BLE001 # ANY Redis fault fails the single-use claim closed + verbose_logger.warning( + "mcp gateway single-use claim: shared cache backend unavailable, failing closed: %s", e + ) + return False + return count == 1 + # No shared backend configured (single-replica): the in-memory increment is authoritative. + count = await self._cache.async_increment_cache(key, 1, ttl=ttl_seconds, local_only=True) + return count == 1 + + +def _session_token_pair(principal: SessionPrincipal, keys: SessionKeys, now: datetime) -> Response: + access = mint_session_token(principal, keys, now) + refresh = mint_session_refresh_token(principal, keys, now) + if not isinstance(access, MintedSessionToken) or not isinstance(refresh, MintedSessionToken): + return _oauth_error(500, "server_error", "failed to mint the session credential") + return JSONResponse( + status_code=200, + content={ + "access_token": access.token.get_secret_value(), + "token_type": "Bearer", + "expires_in": int((access.expires_at - now).total_seconds()), + "refresh_token": refresh.token.get_secret_value(), + }, + headers=TOKEN_NO_CACHE_HEADERS, + ) + + +def _reload_failure_response(failure: ReloadUserFailure) -> Response: + """Map the live-user revalidation failure onto its OAuth error, exhaustively, so a new + ``ReloadUserFailure`` member is a type error here rather than silently 400ing.""" + match failure: + case "unavailable": + return _oauth_error(503, "temporarily_unavailable", "the gateway database is unavailable; retry") + case "unresolvable": + return _oauth_error(500, "server_error", "the gateway is not configured to resolve users") + case "no_active_key": + return _oauth_error(400, "invalid_grant", "the user for this grant is no longer active") + case _: + assert_never(failure) + + +async def aggregate_token( + request: Request, + grant_type: str, + code: str | None, + redirect_uri: str | None, + client_id: str, + code_verifier: str | None, + refresh_token: str | None, + master_key: str | None, + reload_user: ReloadUser, + cache: DualCache, +) -> Response: + """The aggregate token verb: authorization_code and refresh_token grants for the + identity-only session pair. Every path re-validates the litellm user live before + minting, so a deactivated user cannot obtain or renew a session.""" + if master_key is None: + verbose_logger.error("mcp_gateway_dcr token grant rejected: no master_key configured") + return _oauth_error(500, "server_error", "the gateway has no master key configured") + keys = session_keys_from_master_key(master_key) + now = datetime.now(timezone.utc) + if grant_type == "authorization_code": + return await _authorization_code_grant( + code=code, + redirect_uri=redirect_uri, + client_id=client_id, + code_verifier=code_verifier, + keys=keys, + now=now, + reload_user=reload_user, + guard=_SingleUseGuard(cache), + ) + if grant_type == "refresh_token": + return await _refresh_token_grant( + refresh_token=refresh_token, + client_id=client_id, + keys=keys, + now=now, + reload_user=reload_user, + guard=_SingleUseGuard(cache), + ) + return _oauth_error(400, "unsupported_grant_type", "grant_type must be authorization_code or refresh_token") + + +async def _authorization_code_grant( + code: str | None, + redirect_uri: str | None, + client_id: str, + code_verifier: str | None, + keys: SessionKeys, + now: datetime, + reload_user: ReloadUser, + guard: _SingleUseGuard, +) -> Response: + if not code or not redirect_uri or not code_verifier: + return _oauth_error(400, "invalid_request", "code, redirect_uri, and code_verifier are required") + if not MIN_CODE_VERIFIER_LENGTH <= len(code_verifier) <= MAX_CODE_VERIFIER_LENGTH: + return _oauth_error(400, "invalid_request", "code_verifier must be 43 to 128 characters (RFC 7636)") + parsed = _open_sealed(code, GATEWAY_AUTH_CODE_PREFIX, _GatewayAuthCode, _AUTH_CODE_DEBUG_KEY) + if parsed is None: + return _oauth_error(400, "invalid_grant", "the authorization code is invalid") + if now.timestamp() >= parsed.exp: + return _oauth_error(400, "invalid_grant", "the authorization code has expired") + if client_id != parsed.client_id or redirect_uri != parsed.redirect_uri: + return _oauth_error(400, "invalid_grant", "the authorization code was issued to a different client") + if not _pkce_verifier_matches(code_verifier, parsed.code_challenge): + return _oauth_error(400, "invalid_grant", "PKCE verification failed") + # Revalidate the user BEFORE claiming the code, so a transient DB outage (a retryable + # 503) does not consume a still-valid code and force the client to restart sign-in. + failure = await reload_user(parsed.user_id) + if failure is not None: + return _reload_failure_response(failure) + # Atomic single-use claim is the gate: on a concurrent double-redeem exactly one caller + # wins, and a claim that cannot be recorded fails closed. + if not await guard.claim( + f"{_USED_CODE_CACHE_PREFIX}{parsed.jti}", GATEWAY_AUTH_CODE_TTL_SECONDS + _CLAIM_TTL_BUFFER_SECONDS + ): + return _oauth_error(400, "invalid_grant", "the authorization code was already used") + return _session_token_pair(SessionPrincipal(user_id=parsed.user_id, client_id=client_id), keys, now) + + +async def _refresh_token_grant( + refresh_token: str | None, + client_id: str, + keys: SessionKeys, + now: datetime, + reload_user: ReloadUser, + guard: _SingleUseGuard, +) -> Response: + if not refresh_token: + return _oauth_error(400, "invalid_request", "refresh_token is required") + opened = open_session_refresh_bearer(refresh_token, keys, now, expected_client_id=client_id) + if not isinstance(opened, SessionRefreshOpened): + return _oauth_error(400, "invalid_grant", "the refresh token is invalid for this client") + failure = await reload_user(opened.principal.user_id) + if failure is not None: + return _reload_failure_response(failure) + # Refresh-token rotation (OAuth 2.0 Security BCP section 4.13): the presented refresh token is + # single-use. Claim its jti before issuing the replacement pair, so a captured or replayed + # refresh token cannot mint a second pair after the legitimate holder rotated. Claimed AFTER + # user revalidation so a transient DB 503 does not burn a still-valid token; a claim that + # cannot be recorded fails closed, exactly like the authorization-code path. + if not await guard.claim( + f"{_USED_REFRESH_CACHE_PREFIX}{opened.jti}", SESSION_REFRESH_TTL_SECONDS + _CLAIM_TTL_BUFFER_SECONDS + ): + return _oauth_error(400, "invalid_grant", "the refresh token was already used") + return _session_token_pair(opened.principal, keys, now) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 90b70dd01f2..0ee74960293 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -13,6 +13,7 @@ import json import os import re import time +from collections.abc import Sequence from contextlib import asynccontextmanager from typing import Any, AsyncIterator, Callable, Literal, Optional, Union, cast from urllib.parse import urlparse @@ -49,6 +50,10 @@ from litellm.litellm_core_utils.url_utils import SSRFError, async_safe_get from litellm.llms.custom_httpx.http_handler import get_async_httpx_client from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import ( MCPRequestHandler, + _is_mcp_admitted_user_subject, +) +from litellm.proxy._experimental.mcp_server.elicitation_handler import ( + MCP_ELICITATION_AVAILABLE, ) from litellm.proxy._experimental.mcp_server.exceptions import ( MCPServerListError, @@ -59,17 +64,14 @@ from litellm.proxy._experimental.mcp_server.faults.list_outcomes import ( raise_classified_list_failure, upstream_auth_challenge, ) -from litellm.proxy._experimental.mcp_server.elicitation_handler import ( - MCP_ELICITATION_AVAILABLE, -) -from litellm.proxy._experimental.mcp_server.sampling_handler import ( - MCP_SAMPLING_AVAILABLE, -) from litellm.proxy._experimental.mcp_server.oauth2_token_cache import ( MCPPerUserTokenCache, mcp_per_user_token_cache, resolve_mcp_auth, ) +from litellm.proxy._experimental.mcp_server.oauth_utils import ( + _redact_mcp_resource_url, +) from litellm.proxy._experimental.mcp_server.outbound_credentials import ( Error, Ok, @@ -100,6 +102,9 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.types import ( ServerSpec, TokenExchangeConfig, ) +from litellm.proxy._experimental.mcp_server.sampling_handler import ( + MCP_SAMPLING_AVAILABLE, +) from litellm.proxy._experimental.mcp_server.utils import ( MCP_TOOL_PREFIX_SEPARATOR, MCPMissingUserEnvVarsError, @@ -143,11 +148,9 @@ from litellm.types.mcp_server.mcp_server_manager import ( from litellm.types.utils import CallTypes try: - from mcp.shared.tool_name_validation import ( - validate_tool_name, # pyright: ignore[reportAssignmentType] - ) from mcp.shared.tool_name_validation import ( SEP_986_URL, + validate_tool_name, # pyright: ignore[reportAssignmentType] ) except ImportError: from pydantic import BaseModel @@ -408,6 +411,88 @@ def _restrict_discovery_to_corroborated_authorization_server( return metadata.model_copy(update={"token_url": None, "registration_url": None}) +def _redacted_origin_list(urls: Sequence[str]) -> str: + return ", ".join(_redact_mcp_resource_url(url) or "" for url in urls) + + +def _sanitized_error_text(exc: Exception) -> str: + return re.sub(r"https?://\S+", "", str(exc))[:200] + + +def _discovery_failure_leaves_needs_unresolved( + *, + needs_authorization_url: bool, + needs_token_url: bool, + manual_authorization_url: str | None, + manual_token_url: str | None, +) -> bool: + return (needs_authorization_url and not manual_authorization_url) or (needs_token_url and not manual_token_url) + + +def _warn_oauth_endpoints_unresolved( + *, + server_ref: str, + server_url: str | None, + discovery_attempted: bool, + issuer_anchored: bool, + metadata: MCPOAuthMetadata | None, + needs_authorization_url: bool, + needs_token_url: bool, + manual_authorization_url: str | None, + manual_token_url: str | None, +) -> None: + """Log one actionable warning when a server that depends on OAuth endpoint discovery finishes a + build without the endpoints that its flows need (LIT-4658). + + This is the operator-facing signal for a misconfigured server url: discovery failures themselves + are logged where they happen (``_descovery_metadata``), and this names WHICH server is affected, + which endpoints stayed unresolved after manual configuration was considered, and the remedies. + Scopes never trigger the warning on their own: scope-less metadata is normal for many servers and + warning on it every rebuild would be noise. Callers own the per-flow policy of which endpoints + are needed (client_credentials never needs authorization_url; OBO needs only token_url); the + issuer-anchored arm is excluded here because it has its own RFC 8414 §3.3 warning. + """ + if issuer_anchored: + return + unresolved = tuple( + field + for field, needed, value in ( + ( + "authorization_url", + needs_authorization_url, + manual_authorization_url or (metadata.authorization_url if metadata else None), + ), + ( + "token_url", + needs_token_url, + manual_token_url or (metadata.token_url if metadata else None), + ), + ) + if needed and not value + ) + if not unresolved: + return + if discovery_attempted: + verbose_logger.warning( + "MCP server %s: OAuth endpoint discovery left %s unresolved (server url origin: %s). OAuth flows " + "that need them will fail with 'not configured' errors until they resolve. Check the preceding " + "'MCP OAuth' log lines for why discovery failed, verify the configured server url, or set the " + "unresolved endpoint urls manually, or set issuer to discover them from the identity provider " + "(RFC 8414)", + server_ref, + ", ".join(unresolved), + _redact_mcp_resource_url(server_url) or "", + ) + return + verbose_logger.warning( + "MCP server %s uses OAuth but has no discovery source (no server url or pinned issuer), and %s not " + "set manually. Set the missing endpoint urls on the server, or set issuer to discover them from the " + "identity provider (RFC 8414)", + server_ref, + " and ".join(unresolved) + (" is" if len(unresolved) == 1 else " are"), + ) + + def invalidate_user_env_vars_cache(user_id: str, server_id: str) -> None: """Drop a cached entry after the user stores or clears their env var values so the next request reads the fresh value instead of a stale one.""" @@ -884,10 +969,10 @@ def _create_sampling_callback(user_api_key_auth: Optional[Any] = None): return None async def _sampling_callback(context, params): + import litellm from litellm.proxy._experimental.mcp_server.sampling_handler import ( handle_sampling_create_message, ) - import litellm from litellm.proxy._experimental.mcp_server.server import ( get_active_auth_context, ) @@ -1284,6 +1369,15 @@ class MCPServerManager: should_discover = _has_oauth_discovery_source(server_url, use_issuer_anchor) and ( is_discovery_auth_type or obo_needs_discovery ) + config_oauth2_flow = server_config.get("oauth2_flow", None) + needs_authorization_url = is_discovery_auth_type and config_oauth2_flow != "client_credentials" + needs_token_url = is_discovery_auth_type or obo_needs_discovery + warn_on_empty_discovery = _discovery_failure_leaves_needs_unresolved( + needs_authorization_url=needs_authorization_url, + needs_token_url=needs_token_url, + manual_authorization_url=manual_authorization_url, + manual_token_url=manual_token_url, + ) if not should_discover: mcp_oauth_metadata = None elif use_issuer_anchor and manual_issuer is not None: @@ -1292,6 +1386,7 @@ class MCPServerManager: mcp_oauth_metadata = await self._descovery_metadata( server_url=server_url, allow_origin_fallback=is_discovery_auth_type, + warn_when_no_metadata=warn_on_empty_discovery, ) if use_issuer_anchor: @@ -1326,7 +1421,6 @@ class MCPServerManager: ) effective_issuer = manual_issuer or discovered_issuer - config_oauth2_flow = server_config.get("oauth2_flow", None) if auth_type == MCPAuth.oauth2 and config_oauth2_flow not in ( "client_credentials", "authorization_code", @@ -1358,6 +1452,18 @@ class MCPServerManager: "authorization-code flow." ) + _warn_oauth_endpoints_unresolved( + server_ref=server_name or server_id, + server_url=server_url, + discovery_attempted=should_discover, + issuer_anchored=use_issuer_anchor, + metadata=gated_oauth_metadata, + needs_authorization_url=needs_authorization_url, + needs_token_url=needs_token_url, + manual_authorization_url=manual_authorization_url, + manual_token_url=manual_token_url, + ) + new_server = MCPServer( server_id=server_id, name=name_for_prefix, @@ -1485,14 +1591,12 @@ class MCPServerManager: from litellm.proxy._experimental.mcp_server.openapi_to_mcp_generator import ( build_input_schema, create_tool_function, + load_openapi_spec_async, + resolve_operation_params, ) from litellm.proxy._experimental.mcp_server.openapi_to_mcp_generator import ( get_base_url as get_openapi_base_url, ) - from litellm.proxy._experimental.mcp_server.openapi_to_mcp_generator import ( - load_openapi_spec_async, - resolve_operation_params, - ) from litellm.proxy._experimental.mcp_server.tool_registry import ( global_mcp_tool_registry, ) @@ -1681,10 +1785,20 @@ class MCPServerManager: scopes: Optional[list[str]], token_exchange_endpoint: Optional[str], ) -> Optional[MCPOAuthMetadata]: + obo_needs_discovery = self._obo_needs_endpoint_discovery(auth_type, token_exchange_endpoint, manual_token_url) + needs_authorization_url = ( + is_discovery_auth_type and getattr(mcp_server, "oauth2_flow", None) != "client_credentials" + ) + needs_token_url = is_discovery_auth_type or obo_needs_discovery + warn_on_empty_discovery = _discovery_failure_leaves_needs_unresolved( + needs_authorization_url=needs_authorization_url, + needs_token_url=needs_token_url, + manual_authorization_url=manual_authorization_url, + manual_token_url=manual_token_url, + ) has_all_upstream_oauth_fields = bool(manual_authorization_url and manual_token_url and scopes) needs_discovery = _has_oauth_discovery_source(server_url, use_issuer_anchor) and ( - (is_discovery_auth_type and not has_all_upstream_oauth_fields) - or self._obo_needs_endpoint_discovery(auth_type, token_exchange_endpoint, manual_token_url) + (is_discovery_auth_type and not has_all_upstream_oauth_fields) or obo_needs_discovery ) if not needs_discovery: mcp_oauth_metadata: Optional[MCPOAuthMetadata] = None @@ -1694,24 +1808,32 @@ class MCPServerManager: mcp_oauth_metadata = await self._descovery_metadata( server_url=server_url, # type: ignore[arg-type] allow_origin_fallback=is_discovery_auth_type, - ) - if needs_discovery and not use_issuer_anchor and mcp_oauth_metadata is None: - verbose_logger.warning( - "MCP OAuth discovery yielded no metadata for server %s (%s); " - "OAuth endpoints/scopes stay unresolved until a rebuild succeeds", - mcp_server.server_id, - server_url, + warn_when_no_metadata=warn_on_empty_discovery, ) if use_issuer_anchor: return mcp_oauth_metadata - if is_discovery_auth_type: - return _restrict_discovery_to_corroborated_authorization_server( + gated_metadata = ( + _restrict_discovery_to_corroborated_authorization_server( mcp_oauth_metadata, manual_authorization_url, mcp_server.server_id, bool(getattr(mcp_server, "dcr_bridge", None)), ) - return mcp_oauth_metadata + if is_discovery_auth_type + else mcp_oauth_metadata + ) + _warn_oauth_endpoints_unresolved( + server_ref=mcp_server.alias or mcp_server.server_name or mcp_server.server_id, + server_url=server_url, + discovery_attempted=needs_discovery, + issuer_anchored=False, + metadata=gated_metadata, + needs_authorization_url=needs_authorization_url, + needs_token_url=needs_token_url, + manual_authorization_url=manual_authorization_url, + manual_token_url=manual_token_url, + ) + return gated_metadata async def build_mcp_server_from_table( self, @@ -2197,6 +2319,56 @@ class MCPServerManager: return [server_id for server_id in submitted_server_ids if self.get_mcp_server_by_id(server_id) is not None] + async def operator_open_server_ids( + self, + user_api_key_auth: UserAPIKeyAuth | None = None, + *, + allow_all_server_ids: list[str] | None = None, + submitted_server_ids: list[str] | None = None, + ) -> set: + """Servers reachable through OPEN channels rather than a grant: operator-opened + ``allow_all_keys`` servers, plus the caller's own active BYOM submissions when the caller + carries no explicit ``mcp_servers`` scope. + + The single owner of that question for BOTH axes. The server union in + ``get_allowed_mcp_servers`` adds these ids, and the admitted subject's tool resolution asks + the same question to treat an open-channel server as default-open for tools — exactly how a + virtual key experiences it. Encoding the channel membership twice is how a server ends up + listable but uninvokable. + + Empty inside a toolset scope: toolset_mcp_route / dynamic_mcp_route set + ``_mcp_active_toolset_id`` before calling the handler, pinning the request to the toolset's + own servers (checking op.mcp_toolsets==[] instead would false-positive on DB-default rows + where Postgres initialises the column to ARRAY[]::TEXT[]). + + ``allow_all_server_ids`` / ``submitted_server_ids`` are injectable so the server union, + which precomputes both for its fallback path, does not compute them twice.""" + from litellm.proxy._experimental.mcp_server.mcp_context import ( # noqa: PLC0415 + _mcp_active_toolset_id, + ) + + if _mcp_active_toolset_id.get() is not None: + return set() + if allow_all_server_ids is None: + allow_all_server_ids = self.get_allow_all_keys_server_ids() + open_ids = set(allow_all_server_ids) + key_object_permission = user_api_key_auth.object_permission if user_api_key_auth else None + # "Explicitly scoped, so do not widen with BYOM" is a rule about a CREDENTIAL that carries + # its own mcp_servers list. It does not describe a keyless admitted subject: its + # object_permission is the user's own row, whose mcp_servers column is [] by DB default, so + # applying this rule would hide almost every admitted user's OWN submitted servers. Their + # submissions are theirs by authorship, and their scope comes from the per-source union. + has_explicit_object_permission = ( + not _is_mcp_admitted_user_subject(user_api_key_auth) + and key_object_permission is not None + and (key_object_permission.mcp_servers is not None) + ) + if not has_explicit_object_permission: + if submitted_server_ids is None: + submitted_server_ids = await self._get_active_submitted_mcp_server_ids_for_user(user_api_key_auth) + open_ids.update(submitted_server_ids) + return open_ids + async def get_allowed_mcp_servers(self, user_api_key_auth: Optional[UserAPIKeyAuth] = None) -> list[str]: """ Get the allowed MCP Servers for the user. @@ -2210,11 +2382,22 @@ class MCPServerManager: allow_all_server_ids = self.get_allow_all_keys_server_ids() + # A keyless admitted subject is resolved per grant source, and channel decisions that are + # absolute for a scoped KEY credential are not absolute for it: its own opt-out silences its + # own source (handled per source in the resolver), never its teams' grants, and its admin + # role does not swallow the grant model — a session bearer is a third-party client + # credential, not the dashboard, so an admin signing in through the connect flow gets their + # grants like anyone else rather than handing the client the full registry ahead of every + # per-team org ceiling. + is_admitted_subject = _is_mcp_admitted_user_subject(user_api_key_auth) + # The key explicitly opted out of every MCP server. Return zero before # layering on allow_all_keys or submitted servers so the opt-out is absolute. key_object_permission = user_api_key_auth.object_permission if user_api_key_auth else None - if key_object_permission is not None and ( - SpecialMCPServerNames.no_mcp_servers.value in (key_object_permission.mcp_servers or []) + if ( + not is_admitted_subject + and key_object_permission is not None + and (SpecialMCPServerNames.no_mcp_servers.value in (key_object_permission.mcp_servers or [])) ): return [] @@ -2234,8 +2417,14 @@ class MCPServerManager: ) try: - # If admin but NO explicit object permission, get all servers - if user_api_key_auth and _user_has_admin_view(user_api_key_auth) and not has_explicit_object_permission: + # If admin but NO explicit object permission, get all servers (never for an admitted + # subject — see is_admitted_subject above) + if ( + user_api_key_auth + and not is_admitted_subject + and _user_has_admin_view(user_api_key_auth) + and not has_explicit_object_permission + ): verbose_logger.debug("Admin user without explicit object_permission - returning all servers") return list(self.get_registry().keys()) @@ -2243,20 +2432,14 @@ class MCPServerManager: allowed_mcp_servers = await MCPRequestHandler.get_allowed_mcp_servers(user_api_key_auth) verbose_logger.debug(f"Allowed MCP Servers for user api key auth: {allowed_mcp_servers}") combined_servers = set(allowed_mcp_servers) - # Only skip allow_all_keys servers when the request is inside a toolset - # scope. toolset_mcp_route / dynamic_mcp_route set _mcp_active_toolset_id - # before calling the handler — that ContextVar is the reliable signal. - # Using op.mcp_toolsets==[] would false-positive on DB-default rows where - # Postgres initialises the column to ARRAY[]::TEXT[]. - from litellm.proxy._experimental.mcp_server.mcp_context import ( # noqa: PLC0415 - _mcp_active_toolset_id, + combined_servers.update( + await self.operator_open_server_ids( + user_api_key_auth, + allow_all_server_ids=allow_all_server_ids, + submitted_server_ids=submitted_server_ids, + ) ) - in_toolset_scope = _mcp_active_toolset_id.get() is not None - if not in_toolset_scope: - combined_servers.update(allow_all_server_ids) - combined_servers.update(submitted_server_ids) - # For anonymous callers (no user_id, no role), also surface any # servers the operator has opted into upstream-delegated auth. # These servers handle their own auth at the upstream level, so @@ -2903,9 +3086,7 @@ class MCPServerManager: ) ): spec = None - auth_value = ( - await resolve_mcp_auth(server, mcp_auth_header, subject_token=subject_token) if spec is None else None - ) + auth_value = await resolve_mcp_auth(server, mcp_auth_header) if spec is None else None # Create sampling and elicitation callbacks for this client sampling_cb = _create_sampling_callback(user_api_key_auth=user_api_key_auth) if server.allow_sampling else None @@ -3430,6 +3611,7 @@ class MCPServerManager: server_url: str, *, allow_origin_fallback: bool = True, + warn_when_no_metadata: bool = False, ) -> Optional[MCPOAuthMetadata]: """Discover OAuth metadata by following RFC 9728 (protected resource metadata discovery). @@ -3438,8 +3620,32 @@ class MCPServerManager: it (a human sees the redirect), but token_exchange (OBO) sets it False so the gateway never exchanges a subject token against an endpoint it inferred rather than one explicitly configured or authoritatively advertised via RFC 9728 / RFC 8414. - """ + ``warn_when_no_metadata`` makes an all-empty result log one WARNING with the per-step attempt + outcomes (LIT-4658), so a misconfigured server url is diagnosable from default-level logs. The + server loaders set it; the issuer-anchored resource-scopes lookup keeps it off because empty + scopes are not a fault there. + """ + metadata, attempts = await self._discover_metadata_recording_attempts( + server_url, allow_origin_fallback=allow_origin_fallback + ) + if metadata is None and warn_when_no_metadata: + verbose_logger.warning( + "MCP OAuth endpoint discovery against %s found no authorization server metadata. Attempts: %s. " + "The MCP server url may be misconfigured, or the upstream may not support OAuth discovery " + "(RFC 9728 / RFC 8414)", + _redact_mcp_resource_url(server_url) or "", + "; ".join(attempts) if attempts else "none recorded", + ) + return metadata + + async def _discover_metadata_recording_attempts( + self, + server_url: str, + *, + allow_origin_fallback: bool, + ) -> tuple[MCPOAuthMetadata | None, tuple[str, ...]]: + origin = _redact_mcp_resource_url(server_url) or "" try: client = get_async_httpx_client(llm_provider=httpxSpecialProvider.MCP) response = await client.get(server_url) @@ -3452,67 +3658,112 @@ class MCPServerManager: if metadata is None and not resource_scopes and authorization_servers and response.status_code == 200: verbose_logger.warning( "MCP OAuth discovery for %s received 200 OK without RFC 9728 challenge and no discoverable authorization metadata.", - server_url, + origin, ) + attempts = ( + f"GET {origin}: HTTP {response.status_code} (no RFC 9728 challenge)", + *( + ("well-known protected-resource lookup found no authorization servers",) + if not authorization_servers + else () + ), + *( + (f"authorization server metadata fetch failed for: {_redacted_origin_list(authorization_servers)}",) + if authorization_servers and metadata is None + else () + ), + ) if metadata is None and resource_scopes: - return MCPOAuthMetadata(scopes=resource_scopes) + return MCPOAuthMetadata(scopes=resource_scopes), attempts if metadata is not None and resource_scopes: metadata.scopes = resource_scopes - return metadata + return metadata, attempts except HTTPStatusError as exc: - verbose_logger.debug( - "MCP OAuth discovery for %s received status error: %s", - server_url, - exc, - ) - - header_value: Optional[str] = None - if exc.response is not None: - header_value = exc.response.headers.get("WWW-Authenticate") or exc.response.headers.get( - "www-authenticate" - ) - - resource_metadata_url, scopes = self._parse_www_authenticate_header(header_value) - - authorization_servers = [] - resource_scopes = None - if resource_metadata_url: - ( - authorization_servers, - resource_scopes, - ) = await self._fetch_oauth_metadata_from_resource(resource_metadata_url, server_url) - else: - ( - authorization_servers, - resource_scopes, - ) = await self._attempt_well_known_discovery(server_url) - - metadata = None - used_origin_fallback = False - if allow_origin_fallback and not authorization_servers: - try: - parsed_url = urlparse(server_url) - if parsed_url.scheme and parsed_url.netloc: - authorization_servers = [f"{parsed_url.scheme}://{parsed_url.netloc}"] - used_origin_fallback = True - except Exception: - authorization_servers = [] - - if authorization_servers: - metadata = await self._fetch_authorization_server_metadata(authorization_servers, server_url) - if metadata is not None and used_origin_fallback: - metadata.from_origin_fallback = True - - preferred_scopes = scopes or resource_scopes - if metadata is None and preferred_scopes: - metadata = MCPOAuthMetadata(scopes=preferred_scopes) - elif metadata is not None and preferred_scopes: - metadata.scopes = preferred_scopes - - return metadata + return await self._discover_after_status_error(server_url, exc, allow_origin_fallback=allow_origin_fallback) except Exception as exc: # pragma: no cover - network/transient issues verbose_logger.debug("MCP OAuth discovery failed for %s: %s", server_url, exc) - return None + return None, (f"GET {origin}: {type(exc).__name__}: {_sanitized_error_text(exc)}",) + + async def _discover_after_status_error( + self, + server_url: str, + exc: HTTPStatusError, + *, + allow_origin_fallback: bool, + ) -> tuple[MCPOAuthMetadata | None, tuple[str, ...]]: + origin = _redact_mcp_resource_url(server_url) or "" + verbose_logger.debug( + "MCP OAuth discovery for %s received status error: %s", + server_url, + exc, + ) + + header_value: Optional[str] = None + if exc.response is not None: + header_value = exc.response.headers.get("WWW-Authenticate") or exc.response.headers.get("www-authenticate") + status_attempt = ( + f"GET {origin}: HTTP {exc.response.status_code}" + if exc.response is not None + else f"GET {origin}: status error" + ) + + resource_metadata_url, scopes = self._parse_www_authenticate_header(header_value) + + authorization_servers = [] + resource_scopes = None + if resource_metadata_url: + ( + authorization_servers, + resource_scopes, + ) = await self._fetch_oauth_metadata_from_resource(resource_metadata_url, server_url) + lookup_attempt = ( + None + if authorization_servers + else "challenge-advertised resource metadata yielded no authorization servers" + ) + else: + ( + authorization_servers, + resource_scopes, + ) = await self._attempt_well_known_discovery(server_url) + lookup_attempt = ( + None + if authorization_servers + else "no challenge-advertised resource metadata; well-known protected-resource lookup found no authorization servers" + ) + + metadata = None + used_origin_fallback = False + if allow_origin_fallback and not authorization_servers: + try: + parsed_url = urlparse(server_url) + if parsed_url.scheme and parsed_url.netloc: + authorization_servers = [f"{parsed_url.scheme}://{parsed_url.netloc}"] + used_origin_fallback = True + except Exception: + authorization_servers = [] + + fallback_attempt = None + if authorization_servers: + metadata = await self._fetch_authorization_server_metadata(authorization_servers, server_url) + if metadata is not None and used_origin_fallback: + metadata.from_origin_fallback = True + if metadata is None: + fallback_attempt = ( + f"origin fallback: no authorization server metadata at {origin}" + if used_origin_fallback + else f"authorization server metadata fetch failed for: {_redacted_origin_list(authorization_servers)}" + ) + + attempts = tuple(entry for entry in (status_attempt, lookup_attempt, fallback_attempt) if entry) + + preferred_scopes = scopes or resource_scopes + if metadata is None and preferred_scopes: + return MCPOAuthMetadata(scopes=preferred_scopes), attempts + if metadata is not None and preferred_scopes: + metadata.scopes = preferred_scopes + + return metadata, attempts def _parse_www_authenticate_header(self, header_value: Optional[str]) -> tuple[Optional[str], Optional[list[str]]]: if not header_value: diff --git a/litellm/proxy/_experimental/mcp_server/oauth2_token_cache.py b/litellm/proxy/_experimental/mcp_server/oauth2_token_cache.py index 43fe3999291..a6acaf8e1d6 100644 --- a/litellm/proxy/_experimental/mcp_server/oauth2_token_cache.py +++ b/litellm/proxy/_experimental/mcp_server/oauth2_token_cache.py @@ -26,7 +26,6 @@ from litellm.proxy.common_utils.encrypt_decrypt_utils import ( decrypt_value_helper, encrypt_value_helper, ) -from litellm.proxy._experimental.mcp_server.auth import token_exchange from litellm.proxy._experimental.mcp_server.auth.token_endpoint_auth import ( build_token_endpoint_client_auth, ) @@ -58,17 +57,12 @@ class MCPOAuth2TokenCache(InMemoryCache): def _has_client_credentials_config(server: "MCPServer") -> bool: return bool(server.client_id and server.client_secret and server.token_url) - async def async_get_token( - self, - server: "MCPServer", - *, - require_client_credentials_flow: bool = True, - ) -> Optional[str]: + async def async_get_token(self, server: "MCPServer") -> Optional[str]: """Return a valid access token, fetching or refreshing as needed. Returns ``None`` when the server lacks client credentials config. """ - if require_client_credentials_flow and not server.has_client_credentials: + if not server.has_client_credentials: return None if not self._has_client_credentials_config(server): return None @@ -278,36 +272,16 @@ mcp_per_user_token_cache = MCPPerUserTokenCache() async def resolve_mcp_auth( server: "MCPServer", mcp_auth_header: Optional[Union[str, Dict[str, str]]] = None, - subject_token: Optional[str] = None, ) -> Optional[Union[str, Dict[str, str]]]: """Resolve the auth value for an MCP server. Priority: 1. ``mcp_auth_header`` — per-request/per-user override - 2. OAuth2 Token Exchange (OBO / RFC 8693) — exchange user token for scoped token - 3. OAuth2 client_credentials token — auto-fetched and cached - 4. ``server.authentication_token`` — static token from config/DB + 2. OAuth2 client_credentials token — auto-fetched and cached + 3. ``server.authentication_token`` — static token from config/DB """ if mcp_auth_header: return mcp_auth_header - if server.has_token_exchange_config: - if subject_token: - return await token_exchange.mcp_token_exchange_handler.exchange_token(subject_token, server) - # No subject_token — fall back to client_credentials using the same client - # credentials and token_url so M2M scenarios still work. - if server.client_id and server.client_secret and server.token_url: - return await mcp_oauth2_token_cache.async_get_token( - server, - require_client_credentials_flow=False, - ) - # OBO configured but no subject_token and missing client credentials — warn - # rather than silently proceeding unauthenticated. - verbose_logger.warning( - "MCP server '%s' is configured for token exchange (OBO) but no subject_token " - "was provided and client credentials (client_id/client_secret/token_url) are " - "incomplete. The request will proceed without authentication.", - server.server_id, - ) if server.has_client_credentials: return await mcp_oauth2_token_cache.async_get_token(server) return server.authentication_token diff --git a/litellm/proxy/_experimental/mcp_server/oauth_utils.py b/litellm/proxy/_experimental/mcp_server/oauth_utils.py index 53686e329bb..9b7760a30d7 100644 --- a/litellm/proxy/_experimental/mcp_server/oauth_utils.py +++ b/litellm/proxy/_experimental/mcp_server/oauth_utils.py @@ -4,7 +4,7 @@ import os from ipaddress import ip_address from typing import Any, Dict, List, NoReturn, Optional -from urllib.parse import ParseResult, urlparse, urlunparse +from urllib.parse import ParseResult, urlparse, urlsplit, urlunparse, urlunsplit from fastapi import HTTPException, Request @@ -70,6 +70,29 @@ def _origin_label(scheme: str, netloc: str) -> str: return f"{scheme}://{netloc}" if netloc else f"{scheme}://" +def _redact_mcp_resource_url(url: Optional[str]) -> Optional[str]: + """Reduce an MCP server URL to its origin (scheme + host + port) for logging. + + Everything else is dropped: userinfo (``user:pass@``), the query string, the + fragment, and the path, because hosted MCP servers routinely embed the + credential in the path (e.g. ``/mcp/s/``) and this value is persisted + in spend-log metadata that a caller who can invoke the tool can read back. + Returns None when the URL has no host to identify (nothing safe to log). + """ + if not isinstance(url, str) or not url: + return None + try: + parts = urlsplit(url) + hostname = parts.hostname + port = parts.port + except ValueError: + return None + if not hostname: + return None + netloc = f"{hostname}:{port}" if port else hostname + return urlunsplit((parts.scheme, netloc, "", "", "")) or None + + def _resolve_proxy_base_url_env() -> Optional[str]: global _warned_invalid_proxy_base_url configured = os.environ.get("PROXY_BASE_URL", "").strip() @@ -343,8 +366,36 @@ def _parse_redirect_uri_for_validation(redirect_uri: str) -> ParseResult: ) -def _validate_trusted_http_redirect_shape(parsed: ParseResult) -> bool: - """Return True when ``parsed`` is an allowlisted native callback (caller may return).""" +def is_loopback_redirect_host(parsed: ParseResult) -> bool: + """True when the redirect host is loopback (RFC 8252 section 7.3). + + Shared by every redirect-URI policy in the MCP OAuth surface so that none of them + hand-rolls its own host list: a literal ``("localhost", "127.0.0.1", "::1")`` tuple + silently misses the rest of 127.0.0.0/8 and IPv6-mapped forms. + """ + host = (parsed.hostname or "").lower() + if host == "localhost": + return True + try: + return ip_address(host).is_loopback + except ValueError: + return False + + +def validate_redirect_uri_shape(parsed: ParseResult) -> bool: + """Validate redirect-URI *hygiene* and resolve allowlisted native callbacks. + + Returns True when ``parsed`` is an allowlisted native callback (the caller may accept + it outright); returns False for http/https, leaving the trust decision to the caller; + raises for a URI that no policy should ever accept (bad scheme, fragment, missing + host, userinfo, backslash in the host). + + This is deliberately separate from :func:`validate_trusted_redirect_uri`, which adds + the *first-party* trust policy (same-origin, loopback, ops allowlist) appropriate to + the proxy's own OAuth endpoints. Public dynamic-client registration accepts any https + client and relies on PKCE plus the consent screen instead, so it shares this hygiene + rule but not that trust policy. + """ if parsed.scheme not in ("http", "https"): if _matches_trusted_native_redirect_uri(parsed): return True @@ -396,14 +447,8 @@ def _trusted_redirect_uri_is_allowed( ): return True - host = (parsed.hostname or "").lower() - if host == "localhost": + if is_loopback_redirect_host(parsed): return True - try: - if ip_address(host).is_loopback: - return True - except ValueError: - pass if parsed.scheme == "https": for entry in _parse_trusted_redirect_origins(): @@ -522,7 +567,7 @@ def validate_trusted_redirect_uri(request: Request, redirect_uri: str) -> None: :func:`validate_loopback_redirect_uri`. """ parsed = _parse_redirect_uri_for_validation(redirect_uri) - if _validate_trusted_http_redirect_shape(parsed): + if validate_redirect_uri_shape(parsed): return redirect_netloc = _strip_default_port(parsed.scheme, parsed.netloc) proxy_base = _resolve_proxy_base_for_redirect(request) diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/session_credentials.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/session_credentials.py index 08d5cc8b1f1..8844d8c8ad0 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/session_credentials.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/session_credentials.py @@ -149,6 +149,7 @@ class SessionRefreshOpened(BaseModel): model_config = ConfigDict(frozen=True) tag: Literal["opened"] = "opened" principal: SessionPrincipal + jti: str class SessionRefreshInvalid(BaseModel): @@ -187,4 +188,4 @@ def open_session_refresh_bearer( return SessionRefreshInvalid() if opened.principal.client_id != expected_client_id: return SessionRefreshInvalid() - return SessionRefreshOpened(principal=opened.principal) + return SessionRefreshOpened(principal=opened.principal, jti=opened.jti) diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/session_token.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/session_token.py index 9325428f049..4ccbcd1a511 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/session_token.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/session_token.py @@ -113,10 +113,12 @@ class MintedSessionToken(BaseModel): class OpenedSessionToken(BaseModel): - """A validated session token of either kind: the principal it was minted for.""" + """A validated session token of either kind: the principal it was minted for, plus the + ``jti`` so the token endpoint can enforce single-use rotation on a refresh token.""" model_config = ConfigDict(frozen=True) principal: SessionPrincipal + jti: str class SessionTokenTooLarge(BaseModel): @@ -320,7 +322,9 @@ def _open( return SessionMalformed() if now.timestamp() >= claims.exp: return SessionExpired() - return OpenedSessionToken(principal=SessionPrincipal(user_id=claims.user_id, client_id=claims.client_id)) + return OpenedSessionToken( + principal=SessionPrincipal(user_id=claims.user_id, client_id=claims.client_id), jti=claims.jti + ) def _decode_claims( diff --git a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py index 94271c54f4b..26e4176e09b 100644 --- a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py @@ -230,16 +230,33 @@ if MCP_AVAILABLE: return server_auth return mcp_auth_header - def _get_oauth2_server_ids(allowed_server_ids: List[str]) -> Set[str]: - """Return the subset of *allowed_server_ids* whose servers use OAuth2 auth. + def _is_v1_resolved_oauth2_server(server: Optional[MCPServer]) -> bool: + """Whether this server's per-user OAuth2 token is still resolved by v1. - Used as a cheap pre-flight check to skip bulk credential fetching when no - OAuth2 servers are involved in the current request. + A server the v2 resolver owns reads its stored token from the resolver at connect + time and drops any Authorization built for it here, so the v1 lookup would be a DB + round-trip whose result is discarded. Mirrors the same guard on the protocol listing + path and in ``_resolve_oauth2_headers_for_tool_call``. + """ + from litellm.proxy._experimental.mcp_server.outbound_credentials.adapter import ( + to_server_spec, + ) + + if getattr(server, "auth_type", None) != MCPAuth.oauth2: + return False + return to_server_spec(server) is None + + def _v1_resolved_oauth2_server_ids(allowed_server_ids: List[str]) -> Set[str]: + """Return the subset of *allowed_server_ids* whose per-user OAuth2 token is still + resolved by v1. + + Used as a cheap pre-flight check to skip bulk credential fetching when no such + server is involved in the current request. """ return { sid for sid in allowed_server_ids - if getattr(global_mcp_server_manager.get_mcp_server_by_id(sid), "auth_type", None) == MCPAuth.oauth2 + if _is_v1_resolved_oauth2_server(global_mcp_server_manager.get_mcp_server_by_id(sid)) } async def _get_user_oauth_extra_headers( @@ -253,11 +270,13 @@ if MCP_AVAILABLE: the MCP server the same way the admin "Add MCP / Authorize and Fetch" flow does. Returns None for non-OAuth2 servers or when no credential is stored. + A server the v2 resolver owns is skipped; see ``_is_v1_resolved_oauth2_server``. + Args: prefetched_creds: Optional dict keyed by server_id with credential payloads. When provided, avoids a per-server DB round-trip. """ - if getattr(server, "auth_type", None) != MCPAuth.oauth2: + if not _is_v1_resolved_oauth2_server(server): return None user_id = getattr(user_api_key_dict, "user_id", None) server_id = getattr(server, "server_id", None) @@ -320,38 +339,6 @@ if MCP_AVAILABLE: verbose_logger.warning(f"_prefetch_user_oauth_creds: failed to prefetch for user={user_id}: {e}") return {} - async def _get_bulk_user_oauth_headers( - user_api_key_dict: UserAPIKeyAuth, - ) -> Dict[str, Dict[str, str]]: - """ - Fetch ALL OAuth2 credentials for the current user in a single DB query and - return a mapping of server_id → {"Authorization": "Bearer "}. - - This is the batch alternative to calling _get_user_oauth_extra_headers - per-server inside a loop (N+1 DB queries). - """ - user_id = getattr(user_api_key_dict, "user_id", None) - if not user_id: - return {} - try: - from litellm.proxy._experimental.mcp_server.db import ( - list_user_oauth_credentials, - ) - from litellm.proxy.utils import get_prisma_client_or_throw - - prisma_client = get_prisma_client_or_throw( - "Database not connected. Connect a database to use OAuth2 MCP tools." - ) - creds = await list_user_oauth_credentials(prisma_client, user_id) - return { - c["server_id"]: {"Authorization": f"Bearer {c['access_token']}"} - for c in creds - if c.get("access_token") and c.get("server_id") - } - except Exception: - verbose_logger.debug("Failed to bulk-fetch OAuth credentials", exc_info=True) - return {} - def _create_tool_response_objects(tools, server: MCPServer): """Helper function to create tool response objects. @@ -825,7 +812,7 @@ if MCP_AVAILABLE: # to avoid an unnecessary DB round-trip on requests with no OAuth2 MCP servers. prefetched_oauth_creds = ( await _prefetch_user_oauth_creds(user_api_key_dict) - if _get_oauth2_server_ids(allowed_server_ids) + if _v1_resolved_oauth2_server_ids(allowed_server_ids) else {} ) diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 396dd6c7dc7..4fca4406a6f 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -27,7 +27,6 @@ from typing import ( Union, cast, ) -from urllib.parse import urlsplit, urlunsplit import httpx from fastapi import FastAPI, HTTPException @@ -59,6 +58,9 @@ from litellm.proxy._experimental.mcp_server.mcp_context import ( _mcp_gateway_server_name, ) from litellm.proxy._experimental.mcp_server.mcp_debug import MCPDebug +from litellm.proxy._experimental.mcp_server.oauth_utils import ( + _redact_mcp_resource_url, +) from litellm.proxy._experimental.mcp_server.utils import ( LITELLM_MCP_SERVER_DESCRIPTION, LITELLM_MCP_SERVER_NAME, @@ -106,27 +108,6 @@ _MAX_STATEFUL_SESSIONS_PER_OWNER = 100 _MCP_ROUTING_PEEK_MAX_BYTES = 4096 -def _redact_mcp_resource_url(url: Optional[str]) -> Optional[str]: - """Reduce an MCP server URL to its origin (scheme + host + port) for logging. - - Everything else is dropped: userinfo (``user:pass@``), the query string, the - fragment, and the path, because hosted MCP servers routinely embed the - credential in the path (e.g. ``/mcp/s/``) and this value is persisted - in spend-log metadata that a caller who can invoke the tool can read back. - Returns None when the URL has no host to identify (nothing safe to log). - """ - if not isinstance(url, str) or not url: - return None - try: - parts = urlsplit(url) - except ValueError: - return None - if not parts.hostname: - return None - netloc = f"{parts.hostname}:{parts.port}" if parts.port else parts.hostname - return urlunsplit((parts.scheme, netloc, "", "", "")) or None - - def _invalidate_byok_cred_cache(user_id: str, server_id: str) -> None: """Remove a (user_id, server_id) entry from the BYOK credential cache. @@ -977,7 +958,17 @@ if MCP_AVAILABLE: data = await add_litellm_data_to_request( data=body_data, request=request, - user_api_key_dict=user_api_key_auth, + # Bill a team-derived call to the team that granted it. A keyless admitted + # subject carries no team_id, so spend skipped team updates entirely and + # charged the user's PRIMARY org — the granting team's budget never + # accumulated (so it could never begin to block) and, cross-org, the wrong + # organization was charged. This is the ACCOUNTING half; the enforcement + # half (an already-over-budget team stops granting) lives in the source gate. + # Authorization is unaffected: it ran before this, and the union is resolved + # from the untouched auth object passed to call_mcp_tool below. + user_api_key_dict=await MCPRequestHandler.billing_auth_for_tool_call( + user_api_key_auth, tool_name=name + ), proxy_config=proxy_config, ) else: diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index 67d935e34e3..7e8d08e7cad 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -26762,6 +26762,113 @@ "title": "ToolPolicyUpdateResponse", "type": "object" }, + "ToolSpendDailyEntry": { + "description": "Spend attributed to one tool on one UTC day.", + "properties": { + "call_count": { + "default": 0, + "title": "Call Count", + "type": "integer" + }, + "date": { + "title": "Date", + "type": "string" + }, + "spend": { + "default": 0.0, + "title": "Spend", + "type": "number" + }, + "tool_name": { + "title": "Tool Name", + "type": "string" + } + }, + "required": [ + "date", + "tool_name" + ], + "title": "ToolSpendDailyEntry", + "type": "object" + }, + "ToolSpendEntry": { + "description": "Total spend attributed to one tool over the requested window.", + "properties": { + "call_count": { + "default": 0, + "title": "Call Count", + "type": "integer" + }, + "spend": { + "default": 0.0, + "description": "Attributed spend: a request that used several tools counts its full spend toward each of them", + "title": "Spend", + "type": "number" + }, + "tool_name": { + "title": "Tool Name", + "type": "string" + }, + "total_tokens": { + "default": 0, + "title": "Total Tokens", + "type": "integer" + } + }, + "required": [ + "tool_name" + ], + "title": "ToolSpendEntry", + "type": "object" + }, + "ToolSpendResponse": { + "properties": { + "by_tool": { + "items": { + "$ref": "#/components/schemas/ToolSpendEntry" + }, + "title": "By Tool", + "type": "array" + }, + "daily": { + "items": { + "$ref": "#/components/schemas/ToolSpendDailyEntry" + }, + "title": "Daily", + "type": "array" + }, + "end_date": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "End Date" + }, + "start_date": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Start Date" + }, + "total_spend": { + "default": 0.0, + "description": "Deduplicated spend of every request that called at least one tool in the window; less than the sum of per-tool attributed spend whenever multi-tool requests exist", + "title": "Total Spend", + "type": "number" + } + }, + "title": "ToolSpendResponse", + "type": "object" + }, "ToolUsageLogEntry": { "description": "One spend log row for a tool call (for UI \"recent logs\" table).", "properties": { @@ -26858,6 +26965,13 @@ }, "ValidationError": { "properties": { + "ctx": { + "title": "Context", + "type": "object" + }, + "input": { + "title": "Input" + }, "loc": { "items": { "anyOf": [ @@ -27301,6 +27415,81 @@ ] } }, + "/v1/tool/spend": { + "get": { + "description": "Spend attributed to each tool over a date range, for the Cost Optimization dashboard.\n\nJoins ``LiteLLM_SpendLogToolIndex`` (which tool names ran on which request) to\n``LiteLLM_SpendLogs`` (what the request cost). A request that used multiple tools\ncounts its full spend toward each of those tools, so per-tool numbers are\nattributions. ``total_spend`` is the deduplicated spend of every request that\ncalled at least one tool in the window, so it never double counts.", + "operationId": "get_tool_spend_v1_tool_spend_get", + "parameters": [ + { + "description": "YYYY-MM-DD (defaults to 30 days ago)", + "in": "query", + "name": "start_date", + "required": false, + "schema": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "description": "YYYY-MM-DD (defaults to 30 days ago)", + "title": "Start Date" + } + }, + { + "description": "YYYY-MM-DD (defaults to today)", + "in": "query", + "name": "end_date", + "required": false, + "schema": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "description": "YYYY-MM-DD (defaults to today)", + "title": "End Date" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/ToolSpendResponse" + } + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Get Tool Spend", + "tags": [ + "tools" + ] + } + }, "/v1/tool/{tool_name}": { "get": { "description": "Get details for a single tool.", diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 7a7de639d33..0ae93646ca6 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -1868,6 +1868,17 @@ class UpdateTeamRequest(LiteLLMPydanticObjectBase): default_team_member_models: Optional[List[str]] = None # default allowed_models seeded onto new team members +class PatchTeamRequest(UpdateTeamRequest): + """ + Body of PATCH /team/{team_id}. + + Identical to UpdateTeamRequest except team_id is optional, because PATCH takes it + from the path. A team_id in the body is still accepted when it matches the path. + """ + + team_id: str | None = None + + class ResetTeamBudgetRequest(LiteLLMPydanticObjectBase): """ internal type used to reset the budget on a team @@ -2607,6 +2618,28 @@ class UserAPIKeyAuth(LiteLLM_VerificationTokenView): # the expected response ob user_max_budget: Optional[float] = None request_route: Optional[str] = None is_session_token: bool = False + # Server-only marker set exclusively by the MCP gateway admission path + # (_reload_admitted_user) for a keyless user-subject admitted via a gateway DCR session + # bearer or bridge envelope. Not a DB column and never populated from caller-controlled key + # metadata or JWT claims, so it cannot be forged to gain the team-inherited MCP grant union + # or to escape the caller-Authorization egress scrub. exclude=True keeps it out of serialization. + mcp_admitted_user_subject: bool = Field(default=False, exclude=True) + # team_id -> that team's mcp_rpm_limit map, for a keyless admitted subject that reaches MCP + # servers through several teams at once and therefore has no single team_id for the limiter to + # key off. Server-only and stripped from validated input for the same reason as the marker + # above: a forged entry would let a caller pick which team's rpm bucket it is charged against. + mcp_source_team_rpm_limits: dict[str, dict[str, int]] | None = Field(default=None, exclude=True) + via_virtual_key: bool = Field( + default=False, + exclude=True, + description=( + "Server-only marker set exclusively by the DB virtual-key and master-key auth paths via " + "post-construction assignment. Stripped from validated input so custom auth handlers, JWT " + "claims, or key metadata cannot forge it. Gates overwrite_user_with_key_hash stamping: only " + "a credential the proxy itself validated as a key may be forwarded as the provider-facing " + "user id." + ), + ) budget_reservation: Optional[Dict[str, Any]] = Field(default=None, exclude=True) budget_throttle_pct: Optional[float] = Field(default=None, exclude=True) user: Optional[Any] = None # Expanded user object when expand=user is used @@ -2627,6 +2660,12 @@ class UserAPIKeyAuth(LiteLLM_VerificationTokenView): # the expected response ob # If values is already an instance (not a dict), return it as-is if not isinstance(values, dict): return values + # mcp_admitted_user_subject is a server-only marker, set ONLY by the MCP gateway admission + # path via post-construction assignment. Strip it from any validated input (constructor + # kwargs, model_validate, a JWT/key claim splat) so it can never be forged from caller data. + values.pop("mcp_admitted_user_subject", None) + values.pop("mcp_source_team_rpm_limits", None) + values.pop("via_virtual_key", None) if values.get("api_key") is not None: values.update({"token": cls._safe_hash_litellm_api_key(values.get("api_key"))}) if isinstance(values.get("api_key"), str): @@ -2781,6 +2820,31 @@ class LiteLLM_OrganizationTableUpdate(LiteLLM_BudgetTable): return values +class OrganizationUpdateRequestV2(LiteLLMPydanticObjectBase): + """ + Typed PATCH body for ``/v2/organization/{organization_id}`` (RFC 7396 merge-patch). + + Presence is read from ``model_fields_set``, so a sent field is written and an omitted one is + left untouched. ``extra="forbid"`` makes an unknown key a 422 rather than a silent no-op, since + the contract hinges on which keys are present. See the endpoint for the per-field clear tokens. + """ + + model_config = ConfigDict(extra="forbid") + + organization_alias: str | None = None + models: list[str] | None = None + metadata: dict | None = None + tpm_limit: int | None = None + rpm_limit: int | None = None + max_budget: float | None = None + soft_budget: float | None = None + max_parallel_requests: int | None = None + model_max_budget: dict | None = None + budget_duration: str | None = None + logging_exporters: list[str] | None = None + object_permission: LiteLLM_ObjectPermissionBase | None = None + + from litellm.models.organization import ( # noqa: E402 LiteLLM_OrganizationTable as LiteLLM_OrganizationTable, ) diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 99a867a5d07..ce82ca74267 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -2661,6 +2661,15 @@ async def get_managed_vector_store_rows_by_uuids( return result +class OrganizationNotFoundError(Exception): + """The organization row is CONFIRMED absent, as opposed to a lookup that failed. + + Subclasses Exception so every existing except Exception caller keeps its current + behavior; it exists so a caller that wants to treat "no such org" as "no restriction" can do + that WITHOUT also swallowing an outage and silently dropping a real org ceiling. + """ + + @log_db_metrics async def get_org_object( org_id: str, @@ -2707,25 +2716,30 @@ async def get_org_object( query_kwargs["include"] = {"litellm_budget_table": True} response = await OrganizationRepository(prisma_client).table.find_unique(**query_kwargs) - - if response is None: - raise Exception - - _org_obj = LiteLLM_OrganizationTable(**response.model_dump()) - # Cache the result - await user_api_key_cache.async_set_cache( - key=cache_key, - value=_org_obj, - model_type=LiteLLM_OrganizationTable, - ttl=DEFAULT_IN_MEMORY_TTL, - ) - - return _org_obj except Exception: - raise Exception( + # An operational failure (DB down, timeout, cache fault) is NOT the same fact as a confirmed + # missing row, and relabelling it as "doesn't exist" made every caller unable to tell them + # apart — a caller that treats absence as "this org places no restriction" then drops a real + # org ceiling during an outage. Propagate the real error; callers that already catch + # Exception are unaffected. + raise + + if response is None: + raise OrganizationNotFoundError( f"Organization doesn't exist in db. Organization={org_id}. Create organization via `/organization/new` call." ) + _org_obj = LiteLLM_OrganizationTable(**response.model_dump()) + # Cache the result + await user_api_key_cache.async_set_cache( + key=cache_key, + value=_org_obj, + model_type=LiteLLM_OrganizationTable, + ttl=DEFAULT_IN_MEMORY_TTL, + ) + + return _org_obj + async def _get_resources_from_access_groups( access_group_ids: List[str], diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index 293bb74e211..ecb37e67c14 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -3,7 +3,7 @@ import re import sys from functools import lru_cache from logging import Logger -from typing import Any, Dict, FrozenSet, List, Mapping, Optional, Tuple, Union +from typing import Any, Dict, FrozenSet, Iterator, List, Mapping, Optional, Tuple, Union from fastapi import HTTPException, Request, status @@ -12,7 +12,12 @@ from litellm import Router, provider_list from litellm._logging import verbose_proxy_logger from litellm.constants import MINIMUM_CUSTOM_KEY_LENGTH, STANDARD_CUSTOMER_ID_HEADERS from litellm.litellm_core_utils.safe_json_loads import safe_json_loads -from litellm.litellm_core_utils.url_utils import SSRFError, validate_url +from litellm.litellm_core_utils.url_utils import ( + SSRFError, + is_url_destination_allowed_by_host, + provider_url_destination_candidates, + validate_url, +) from litellm.proxy._types import * from litellm.types.passthrough_endpoints.pass_through_endpoints import ( LITELLM_PASS_THROUGH_ENDPOINT_MARKER, @@ -290,6 +295,7 @@ _BANNED_REQUEST_BODY_PARAMS: Tuple[str, ...] = ( "use_ssl", # SDK-only field; also rejected outright in is_request_body_safe. "model_list", + "vertex_ai_credentials", # Observability credentials, hosts, and project identifiers: derived # from the canonical ``_supported_callback_params`` allowlist so new # integrations are covered automatically. Sorted for stable iteration @@ -342,6 +348,60 @@ def _check_banned_params( ) +_FALLBACK_FIELDS: tuple[str, ...] = ( + "fallbacks", + "context_window_fallbacks", + "content_policy_fallbacks", +) + + +def _iter_fallback_field_values(request_body: Mapping[str, object]) -> Iterator[object]: + override = request_body.get("router_settings_override") + for source in (request_body, override): + if isinstance(source, Mapping): + for field in _FALLBACK_FIELDS: + yield source.get(field) + + +def _iter_fallback_targets(value: object, depth: int) -> Iterator[str | Mapping[str, object]]: + if depth > 2 * litellm.ROUTER_MAX_FALLBACKS: + raise ValueError("Rejected Request: fallback nesting exceeds the allowed validation depth.") + if not isinstance(value, list): + return + for item in value: + if isinstance(item, str): + yield item + elif isinstance(item, Mapping): + values = tuple(item.values()) + if not (values and all(isinstance(v, list) for v in values)): + yield item + if isinstance(item.get("model"), str): + for field in _FALLBACK_FIELDS: + yield from _iter_fallback_targets(item.get(field), depth + 1) + else: + for target_list in values: + yield from _iter_fallback_targets(target_list, depth + 1) + + +def iter_request_fallback_targets(request_body: Mapping[str, object]) -> Iterator[str | Mapping[str, object]]: + for value in _iter_fallback_field_values(request_body): + yield from _iter_fallback_targets(value, 0) + + +def _reject_url_valued_fallback_target(value: str) -> None: + allowed_hosts = getattr(litellm, "provider_url_destination_allowed_hosts", []) or [] + for candidate in provider_url_destination_candidates(value): + if not candidate.lower().startswith(("http://", "https://")): + continue + if is_url_destination_allowed_by_host(candidate, allowed_hosts): + continue + raise ValueError( + f"Rejected Request: URL-valued fallback destination '{value}' is not allowed. " + "Configure custom endpoints with api_base instead, or add the destination host to " + "`provider_url_destination_allowed_hosts` in litellm_settings." + ) + + def is_request_body_safe(request_body: dict, general_settings: dict, llm_router: Optional[Router], model: str) -> bool: """ Check if the request body is safe. @@ -379,6 +439,14 @@ def is_request_body_safe(request_body: dict, general_settings: dict, llm_router: metadata = _coerce_metadata_to_dict(request_body.get(metadata_key)) if metadata is not None: _check_banned_params(metadata, general_settings, llm_router, model) + for target in iter_request_fallback_targets(request_body): + if isinstance(target, dict): + _check_banned_params(target, general_settings, llm_router, model) + target_model = target.get("model") + if isinstance(target_model, str): + _reject_url_valued_fallback_target(target_model) + elif isinstance(target, str): + _reject_url_valued_fallback_target(target) litellm_params = _coerce_metadata_to_dict(request_body.get("litellm_params")) if litellm_params is not None: litellm_params_metadata = _coerce_metadata_to_dict(litellm_params.get("metadata")) diff --git a/litellm/proxy/auth/login_utils.py b/litellm/proxy/auth/login_utils.py index 11f12e597b9..f35d94c986e 100644 --- a/litellm/proxy/auth/login_utils.py +++ b/litellm/proxy/auth/login_utils.py @@ -7,12 +7,15 @@ login endpoints (e.g., /login and /v2/login). import os import secrets +from datetime import datetime, timedelta, timezone from typing import Literal, Optional, cast +import jwt from fastapi import HTTPException import litellm from litellm.constants import LITELLM_PROXY_ADMIN_NAME, LITELLM_UI_SESSION_DURATION +from litellm.litellm_core_utils.duration_parser import duration_in_seconds from litellm.proxy._types import ( LiteLLM_UserTable, LitellmUserRoles, @@ -313,6 +316,29 @@ async def authenticate_user( ) +def _ui_session_exp_timestamp() -> int: + """The ``exp`` claim (unix seconds) for a UI session cookie, ``LITELLM_UI_SESSION_DURATION`` + from now. The virtual key sealed inside the cookie already expires after this same + duration; stamping the JWT itself gives the cookie the bounded lifetime the dashboard's + client-side expiry check and the server-side session-cookie readers both assume, instead + of a token that stays signature-valid until the master key rotates.""" + ttl_seconds = duration_in_seconds(LITELLM_UI_SESSION_DURATION) + return int((datetime.now(timezone.utc) + timedelta(seconds=ttl_seconds)).timestamp()) + + +def encode_ui_session_jwt(returned_ui_token_object: ReturnedUITokenObject, master_key: str) -> str: + """Encode a UI session cookie JWT with a bounded ``exp``. + + The single choke point every UI login path (SSO and username/password /login, /v2, + /v3) uses to mint the ``token`` cookie, so the cookie's lifetime is set in exactly one + place and cannot drift between paths. Without the ``exp`` the cookie is valid until the + master key rotates, and the session-cookie readers that require a bounded lifetime + (the MCP interactive sign-in) reject it. + """ + claims = {**cast(dict, returned_ui_token_object), "exp": _ui_session_exp_timestamp()} + return jwt.encode(claims, master_key, algorithm="HS256") + + def create_ui_token_object( login_result: LoginResult, general_settings: dict, diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 1af1c05a47a..b0e59e3d6a8 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -14,7 +14,7 @@ import secrets import orjson from datetime import datetime, timezone -from typing import Any, Dict, Iterator, NamedTuple, List, Optional, Protocol, Tuple, Union, cast +from typing import Any, Dict, NamedTuple, List, Optional, Protocol, Tuple, Union, cast import fastapi from fastapi import HTTPException, Request, WebSocket, status @@ -58,6 +58,7 @@ from litellm.proxy.auth.auth_utils import ( get_model_from_request, get_request_route, get_request_route_template, + iter_request_fallback_targets, normalize_request_route, pre_db_read_auth_checks, route_in_additonal_public_routes, @@ -1538,6 +1539,13 @@ async def _user_api_key_auth_builder( check_cache_only=True, ).resolve(hashed_token=hash_token(api_key)) ) + # Key-cache entries are written only after the proxy validated a + # virtual key or the master key, but via_virtual_key is exclude=True + # so serialization drops it; restore it at this trusted boundary. + # The UI-login JWT fallback below constructs its token from a + # decrypted blob, not this cache, and stays unmarked. + if isinstance(valid_token, UserAPIKeyAuth): + valid_token.via_virtual_key = True except Exception: verbose_logger.debug("api key not found in cache.") valid_token = None @@ -1655,6 +1663,7 @@ async def _user_api_key_auth_builder( _user_api_key_obj = update_valid_token_with_end_user_params( valid_token=_user_api_key_obj, end_user_params=end_user_params ) + _user_api_key_obj.via_virtual_key = True return _user_api_key_obj @@ -2064,7 +2073,7 @@ async def _user_api_key_auth_builder( # No token was found when looking up in the DB raise Exception("Invalid proxy server token passed") if valid_token_dict is not None: - return await _return_user_api_key_auth_obj( + virtual_key_auth_obj = await _return_user_api_key_auth_obj( user_obj=user_obj, api_key=api_key, parent_otel_span=parent_otel_span, @@ -2072,6 +2081,8 @@ async def _user_api_key_auth_builder( route=route, start_time=start_time, ) + virtual_key_auth_obj.via_virtual_key = True + return virtual_key_auth_obj except Exception as e: return await UserAPIKeyAuthExceptionHandler._handle_authentication_error( e=e, @@ -2485,6 +2496,7 @@ async def _reserve_budget_after_common_checks( end_user_id=end_user_id, end_user_object=end_user_object, skip_user_budget_on_team_key=general_settings.get("skip_user_budget_on_team_key") is True, + fail_closed_budget_enforcement=general_settings.get("fail_closed_budget_enforcement") is True, ) @@ -2854,19 +2866,11 @@ async def _enforce_key_and_fallback_model_access( llm_router=llm_router, ) - # Validate every fallback model name reachable by this request. - # All three fields (``fallbacks``, ``context_window_fallbacks``, - # ``content_policy_fallbacks``) are forwarded to the router as - # per-request kwargs whether they appear at the top level of - # ``request_data`` or nested under ``router_settings_override``. - # Both surfaces must be validated against the API key's model - # allowlist or a caller can smuggle a restricted model. VERIA-44. - fallback_names: List[str] = [] - override_settings = request_data.get("router_settings_override") - for _fb_key in ROUTER_FALLBACK_FIELDS: - fallback_names.extend(iter_router_fallback_model_names(request_data.get(_fb_key))) - if isinstance(override_settings, dict): - fallback_names.extend(iter_router_fallback_model_names(override_settings.get(_fb_key))) + fallback_names = tuple( + name + for target in iter_request_fallback_targets(request_data) + if (name := _fallback_target_model_name(target)) is not None + ) for _name in dict.fromkeys(fallback_names): # dedupe, preserve order await can_key_call_model( @@ -2882,36 +2886,14 @@ async def _enforce_key_and_fallback_model_access( ) -ROUTER_FALLBACK_FIELDS: Tuple[str, ...] = ( - "fallbacks", - "context_window_fallbacks", - "content_policy_fallbacks", -) - - -def iter_router_fallback_model_names(fallbacks: Any) -> Iterator[str]: - """Yield leaf model names from any of the supported fallbacks shapes. - - Handles the simple top-level shape (``str`` or ``{"model": str}``) and - the nested router-config shape (``[{primary: [fallback_list]}]``). - """ - if not isinstance(fallbacks, list): - return - for entry in fallbacks: - if isinstance(entry, str): - yield entry - elif isinstance(entry, dict): - if isinstance(entry.get("model"), str): - yield entry["model"] - continue - for fallback_list in entry.values(): - if not isinstance(fallback_list, list): - continue - for m in fallback_list: - if isinstance(m, str): - yield m - elif isinstance(m, dict) and isinstance(m.get("model"), str): - yield m["model"] +def _fallback_target_model_name(target: object) -> str | None: + if isinstance(target, str): + return target + if isinstance(target, dict): + model = target.get("model") + if isinstance(model, str): + return model + return None async def _run_post_custom_auth_checks( diff --git a/litellm/proxy/client/cli/README.md b/litellm/proxy/client/cli/README.md index 84bc27ef0d4..2ad8a08b8c3 100644 --- a/litellm/proxy/client/cli/README.md +++ b/litellm/proxy/client/cli/README.md @@ -523,7 +523,7 @@ export LITELLM_PROXY_API_KEY=sk-... lite model-groups list [--format table|json] ``` -Lists the model groups your key can reach on the proxy, via `/model_group/info`, along with each group's mode (`chat`, `embedding`, etc.) and per-token pricing. This is also what `lite autoroute configure` uses internally to discover what it can offer you. +Lists the model groups your key can reach on the proxy, via `/model_group/info`, along with each group's mode (`chat`, `embedding`, etc.) and per-token pricing. Note this route needs management access; `lite autoroute configure` instead discovers models through `/v1/models`, so it works with a key scoped to just the AI API routes #### Configure the Auto-Router diff --git a/litellm/proxy/client/cli/commands/autoroute/config.py b/litellm/proxy/client/cli/commands/autoroute/config.py index 603cea38f6f..237705564ff 100644 --- a/litellm/proxy/client/cli/commands/autoroute/config.py +++ b/litellm/proxy/client/cli/commands/autoroute/config.py @@ -15,41 +15,25 @@ class DiscoveredModel(BaseModel): name: str mode: str = "chat" - input_cost_per_token: float | None = None - output_cost_per_token: float | None = None -class _RawModelGroup(BaseModel): +class _RawModelListing(BaseModel): model_config = ConfigDict(extra="ignore") - model_group: str - # Optional: some real deployments return an explicit `"mode": null` for models that - # were registered without a mode (seen for embedding models like voyage-4-large). - # ModelGroupInfo's own "chat" default (litellm/types/router.py) only applies when the - # key is missing entirely, not when it's present as null, so this must tolerate None. - mode: str | None = "chat" - input_cost_per_token: float | None = None - output_cost_per_token: float | None = None + id: str + # /v1/models attaches "mode" (sourced from the cost map) only for models it can resolve; + # a model whose mode is unknown arrives without the field, so default it to chat rather + # than dropping it, which keeps it selectable as a routing target in the wizard. + mode: str = "chat" -_RAW_MODEL_GROUPS_ADAPTER = TypeAdapter(list[_RawModelGroup]) +_RAW_MODEL_LISTING_ADAPTER = TypeAdapter(list[_RawModelListing]) def parse_discovered_models(raw: list[JsonValue]) -> tuple[DiscoveredModel, ...]: - """Validate a raw `/model_group/info` response into typed models.""" - parsed = _RAW_MODEL_GROUPS_ADAPTER.validate_python(raw) - return tuple( - DiscoveredModel( - name=group.model_group, - # A null mode means the server genuinely doesn't know what this model does; - # "unknown" (rather than guessing "chat") keeps it out of both chat_models() - # and embedding_models() instead of risking a wrong-mode deployment. - mode=group.mode or "unknown", - input_cost_per_token=group.input_cost_per_token, - output_cost_per_token=group.output_cost_per_token, - ) - for group in parsed - ) + """Validate a raw `/v1/models` response into typed models.""" + parsed = _RAW_MODEL_LISTING_ADAPTER.validate_python(raw) + return tuple(DiscoveredModel(name=item.id, mode=item.mode) for item in parsed) def chat_models(models: tuple[DiscoveredModel, ...]) -> tuple[DiscoveredModel, ...]: diff --git a/litellm/proxy/client/cli/commands/autoroute/wizard.py b/litellm/proxy/client/cli/commands/autoroute/wizard.py index 8ad87315fb9..d3fe458d9f0 100644 --- a/litellm/proxy/client/cli/commands/autoroute/wizard.py +++ b/litellm/proxy/client/cli/commands/autoroute/wizard.py @@ -111,12 +111,12 @@ def run_configure_wizard(ctx: click.Context) -> Path: api_key = ctx.obj["api_key"] client = Client(base_url=base_url, api_key=api_key) - raw_groups = client.model_groups.info() - if not isinstance(raw_groups, list): + raw_models = client.models.list() + if not isinstance(raw_models, list): raise click.ClickException( - f"Unexpected response from /model_group/info: expected a list, got {type(raw_groups).__name__}" + f"Unexpected response from /v1/models: expected a list, got {type(raw_models).__name__}" ) - discovered = parse_discovered_models(raw_groups) + discovered = parse_discovered_models(raw_models) chat_pool = chat_models(discovered) embedding_pool = embedding_models(discovered) diff --git a/litellm/proxy/common_utils/callback_utils.py b/litellm/proxy/common_utils/callback_utils.py index c644ecc3dae..a9c2a12aff7 100644 --- a/litellm/proxy/common_utils/callback_utils.py +++ b/litellm/proxy/common_utils/callback_utils.py @@ -1,4 +1,5 @@ import copy +import os from typing import TYPE_CHECKING, Any, Callable, Dict, Iterable, List, Literal, Optional import litellm @@ -564,11 +565,8 @@ def process_callback(_callback: str, callback_type: str, environment_variables: env_vars_dict: dict[str, str | None] = {} for _var in env_vars: - env_variable = environment_variables.get(_var, None) - if env_variable is None: - env_vars_dict[_var] = None - else: - env_vars_dict[_var] = env_variable + stored_value = environment_variables.get(_var, None) + env_vars_dict[_var] = stored_value if stored_value is not None else os.getenv(_var) return {"name": _callback, "variables": env_vars_dict, "type": callback_type} diff --git a/litellm/proxy/config_resolvers/__init__.py b/litellm/proxy/config_resolvers/__init__.py new file mode 100644 index 00000000000..88b4c3961f0 --- /dev/null +++ b/litellm/proxy/config_resolvers/__init__.py @@ -0,0 +1,9 @@ +"""Typed, provenance-aware resolution of proxy settings from DB then env.""" + +from litellm.proxy.config_resolvers._descriptors import ( + FieldDescriptor, + FieldSource, + resolve_fields, +) + +__all__ = ["FieldDescriptor", "FieldSource", "resolve_fields"] diff --git a/litellm/proxy/config_resolvers/_descriptors.py b/litellm/proxy/config_resolvers/_descriptors.py new file mode 100644 index 00000000000..f67a690f92f --- /dev/null +++ b/litellm/proxy/config_resolvers/_descriptors.py @@ -0,0 +1,73 @@ +"""Shared primitive for resolving a settings value from its sources. + +A ``FieldDescriptor`` names, for one setting, where it lives in the stored DB +row (``db_key``), which process env var carries it (``env_var``), whether it is +a secret, and its effective default. ``resolve_fields`` reconciles a set of +descriptors against a decrypted DB row and the process environment with a fixed +precedence, returning the resolved values plus per-field provenance so a caller +can tell whether a value came from the database, the environment, a default, or +is unset. +""" + +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from typing import Literal + +FieldSource = Literal["db", "env", "default", "unset"] + + +@dataclass(frozen=True, slots=True) +class FieldDescriptor: + field_name: str + db_key: str + env_var: str + is_secret: bool = False + default: str | None = None + + +def _db_is_set(db_value: object, empty_db_is_set: bool) -> bool: + if empty_db_is_set: + # A stored key that is present, even as "", is an explicit admin choice + # (e.g. clearing an alerting webhook) and must win over a stale env var. + return db_value is not None + # A blank stored value is treated as absent, so it falls through to env. This + # fits settings whose clear path also unsets the env var (e.g. SSO). + return isinstance(db_value, str) and bool(db_value.strip()) + + +def _resolve_one( + descriptor: FieldDescriptor, + db_values: Mapping[str, object], + env: Mapping[str, str], + empty_db_is_set: bool, +) -> tuple[str, str | None, FieldSource]: + db_value = db_values.get(descriptor.db_key) + if _db_is_set(db_value, empty_db_is_set): + return descriptor.field_name, db_value if isinstance(db_value, str) else str(db_value), "db" + env_value = env.get(descriptor.env_var) + if isinstance(env_value, str) and env_value.strip(): + return descriptor.field_name, env_value, "env" + if descriptor.default is not None: + return descriptor.field_name, descriptor.default, "default" + return descriptor.field_name, None, "unset" + + +def resolve_fields( + descriptors: Sequence[FieldDescriptor], + db_values: Mapping[str, object], + env: Mapping[str, str], + empty_db_is_set: bool = False, +) -> tuple[dict[str, str | None], dict[str, FieldSource]]: + """Resolve every descriptor to (values, provenance). + + Precedence per field: a set stored value wins, else a non-blank process env + var, else the descriptor default, else unset. ``empty_db_is_set`` selects + how a present-but-empty stored value is read: ``False`` treats it as absent + so it falls back to env (SSO, whose clear path also unsets the env var); + ``True`` treats it as an explicit clear that wins over env (alerting, whose + clear path stores "" without unsetting the env var). + """ + resolved = tuple(_resolve_one(descriptor, db_values, env, empty_db_is_set) for descriptor in descriptors) + values = {field_name: value for field_name, value, _ in resolved} + provenance = {field_name: source for field_name, _, source in resolved} + return values, provenance diff --git a/litellm/proxy/config_resolvers/alerting.py b/litellm/proxy/config_resolvers/alerting.py new file mode 100644 index 00000000000..3704ec09355 --- /dev/null +++ b/litellm/proxy/config_resolvers/alerting.py @@ -0,0 +1,25 @@ +"""Descriptor tables for the alerting settings surfaced by /get/config/callbacks. + +These reconcile the stored ``environment_variables`` blob (keyed by the +uppercase env-var names) with the process environment. SMTP_PORT and SMTP_TLS +carry the same effective defaults the mail-send path applies, so the settings +page shows the config that mail would actually use rather than a blank. +""" + +from litellm.proxy.config_resolvers._descriptors import FieldDescriptor + +EMAIL_DESCRIPTORS: tuple[FieldDescriptor, ...] = ( + FieldDescriptor("SMTP_HOST", "SMTP_HOST", "SMTP_HOST"), + FieldDescriptor("SMTP_PORT", "SMTP_PORT", "SMTP_PORT", default="587"), + FieldDescriptor("SMTP_TLS", "SMTP_TLS", "SMTP_TLS", default="True"), + FieldDescriptor("SMTP_USERNAME", "SMTP_USERNAME", "SMTP_USERNAME", is_secret=True), + FieldDescriptor("SMTP_PASSWORD", "SMTP_PASSWORD", "SMTP_PASSWORD", is_secret=True), + FieldDescriptor("SMTP_SENDER_EMAIL", "SMTP_SENDER_EMAIL", "SMTP_SENDER_EMAIL"), + FieldDescriptor("TEST_EMAIL_ADDRESS", "TEST_EMAIL_ADDRESS", "TEST_EMAIL_ADDRESS"), + FieldDescriptor("EMAIL_LOGO_URL", "EMAIL_LOGO_URL", "EMAIL_LOGO_URL"), + FieldDescriptor("EMAIL_SUPPORT_CONTACT", "EMAIL_SUPPORT_CONTACT", "EMAIL_SUPPORT_CONTACT"), +) + +SLACK_DESCRIPTORS: tuple[FieldDescriptor, ...] = ( + FieldDescriptor("SLACK_WEBHOOK_URL", "SLACK_WEBHOOK_URL", "SLACK_WEBHOOK_URL", is_secret=True), +) diff --git a/litellm/proxy/config_resolvers/sso.py b/litellm/proxy/config_resolvers/sso.py new file mode 100644 index 00000000000..97c42106018 --- /dev/null +++ b/litellm/proxy/config_resolvers/sso.py @@ -0,0 +1,98 @@ +"""Resolved SSO config object. + +Reconciles the dedicated ``sso_config`` DB row (lowercase, per-value encrypted +keys) with the process environment (uppercase env vars) into a typed +``SSOConfig`` plus per-field provenance. This is the single source of truth for +the SSO field -> env-var mapping, used by both the read-back endpoint and the +save endpoint so the two can never drift. +""" + +from collections.abc import Mapping +from dataclasses import dataclass + +from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value_helper +from litellm.proxy.config_resolvers._descriptors import ( + FieldDescriptor, + FieldSource, + resolve_fields, +) +from litellm.types.proxy.management_endpoints.ui_sso import ( + RoleMappings, + SSOConfig, + TeamMappings, +) + +SSO_DESCRIPTORS: tuple[FieldDescriptor, ...] = ( + FieldDescriptor("google_client_id", "google_client_id", "GOOGLE_CLIENT_ID"), + FieldDescriptor("google_client_secret", "google_client_secret", "GOOGLE_CLIENT_SECRET", is_secret=True), + FieldDescriptor("microsoft_client_id", "microsoft_client_id", "MICROSOFT_CLIENT_ID"), + FieldDescriptor("microsoft_client_secret", "microsoft_client_secret", "MICROSOFT_CLIENT_SECRET", is_secret=True), + FieldDescriptor("microsoft_tenant", "microsoft_tenant", "MICROSOFT_TENANT"), + FieldDescriptor("generic_client_id", "generic_client_id", "GENERIC_CLIENT_ID"), + FieldDescriptor("generic_client_secret", "generic_client_secret", "GENERIC_CLIENT_SECRET", is_secret=True), + FieldDescriptor( + "generic_authorization_endpoint", "generic_authorization_endpoint", "GENERIC_AUTHORIZATION_ENDPOINT" + ), + FieldDescriptor("generic_token_endpoint", "generic_token_endpoint", "GENERIC_TOKEN_ENDPOINT"), + FieldDescriptor("generic_userinfo_endpoint", "generic_userinfo_endpoint", "GENERIC_USERINFO_ENDPOINT"), + FieldDescriptor("generic_scope", "generic_scope", "GENERIC_SCOPE", default="openid email profile"), + FieldDescriptor("saml_idp_metadata_url", "saml_idp_metadata_url", "SAML_IDP_METADATA_URL"), + FieldDescriptor("saml_idp_metadata_xml", "saml_idp_metadata_xml", "SAML_IDP_METADATA_XML"), + FieldDescriptor("saml_sp_entity_id", "saml_sp_entity_id", "SAML_SP_ENTITY_ID"), + FieldDescriptor("saml_allow_unsolicited", "saml_allow_unsolicited", "SAML_ALLOW_UNSOLICITED"), + FieldDescriptor("proxy_base_url", "proxy_base_url", "PROXY_BASE_URL"), +) + +# Derived from the descriptor table so read (masking) and the field->env mapping +# never diverge from the resolver. +SSO_SECRET_FIELDS: frozenset[str] = frozenset(d.field_name for d in SSO_DESCRIPTORS if d.is_secret) +SSO_FIELD_ENV_VARS: dict[str, str] = {d.field_name: d.env_var for d in SSO_DESCRIPTORS} + +# Structured sub-objects stored on the SSO row that are not simple env-backed +# scalars; handled outside the descriptor resolution. +_STRUCTURED_KEYS = ("role_mappings", "team_mappings") + + +@dataclass(frozen=True, slots=True) +class ResolvedSSOConfig: + config: SSOConfig + provenance: dict[str, FieldSource] + + +def _decrypt(raw: Mapping[str, object]) -> dict[str, object]: + return { + key: ( + decrypt_value_helper(value=value, key=key, return_original_value=True) if isinstance(value, str) else value + ) + for key, value in raw.items() + } + + +def _parse_role_mappings(data: object) -> RoleMappings | None: + # The stored row is JSON, so mappings arrive as a dict (or are absent). + return RoleMappings(**data) if isinstance(data, dict) else None + + +def _parse_team_mappings(data: object) -> TeamMappings | None: + return TeamMappings(**data) if isinstance(data, dict) else None + + +def resolve_sso_config(sso_db_settings: Mapping[str, object] | None, env: Mapping[str, str]) -> ResolvedSSOConfig: + """Resolve the effective SSO config: stored row first, then process env. + + Decryption happens here, once, via the pure ``decrypt_value_helper``; this + function never writes ``os.environ`` (unlike the legacy read path). Values + are returned unmasked so the login path could consume them; the read-back + endpoint is responsible for masking secrets before responding to the UI. + """ + raw = dict(sso_db_settings) if sso_db_settings else {} + decrypted = _decrypt({key: value for key, value in raw.items() if key not in _STRUCTURED_KEYS}) + values, provenance = resolve_fields(SSO_DESCRIPTORS, decrypted, env) + structured = { + "user_email": decrypted.get("user_email"), + "ui_access_mode": decrypted.get("ui_access_mode"), + "role_mappings": _parse_role_mappings(raw.get("role_mappings")), + "team_mappings": _parse_team_mappings(raw.get("team_mappings")), + } + config = SSOConfig(**{**values, **structured}) + return ResolvedSSOConfig(config=config, provenance=provenance) diff --git a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py index 54156715da8..cec682d772a 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py +++ b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py @@ -2046,6 +2046,90 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): masking_index += 1 verbose_proxy_logger.debug("Applied masking to choice text content") + @staticmethod + def _incremental_scan_cache() -> DualCache: + """Resolve the cache used to remember which segments a session already scanned. + + Prefers the proxy's shared cache (``internal_usage_cache.dual_cache``), which is + backed by Redis when the deployment configures it, so incremental state is shared + across proxy instances. Falls back to a process-local ``DualCache`` singleton when + the proxy is not running (e.g. unit tests), where sharing does not apply. + """ + from litellm.integrations.custom_guardrail import dc as fallback_cache + + try: + from litellm.proxy.proxy_server import proxy_logging_obj as _proxy_logging + except Exception: # noqa: BLE001 # proxy not importable outside the server; use local fallback + return fallback_cache + if _proxy_logging is not None: + return _proxy_logging.internal_usage_cache.dual_cache + return fallback_cache + + def _bedrock_response_has_masked_output(self, response: BedrockGuardrailResponse) -> bool: + """Return True if the guardrail rewrote (masked/anonymized) any scanned text. + + Bedrock returns non-empty ``output``/``outputs`` text only when it changed the + content; an ``action == "NONE"`` response leaves both empty. + """ + for field in ("output", "outputs"): + items = response.get(field) or [] + if any(isinstance(item, dict) and item.get("text") for item in items): + return True + return False + + async def _apply_incremental_request_scan( + self, + texts: list[str], + inputs: "GenericGuardrailAPIInputs", + request_data: dict, + ) -> Optional["GenericGuardrailAPIInputs"]: + """Scan only the text segments not already seen earlier in this session. + + Returns ``None`` when incremental scanning is inactive (feature off, no + session id, masking enabled, or cache unavailable) or when the guardrail + turns out to mask content, telling the caller to run the normal full scan. + Otherwise scans only the new segments and skips the Bedrock call entirely + when nothing is new. Incremental mode is for blocking/detection guardrails + only: if the guardrail returns masked output it cannot be applied to the + skipped context, so the scan falls back to the full path and no session + state is recorded. + """ + cache = self._incremental_scan_cache() + + new_texts = await self.filter_new_texts_for_session( + texts=texts, + request_data=request_data, + cache=cache, + ) + if new_texts is None: + return None + + if not new_texts: + verbose_proxy_logger.debug("Bedrock Guardrail: no new messages to scan for this session, skipping API call") + return inputs + + bedrock_response = await self.make_bedrock_api_request( + source="INPUT", + messages=[ChatCompletionUserMessage(role="user", content=text) for text in new_texts], + request_data=request_data, + logging_event_type=GuardrailEventHooks.pre_call, + ) + + if self._bedrock_response_has_masked_output(bedrock_response): + verbose_proxy_logger.warning( + "Bedrock Guardrail %s: guardrail returned masked/anonymized content; " + "only_scan_new_messages cannot apply masking to skipped context, falling back to a full-context scan", + self.guardrail_name, + ) + return None + + await self.mark_texts_scanned( + texts=texts, + request_data=request_data, + cache=cache, + ) + return inputs + async def apply_guardrail( self, inputs: "GenericGuardrailAPIInputs", @@ -2077,6 +2161,15 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): try: verbose_proxy_logger.debug(f"Bedrock Guardrail: Applying guardrail to {len(texts)} text(s)") + if input_type == "request": + incremental_result = await self._apply_incremental_request_scan( + texts=texts, + inputs=inputs, + request_data=request_data, + ) + if incremental_result is not None: + return incremental_result + masked_texts = [] selection = self._select_messages_for_apply_guardrail( diff --git a/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py b/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py index 31535a5b569..28b9dec100f 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py +++ b/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py @@ -432,7 +432,7 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): Override to store only the Model Armor API response, not the entire data dict. This prevents circular references in logging. """ - metadata = request_data.get("metadata", {}) if isinstance(request_data, dict) else {} + metadata = (request_data.get("metadata") or {}) if isinstance(request_data, dict) else {} guardrail_response = metadata.get("_model_armor_response", {}) # Determine status – default to "success" but prefer the explicit value if present. diff --git a/litellm/proxy/guardrails/guardrail_initializers.py b/litellm/proxy/guardrails/guardrail_initializers.py index 14e76a21093..e909c15382b 100644 --- a/litellm/proxy/guardrails/guardrail_initializers.py +++ b/litellm/proxy/guardrails/guardrail_initializers.py @@ -35,6 +35,7 @@ def initialize_bedrock(litellm_params: LitellmParams, guardrail: Guardrail): aws_sts_endpoint=litellm_params.aws_sts_endpoint, aws_bedrock_runtime_endpoint=litellm_params.aws_bedrock_runtime_endpoint, experimental_use_latest_role_message_only=litellm_params.experimental_use_latest_role_message_only, + only_scan_new_messages=litellm_params.only_scan_new_messages or False, ) litellm.logging_callback_manager.add_litellm_callback(_bedrock_callback) return _bedrock_callback diff --git a/litellm/proxy/guardrails/guardrail_registry.py b/litellm/proxy/guardrails/guardrail_registry.py index 1f2c9e0c182..bd00e9815a8 100644 --- a/litellm/proxy/guardrails/guardrail_registry.py +++ b/litellm/proxy/guardrails/guardrail_registry.py @@ -489,6 +489,9 @@ class InMemoryGuardrailHandler: "skip_tool_message_in_guardrail", getattr(litellm_params, "skip_tool_message_in_guardrail", None), ) + configured_run_in_parallel = getattr(litellm_params, "run_in_parallel", None) + if configured_run_in_parallel is not None: + custom_guardrail_callback.run_in_parallel = bool(configured_run_in_parallel) parsed_guardrail = Guardrail( guardrail_id=guardrail.get("guardrail_id"), diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index 22ea9fe176a..b2216488db2 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -1781,28 +1781,38 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): """ from litellm.proxy.auth.auth_utils import get_team_mcp_rpm_limit - if not mcp_server_name or not user_api_key_dict.team_id: + if not mcp_server_name: return - mcp_rpm_limit = get_team_mcp_rpm_limit(user_api_key_dict) - if not mcp_rpm_limit: - return + # Which teams' buckets does this call charge? A key is pinned to exactly one team. A keyless + # MCP-admitted subject reaches servers through SEVERAL teams at once and has no team_id, so + # without the second source below its calls charged no team bucket at all and it outran every + # team's mcp_rpm_limit. Every applicable team is charged rather than one being picked: the + # limiter enforces all descriptors, so each team's own ceiling binds on a call made through + # its grant, and there is no arbitrary attribution when several teams grant the same server. + team_limits: list[tuple[str | None, dict[str, int] | None]] = [] + if user_api_key_dict.team_id: + team_limits.append((user_api_key_dict.team_id, get_team_mcp_rpm_limit(user_api_key_dict))) + for source_team_id, source_limit in (user_api_key_dict.mcp_source_team_rpm_limits or {}).items(): + team_limits.append((source_team_id, source_limit)) - server_rpm_limit = mcp_rpm_limit.get(mcp_server_name) - if server_rpm_limit is None: - return - - descriptors.append( - RateLimitDescriptor( - key="mcp_per_team", - value=f"{user_api_key_dict.team_id}:{mcp_server_name}", - rate_limit={ - "requests_per_unit": server_rpm_limit, - "tokens_per_unit": None, - "window_size": self.window_size, - }, + for team_id, mcp_rpm_limit in team_limits: + if not team_id or not mcp_rpm_limit: + continue + server_rpm_limit = mcp_rpm_limit.get(mcp_server_name) + if server_rpm_limit is None: + continue + descriptors.append( + RateLimitDescriptor( + key="mcp_per_team", + value=f"{team_id}:{mcp_server_name}", + rate_limit={ + "requests_per_unit": server_rpm_limit, + "tokens_per_unit": None, + "window_size": self.window_size, + }, + ) ) - ) def _should_enforce_rate_limit( self, diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index eea2e58181c..1bdb34a7481 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -13,13 +13,16 @@ from starlette.datastructures import Headers import litellm from litellm._logging import verbose_logger, verbose_proxy_logger from litellm._service_logger import ServiceLogging -from litellm.constants import PRE_CALL_EXECUTED_GUARDRAILS_KEY +from litellm.constants import LITELLM_PROXY_MASTER_KEY_ALIAS, PRE_CALL_EXECUTED_GUARDRAILS_KEY from litellm.litellm_core_utils.credential_accessor import CredentialAccessor from litellm.litellm_core_utils.initialize_dynamic_callback_params import ( iter_client_callback_metadata_dicts, ) from litellm.litellm_core_utils.safe_json_loads import safe_json_loads -from litellm.litellm_core_utils.url_utils import is_url_destination_allowed_by_host +from litellm.litellm_core_utils.url_utils import ( + is_url_destination_allowed_by_host, + provider_url_destination_candidates, +) from litellm.proxy._types import ( AddTeamCallback, CommonProxyErrors, @@ -45,6 +48,24 @@ _EXPLICIT_SESSION_HEADERS = frozenset({"x-litellm-trace-id", "x-litellm-session- # Session-id values must be non-empty strings of alphanumerics, hyphens, or underscores # (covers UUIDs and most common session-id formats). _SESSION_ID_VALUE_RE = re.compile(r"^[a-zA-Z0-9_\-]{8,}$") + +_SHA256_HEX_RE = re.compile(r"^[0-9a-f]{64}$") + + +def _stampable_key_hash(user_api_key_dict: UserAPIKeyAuth) -> str | None: + """Only proxy-validated keys are stamped, proven by the unforgeable + via_virtual_key marker AND a known non-secret shape: the sha256 hex digest + UserAPIKeyAuth stores virtual keys in, or the master key's stable alias. + Custom-auth credentials arrive raw (never forward auth material) and hashed + JWTs rotate on re-issue (useless as a stable ban id), so both are skipped.""" + api_key = user_api_key_dict.api_key + if not user_api_key_dict.via_virtual_key or api_key is None: + return None + if api_key == LITELLM_PROXY_MASTER_KEY_ALIAS or _SHA256_HEX_RE.fullmatch(api_key): + return api_key + return None + + _ANTHROPIC_SESSION_ID_VALUE_RE = re.compile(r"^[a-zA-Z0-9_\-]+$") @@ -229,23 +250,26 @@ def _reject_url_valued_destinations(data: Dict[str, Any]) -> None: allowed_hosts = getattr(litellm, "provider_url_destination_allowed_hosts", []) or [] for field in _URL_DESTINATION_REQUEST_FIELDS: value = data.get(field) - if not isinstance(value, str) or not value.startswith(("http://", "https://")): + if not isinstance(value, str): continue - if is_url_destination_allowed_by_host(value, allowed_hosts): - continue - raise HTTPException( - status_code=400, - detail={ - "error": "invalid_request", - "param": field, - "message": ( - f"URL-valued '{field}' is not allowed. Configure custom " - "endpoints with api_base instead, or add the destination " - "host to `provider_url_destination_allowed_hosts` in " - "litellm_settings." - ), - }, - ) + for candidate in provider_url_destination_candidates(value): + if not candidate.lower().startswith(("http://", "https://")): + continue + if is_url_destination_allowed_by_host(candidate, allowed_hosts): + continue + raise HTTPException( + status_code=400, + detail={ + "error": "invalid_request", + "param": field, + "message": ( + f"URL-valued '{field}' is not allowed. Configure custom " + "endpoints with api_base instead, or add the destination " + "host to `provider_url_destination_allowed_hosts` in " + "litellm_settings." + ), + }, + ) def _strip_untrusted_request_header_controls( @@ -1693,6 +1717,11 @@ async def add_litellm_data_to_request( if "user" not in data: data["user"] = user + if litellm.overwrite_user_with_key_hash is True: + stampable_hash = _stampable_key_hash(user_api_key_dict) + if stampable_hash is not None: + data["user"] = stampable_hash + data["secret_fields"] = SecretFields(raw_headers=_raw_headers) ## Dynamic api version (Azure OpenAI endpoints) ## diff --git a/litellm/proxy/management_endpoints/cache_settings_endpoints.py b/litellm/proxy/management_endpoints/cache_settings_endpoints.py index 9f45cb619aa..7c0d8958a28 100644 --- a/litellm/proxy/management_endpoints/cache_settings_endpoints.py +++ b/litellm/proxy/management_endpoints/cache_settings_endpoints.py @@ -18,8 +18,9 @@ from pydantic import BaseModel, Field import litellm from litellm._logging import verbose_proxy_logger +from litellm._redis import _redis_kwargs_from_environment from litellm._uuid import uuid -from litellm.litellm_core_utils.sensitive_data_masker import mask_sensitive_keys +from litellm.litellm_core_utils.sensitive_data_masker import SensitiveDataMasker from litellm.proxy._types import ( AUDIT_ACTIONS, LiteLLM_AuditLogs, @@ -43,6 +44,17 @@ router = APIRouter() # (e.g. redis://:secret@host:6379/1). _CACHE_SENSITIVE_FIELDS: set = {"password", "sentinel_password", "url"} +# The env fallback resolves the full set of redis.Redis kwargs, which includes +# credential-bearing params (azure_client_secret, ssl_password, ...) that are +# not cache UI fields. Only overlay fields the settings page actually renders, +# so the read never surfaces a credential the UI does not manage. +_CACHE_SETTINGS_FIELD_NAMES: frozenset = frozenset(field.field_name for field in CACHE_SETTINGS_FIELDS) + +# Classifier used, alongside _CACHE_SENSITIVE_FIELDS, to redact any +# credential-bearing key before it leaves the server (`url` is kept in the +# explicit set because its name carries no sensitive segment). +_CREDENTIAL_CLASSIFIER = SensitiveDataMasker() + _REDACTED_VALUE = "***REDACTED***" @@ -67,6 +79,165 @@ def _resolve_cache_url_precedence(settings: Mapping[str, Any]) -> dict[str, Any] return {k: v for k, v in settings.items() if k not in _URL_OVERRIDDEN_CONNECTION_FIELDS} +def _parse_stored_settings(cache_settings_value: object) -> dict[str, Any]: + """Normalize a stored cache_settings blob to a dict. + + The prisma column comes back as either a JSON string or an already-parsed + dict depending on the client, so callers that json.loads unconditionally + silently drop the whole (still-encrypted) row on the dict path. + """ + parsed = json.loads(cache_settings_value) if isinstance(cache_settings_value, str) else cache_settings_value + return parsed if isinstance(parsed, dict) else {} + + +def _overlay_environment(stored: Mapping[str, Any]) -> dict[str, Any]: + """Fill connection fields from the REDIS_* environment the cache actually reads. + + A response cache pointed at Redis resolves host/port/password/etc. from the + REDIS_* env vars when the stored config leaves them unset, so a cache + configured purely through the environment works while its settings page, + which reads only the database row, shows blank. Overlaying the same env + kwargs the runtime uses makes the page reflect the effective connection. + Stored values win; the environment only fills what the stored config omits. + """ + env_kwargs = { + key: value for key, value in _redis_kwargs_from_environment().items() if key in _CACHE_SETTINGS_FIELD_NAMES + } + if not env_kwargs: + return dict(stored) + effective = {**env_kwargs, **stored} + # the env fallback is a Redis connection, so name the type when the stored + # config did not, letting the UI render the Redis fields it just populated + effective.setdefault("type", "redis") + return effective + + +def _redact_credentials(settings: Mapping[str, Any]) -> dict[str, Any]: + """Replace credential-bearing values with a fixed marker, keeping the rest. + + The marker is unambiguous on the way back in: an admin who edits an + unrelated field and re-submits sends the marker for the untouched secret, + which the update path maps back to the stored value rather than persisting + the marker over a working password. + """ + return { + key: (_REDACTED_VALUE if value is not None and _is_credential_field(key) else value) + for key, value in settings.items() + } + + +def _is_credential_field(key: str) -> bool: + """Whether a cache setting carries a credential and must be redacted on read.""" + return key in _CACHE_SENSITIVE_FIELDS or _CREDENTIAL_CLASSIFIER.is_sensitive_key(key) + + +def _has_connection_target(value: object) -> bool: + """Whether a payload value names a live discrete connection target.""" + if isinstance(value, str): + return value.strip() != "" and value != _REDACTED_VALUE + return value not in (None, [], {}) + + +# Every field that identifies which Redis a credential belongs to, across node +# (host/port/url), cluster (redis_startup_nodes), and sentinel +# (sentinel_nodes/service_name) modes. A stored secret is bound to these. +_CONNECTION_TARGET_FIELDS: tuple = ( + "host", + "port", + "url", + "redis_startup_nodes", + "sentinel_nodes", + "service_name", +) + + +def _target_repr(value: object) -> str: + """Canonical string form of a connection-target value for equality checks. + + The client may serialize the same target differently from storage (a port as + "6379" vs 6379, node lists round-tripped through JSON), so compare normalized + forms rather than raw values to avoid treating an unchanged target as a change. + """ + if isinstance(value, (list, dict)): + return json.dumps(value, sort_keys=True, default=str) + return str(value) + + +def _saved_secret_is_reusable(incoming: Mapping[str, Any], saved: Mapping[str, Any]) -> bool: + """Whether a stored credential may be restored for this request. + + A stored secret belongs to the stored connection target, so it is reused only + when the request describes that same target on every dimension the stored + config pins (host/port, url, cluster nodes, sentinel nodes/service). This + prevents credential replay: a caller cannot omit the credential, point at a + different (or incomplete) target, and have the proxy send the stored secret + to a Redis of their choosing. + + Non-secret target fields (host/port/nodes/service) must be supplied and match + in normalized form, so equivalent representations (port "6379" vs 6379) are + not seen as a change while an omitted or different value is. ``url`` is the + exception: it is itself the secret and the form never re-prefills it, so a + redacted or omitted url means "keep the stored url" (same target) and only a + different supplied url blocks reuse. + """ + for field in _CONNECTION_TARGET_FIELDS: + saved_value = saved.get(field) + if saved_value in (None, "", [], {}): + continue # the stored config does not pin this dimension + incoming_value = incoming.get(field) + if field == "url": + if incoming_value in (None, "", _REDACTED_VALUE): + continue # url kept as-is (same target) + if _target_repr(incoming_value) != _target_repr(saved_value): + return False + continue + if _target_repr(incoming_value) != _target_repr(saved_value): + return False # a pinned target field is missing or different + return True + + +def _merge_over_saved(incoming: Mapping[str, Any], saved: Mapping[str, Any]) -> dict[str, Any]: + """Keep the stored secret behind any credential the caller echoed back redacted or omitted. + + GET returns credentials as the marker and the form never re-prefills a + secret, so a save that does not touch a credential arrives with the marker + or with the field absent. Either way the real secret must survive: it is + restored from the stored row, or dropped when there is no stored row (the + value is env-sourced and the marker must never be persisted). Non-secret + fields are taken from the incoming payload as-is, so clearing one still works. + + ``url`` is the exception: it is credential-bearing (redacted) yet also a + connection-mode selector that url-precedence resolves against host/port. If + the caller supplies a discrete target (host, cluster, or sentinel nodes), a + stored url is a stale mode the caller is leaving, so it is dropped rather + than restored, otherwise url-precedence would resurrect it and discard the + submitted host/port. + """ + switching_to_discrete_target = ( + _has_connection_target(incoming.get("host")) + or _has_connection_target(incoming.get("redis_startup_nodes")) + or _has_connection_target(incoming.get("sentinel_nodes")) + ) + reuse_saved_secret = _saved_secret_is_reusable(incoming, saved) + merged = dict(incoming) + for field in _CACHE_SENSITIVE_FIELDS: + # A value the caller explicitly supplied is honored verbatim: a new + # secret, or an empty string / null to clear the stored one. Only an + # omitted field or the echoed-back marker triggers preserve-or-drop. + if field in incoming and incoming[field] != _REDACTED_VALUE: + continue + if field == "url" and switching_to_discrete_target: + merged.pop(field, None) + continue + if field in saved and reuse_saved_secret: + merged[field] = saved[field] + else: + # nothing stored to reuse, or the caller is pointing at a different + # target: never persist/replay the marker or the stored secret + merged.pop(field, None) + return merged + + def _redact_settings(settings: Optional[Mapping[str, Any]]) -> Dict[str, Any]: """Replace every value in a settings map with a fixed marker. @@ -270,34 +441,34 @@ async def get_cache_settings( # Get cache settings fields from types file cache_fields = [field.model_copy(deep=True) for field in CACHE_SETTINGS_FIELDS] - # Try to get cache settings from database - current_values = {} + # Read the stored settings (decrypted); an env-only cache has none. + stored: dict[str, Any] = {} if prisma_client is not None: cache_config = await CacheConfigRepository(prisma_client).table.find_unique(where={"id": "cache_config"}) if cache_config is not None and cache_config.cache_settings: - # Decrypt cache settings - cache_settings_json = cache_config.cache_settings - if isinstance(cache_settings_json, str): - cache_settings_dict = json.loads(cache_settings_json) - else: - cache_settings_dict = cache_settings_json + stored = proxy_config._decrypt_db_variables( + variables_dict=_parse_stored_settings(cache_config.cache_settings) + ) - # Decrypt environment variables - decrypted_settings = proxy_config._decrypt_db_variables(variables_dict=cache_settings_dict) + # Fill connection fields from the REDIS_* environment the cache resolves + # from when the stored config leaves them unset, then apply url precedence + # so a url-mode config does not surface conflicting discrete fields (which + # would otherwise let a no-op save silently switch it to host/port). + effective = _resolve_cache_url_precedence(_overlay_environment(stored)) - # Derive redis_type for UI based on settings - # UI uses redis_type to show/hide fields, backend only stores 'type' - if decrypted_settings.get("type") == "redis": - if decrypted_settings.get("redis_startup_nodes"): - decrypted_settings["redis_type"] = "cluster" - elif decrypted_settings.get("sentinel_nodes"): - decrypted_settings["redis_type"] = "sentinel" - else: - decrypted_settings["redis_type"] = "node" + # Derive redis_type for UI based on settings + # UI uses redis_type to show/hide fields, backend only stores 'type' + if effective.get("type") == "redis": + if effective.get("redis_startup_nodes"): + effective["redis_type"] = "cluster" + elif effective.get("sentinel_nodes"): + effective["redis_type"] = "sentinel" + else: + effective["redis_type"] = "node" - # Mask credential fields so the GET response never carries - # plaintext Redis / Sentinel passwords off the server. - current_values = mask_sensitive_keys(decrypted_settings, _CACHE_SENSITIVE_FIELDS) + # Redact credential fields so the GET response never carries a plaintext + # Redis / Sentinel password off the server. + current_values = _redact_credentials(effective) # Update field values with current values for field in cache_fields: @@ -331,10 +502,27 @@ async def test_cache_connection( to verify the credentials work without affecting global state. """ from litellm import Cache + from litellm.proxy.proxy_server import prisma_client, proxy_config try: - cache_settings = _resolve_cache_url_precedence(request.cache_settings) - verbose_proxy_logger.debug("Testing cache connection with settings: %s", cache_settings) + # A credential the form left untouched arrives redacted; resolve it back + # to the stored secret so the test connects with the real password. A + # lookup failure must not block the test, so fall back to no stored row. + saved_settings: dict[str, Any] = {} + if prisma_client is not None: + try: + existing_row = await CacheConfigRepository(prisma_client).table.find_unique( + where={"id": "cache_config"} + ) + if existing_row is not None and existing_row.cache_settings: + saved_settings = proxy_config._decrypt_db_variables( + variables_dict=_parse_stored_settings(existing_row.cache_settings) + ) + except Exception: # noqa: BLE001 - a saved-settings lookup failure must not block a connection test + saved_settings = {} + cache_settings = _resolve_cache_url_precedence(_merge_over_saved(request.cache_settings, saved_settings)) + # cache_settings now carries the resolved plaintext credential; never log it raw + verbose_proxy_logger.debug("Testing cache connection with settings: %s", _redact_credentials(cache_settings)) # Only support Redis for now if cache_settings.get("type") != "redis": @@ -400,19 +588,20 @@ async def update_cache_settings( ) try: - cache_settings = _resolve_cache_url_precedence(request.cache_settings) - - # Snapshot the prior settings (key set only — values get redacted in - # the audit row) so the audit-log entry shows which fields changed. + # Read the stored row first: its decrypted values back any credential the + # caller echoed back redacted, and its key set drives the audit diff. existing_row = await CacheConfigRepository(prisma_client).table.find_unique(where={"id": "cache_config"}) before_settings: Optional[Dict[str, Any]] = None + saved_settings: dict[str, Any] = {} if existing_row is not None and existing_row.cache_settings: - try: - before_settings = json.loads(existing_row.cache_settings) - except (TypeError, ValueError): - before_settings = None + before_settings = _parse_stored_settings(existing_row.cache_settings) + saved_settings = proxy_config._decrypt_db_variables(variables_dict=before_settings) action: AUDIT_ACTIONS = "updated" if existing_row is not None else "created" + # Preserve stored secrets behind any redacted or omitted credential, then + # resolve the url-vs-discrete-fields precedence. + cache_settings = _resolve_cache_url_precedence(_merge_over_saved(request.cache_settings, saved_settings)) + # Encrypt sensitive fields (keep redis_type for storage) encrypted_settings = proxy_config._encrypt_env_variables(environment_variables=cache_settings) @@ -461,7 +650,7 @@ async def update_cache_settings( return { "message": "Cache settings updated successfully", "status": "success", - "settings": cache_settings, + "settings": _redact_credentials(cache_settings), } except Exception as e: verbose_proxy_logger.error(f"Error updating cache settings: {str(e)}") diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index d50db8324ef..89f28a30a84 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -136,9 +136,12 @@ if MCP_AVAILABLE: from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( _raise_if_not_oauth2, authorize_with_server, + client_supplied_redirect_uris, exchange_token_with_server, get_request_base_url, + redeem_passthrough_authorization_code, register_client_with_server, + resolve_ephemeral_dcr_client, ) from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( global_mcp_server_manager, @@ -1661,7 +1664,21 @@ if MCP_AVAILABLE: mcp_server = await _get_cached_temporary_mcp_server_or_404(server_id, user_api_key_dict, request=request) _raise_if_not_oauth2(mcp_server) # Use the server's stored client_id when the caller doesn't supply one - resolved_client_id = mcp_server.client_id or client_id or "" + stored_or_supplied_client_id = mcp_server.client_id or client_id or "" + ephemeral_dcr_client = ( + await resolve_ephemeral_dcr_client( + request=request, + mcp_server=mcp_server, + code_challenge=code_challenge, + code_challenge_method=code_challenge_method, + redirect_uri=redirect_uri, + ) + if not stored_or_supplied_client_id + else None + ) + resolved_client_id = stored_or_supplied_client_id or ( + ephemeral_dcr_client.client_id if ephemeral_dcr_client else "" + ) if not resolved_client_id: raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, @@ -1683,6 +1700,7 @@ if MCP_AVAILABLE: code_challenge_method=code_challenge_method, response_type=response_type, scope=scope, + ephemeral_dcr_client=ephemeral_dcr_client, ) @router.post( @@ -1705,7 +1723,21 @@ if MCP_AVAILABLE: ): mcp_server = await _get_cached_temporary_mcp_server_or_404(server_id, user_api_key_dict, request=request) _raise_if_not_oauth2(mcp_server) - resolved_client_id = mcp_server.client_id or client_id or "" + # Sealed passthrough codes exist only for the authorization_code grant. A refresh_token + # grant must never open one: the minted client is unrecoverable after the single flow by + # contract, so an expired browser-held token re-runs authorize instead. + sealed_code = ( + redeem_passthrough_authorization_code(code=code, mcp_server=mcp_server, code_verifier=code_verifier) + if grant_type == "authorization_code" + else None + ) + resolved_code = sealed_code.upstream_code if sealed_code else code + # A sealed flow ran the gateway /callback as its upstream redirect (bridge short-circuit + # or plain flow alike), so the exchange must present that binding, not the browser page. + resolved_redirect_uri = f"{get_request_base_url(request)}/callback" if sealed_code else redirect_uri + caller_client_id = sealed_code.client_id if sealed_code else client_id + caller_client_secret = sealed_code.client_secret if sealed_code else client_secret + resolved_client_id = mcp_server.client_id or caller_client_id or "" if not resolved_client_id: raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, @@ -1721,13 +1753,14 @@ if MCP_AVAILABLE: request=request, mcp_server=mcp_server, grant_type=grant_type, - code=code, - redirect_uri=redirect_uri, + code=resolved_code, + redirect_uri=resolved_redirect_uri, client_id=resolved_client_id, - client_secret=client_secret, + client_secret=caller_client_secret, code_verifier=code_verifier, refresh_token=refresh_token, scope=scope, + client_token_endpoint_auth_method=sealed_code.token_endpoint_auth_method if sealed_code else None, ) @router.post( @@ -1743,6 +1776,7 @@ if MCP_AVAILABLE: mcp_server = await _get_cached_temporary_mcp_server_or_404(server_id, user_api_key_dict, request=request) request_data = await _read_request_body(request=request) data: dict = {**request_data} + client_redirect_uris = client_supplied_redirect_uris(data.get("redirect_uris")) return await register_client_with_server( request=request, @@ -1753,6 +1787,7 @@ if MCP_AVAILABLE: token_endpoint_auth_method=data.get("token_endpoint_auth_method", ""), fallback_client_id=server_id, persist_credentials=_user_is_full_admin(user_api_key_dict), + client_redirect_uris=client_redirect_uris, ) @router.delete( diff --git a/litellm/proxy/management_endpoints/organization_endpoints.py b/litellm/proxy/management_endpoints/organization_endpoints.py index 5fe27e9dc13..df3c9c3c17b 100644 --- a/litellm/proxy/management_endpoints/organization_endpoints.py +++ b/litellm/proxy/management_endpoints/organization_endpoints.py @@ -13,16 +13,18 @@ Endpoints for /organization operations #### ORGANIZATION MANAGEMENT #### -from typing import Any, Dict, List, Optional, Tuple +from typing import Annotated, Any, Dict, List, Mapping, Optional, Tuple import fastapi from fastapi import APIRouter, Depends, HTTPException, Request, status +from pydantic import TypeAdapter from litellm._logging import verbose_proxy_logger from litellm._uuid import uuid from litellm.proxy._types import * from litellm.proxy.auth.auth_checks import can_user_call_model, get_user_object from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time from litellm.proxy.management_endpoints.budget_management_endpoints import ( new_budget, update_budget, @@ -34,6 +36,7 @@ from litellm.proxy.management_endpoints.common_utils import ( ) from litellm.proxy.management_helpers.object_permission_utils import ( handle_update_object_permission_common, + prepare_object_permission_upsert, ) from litellm.proxy.management_helpers.utils import ( get_new_internal_user_defaults, @@ -101,6 +104,30 @@ async def _verify_org_access( ) +_STR_OBJECT_DICT_ADAPTER = TypeAdapter(dict[str, object]) +_BUDGET_SETTABLE_FIELDS = frozenset(LiteLLM_BudgetTable.model_fields.keys()) - {"budget_id"} +_ORG_COLUMN_FIELDS = frozenset({"organization_alias", "models"}) + + +def build_budget_write_data(budget_updates: Mapping[str, object], updated_by: str) -> Mapping[str, object]: + """ + Budget-row columns to write. ``budget_reset_at`` tracks any sent ``budget_duration``: + recomputed for a new duration, cleared alongside a ``None`` duration so no stale reset + timestamp survives. Other sent fields (including a ``None`` clear) are written as-is. + """ + budget_duration = budget_updates.get("budget_duration") + recomputed_reset_at: Mapping[str, object] = ( + { + "budget_reset_at": ( + get_budget_reset_time(budget_duration=budget_duration) if isinstance(budget_duration, str) else None + ) + } + if "budget_duration" in budget_updates + else {} + ) + return {**budget_updates, **recomputed_reset_at, "updated_by": updated_by} + + def handle_nested_budget_structure_in_organization_update_request( raw_data: dict, ) -> dict: @@ -568,6 +595,155 @@ async def handle_update_object_permission( return data_json +@router.patch( + "/v2/organization/{organization_id}", + tags=["organization management"], + dependencies=[Depends(user_api_key_auth)], + response_model=LiteLLM_OrganizationTableWithMembers, + include_in_schema=False, +) +async def update_organization_v2( + organization_id: str, + data: OrganizationUpdateRequestV2, + user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)], +): + """ + Partial update of an organization (RESTful PATCH, RFC 7396 merge-patch semantics). + + A sent field is written and an omitted one is left untouched (presence is read from + ``model_fields_set``). Clear tokens are per field: budget limits and ``metadata`` clear with + ``null``, ``models`` with ``[]``, and ``object_permission`` with ``null`` (it merges when sent, + so an empty ``{}`` is rejected). ``organization_alias`` is required and cannot be cleared. + Validation failures return 422; the object-permission upsert, budget-row write, and + org-row write are one transaction. + """ + from litellm.proxy.proxy_server import prisma_client + + if prisma_client is None: + raise HTTPException( + status_code=500, + detail={"error": CommonProxyErrors.db_not_connected_error.value}, + ) + + if user_api_key_dict.user_id is None: + raise HTTPException( + status_code=400, + detail={ + "error": "Cannot associate a user_id to this action. Check `/key/info` to validate if 'user_id' is set." + }, + ) + + if data.max_budget is not None and (not math.isfinite(data.max_budget) or data.max_budget < 0): + raise HTTPException( + status_code=422, + detail={"error": f"max_budget must be a non-negative finite number. Received: {data.max_budget}"}, + ) + if data.soft_budget is not None and (not math.isfinite(data.soft_budget) or data.soft_budget < 0): + raise HTTPException( + status_code=422, + detail={"error": f"soft_budget must be a non-negative finite number. Received: {data.soft_budget}"}, + ) + if data.model_max_budget: + from litellm.proxy.management_endpoints.key_management_endpoints import ( + validate_model_max_budget, + ) + + try: + validate_model_max_budget(data.model_max_budget) + except ValueError as e: + raise HTTPException(status_code=422, detail={"error": str(e)}) + + if "organization_alias" in data.model_fields_set and data.organization_alias is None: + raise HTTPException( + status_code=422, + detail={"error": "organization_alias cannot be cleared; it is required"}, + ) + if "models" in data.model_fields_set and data.models is None: + raise HTTPException( + status_code=422, + detail={"error": "models cannot be set to null; send [] to clear it"}, + ) + if data.object_permission is not None and not data.object_permission.model_dump(exclude_none=True): + raise HTTPException( + status_code=422, + detail={ + "error": "object_permission cannot be an empty object; send null to clear it, or a non-empty object to set grants" + }, + ) + + await _verify_org_access( + organization_id=organization_id, + user_api_key_dict=user_api_key_dict, + prisma_client=prisma_client, + ) + + existing_organization_row = await OrganizationRepository(prisma_client).table.find_unique( + where={"organization_id": organization_id}, + ) + if existing_organization_row is None: + raise HTTPException( + status_code=404, + detail={"error": f"Organization not found for organization_id={organization_id}"}, + ) + + field_values = _STR_OBJECT_DICT_ADAPTER.validate_python(data.model_dump()) + present_fields = data.model_fields_set + budget_updates = {field: field_values[field] for field in present_fields if field in _BUDGET_SETTABLE_FIELDS} + org_column_updates: Mapping[str, object] = { + **{field: field_values[field] for field in present_fields if field in _ORG_COLUMN_FIELDS}, + **({"metadata": data.metadata or {}} if "metadata" in present_fields else {}), + **({"logging_exporters": data.logging_exporters or []} if "logging_exporters" in present_fields else {}), + } + + object_permission_cleared = "object_permission" in present_fields and data.object_permission is None + object_permission_upsert = ( + await prepare_object_permission_upsert( + new_object_permission=data.object_permission.model_dump(exclude_none=True), + existing_object_permission_id=existing_organization_row.object_permission_id, + prisma_client=prisma_client, + ) + if data.object_permission is not None + else None + ) + object_permission_write: Mapping[str, object] = ( + {"object_permission_id": object_permission_upsert.object_permission_id} + if object_permission_upsert is not None + else ({"object_permission_id": None} if object_permission_cleared else {}) + ) + + organization_write_data = prisma_client.jsonify_object( + { + **org_column_updates, + **object_permission_write, + "updated_by": user_api_key_dict.user_id, + } + ) + + async with prisma_client.db.tx() as tx: + if object_permission_upsert is not None: + await tx.litellm_objectpermissiontable.upsert( + where={"object_permission_id": object_permission_upsert.object_permission_id}, + data={ + "create": object_permission_upsert.record, + "update": object_permission_upsert.record, + }, + ) + if budget_updates: + await tx.litellm_budgettable.update( + where={"budget_id": existing_organization_row.budget_id}, + data=prisma_client.jsonify_object( + dict(build_budget_write_data(budget_updates, user_api_key_dict.user_id)) + ), + ) + response = await tx.litellm_organizationtable.update( + where={"organization_id": organization_id}, + data=organization_write_data, + include={"members": True, "teams": True, "litellm_budget_table": True}, + ) + + return response + + @router.delete( "/organization/delete", tags=["organization management"], diff --git a/litellm/proxy/management_endpoints/scim/scim_v2.py b/litellm/proxy/management_endpoints/scim/scim_v2.py index fa123b7d76c..90eae5bbb21 100644 --- a/litellm/proxy/management_endpoints/scim/scim_v2.py +++ b/litellm/proxy/management_endpoints/scim/scim_v2.py @@ -98,14 +98,27 @@ class UserProvisionerHelpers: if not existing_user: return None - # Update the user + new_teams = list(dict.fromkeys(new_user_request.teams or [])) + + if new_user_request.user_id != existing_user.user_id: + await UserRepository(prisma_client).table.update( + where={"user_id": existing_user.user_id}, + data={"user_id": new_user_request.user_id}, + ) + + await _handle_team_membership_changes( + user_id=new_user_request.user_id, + existing_teams=existing_user.teams or [], + new_teams=new_teams, + raise_on_error=True, + ) + updated_user = await UserRepository(prisma_client).table.update( - where={"user_id": existing_user.user_id}, + where={"user_id": new_user_request.user_id}, data={ - "user_id": new_user_request.user_id, "user_email": new_user_request.user_email, "user_alias": new_user_request.user_alias, - "teams": new_user_request.teams, + "teams": new_teams, "metadata": safe_dumps(new_user_request.metadata), **({"user_role": new_user_request.user_role} if admin_group is not None else {}), }, @@ -440,7 +453,12 @@ async def _get_team_members_display(member_ids: List[str]) -> List[SCIMMember]: return members -async def _handle_team_membership_changes(user_id: str, existing_teams: List[str], new_teams: List[str]) -> None: +async def _handle_team_membership_changes( + user_id: str, + existing_teams: List[str], + new_teams: List[str], + raise_on_error: bool = False, +) -> None: """Handle adding/removing user from teams based on changes.""" existing_teams_set = set(existing_teams) new_teams_set = set(new_teams) @@ -453,6 +471,7 @@ async def _handle_team_membership_changes(user_id: str, existing_teams: List[str user_id=user_id, teams_ids_to_add_user_to=list(teams_to_add), teams_ids_to_remove_user_from=list(teams_to_remove), + raise_on_error=raise_on_error, ) @@ -1298,6 +1317,13 @@ async def delete_user( where={"team_id": team.team_id}, data={"members": new_members} ) + team_row = LiteLLM_TeamTable(**team.model_dump()) + if any(member.user_id == user_id for member in team_row.members_with_roles or []): + await team_member_delete( + data=TeamMemberDeleteRequest(team_id=team_row.team_id, user_id=user_id), + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), + ) + await _set_user_keys_blocked(user_id=user_id, blocked=True) await _delete_rows_referencing_user(prisma_client, user_id=user_id) @@ -1327,6 +1353,31 @@ def _extract_group_values(value: Any) -> List[str]: return group_values +def _extract_ids_from_path_filter(path: str | None, attribute: str) -> List[str]: + """Return ids from a SCIM filtered path like ``members[value eq "id"]``. + + Okta commonly sends membership removals as a filtered path and omits the + request body ``value``, so the id lives only inside the ``[value eq "..."]`` + filter. The ``eq`` operator is matched case-insensitively per the SCIM + spec; the id keeps its original case. Per the SCIM filter grammar the + compared value must be quoted (single or double), so malformed unquoted + filters yield no id. A quoted id may contain escaped quotes and + backslashes (``\\"`` and ``\\\\``), which are unescaped before use. + ``path`` must be the raw, case-preserving path from the patch op. + """ + if not path: + return [] + match = re.match( + rf"""\s*{re.escape(attribute)}\s*\[\s*value\s+eq\s+(['"])((?:\\.|[^\\])*?)\1\s*\]\s*$""", + path, + flags=re.IGNORECASE, + ) + if not match: + return [] + extracted = re.sub(r"\\(.)", r"\1", match.group(2)) + return [extracted] if extracted else [] + + def _handle_displayname_update(op_type: str, value: Any, update_data: Dict[str, Any]) -> None: """Handle displayname updates.""" if op_type == "remove": @@ -1370,9 +1421,11 @@ def _handle_name_update(path: str, op_type: str, value: Any, scim_metadata: Dict scim_metadata["familyName"] = str(value) -def _handle_group_operations(op_type: str, value: Any, teams_set: Set[str]) -> Optional[Set[str]]: +def _handle_group_operations(op_type: str, value: Any, teams_set: Set[str], path: str | None) -> Set[str] | None: """Handle group/team membership operations.""" group_values = _extract_group_values(value) + if not group_values and value is None: + group_values = _extract_ids_from_path_filter(path, "groups") if op_type == "replace": return set(group_values) elif op_type == "add": @@ -1485,7 +1538,7 @@ def _apply_patch_ops( elif _multi_valued_attribute_base(path) in SCIM_MULTI_VALUED_ATTRIBUTE_METADATA_KEYS: _handle_multi_valued_attribute_update(path, op_type, value, metadata) elif path.startswith("groups"): - new_replace_set = _handle_group_operations(op_type, value, teams_set) + new_replace_set = _handle_group_operations(op_type, value, teams_set, op.path) if new_replace_set is not None: replace_team_set = new_replace_set else: @@ -1497,16 +1550,29 @@ def _apply_patch_ops( return update_data, final_team_set +def _is_user_not_in_team_error(exc: HTTPException) -> bool: + """True when team_member_delete reports the user was already absent from the + team, which is the idempotent no-op case for a removal.""" + detail = exc.detail + return isinstance(detail, dict) and detail.get("error") == "User not found in team" + + async def patch_team_membership( user_id: str, teams_ids_to_add_user_to: List[str], teams_ids_to_remove_user_from: List[str], + raise_on_error: bool = False, ) -> bool: """ Add or remove user from teams Handles duplicate membership gracefully (idempotent operation). - If a user is already in a team, that's fine - we don't treat it as an error. + A user already being in a team (on add) or already absent from it (on + remove) is treated as a no-op, not an error. + + When ``raise_on_error`` is True a genuine add or remove failure (anything + other than those idempotent no-ops) propagates instead of being swallowed, + so a caller can avoid persisting a teams array the roster never received. """ for _team_id in teams_ids_to_add_user_to: try: @@ -1521,9 +1587,13 @@ async def patch_team_membership( # Handle duplicate membership gracefully - this is idempotent if e.type == ProxyErrorTypes.team_member_already_in_team: verbose_proxy_logger.debug(f"User {user_id} is already in team {_team_id}, skipping add") + elif raise_on_error: + raise else: verbose_proxy_logger.exception(f"Error adding user to team {_team_id}: {e}") except Exception as e: + if raise_on_error: + raise verbose_proxy_logger.exception(f"Error adding user to team {_team_id}: {e}") for _team_id in teams_ids_to_remove_user_from: @@ -1532,7 +1602,16 @@ async def patch_team_membership( data=TeamMemberDeleteRequest(team_id=_team_id, user_id=user_id), user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), ) + except HTTPException as e: + if _is_user_not_in_team_error(e): + verbose_proxy_logger.debug(f"User {user_id} is not in team {_team_id}, skipping remove") + elif raise_on_error: + raise + else: + verbose_proxy_logger.exception(f"Error removing user from team {_team_id}: {e}") except Exception as e: + if raise_on_error: + raise verbose_proxy_logger.exception(f"Error removing user from team {_team_id}: {e}") return True @@ -1654,8 +1733,11 @@ async def get_groups( # Convert to SCIM format scim_groups = [] for team in teams: - # Get team members with display names - members = await _get_team_members_display(team.members or []) + # Get team members with display names. members_with_roles is the + # source of truth; the legacy `members` column is not populated by + # team creation, so reading it here would report an empty member + # list to the IdP and trigger repeated re-provisioning. + members = await _get_team_members_display(await _get_team_member_user_ids_from_team(team)) verbose_proxy_logger.debug(f"SCIM GET GROUPS members: {members}") team_alias = getattr(team, "team_alias", team.team_id) team_created_at = team.created_at.isoformat() if team.created_at else None @@ -1877,16 +1959,28 @@ async def delete_group( async def _process_group_patch_operations( patch_ops: SCIMPatchOp, existing_team, prisma_client -) -> Tuple[Dict[str, Any], Set[str]]: - """Process patch operations for a group and return update data and final members.""" +) -> Tuple[Dict[str, Any], Set[str], Set[str] | None]: + """Process patch operations for a group and return update data, final members + and, when the request contained a member ``replace`` op, the absolute target + roster it declared (``None`` otherwise). + + ``add``/``remove`` are deltas relative to the current roster, but ``replace`` + is absolute: it declares the roster is exactly this set, so the caller must + reconcile against it as a set-to-target rather than rebasing it onto a + concurrently-mutated roster. + """ update_data: Dict[str, Any] = {} # Create a fresh copy of existing metadata to avoid Prisma issues existing_metadata = existing_team.metadata or {} metadata = dict(existing_metadata) if existing_metadata else {} - # Track member changes - current_members = set(existing_team.members or []) + # Track member changes. members_with_roles is the source of truth for team + # membership; the legacy `members` column is not populated by team creation + # or the real team endpoints, so seeding from it would make an `add`/`remove` + # operation recompute the member set from an empty base and silently drop + # everyone already in the team. + current_members = set(await _get_team_member_user_ids_from_team(existing_team)) final_members = current_members.copy() # Process each patch operation @@ -1908,6 +2002,8 @@ async def _process_group_patch_operations( elif path.startswith("members"): # Handle member operations member_values = _extract_group_values(value) + if not member_values and value is None: + member_values = _extract_ids_from_path_filter(op.path, "members") # Check the feature flag scim_upsert_user = await _get_scim_upsert_user_setting() # Validate all users exist or create them based on feature flag @@ -1960,27 +2056,32 @@ async def _process_group_patch_operations( if metadata: update_data["metadata"] = metadata - return update_data, final_members + member_replace_present = any( + op.op == "replace" and (op.path or "").lower().startswith("members") for op in patch_ops.Operations + ) + replace_target = set(final_members) if member_replace_present else None + + return update_data, final_members, replace_target -async def _apply_group_patch_updates( - group_id: str, update_data: Dict[str, Any], final_members: Set[str], prisma_client -): - """Apply patch updates to the group in the database.""" - # Serialize metadata if present +async def _apply_group_patch_updates(group_id: str, update_data: Dict[str, Any], prisma_client): + """Apply the group's metadata/displayName patch updates to the database. + + Membership itself is not written here; it is reconciled onto the source of + truth (members_with_roles and each member's user.teams) by + _handle_group_membership_changes via team_member_add/team_member_delete. + Writing the legacy `members` column here too would create a second, unread + copy of membership that could drift from the source of truth. + """ if "metadata" in update_data and isinstance(update_data["metadata"], dict): update_data["metadata"] = safe_dumps(update_data["metadata"]) - # Update members list - update_data["members"] = list(final_members) - - # Update team in database - updated_team = await TeamRepository(prisma_client).table.update( - where={"team_id": group_id}, - data=update_data, - ) - - return updated_team + if update_data: + return await TeamRepository(prisma_client).table.update( + where={"team_id": group_id}, + data=update_data, + ) + return await TeamRepository(prisma_client).table.find_unique(where={"team_id": group_id}) async def _handle_group_membership_changes(group_id: str, current_members: Set[str], final_members: Set[str]): @@ -2031,27 +2132,29 @@ async def patch_group( existing_team = await _check_team_exists(group_id) # Process patch operations - update_data, final_members = await _process_group_patch_operations(patch_ops, existing_team, prisma_client) + update_data, final_members, replace_target = await _process_group_patch_operations( + patch_ops, existing_team, prisma_client + ) - # Track current members BEFORE update for comparison - current_members = set(await _get_team_member_user_ids_from_team(existing_team)) + snapshot_members = set(await _get_team_member_user_ids_from_team(existing_team)) + intended_add = final_members - snapshot_members + intended_remove = snapshot_members - final_members - # Apply updates to the database - updated_team = await _apply_group_patch_updates(group_id, update_data, final_members, prisma_client) + # Apply the metadata/displayName updates to the database + updated_team = await _apply_group_patch_updates(group_id, update_data, prisma_client) - # Refresh team data from database to get the latest state after concurrent updates - # This prevents race conditions when multiple PATCH requests come in simultaneously refreshed_team = await TeamRepository(prisma_client).table.find_unique(where={"team_id": group_id}) - if refreshed_team: - # Re-read current members from refreshed team to account for concurrent updates - refreshed_current_members = set( - await _get_team_member_user_ids_from_team(LiteLLM_TeamTable(**refreshed_team.model_dump())) - ) - # Use the refreshed members for comparison - current_members = refreshed_current_members + refreshed_current = ( + set(await _get_team_member_user_ids_from_team(LiteLLM_TeamTable(**refreshed_team.model_dump()))) + if refreshed_team + else snapshot_members + ) - # Handle user-team relationship changes - await _handle_group_membership_changes(group_id, current_members, final_members) + effective_final = ( + replace_target if replace_target is not None else (refreshed_current | intended_add) - intended_remove + ) + + await _handle_group_membership_changes(group_id, refreshed_current, effective_final) # A rename can flip whether this group matches scim_admin_group by display # name, so retained members must be re-resolved too, not just the ones whose @@ -2060,7 +2163,7 @@ async def patch_group( alias_changed = new_alias != existing_team.team_alias await _recompute_scim_member_roles( prisma_client, - (current_members | final_members if alias_changed else current_members ^ final_members), + (refreshed_current | effective_final if alias_changed else refreshed_current ^ effective_final), ) # Refresh team one more time to get final state after membership changes diff --git a/litellm/proxy/management_endpoints/sso/saml_sso.py b/litellm/proxy/management_endpoints/sso/saml_sso.py new file mode 100644 index 00000000000..37b641ca123 --- /dev/null +++ b/litellm/proxy/management_endpoints/sso/saml_sso.py @@ -0,0 +1,493 @@ +""" +SAML 2.0 SSO for the LiteLLM proxy admin UI. + +Supports both SP-initiated and IdP-initiated login via the HTTP-POST binding, +using the OneLogin python3-saml toolkit for signature, audience and time +validation. The IdP is configured from its metadata (``SAML_IDP_METADATA_URL`` +or inline ``SAML_IDP_METADATA_XML``); a successful login is mapped to a +``CustomOpenID`` and handed to the shared post-login path used by every other +SSO provider. + +python3-saml pulls in the native ``xmlsec``/``libxml2`` libraries, so it is an +optional dependency. When it is not installed the SAML routes return a clear +error instead of breaking proxy startup. +""" + +# python3-saml ships no type stubs, so the type checker sees every onelogin call +# as Unknown and the guarded optional import as possibly-unbound. Values crossing +# that boundary are cast() to concrete types at each use site; these directives +# silence only the unavoidable noise from the untyped dependency in this module. +# pyright: reportUnknownMemberType=false, reportUnknownVariableType=false +# pyright: reportUnknownArgumentType=false, reportUnknownParameterType=false +# pyright: reportMissingTypeStubs=false, reportPossiblyUnboundVariable=false +# pyright: reportConstantRedefinition=false + +import asyncio +import hashlib +import os +import secrets +import time +from typing import cast +from urllib.parse import parse_qsl + +from fastapi import HTTPException, Request, status +from fastapi.responses import RedirectResponse +from pydantic import ValidationError + +from litellm._logging import verbose_proxy_logger +from litellm.caching.dual_cache import DualCache +from litellm.proxy.management_endpoints.types import CustomOpenID, get_litellm_user_role +from litellm.proxy.utils import get_custom_url + +try: + from onelogin.saml2.auth import OneLogin_Saml2_Auth + from onelogin.saml2.idp_metadata_parser import OneLogin_Saml2_IdPMetadataParser + from onelogin.saml2.settings import OneLogin_Saml2_Settings + from onelogin.saml2.xml_utils import OneLogin_Saml2_XML + + SAML_AVAILABLE = True +except ImportError: + SAML_AVAILABLE = False + +SAML_LOGIN_ROUTE = "sso/saml/login" +SAML_CALLBACK_ROUTE = "sso/saml/callback" +SAML_METADATA_ROUTE = "sso/saml/metadata" + +_SAML_AUTHN_STATE_COOKIE = "litellm_saml_authn" +_SAML_IDP_SETTINGS_CACHE_PREFIX = "saml_idp_settings" +_SAML_AUTHN_REQUEST_CACHE_PREFIX = "saml_authn_request" +_SAML_CONSUMED_ASSERTION_CACHE_PREFIX = "saml_consumed_assertion" +_SAML_AUTHN_REQUEST_TTL_SECONDS = 600 +_SAML_IDP_METADATA_TTL_SECONDS = 3600 +_SAML_METADATA_FETCH_TIMEOUT_SECONDS = 10 +_SAML_MAX_POST_BYTES = 5 * 1024 * 1024 +# The replay guard tracks each assertion's NotOnOrAfter so it spans the full +# validity window; the floor covers IdPs that issue hour-long assertions or omit +# the timestamp, and the cap bounds cache growth. +_SAML_REPLAY_GUARD_DEFAULT_TTL_SECONDS = 3600 +_SAML_REPLAY_GUARD_MAX_TTL_SECONDS = 86400 + +_EMAIL_ATTRIBUTE_CANDIDATES = ( + "urn:oid:0.9.2342.19200300.100.1.3", + "http://schemas.xmlsoap.org/ws/2005/05/identity/claims/emailaddress", + "email", + "emailAddress", + "mail", + "Email", +) +_FIRST_NAME_ATTRIBUTE_CANDIDATES = ( + "urn:oid:2.5.4.42", + "http://schemas.xmlsoap.org/ws/2005/05/identity/claims/givenname", + "givenName", + "first_name", + "firstName", +) +_LAST_NAME_ATTRIBUTE_CANDIDATES = ( + "urn:oid:2.5.4.4", + "http://schemas.xmlsoap.org/ws/2005/05/identity/claims/surname", + "sn", + "surname", + "last_name", + "lastName", +) +_ROLE_ATTRIBUTE_CANDIDATES = ("role", "roles", "litellm_role") +_TEAM_IDS_ATTRIBUTE_CANDIDATES = ("teams", "team_ids", "groups") + + +def _saml_unavailable_error() -> HTTPException: + return HTTPException( + status_code=status.HTTP_501_NOT_IMPLEMENTED, + detail=( + "SAML SSO requires the optional 'python3-saml' dependency, which is " + "not installed. Re-install litellm with the saml extra: " + "'pip install litellm[saml]'. The saml extra bundles the native " + "xmlsec/libxml2 libraries, so no system packages are required." + ), + ) + + +class SAMLAuthHandler: + """SP- and IdP-initiated SAML 2.0 login for the admin UI.""" + + @staticmethod + def _env(name: str, default: str | None = None) -> str | None: + return os.getenv(name, default) + + @staticmethod + def is_saml_configured() -> bool: + return bool(SAMLAuthHandler._env("SAML_IDP_METADATA_URL") or SAMLAuthHandler._env("SAML_IDP_METADATA_XML")) + + @staticmethod + def _bool_env(name: str, default: bool) -> bool: + raw = SAMLAuthHandler._env(name) + if raw is None: + return default + return raw.strip().lower() in ("true", "1", "yes", "on") + + @staticmethod + def _base_url(request: Request) -> str: + base = get_custom_url(request_base_url=str(request.base_url)) + return base if base.endswith("/") else base + "/" + + @staticmethod + def _is_https(request: Request) -> bool: + return SAMLAuthHandler._base_url(request).startswith("https") + + @staticmethod + def _acs_url(request: Request) -> str: + return SAMLAuthHandler._base_url(request) + SAML_CALLBACK_ROUTE + + @staticmethod + def _metadata_url(request: Request) -> str: + return SAMLAuthHandler._base_url(request) + SAML_METADATA_ROUTE + + @staticmethod + def _sp_entity_id(request: Request) -> str: + return SAMLAuthHandler._env("SAML_SP_ENTITY_ID") or SAMLAuthHandler._metadata_url(request) + + @staticmethod + async def _load_idp_settings(cache: DualCache) -> dict[str, object]: + metadata_url = SAMLAuthHandler._env("SAML_IDP_METADATA_URL") + metadata_xml = SAMLAuthHandler._env("SAML_IDP_METADATA_XML") + source = metadata_url or metadata_xml + if source is None: + raise HTTPException( + status_code=status.HTTP_501_NOT_IMPLEMENTED, + detail="SAML SSO is not configured. Set SAML_IDP_METADATA_URL or SAML_IDP_METADATA_XML.", + ) + + cache_key = f"{_SAML_IDP_SETTINGS_CACHE_PREFIX}:{hashlib.sha256(source.encode()).hexdigest()}" + cached = cache.get_cache(key=cache_key) + if isinstance(cached, dict): + return cast(dict[str, object], cached) # cast-ok: untyped python3-saml + + if metadata_url is not None: + parsed = await asyncio.to_thread( + OneLogin_Saml2_IdPMetadataParser.parse_remote, + metadata_url, + validate_cert=SAMLAuthHandler._bool_env("SAML_IDP_METADATA_VALIDATE_CERT", True), + timeout=_SAML_METADATA_FETCH_TIMEOUT_SECONDS, + ) + else: + parsed = OneLogin_Saml2_IdPMetadataParser.parse(cast(str, metadata_xml)) # cast-ok: untyped python3-saml + + idp_settings = cast(dict[str, object], parsed) # cast-ok: untyped python3-saml + if not idp_settings.get("idp"): + raise HTTPException( + status_code=status.HTTP_502_BAD_GATEWAY, + detail="Could not parse an IdP entityID/SSO URL/certificate from the SAML metadata.", + ) + cache.set_cache(key=cache_key, value=idp_settings, ttl=_SAML_IDP_METADATA_TTL_SECONDS) + return idp_settings + + @staticmethod + def _build_settings(request: Request, idp_settings: dict[str, object]) -> dict[str, object]: + sp_settings: dict[str, object] = { + "strict": SAMLAuthHandler._bool_env("SAML_STRICT", True), + "debug": False, + "sp": { + "entityId": SAMLAuthHandler._sp_entity_id(request), + "assertionConsumerService": { + "url": SAMLAuthHandler._acs_url(request), + "binding": "urn:oasis:names:tc:SAML:2.0:bindings:HTTP-POST", + }, + "NameIDFormat": SAMLAuthHandler._env( + "SAML_SP_NAME_ID_FORMAT", + "urn:oasis:names:tc:SAML:1.1:nameid-format:emailAddress", + ), + }, + "security": { + "wantAssertionsSigned": SAMLAuthHandler._bool_env("SAML_WANT_ASSERTIONS_SIGNED", True), + "wantMessagesSigned": SAMLAuthHandler._bool_env("SAML_WANT_MESSAGES_SIGNED", False), + "authnRequestsSigned": SAMLAuthHandler._bool_env("SAML_AUTHN_REQUESTS_SIGNED", False), + "wantNameId": True, + "requestedAuthnContext": False, + "rejectUnsolicitedResponsesWithInResponseTo": False, + }, + } + return OneLogin_Saml2_IdPMetadataParser.merge_settings(sp_settings, idp_settings) + + @staticmethod + def _prepare_request_data(request: Request, post_data: dict[str, str] | None = None) -> dict[str, object]: + base = SAMLAuthHandler._base_url(request) + scheme, _, host_part = base.partition("://") + host = host_part.split("/", 1)[0] + return { + "https": "on" if scheme == "https" else "off", + "http_host": host, + "script_name": "/" + SAML_CALLBACK_ROUTE, + "get_data": dict(request.query_params), + "post_data": post_data or {}, + } + + @staticmethod + async def _build_auth( + request: Request, + cache: DualCache, + post_data: dict[str, str] | None = None, + ) -> "OneLogin_Saml2_Auth": + if not SAML_AVAILABLE: + raise _saml_unavailable_error() + idp_settings = await SAMLAuthHandler._load_idp_settings(cache) + settings = SAMLAuthHandler._build_settings(request, idp_settings) + request_data = SAMLAuthHandler._prepare_request_data(request, post_data) + try: + return OneLogin_Saml2_Auth(request_data, old_settings=settings) + except Exception as e: # noqa: BLE001 - toolkit exposes no common exception base; fail closed + raise HTTPException( + status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, + detail=f"Invalid SAML configuration: {e}", + ) + + @staticmethod + async def build_login_redirect( + request: Request, cache: DualCache, relay_state: str | None = None + ) -> RedirectResponse: + auth = await SAMLAuthHandler._build_auth(request, cache) + redirect_url = cast(str, auth.login(return_to=relay_state)) # cast-ok: untyped python3-saml + response = RedirectResponse(url=redirect_url, status_code=303) + request_id = cast(str | None, auth.get_last_request_id()) # cast-ok: untyped python3-saml + if request_id is not None: + cache.set_cache( + key=f"{_SAML_AUTHN_REQUEST_CACHE_PREFIX}:{request_id}", + value="1", + ttl=_SAML_AUTHN_REQUEST_TTL_SECONDS, + ) + secure = SAMLAuthHandler._is_https(request) + response.set_cookie( + key=_SAML_AUTHN_STATE_COOKIE, + value=request_id, + max_age=_SAML_AUTHN_REQUEST_TTL_SECONDS, + httponly=True, + secure=secure, + samesite="none" if secure else "lax", + ) + return response + + @staticmethod + async def build_sp_metadata(request: Request, cache: DualCache) -> str: + if not SAML_AVAILABLE: + raise _saml_unavailable_error() + idp_settings = await SAMLAuthHandler._load_idp_settings(cache) + settings = SAMLAuthHandler._build_settings(request, idp_settings) + saml_settings = OneLogin_Saml2_Settings(settings, sp_validation_only=True) + metadata = cast(str, saml_settings.get_sp_metadata()) # cast-ok: untyped python3-saml + errors = cast(list[str], saml_settings.validate_metadata(metadata)) # cast-ok: untyped python3-saml + if errors: + raise HTTPException( + status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, + detail=f"Invalid SP metadata: {', '.join(errors)}", + ) + return metadata + + @staticmethod + async def read_acs_post_data(request: Request) -> dict[str, str]: + """Read the ACS POST form under a hard size cap before any base64/XML decoding. + + Bounds both Content-Length-declared and chunked requests so an unauthenticated + caller cannot force unbounded buffering while decoding the SAMLResponse.""" + declared = request.headers.get("content-length") + if declared is not None and declared.isdigit() and int(declared) > _SAML_MAX_POST_BYTES: + raise HTTPException( + status_code=status.HTTP_413_CONTENT_TOO_LARGE, + detail="SAML response exceeds the maximum allowed size.", + ) + + body = bytearray() + async for chunk in request.stream(): + body += chunk + if len(body) > _SAML_MAX_POST_BYTES: + raise HTTPException( + status_code=status.HTTP_413_CONTENT_TOO_LARGE, + detail="SAML response exceeds the maximum allowed size.", + ) + + return dict(parse_qsl(body.decode("utf-8", "replace"))) + + @staticmethod + async def handle_acs(request: Request, cache: DualCache, post_data: dict[str, str]) -> CustomOpenID: + auth = await SAMLAuthHandler._build_auth(request, cache, post_data=post_data) + browser_request_id = request.cookies.get(_SAML_AUTHN_STATE_COOKIE) + try: + auth.process_response(request_id=browser_request_id) + except Exception as e: # noqa: BLE001 - toolkit exposes no common exception base; fail closed + raise HTTPException( + status_code=status.HTTP_401_UNAUTHORIZED, + detail=f"Could not process SAML response: {e}", + ) + + errors = cast(list[str], auth.get_errors()) # cast-ok: untyped python3-saml + if errors or not auth.is_authenticated(): + reason = auth.get_last_error_reason() + raise HTTPException( + status_code=status.HTTP_401_UNAUTHORIZED, + detail=f"SAML authentication failed: {reason or ', '.join(errors)}", + ) + + await SAMLAuthHandler._enforce_response_binding(auth, cache, browser_request_id) + return SAMLAuthHandler._result_from_auth(auth) + + @staticmethod + def _replay_guard_ttl(auth: "OneLogin_Saml2_Auth") -> int: + not_on_or_after = auth.get_last_assertion_not_on_or_after() + if not isinstance(not_on_or_after, int): + return _SAML_REPLAY_GUARD_DEFAULT_TTL_SECONDS + remaining = not_on_or_after - int(time.time()) + return min( + max(remaining, _SAML_REPLAY_GUARD_DEFAULT_TTL_SECONDS), + _SAML_REPLAY_GUARD_MAX_TTL_SECONDS, + ) + + @staticmethod + def _response_in_response_to(auth: "OneLogin_Saml2_Auth") -> str | None: + """The request id this response answers, read from the Response element or, when the + IdP only stamps it on the bearer SubjectConfirmationData, from there. A non-None value + marks the response as solicited (SP-initiated) and so requiring browser binding.""" + value = cast(str | None, auth.get_last_response_in_response_to()) # cast-ok: untyped python3-saml + if value: + return value + xml = cast(bytes | None, auth.get_last_response_xml()) # cast-ok: untyped python3-saml + if not xml: + return None + root = OneLogin_Saml2_XML.to_etree(xml) + for node in OneLogin_Saml2_XML.query(root, "//saml:SubjectConfirmationData[@InResponseTo]"): + irt = cast(str | None, node.get("InResponseTo")) # cast-ok: untyped python3-saml + if irt: + return irt + return None + + @staticmethod + async def _enforce_response_binding( + auth: "OneLogin_Saml2_Auth", + cache: DualCache, + browser_request_id: str | None, + ) -> None: + in_response_to = SAMLAuthHandler._response_in_response_to(auth) + + if in_response_to is not None: + authn_key = f"{_SAML_AUTHN_REQUEST_CACHE_PREFIX}:{in_response_to}" + if cache.get_cache(key=authn_key) is None: + raise HTTPException( + status_code=status.HTTP_401_UNAUTHORIZED, + detail="SAML response references an unknown or already-used login request.", + ) + if browser_request_id is None or not secrets.compare_digest(browser_request_id, in_response_to): + raise HTTPException( + status_code=status.HTTP_401_UNAUTHORIZED, + detail="SAML response is not bound to this browser's login request.", + ) + elif browser_request_id is not None: + raise HTTPException( + status_code=status.HTTP_401_UNAUTHORIZED, + detail="SAML response is not bound to this browser's login request.", + ) + elif not SAMLAuthHandler._bool_env("SAML_ALLOW_UNSOLICITED", False): + raise HTTPException( + status_code=status.HTTP_401_UNAUTHORIZED, + detail="Unsolicited (IdP-initiated) SAML responses are disabled.", + ) + elif cache.redis_cache is None: + raise HTTPException( + status_code=status.HTTP_401_UNAUTHORIZED, + detail=( + "Unsolicited (IdP-initiated) SAML responses require a shared Redis cache " + "so the replay guard is enforced across every worker." + ), + ) + + assertion_id = cast(str | None, auth.get_last_assertion_id()) # cast-ok: untyped python3-saml + if assertion_id is None: + raise HTTPException( + status_code=status.HTTP_401_UNAUTHORIZED, + detail="SAML assertion is missing the required ID attribute.", + ) + consumed_key = f"{_SAML_CONSUMED_ASSERTION_CACHE_PREFIX}:{assertion_id}" + consumed_count = await cache.async_increment_cache( + key=consumed_key, value=1, ttl=SAMLAuthHandler._replay_guard_ttl(auth) + ) + if consumed_count is not None and consumed_count > 1: + raise HTTPException( + status_code=status.HTTP_401_UNAUTHORIZED, + detail="SAML assertion has already been used (replay detected).", + ) + + @staticmethod + def _result_from_auth(auth: "OneLogin_Saml2_Auth") -> CustomOpenID: + attributes = cast(dict[str, list[str]], auth.get_attributes()) # cast-ok: untyped python3-saml + name_id = cast(str | None, auth.get_nameid()) # cast-ok: untyped python3-saml + + email = SAMLAuthHandler._attribute_value(attributes, "SAML_ATTRIBUTE_EMAIL", _EMAIL_ATTRIBUTE_CANDIDATES) + if email is None and name_id is not None and "@" in name_id: + email = name_id + + if email is None and SAMLAuthHandler._env("ALLOWED_EMAIL_DOMAINS") is not None: + raise HTTPException( + status_code=status.HTTP_401_UNAUTHORIZED, + detail=( + "SAML assertion did not contain an email address, but ALLOWED_EMAIL_DOMAINS " + "restricts sign-in by email domain." + ), + ) + + user_id = SAMLAuthHandler._attribute_value(attributes, "SAML_ATTRIBUTE_USER_ID", ()) or name_id or email + if user_id is None: + raise HTTPException( + status_code=status.HTTP_401_UNAUTHORIZED, + detail="SAML assertion did not contain a usable subject (NameID) or email.", + ) + + first_name = SAMLAuthHandler._attribute_value( + attributes, "SAML_ATTRIBUTE_FIRST_NAME", _FIRST_NAME_ATTRIBUTE_CANDIDATES + ) + last_name = SAMLAuthHandler._attribute_value( + attributes, "SAML_ATTRIBUTE_LAST_NAME", _LAST_NAME_ATTRIBUTE_CANDIDATES + ) + role_value = SAMLAuthHandler._attribute_value(attributes, "SAML_ATTRIBUTE_ROLE", _ROLE_ATTRIBUTE_CANDIDATES) + team_ids = SAMLAuthHandler._attribute_values( + attributes, "SAML_ATTRIBUTE_TEAM_IDS", _TEAM_IDS_ATTRIBUTE_CANDIDATES + ) + + display_name = " ".join(part for part in (first_name, last_name) if part) or email + + verbose_proxy_logger.info(f"SAML login: subject={user_id}, email={email}, attributes={list(attributes.keys())}") + + try: + return CustomOpenID( + id=user_id, + email=email, + first_name=first_name, + last_name=last_name, + display_name=display_name, + picture=None, + provider="saml", + team_ids=team_ids, + user_role=get_litellm_user_role(role_value) if role_value else None, + ) + except ValidationError as e: + raise HTTPException( + status_code=status.HTTP_401_UNAUTHORIZED, + detail=f"SAML assertion contained an invalid subject or email: {e}", + ) + + @staticmethod + def _attribute_value( + attributes: dict[str, list[str]], + env_override: str, + candidates: tuple[str, ...], + ) -> str | None: + values = SAMLAuthHandler._attribute_values(attributes, env_override, candidates) + return values[0] if values else None + + @staticmethod + def _attribute_values( + attributes: dict[str, list[str]], + env_override: str, + candidates: tuple[str, ...], + ) -> list[str]: + override = SAMLAuthHandler._env(env_override) + keys = (override, *candidates) if override else candidates + for key in keys: + values = attributes.get(key) + if values: + return [v for v in values if v] + return [] diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index a3ab210ebab..50f4437445c 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -47,6 +47,7 @@ from litellm.proxy._types import ( Member, NewTeamRequest, OrgMember, + PatchTeamRequest, ProxyErrorTypes, ProxyException, SpecialManagementEndpointEnums, @@ -2005,6 +2006,7 @@ async def update_team( ) async def patch_team( team_id: str, + data: PatchTeamRequest, http_request: Request, user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)], litellm_changed_by: Annotated[ @@ -2017,11 +2019,12 @@ async def patch_team( """ Partially update a team using RFC 7386 JSON Merge Patch semantics. - `team_id` is taken from the path. `metadata` is merged with the team's stored - metadata rather than replacing it: an omitted key is preserved, `key: null` - deletes it, and any other value overwrites (recursing into nested objects). - Every other field behaves exactly like `POST /team/update` (omitted preserves, - a value overwrites). Returns the full updated team. + `team_id` is taken from the path; a `team_id` in the body is accepted only when it + matches. `metadata` is merged with the team's stored metadata rather than replacing + it: an omitted key is preserved, `key: null` deletes it, and any other value + overwrites (recursing into nested objects). Every other field behaves exactly like + `POST /team/update` (omitted preserves, a value overwrites). Returns the full + updated team. ``` curl --location --request PATCH 'http://0.0.0.0:4000/team/8d916b1c-510d-4894-a334-1c16a93344f5' \ @@ -2041,21 +2044,15 @@ async def patch_team( detail={"error": CommonProxyErrors.db_not_connected_error.value}, ) - try: - body = await http_request.json() - except (json.JSONDecodeError, ValueError): - raise HTTPException(status_code=400, detail={"error": "Request body must be a JSON object"}) - if not isinstance(body, dict): - raise HTTPException(status_code=400, detail={"error": "Request body must be a JSON object"}) - - body_team_id = body.pop("team_id", None) - if body_team_id is not None and body_team_id != team_id: + if data.team_id is not None and data.team_id != team_id: raise HTTPException( status_code=400, - detail={"error": f"team_id in body ({body_team_id}) does not match team_id in path ({team_id})"}, + detail={"error": f"team_id in body ({data.team_id}) does not match team_id in path ({team_id})"}, ) - if "metadata" in body: + patch_fields = data.model_dump(exclude_unset=True, exclude={"team_id"}) + + if "metadata" in patch_fields: existing_team_row = await TeamRepository(prisma_client).table.find_unique(where={"team_id": team_id}) if existing_team_row is None: raise HTTPException( @@ -2063,9 +2060,9 @@ async def patch_team( detail={"error": f"Team not found, passed team_id={team_id}"}, ) existing_metadata = existing_team_row.metadata if isinstance(existing_team_row.metadata, dict) else {} - body["metadata"] = apply_json_merge_patch(existing_metadata, body["metadata"]) + patch_fields["metadata"] = apply_json_merge_patch(existing_metadata, patch_fields["metadata"]) - update_request = UpdateTeamRequest(team_id=team_id, **body) + update_request = UpdateTeamRequest(team_id=team_id, **patch_fields) result = await update_team( data=update_request, @@ -2424,7 +2421,15 @@ async def _add_team_members_to_team( user_api_key_dict: UserAPIKeyAuth, litellm_proxy_admin_name: str, ) -> Tuple[LiteLLM_TeamTable, List[LiteLLM_UserTable], List[LiteLLM_TeamMembership]]: - """Add team members to the team.""" + """Add team members to the team. + + The members_with_roles reconciliation runs inside a transaction that locks + the team row with ``SELECT ... FOR UPDATE`` before reading the current + membership. Concurrent /team/member_add calls for the same team therefore + serialize on the row lock and each appends onto the other's committed + result, instead of both rewriting the whole JSON array from a stale + snapshot (which silently drops one member on the losing write). + """ # Process and add new members updated_users, updated_team_memberships = await _process_team_members( data=data, @@ -2434,19 +2439,22 @@ async def _add_team_members_to_team( litellm_proxy_admin_name=litellm_proxy_admin_name, ) - # Update team members list - await _update_team_members_list( - data=data, - complete_team_data=complete_team_data, - updated_users=updated_users, - ) + async with prisma_client.tx() as tx: + complete_team_data.members_with_roles = await TeamRepository(prisma_client).get_members_with_roles_locked( + tx, data.team_id + ) - # ADD MEMBER TO TEAM - _db_team_members = [m.model_dump() for m in complete_team_data.members_with_roles] - updated_team = await TeamRepository(prisma_client).table.update( - where={"team_id": data.team_id}, - data={"members_with_roles": json.dumps(_db_team_members)}, # type: ignore - ) + await _update_team_members_list( + data=data, + complete_team_data=complete_team_data, + updated_users=updated_users, + ) + + _db_team_members = [m.model_dump() for m in complete_team_data.members_with_roles] + updated_team = await tx.litellm_teamtable.update( + where={"team_id": data.team_id}, + data={"members_with_roles": json.dumps(_db_team_members)}, + ) return updated_team, updated_users, updated_team_memberships diff --git a/litellm/proxy/management_endpoints/tool_management_endpoints.py b/litellm/proxy/management_endpoints/tool_management_endpoints.py index 9d71761f115..ca606e07cee 100644 --- a/litellm/proxy/management_endpoints/tool_management_endpoints.py +++ b/litellm/proxy/management_endpoints/tool_management_endpoints.py @@ -10,16 +10,18 @@ POST /v1/tool/policy - Update the input_policy / output_policy for a """ import uuid -from datetime import datetime, timezone -from typing import TYPE_CHECKING, Any, List, Optional +from datetime import datetime, timedelta, timezone +from itertools import groupby +from typing import TYPE_CHECKING, Annotated, Any, List, Optional from fastapi import APIRouter, Depends, HTTPException, Query +from pydantic import BaseModel, TypeAdapter if TYPE_CHECKING: from litellm.proxy.utils import PrismaClient from litellm._logging import verbose_proxy_logger -from litellm.proxy._types import CommonProxyErrors, UserAPIKeyAuth +from litellm.proxy._types import CommonProxyErrors, LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.repositories.object_permission_repository import ObjectPermissionRepository from litellm.repositories.table_repositories import ( @@ -39,6 +41,9 @@ from litellm.types.tool_management import ( ToolPolicyOptionsResponse, ToolPolicyUpdateRequest, ToolPolicyUpdateResponse, + ToolSpendDailyEntry, + ToolSpendEntry, + ToolSpendResponse, ToolUsageLogEntry, ToolUsageLogsResponse, ) @@ -124,6 +129,147 @@ async def list_tools( raise HTTPException(status_code=500, detail=str(e)) +def _parse_day_start(value: str | None) -> datetime | None: + if not value: + return None + try: + return datetime.strptime(value.strip(), "%Y-%m-%d").replace(tzinfo=timezone.utc) + except ValueError: + raise HTTPException( + status_code=400, + detail=f"Invalid date format: {value}. Expected: 'YYYY-MM-DD'", + ) + + +class _ToolSpendRow(BaseModel): + date: str + tool_name: str + call_count: int + spend: float + total_tokens: int + + +class _RequestTotalRow(BaseModel): + total_spend: float + + +_TOOL_SPEND_ROWS = TypeAdapter(list[_ToolSpendRow]) +_REQUEST_TOTAL_ROWS = TypeAdapter(list[_RequestTotalRow]) + + +def _summarize_tool(name: str, grp: tuple[_ToolSpendRow, ...]) -> ToolSpendEntry: + return ToolSpendEntry( + tool_name=name, + spend=sum(r.spend for r in grp), + call_count=sum(r.call_count for r in grp), + total_tokens=sum(r.total_tokens for r in grp), + ) + + +def _build_tool_spend_response( + rows: list[_ToolSpendRow], + total_spend: float, + start_date: str, + end_date: str, +) -> ToolSpendResponse: + daily = [ + ToolSpendDailyEntry(date=r.date, tool_name=r.tool_name, spend=r.spend, call_count=r.call_count) for r in rows + ] + grouped = groupby(sorted(rows, key=lambda r: r.tool_name), key=lambda r: r.tool_name) + by_tool = sorted( + (_summarize_tool(name, tuple(grp)) for name, grp in grouped), + key=lambda e: e.spend, + reverse=True, + ) + return ToolSpendResponse( + by_tool=by_tool, + daily=daily, + total_spend=total_spend, + start_date=start_date, + end_date=end_date, + ) + + +@router.get( + "/v1/tool/spend", + tags=["tool management"], + dependencies=[Depends(user_api_key_auth)], + response_model=ToolSpendResponse, +) +async def get_tool_spend( + user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)], + start_date: Annotated[str | None, Query(description="YYYY-MM-DD (defaults to 30 days ago)")] = None, + end_date: Annotated[str | None, Query(description="YYYY-MM-DD (defaults to today)")] = None, +): + """ + Spend attributed to each tool over a date range, for the Cost Optimization dashboard. + + Joins ``LiteLLM_SpendLogToolIndex`` (which tool names ran on which request) to + ``LiteLLM_SpendLogs`` (what the request cost). A request that used multiple tools + counts its full spend toward each of those tools, so per-tool numbers are + attributions. ``total_spend`` is the deduplicated spend of every request that + called at least one tool in the window, so it never double counts. + """ + from litellm.proxy.proxy_server import prisma_client + + if user_api_key_dict.user_role not in ( + LitellmUserRoles.PROXY_ADMIN, + LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY, + ): + raise HTTPException( + status_code=403, + detail="Only proxy admin roles can view tool spend across the deployment", + ) + + if prisma_client is None: + raise HTTPException(status_code=500, detail=CommonProxyErrors.db_not_connected_error.value) + + now = datetime.now(timezone.utc) + end_day = _parse_day_start(end_date) + start_dt = _parse_day_start(start_date) or ((end_day or now) - timedelta(days=30)) + end_exclusive = (end_day + timedelta(days=1)) if end_day else now + + rows = await prisma_client.db.query_raw( + """ + SELECT to_char(ti.start_time, 'YYYY-MM-DD') AS date, + ti.tool_name AS tool_name, + COUNT(*)::int AS call_count, + COALESCE(SUM(sl.spend), 0)::double precision AS spend, + COALESCE(SUM(sl.total_tokens), 0)::bigint AS total_tokens + FROM "LiteLLM_SpendLogToolIndex" ti + JOIN "LiteLLM_SpendLogs" sl ON sl.request_id = ti.request_id + WHERE ti.start_time >= ($1::timestamptz AT TIME ZONE 'UTC') + AND ti.start_time < ($2::timestamptz AT TIME ZONE 'UTC') + GROUP BY date, ti.tool_name + ORDER BY date ASC, spend DESC + """, + start_dt.isoformat(), + end_exclusive.isoformat(), + ) + totals = await prisma_client.db.query_raw( + """ + SELECT COALESCE(SUM(sl.spend), 0)::double precision AS total_spend + FROM "LiteLLM_SpendLogs" sl + WHERE EXISTS ( + SELECT 1 + FROM "LiteLLM_SpendLogToolIndex" ti + WHERE ti.request_id = sl.request_id + AND ti.start_time >= ($1::timestamptz AT TIME ZONE 'UTC') + AND ti.start_time < ($2::timestamptz AT TIME ZONE 'UTC') + ) + """, + start_dt.isoformat(), + end_exclusive.isoformat(), + ) + total_rows = _REQUEST_TOTAL_ROWS.validate_python(totals or []) + return _build_tool_spend_response( + rows=_TOOL_SPEND_ROWS.validate_python(rows or []), + total_spend=total_rows[0].total_spend if total_rows else 0.0, + start_date=start_dt.strftime("%Y-%m-%d"), + end_date=(end_day or now).strftime("%Y-%m-%d"), + ) + + @router.get( "/v1/tool/{tool_name:path}/detail", tags=["tool management"], diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index 3c8444ecf26..8682b61f910 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -12,6 +12,7 @@ import asyncio import base64 import hashlib import inspect +import json import os import re import secrets @@ -35,7 +36,7 @@ if TYPE_CHECKING: import httpx import jwt -from fastapi import APIRouter, Depends, Header, HTTPException, Request, status +from fastapi import APIRouter, Depends, Header, HTTPException, Request, Response, status from fastapi.responses import RedirectResponse import litellm @@ -99,6 +100,7 @@ from litellm.proxy.common_utils.html_forms.ui_login import build_ui_login_form from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.proxy.management_endpoints.internal_user_endpoints import new_user from litellm.proxy.management_endpoints.sso import CustomMicrosoftSSO +from litellm.proxy.management_endpoints.sso.saml_sso import SAMLAuthHandler from litellm.proxy.management_endpoints.sso_helper_utils import ( check_is_admin_only_access, has_admin_ui_access, @@ -258,11 +260,20 @@ def _get_cli_sso_flow_or_raise(login_id: Optional[str], cache: DualCache) -> dic raise HTTPException(status_code=400, detail="Invalid CLI login session id") cache_key = _get_cli_sso_flow_cache_key(cast(str, login_id)) - flow = cache.get_cache(key=cache_key) + redis_cache = cache.redis_cache + if redis_cache is not None: + flow = redis_cache.get_cache(key=cache_key) + else: + flow = cache.get_cache(key=cache_key) + if isinstance(flow, str): + try: + flow = json.loads(flow) + except ValueError: + flow = None if not isinstance(flow, dict) or "poll_secret_hash" not in flow: verbose_proxy_logger.warning( "CLI SSO login session not found in cache for login_id=%s. If the proxy runs multiple replicas, " - "a shared Redis cache (enable_redis_auth_cache: true) is required for CLI login to work.", + "a shared Redis cache is required for CLI login to work.", login_id, ) raise HTTPException( @@ -270,7 +281,7 @@ def _get_cli_sso_flow_or_raise(login_id: Optional[str], cache: DualCache) -> dic detail=( "CLI login session not found or expired. Run `litellm-proxy login` again. " "If this happens immediately after starting a login, the proxy is likely running multiple " - "replicas without a shared cache; configure Redis with `enable_redis_auth_cache: true` " + "replicas without a shared cache; configure a Redis cache " "so every replica can see the login session." ), ) @@ -278,11 +289,12 @@ def _get_cli_sso_flow_or_raise(login_id: Optional[str], cache: DualCache) -> dic def _set_cli_sso_flow(login_id: str, cache: DualCache, flow: dict) -> None: - cache.set_cache( - key=_get_cli_sso_flow_cache_key(login_id), - value=flow, - ttl=CLI_SSO_SESSION_TTL_SECONDS, - ) + cache_key = _get_cli_sso_flow_cache_key(login_id) + redis_cache = cache.redis_cache + if redis_cache is not None: + redis_cache.set_cache(key=cache_key, value=json.dumps(flow), ttl=CLI_SSO_SESSION_TTL_SECONDS) + else: + cache.set_cache(key=cache_key, value=flow, ttl=CLI_SSO_SESSION_TTL_SECONDS) def _verify_cli_sso_poll_secret(flow: dict, poll_secret: Optional[str]) -> bool: @@ -593,11 +605,11 @@ def _render_cli_sso_verification_page( @router.post("/sso/cli/start", tags=["experimental"], include_in_schema=False) async def cli_sso_start(request: Request): - from litellm.proxy.proxy_server import general_settings, user_api_key_cache + from litellm.proxy.proxy_server import cli_sso_session_cache, general_settings _check_cli_sso_start_rate_limit( request=request, - cache=user_api_key_cache, + cache=cli_sso_session_cache, use_x_forwarded_for=bool((general_settings or {}).get("use_x_forwarded_for", False)), ) @@ -612,7 +624,7 @@ async def cli_sso_start(request: Request): "user_code_verified": False, "session_data": None, } - _set_cli_sso_flow(login_id=login_id, cache=user_api_key_cache, flow=flow) + _set_cli_sso_flow(login_id=login_id, cache=cli_sso_session_cache, flow=flow) verification_uri_complete: str | None = ( ( @@ -644,9 +656,9 @@ async def cli_sso_complete(request: Request, login_id: str): from litellm.proxy.common_utils.html_forms.cli_sso_success import ( render_cli_sso_success_page, ) - from litellm.proxy.proxy_server import user_api_key_cache + from litellm.proxy.proxy_server import cli_sso_session_cache - flow = _get_cli_sso_flow_or_raise(login_id=login_id, cache=user_api_key_cache) + flow = _get_cli_sso_flow_or_raise(login_id=login_id, cache=cli_sso_session_cache) if not flow.get("sso_complete") or not flow.get("session_data"): raise HTTPException(status_code=400, detail="CLI login is not ready") @@ -670,7 +682,7 @@ async def cli_sso_complete(request: Request, login_id: str): raise HTTPException(status_code=400, detail="Invalid verification code") flow["user_code_verified"] = True - _set_cli_sso_flow(login_id=login_id, cache=user_api_key_cache, flow=flow) + _set_cli_sso_flow(login_id=login_id, cache=cli_sso_session_cache, flow=flow) html_content = render_cli_sso_success_page() return HTMLResponse(content=html_content, status_code=200) @@ -846,6 +858,27 @@ def process_sso_jwt_access_token( return None +async def _raise_if_sso_exceeds_free_user_limit(premium_user: bool, prisma_client: PrismaClient | None) -> None: + """Free tier allows SSO for up to 5 billable users; beyond that requires an Enterprise license.""" + if premium_user is True: + return + if prisma_client is None: + raise ProxyException( + message=CommonProxyErrors.db_not_connected_error.value, + type=ProxyErrorTypes.auth_error, + param="premium_user", + code=status.HTTP_403_FORBIDDEN, + ) + billable_users = await UserRepository(prisma_client).count_billable_users() + if billable_users and billable_users > 5: + raise ProxyException( + message="You must be a LiteLLM Enterprise user to use SSO for more than 5 users. If you have a license please set `LITELLM_LICENSE` in your env. If you want to obtain a license meet with us here: https://enterprise.litellm.ai/demo You are seeing this error message because You configured SSO (one of `MICROSOFT_CLIENT_ID`, `GOOGLE_CLIENT_ID`, `GENERIC_CLIENT_ID`, or SAML) in your env. Please unset it", + type=ProxyErrorTypes.auth_error, + param="premium_user", + code=status.HTTP_403_FORBIDDEN, + ) + + @router.get("/sso/key/generate", tags=["experimental"], include_in_schema=False) async def google_login( request: Request, @@ -861,6 +894,7 @@ async def google_login( Example: """ from litellm.proxy.proxy_server import ( + cli_sso_session_cache, general_settings, premium_user, prisma_client, @@ -880,25 +914,13 @@ async def google_login( return admin_ui_disabled() ####### Check if user is a Enterprise / Premium User ####### - if microsoft_client_id is not None or google_client_id is not None or generic_client_id is not None: - if premium_user is not True: - # Check if under 'free SSO user' limit - if prisma_client is not None: - billable_users = await UserRepository(prisma_client).count_billable_users() - if billable_users and billable_users > 5: - raise ProxyException( - message="You must be a LiteLLM Enterprise user to use SSO for more than 5 users. If you have a license please set `LITELLM_LICENSE` in your env. If you want to obtain a license meet with us here: https://enterprise.litellm.ai/demo You are seeing this error message because You set one of `MICROSOFT_CLIENT_ID`, `GOOGLE_CLIENT_ID`, or `GENERIC_CLIENT_ID` in your env. Please unset this", - type=ProxyErrorTypes.auth_error, - param="premium_user", - code=status.HTTP_403_FORBIDDEN, - ) - else: - raise ProxyException( - message=CommonProxyErrors.db_not_connected_error.value, - type=ProxyErrorTypes.auth_error, - param="premium_user", - code=status.HTTP_403_FORBIDDEN, - ) + if ( + microsoft_client_id is not None + or google_client_id is not None + or generic_client_id is not None + or SAMLAuthHandler.is_saml_configured() + ): + await _raise_if_sso_exceeds_free_user_limit(premium_user, prisma_client) ####### Detect DB + MASTER KEY in .env ####### missing_env_vars = show_missing_vars_in_env() @@ -912,7 +934,7 @@ async def google_login( ) if source == LITELLM_CLI_SOURCE_IDENTIFIER: - _get_cli_sso_flow_or_raise(login_id=key, cache=user_api_key_cache) + _get_cli_sso_flow_or_raise(login_id=key, cache=cli_sso_session_cache) # Store CLI login handle in state for OAuth flow cli_state: Optional[str] = SSOAuthenticationHandler._get_cli_state( @@ -936,6 +958,19 @@ async def google_login( "Enterprise features are not available. Custom UI SSO sign-in requires LiteLLM Enterprise." ) + if ( + microsoft_client_id is None + and google_client_id is None + and generic_client_id is None + and SAMLAuthHandler.is_saml_configured() + ): + verbose_proxy_logger.info("Redirecting to SAML SSO login") + return await SAMLAuthHandler.build_login_redirect( + request=request, + cache=user_api_key_cache, + relay_state=return_to, + ) + # Check if we should use SSO handler if ( SSOAuthenticationHandler.should_use_sso_handler( @@ -954,15 +989,8 @@ async def google_login( state=cli_state, request=request, ) - if return_to is not None and sso_redirect is not None: - if SSOAuthenticationHandler._validate_return_to(return_to): - sso_redirect.set_cookie( - key="litellm_cp_return_to", - value=return_to, - max_age=600, - httponly=True, - samesite="lax", - ) + if sso_redirect is not None: + _persist_return_to_cookie(sso_redirect, return_to) return sso_redirect from fastapi.responses import HTMLResponse @@ -971,13 +999,19 @@ async def google_login( os.getenv("LITELLM_HIDE_DEFAULT_CREDENTIALS_HINT", "false").lower() == "true" or general_settings.get("hide_default_credentials_hint", False) is True ) - return HTMLResponse( + form_response = HTMLResponse( content=build_ui_login_form( show_deprecation_banner=True, hide_default_credentials_hint=hide_default_credentials_hint, ), status_code=200, ) + # Preserve return_to across the username/password sign-in too, via the SAME shared, never-raising + # helper the SSO branch uses, so /login can resume the connect flow instead of dead-ending at the + # dashboard. One implementation → the two sign-in branches cannot diverge (and the login form always + # renders, since the helper never raises on a bad return_to). + _persist_return_to_cookie(form_response, return_to) + return form_response def generic_response_convertor( @@ -1903,6 +1937,81 @@ async def auth_callback(request: Request, state: Optional[str] = None): ) +@router.get("/sso/saml/login", tags=["experimental"], include_in_schema=False) +async def saml_login(request: Request, return_to: str | None = None): + """SP-initiated SAML login. Redirects the user to the configured IdP.""" + from litellm.proxy.proxy_server import user_api_key_cache + + _disable_ui_flag = os.getenv("DISABLE_ADMIN_UI") + if _disable_ui_flag is not None and str_to_bool(value=_disable_ui_flag): + return admin_ui_disabled() + + return await SAMLAuthHandler.build_login_redirect(request=request, cache=user_api_key_cache, relay_state=return_to) + + +@router.get("/sso/saml/metadata", tags=["experimental"], include_in_schema=False) +async def saml_metadata(request: Request): + """Service Provider metadata XML, for registering this proxy at the IdP.""" + from litellm.proxy.proxy_server import user_api_key_cache + + metadata = await SAMLAuthHandler.build_sp_metadata(request=request, cache=user_api_key_cache) + return Response(content=metadata, media_type="application/xml") + + +@router.post("/sso/saml/callback", tags=["experimental"], include_in_schema=False) +async def saml_callback(request: Request): + """Assertion Consumer Service. Validates the IdP assertion and issues a UI session.""" + from litellm.proxy.proxy_server import ( + general_settings, + jwt_handler, + master_key, + premium_user, + prisma_client, + user_api_key_cache, + ) + + _disable_ui_flag = os.getenv("DISABLE_ADMIN_UI") + if _disable_ui_flag is not None and str_to_bool(value=_disable_ui_flag): + return admin_ui_disabled() + + if prisma_client is None: + raise HTTPException(status_code=500, detail=CommonProxyErrors.db_not_connected_error.value) + if master_key is None: + raise ProxyException( + message="Master Key not set for Proxy. Set `LITELLM_MASTER_KEY` in .env or general_settings:master_key in config.yaml.", + type=ProxyErrorTypes.auth_error, + param="master_key", + code=status.HTTP_500_INTERNAL_SERVER_ERROR, + ) + + post_data = await SAMLAuthHandler.read_acs_post_data(request) + if "SAMLResponse" not in post_data: + raise HTTPException(status_code=400, detail="Missing SAMLResponse in callback request.") + + result = await SAMLAuthHandler.handle_acs(request=request, cache=user_api_key_cache, post_data=post_data) + + await _raise_if_sso_exceeds_free_user_limit(premium_user, prisma_client) + + ui_access_mode = general_settings.get("ui_access_mode", None) + relay_state = post_data.get("RelayState") + cp_return_to: str | None = ( + relay_state + if isinstance(relay_state, str) and SSOAuthenticationHandler._validate_return_to(relay_state) + else None + ) + + return await SSOAuthenticationHandler.get_redirect_response_from_openid( + result=result, + request=request, + received_response=None, + generic_client_id=None, + ui_access_mode=ui_access_mode, + access_token_payload=None, + jwt_handler=jwt_handler, + return_to=cp_return_to, + ) + + async def _build_cli_sso_user_defined_values( result: Union[OpenID, dict], parsed_openid_result: ParsedOpenIDResult, @@ -1957,6 +2066,7 @@ async def _complete_cli_sso_callback_session( user_defined_values: Optional[SSOUserDefinedValues], prisma_client: PrismaClient, user_api_key_cache: UserApiKeyCache, + cli_sso_session_cache: DualCache, proxy_logging_obj: ProxyLogging, prefill_user_code: str | None = None, sso_assertion: SSOIdentityAssertion | None = None, @@ -2006,7 +2116,7 @@ async def _complete_cli_sso_callback_session( flow["sso_complete"] = True browser_complete_token = secrets.token_urlsafe(32) flow["browser_complete_token_hash"] = _hash_cli_sso_secret(browser_complete_token) - _set_cli_sso_flow(login_id=key, cache=user_api_key_cache, flow=flow) + _set_cli_sso_flow(login_id=key, cache=cli_sso_session_cache, flow=flow) verbose_proxy_logger.info( f"Stored CLI SSO session for user: {user_info.user_id}, teams: {teams}, num_teams: {len(teams)}" @@ -2037,13 +2147,14 @@ async def cli_sso_callback( verbose_proxy_logger.info("CLI SSO callback") from litellm.proxy.proxy_server import ( + cli_sso_session_cache, general_settings, prisma_client, proxy_logging_obj, user_api_key_cache, ) - flow = _get_cli_sso_flow_or_raise(login_id=key, cache=user_api_key_cache) + flow = _get_cli_sso_flow_or_raise(login_id=key, cache=cli_sso_session_cache) if prisma_client is None: raise HTTPException(status_code=500, detail=CommonProxyErrors.db_not_connected_error.value) @@ -2083,6 +2194,7 @@ async def cli_sso_callback( user_defined_values=user_defined_values, prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, + cli_sso_session_cache=cli_sso_session_cache, proxy_logging_obj=proxy_logging_obj, prefill_user_code=prefill_user_code, sso_assertion=sso_assertion, @@ -2114,10 +2226,10 @@ async def cli_poll_key( team_id: Optional team ID to assign to the JWT. If provided, must be one of user's teams. """ from litellm.proxy.auth.auth_checks import ExperimentalUIJWTToken - from litellm.proxy.proxy_server import user_api_key_cache + from litellm.proxy.proxy_server import cli_sso_session_cache try: - flow = _get_cli_sso_flow_or_raise(login_id=key_id, cache=user_api_key_cache) + flow = _get_cli_sso_flow_or_raise(login_id=key_id, cache=cli_sso_session_cache) if not _verify_cli_sso_poll_secret(flow=flow, poll_secret=x_litellm_cli_poll_secret): raise HTTPException(status_code=403, detail="Invalid CLI polling secret") @@ -2192,7 +2304,7 @@ async def cli_poll_key( ) # Delete cache entry (single-use) - user_api_key_cache.delete_cache(key=_get_cli_sso_flow_cache_key(key_id)) + cli_sso_session_cache.delete_cache(key=_get_cli_sso_flow_cache_key(key_id)) verbose_proxy_logger.info(f"CLI JWT generated for user: {user_id}, team: {team_id}") poll_response = { @@ -2404,6 +2516,92 @@ async def sso_readiness(): ) +def _is_same_origin_return_path(return_to: str) -> bool: + """True for a strictly relative return path that stays on the gateway's own origin by + construction, and is therefore safe to honor without a configured ``control_plane_url``. + Used by the MCP gateway DCR authorize round-trip so a browser sent through login lands + back on the authorize request. + + Requires a single leading ``/`` (not protocol-relative ``//``), no backslash (browsers + fold ``\\`` to ``/``, so ``/\\evil.com`` would escape the origin), and no control or + whitespace characters. Rejecting control chars keeps a ``\\r\\n``/tab-bearing value out + of the redirect ``Location`` and the ``litellm_cp_return_to`` cookie entirely, rather + than relying on downstream header encoding to neutralize it.""" + if not return_to.startswith("/") or return_to.startswith("//") or "\\" in return_to: + return False + return not any(ord(ch) < 0x20 or ch in (" ", "\x7f") for ch in return_to) + + +async def _sso_return_to_redirect( + return_to: str | None, + jwt_token: str, + redis_usage_cache, + user_api_key_cache, +) -> RedirectResponse | None: + """Resolve the post-SSO redirect for a ``return_to``, or None to fall through to the dashboard. + + Two arms, both clearing the one-shot ``litellm_cp_return_to`` cookie: + - **Same-origin relative path** (the MCP gateway DCR authorize round-trip): set the session cookie + exactly like the dashboard path, then send the browser back where it came from. + - **Control-plane cross-origin** (``control_plane_url``): stash the JWT behind a single-use opaque + code (60s TTL) so the token never lands in browser history/logs; the control plane redeems it via + ``POST /v3/login/exchange``. + + Extracted from ``get_redirect_response_from_openid`` to keep that method inside the complexity + budget; behavior is identical to the inline arms it replaces (including letting + ``_validate_return_to`` raise for a mismatched absolute return_to, as before).""" + if return_to is None: + return None + + if _is_same_origin_return_path(return_to): + redirect_response = RedirectResponse(url=return_to, status_code=303) + redirect_response.set_cookie(key="token", value=jwt_token) + redirect_response.delete_cookie("litellm_cp_return_to") + return redirect_response + + if SSOAuthenticationHandler._validate_return_to(return_to): + code = secrets.token_urlsafe(32) + cache_key = f"login_code:{code}" + cache_value = {"token": jwt_token, "redirect_url": return_to} + if redis_usage_cache is not None: + await redis_usage_cache.async_set_cache(key=cache_key, value=cache_value, ttl=60) + else: + await user_api_key_cache.async_set_cache(key=cache_key, value=cache_value, ttl=60) + + separator = "&" if "?" in return_to else "?" + redirect_url = return_to + separator + urlencode({"login": "success", "code": code}) + verbose_proxy_logger.info("Cross-origin SSO: redirecting to control plane with login code") + redirect_response = RedirectResponse(url=redirect_url, status_code=303) + redirect_response.delete_cookie("litellm_cp_return_to") + return redirect_response + + return None + + +def _persist_return_to_cookie(response: Response, return_to: str | None) -> None: + """Best-effort: persist a SAFE ``return_to`` on ``response`` as the one-shot ``litellm_cp_return_to`` + cookie so ANY sign-in path — SSO / Okta / generic OR the username/password form — can resume there + afterwards. THIS is the single source of truth, called by every sign-in branch so they cannot + diverge (a per-branch reimplementation is exactly how the two drifted before). Honors a strictly + relative same-origin path, and (when ``control_plane_url`` is configured) a return_to matching that + origin. It NEVER raises: a mismatched or invalid ``return_to`` is simply not stored, so it can never + block sign-in — the login entrypoint must always render.""" + if return_to is None: + return + try: + safe = _is_same_origin_return_path(return_to) or SSOAuthenticationHandler._validate_return_to(return_to) + except HTTPException: + return # a non-matching absolute return_to is ignored, never blocks sign-in + if safe: + response.set_cookie( + key="litellm_cp_return_to", + value=return_to, + max_age=600, + httponly=True, + samesite="lax", + ) + + class SSOAuthenticationHandler: """ Handler for SSO Authentication across all SSO providers @@ -3041,7 +3239,6 @@ class SSOAuthenticationHandler: return_to: Optional[str] = None, sso_assertion: SSOIdentityAssertion | None = None, ) -> RedirectResponse: - import jwt from litellm.proxy.proxy_server import ( general_settings, @@ -3205,30 +3402,21 @@ class SSOAuthenticationHandler: server_root_path=get_server_root_path(), ) - jwt_token = jwt.encode( - cast(dict, returned_ui_token_object), - master_key or "", - algorithm="HS256", + from litellm.proxy.auth.login_utils import encode_ui_session_jwt + + jwt_token = encode_ui_session_jwt(returned_ui_token_object, master_key or "") + + # Post-SSO return_to handling (the same-origin DCR round-trip and the control-plane + # cross-origin code exchange) lives in one shared helper so this method stays inside the + # complexity budget. None falls through to the dashboard redirect below. + return_to_redirect = await _sso_return_to_redirect( + return_to=return_to, + jwt_token=jwt_token, + redis_usage_cache=redis_usage_cache, + user_api_key_cache=user_api_key_cache, ) - - # Control-plane cross-origin: store JWT behind a single-use opaque - # code (60s TTL) so the token never appears in browser history / logs. - # The control plane redeems it via POST /v3/login/exchange. - if return_to is not None and SSOAuthenticationHandler._validate_return_to(return_to): - code = secrets.token_urlsafe(32) - cache_key = f"login_code:{code}" - cache_value = {"token": jwt_token, "redirect_url": return_to} - if redis_usage_cache is not None: - await redis_usage_cache.async_set_cache(key=cache_key, value=cache_value, ttl=60) - else: - await user_api_key_cache.async_set_cache(key=cache_key, value=cache_value, ttl=60) - - separator = "&" if "?" in return_to else "?" - redirect_url = return_to + separator + urlencode({"login": "success", "code": code}) - verbose_proxy_logger.info("Cross-origin SSO: redirecting to control plane with login code") - redirect_response = RedirectResponse(url=redirect_url, status_code=303) - redirect_response.delete_cookie("litellm_cp_return_to") - return redirect_response + if return_to_redirect is not None: + return return_to_redirect if user_id is not None and isinstance(user_id, str): litellm_dashboard_ui += "?login=success" diff --git a/litellm/proxy/management_helpers/object_permission_utils.py b/litellm/proxy/management_helpers/object_permission_utils.py index 6bbd41b93ed..015ab1d6df6 100644 --- a/litellm/proxy/management_helpers/object_permission_utils.py +++ b/litellm/proxy/management_helpers/object_permission_utils.py @@ -4,7 +4,8 @@ organizations, teams, and keys. """ import json -from typing import TYPE_CHECKING, Any, Dict, List, Optional, Set, Union +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any, Dict, List, Mapping, Optional, Set, Union from fastapi import HTTPException, status @@ -64,6 +65,57 @@ async def attach_object_permission_to_dict( return data_dict +@dataclass(frozen=True, slots=True) +class ObjectPermissionUpsert: + object_permission_id: str + record: dict[str, object] + + +async def prepare_object_permission_upsert( + new_object_permission: Mapping[str, object], + existing_object_permission_id: str | None, + prisma_client: PrismaClient, +) -> ObjectPermissionUpsert: + """ + Read-and-merge half of an object permission upsert; performs no writes. + + Merges the sent grants over the existing row (looked up by + ``existing_object_permission_id``, or a fresh uuid when the entity has none) and + returns the id plus the full record to upsert. The id is pinned inside the record + because the column has ``@default(uuid())``, so a create without it would mint a + different id than the one the caller links. ``mcp_tool_permissions`` is serialized + to a JSON string to avoid GraphQL parsing issues (e.g. server IDs starting with + "3e64" being interpreted as floats). + + Keeping this separate from the write lets callers run the upsert inside the same + transaction as the row that links ``object_permission_id``, so a rolled-back + update cannot leave permission changes live. + """ + object_permission_id = existing_object_permission_id or str(uuid.uuid4()) + existing_object_permission = await ObjectPermissionRepository(prisma_client).table.find_unique( + where={"object_permission_id": object_permission_id}, + ) + existing_fields: dict[str, object] = ( + existing_object_permission.model_dump(exclude_unset=True, exclude_none=True) + if existing_object_permission is not None + else {} + ) + merged: dict[str, object] = { + **existing_fields, + **new_object_permission, + "object_permission_id": object_permission_id, + } + record: dict[str, object] = { + **merged, + **( + {"mcp_tool_permissions": safe_dumps(merged["mcp_tool_permissions"])} + if "mcp_tool_permissions" in merged + else {} + ), + } + return ObjectPermissionUpsert(object_permission_id=object_permission_id, record=record) + + async def handle_update_object_permission_common( data_json: Dict, existing_object_permission_id: Optional[str], @@ -93,50 +145,23 @@ async def handle_update_object_permission_common( if prisma_client is None: raise ValueError("Prisma client not found") - ######################################################### - # Ensure `object_permission` is not added to the data_json - # We need to update the entity at the object_permission_id level in the LiteLLM_ObjectPermissionTable - ######################################################### - new_object_permission: Union[dict, str] = data_json.pop("object_permission", None) + new_object_permission: Union[dict, str, None] = data_json.pop("object_permission", None) if new_object_permission is None: return None - # Lookup existing object permission ID and update that entry - object_permission_id_to_use: str = existing_object_permission_id or str(uuid.uuid4()) - existing_object_permissions_dict: Dict = {} - - existing_object_permission = await ObjectPermissionRepository(prisma_client).table.find_unique( - where={"object_permission_id": object_permission_id_to_use}, - ) - - # Update the object permission - if existing_object_permission is not None: - existing_object_permissions_dict = existing_object_permission.model_dump(exclude_unset=True, exclude_none=True) - - # Handle string JSON object permission if isinstance(new_object_permission, str): new_object_permission = json.loads(new_object_permission) - if isinstance(new_object_permission, dict): - existing_object_permissions_dict.update(new_object_permission) - - ######################################################### - # Serialize mcp_tool_permissions JSON field to avoid GraphQL parsing issues - # (e.g., server IDs starting with "3e64" being interpreted as floats) - ######################################################### - if "mcp_tool_permissions" in existing_object_permissions_dict: - existing_object_permissions_dict["mcp_tool_permissions"] = safe_dumps( - existing_object_permissions_dict["mcp_tool_permissions"] - ) - - ######################################################### - # Commit the update to the LiteLLM_ObjectPermissionTable - ######################################################### + upsert = await prepare_object_permission_upsert( + new_object_permission=new_object_permission if isinstance(new_object_permission, dict) else {}, + existing_object_permission_id=existing_object_permission_id, + prisma_client=prisma_client, + ) created_object_permission_row = await ObjectPermissionRepository(prisma_client).table.upsert( - where={"object_permission_id": object_permission_id_to_use}, + where={"object_permission_id": upsert.object_permission_id}, data={ - "create": existing_object_permissions_dict, - "update": existing_object_permissions_dict, + "create": upsert.record, + "update": upsert.record, }, ) diff --git a/litellm/proxy/management_helpers/utils.py b/litellm/proxy/management_helpers/utils.py index 11a99caebf5..7c52b04c4eb 100644 --- a/litellm/proxy/management_helpers/utils.py +++ b/litellm/proxy/management_helpers/utils.py @@ -252,6 +252,21 @@ async def _resolve_member_budget_id( return response.budget_id +async def _append_team_id_if_absent(prisma_client: PrismaClient, user_id: str, team_id: str) -> None: + """Append team_id to a user's teams array, only if it is not already present. + + The row-level filter makes the append a no-op once the team is present, so + repeated or concurrent adds of the same team cannot accumulate duplicate + team ids in user.teams (a duplicate also breaks auth logic that keys off the + number of teams a user belongs to). Teams added concurrently for a different + team id are unaffected, since each update filters on its own team id. + """ + await UserRepository(prisma_client).table.update_many( + where={"user_id": user_id, "NOT": {"teams": {"has": team_id}}}, + data={"teams": {"push": [team_id]}}, + ) + + async def add_new_member( new_member: Member, max_budget_in_team: Optional[float], @@ -276,13 +291,22 @@ async def add_new_member( ## ADD TEAM ID, to USER TABLE IF NEW ## if new_member.user_id is not None: new_user_defaults = get_new_internal_user_defaults(user_id=new_member.user_id) + # Upsert ensures the user row exists atomically (no create race when the + # same new user is provisioned concurrently), seeding teams on create. + # The teams append lives in the filtered update below rather than the + # upsert's update branch so an already-existing user does not get a + # duplicate team id. The update branch still has to write something: + # Prisma only compiles an upsert down to INSERT ... ON CONFLICT when it + # is non-empty, and falls back to a racy SELECT-then-INSERT when it is + # not, so this re-states user_id as a no-op rather than being empty. _returned_user = await UserRepository(prisma_client).table.upsert( where={"user_id": new_member.user_id}, data={ - "update": {"teams": {"push": [team_id]}}, - "create": {"teams": [team_id], **new_user_defaults}, # type: ignore + "create": {"teams": [team_id], **new_user_defaults}, + "update": {"user_id": new_member.user_id}, }, ) + await _append_team_id_if_absent(prisma_client, new_member.user_id, team_id) if _returned_user is not None: returned_user = LiteLLM_UserTable(**_returned_user.model_dump()) elif new_member.user_email is not None: @@ -302,12 +326,8 @@ async def add_new_member( returned_user = LiteLLM_UserTable(**_returned_user.model_dump()) elif len(existing_user_row) == 1: user_info = existing_user_row[0] - _returned_user = await UserRepository(prisma_client).table.update( - where={"user_id": user_info.user_id}, # type: ignore - data={"teams": {"push": [team_id]}}, - ) - if _returned_user is not None: - returned_user = LiteLLM_UserTable(**_returned_user.model_dump()) + await _append_team_id_if_absent(prisma_client, user_info.user_id, team_id) + returned_user = LiteLLM_UserTable(**user_info.model_dump()) elif len(existing_user_row) > 1: raise HTTPException( status_code=400, diff --git a/litellm/proxy/mcp_registry.json b/litellm/proxy/mcp_registry.json index 84431634e24..f37fc39813e 100644 --- a/litellm/proxy/mcp_registry.json +++ b/litellm/proxy/mcp_registry.json @@ -198,13 +198,9 @@ "icon_url": "https://cdn.simpleicons.org/googledrive", "category": "Productivity", "registry_url": null, - "transport": "stdio", - "command": "npx", - "args": ["-y", "@modelcontextprotocol/server-gdrive"], - "env_vars": [ - {"name": "GOOGLE_CLIENT_ID", "description": "Google OAuth Client ID", "secret": false}, - {"name": "GOOGLE_CLIENT_SECRET", "description": "Google OAuth Client Secret", "secret": true} - ] + "transport": "http", + "url": "https://drivemcp.googleapis.com/mcp/v1", + "env_vars": [] }, { "name": "google_calendar", diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 6de3e43fc1a..a20b557e38b 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -226,6 +226,7 @@ from litellm.constants import ( APSCHEDULER_MAX_INSTANCES, APSCHEDULER_MISFIRE_GRACE_TIME, APSCHEDULER_REPLACE_EXISTING, + CLI_SSO_SESSION_TTL_SECONDS, DAYS_IN_A_MONTH, DEFAULT_HEALTH_CHECK_INTERVAL, DEFAULT_MODEL_CREATED_AT_TIME, @@ -303,6 +304,11 @@ from litellm.proxy.common_utils.encrypt_decrypt_utils import ( encrypt_value_helper, ) from litellm.proxy.common_utils.html_forms.ui_login import build_ui_login_form +from litellm.proxy.config_resolvers import resolve_fields +from litellm.proxy.config_resolvers.alerting import ( + EMAIL_DESCRIPTORS, + SLACK_DESCRIPTORS, +) from litellm.proxy.common_utils.http_parsing_utils import ( _read_request_body, _safe_get_request_headers, @@ -1158,9 +1164,9 @@ _OPENAPI_HTTP_METHODS = { # Credentials surfaced by `/get/config/callbacks` in the alerting block: the # full Slack incoming-webhook URL is itself a credential, and the SMTP # password is a service password. Masked on read so plaintext never reaches -# the UI. Kept here at module scope to match the analogous -# `_SSO_SENSITIVE_FIELDS` / `_CACHE_SENSITIVE_FIELDS` constants in the SSO -# and cache endpoint files. +# the UI. Kept here at module scope to match the analogous descriptor +# `is_secret` flags in litellm.proxy.config_resolvers and the +# `_CACHE_SENSITIVE_FIELDS` constant in the cache endpoint file. _ALERTING_SENSITIVE_VARS: Set[str] = {"SLACK_WEBHOOK_URL", "SMTP_PASSWORD"} @@ -1970,6 +1976,7 @@ user_api_key_cache: UserApiKeyCache = UserApiKeyCache( default_in_memory_ttl=UserAPIKeyCacheTTLEnum.in_memory_cache_ttl.value ) spend_counter_cache = DualCache(default_in_memory_ttl=UserAPIKeyCacheTTLEnum.in_memory_cache_ttl.value) +cli_sso_session_cache = DualCache(default_in_memory_ttl=CLI_SSO_SESSION_TTL_SECONDS) model_max_budget_limiter = _PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=user_api_key_cache) litellm.logging_callback_manager.add_litellm_callback(model_max_budget_limiter) redis_usage_cache: Optional[RedisCache] = None # redis cache used for tracking spend, tpm/rpm limits @@ -3696,13 +3703,22 @@ def _build_redis_usage_cache_from_environment() -> RedisCache | None: def _attach_redis_usage_cache(redis_cache: RedisCache, enable_redis_auth_cache: bool) -> None: """ Wires an established coordination Redis into the proxy-level caches that - consume it directly: the spend counter cache, the cluster-wide config - cache, and (only when opted in) the virtual-key auth cache. + consume it directly: the spend counter cache, the CLI SSO login-session + cache, the cluster-wide config cache, and (only when opted in) the + virtual-key auth cache. + + The CLI SSO login-session cache is always backed by Redis when available so + that the browser SSO flow behind `lite login` survives landing on different + workers; it must not be gated behind enable_redis_auth_cache. """ spend_counter_cache.attach_redis_cache( redis_cache, default_redis_ttl=litellm.default_redis_ttl, ) + cli_sso_session_cache.attach_redis_cache( + redis_cache, + default_redis_ttl=CLI_SSO_SESSION_TTL_SECONDS, + ) if enable_redis_auth_cache is True: user_api_key_cache.attach_redis_cache( redis_cache, @@ -4618,6 +4634,11 @@ class ProxyConfig: verbose_proxy_logger.debug( f"{blue_color_code} Initialized polling via cache: enabled={polling_via_cache_enabled}, native_background_mode={native_background_mode}, ttl={polling_cache_ttl}{reset_color_code}" ) + elif key == "max_ui_session_budget": + litellm.max_ui_session_budget = float(value) if value is not None else None + verbose_proxy_logger.debug( + f"{blue_color_code} setting litellm.max_ui_session_budget={litellm.max_ui_session_budget}{reset_color_code}" + ) elif key == "default_team_settings": for idx, team_setting in enumerate(value): # run through pydantic validation try: @@ -13451,7 +13472,7 @@ async def fallback_login(request: Request): @router.post("/login", include_in_schema=False) # hidden since this is a helper for UI sso login async def login(request: Request): global premium_user, general_settings, master_key - from litellm.proxy.auth.login_utils import authenticate_user, create_ui_token_object + from litellm.proxy.auth.login_utils import authenticate_user, create_ui_token_object, encode_ui_session_jwt from litellm.proxy.utils import get_custom_url form = await request.form() @@ -13474,13 +13495,7 @@ async def login(request: Request): ) # Generate JWT token - import jwt - - jwt_token = jwt.encode( - cast(dict, returned_ui_token_object), - cast(str, master_key), - algorithm="HS256", - ) + jwt_token = encode_ui_session_jwt(returned_ui_token_object, cast(str, master_key)) # Build redirect URL litellm_dashboard_ui = get_custom_url(str(request.base_url)) @@ -13490,16 +13505,51 @@ async def login(request: Request): litellm_dashboard_ui += "/ui/" litellm_dashboard_ui += "?login=success" + # Honor a same-origin return_to preserved by the sign-in page (e.g. the aggregate DCR connect flow's + # authorize round-trip), mirroring the SSO callback; otherwise land on the dashboard. Gated by + # _is_same_origin_return_path (strictly relative path) so it can never be an open redirect, and the + # one-shot cookie is cleared after use. + from litellm.proxy.management_endpoints.ui_sso import _sso_return_to_redirect + + # Resume through the SAME resumer the SSO callback uses, rather than a second, narrower arm. + # _persist_return_to_cookie stores both shapes it accepts (a relative same-origin path AND a + # control_plane_url-matching absolute URL); honoring only the relative one here silently dropped + # the control-plane case, landing the user on the dashboard. One function decides how a stored + # return_to is honored for EVERY sign-in branch, so the write and read sets cannot diverge: it + # sets the token cookie on the same-origin arm and hands off via a one-time login code on the + # cross-origin arm, and clears the one-shot cookie in both. + cp_return_to = request.cookies.get("litellm_cp_return_to") + if cp_return_to: + try: + resumed = await _sso_return_to_redirect( + return_to=cp_return_to, + jwt_token=jwt_token, + redis_usage_cache=redis_usage_cache, + user_api_key_cache=user_api_key_cache, + ) + except Exception: # noqa: BLE001 # resuming must NEVER block a completed sign-in + # The symmetric half of _persist_return_to_cookie's "never raises" contract. The resumer + # rejects a return_to that no longer matches control_plane_url (a config change between + # the cookie's write and this read), and the user has ALREADY authenticated here — + # failing their login over a stale one-shot cookie is the worst possible outcome. Land + # on the dashboard instead; the cookie is cleared below either way. + verbose_proxy_logger.info("Ignoring stale litellm_cp_return_to cookie; landing on dashboard") + resumed = None + if resumed is not None: + return resumed + # Create redirect response with cookie redirect_response = RedirectResponse(url=litellm_dashboard_ui, status_code=303) redirect_response.set_cookie(key="token", value=jwt_token) + if cp_return_to: + redirect_response.delete_cookie(key="litellm_cp_return_to") return redirect_response @router.post("/v2/login", include_in_schema=False) # hidden helper for UI logins via API async def login_v2(request: Request): global premium_user, general_settings, master_key - from litellm.proxy.auth.login_utils import authenticate_user, create_ui_token_object + from litellm.proxy.auth.login_utils import authenticate_user, create_ui_token_object, encode_ui_session_jwt from litellm.proxy.utils import get_custom_url try: @@ -13520,13 +13570,7 @@ async def login_v2(request: Request): premium_user=premium_user, ) - import jwt - - jwt_token = jwt.encode( - cast(dict, returned_ui_token_object), - cast(str, master_key), - algorithm="HS256", - ) + jwt_token = encode_ui_session_jwt(returned_ui_token_object, cast(str, master_key)) litellm_dashboard_ui = get_custom_url(str(request.base_url)) if litellm_dashboard_ui.endswith("/"): @@ -13570,7 +13614,7 @@ async def login_v2(request: Request): ) # control-plane login — always returns token in body for cross-origin use async def login_v3(request: Request): global premium_user, general_settings, master_key - from litellm.proxy.auth.login_utils import authenticate_user, create_ui_token_object + from litellm.proxy.auth.login_utils import authenticate_user, create_ui_token_object, encode_ui_session_jwt from litellm.proxy.utils import get_custom_url try: @@ -13599,13 +13643,7 @@ async def login_v3(request: Request): premium_user=premium_user, ) - import jwt - - jwt_token = jwt.encode( - cast(dict, returned_ui_token_object), - cast(str, master_key), - algorithm="HS256", - ) + jwt_token = encode_ui_session_jwt(returned_ui_token_object, cast(str, master_key)) litellm_dashboard_ui = get_custom_url(str(request.base_url)) if litellm_dashboard_ui.endswith("/"): @@ -14914,10 +14952,11 @@ GeneralSettingsUILiteLLMValue = Union[float, bool, str, None] class GeneralSettingsUILiteLLMFieldSpec(TypedDict): - type: Literal["Float", "Boolean", "Select"] + type: Literal["Float", "Dollar", "Boolean", "Select"] description: str options: NotRequired[tuple[str, ...]] tab: NotRequired[str] # Admin UI sub-tab this field renders under; None groups it with the rest + default: NotRequired[float] # reset/clear restores this instead of None; fields whose None means fail-open set it _GENERAL_SETTINGS_UI_LITELLM_FIELDS: dict[str, GeneralSettingsUILiteLLMFieldSpec] = { @@ -14943,21 +14982,32 @@ _GENERAL_SETTINGS_UI_LITELLM_FIELDS: dict[str, GeneralSettingsUILiteLLMFieldSpec "tab": "prompt_caching", "description": "Empty uses Anthropic's 5m default. 1h suits long sessions but doubles the cache write cost.", }, + "max_ui_session_budget": { + "type": "Dollar", + "default": 1.0, + "description": ( + "USD spend cap for each dashboard login session; covers LLM calls made from the dashboard " + "such as the playground and auto router Test Connection. Each login starts a fresh session " + "with this budget. Clearing restores the $1 default." + ), + }, } def _general_settings_ui_litellm_default( - field_type: Literal["Float", "Boolean", "Select"], + spec: GeneralSettingsUILiteLLMFieldSpec, ) -> GeneralSettingsUILiteLLMValue: """The value a field falls back to when it is cleared or reset.""" - return False if field_type == "Boolean" else None + if "default" in spec: + return spec["default"] + return False if spec["type"] == "Boolean" else None def _validate_general_settings_ui_litellm_value(field_name: str, value: Any) -> GeneralSettingsUILiteLLMValue: spec = _GENERAL_SETTINGS_UI_LITELLM_FIELDS[field_name] field_type = spec["type"] if value is None or value == "": - return _general_settings_ui_litellm_default(field_type) + return _general_settings_ui_litellm_default(spec) match field_type: case "Boolean": if not isinstance(value, bool): @@ -14981,6 +15031,13 @@ def _validate_general_settings_ui_litellm_value(field_name: str, value: Any) -> detail={"error": f"{field_name} must be a number in (0, 1] or empty"}, ) return float(value) + case "Dollar": + if isinstance(value, bool) or not isinstance(value, (int, float)) or float(value) <= 0: + raise HTTPException( + status_code=400, + detail={"error": f"{field_name} must be a positive dollar amount or empty"}, + ) + return float(value) case _: assert_never(field_type) @@ -15003,7 +15060,7 @@ async def _persist_general_settings_ui_litellm_field( async def _reset_general_settings_ui_litellm_field(field_name: str, user_api_key_dict: UserAPIKeyAuth) -> dict: config = await proxy_config.get_config() before_value = config.get("litellm_settings", {}).get(field_name) - default_value = _general_settings_ui_litellm_default(_GENERAL_SETTINGS_UI_LITELLM_FIELDS[field_name]["type"]) + default_value = _general_settings_ui_litellm_default(_GENERAL_SETTINGS_UI_LITELLM_FIELDS[field_name]) setattr(litellm, field_name, default_value) if "litellm_settings" in config: config["litellm_settings"].pop(field_name, None) @@ -15178,7 +15235,7 @@ async def get_config_list( ) for litellm_field_name, spec in _GENERAL_SETTINGS_UI_LITELLM_FIELDS.items(): current_value: GeneralSettingsUILiteLLMValue = getattr(litellm, litellm_field_name, None) - default_value = _general_settings_ui_litellm_default(spec["type"]) + default_value = _general_settings_ui_litellm_default(spec) stored_in_db_litellm: Optional[bool] if litellm_field_name in db_litellm_settings: stored_in_db_litellm = True @@ -15456,14 +15513,10 @@ async def get_config( _alerting = _general_settings.get("alerting", []) alerting_data = [] if "slack" in _alerting: - _slack_vars = [ - "SLACK_WEBHOOK_URL", - ] - _slack_env_vars = { - _var: (value if (value := environment_variables.get(_var)) is not None else os.getenv(_var)) - for _var in _slack_vars - } - _slack_env_vars = _apply_alerting_env_role_gate(_slack_env_vars, is_full_admin) + _slack_values, _ = resolve_fields( + SLACK_DESCRIPTORS, environment_variables, os.environ, empty_db_is_set=True + ) + _slack_env_vars = _apply_alerting_env_role_gate(_slack_values, is_full_admin) _alerting_types = proxy_logging_obj.slack_alerting_instance.alert_types _all_alert_types = proxy_logging_obj.slack_alerting_instance._all_possible_alert_types() @@ -15479,19 +15532,8 @@ async def get_config( } ) # pass email alerting vars - _email_vars = [ - "SMTP_HOST", - "SMTP_PORT", - "SMTP_USERNAME", - "SMTP_PASSWORD", - "SMTP_SENDER_EMAIL", - "TEST_EMAIL_ADDRESS", - "EMAIL_LOGO_URL", - "EMAIL_SUPPORT_CONTACT", - ] - _email_env_vars = _apply_alerting_env_role_gate( - {_var: environment_variables.get(_var) for _var in _email_vars}, is_full_admin - ) + _email_values, _ = resolve_fields(EMAIL_DESCRIPTORS, environment_variables, os.environ, empty_db_is_set=True) + _email_env_vars = _apply_alerting_env_role_gate(_email_values, is_full_admin) alerting_data.append( { diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index 80fd8a1594e..013873179c6 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -4,7 +4,9 @@ import asyncio import json from dataclasses import dataclass from datetime import datetime, timedelta, timezone -from typing import Any, Dict, List, Mapping, Optional, Sequence, cast +from typing import Any, Dict, List, Mapping, NoReturn, Optional, Sequence, cast + +from fastapi import HTTPException, status import litellm from litellm._logging import verbose_proxy_logger @@ -59,6 +61,22 @@ class _CounterReservationUnavailable(Exception): super().__init__("Counter reservation unavailable") +def _raise_reservation_unavailable(counter_key: str) -> NoReturn: + verbose_proxy_logger.warning( + "fail_closed_budget_enforcement: rejecting request — budget reservation for %s could not be written", + counter_key, + ) + raise HTTPException( + status_code=status.HTTP_503_SERVICE_UNAVAILABLE, + detail=( + "Budget enforcement unavailable: the budget reservation could not " + "be written to the spend counter backend, and " + "fail_closed_budget_enforcement is enabled, so the request was " + "rejected to avoid exceeding the configured budget. Retry shortly." + ), + ) + + def get_reserved_counter_keys(budget_reservation: Optional[dict]) -> set: if not budget_reservation: return set() @@ -138,6 +156,7 @@ async def reserve_budget_for_request( end_user_id: Optional[str] = None, end_user_object: Optional[Any] = None, skip_user_budget_on_team_key: bool = False, + fail_closed_budget_enforcement: bool = False, ) -> Optional[dict]: if valid_token is None or not RouteChecks.is_llm_api_route(route=route): return None @@ -193,6 +212,8 @@ async def reserve_budget_for_request( default_reserved_cost=reservation_cost, ) applied_entries.remove(entry) + if fail_closed_budget_enforcement: + _raise_reservation_unavailable(counter_key=counter.counter_key) continue if reserved_value is not None: diff --git a/litellm/proxy/spend_tracking/spend_management_endpoints.py b/litellm/proxy/spend_tracking/spend_management_endpoints.py index 55b50e7d9ff..0c525ee9466 100644 --- a/litellm/proxy/spend_tracking/spend_management_endpoints.py +++ b/litellm/proxy/spend_tracking/spend_management_endpoints.py @@ -1695,13 +1695,10 @@ async def ui_view_spend_logs( code=status.HTTP_401_UNAUTHORIZED, ) - if start_date is None or end_date is None: - raise ProxyException( - message="Start date and end date are required", - type="bad_request", - param="None", - code=status.HTTP_400_BAD_REQUEST, - ) + # Inline import — auth_utils participates in a proxy import cycle. + from litellm.proxy.auth.auth_utils import get_request_route # noqa: PLC0415 + + is_v2 = "/spend/logs/v2" in get_request_route(request) # Validate sort_by and sort_order valid_sort_fields = { @@ -1729,36 +1726,50 @@ async def ui_view_spend_logs( ) try: - # Inline import — auth_utils participates in a proxy import cycle. - from litellm.proxy.auth.auth_utils import get_request_route # noqa: PLC0415 + is_admin_view = _is_admin_view_safe(user_api_key_dict=user_api_key_dict) + is_request_id_lookup = request_id is not None and not is_v2 - is_v2 = "/spend/logs/v2" in get_request_route(request) - formats = ["%Y-%m-%d %H:%M:%S", "%Y-%m-%d"] if is_v2 else ["%Y-%m-%d %H:%M:%S"] + if is_request_id_lookup: + # request_id is the @id primary key: it identifies a single row, so a + # time window is meaningless. The dashboard always sends a default 24h + # window, which hid ids copied from an older page (LIT-3981). Drop the + # window for the id lookup so it resolves across all time; every other + # query, including the public v2 route, still requires one (below). + start_date_obj: datetime | None = None + end_date_obj: datetime | None = None + else: + if start_date is None or end_date is None: + raise ProxyException( + message="Start date and end date are required", + type="bad_request", + param="None", + code=status.HTTP_400_BAD_REQUEST, + ) + formats = ["%Y-%m-%d %H:%M:%S", "%Y-%m-%d"] if is_v2 else ["%Y-%m-%d %H:%M:%S"] - def parse_date(date_str: str) -> datetime: - date_str = date_str.strip() - for fmt in formats: - try: - return datetime.strptime(date_str, fmt).replace(tzinfo=timezone.utc) - except ValueError: - continue - expected = "'YYYY-MM-DD' or 'YYYY-MM-DD HH:MM:SS'" if is_v2 else "'YYYY-MM-DD HH:MM:SS'" - raise HTTPException( - status_code=status.HTTP_400_BAD_REQUEST, - detail=f"Invalid date format: {date_str}. Expected: {expected}", - ) + def parse_date(date_str: str) -> datetime: + date_str = date_str.strip() + for fmt in formats: + try: + return datetime.strptime(date_str, fmt).replace(tzinfo=timezone.utc) + except ValueError: + continue + expected = "'YYYY-MM-DD' or 'YYYY-MM-DD HH:MM:SS'" if is_v2 else "'YYYY-MM-DD HH:MM:SS'" + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail=f"Invalid date format: {date_str}. Expected: {expected}", + ) - start_date_obj = parse_date(start_date) - end_date_obj = parse_date(end_date) - - # Convert to ISO format strings for Prisma - start_date_iso = start_date_obj.isoformat() # Already in UTC, no need to add Z - end_date_iso = end_date_obj.isoformat() # Already in UTC, no need to add Z + start_date_obj = parse_date(start_date) + end_date_obj = parse_date(end_date) # Build where conditions - where_conditions: dict[str, Any] = { - "startTime": {"gte": start_date_iso, "lte": end_date_iso}, - } + where_conditions: dict[str, Any] = {} + if start_date_obj is not None and end_date_obj is not None: + where_conditions["startTime"] = { + "gte": start_date_obj.isoformat(), # Already in UTC, no need to add Z + "lte": end_date_obj.isoformat(), + } if team_id is not None: where_conditions["team_id"] = team_id @@ -1827,9 +1838,19 @@ async def ui_view_spend_logs( where_conditions["spend"]["gte"] = min_spend if max_spend is not None: where_conditions["spend"]["lte"] = max_spend - is_admin_view = _is_admin_view_safe(user_api_key_dict=user_api_key_dict) + # A request_id lookup drops the date window, so a non-admin could otherwise + # reach any single row by id; require they own it, mirroring the detail + # endpoint. That ownership check fully authorizes the one row, so the + # general scoping below is skipped for id lookups. Scoped to the UI route + # so the public v2 contract is unchanged. + if request_id is not None and not is_v2 and not is_admin_view: + await _assert_user_can_view_request_id( + prisma_client=prisma_client, + user_api_key_dict=user_api_key_dict, + request_id=request_id, + ) permitted_team_ids: List[str] | None = None - if not is_admin_view: + if not is_request_id_lookup and not is_admin_view: if team_id is not None: can_view_team = await _can_team_member_view_log( prisma_client=prisma_client, @@ -1875,15 +1896,16 @@ async def ui_view_spend_logs( sql_params: List[Any] = [] p = 1 # parameter index counter - # Date range (always present). Wrap the param side with - # `AT TIME ZONE 'UTC'` so comparison against the plain `timestamp` - # column does not depend on the DB session timezone (see #22529). - sql_conditions.append(f"\"startTime\" >= (${p}::timestamptz AT TIME ZONE 'UTC')") - sql_params.append(start_date_obj) - p += 1 - sql_conditions.append(f"\"startTime\" <= (${p}::timestamptz AT TIME ZONE 'UTC')") - sql_params.append(end_date_obj) - p += 1 + # Date range. Wrap the param side with `AT TIME ZONE 'UTC'` so comparison + # against the plain `timestamp` column does not depend on the DB session + # timezone (see #22529). Absent for a request_id-only lookup (see above). + if start_date_obj is not None and end_date_obj is not None: + sql_conditions.append(f"\"startTime\" >= (${p}::timestamptz AT TIME ZONE 'UTC')") + sql_params.append(start_date_obj) + p += 1 + sql_conditions.append(f"\"startTime\" <= (${p}::timestamptz AT TIME ZONE 'UTC')") + sql_params.append(end_date_obj) + p += 1 # Equality filters - read effective values from where_conditions (post-authorization) for sql_col, wc_key in [ diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index 50f6f791bc2..a6105b6dff9 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -374,12 +374,22 @@ def get_logging_payload(kwargs, response_obj, start_time, end_time) -> SpendLogs if isinstance(v, BaseModel): v = v.model_dump() additional_usage_values.update({k: v}) - if "cache_read_input_tokens" not in additional_usage_values: - prompt_tokens_details = additional_usage_values.get("prompt_tokens_details") - if isinstance(prompt_tokens_details, dict): + prompt_tokens_details = additional_usage_values.get("prompt_tokens_details") + if not isinstance(prompt_tokens_details, dict): + usage_object = clean_metadata.get("usage_object") + if isinstance(usage_object, dict): + prompt_tokens_details = usage_object.get("prompt_tokens_details") + if isinstance(prompt_tokens_details, dict): + if "cache_read_input_tokens" not in additional_usage_values: cached_tokens = prompt_tokens_details.get("cached_tokens") if isinstance(cached_tokens, int) and cached_tokens > 0: additional_usage_values["cache_read_input_tokens"] = cached_tokens + if "cache_creation_input_tokens" not in additional_usage_values: + cache_write_tokens = prompt_tokens_details.get("cache_write_tokens") or prompt_tokens_details.get( + "cache_creation_tokens" + ) + if isinstance(cache_write_tokens, int) and cache_write_tokens > 0: + additional_usage_values["cache_creation_input_tokens"] = cache_write_tokens clean_metadata["additional_usage_values"] = additional_usage_values if litellm.cache is not None: diff --git a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py index 48fa4bebaa3..10c71c00110 100644 --- a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py +++ b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py @@ -1,6 +1,8 @@ #### CRUD ENDPOINTS for UI Settings ##### import asyncio import json +import os +from collections.abc import Mapping from typing import Any, Dict, List, Optional, Set, Tuple, Type, Union from urllib.parse import urlparse @@ -13,6 +15,11 @@ from litellm._logging import verbose_proxy_logger from litellm.litellm_core_utils.sensitive_data_masker import mask_sensitive_keys from litellm.proxy._types import * from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.config_resolvers.sso import ( + SSO_FIELD_ENV_VARS, + SSO_SECRET_FIELDS, + resolve_sso_config, +) from litellm.repositories.config_repository import ConfigRepository from litellm.repositories.table_repositories import ( SSOConfigRepository, @@ -25,17 +32,45 @@ from litellm.types.proxy.management_endpoints.ui_sso import ( router = APIRouter() -# SSO secret fields returned by /get/sso_settings. These are masked on read so -# the UI can show "(set)" without ever transporting the plaintext OAuth secret -# off the server, matching the write-once + masked-on-read contract used for -# the HashiCorp Vault config override. -_SSO_SENSITIVE_FIELDS: Set[str] = { - "google_client_secret", - "microsoft_client_secret", - "generic_client_secret", +# Maps each UIThemeConfig field to the env var the UI branding path reads it +# from. /update/ui_theme_settings writes both the stored ui_theme_config and +# these env vars, so /get/ui_theme_settings resolves the same env vars to +# reflect a deployment branded purely through process env. +_UI_THEME_FIELD_ENV_VARS: dict[str, str] = { + "logo_url": "UI_LOGO_PATH", + "favicon_url": "LITELLM_FAVICON_URL", } +def _is_public_http_url(value: str | None) -> bool: + """Whether a value is a plain http(s) URL with a host, safe to disclose publicly.""" + if not isinstance(value, str) or not value.strip(): + return False + parsed = urlparse(value.strip()) + return parsed.scheme in ("http", "https") and bool(parsed.netloc) + + +def _resolve_ui_theme_field(stored_values: Mapping[str, Any], field_name: str) -> str | None: + """Resolve one UI theme field to the value the branding path actually uses. + + The stored ui_theme_config wins; a field absent or blank there falls back to + the process environment. The branding path reads the env var, and stored + settings reach it by being pushed into the environment on save, so a value + supplied only as a process env var is live even though no stored entry exists. + + This endpoint is unauthenticated, so the env fallback only surfaces a public + http(s) URL: an operator can point UI_LOGO_PATH at a local filesystem path + (the branding path serves it server-side), and that path must not be + disclosed to anonymous callers. A stored value is already validated as a + public URL on write, so it passes through. + """ + stored = stored_values.get(field_name) + if isinstance(stored, str) and stored.strip(): + return stored + env_value = os.environ.get(_UI_THEME_FIELD_ENV_VARS[field_name]) + return env_value if _is_public_http_url(env_value) else None + + class IPAddress(BaseModel): ip: str @@ -69,7 +104,8 @@ class SettingsResponse(BaseModel): class SSOSettingsResponse(SettingsResponse): """Response model for SSO settings""" - pass + provenance: Dict[str, str] = Field(default_factory=dict) + """Per-field source of each value: 'db', 'env', 'default', or 'unset'.""" class InternalUserSettingsResponse(SettingsResponse): @@ -717,7 +753,7 @@ async def get_sso_settings(): Returns a structured object with values and descriptions for UI display. """ - from litellm.proxy.proxy_server import prisma_client, proxy_config + from litellm.proxy.proxy_server import prisma_client if prisma_client is None: raise HTTPException( @@ -725,59 +761,12 @@ async def get_sso_settings(): detail={"error": "Database not connected. Please connect a database."}, ) - # Get SSO config from dedicated table + # Resolve the effective SSO config: the stored row wins, else the process + # environment, else each field's default. Unlike the legacy read path this + # does not write os.environ; a GET has no business mutating the environment. sso_db_record = await SSOConfigRepository(prisma_client).table.find_unique(where={"id": "sso_config"}) - - # Initialize with defaults - sso_settings_dict = {} - - if sso_db_record and sso_db_record.sso_settings: - # Load settings from database - sso_settings_dict = dict(sso_db_record.sso_settings) - - role_mappings_data = sso_settings_dict.pop("role_mappings", None) - role_mappings = None - if role_mappings_data: - from litellm.types.proxy.management_endpoints.ui_sso import RoleMappings - - if isinstance(role_mappings_data, dict): - role_mappings = RoleMappings(**role_mappings_data) - elif isinstance(role_mappings_data, RoleMappings): - role_mappings = role_mappings_data - - team_mappings_data = sso_settings_dict.pop("team_mappings", None) - team_mappings = None - if team_mappings_data: - from litellm.types.proxy.management_endpoints.ui_sso import TeamMappings - - if isinstance(team_mappings_data, dict): - team_mappings = TeamMappings(**team_mappings_data) - elif isinstance(team_mappings_data, TeamMappings): - team_mappings = team_mappings_data - - decrypted_sso_settings_dict = proxy_config._decrypt_and_set_db_env_variables( - environment_variables=sso_settings_dict - ) - - # Build SSO config with database values or environment fallback - - sso_config = SSOConfig( - google_client_id=decrypted_sso_settings_dict.get("google_client_id", None), - google_client_secret=decrypted_sso_settings_dict.get("google_client_secret", None), - microsoft_client_id=decrypted_sso_settings_dict.get("microsoft_client_id", None), - microsoft_client_secret=decrypted_sso_settings_dict.get("microsoft_client_secret", None), - microsoft_tenant=decrypted_sso_settings_dict.get("microsoft_tenant", None), - generic_client_id=decrypted_sso_settings_dict.get("generic_client_id", None), - generic_client_secret=decrypted_sso_settings_dict.get("generic_client_secret", None), - generic_authorization_endpoint=decrypted_sso_settings_dict.get("generic_authorization_endpoint", None), - generic_token_endpoint=decrypted_sso_settings_dict.get("generic_token_endpoint", None), - generic_userinfo_endpoint=decrypted_sso_settings_dict.get("generic_userinfo_endpoint", None), - proxy_base_url=decrypted_sso_settings_dict.get("proxy_base_url", None), - user_email=decrypted_sso_settings_dict.get("user_email"), - ui_access_mode=decrypted_sso_settings_dict.get("ui_access_mode"), - role_mappings=role_mappings, - team_mappings=team_mappings, - ) + sso_db_settings = dict(sso_db_record.sso_settings) if sso_db_record and sso_db_record.sso_settings else None + resolved = resolve_sso_config(sso_db_settings, os.environ) # Get the schema for UI display from pydantic import TypeAdapter @@ -786,11 +775,12 @@ async def get_sso_settings(): # Convert to dict for response, masking OAuth client secrets so plaintext # is never sent to the UI. - sso_dict = mask_sensitive_keys(sso_config.model_dump(), _SSO_SENSITIVE_FIELDS) + sso_dict = mask_sensitive_keys(resolved.config.model_dump(), set(SSO_SECRET_FIELDS)) # Add descriptions to the response result = { "values": sso_dict, + "provenance": resolved.provenance, "field_schema": { "description": schema.get("description", ""), "properties": {}, @@ -841,21 +831,6 @@ async def update_sso_settings( detail={"error": "Set `'STORE_MODEL_IN_DB='True'` in your env to enable this feature."}, ) - # Update environment variables - env_var_mapping = { - "google_client_id": "GOOGLE_CLIENT_ID", - "google_client_secret": "GOOGLE_CLIENT_SECRET", - "microsoft_client_id": "MICROSOFT_CLIENT_ID", - "microsoft_client_secret": "MICROSOFT_CLIENT_SECRET", - "microsoft_tenant": "MICROSOFT_TENANT", - "generic_client_id": "GENERIC_CLIENT_ID", - "generic_client_secret": "GENERIC_CLIENT_SECRET", - "generic_authorization_endpoint": "GENERIC_AUTHORIZATION_ENDPOINT", - "generic_token_endpoint": "GENERIC_TOKEN_ENDPOINT", - "generic_userinfo_endpoint": "GENERIC_USERINFO_ENDPOINT", - "proxy_base_url": "PROXY_BASE_URL", - } - # Read the existing SSO row first so the audit log captures a real # before/after diff. Stored values are encrypted; decrypt them so the # before-snapshot has the same shape as after_value, and rely on @@ -884,8 +859,8 @@ async def update_sso_settings( # Update environment variables in config and in memory sso_data = sso_config.model_dump() for field_name, value in sso_data.items(): - if field_name in env_var_mapping: - env_var_name = env_var_mapping[field_name] + if field_name in SSO_FIELD_ENV_VARS: + env_var_name = SSO_FIELD_ENV_VARS[field_name] if value: os.environ[env_var_name] = value else: @@ -935,7 +910,7 @@ async def update_sso_settings( else: environment_variables = {} - env_vars_to_remove = set(env_var_mapping.values()) + env_vars_to_remove = set(SSO_FIELD_ENV_VARS.values()) filtered_env_vars = { key: value for key, value in environment_variables.items() if key not in env_vars_to_remove } @@ -977,12 +952,19 @@ async def get_ui_theme_settings(): # Load existing config config = await proxy_config.get_config() - return await _get_settings_with_schema( + result = await _get_settings_with_schema( settings_key="ui_theme_config", settings_class=UIThemeConfig, config=config, ) + stored_values = result.get("values", {}) + result["values"] = { + **stored_values, + **{field: _resolve_ui_theme_field(stored_values, field) for field in _UI_THEME_FIELD_ENV_VARS}, + } + return result + def _validate_public_image_url(value: Optional[str], field_name: str) -> None: """ diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 43921a847a9..171e17ce650 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -172,6 +172,7 @@ from litellm.types.utils import LLMResponseTypes, LoggedLiteLLMParams if TYPE_CHECKING: from opentelemetry.trace import Span as _Span + from prisma.client import TransactionManager from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj @@ -1399,6 +1400,14 @@ class ProxyLogging: self._process_guardrail_metadata(data) return data + parallel_guardrails: tuple[CustomGuardrail, ...] = tuple( + cb + for cb in caps.resolved_callbacks + if isinstance(cb, CustomGuardrail) + and getattr(cb, "run_in_parallel", False) + and not (cb.guardrail_name and cb.guardrail_name in pipeline_managed) + ) + deferred_route_exc: Optional[SensitiveDataRouteException] = None for _callback in caps.resolved_callbacks: start_time = time.time() @@ -1408,6 +1417,9 @@ class ProxyLogging: if _callback.guardrail_name and _callback.guardrail_name in pipeline_managed: continue + if getattr(_callback, "run_in_parallel", False): + continue + result = await self._process_guardrail_callback( callback=_callback, data=data, # type: ignore @@ -1464,6 +1476,14 @@ class ProxyLogging: if deferred_route_exc is not None and data is not None: data = await self._handle_sensitive_data_route_exception(deferred_route_exc, data, user_api_key_dict) + if parallel_guardrails and data is not None: + await self._run_parallel_pre_call_guardrails( + guardrails=parallel_guardrails, + data=data, + user_api_key_dict=user_api_key_dict, + call_type=call_type, + ) + if data is not None: self._process_guardrail_metadata(data) @@ -1476,6 +1496,47 @@ class ProxyLogging: except Exception as e: raise e + async def _run_parallel_pre_call_guardrails( + self, + guardrails: tuple[CustomGuardrail, ...], + data: dict, + user_api_key_dict: UserAPIKeyAuth, + call_type: CallTypesLiteral, + ) -> None: + """ + Run opted-in pre_call guardrails concurrently against one shared payload + snapshot. These guardrails are declared block-only, so any modified data + they return is discarded; they run for their blocking side effect (raising + to reject the request before it reaches the LLM). Every guardrail is + awaited to completion (``return_exceptions=True``) so a raise by one never + leaves the others running as unobserved background tasks. A guardrail that + blocks (any exception other than a reroute or passthrough) takes precedence + over one that only changes the request flow, so a fast reroute can never + let a slower block be bypassed; the request is rejected before it reaches + the LLM, preserving the pre-call barrier that ``during_call`` guardrails + cannot provide. Per-guardrail latency is recorded by + ``_process_guardrail_callback``'s own metrics. + """ + results = await asyncio.gather( + *( + self._process_guardrail_callback( + callback=callback, + data=data, + user_api_key_dict=user_api_key_dict, + call_type=call_type, + event_type=GuardrailEventHooks.pre_call, + ) + for callback in guardrails + ), + return_exceptions=True, + ) + raised = tuple(result for result in results if isinstance(result, BaseException)) + blocking = next((exc for exc in raised if not _exception_changes_request_flow(exc)), None) + if blocking is not None: + raise blocking + if raised: + raise raised[0] + async def _handle_sensitive_data_route_exception( self, exc: SensitiveDataRouteException, @@ -2276,9 +2337,16 @@ class ProxyLogging: # Merge model-level guardrails before checking which guardrails to run guardrail_data = _check_and_merge_model_level_guardrails(data=data, llm_router=llm_router) + parallel_guardrails: tuple[CustomGuardrail, ...] = tuple( + callback for callback in guardrail_callbacks if getattr(callback, "run_in_parallel", False) + ) + for callback in guardrail_callbacks: # Main - V2 Guardrails implementation + if getattr(callback, "run_in_parallel", False): + continue + if ( callback.should_run_guardrail( data=guardrail_data, @@ -2315,6 +2383,15 @@ class ProxyLogging: if guardrail_response is not None: response = guardrail_response + if parallel_guardrails: + await self._run_parallel_post_call_guardrails( + guardrails=parallel_guardrails, + data=data, + guardrail_data=guardrail_data, + response=response, + user_api_key_dict=user_api_key_dict, + ) + ############ Handle CustomLogger ############################### ################################################################# @@ -2328,6 +2405,65 @@ class ProxyLogging: raise e return response + async def _run_parallel_post_call_guardrails( + self, + guardrails: tuple[CustomGuardrail, ...], + data: dict, + guardrail_data: dict, + response: LLMResponseTypes, + user_api_key_dict: UserAPIKeyAuth, + ) -> None: + """ + Run opted-in post_call guardrails concurrently against the response + produced by the sequential guardrails. These guardrails are declared + block-only, so any modified response they return is discarded; they run + for their blocking side effect (raising to reject the response before it + reaches the client). Every guardrail is awaited to completion + (``return_exceptions=True``) so a raise by one never leaves the others + running as unobserved background tasks. A guardrail that blocks (any + exception other than a passthrough) takes precedence over one that only + changes the response flow, so a fast passthrough can never let a slower + block be bypassed. Each per-guardrail coroutine sets ``guardrail_to_apply`` + immediately before awaiting, and the unified hook pops it before its first + suspension point, so concurrent guardrails never race on that key. + """ + + async def _run_one(callback: CustomGuardrail) -> None: + if callback.should_run_guardrail(data=guardrail_data, event_type=GuardrailEventHooks.post_call) is not True: + return + if "apply_guardrail" in type(callback).__dict__: + data["guardrail_to_apply"] = callback + await self._run_guardrail_with_metrics( + callback, + unified_guardrail.async_post_call_success_hook( + user_api_key_dict=user_api_key_dict, + data=data, + response=response, + ), + "post_call", + ) + else: + await self._run_guardrail_with_metrics( + callback, + callback.async_post_call_success_hook( + user_api_key_dict=user_api_key_dict, + data=data, + response=response, + ), + "post_call", + ) + + results = await asyncio.gather( + *(_run_one(callback) for callback in guardrails), + return_exceptions=True, + ) + raised = tuple(result for result in results if isinstance(result, BaseException)) + blocking = next((exc for exc in raised if not _exception_changes_request_flow(exc)), None) + if blocking is not None: + raise blocking + if raised: + raise raised[0] + async def post_call_response_headers_hook( self, data: dict, @@ -2922,6 +3058,14 @@ class PrismaClient: return self.db.writer return self.db + def tx(self) -> "TransactionManager": + """Open an interactive transaction on the writer. + + Callers go through this instead of reaching into ``self.db`` so writer + selection and read-replica routing stay encapsulated in the wrapper. + """ + return cast("TransactionManager", self.db.tx()) # cast-ok: wrappers delegate tx via __getattr__ (untyped) + def get_request_status(self, payload: Union[dict, SpendLogsPayload]) -> Literal["success", "failure"]: """ Determine if a request was successful or failed based on payload metadata. @@ -6159,6 +6303,9 @@ def create_model_info_response( if model_cost_info is not None: max_input_tokens = coerce_token_limit(model_cost_info.get("max_input_tokens")) max_output_tokens = coerce_token_limit(model_cost_info.get("max_output_tokens")) + mode = model_cost_info.get("mode") + if isinstance(mode, str): + base["mode"] = mode if llm_router is not None: configured_input, configured_output = llm_router.get_configured_token_limits(model_id) diff --git a/litellm/repositories/team_repository.py b/litellm/repositories/team_repository.py index 3227aa812ca..68875bd7972 100644 --- a/litellm/repositories/team_repository.py +++ b/litellm/repositories/team_repository.py @@ -4,11 +4,18 @@ Team repository for database operations on LiteLLM_TeamTable. import json from datetime import datetime -from typing import Any, Dict, List, Optional, Type +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Type -from litellm.models.team import LiteLLM_TeamTable +from pydantic import TypeAdapter + +from litellm.models.team import LiteLLM_TeamTable, Member from litellm.repositories.base_repository import BaseRepository +if TYPE_CHECKING: + from prisma import Prisma + +_MEMBERS_WITH_ROLES_ADAPTER = TypeAdapter(list[Member]) + class TeamRepository(BaseRepository[LiteLLM_TeamTable]): """Repository for team database operations.""" @@ -46,6 +53,24 @@ class TeamRepository(BaseRepository[LiteLLM_TeamTable]): return LiteLLM_TeamTable(**data) + async def get_members_with_roles_locked(self, tx: "Prisma", team_id: str) -> List[Member]: + """Return the team's members_with_roles, locking the row FOR UPDATE. + + Must be called inside a transaction so the row lock is held until + commit. This serializes concurrent membership writers on the team row + so the losing writer appends onto the winner's committed result instead + of overwriting it from a stale snapshot. + """ + rows = await tx.query_raw( + 'SELECT members_with_roles FROM "LiteLLM_TeamTable" WHERE team_id = $1 FOR UPDATE', + team_id, + ) + raw_value = rows[0]["members_with_roles"] if rows else None + parsed = json.loads(raw_value) if isinstance(raw_value, str) else raw_value + if not parsed: + return [] + return _MEMBERS_WITH_ROLES_ADAPTER.validate_python(parsed) + async def find_by_id(self, team_id: str, id_field: str = "team_id") -> Optional[LiteLLM_TeamTable]: return await super().find_by_id(team_id, id_field) diff --git a/litellm/responses/main.py b/litellm/responses/main.py index 206736f501a..944cf58df1c 100644 --- a/litellm/responses/main.py +++ b/litellm/responses/main.py @@ -494,7 +494,14 @@ async def aresponses( prompt_label=kwargs.get("prompt_label", None), prompt_version=kwargs.get("prompt_version", None), ) - input = cast(Union[str, ResponseInputParam], merged_input) + input = cast( + Union[str, ResponseInputParam], + ResponsesAPIRequestUtils.merge_prompt_management_input( + original_input=input, + client_input=client_input, + merged_input=merged_input, + ), + ) if model != original_model: _, custom_llm_provider, _, _ = litellm.get_llm_provider(model=model) kwargs.pop("prompt_id", None) @@ -609,7 +616,14 @@ def _apply_prompt_management_to_responses_call( prompt_label=kwargs.get("prompt_label", None), prompt_version=kwargs.get("prompt_version", None), ) - input = cast(Union[str, ResponseInputParam], merged_input) + input = cast( + Union[str, ResponseInputParam], + ResponsesAPIRequestUtils.merge_prompt_management_input( + original_input=input, + client_input=client_input, + merged_input=merged_input, + ), + ) local_vars["input"] = input local_vars["model"] = model if model != original_model: diff --git a/litellm/responses/utils.py b/litellm/responses/utils.py index 7a42cb96566..12c890ec91d 100644 --- a/litellm/responses/utils.py +++ b/litellm/responses/utils.py @@ -19,7 +19,9 @@ import litellm from litellm._logging import verbose_logger from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig from litellm.types.llms.openai import ( + AllMessageValues, ResponseAPIUsage, + ResponseInputParam, ResponsesAPIOptionalRequestParams, ResponsesAPIResponse, ResponseText, @@ -36,6 +38,57 @@ from litellm.types.utils import ( class ResponsesAPIRequestUtils: """Helper utils for constructing ResponseAPI requests""" + @staticmethod + def merge_prompt_management_input( + original_input: str | ResponseInputParam, + client_input: list[AllMessageValues], + merged_input: list[AllMessageValues], + ) -> list[object]: + if isinstance(original_input, str): + return [*merged_input] + + original_items = tuple(original_input) + client_item_ids = frozenset(id(item) for item in client_input) + message_positions = tuple(index for index, item in enumerate(original_items) if id(item) in client_item_ids) + + if len(message_positions) == len(original_items): + return [*merged_input] + if not message_positions: + verbose_logger.warning( + "Prompt management hook returned messages without Responses API input messages; merged messages were ignored" + ) + return [*original_items] + + corresponding_messages = len(client_input) == len(merged_input) and all( + original.get("role") == merged.get("role") + and (not isinstance(original.get("id"), str) or original.get("id") == merged.get("id")) + for original, merged in zip(client_input, merged_input) + ) + if corresponding_messages: + merged_by_position = dict(zip(message_positions, merged_input)) + return [ + merged_by_position[index] if index in merged_by_position else item + for index, item in enumerate(original_items) + ] + + all_messages_preserved = all(any(original is merged for merged in merged_input) for original in client_input) + if all_messages_preserved: + prefixes = { + id(original_items[position]): original_items[ + message_positions[index - 1] + 1 if index else 0 : position + ] + for index, position in enumerate(message_positions) + } + trailing_items = original_items[message_positions[-1] + 1 :] + return [item for merged in merged_input for item in (*prefixes.get(id(merged), ()), merged)] + list( + trailing_items + ) + + verbose_logger.warning( + "Prompt management hook replaced Responses API messages; non-message input items were dropped" + ) + return [*merged_input] + @staticmethod def _check_valid_arg( supported_params: Optional[List[str]], @@ -996,6 +1049,7 @@ class ResponseAPILoggingUtils: audio_tokens=getattr(response_api_usage.input_tokens_details, "audio_tokens", None), text_tokens=getattr(response_api_usage.input_tokens_details, "text_tokens", None), image_tokens=getattr(response_api_usage.input_tokens_details, "image_tokens", None), + cache_write_tokens=getattr(response_api_usage.input_tokens_details, "cache_write_tokens", None), ) completion_tokens_details: Optional[CompletionTokensDetailsWrapper] = None output_tokens_details = getattr(response_api_usage, "output_tokens_details", None) diff --git a/litellm/setup_wizard.py b/litellm/setup_wizard.py index 10b4fb30f22..c6d0c1717a9 100644 --- a/litellm/setup_wizard.py +++ b/litellm/setup_wizard.py @@ -52,12 +52,13 @@ PROVIDERS: List[Dict] = [ { "id": "anthropic", "name": "Anthropic", - "description": "Claude Fable 5, Opus 4.8, Opus 4.7, Opus 4.6, Sonnet 4.6, Haiku 4.5", + "description": "Claude Fable 5, Opus 5, Opus 4.8, Opus 4.7, Opus 4.6, Sonnet 5, Sonnet 4.6, Haiku 4.5", "env_key": "ANTHROPIC_API_KEY", "key_hint": "sk-ant-...", "test_model": "claude-haiku-4-5-20251001", "models": [ "claude-fable-5", + "claude-opus-5", "claude-sonnet-5", "claude-opus-4-8", "claude-opus-4-7", diff --git a/litellm/types/guardrails.py b/litellm/types/guardrails.py index 47d93fc2d7a..a324e71e289 100644 --- a/litellm/types/guardrails.py +++ b/litellm/types/guardrails.py @@ -725,6 +725,18 @@ class BaseLitellmParams(ContentFilterConfigModel): # works for new and patch up description="When True, guardrails only receive the latest message for the relevant role (e.g., newest user input pre-call, newest assistant output post-call)", ) + only_scan_new_messages: Optional[bool] = Field( + default=False, + description=( + "When True, the guardrail only scans messages that have not already been scanned " + "earlier in the same session (identified by litellm_session_id / session_id). " + "Message content is hashed per session and cached; only the diff (new or edited " + "messages) is sent to the guardrail provider on follow-up calls. Falls back to a " + "full scan when the request has no session id or the cache is unavailable. Intended " + "for blocking/detection guardrails; not applied when mask_request_content is set." + ), + ) + skip_system_message_in_guardrail: Optional[bool] = Field( default=None, description=( @@ -898,6 +910,17 @@ class BaseLitellmParams(ContentFilterConfigModel): # works for new and patch up ), ) + run_in_parallel: Optional[bool] = Field( + default=None, + description=( + "When True, this pre_call or post_call guardrail runs concurrently with other opted-in " + "guardrails of the same hook, after the sequential guardrails have run. Use only for " + "block-only guardrails that inspect and reject; do not enable it for guardrails that " + "modify the request or response (e.g. PII masking or sensitive-data routing), since " + "parallel runs share one snapshot and their mutations would race." + ), + ) + @field_validator( "mode", "default_action", diff --git a/litellm/types/mcp_server/mcp_server_manager.py b/litellm/types/mcp_server/mcp_server_manager.py index 8ae974b19a6..b0af22e7c3f 100644 --- a/litellm/types/mcp_server/mcp_server_manager.py +++ b/litellm/types/mcp_server/mcp_server_manager.py @@ -261,12 +261,3 @@ class MCPServer(BaseModel): if self.oauth_passthrough is not True: return False return any(h.lower() == "authorization" for h in self.extra_headers) - - @property - def has_token_exchange_config(self) -> bool: - """True if this server is configured for OAuth2 token exchange (OBO / RFC 8693).""" - return ( - self.auth_type == MCPAuth.oauth2_token_exchange - and bool(self.client_id and self.client_secret) - and bool(self.token_exchange_endpoint or self.token_url) - ) diff --git a/litellm/types/proxy/management_endpoints/ui_sso.py b/litellm/types/proxy/management_endpoints/ui_sso.py index 7234cc2650f..d4b1d98f957 100644 --- a/litellm/types/proxy/management_endpoints/ui_sso.py +++ b/litellm/types/proxy/management_endpoints/ui_sso.py @@ -148,6 +148,28 @@ class SSOConfig(LiteLLMPydanticObjectBase): default=None, description="User info endpoint URL for generic OAuth provider", ) + generic_scope: Optional[str] = Field( + default=None, + description="Space-separated OAuth scopes requested from the generic provider, e.g. 'openid email profile'", + ) + + # SAML SSO + saml_idp_metadata_url: Optional[str] = Field( + default=None, + description="URL of the SAML IdP metadata to fetch and parse for SSO authentication", + ) + saml_idp_metadata_xml: Optional[str] = Field( + default=None, + description="Inline SAML IdP metadata XML, used when a metadata URL is not available", + ) + saml_sp_entity_id: Optional[str] = Field( + default=None, + description="SAML Service Provider entityID; defaults to the proxy's /sso/saml/metadata URL", + ) + saml_allow_unsolicited: Optional[str] = Field( + default=None, + description="'true' to accept IdP-initiated (unsolicited) SAML responses, which cannot be browser-bound against login CSRF", + ) # Common settings proxy_base_url: Optional[str] = Field( diff --git a/litellm/types/proxy/model_listing.py b/litellm/types/proxy/model_listing.py index c3330da0d66..b59c0f2cf19 100644 --- a/litellm/types/proxy/model_listing.py +++ b/litellm/types/proxy/model_listing.py @@ -10,12 +10,16 @@ class ModelInfoMetadata(TypedDict): class ModelInfoResponse(TypedDict): - """OpenAI-compatible model object. `metadata` is present only when the - endpoint is called with include_metadata=true. + """OpenAI-compatible model object. `mode`, `max_input_tokens`, and + `max_output_tokens` are attached when the cost map knows them; `metadata` + is present only when the endpoint is called with include_metadata=true. """ id: str object: Literal["model"] created: int owned_by: str + mode: NotRequired[str] + max_input_tokens: NotRequired[int] + max_output_tokens: NotRequired[int] metadata: NotRequired[ModelInfoMetadata] diff --git a/litellm/types/tool_management.py b/litellm/types/tool_management.py index 1c5e1df9e9a..71ec412e8ef 100644 --- a/litellm/types/tool_management.py +++ b/litellm/types/tool_management.py @@ -98,3 +98,38 @@ class ToolUsageLogsResponse(BaseModel): total: int page: int page_size: int + + +class ToolSpendEntry(BaseModel): + """Total spend attributed to one tool over the requested window.""" + + tool_name: str + spend: float = Field( + 0.0, + description="Attributed spend: a request that used several tools counts its full spend toward each of them", + ) + call_count: int = 0 + total_tokens: int = 0 + + +class ToolSpendDailyEntry(BaseModel): + """Spend attributed to one tool on one UTC day.""" + + date: str + tool_name: str + spend: float = 0.0 + call_count: int = 0 + + +class ToolSpendResponse(BaseModel): + by_tool: List[ToolSpendEntry] = Field(default_factory=list) + daily: List[ToolSpendDailyEntry] = Field(default_factory=list) + total_spend: float = Field( + 0.0, + description=( + "Deduplicated spend of every request that called at least one tool in the window; " + "less than the sum of per-tool attributed spend whenever multi-tool requests exist" + ), + ) + start_date: str | None = None + end_date: str | None = None diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 2f4d2c98b92..808ca7c03a5 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -1534,14 +1534,27 @@ class PromptTokensDetailsWrapper( audio_length_seconds: Optional[float] = None """Length of audio sent to the model. Used for multimodal embeddings priced per audio-second.""" + cache_write_tokens: Optional[int] = None + """Number of cache write (creation) tokens sent to the model. OpenAI naming (prompt_tokens_details.cache_write_tokens); this is the canonical field.""" + cache_creation_tokens: Optional[int] = None - """Number of cache creation tokens sent to the model. Used for Anthropic prompt caching.""" + """Number of cache creation tokens sent to the model. Anthropic/Bedrock naming; kept in sync with cache_write_tokens (assigning either mirrors to the other).""" cache_creation_token_details: Optional[CacheCreationTokenDetails] = None """Details of cache creation tokens sent to the model. Used for tracking 5m/1h cache creation tokens for Anthropic prompt caching.""" + def __setattr__(self, name: str, value: object) -> None: + super().__setattr__(name, value) + if name == "cache_write_tokens": + super().__setattr__("cache_creation_tokens", value) + elif name == "cache_creation_tokens": + super().__setattr__("cache_write_tokens", value) + def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) + self.cache_write_tokens = ( + self.cache_write_tokens if self.cache_write_tokens is not None else self.cache_creation_tokens + ) if self.character_count is None: del self.character_count if self.image_count is None: @@ -1554,6 +1567,8 @@ class PromptTokensDetailsWrapper( del self.web_search_requests if self.tool_use_tokens is None: del self.tool_use_tokens + if self.cache_write_tokens is None: + del self.cache_write_tokens if self.cache_creation_tokens is None: del self.cache_creation_tokens if self.cache_creation_token_details is None: @@ -1662,10 +1677,10 @@ class Usage(SafeAttributeModel, CompletionUsage): if "cache_creation_input_tokens" in params and isinstance(params["cache_creation_input_tokens"], int): if _prompt_tokens_details is None: _prompt_tokens_details = PromptTokensDetailsWrapper( - cache_creation_tokens=params["cache_creation_input_tokens"] + cache_write_tokens=params["cache_creation_input_tokens"] ) else: - _prompt_tokens_details.cache_creation_tokens = params["cache_creation_input_tokens"] + _prompt_tokens_details.cache_write_tokens = params["cache_creation_input_tokens"] super().__init__( prompt_tokens=prompt_tokens or 0, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index c9d871fc41d..749b2566c2a 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -1502,6 +1502,222 @@ "supports_parallel_tool_use_config": true, "prompt_cache_min_tokens": 1024 }, + "anthropic.claude-opus-5": { + "bedrock_converse_supports_strict_tools": false, + "supports_adaptive_thinking": true, + "supports_mid_conversation_system": true, + "cache_creation_input_token_cost": 6.25e-06, + "cache_creation_input_token_cost_above_1hr": 1e-05, + "cache_read_input_token_cost": 5e-07, + "input_cost_per_token": 5e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 2.5e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_native_structured_output": true, + "supports_max_reasoning_effort": true, + "supports_output_config": true, + "supports_parallel_tool_use_config": true, + "prompt_cache_min_tokens": 512 + }, + "global.anthropic.claude-opus-5": { + "bedrock_converse_supports_strict_tools": false, + "supports_adaptive_thinking": true, + "supports_mid_conversation_system": true, + "cache_creation_input_token_cost": 6.25e-06, + "cache_creation_input_token_cost_above_1hr": 1e-05, + "cache_read_input_token_cost": 5e-07, + "input_cost_per_token": 5e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 2.5e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_native_structured_output": true, + "supports_max_reasoning_effort": true, + "supports_output_config": true, + "supports_parallel_tool_use_config": true, + "prompt_cache_min_tokens": 512 + }, + "us.anthropic.claude-opus-5": { + "bedrock_converse_supports_strict_tools": false, + "supports_adaptive_thinking": true, + "supports_mid_conversation_system": true, + "cache_creation_input_token_cost": 6.875e-06, + "cache_creation_input_token_cost_above_1hr": 1.1e-05, + "cache_read_input_token_cost": 5.5e-07, + "input_cost_per_token": 5.5e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 2.75e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_native_structured_output": true, + "supports_max_reasoning_effort": true, + "supports_output_config": true, + "supports_parallel_tool_use_config": true, + "prompt_cache_min_tokens": 512 + }, + "eu.anthropic.claude-opus-5": { + "bedrock_converse_supports_strict_tools": false, + "supports_adaptive_thinking": true, + "supports_mid_conversation_system": true, + "cache_creation_input_token_cost": 6.875e-06, + "cache_creation_input_token_cost_above_1hr": 1.1e-05, + "cache_read_input_token_cost": 5.5e-07, + "input_cost_per_token": 5.5e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 2.75e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_native_structured_output": true, + "supports_max_reasoning_effort": true, + "supports_output_config": true, + "supports_parallel_tool_use_config": true, + "prompt_cache_min_tokens": 512 + }, + "au.anthropic.claude-opus-5": { + "bedrock_converse_supports_strict_tools": false, + "supports_adaptive_thinking": true, + "supports_mid_conversation_system": true, + "cache_creation_input_token_cost": 6.875e-06, + "cache_creation_input_token_cost_above_1hr": 1.1e-05, + "cache_read_input_token_cost": 5.5e-07, + "input_cost_per_token": 5.5e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 2.75e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_native_structured_output": true, + "supports_max_reasoning_effort": true, + "supports_output_config": true, + "supports_parallel_tool_use_config": true, + "prompt_cache_min_tokens": 512 + }, + "jp.anthropic.claude-opus-5": { + "bedrock_converse_supports_strict_tools": false, + "supports_adaptive_thinking": true, + "supports_mid_conversation_system": true, + "cache_creation_input_token_cost": 6.875e-06, + "cache_creation_input_token_cost_above_1hr": 1.1e-05, + "cache_read_input_token_cost": 5.5e-07, + "input_cost_per_token": 5.5e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 2.75e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_native_structured_output": true, + "supports_max_reasoning_effort": true, + "supports_output_config": true, + "supports_parallel_tool_use_config": true, + "prompt_cache_min_tokens": 512 + }, "anthropic.claude-opus-4-8": { "bedrock_converse_supports_strict_tools": false, "supports_adaptive_thinking": true, @@ -2756,6 +2972,38 @@ "supports_xhigh_reasoning_effort": true, "supports_max_reasoning_effort": true }, + "azure_ai/claude-opus-5": { + "supports_mid_conversation_system": true, + "supports_adaptive_thinking": true, + "input_cost_per_token": 5e-06, + "output_cost_per_token": 2.5e-05, + "litellm_provider": "azure_ai", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "cache_creation_input_token_cost": 6.25e-06, + "cache_creation_input_token_cost_above_1hr": 1e-05, + "cache_read_input_token_cost": 5e-07, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_max_reasoning_effort": true, + "prompt_cache_min_tokens": 512 + }, "azure_ai/claude-opus-4-8": { "supports_mid_conversation_system": true, "supports_adaptive_thinking": true, @@ -11846,6 +12094,44 @@ "supports_output_config": true, "prompt_cache_min_tokens": 512 }, + "claude-opus-5": { + "cache_creation_input_token_cost": 6.25e-06, + "cache_creation_input_token_cost_above_1hr": 1e-05, + "cache_read_input_token_cost": 5e-07, + "input_cost_per_token": 5e-06, + "litellm_provider": "anthropic", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 2.5e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_native_structured_output": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_max_reasoning_effort": true, + "provider_specific_entry": { + "us": 1.1, + "fast": 2.0 + }, + "supports_output_config": true, + "supports_speed": true, + "prompt_cache_min_tokens": 512 + }, "claude-opus-4-8": { "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, @@ -36987,6 +37273,70 @@ "supports_xhigh_reasoning_effort": true, "supports_max_reasoning_effort": true }, + "vertex_ai/claude-opus-5": { + "supports_mid_conversation_system": true, + "supports_adaptive_thinking": true, + "cache_creation_input_token_cost": 6.25e-06, + "cache_creation_input_token_cost_above_1hr": 1e-05, + "cache_read_input_token_cost": 5e-07, + "input_cost_per_token": 5e-06, + "litellm_provider": "vertex_ai-anthropic_models", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 2.5e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_max_reasoning_effort": true, + "prompt_cache_min_tokens": 512 + }, + "vertex_ai/claude-opus-5@default": { + "supports_mid_conversation_system": true, + "supports_adaptive_thinking": true, + "cache_creation_input_token_cost": 6.25e-06, + "cache_creation_input_token_cost_above_1hr": 1e-05, + "cache_read_input_token_cost": 5e-07, + "input_cost_per_token": 5e-06, + "litellm_provider": "vertex_ai-anthropic_models", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 2.5e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_max_reasoning_effort": true, + "prompt_cache_min_tokens": 512 + }, "vertex_ai/claude-opus-4-8": { "supports_mid_conversation_system": true, "supports_adaptive_thinking": true, diff --git a/pyproject.toml b/pyproject.toml index 080d06258ed..44c1967ad9b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -99,6 +99,10 @@ utils = [ "numpydoc>=1.8.0,<2.0", ] caching = ["diskcache>=5.6.3,<6.0"] +# SAML SSO for the admin UI. python3-saml pulls in xmlsec/lxml, whose wheels +# bundle the native libxmlsec1/libxml2 libraries, so no system packages are +# required. Kept out of the base `proxy` extra so it stays optional. +saml = ["python3-saml>=1.16.0,<2.0"] semantic-router = [ "semantic-router>=0.1.15,<1.0; python_version < '3.14'", "aurelio-sdk>=0.0.19,<1.0; python_version < '3.14'", @@ -117,6 +121,14 @@ stt-nvidia-riva = [ "numpy>=1.26.0", ] google = ["google-cloud-aiplatform>=1.133.0,<2.0"] +bedrock-realtime = [ + # Bedrock Nova Sonic realtime (speech-to-speech) uses the + # InvokeModelWithBidirectionalStream API, which boto3 cannot do. This + # experimental AWS SDK (with its smithy-* deps, pulled transitively) + # provides the bidirectional stream; imported lazily in the realtime + # handler so litellm core stays usable without it. + "aws-sdk-bedrock-runtime>=0.7.0,<0.8.0; python_version >= '3.12'", +] proxy-runtime = [ # Historically bundled in the proxy Docker images via requirements.txt. # Keep these in a dedicated extra so uv-based images preserve the same @@ -190,6 +202,7 @@ e2e-dev = [ "playwright==1.61.0", "websockets>=15.0.1,<16.0", "locust==2.45.0", + "mcp>=1.28.1,<2.0", ] proxy-dev = [ "prisma==0.11.0", diff --git a/pyrightconfig.json b/pyrightconfig.json index eabfbf515c4..2686ccd73d9 100644 --- a/pyrightconfig.json +++ b/pyrightconfig.json @@ -1,7 +1,7 @@ { "include": ["litellm"], "ignore": [], - "exclude": ["**/node_modules", "**/__pycache__", "tests/e2e/claude_code", "litellm/types/utils.py", "litellm/proxy/_types.py"], + "exclude": ["**/node_modules", "**/__pycache__", "tests/e2e/claude_code", "tests/e2e/ui", "litellm/types/utils.py", "litellm/proxy/_types.py"], "pythonVersion": "3.12", "typeCheckingMode": "strict", "enableTypeIgnoreComments": false, diff --git a/ruff-strict-budget.json b/ruff-strict-budget.json index a5950677739..437d09b726b 100644 --- a/ruff-strict-budget.json +++ b/ruff-strict-budget.json @@ -222,7 +222,7 @@ "limit": 38 }, "RET504": { - "limit": 721 + "limit": 719 }, "RUF010": { "limit": 874 @@ -324,7 +324,7 @@ "limit": 883 }, "UP006": { - "limit": 12792 + "limit": 12789 }, "UP007": { "limit": 2570 diff --git a/scripts/pre_commit_lint.sh b/scripts/pre_commit_lint.sh index cce0cb61c1e..150a4bbf9de 100755 --- a/scripts/pre_commit_lint.sh +++ b/scripts/pre_commit_lint.sh @@ -5,6 +5,7 @@ # gating CI checks, so a clean run means a green CI lint: # - litellm/ Python staged -> `make lint` (test-linting.yml's lint job) # - tests/e2e Python staged -> `make lint-e2e-basedpyright` (test-linting.yml's e2e type-check step) +# + raw HTTP client ban (test-code-quality.yml's check_e2e_no_raw_requests) # - dashboard staged -> prettier + eslint + lint budgets (test-litellm-ui-build.yml's frontend-lint) # - proxy/types staged -> regenerate dashboard API types and fail on drift (check-ui-api-types.yml) # @@ -112,6 +113,12 @@ if [ -n "$e2e_py_files" ] && [ -z "$litellm_py_files" ]; then make lint-e2e-basedpyright || { echo "✗ tests/e2e basedpyright failed. Fix the errors above, then re-run make pre-commit." >&2; status=1; } fi +if [ -n "$e2e_py_files" ]; then + echo "pre-commit: checking tests/e2e raw HTTP client ban (check_e2e_no_raw_requests)" + uv run --no-sync python tests/code_coverage_tests/check_e2e_no_raw_requests.py \ + || { echo "✗ Raw HTTP client import in tests/e2e. Route the call through tests/e2e/e2e_http.py, then re-run make pre-commit." >&2; status=1; } +fi + if [ -n "$ui_prettier_files" ] || [ -n "$ui_eslint_files" ]; then echo "pre-commit: linting dashboard (prettier + eslint + lint budgets)" if [ ! -d ui/litellm-dashboard/node_modules ]; then diff --git a/tests/code_coverage_tests/check_e2e_no_raw_requests.py b/tests/code_coverage_tests/check_e2e_no_raw_requests.py new file mode 100644 index 00000000000..fe6a77fc26c --- /dev/null +++ b/tests/code_coverage_tests/check_e2e_no_raw_requests.py @@ -0,0 +1,84 @@ +"""tests/e2e routes every HTTP call through the typed transport (e2e_http.py), so +raw HTTP client imports (requests, urllib.request, httpx, aiohttp, http.client) are +banned in suite code. Importing requests' exception types for catching is fine +anywhere; a small allowlist grandfathers the files that legitimately make raw calls +(the transport itself, the root conftest liveness probe, the claude_code version +resolver's constant registry URL fetch, and the mcp OAuth client, whose httpx +client is the object the official mcp SDK's streamable_http_client requires and so +cannot go through the sync requests transport). Referenced by tests/e2e/CLAUDE.md.""" + +from __future__ import annotations + +import ast +import sys +from pathlib import Path + +E2E_DIR = Path(__file__).resolve().parents[1] / "e2e" + +BANNED_MODULES = ("requests", "urllib.request", "http.client", "httpx", "aiohttp") + +ALLOWED_RAW_CLIENT_FILES = { + "e2e_http.py": ("requests",), + "conftest.py": ("requests",), + "claude_code/pr_gate_version_resolver.py": ("urllib.request",), + "mcp/oauth_chat_client.py": ("httpx",), +} + +EXCEPTION_ONLY_NAMES = frozenset({"RequestException", "ConnectionError", "Timeout", "HTTPError"}) + + +def _is_banned(module: str) -> bool: + return any(module == banned or module.startswith(banned + ".") for banned in BANNED_MODULES) + + +def _banned_imports(tree: ast.Module) -> tuple[tuple[str, int], ...]: + plain = tuple( + (alias.name, node.lineno) + for node in ast.walk(tree) + if isinstance(node, ast.Import) + for alias in node.names + if _is_banned(alias.name) + ) + from_imports = tuple( + (node.module, node.lineno) + for node in ast.walk(tree) + if isinstance(node, ast.ImportFrom) + and node.module is not None + and _is_banned(node.module) + and not all(alias.name in EXCEPTION_ONLY_NAMES for alias in node.names) + ) + return plain + from_imports + + +def _violations_in(path: Path) -> tuple[str, ...]: + relative = path.relative_to(E2E_DIR).as_posix() + allowed = ALLOWED_RAW_CLIENT_FILES.get(relative, ()) + tree = ast.parse(path.read_text(encoding="utf-8"), filename=str(path)) + return tuple( + f"tests/e2e/{relative}:{lineno}: raw HTTP client import '{module}'" + for module, lineno in _banned_imports(tree) + if module not in allowed + ) + + +def main() -> int: + violations = tuple( + violation + for path in sorted(E2E_DIR.rglob("*.py")) + for violation in _violations_in(path) + ) + for violation in violations: + print(violation) + if violations: + print( + f"\n{len(violations)} raw HTTP client import(s) in tests/e2e. " + "Route the call through tests/e2e/e2e_http.py (get_external for absolute " + "third-party URLs) so it gets the typed Result handling." + ) + return 1 + print("tests/e2e raw HTTP client check passed") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tests/code_coverage_tests/recursive_detector.py b/tests/code_coverage_tests/recursive_detector.py index fa81efde5db..0bc3cebdd5a 100644 --- a/tests/code_coverage_tests/recursive_detector.py +++ b/tests/code_coverage_tests/recursive_detector.py @@ -56,6 +56,7 @@ IGNORE_FUNCTIONS = [ "apply_json_merge_patch", # max depth set (_MAX_MERGE_DEPTH=64); fails closed by raising ValueError at the cap. "_filter_mcp_argument_value", # max depth set (DEFAULT_MAX_RECURSE_DEPTH); fails closed by blocking the MCP call at the cap. "_redact_scanned_content", # max depth set (DEFAULT_MAX_RECURSE_DEPTH); fails closed by returning "[REDACTED]" at the cap. + "_iter_fallback_targets", # max depth set (2 * ROUTER_MAX_FALLBACKS); fails closed by raising ValueError at the cap. ] diff --git a/tests/e2e/CLAUDE.md b/tests/e2e/CLAUDE.md index bf34a896771..17aee22560c 100644 --- a/tests/e2e/CLAUDE.md +++ b/tests/e2e/CLAUDE.md @@ -13,14 +13,16 @@ Each subdirectory under `tests/e2e/` is one suite, scoped to an endpoint family - `realtime/` - realtime websocket sessions, including the pipecat audio path - `quota_management/` - quota enforcement and accounting, one subfolder per behavior: `ratelimit/` (rpm/tpm blocks, window reset, pacing headers on live traffic), `budgets/` (budget definition, enforcement, and reset windows: key, team, tag, soft, multi-window), and `spend_tracking/` (spend logging and cost attribution on `/spend/*`) - `management/` - key/team/user/organization management routes: create/update/delete persistence via the info routes, team membership, and llm-only-key route denials (API surface; not Playwright) -- `mcp/` - the MCP server surface over api_key auth against the real Datadog remote MCP server only (see "MCP suite: real Datadog only" below) +- `a2a/` - the A2A (agent-to-agent) surface: admin registration via `/v1/agents`, proxy-fronted card discovery at `/.well-known/agent-card.json`, and JSON-RPC `message/send` invocation, driving agents backed by the litellm completion bridge (a real provider) and asserting protocol-version normalization (0.3 vs 1.0) +- `mcp/` - the MCP server surface over api_key auth against the real Datadog remote MCP server (see "MCP suite: real Datadog only" below); plus the gateway-managed OAuth (authorization_code) path exercised through `/chat/completions`, the one behavior Datadog's static-header auth cannot reach, seeding the per-user upstream token via the interactive authorize dance driven with the mcp SDK's own OAuth client (headless-browser consent from a saved session) and asserting the completion lists and executes the server's tools with the stored per-user token - `logging/` - logging-integration delivery (datadog and friends) - `security/` - secret handling and log-leak protection - `router/` - routing and reliability behavior (fallbacks, cooldowns) -- `load/` - throughput/performance under concurrency: drives real concurrent traffic through the whole stack with Locust and asserts a throughput SLO; marked `load` so the parent conftest collects it last and it never perturbs latency-sensitive suites +- `load/` - throughput/performance under concurrency: drives real concurrent traffic through the whole stack with Locust and asserts a throughput SLO; marked `load` so the parent conftest collects it last and it never perturbs latency-sensitive suites. Also home of the weekly session-anomaly test (`test_weekly_session_anomaly_e2e.py`): Claude Code-shaped multi-turn sessions against real providers with ceilings on error rate, cache read/write, turn time, and spend; additionally marked `weekly` and deselected unless `E2E_WEEKLY_ANOMALY` is set, because it spends real provider money (driven by `.github/workflows/weekly_load_anomaly.yml`) - `other/` - the holding-pen suite for the `other.*` registry cluster with no home of its own yet: the master-key auth gate and the process-lifecycle health probes (liveness, public readiness, authenticated readiness diagnostics). Promote a cluster out once it is large/stable enough for its own suite - `gateway/` - proxy configuration only (`litellm-config.yml`); no tests - `claude_code/` - the Claude Code compatibility matrix: drives the real `claude` CLI (and HTTP probes) against a proxy for each feature x provider cell, reporting tagged-union outcomes via the `compat_result` fixture; ships its own driver/builder/publisher plus `_*_unit_tests/` trees. The HTTP probes ride the shared transport (`ProxyClient.count_tokens` / `ProxyClient.messages`); the CLI-driving path stays bespoke +- `ui/` - the Admin UI browser suite: Playwright in TypeScript, driving the dashboard served by a live proxy on port 4000 (seeded postgres + mock LLM upstream; see its `run_e2e.sh`). It is a self-contained npm package with its own lockfile and does not use the Python harness, pytest markers, or the shared transport; the Python rules in this file (typed models, `Result` unions, basedpyright zero-error gate) do not apply inside it. Its only Python file, `fixtures/mock_llm_server/server.py`, is excluded from the e2e basedpyright gate via the root `pyrightconfig.json` ## MCP suite: real Datadog only @@ -31,6 +33,7 @@ Every test under `tests/e2e/mcp/` must exercise the proxy against the real Datad - Prefer calling real Datadog tools that prove the product path (e.g. `search_datadog_logs` for list/call and permission denials). Seed a unique marker (`e2e-datadog-mcp-*`) in a chat completion when you need a log the tool can find; dual-read with `dd_logs` from conftest when delivery matters - Delete the MCP server (and any keys) through `resources.defer` the same way every other suite tears down - If a new MCP behavior cannot be covered with Datadog's tool surface, say so in the PR and get agreement before inventing another upstream; the default is always Datadog +- The one standing exception is `test_mcp_chat_completion_oauth_e2e.py`. Datadog authenticates with the static `DD-API-KEY` / `DD-APPLICATION-KEY` headers and exposes no authorize/token dance at all, so it cannot exercise gateway-managed OAuth or per-user token seeding in any form. That test drives a real Linear MCP server instead; it is still a real remote upstream, so the no-mock, no-fixture rule above holds unchanged ## Lay the pattern down in a class @@ -131,7 +134,7 @@ reliability... behavior : fallback | retry | cooldown | timeout | routing | cache | circuit_breaker | perf variant : 5xx | context_window | content_policy | 429 | timeout simple_shuffle | usage_based | latency_based | cost_based | least_busy - latency | throughput (perf only; SLO/threshold assertion, not binary) + latency | throughput | session_anomaly (perf only; SLO/threshold assertion, not binary) assertion : routes_to_fallback | succeeds_within_retries | picks_under_tpm | returns_cached | trips_then_recovers | under_slo e.g. reliability.fallback.context_window.routes_to_fallback exercised_on=[chat_completions] diff --git a/tests/e2e/a2a/a2a_client.py b/tests/e2e/a2a/a2a_client.py new file mode 100644 index 00000000000..916ef623d3a --- /dev/null +++ b/tests/e2e/a2a/a2a_client.py @@ -0,0 +1,343 @@ +"""Client for the proxy's A2A (agent-to-agent) surface. + +An A2A agent is registered admin-side via POST /v1/agents with an agent card and +litellm_params; the proxy fronts it at /a2a/{id}, serving a proxy-owned agent card +at /.well-known/agent-card.json and accepting A2A JSON-RPC calls at /a2a/{id}. This +suite registers agents backed by the litellm_completion_bridge (custom_llm_provider ++ model), so message/send runs a real provider completion and comes back in the +agent's pinned A2A protocol version. The A2A request/response models are co-located +here because only this suite uses them. +""" + +from __future__ import annotations + +import warnings +from dataclasses import dataclass + +from pydantic import BaseModel, ConfigDict, Field + +from e2e_http import NoBody, Result, get_external, is_ok +from proxy_client import ProxyClient + + +class A2ACapabilities(BaseModel): + streaming: bool | None = None + push_notifications: bool | None = Field(default=None, serialization_alias="pushNotifications") + + +class A2ASkill(BaseModel): + id: str + name: str + description: str + tags: list[str] + examples: list[str] | None = None + + +class A2AProvider(BaseModel): + organization: str + url: str + + +class AgentCardParams(BaseModel): + """The upstream agent card an admin registers. `protocolVersion` is the field the + proxy validates against SUPPORTED_A2A_PROTOCOL_VERSIONS on registration.""" + + protocol_version: str = Field(serialization_alias="protocolVersion") + name: str + description: str + version: str + url: str | None = None + capabilities: A2ACapabilities = A2ACapabilities() + skills: list[A2ASkill] + default_input_modes: list[str] = Field(default=["text"], serialization_alias="defaultInputModes") + default_output_modes: list[str] = Field(default=["text"], serialization_alias="defaultOutputModes") + preferred_transport: str | None = Field(default=None, serialization_alias="preferredTransport") + + +class UpstreamAgentCard(BaseModel): + """A real published agent card parsed from a public /.well-known endpoint. Keys on + the A2A wire aliases so `model_validate_json` reads the served JSON and + `model_dump(by_alias=True)` re-emits it unchanged for verbatim registration; it is + only ever fetched-and-validated, never hand-constructed, so aliasing on the wire + names does not affect any call site.""" + + model_config = ConfigDict(populate_by_name=True) + + protocol_version: str = Field(alias="protocolVersion") + name: str + description: str + version: str + url: str + provider: A2AProvider | None = None + documentation_url: str | None = Field(default=None, alias="documentationUrl") + capabilities: A2ACapabilities = A2ACapabilities() + skills: list[A2ASkill] + default_input_modes: list[str] = Field(default=["text"], alias="defaultInputModes") + default_output_modes: list[str] = Field(default=["text"], alias="defaultOutputModes") + preferred_transport: str | None = Field(default=None, alias="preferredTransport") + + +class A2ABridgeParams(BaseModel): + """litellm_params that route the agent through the completion bridge: an A2A + message/send is transformed into a litellm.acompletion against this provider.""" + + model_config = ConfigDict(protected_namespaces=()) + + custom_llm_provider: str + model: str + api_key: str | None = None + + +class AgentRegisterBody(BaseModel): + agent_name: str + agent_card_params: AgentCardParams | UpstreamAgentCard + litellm_params: A2ABridgeParams | None = None + + +class A2ASecurityScheme(BaseModel): + type: str + scheme: str + + +class A2AInterface(BaseModel): + model_config = ConfigDict(populate_by_name=True) + + url: str + protocol_version: str | None = Field(default=None, alias="protocolVersion") + + +class ServedAgentCard(BaseModel): + """The proxy-owned card, either nested under a registration response's + `agent_card_params` or served raw at /.well-known/agent-card.json. The proxy + rewrites `url`/`supportedInterfaces` to itself and replaces the security scheme + with its own virtual-key bearer scheme.""" + + model_config = ConfigDict(populate_by_name=True) + + protocol_version: str = Field(alias="protocolVersion") + name: str + url: str | None = None + security_schemes: dict[str, A2ASecurityScheme] | None = Field(default=None, alias="securitySchemes") + security: list[dict[str, list[str]]] | None = None + supported_interfaces: list[A2AInterface] | None = Field(default=None, alias="supportedInterfaces") + + +class AgentResponse(BaseModel): + agent_id: str + agent_name: str + agent_card_params: ServedAgentCard + + +class A2ATextPart(BaseModel): + kind: str = "text" + text: str + + +class A2ASearchPropertiesParams(BaseModel): + """The strict param schema of the published property agent's `search_properties` + skill (unknown keys are rejected upstream), so a natural-language query like + "properties for sale in SF under $2M" is expressed as typed fields.""" + + un_locode: str | None = None + service_type: str | None = None + property_type: str | None = None + bedrooms_min: int | None = None + asking_price_max: float | None = None + limit: int | None = None + + +class A2ASkillInvocation(BaseModel): + skill: str + params: A2ASearchPropertiesParams + + +class A2ADataPart(BaseModel): + kind: str = "data" + data: A2ASkillInvocation + + +class A2AOutboundMessage(BaseModel): + role: str = "user" + parts: list[A2ATextPart | A2ADataPart] + message_id: str = Field(serialization_alias="messageId") + + +class A2AMessageSendParams(BaseModel): + message: A2AOutboundMessage + + +class A2AJsonRpcRequest(BaseModel): + jsonrpc: str = "2.0" + id: str + method: str = "message/send" + params: A2AMessageSendParams + + +class A2AResponsePart(BaseModel): + kind: str | None = None + text: str | None = None + + +class A2AResponseMessage(BaseModel): + model_config = ConfigDict(populate_by_name=True) + + message_id: str | None = Field(default=None, alias="messageId") + role: str | None = None + parts: list[A2AResponsePart] = [] + + +class A2ATaskStatus(BaseModel): + state: str | None = None + message: A2AResponseMessage | None = None + + +class A2AListingLocation(BaseModel): + """Only the location fields a test reads back off a returned listing.""" + + un_locode: str | None = None + + +class A2AListing(BaseModel): + """A single property card from the agent's `search_results` artifact; only the + identity/location fields a test asserts on are modelled.""" + + raia_id: str + property_type: str | None = None + service_type: str | None = None + location: A2AListingLocation = A2AListingLocation() + + +class A2ASearchResults(BaseModel): + """The DataPart payload the property agent's `search_properties` skill returns: + the run count plus the listing cards themselves. Proof the tool actually ran and + matched, not just that the task completed with some text.""" + + total: int + count: int + listings: list[A2AListing] = [] + + +class A2AArtifactPart(BaseModel): + kind: str | None = None + data: A2ASearchResults | None = None + + +class A2AArtifact(BaseModel): + model_config = ConfigDict(populate_by_name=True) + + artifact_id: str | None = Field(default=None, alias="artifactId") + name: str | None = None + parts: list[A2AArtifactPart] = [] + + +class A2AResult(BaseModel): + """A message/send result. In 0.3 the message fields sit directly on the result + (`kind`/`role`/`parts`); in 1.0 they are nested under `message`; a real agent that + runs a task replies with a `task` whose agent text lives on `status.message` and + whose tool output lives on `artifacts`. `text` reads the agent's reply from + whichever shape the served version produced; `search_results` reads the tool's + structured output when the agent ran a skill.""" + + model_config = ConfigDict(populate_by_name=True) + + kind: str | None = None + role: str | None = None + message_id: str | None = Field(default=None, alias="messageId") + parts: list[A2AResponsePart] = [] + message: A2AResponseMessage | None = None + status: A2ATaskStatus | None = None + artifacts: list[A2AArtifact] = [] + + @property + def text(self) -> str: + if self.message is not None: + parts = self.message.parts + elif self.parts: + parts = self.parts + elif self.status is not None and self.status.message is not None: + parts = self.status.message.parts + else: + parts = [] + return "".join(part.text or "" for part in parts) + + @property + def is_nested_v1_shape(self) -> bool: + return self.message is not None + + @property + def search_results(self) -> A2ASearchResults | None: + for artifact in self.artifacts: + for part in artifact.parts: + if part.data is not None: + return part.data + return None + + +class A2AError(BaseModel): + code: int + message: str + + +class A2AResponse(BaseModel): + jsonrpc: str + id: str | None = None + result: A2AResult | None = None + error: A2AError | None = None + + +@dataclass(frozen=True, slots=True) +class A2AClient: + proxy: ProxyClient + + def register_agent(self, body: AgentRegisterBody) -> Result[AgentResponse]: + return self.proxy.transport.post( + "/v1/agents", + headers=self.proxy.transport.master, + json=body, + response_type=AgentResponse, + ) + + def get_agent(self, agent_id: str) -> Result[AgentResponse]: + return self.proxy.transport.get( + f"/v1/agents/{agent_id}", + headers=self.proxy.transport.master, + params=NoBody(), + response_type=AgentResponse, + ) + + def delete_agent(self, agent_id: str) -> None: + result = self.proxy.transport.delete( + f"/v1/agents/{agent_id}", + headers=self.proxy.transport.master, + json=NoBody(), + response_type=NoBody, + ) + if not is_ok(result): + warnings.warn(f"delete_agent({agent_id!r}) failed: {result}", stacklevel=2) + + def agent_card(self, agent_id: str, key: str) -> Result[ServedAgentCard]: + return self.proxy.transport.get( + f"/a2a/{agent_id}/.well-known/agent-card.json", + headers=self.proxy.transport.bearer(key), + params=NoBody(), + response_type=ServedAgentCard, + ) + + def send_message(self, agent_id: str, key: str, body: A2AJsonRpcRequest) -> Result[A2AResponse]: + return self.proxy.transport.post( + f"/a2a/{agent_id}", + headers=self.proxy.transport.bearer(key), + json=body, + response_type=A2AResponse, + ) + + +def build_a2a_client(proxy: ProxyClient) -> A2AClient: + return A2AClient(proxy=proxy) + + +def fetch_agent_card(url: str, *, timeout: float = 20.0) -> Result[UpstreamAgentCard]: + """Fetch a live A2A agent card from its /.well-known endpoint and parse it into the + registration model, so a test can register a real published card verbatim rather + than a hand-rolled one.""" + return get_external(url, response_type=UpstreamAgentCard, timeout=timeout) diff --git a/tests/e2e/a2a/conftest.py b/tests/e2e/a2a/conftest.py new file mode 100644 index 00000000000..93f3b56c8f7 --- /dev/null +++ b/tests/e2e/a2a/conftest.py @@ -0,0 +1,17 @@ +"""A2A suite's `client` fixture. + +The shared lifecycle (resources/scoped_key), proxy liveness gate, and e2e marker +live in the parent tests/e2e/conftest.py. A2AClient holds the shared ProxyClient, +so the `resources` fixture cleans up keys this suite creates; agents are torn down +via `resources.defer(...)` in each test. +""" + +import pytest + +from a2a_client import A2AClient, build_a2a_client +from proxy_client import ProxyClient + + +@pytest.fixture(scope="session") +def client(proxy: ProxyClient) -> A2AClient: + return build_a2a_client(proxy) diff --git a/tests/e2e/a2a/test_a2a_agent_e2e.py b/tests/e2e/a2a/test_a2a_agent_e2e.py new file mode 100644 index 00000000000..802b01455b1 --- /dev/null +++ b/tests/e2e/a2a/test_a2a_agent_e2e.py @@ -0,0 +1,212 @@ +"""A2A agents end to end, against a live proxy. + +An admin registers an agent whose card pins an A2A protocol version and whose +litellm_params route it through the completion bridge; a caller then discovers the +proxy-owned card and drives it over A2A JSON-RPC. These tests assert the recorded +state (the agent persists, a spend row lands) and the enforced behavior (the served +card points back at the proxy, message/send returns a real completion in the pinned +protocol version, and an unsupported version is refused at registration). +""" + +from __future__ import annotations + +import pytest + +from a2a_client import ( + A2ABridgeParams, + A2AClient, + A2ADataPart, + A2AJsonRpcRequest, + A2AMessageSendParams, + A2AOutboundMessage, + A2ASearchPropertiesParams, + A2ASkill, + A2ASkillInvocation, + A2ATextPart, + AgentCardParams, + AgentRegisterBody, + AgentResponse, + fetch_agent_card, +) +from e2e_config import unique_marker +from e2e_http import Result, UnknownApiError, unwrap +from lifecycle import ResourceManager + +BRIDGE = A2ABridgeParams( + custom_llm_provider="anthropic", + model="claude-haiku-4-5", + api_key="os.environ/ANTHROPIC_API_KEY", +) + +MOVEHOME_AGENT_CARD_URL = "https://movehome.org/.well-known/agent.json" +MOVEHOME_ORIGIN = "https://movehome.org" + +pytestmark = pytest.mark.e2e + + +def _register(client: A2AClient, resources: ResourceManager, protocol_version: str) -> AgentResponse: + marker = unique_marker() + body = AgentRegisterBody( + agent_name=f"e2e-a2a-{marker}", + agent_card_params=AgentCardParams( + protocol_version=protocol_version, + name=f"E2E A2A {marker}", + description="e2e agent backed by the litellm completion bridge", + version="1.0.0", + skills=[A2ASkill(id="chat", name="Chat", description="general chat", tags=["chat"])], + ), + litellm_params=BRIDGE, + ) + agent = unwrap(client.register_agent(body)) + resources.defer(lambda: client.delete_agent(agent.agent_id)) + return agent + + +def _register_rejection(client: A2AClient, protocol_version: str) -> Result[AgentResponse]: + marker = unique_marker() + body = AgentRegisterBody( + agent_name=f"e2e-a2a-bad-{marker}", + agent_card_params=AgentCardParams( + protocol_version=protocol_version, + name=f"E2E A2A bad {marker}", + description="rejected at registration", + version="1.0.0", + skills=[A2ASkill(id="chat", name="Chat", description="c", tags=["chat"])], + ), + litellm_params=BRIDGE, + ) + return client.register_agent(body) + + +def _ask(text: str) -> A2AJsonRpcRequest: + return A2AJsonRpcRequest( + id=f"e2e-{unique_marker()}", + params=A2AMessageSendParams( + message=A2AOutboundMessage(parts=[A2ATextPart(text=text)], message_id=unique_marker()) + ), + ) + + +class TestA2AAgentLifecycle: + @pytest.mark.covers("other.a2a.register.persists") + def test_register_persists(self, client: A2AClient, resources: ResourceManager) -> None: + agent = _register(client, resources, "0.3") + fetched = unwrap(client.get_agent(agent.agent_id)) + assert fetched.agent_id == agent.agent_id + assert fetched.agent_name == agent.agent_name + assert fetched.agent_card_params.protocol_version == "0.3" + + @pytest.mark.covers("other.a2a.register.semver_version_accepted") + def test_semver_protocol_version_registers_and_serves(self, client: A2AClient, resources: ResourceManager, scoped_key: str) -> None: + agent = _register(client, resources, "0.3.0") + assert agent.agent_card_params.protocol_version == "0.3" + card = unwrap(client.agent_card(agent.agent_id, scoped_key)) + assert card.protocol_version == "0.3" + assert card.supported_interfaces is not None + assert card.supported_interfaces[0].protocol_version == "0.3" + result = unwrap(client.send_message(agent.agent_id, scoped_key, _ask("Say hi in one word"))).result + assert result is not None + assert result.text != "" + + @pytest.mark.covers("other.a2a.message_send.real_world_agent_replies") + def test_real_world_agent_replies_to_property_query(self, client: A2AClient, resources: ResourceManager, scoped_key: str) -> None: + upstream = unwrap(fetch_agent_card(MOVEHOME_AGENT_CARD_URL)).model_copy(update={"url": MOVEHOME_ORIGIN}) + assert upstream.protocol_version == "0.3.0" + marker = unique_marker() + body = AgentRegisterBody(agent_name=f"e2e-a2a-real-{marker}", agent_card_params=upstream) + agent = unwrap(client.register_agent(body)) + resources.defer(lambda: client.delete_agent(agent.agent_id)) + assert agent.agent_card_params.protocol_version == "0.3" + location = "GBLON" + request = A2AJsonRpcRequest( + id=f"e2e-{unique_marker()}", + params=A2AMessageSendParams( + message=A2AOutboundMessage( + parts=[ + A2ADataPart( + data=A2ASkillInvocation( + skill="search_properties", + params=A2ASearchPropertiesParams(un_locode=location, service_type="long_term", limit=3), + ) + ) + ], + message_id=unique_marker(), + ) + ), + ) + response = unwrap(client.send_message(agent.agent_id, scoped_key, request)) + assert response.error is None + assert response.result is not None + results = response.result.search_results + assert results is not None, "agent returned no search_results artifact; skill did not run" + assert results.total > 0 + assert results.listings, "search_properties matched nothing; agent returned no property cards" + assert all(listing.raia_id for listing in results.listings) + assert all(listing.location.un_locode == location for listing in results.listings) + + @pytest.mark.covers("other.a2a.discovery.proxy_fronted_card") + def test_discovery_card_is_proxy_fronted(self, client: A2AClient, resources: ResourceManager, scoped_key: str) -> None: + agent = _register(client, resources, "0.3") + card = unwrap(client.agent_card(agent.agent_id, scoped_key)) + assert card.url is not None and card.url.endswith(f"/a2a/{agent.agent_id}") + assert card.security_schemes is not None + scheme = next(iter(card.security_schemes.values())) + assert scheme.scheme == "bearer" + assert card.supported_interfaces is not None + assert card.supported_interfaces[0].url == card.url + + @pytest.mark.covers("other.a2a.message_send.bridge_invokes") + def test_message_send_runs_completion_bridge(self, client: A2AClient, resources: ResourceManager, scoped_key: str) -> None: + agent = _register(client, resources, "0.3") + request = _ask("Reply with exactly the word PONG and nothing else") + response = unwrap(client.send_message(agent.agent_id, scoped_key, request)) + assert response.error is None + assert response.result is not None + assert "PONG" in response.result.text.upper() + + rows = client.proxy.poll_logs_for_request_id(request.id) + assert rows, f"no spend log row landed for a2a request {request.id}" + assert rows[0].call_type == "asend_message" + assert rows[0].model == f"a2a_agent/{agent.agent_card_params.name}" + + @pytest.mark.covers("other.a2a.version.serves_pinned_0_3") + def test_pinned_v0_3_serves_flat_message_shape(self, client: A2AClient, resources: ResourceManager, scoped_key: str) -> None: + agent = _register(client, resources, "0.3") + request = _ask("Say hi in one word") + result = unwrap(client.send_message(agent.agent_id, scoped_key, request)).result + assert result is not None + assert not result.is_nested_v1_shape + assert result.kind == "message" + assert result.role == "agent" + assert result.text != "" + + @pytest.mark.covers("other.a2a.version.serves_pinned_1_0") + def test_pinned_v1_0_serves_nested_message_shape(self, client: A2AClient, resources: ResourceManager, scoped_key: str) -> None: + agent = _register(client, resources, "1.0") + request = _ask("Say hi in one word") + result = unwrap(client.send_message(agent.agent_id, scoped_key, request)).result + assert result is not None + assert result.is_nested_v1_shape + assert result.message is not None + assert result.message.role == "ROLE_AGENT" + assert result.text != "" + + @pytest.mark.covers("other.a2a.register.unsupported_version_rejected") + def test_unsupported_protocol_version_rejected(self, client: A2AClient) -> None: + result = _register_rejection(client, "9.9") + match result: + case UnknownApiError(status_code=status, body=detail): + assert status == 400 + assert "protocolVersion" in detail + case _: + pytest.fail(f"expected 400 for unsupported protocolVersion, got {result}") + + @pytest.mark.covers("other.a2a.register.malformed_version_rejected") + def test_malformed_protocol_version_rejected(self, client: A2AClient) -> None: + result = _register_rejection(client, "0.3.garbage") + match result: + case UnknownApiError(status_code=status, body=detail): + assert status == 400 + assert "Unsupported protocolVersion '0.3.garbage'" in detail + case _: + pytest.fail(f"expected 400 for malformed protocolVersion, got {result}") diff --git a/tests/e2e/batches/batch_client.py b/tests/e2e/batches/batch_client.py index 7db5d0b6beb..5cc5d1dae3b 100644 --- a/tests/e2e/batches/batch_client.py +++ b/tests/e2e/batches/batch_client.py @@ -26,16 +26,24 @@ from e2e_http import ( ) from models import LiteLLMParamsBody +UPLOAD_FILENAME = "batch_input.jsonl" + class FileObject(BaseModel): id: str object: str | None = None purpose: str | None = None + filename: str | None = None bytes: int | None = None status: str | None = None created_at: int | None = None +class FileList(BaseModel): + object: str | None = None + data: list[FileObject] = [] + + class BatchObject(BaseModel): id: str object: str | None = None @@ -106,12 +114,30 @@ class BatchClient: _files_path(provider), headers=self.proxy.transport.bearer(key), form=form, - filename="batch_input.jsonl", + filename=UPLOAD_FILENAME, content=content, params=ModelQuery(model=model), response_type=FileObject, ) + def retrieve_file( + self, file_id: str, *, key: str, provider: str | None = None + ) -> Result[FileObject]: + return self.proxy.transport.get( + f"{_files_path(provider)}/{file_id}", + headers=self.proxy.transport.bearer(key), + params=NoBody(), + response_type=FileObject, + ) + + def list_files(self, *, key: str, provider: str | None = None) -> Result[FileList]: + return self.proxy.transport.get( + _files_path(provider), + headers=self.proxy.transport.bearer(key), + params=NoBody(), + response_type=FileList, + ) + def create_batch( self, *, body: BatchCreateBody, key: str, provider: str | None = None ) -> StreamingResponse: diff --git a/tests/e2e/batches/test_batches_e2e.py b/tests/e2e/batches/test_batches_e2e.py index f886c09b705..f9cd2a3f15f 100644 --- a/tests/e2e/batches/test_batches_e2e.py +++ b/tests/e2e/batches/test_batches_e2e.py @@ -23,9 +23,10 @@ from typing import Callable import pytest -from e2e_config import require_env, unique_marker +from e2e_config import unique_marker from batch_client import ( + UPLOAD_FILENAME, BatchClient, BatchCreateBody, BatchObject, @@ -511,6 +512,72 @@ class TestBatchFileContent: ) +class TestOpenAIFiles: + """GET /v1/files (list) and GET /v1/files/{id} (retrieve) over the OpenAI route. + + The proxy lists the OpenAI org's raw file ids, so the list case uploads a raw + (provider-routed) file whose id matches what list returns; retrieve re-encodes + the id it was called with, so the model-encoded upload round-trips unchanged. + """ + + @pytest.mark.covers( + "llm.files.openai.list.nonstream.works", + exercised_on=["files"], + ) + def test_uploaded_file_appears_in_list( + self, client: BatchClient, resources: ResourceManager, batch_deployments: None + ) -> None: + key = resources.key() + file = unwrap( + client.upload_file( + content=render_jsonl(OPENAI_BATCH_MODEL), + form=FileUploadForm(purpose="batch"), + key=key, + provider="openai", + ) + ) + resources.defer( + quietly(lambda: client.delete_file(file.id, key=key, provider="openai")) + ) + + listed = unwrap(client.list_files(key=key)) + assert listed.object is None or listed.object == "list", ( + f"list envelope object={listed.object!r}" + ) + match = next((entry for entry in listed.data if entry.id == file.id), None) + assert match is not None, f"uploaded file {file.id!r} absent from GET /v1/files" + assert match.purpose == "batch", ( + f"listed file must round-trip the upload purpose, got {match.purpose!r}" + ) + + @pytest.mark.covers( + "llm.files.openai.retrieve.nonstream.works", + exercised_on=["files"], + ) + def test_retrieve_round_trips_metadata( + self, client: BatchClient, resources: ResourceManager, batch_deployments: None + ) -> None: + key = resources.key() + file = unwrap( + client.upload_file( + content=render_jsonl(OPENAI_BATCH_MODEL), + form=FileUploadForm(purpose="batch"), + model=OPENAI_BATCH_MODEL, + key=key, + ) + ) + resources.defer(quietly(lambda: client.delete_file(file.id, key=key))) + + fetched = unwrap(client.retrieve_file(file.id, key=key)) + assert fetched.id == file.id, "retrieve must echo the uploaded file id" + assert fetched.purpose == "batch", ( + f"retrieve must round-trip purpose, got {fetched.purpose!r}" + ) + assert fetched.filename == UPLOAD_FILENAME, ( + f"retrieve must round-trip filename, got {fetched.filename!r}" + ) + + BATCH_RL_REQUEST_LINES = 3 BATCH_RL_RPM_LIMIT = 2 @@ -635,14 +702,7 @@ class TestBedrockBatchAssumeRole: def test_unified_batch_create_with_assume_role( self, client: BatchClient, resources: ResourceManager ) -> None: - (role_arn,) = require_env("AWS_ROLE_NAME") - require_env( - "AWS_ACCESS_KEY_ID", - "AWS_SECRET_ACCESS_KEY", - "AWS_REGION", - "AWS_BATCH_S3_BUCKET", - "AWS_BATCH_ROLE_ARN", - ) + role_arn = os.environ["AWS_ROLE_NAME"] session_name = f"e2e-batch-sts-{unique_marker()}"[:64] model_name = batch_model_name("bedrock-sts-batch") @@ -752,7 +812,7 @@ class TestHostedVllmBatch: def test_unified_file_and_batch_create( self, client: BatchClient, resources: ResourceManager ) -> None: - (api_base,) = require_env("HOSTED_VLLM_API_BASE") + api_base = os.environ["HOSTED_VLLM_API_BASE"] api_key = (os.environ.get("HOSTED_VLLM_API_KEY") or "").strip() or None model_id = ( os.environ.get("HOSTED_VLLM_MODEL") or "meta-llama/Llama-3.2-3B-Instruct" diff --git a/tests/e2e/conftest.py b/tests/e2e/conftest.py index 609da6a9b07..eff3b4ddf58 100644 --- a/tests/e2e/conftest.py +++ b/tests/e2e/conftest.py @@ -5,9 +5,9 @@ answers or when credentials/env are missing; they never skip. Pure unit coverage of the harness itself carries no `e2e` marker and runs regardless of whether a proxy is up. -Lifecycle: the `resources` fixture maps the init -> run -> teardown contract -(lifecycle.E2ECase) onto pytest - setup is init(), the test body is run(), and -teardown deletes every resource the test created on the long-lived proxy. +Lifecycle: the `resources` fixture hands each test a lifecycle.ResourceManager - +the test registers a cleanup for every resource it creates, and the fixture's +teardown deletes them all on the long-lived proxy, even when the test fails. Each suite provides its own `client` fixture (a lifecycle.ResourceClient); these shared fixtures build on it. @@ -43,6 +43,10 @@ def pytest_configure(config: pytest.Config) -> None: "markers", "load: heavy throughput/load test; collected last so it never perturbs latency-sensitive suites", ) + config.addinivalue_line( + "markers", + "weekly: real-provider anomaly load test that spends real money; deselected unless E2E_WEEKLY_ANOMALY is set", + ) def pytest_collection_modifyitems(items: list[pytest.Item]) -> None: diff --git a/tests/e2e/coverage_registry/llm_nonconversational.yaml b/tests/e2e/coverage_registry/llm_nonconversational.yaml index 2b456aacefc..63e6fde14a3 100644 --- a/tests/e2e/coverage_registry/llm_nonconversational.yaml +++ b/tests/e2e/coverage_registry/llm_nonconversational.yaml @@ -32,7 +32,7 @@ - {id: llm.files.hosted_vllm.upload.nonstream.works, module: llm, tier: P1, subject_endpoint: files, route: hosted_vllm, capability: basic, streaming: nonstream, assertions: [works], source: "test_batches_e2e.py", rationale: "hosted_vllm OpenAI-compatible file upload"} - {id: llm.rerank.cohere.basic.nonstream.works, module: llm, tier: P1, subject_endpoint: rerank, route: cohere, capability: basic, streaming: nonstream, assertions: [works], source: "test_rerank_e2e.py:29", rationale: "Cohere rerank, top_n + relevance_score"} - {id: llm.files.openai.content.nonstream.works, module: llm, tier: P0, subject_endpoint: files, route: openai, capability: basic, streaming: nonstream, assertions: [works], source: "test_batches_e2e.py", rationale: "GET /v1/files/{id}/content returns uploaded batch JSONL bytes"} -- {id: llm.realtime.bedrock_converse.basic.stream.works, module: llm, tier: P0, subject_endpoint: realtime, route: bedrock_converse, capability: basic, streaming: stream, assertions: [works], source: "test_nova_sonic_realtime_e2e.py", rationale: "Nova Sonic realtime session emits response.done (LIT-2239)"} +- {id: llm.realtime.bedrock_converse.basic.stream.works, module: llm, tier: P0, subject_endpoint: realtime, route: bedrock_converse, capability: basic, streaming: stream, assertions: [works], source: "test_realtime_bedrock_e2e.py", rationale: "Nova Sonic realtime session emits response.done (LIT-2239)"} - {id: llm.rerank.bedrock.basic.nonstream.works, module: llm, tier: P1, subject_endpoint: rerank, route: bedrock_converse, capability: basic, streaming: nonstream, assertions: [works], source: "llms/bedrock/rerank/handler.py", rationale: "Bedrock rerank"} - {id: llm.rerank.together_ai.basic.nonstream.works, module: llm, tier: P1, subject_endpoint: rerank, route: together_ai, capability: basic, streaming: nonstream, assertions: [works], source: "llms/together_ai/rerank/handler.py", rationale: "Together rerank"} - {id: llm.images_generations.openai.basic.nonstream.works, module: llm, tier: P1, subject_endpoint: images_generations, route: openai, capability: basic, streaming: nonstream, assertions: [works], source: "test_image_generation_e2e.py:22", rationale: "OpenAI image gen, b64/url"} diff --git a/tests/e2e/coverage_registry/other.yaml b/tests/e2e/coverage_registry/other.yaml index 6b183cbf9f3..ace4f8bcdc9 100644 --- a/tests/e2e/coverage_registry/other.yaml +++ b/tests/e2e/coverage_registry/other.yaml @@ -28,3 +28,12 @@ - {id: other.config.overrides.audit_logged, module: other, tier: P1, area: config, assertions: [audit_logged], source: "config_override_endpoints.py:67-100", rationale: "Config override mutations audit-logged, values redacted"} - {id: other.key_mgmt.regenerate.grace_period_honored, module: other, tier: P1, area: auth, assertions: [grace_period_honored], source: "key_management_endpoints.py:4503-4560", rationale: "Old key valid during grace_period then revoked"} - {id: other.key_mgmt.spend_reset.resets_to_value, module: other, tier: P1, area: auth, assertions: [resets_to_value], source: "key_management_endpoints.py:4841", rationale: "reset_spend resets accumulated spend"} +- {id: other.a2a.register.persists, module: other, tier: P1, area: a2a, assertions: [persists], source: "agent_endpoints/endpoints.py:325-443", rationale: "POST /v1/agents registers an agent card; GET /v1/agents/{id} reads it back"} +- {id: other.a2a.register.unsupported_version_rejected, module: other, tier: P1, area: a2a, assertions: [unsupported_version_rejected], source: "agent_endpoints/endpoints.py _validate_protocol_version", rationale: "A card pinning a protocolVersion outside SUPPORTED_A2A_PROTOCOL_VERSIONS is refused with 400"} +- {id: other.a2a.register.semver_version_accepted, module: other, tier: P1, area: a2a, assertions: [semver_version_accepted], source: "agent_endpoints/endpoints.py _validate_protocol_version", rationale: "A card pinning a patch-level semver like 0.3.0 (what the Google A2A SDK emits) registers, stores and serves the canonical 0.3 rather than 400ing; regression guard for the v1.92 report"} +- {id: other.a2a.message_send.real_world_agent_replies, module: other, tier: P1, area: a2a, assertions: [real_world_agent_replies], source: "agent_endpoints/a2a_endpoints.py asend_message", rationale: "A real published a2a agent fetched live from a public /.well-known endpoint (pinning the full semver 0.3.0 the a2a-sdk emits) registers, serves the canonical 0.3, and a message/send skill invocation proxies to the live upstream and returns the agent's reply"} +- {id: other.a2a.register.malformed_version_rejected, module: other, tier: P1, area: a2a, assertions: [malformed_version_rejected], source: "a2a/agent_card.py normalize_protocol_version", rationale: "A malformed protocolVersion like 0.3.garbage fails full-string semver validation and is refused with 400 instead of truncating to a supported family"} +- {id: other.a2a.discovery.proxy_fronted_card, module: other, tier: P1, area: a2a, assertions: [proxy_fronted_card], source: "agent_endpoints/a2a_endpoints.py get_agent_card", rationale: "/.well-known/agent-card.json serves the proxy url + supportedInterfaces and the LiteLLM virtual-key bearer scheme, not the upstream"} +- {id: other.a2a.message_send.bridge_invokes, module: other, tier: P1, area: a2a, assertions: [bridge_invokes], source: "a2a_protocol/litellm_completion_bridge/handler.py", rationale: "A2A message/send routes through the completion bridge to a real provider and logs an asend_message spend row"} +- {id: other.a2a.version.serves_pinned_0_3, module: other, tier: P1, area: a2a, assertions: [serves_pinned_0_3], source: "agent_endpoints/a2a_endpoints.py _served_version", rationale: "An agent pinning 0.3 returns the flat 0.3 message shape (parts on the result)"} +- {id: other.a2a.version.serves_pinned_1_0, module: other, tier: P1, area: a2a, assertions: [serves_pinned_1_0], source: "agent_endpoints/a2a_endpoints.py _served_version", rationale: "An agent pinning 1.0 returns the nested 1.0 message shape (result.message with ROLE_AGENT)"} diff --git a/tests/e2e/coverage_registry/reliability.yaml b/tests/e2e/coverage_registry/reliability.yaml index 1538d3f3cda..ebbfd3415a5 100644 --- a/tests/e2e/coverage_registry/reliability.yaml +++ b/tests/e2e/coverage_registry/reliability.yaml @@ -24,3 +24,4 @@ - {id: reliability.timeout.request_timeout.exceeds_deadline, module: reliability, tier: P1, behavior: timeout, variant: request_timeout, assertions: [exceeds_deadline], exercised_on: [chat_completions, messages], source: "litellm/router.py:545-551", rationale: "Per-request timeout raises Timeout"} - {id: reliability.timeout.stream_timeout.exceeds_deadline, module: reliability, tier: P1, behavior: timeout, variant: stream_timeout, assertions: [exceeds_deadline], exercised_on: [chat_completions], source: "litellm/router.py:551", rationale: "Streaming chunk-delivery timeout"} - {id: reliability.perf.throughput.under_slo, module: reliability, tier: P1, behavior: perf, variant: throughput, assertions: [under_slo], exercised_on: [chat_completions, messages], source: grammar, rationale: "Throughput SLO under load"} +- {id: reliability.perf.session_anomaly.under_slo, module: reliability, tier: P1, behavior: perf, variant: session_anomaly, assertions: [under_slo], exercised_on: [messages], source: grammar, rationale: "Weekly Claude Code-shaped multi-turn session load against real providers; ceilings on error rate, warm-turn cache read/write, p95 turn time, and gateway-recorded spend (LIT-4562)"} diff --git a/tests/e2e/e2e_config.py b/tests/e2e/e2e_config.py index 3be339d28a0..feed680bd4a 100644 --- a/tests/e2e/e2e_config.py +++ b/tests/e2e/e2e_config.py @@ -38,6 +38,9 @@ UI_BASE_URL = os.environ.get("E2E_UI_BASE_URL", PROXY_BASE_URL).rstrip("/") CHEAP_ANTHROPIC_MODEL = os.environ.get("E2E_CHEAP_ANTHROPIC_MODEL", "claude-haiku-4-5") CHEAP_OPENAI_MODEL = os.environ.get("E2E_CHEAP_OPENAI_MODEL", "gpt-5.5") +LINEAR_MCP_URL = os.environ.get("E2E_LINEAR_MCP_URL", "https://mcp.linear.app/mcp") +LINEAR_STORAGE_STATE = os.environ.get("E2E_LINEAR_STORAGE_STATE", "") + # Jaeger query API of the compose stack's OTEL trace destination (the `jaeger` # service in docker-compose.yml maps it to host 16686). Trace-completeness tests # read exported spans back through it. @@ -72,27 +75,31 @@ POLL_TIMEOUT = float(os.environ.get("E2E_POLL_TIMEOUT", "120")) POLL_INTERVAL = float(os.environ.get("E2E_POLL_INTERVAL", "5")) REQUEST_TIMEOUT = float(os.environ.get("E2E_REQUEST_TIMEOUT", "60")) +EXPECT_RUST = os.environ.get("E2E_EXPECT_RUST", "").strip().lower() in ("1", "true", "yes") + LOAD_USERS = int(os.environ.get("E2E_LOAD_USERS", "750")) LOAD_SPAWN_RATE = float(os.environ.get("E2E_LOAD_SPAWN_RATE", "50")) LOAD_DURATION_SECONDS = float(os.environ.get("E2E_LOAD_DURATION_SECONDS", "60")) LOAD_MIN_RPS = float(os.environ.get("E2E_LOAD_MIN_RPS", "355")) LOAD_MAX_FAILURE_RATIO = float(os.environ.get("E2E_LOAD_MAX_FAILURE_RATIO", "0.01")) - -def require_env(*names: str) -> tuple[str, ...]: - """Return the non-empty values for each env name, or hard-fail naming which are missing. - - Live e2e never skips for missing credentials: a missing key is a red run so - ops knows the suite cannot prove the product path. - """ - missing = tuple(name for name in names if not (os.environ.get(name) or "").strip()) - if missing: - joined = ", ".join(missing) - raise AssertionError( - f"missing required env for e2e: {joined}. " - "Add them to tests/e2e/.env locally and to litellm ops for stage/CI." - ) - return tuple((os.environ.get(name) or "").strip() for name in names) +WEEKLY_ANOMALY_OPT_IN_ENV = "E2E_WEEKLY_ANOMALY" +ANOMALY_SESSIONS = int(os.environ.get("E2E_ANOMALY_SESSIONS", "6")) +ANOMALY_TURNS_PER_SESSION = int(os.environ.get("E2E_ANOMALY_TURNS_PER_SESSION", "6")) +ANOMALY_TURN_ATTEMPTS = int(os.environ.get("E2E_ANOMALY_TURN_ATTEMPTS", "3")) +ANOMALY_MAX_ERROR_RATIO = float(os.environ.get("E2E_ANOMALY_MAX_ERROR_RATIO", "0.05")) +ANOMALY_MIN_WARM_CACHE_READ_SHARE = float( + os.environ.get("E2E_ANOMALY_MIN_WARM_CACHE_READ_SHARE", "0.65") +) +ANOMALY_MAX_P95_TURN_SECONDS = float( + os.environ.get("E2E_ANOMALY_MAX_P95_TURN_SECONDS", "30") +) +ANOMALY_MAX_KEY_SPEND_USD = float( + os.environ.get("E2E_ANOMALY_MAX_KEY_SPEND_USD", "0.60") +) +ANOMALY_SPEND_SETTLE_SECONDS = float( + os.environ.get("E2E_ANOMALY_SPEND_SETTLE_SECONDS", "75") +) def datadog_mcp_url(*, toolsets: str = "core") -> str: diff --git a/tests/e2e/e2e_http.py b/tests/e2e/e2e_http.py index 7ec4a439e10..03d7b5d051a 100644 --- a/tests/e2e/e2e_http.py +++ b/tests/e2e/e2e_http.py @@ -147,6 +147,32 @@ class StreamingResponse(BaseModel): return "text/event-stream" in (self.content_type or "") +class BinaryStream(BaseModel): + """Outcome of consuming a binary chunked response (e.g. TTS audio) as a stream. + + Unlike StreamingResponse, which line-splits an SSE text body, this iterates the + raw bytes with iter_content and reports how many non-empty chunks arrived and + the total byte count, so a caller can assert customer-observable streaming + (multiple chunks, real bytes) without decoding the payload.""" + + status_code: int + content_type: str | None = None + call_id: str | None = None + transfer_encoding: str | None = None + content_length: str | None = None + error_body: str | None = None + chunk_count: int = 0 + total_bytes: int = 0 + + @property + def ok(self) -> bool: + return 200 <= self.status_code < 300 + + @property + def chunked(self) -> bool: + return "chunked" in (self.transfer_encoding or "") + + def _hdr(resp: requests.Response, name: str) -> str | None: value = resp.headers.get(name) return value if isinstance(value, str) else None @@ -257,6 +283,26 @@ def get[R: BaseModel]( return _classify(resp, response_type) +def get_external[R: BaseModel]( + url: str, + *, + response_type: type[R], + timeout: float = 30.0, +) -> Result[R]: + """GET an absolute URL outside the proxy (e.g. a public /.well-known document). + Unlike the transport wrappers there is no proxy base url and no proxy auth; the + response still gets the same tagged-union classification as every other call.""" + try: + resp = requests.get( + url, + headers={"Accept": "application/json"}, + timeout=timeout, + ) + except requests.RequestException as exc: + return NetworkError(message=str(exc)) + return _classify(resp, response_type) + + def delete[R: BaseModel]( url: URL, *, @@ -430,16 +476,18 @@ def upload[R: BaseModel]( url: URL, *, headers: BaseModel, - form: FileUploadForm, + form: BaseModel, filename: str, content: bytes, + file_content_type: str = "application/jsonl", params: BaseModel | None = None, response_type: type[R], timeout: float = 60.0, ) -> Result[R]: - """Multipart POST for file uploads (/v1/files). Form fields come from `form`, - the file bytes are sent as the `file` part, and `params` carries any query - routing (e.g. ?model=). requests sets the multipart Content-Type itself.""" + """Multipart POST for file-bearing routes (/v1/files, /v1/audio/transcriptions). + Form fields come from `form`, the file bytes are sent as the `file` part with + `file_content_type`, and `params` carries any query routing (e.g. ?model=). + requests sets the multipart Content-Type itself.""" dumped: dict[str, object] = form.model_dump(by_alias=True, exclude_none=True) data = {key: str(value) for key, value in dumped.items()} try: @@ -448,7 +496,7 @@ def upload[R: BaseModel]( headers=_headers(headers), params=_params(params), data=data, - files={"file": (filename, content, "application/jsonl")}, + files={"file": (filename, content, file_content_type)}, timeout=timeout, ) except requests.RequestException as exc: @@ -456,6 +504,54 @@ def upload[R: BaseModel]( return _classify(resp, response_type) +def stream_binary( + url: URL, + *, + headers: BaseModel, + json: BaseModel, + chunk_size: int = 8192, + timeout: float = 60.0, +) -> BinaryStream: + """POST that consumes a binary chunked response (e.g. TTS audio) as a stream, + counting non-empty chunks and total bytes with iter_content. A non-2xx status + short-circuits with the counts left at zero so the caller can fail loudly.""" + try: + resp = requests.post( + str(url), + headers=_headers(headers), + json=json.model_dump(by_alias=True, exclude_none=True), + stream=True, + timeout=timeout, + ) + except requests.RequestException as exc: + return BinaryStream(status_code=-1, error_body=str(exc)[:300]) + with resp: + content_type = _hdr(resp, "content-type") + call_id = _hdr(resp, "x-litellm-call-id") + transfer_encoding = _hdr(resp, "transfer-encoding") + content_length = _hdr(resp, "content-length") + if not (200 <= resp.status_code < 300): + return BinaryStream( + status_code=resp.status_code, + content_type=content_type, + call_id=call_id, + transfer_encoding=transfer_encoding, + content_length=content_length, + error_body=resp.text[:300], + ) + raw_chunks = cast("Iterator[bytes]", resp.iter_content(chunk_size=chunk_size)) + chunks = tuple(chunk for chunk in raw_chunks if chunk) + return BinaryStream( + status_code=resp.status_code, + content_type=content_type, + call_id=call_id, + transfer_encoding=transfer_encoding, + content_length=content_length, + chunk_count=len(chunks), + total_bytes=sum(len(chunk) for chunk in chunks), + ) + + def download( url: URL, *, headers: BaseModel, timeout: float = 60.0 ) -> StreamingResponse: diff --git a/tests/e2e/guardrails/test_bedrock_guardrail_e2e.py b/tests/e2e/guardrails/test_bedrock_guardrail_e2e.py index 9e41b8808e8..a2408f0021e 100644 --- a/tests/e2e/guardrails/test_bedrock_guardrail_e2e.py +++ b/tests/e2e/guardrails/test_bedrock_guardrail_e2e.py @@ -8,9 +8,11 @@ a 200 means the guardrail never ran. from __future__ import annotations +import os + import pytest -from e2e_config import require_env, unique_marker +from e2e_config import unique_marker from e2e_http import UnknownApiError from guardrails_client import GuardrailsClient from lifecycle import ResourceManager @@ -33,11 +35,8 @@ class TestBedrockGuardrail: def test_bedrock_pre_call_blocks_harmful_prompt( self, client: GuardrailsClient, resources: ResourceManager, scoped_key: str ) -> None: - (identifier, version) = require_env( - "BEDROCK_GUARDRAIL_IDENTIFIER", - "BEDROCK_GUARDRAIL_VERSION", - ) - require_env("AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY", "AWS_REGION") + identifier = os.environ["BEDROCK_GUARDRAIL_IDENTIFIER"] + version = os.environ["BEDROCK_GUARDRAIL_VERSION"] name = f"e2e-bedrock-guard-{unique_marker()}" guardrail_id = client.create_bedrock_guardrail( diff --git a/tests/e2e/guardrails/test_block_code_execution_guardrail_e2e.py b/tests/e2e/guardrails/test_block_code_execution_guardrail_e2e.py index e36fc7c3f9d..de087b190d0 100644 --- a/tests/e2e/guardrails/test_block_code_execution_guardrail_e2e.py +++ b/tests/e2e/guardrails/test_block_code_execution_guardrail_e2e.py @@ -16,7 +16,7 @@ from __future__ import annotations import pytest -from e2e_config import require_env, unique_marker +from e2e_config import unique_marker from e2e_http import unwrap from guardrails_client import BlockCodeExecutionParamsBody, GuardrailsClient from lifecycle import ResourceManager @@ -46,7 +46,6 @@ class TestBlockCodeExecutionGuardrail: def test_blocks_execution_request_but_allows_explanation( self, client: GuardrailsClient, resources: ResourceManager, scoped_key: str ) -> None: - require_env("GEMINI_API_KEY") model = client.create_backend_model(resources, prefix="e2e-blockcode-backend") name = f"e2e-block-code-{unique_marker()}" diff --git a/tests/e2e/guardrails/test_openai_moderation_guardrail_e2e.py b/tests/e2e/guardrails/test_openai_moderation_guardrail_e2e.py index 4e2fcbf8fba..39950259fb5 100644 --- a/tests/e2e/guardrails/test_openai_moderation_guardrail_e2e.py +++ b/tests/e2e/guardrails/test_openai_moderation_guardrail_e2e.py @@ -14,7 +14,7 @@ from __future__ import annotations import pytest -from e2e_config import require_env, unique_marker +from e2e_config import unique_marker from e2e_http import UnknownApiError, unwrap from guardrails_client import GuardrailsClient, OpenAIModerationParamsBody from lifecycle import ResourceManager @@ -34,7 +34,6 @@ class TestOpenAIModerationGuardrail: def test_moderation_blocks_flagged_input( self, client: GuardrailsClient, resources: ResourceManager, scoped_key: str ) -> None: - require_env("OPENAI_API_KEY", "GEMINI_API_KEY") model = client.create_backend_model(resources, prefix="e2e-moderation-backend") name = f"e2e-openai-moderation-{unique_marker()}" diff --git a/tests/e2e/guardrails/test_presidio_guardrail_e2e.py b/tests/e2e/guardrails/test_presidio_guardrail_e2e.py index a911f387382..d103714b1dd 100644 --- a/tests/e2e/guardrails/test_presidio_guardrail_e2e.py +++ b/tests/e2e/guardrails/test_presidio_guardrail_e2e.py @@ -25,11 +25,12 @@ The chat backend is a gemini deployment created for the test. from __future__ import annotations +import os import time import pytest -from e2e_config import POLL_INTERVAL, POLL_TIMEOUT, require_env, unique_marker +from e2e_config import POLL_INTERVAL, POLL_TIMEOUT, unique_marker from e2e_http import NoBody, require_successful_call, unwrap from guardrails_client import GuardrailMode, GuardrailsClient, PresidioParamsBody from lifecycle import ResourceManager @@ -88,9 +89,8 @@ def _poll_logged_prompt(reader: OtelReader, *, call_id: str, genai_span: str) -> def _presidio_params( mode: GuardrailMode, *, apply_to_output: bool = False, logging_only: bool = False ) -> PresidioParamsBody: - analyzer, anonymizer = require_env( - "PRESIDIO_ANALYZER_API_BASE", "PRESIDIO_ANONYMIZER_API_BASE" - ) + analyzer = os.environ["PRESIDIO_ANALYZER_API_BASE"] + anonymizer = os.environ["PRESIDIO_ANONYMIZER_API_BASE"] return PresidioParamsBody( mode=mode, default_on=False, @@ -124,7 +124,6 @@ class TestPresidioGuardrail: def test_pre_call_masks_pii_before_the_model_sees_it( self, client: GuardrailsClient, resources: ResourceManager, scoped_key: str ) -> None: - require_env("GEMINI_API_KEY") model = client.create_backend_model(resources, prefix="e2e-presidio-pre") name = f"e2e-presidio-pre-{unique_marker()}" guardrail_id = client.register(name, _presidio_params("pre_call")) @@ -149,7 +148,6 @@ class TestPresidioGuardrail: def test_post_call_masks_pii_in_model_output( self, client: GuardrailsClient, resources: ResourceManager, scoped_key: str ) -> None: - require_env("GEMINI_API_KEY") model = client.create_backend_model(resources, prefix="e2e-presidio-post") name = f"e2e-presidio-post-{unique_marker()}" guardrail_id = client.register(name, _presidio_params("post_call", apply_to_output=True)) @@ -173,7 +171,6 @@ class TestPresidioGuardrail: def test_logging_only_masks_the_logged_prompt( self, client: GuardrailsClient, resources: ResourceManager, scoped_key: str ) -> None: - require_env("GEMINI_API_KEY") _require_otel_v2_active(client) reader = build_otel_reader() diff --git a/tests/e2e/lifecycle.py b/tests/e2e/lifecycle.py index c0de074f0d2..4ef25509905 100644 --- a/tests/e2e/lifecycle.py +++ b/tests/e2e/lifecycle.py @@ -1,13 +1,11 @@ -"""Lifecycle contract and resource cleanup for stateful e2e tests. +"""Resource cleanup for stateful e2e tests. Shared by every e2e suite under tests/e2e/. The proxy under test is long-lived and never reset between tests, so anything a test creates (keys, customers, teams, orgs, users, guardrails, budgets, ...) persists unless -explicitly deleted. Every check follows an init -> run -> teardown lifecycle; -teardown releases each resource init() created, even when run() raises. - -In pytest terms (see conftest.py): the `resources` fixture's setup is init(), -the test body is run(), and the fixture's teardown is teardown(). +explicitly deleted. The `resources` fixture (see conftest.py) hands each test a +ResourceManager; the test registers a cleanup for every resource it creates, and +the fixture's teardown releases them all even when the test body raises. """ from dataclasses import dataclass, field @@ -17,38 +15,6 @@ from proxy_client import ProxyClient from models import KeyGenerateBody -@runtime_checkable -class E2ECase(Protocol): - """A stateful e2e check run against a long-lived proxy. - - init() acquires resources, run() exercises behaviour and asserts, teardown() - releases everything init() created. teardown() must run even if init() fails - partway or run() raises. - """ - - def init(self) -> None: ... - - def run(self) -> None: ... - - def teardown(self) -> None: ... - - -def run_case(case: E2ECase) -> None: - """Drive a case through its lifecycle: init -> run -> teardown. - - teardown always runs - even when init() fails partway or run() raises (or - skips) - so resources the case already registered on the long-lived proxy are - released. init() is inside the try because cases register cleanups - progressively (e.g. create team, then user, then key), and a failure after - the first creation must still release what came before. - """ - try: - case.init() - case.run() - finally: - case.teardown() - - @runtime_checkable class ResourceClient(Protocol): """Proxy operations the convenience creators use. Resource types without a diff --git a/tests/e2e/llm_translation/endpoints_client.py b/tests/e2e/llm_translation/endpoints_client.py index 32d9922c775..ace621d03b3 100644 --- a/tests/e2e/llm_translation/endpoints_client.py +++ b/tests/e2e/llm_translation/endpoints_client.py @@ -15,7 +15,7 @@ from typing import Literal from pydantic import BaseModel from proxy_client import ProxyClient -from e2e_http import StreamingResponse +from e2e_http import BinaryStream, Result, StreamingResponse from models import CacheControl, ChatMessage, LiteLLMParamsBody, RichMessage, TextBlock __all__ = [ @@ -110,6 +110,16 @@ class ImageRequest(BaseModel): size: str = "1024x1024" +class TranscriptionForm(BaseModel): + model: str + response_format: str = "json" + + +class ModerationRequest(BaseModel): + model: str + input: str + + class ResponsesOutputContent(BaseModel): type: str | None = None text: str | None = None @@ -213,6 +223,27 @@ class ImagesResult(BaseModel): data: list[ImageItem] = [] +class TranscriptionResult(BaseModel): + text: str = "" + + +class ModerationResultItem(BaseModel): + flagged: bool + categories: dict[str, bool] = {} + + @property + def flagged_categories(self) -> tuple[str, ...]: + return tuple(name for name, hit in self.categories.items() if hit) + + +class ModerationResult(BaseModel): + results: list[ModerationResultItem] = [] + + @property + def first(self) -> ModerationResultItem | None: + return self.results[0] if self.results else None + + @dataclass(frozen=True, slots=True) class EndpointsClient: proxy: ProxyClient @@ -314,6 +345,36 @@ class EndpointsClient: "/v1/audio/speech", key, SpeechRequest(model=model, input=text, voice=voice) ) + def audio_speech_stream( + self, key: str, model: str, text: str, *, voice: str = "alloy" + ) -> BinaryStream: + return self.proxy.transport.stream_binary( + "/v1/audio/speech", + headers=self.proxy.transport.bearer(key), + json=SpeechRequest(model=model, input=text, voice=voice), + ) + + def transcribe( + self, key: str, model: str, *, filename: str, content: bytes + ) -> Result[TranscriptionResult]: + return self.proxy.transport.upload( + "/v1/audio/transcriptions", + headers=self.proxy.transport.bearer(key), + form=TranscriptionForm(model=model), + filename=filename, + content=content, + file_content_type="audio/wav", + response_type=TranscriptionResult, + ) + + def moderations(self, key: str, model: str, text: str) -> Result[ModerationResult]: + return self.proxy.transport.post( + "/v1/moderations", + headers=self.proxy.transport.bearer(key), + json=ModerationRequest(model=model, input=text), + response_type=ModerationResult, + ) + def images(self, key: str, model: str, prompt: str) -> StreamingResponse: return self._send( "/v1/images/generations", key, ImageRequest(model=model, prompt=prompt) diff --git a/tests/e2e/llm_translation/realtime/test_nova_sonic_realtime_e2e.py b/tests/e2e/llm_translation/realtime/test_realtime_bedrock_e2e.py similarity index 100% rename from tests/e2e/llm_translation/realtime/test_nova_sonic_realtime_e2e.py rename to tests/e2e/llm_translation/realtime/test_realtime_bedrock_e2e.py diff --git a/tests/e2e/llm_translation/test_audio_speech_e2e.py b/tests/e2e/llm_translation/test_audio_speech_e2e.py index f7a04d94cb3..b95cef8db4d 100644 --- a/tests/e2e/llm_translation/test_audio_speech_e2e.py +++ b/tests/e2e/llm_translation/test_audio_speech_e2e.py @@ -1,8 +1,9 @@ -"""Live e2e: POST /v1/audio/speech returns audio. +"""Live e2e: POST /v1/audio/speech returns audio, non-streamed and streamed. -Registers an OpenAI text-to-speech deployment at runtime and asserts the response -is an audio body (binary, not JSON). Migrated from -litellm-regression-tests/tests/test_inference_endpoints.py. +The non-streamed call asserts an audio (not JSON) body. The streamed call consumes +the response the way a player would and asserts customer-observable streaming: +chunked transfer encoding (a buffered body would carry a content-length) with +non-zero audio bytes. """ from __future__ import annotations @@ -19,6 +20,7 @@ pytestmark = pytest.mark.e2e class TestAudioSpeech: + @pytest.mark.covers("llm.audio_speech.openai.basic.nonstream.works") def test_audio_speech_returns_audio( self, endpoints_client: EndpointsClient, resources: ResourceManager ) -> None: @@ -38,3 +40,39 @@ class TestAudioSpeech: f"/audio/speech content-type is not audio: {result.content_type!r}" ) assert result.body, "/audio/speech returned an empty body" + + @pytest.mark.covers("llm.audio_speech.openai.basic.stream.works") + def test_audio_speech_streams_audio_chunks( + self, endpoints_client: EndpointsClient, resources: ResourceManager + ) -> None: + model = f"e2e-speech-stream-{unique_marker()}" + model_id = endpoints_client.create_model( + model, + LiteLLMParamsBody( + model="openai/gpt-4o-mini-tts", api_key="os.environ/OPENAI_API_KEY" + ), + ) + resources.defer(lambda: endpoints_client.delete_model(model_id)) + key = resources.key() + + result = endpoints_client.audio_speech_stream( + key, + model, + "Streaming speech should arrive in several audio chunks so a client can " + "begin playback well before the whole clip has finished generating.", + ) + assert result.ok, ( + f"/audio/speech stream failed (status {result.status_code}); body={result.error_body}" + ) + assert "audio" in (result.content_type or ""), ( + f"/audio/speech content-type is not audio: {result.content_type!r}" + ) + assert result.chunked, ( + f"/audio/speech did not stream: transfer-encoding={result.transfer_encoding!r}, " + f"content-length={result.content_length!r} (a buffered body is not a stream)" + ) + assert result.content_length is None, ( + f"/audio/speech advertised content-length={result.content_length!r} on a " + f"streamed response (a buffered body is not a stream)" + ) + assert result.total_bytes > 0, "/audio/speech stream returned no audio bytes" diff --git a/tests/e2e/llm_translation/test_audio_transcriptions_e2e.py b/tests/e2e/llm_translation/test_audio_transcriptions_e2e.py new file mode 100644 index 00000000000..af6123dc46a --- /dev/null +++ b/tests/e2e/llm_translation/test_audio_transcriptions_e2e.py @@ -0,0 +1,51 @@ +"""Live e2e: POST /v1/audio/transcriptions turns speech into text. + +Registers an OpenAI speech-to-text deployment at runtime and uploads a spoken +weather question (the realtime suite's 24kHz WAV fixture) as multipart, asserting +the returned transcript is non-empty and mentions the word it was asked about. +""" + +from __future__ import annotations + +from pathlib import Path + +import pytest + +from e2e_config import unique_marker +from e2e_http import unwrap +from endpoints_client import EndpointsClient +from lifecycle import ResourceManager +from models import LiteLLMParamsBody + +pytestmark = pytest.mark.e2e + +WEATHER_WAV = ( + Path(__file__).resolve().parent / "realtime" / "fixtures" / "weather_question_24k.wav" +) + + +class TestAudioTranscriptions: + @pytest.mark.covers("llm.audio_transcriptions.openai.basic.nonstream.works") + def test_audio_transcriptions_returns_text( + self, endpoints_client: EndpointsClient, resources: ResourceManager + ) -> None: + model = f"e2e-transcribe-{unique_marker()}" + model_id = endpoints_client.create_model( + model, + LiteLLMParamsBody( + model="openai/gpt-4o-mini-transcribe", api_key="os.environ/OPENAI_API_KEY" + ), + ) + resources.defer(lambda: endpoints_client.delete_model(model_id)) + key = resources.key() + + result = unwrap( + endpoints_client.transcribe( + key, model, filename=WEATHER_WAV.name, content=WEATHER_WAV.read_bytes() + ) + ) + text = result.text.strip() + assert text, "/audio/transcriptions returned an empty transcript" + assert "weather" in text.lower(), ( + f"transcript of a spoken weather question does not mention weather: {text!r}" + ) diff --git a/tests/e2e/llm_translation/test_chat_completions_regression_e2e.py b/tests/e2e/llm_translation/test_chat_completions_regression_e2e.py index 13992744f42..af0e782e224 100644 --- a/tests/e2e/llm_translation/test_chat_completions_regression_e2e.py +++ b/tests/e2e/llm_translation/test_chat_completions_regression_e2e.py @@ -19,17 +19,176 @@ from __future__ import annotations import os import pytest +from pydantic import BaseModel -from e2e_config import require_env, unique_marker -from e2e_http import unwrap +from e2e_config import unique_marker +from e2e_http import StreamingResponse, unwrap from lifecycle import ResourceManager -from models import ChatBody, ChatMessage, LiteLLMParamsBody +from models import ( + ChatBody, + ChatMessage, + ChatResponse, + ChatTool, + ChatToolFunction, + ImageContentPart, + ImageUrl, + LiteLLMParamsBody, + TextContentPart, + ThinkingParam, +) from passthrough_client import PassthroughClient pytestmark = pytest.mark.e2e COHERE_BACKEND = "cohere/command-r-08-2024" GEMINI_BACKEND = "gemini/gemini-2.5-flash" +OPENAI_BACKEND = "openai/gpt-5.6" +BEDROCK_CONVERSE_BACKEND = "bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0" + + +class _StreamToolCallFunction(BaseModel): + name: str | None = None + arguments: str | None = None + + +class _StreamToolCall(BaseModel): + function: _StreamToolCallFunction = _StreamToolCallFunction() + + +class _StreamDelta(BaseModel): + content: str | None = None + tool_calls: list[_StreamToolCall] | None = None + + +class _StreamChoice(BaseModel): + delta: _StreamDelta = _StreamDelta() + + +class _StreamChunk(BaseModel): + choices: list[_StreamChoice] = [] + + +def _streamed_tool_call(events: list[str]) -> tuple[str, str]: + """Reassemble the tool call streamed across chunks: the name arrives once and the + arguments arrive as fragments, so concatenating both and parsing the arguments as + JSON catches a stream that never completes the call or splits its argument JSON.""" + chunks = [_StreamChunk.model_validate_json(event) for event in events] + calls = [call for chunk in chunks for choice in chunk.choices for call in (choice.delta.tool_calls or [])] + name = "".join(call.function.name or "" for call in calls) + arguments = "".join(call.function.arguments or "" for call in calls) + return name, arguments + + +CAT_IMAGE_URL = "https://upload.wikimedia.org/wikipedia/commons/3/3a/Cat03.jpg" +OPENAI_VISION_BACKEND = "openai/gpt-4o" + +# OpenAI caches a shared prompt prefix once it exceeds ~1024 tokens; this is well +# past that, so a repeat call reports cached prompt tokens. +CACHE_PREFIX = ( + "You are a meticulous assistant. Follow these standing instructions exactly. " + * 300 +) + + +def _vision_messages() -> list[ChatMessage]: + return [ + ChatMessage( + role="user", + content=[ + TextContentPart(text="What animal is in this image? Answer in one word."), + ImageContentPart(image_url=ImageUrl(url=CAT_IMAGE_URL)), + ], + ) + ] + + +def _assert_describes_cat(response: ChatResponse) -> None: + assert response.choices, f"vision returned no choices: {response}" + message = response.choices[0].message + content = (message.content if message else None) or "" + assert "cat" in content.lower() or "feline" in content.lower(), ( + f"vision response did not describe the image: {content[:200]}" + ) + + +def _streamed_text(events: list[str]) -> str: + """Concatenate the delta content across streamed chunks. Parsing every event as + JSON also fails loudly on a truncated or garbled chunk (the vertex/gemini image + streaming regression class), so an incomplete stream cannot pass as content.""" + chunks = [_StreamChunk.model_validate_json(event) for event in events] + return "".join(choice.delta.content or "" for chunk in chunks for choice in chunk.choices) + + +def _assert_streamed_completion(result: StreamingResponse) -> None: + """A streamed /chat/completions must deliver real content, not a clean-but-empty + stream (the #28991 class on the streaming path).""" + assert result.ok and result.is_streaming, f"stream was not established: {result}" + assert result.stream_error is None, f"stream carried an error event: {result.stream_error}" + assert len(result.stream_events) > 1, f"stream did not deliver multiple data events: {result}" + assert _streamed_text(result.stream_events).strip(), ( + f"stream completed with no content deltas: {result.stream_events[:3]}" + ) + + +def _bedrock_params() -> LiteLLMParamsBody: + return LiteLLMParamsBody( + model=BEDROCK_CONVERSE_BACKEND, + aws_access_key_id="os.environ/AWS_ACCESS_KEY_ID", + aws_secret_access_key="os.environ/AWS_SECRET_ACCESS_KEY", + aws_region_name="os.environ/AWS_REGION", + ) + + +class _WeatherArgs(BaseModel): + location: str + + +_WEATHER_TOOL = ChatTool( + function=ChatToolFunction( + name="get_weather", + description="Get the current weather for a location", + parameters={ + "type": "object", + "properties": {"location": {"type": "string"}}, + "required": ["location"], + }, + ) +) + + +def _assert_weather_tool_call(response: ChatResponse) -> None: + """The model, forced to call the tool, must return a get_weather call whose + arguments parse as JSON and carry a location. A regression that drops tool_calls + or emits malformed argument JSON fails here rather than passing on a 200.""" + assert response.choices, f"chat returned no choices: {response}" + message = response.choices[0].message + calls = message.tool_calls if message else None + assert calls, f"model returned no tool call for a tool-forced prompt: {response}" + weather = next((call for call in calls if call.function.name == "get_weather"), None) + assert weather is not None, f"expected a get_weather call, got {[c.function.name for c in calls]}" + assert weather.function.arguments, f"get_weather call carried no arguments: {weather}" + args = _WeatherArgs.model_validate_json(weather.function.arguments) + assert args.location.strip(), f"get_weather arguments missing location: {weather.function.arguments}" + + +class _Person(BaseModel): + name: str + age: int + + +_PERSON_SCHEMA: dict[str, object] = { + "type": "json_schema", + "json_schema": { + "name": "person", + "strict": True, + "schema": { + "type": "object", + "properties": {"name": {"type": "string"}, "age": {"type": "integer"}}, + "required": ["name", "age"], + "additionalProperties": False, + }, + }, +} CHAT_MODELS: tuple[tuple[str, str], ...] = ( ("gpt-5.5", "openai"), @@ -91,7 +250,7 @@ class TestCohereChat: def test_cohere_chat_returns_content( self, client: PassthroughClient, resources: ResourceManager ) -> None: - (cohere_key,) = require_env("COHERE_API_KEY") + cohere_key = os.environ["COHERE_API_KEY"] model = f"e2e-cohere-chat-{unique_marker()}" model_id = client.proxy.create_model( model, @@ -184,7 +343,7 @@ class TestHostedVllmChat: def test_hosted_vllm_chat_returns_content( self, client: PassthroughClient, resources: ResourceManager ) -> None: - (api_base,) = require_env("HOSTED_VLLM_API_BASE") + api_base = os.environ["HOSTED_VLLM_API_BASE"] api_key = (os.environ.get("HOSTED_VLLM_API_KEY") or "").strip() or None backend = ( os.environ.get("HOSTED_VLLM_MODEL") or "meta-llama/Llama-3.2-3B-Instruct" @@ -219,3 +378,382 @@ class TestHostedVllmChat: assert response.choices, f"hosted_vllm chat returned no choices: {response}" content = response.choices[0].message.content if response.choices[0].message else None assert content and content.strip(), f"hosted_vllm empty content: {response}" + + +class TestOpenAIChatCompletions: + """OpenAI /chat/completions, the SDK path the customer runs against the proxy. + + The streamed call must deliver real content deltas (a clean-but-empty stream is + the regression), and a non-streamed call must be costed so per-request spend and + the response-cost header stay accurate. + """ + + @pytest.mark.covers( + "llm.chat_completions.openai.basic.stream.works", + exercised_on=["chat_completions"], + ) + def test_openai_chat_streams_real_content( + self, client: PassthroughClient, resources: ResourceManager + ) -> None: + model = f"e2e-openai-chat-{unique_marker()}" + model_id = client.proxy.create_model( + model, LiteLLMParamsBody(model=OPENAI_BACKEND, api_key="os.environ/OPENAI_API_KEY") + ) + resources.defer(lambda: client.proxy.delete_model(model_id)) + key = resources.key() + + result = client.proxy.chat_stream( + key, + ChatBody( + model=model, + messages=[ + ChatMessage(role="user", content=f"Count from 1 to 5, one number per line. {unique_marker()}") + ], + max_tokens=64, + stream=True, + ), + ) + _assert_streamed_completion(result) + + @pytest.mark.covers( + "llm.chat_completions.openai.basic.nonstream.cost_logged", + exercised_on=["chat_completions"], + ) + def test_openai_chat_logs_cost( + self, client: PassthroughClient, resources: ResourceManager + ) -> None: + model = f"e2e-openai-cost-{unique_marker()}" + model_id = client.proxy.create_model( + model, LiteLLMParamsBody(model=OPENAI_BACKEND, api_key="os.environ/OPENAI_API_KEY") + ) + resources.defer(lambda: client.proxy.delete_model(model_id)) + key = resources.key() + + response = unwrap( + client.proxy.chat( + key, + ChatBody( + model=model, + messages=[ChatMessage(role="user", content=f"Reply with the single word pong. {unique_marker()}")], + max_tokens=16, + ), + ) + ) + assert response.choices, f"openai chat returned no choices: {response}" + + rows = client.proxy.poll_logs_for_key( + key, min_rows=1, predicate=lambda rs: any((r.spend or 0) > 0 for r in rs) + ) + priced = [r for r in rows if (r.spend or 0) > 0] + assert priced, f"openai chat was not costed on key ...{key[-6:]}: {rows}" + assert priced[0].status == "success", f"openai chat spend status={priced[0].status!r}" + + @pytest.mark.covers( + "llm.chat_completions.openai.tool_use.nonstream.works", + exercised_on=["chat_completions"], + ) + def test_openai_chat_returns_tool_call( + self, client: PassthroughClient, resources: ResourceManager + ) -> None: + model = f"e2e-openai-tool-{unique_marker()}" + model_id = client.proxy.create_model( + model, LiteLLMParamsBody(model=OPENAI_BACKEND, api_key="os.environ/OPENAI_API_KEY") + ) + resources.defer(lambda: client.proxy.delete_model(model_id)) + key = resources.key() + + response = unwrap( + client.proxy.chat( + key, + ChatBody( + model=model, + messages=[ + ChatMessage(role="user", content="What is the weather in San Francisco? Use the get_weather tool.") + ], + tools=[_WEATHER_TOOL], + tool_choice="required", + max_tokens=128, + ), + ) + ) + _assert_weather_tool_call(response) + + @pytest.mark.covers( + "llm.chat_completions.openai.structured_output.nonstream.works", + exercised_on=["chat_completions"], + ) + def test_openai_chat_structured_output_conforms_to_schema( + self, client: PassthroughClient, resources: ResourceManager + ) -> None: + model = f"e2e-openai-schema-{unique_marker()}" + model_id = client.proxy.create_model( + model, LiteLLMParamsBody(model=OPENAI_BACKEND, api_key="os.environ/OPENAI_API_KEY") + ) + resources.defer(lambda: client.proxy.delete_model(model_id)) + key = resources.key() + + response = unwrap( + client.proxy.chat( + key, + ChatBody( + model=model, + messages=[ChatMessage(role="user", content="Extract the person. John Doe is 42 years old.")], + response_format=_PERSON_SCHEMA, + max_tokens=128, + ), + ) + ) + assert response.choices, f"structured output returned no choices: {response}" + content = response.choices[0].message.content if response.choices[0].message else None + assert content, f"structured output returned empty content: {response}" + person = _Person.model_validate_json(content) + assert person.name.strip() and person.age == 42, ( + f"schema-constrained extraction was wrong: {person}" + ) + + @pytest.mark.covers( + "llm.chat_completions.openai.thinking.nonstream.works", + exercised_on=["chat_completions"], + ) + def test_openai_chat_reasoning_reports_reasoning_tokens( + self, client: PassthroughClient, resources: ResourceManager + ) -> None: + model = f"e2e-openai-reasoning-{unique_marker()}" + model_id = client.proxy.create_model( + model, LiteLLMParamsBody(model=OPENAI_BACKEND, api_key="os.environ/OPENAI_API_KEY") + ) + resources.defer(lambda: client.proxy.delete_model(model_id)) + key = resources.key() + + response = unwrap( + client.proxy.chat( + key, + ChatBody( + model=model, + messages=[ + ChatMessage( + role="user", + content="A train travels 60 miles in 1.5 hours. What is its average speed in mph?", + ) + ], + reasoning_effort="low", + max_tokens=2048, + ), + ) + ) + assert response.choices, f"reasoning call returned no choices: {response}" + message = response.choices[0].message + assert message and message.content and message.content.strip(), f"reasoning call had no answer: {response}" + details = response.usage.completion_tokens_details if response.usage else None + assert details and details.reasoning_tokens and details.reasoning_tokens > 0, ( + f"a reasoning model must report reasoning tokens, got usage={response.usage}" + ) + + @pytest.mark.covers( + "llm.chat_completions.openai.vision.nonstream.works", + exercised_on=["chat_completions"], + ) + def test_openai_chat_vision_describes_image( + self, client: PassthroughClient, resources: ResourceManager + ) -> None: + model = f"e2e-openai-vision-{unique_marker()}" + model_id = client.proxy.create_model( + model, LiteLLMParamsBody(model=OPENAI_VISION_BACKEND, api_key="os.environ/OPENAI_API_KEY") + ) + resources.defer(lambda: client.proxy.delete_model(model_id)) + key = resources.key() + + response = unwrap(client.proxy.chat(key, ChatBody(model=model, messages=_vision_messages(), max_tokens=32))) + _assert_describes_cat(response) + + @pytest.mark.covers( + "llm.chat_completions.openai.prompt_cache_5m.nonstream.works", + exercised_on=["chat_completions"], + ) + def test_openai_chat_prompt_cache_hits_on_repeat( + self, client: PassthroughClient, resources: ResourceManager + ) -> None: + model = f"e2e-openai-cache-{unique_marker()}" + model_id = client.proxy.create_model( + model, LiteLLMParamsBody(model=OPENAI_BACKEND, api_key="os.environ/OPENAI_API_KEY") + ) + resources.defer(lambda: client.proxy.delete_model(model_id)) + key = resources.key() + + body = ChatBody( + model=model, + messages=[ + ChatMessage(role="system", content=CACHE_PREFIX), + ChatMessage(role="user", content="Reply with the single word pong."), + ], + max_tokens=16, + ) + unwrap(client.proxy.chat(key, body)) + second = unwrap(client.proxy.chat(key, body)) + + details = second.usage.prompt_tokens_details if second.usage else None + assert details and details.cached_tokens and details.cached_tokens > 0, ( + f"a repeated large-prefix prompt must report cached prompt tokens, got usage={second.usage}" + ) + + @pytest.mark.covers( + "llm.chat_completions.openai.tool_use.stream.works", + exercised_on=["chat_completions"], + ) + def test_openai_chat_streams_tool_call( + self, client: PassthroughClient, resources: ResourceManager + ) -> None: + model = f"e2e-openai-tool-stream-{unique_marker()}" + model_id = client.proxy.create_model( + model, LiteLLMParamsBody(model=OPENAI_BACKEND, api_key="os.environ/OPENAI_API_KEY") + ) + resources.defer(lambda: client.proxy.delete_model(model_id)) + key = resources.key() + + result = client.proxy.chat_stream( + key, + ChatBody( + model=model, + messages=[ + ChatMessage(role="user", content="What is the weather in San Francisco? Use the get_weather tool.") + ], + tools=[_WEATHER_TOOL], + tool_choice="required", + max_tokens=128, + stream=True, + ), + ) + assert result.ok and result.is_streaming, f"tool stream was not established: {result}" + assert result.stream_error is None, f"tool stream carried an error event: {result.stream_error}" + name, arguments = _streamed_tool_call(result.stream_events) + assert name == "get_weather", f"streamed tool call named {name!r}: {result.stream_events[:5]}" + args = _WeatherArgs.model_validate_json(arguments) + assert args.location.strip(), f"streamed tool call arguments missing location: {arguments!r}" + + +class TestBedrockConverseChatCompletions: + """Bedrock Converse via /chat/completions, the customer's AWS stack. A non-OpenAI + provider must return real content on both the non-streamed and streamed paths. + """ + + def _register(self, client: PassthroughClient, resources: ResourceManager, prefix: str) -> str: + model = f"{prefix}-{unique_marker()}" + model_id = client.proxy.create_model(model, _bedrock_params()) + resources.defer(lambda: client.proxy.delete_model(model_id)) + return model + + @pytest.mark.covers( + "llm.chat_completions.bedrock_converse.basic.nonstream.works", + exercised_on=["chat_completions"], + ) + def test_bedrock_converse_chat_returns_content( + self, client: PassthroughClient, resources: ResourceManager + ) -> None: + model = self._register(client, resources, "e2e-bedrock-chat") + key = resources.key() + + response = unwrap( + client.proxy.chat( + key, + ChatBody( + model=model, + messages=[ChatMessage(role="user", content=f"Reply with the single word pong. {unique_marker()}")], + max_tokens=32, + ), + ) + ) + assert response.choices, f"bedrock converse chat returned no choices: {response}" + content = response.choices[0].message.content if response.choices[0].message else None + assert content and content.strip(), f"bedrock converse returned empty content: {response}" + + @pytest.mark.covers( + "llm.chat_completions.bedrock_converse.basic.stream.works", + exercised_on=["chat_completions"], + ) + def test_bedrock_converse_chat_streams_real_content( + self, client: PassthroughClient, resources: ResourceManager + ) -> None: + model = self._register(client, resources, "e2e-bedrock-stream") + key = resources.key() + + result = client.proxy.chat_stream( + key, + ChatBody( + model=model, + messages=[ + ChatMessage(role="user", content=f"Count from 1 to 5, one number per line. {unique_marker()}") + ], + max_tokens=64, + stream=True, + ), + ) + _assert_streamed_completion(result) + + @pytest.mark.covers( + "llm.chat_completions.bedrock_converse.tool_use.nonstream.works", + exercised_on=["chat_completions"], + ) + def test_bedrock_converse_chat_returns_tool_call( + self, client: PassthroughClient, resources: ResourceManager + ) -> None: + model = self._register(client, resources, "e2e-bedrock-tool") + key = resources.key() + + response = unwrap( + client.proxy.chat( + key, + ChatBody( + model=model, + messages=[ + ChatMessage(role="user", content="What is the weather in San Francisco? Use the get_weather tool.") + ], + tools=[_WEATHER_TOOL], + tool_choice="required", + max_tokens=128, + ), + ) + ) + _assert_weather_tool_call(response) + + @pytest.mark.covers( + "llm.chat_completions.bedrock_converse.thinking.nonstream.works", + exercised_on=["chat_completions"], + ) + def test_bedrock_converse_chat_returns_reasoning( + self, client: PassthroughClient, resources: ResourceManager + ) -> None: + model = self._register(client, resources, "e2e-bedrock-thinking") + key = resources.key() + + response = unwrap( + client.proxy.chat( + key, + ChatBody( + model=model, + messages=[ChatMessage(role="user", content="What is 17 times 23? Think it through step by step.")], + thinking=ThinkingParam(type="enabled", budget_tokens=1024), + max_tokens=2048, + ), + ) + ) + assert response.choices, f"bedrock thinking returned no choices: {response}" + message = response.choices[0].message + assert message and message.content and message.content.strip(), ( + f"bedrock thinking returned no answer content: {response}" + ) + assert message.reasoning_content and message.reasoning_content.strip(), ( + "thinking was enabled but no reasoning_content came back on the Bedrock Converse path" + ) + + @pytest.mark.covers( + "llm.chat_completions.bedrock_converse.vision.nonstream.works", + exercised_on=["chat_completions"], + ) + def test_bedrock_converse_chat_vision_describes_image( + self, client: PassthroughClient, resources: ResourceManager + ) -> None: + model = self._register(client, resources, "e2e-bedrock-vision") + key = resources.key() + + response = unwrap(client.proxy.chat(key, ChatBody(model=model, messages=_vision_messages(), max_tokens=32))) + _assert_describes_cat(response) diff --git a/tests/e2e/llm_translation/test_embeddings_endpoint_e2e.py b/tests/e2e/llm_translation/test_embeddings_endpoint_e2e.py index 56f2de8bd4f..157caedd561 100644 --- a/tests/e2e/llm_translation/test_embeddings_endpoint_e2e.py +++ b/tests/e2e/llm_translation/test_embeddings_endpoint_e2e.py @@ -1,9 +1,9 @@ -"""Live e2e: POST /embeddings returns a real vector. +"""Live e2e: POST /embeddings returns a real vector across OpenAI, Bedrock, Vertex. -Registers an OpenAI embedding deployment at runtime and asserts a non-empty, -non-zero vector came back. Migrated from -litellm-regression-tests/tests/test_inference_endpoints.py; the LIT-3167 guard in -tests/e2e/embeddings/ covers the Gemini embedding path. +Each test registers the deployment it needs at runtime (deleted on teardown) and +asserts a non-empty, non-zero vector came back. The LIT-3167 guard in +tests/e2e/embeddings/ covers the Gemini embedding path; embeddings cost tracking is +covered by tests/e2e/quota_management/spend_tracking/. """ from __future__ import annotations @@ -20,6 +20,7 @@ pytestmark = pytest.mark.e2e class TestEmbeddingsEndpoint: + @pytest.mark.covers("llm.embeddings.openai.basic.nonstream.works") def test_embeddings_returns_vector( self, endpoints_client: EndpointsClient, resources: ResourceManager ) -> None: @@ -40,3 +41,49 @@ class TestEmbeddingsEndpoint: assert any(component != 0.0 for component in parsed.first_vector), ( f"embedding vector is all zeros: {result.body[:300]}" ) + + @pytest.mark.covers("llm.embeddings.bedrock.basic.nonstream.works") + def test_bedrock_embeddings_returns_vector( + self, endpoints_client: EndpointsClient, resources: ResourceManager + ) -> None: + model = f"e2e-embeddings-bedrock-{unique_marker()}" + model_id = endpoints_client.create_model( + model, + LiteLLMParamsBody( + model="bedrock/amazon.titan-embed-text-v2:0", aws_region_name="us-west-2" + ), + ) + resources.defer(lambda: endpoints_client.delete_model(model_id)) + key = resources.key() + + result = endpoints_client.embeddings(key, model, "Say this is a test!") + require_successful_call(result) + parsed = EmbeddingsResult.model_validate_json(result.body) + assert parsed.first_vector, f"/embeddings returned no vector: {result.body[:300]}" + assert any(component != 0.0 for component in parsed.first_vector), ( + f"embedding vector is all zeros: {result.body[:300]}" + ) + + @pytest.mark.covers("llm.embeddings.vertex.basic.nonstream.works") + def test_vertex_embeddings_returns_vector( + self, endpoints_client: EndpointsClient, resources: ResourceManager + ) -> None: + model = f"e2e-embeddings-vertex-{unique_marker()}" + model_id = endpoints_client.create_model( + model, + LiteLLMParamsBody( + model="vertex_ai/gemini-embedding-2", + vertex_project="os.environ/VERTEXAI_PROJECT", + vertex_location="us-central1", + ), + ) + resources.defer(lambda: endpoints_client.delete_model(model_id)) + key = resources.key() + + result = endpoints_client.embeddings(key, model, "Say this is a test!") + require_successful_call(result) + parsed = EmbeddingsResult.model_validate_json(result.body) + assert parsed.first_vector, f"/embeddings returned no vector: {result.body[:300]}" + assert any(component != 0.0 for component in parsed.first_vector), ( + f"embedding vector is all zeros: {result.body[:300]}" + ) diff --git a/tests/e2e/llm_translation/test_image_generation_e2e.py b/tests/e2e/llm_translation/test_image_generation_e2e.py index 4d2211f3be4..45861d1e93a 100644 --- a/tests/e2e/llm_translation/test_image_generation_e2e.py +++ b/tests/e2e/llm_translation/test_image_generation_e2e.py @@ -18,7 +18,17 @@ from models import LiteLLMParamsBody pytestmark = pytest.mark.e2e +def _assert_image_returned(body: str) -> None: + parsed = ImagesResult.model_validate_json(body) + assert parsed.data, f"/images/generations returned no data: {body[:300]}" + first = parsed.data[0] + assert first.b64_json or first.url, ( + f"generated image has neither b64_json nor url: {body[:300]}" + ) + + class TestImageGeneration: + @pytest.mark.covers("llm.images_generations.openai.basic.nonstream.works") def test_image_generation_returns_image( self, endpoints_client: EndpointsClient, resources: ResourceManager ) -> None: @@ -34,9 +44,25 @@ class TestImageGeneration: result = endpoints_client.images(key, model, "Draw a cute cat") require_successful_call(result) - parsed = ImagesResult.model_validate_json(result.body) - assert parsed.data, f"/images/generations returned no data: {result.body[:300]}" - first = parsed.data[0] - assert first.b64_json or first.url, ( - f"generated image has neither b64_json nor url: {result.body[:300]}" + _assert_image_returned(result.body) + + @pytest.mark.covers("llm.images_generations.bedrock.basic.nonstream.works", exercised_on=["images_generations"]) + def test_bedrock_image_generation_returns_image( + self, endpoints_client: EndpointsClient, resources: ResourceManager + ) -> None: + model = f"e2e-bedrock-image-{unique_marker()}" + model_id = endpoints_client.create_model( + model, + LiteLLMParamsBody( + model="bedrock/amazon.titan-image-generator-v2:0", + aws_access_key_id="os.environ/AWS_ACCESS_KEY_ID", + aws_secret_access_key="os.environ/AWS_SECRET_ACCESS_KEY", + aws_region_name="os.environ/AWS_REGION", + ), ) + resources.defer(lambda: endpoints_client.delete_model(model_id)) + key = resources.key() + + result = endpoints_client.images(key, model, "Draw a cute cat") + require_successful_call(result) + _assert_image_returned(result.body) diff --git a/tests/e2e/llm_translation/test_messages_azure_foundry_e2e.py b/tests/e2e/llm_translation/test_messages_azure_foundry_e2e.py index 3f2907f202e..d8d44820e80 100644 --- a/tests/e2e/llm_translation/test_messages_azure_foundry_e2e.py +++ b/tests/e2e/llm_translation/test_messages_azure_foundry_e2e.py @@ -12,7 +12,7 @@ from __future__ import annotations import pytest -from e2e_config import unique_marker +from e2e_config import EXPECT_RUST, unique_marker from e2e_http import StreamingResponse, require_successful_call, unwrap from endpoints_client import EndpointsClient from lifecycle import ResourceManager @@ -50,6 +50,13 @@ def _assert_streamed_ok(result: StreamingResponse) -> None: assert any("message_stop" in event for event in result.stream_events), ( "stream never reached message_stop" ) + if EXPECT_RUST: + assert result.headers.get("x-litellm-rust") == "true", ( + "E2E_EXPECT_RUST is set, so this gateway must serve /v1/messages through the " + "Rust path, but the response carried no x-litellm-rust marker. The request " + "still succeeded, which is exactly the failure mode: a gateway whose native " + f"extension is unavailable falls back to Python silently. headers={result.headers}" + ) class TestAzureFoundryMessages: diff --git a/tests/e2e/llm_translation/test_messages_e2e.py b/tests/e2e/llm_translation/test_messages_e2e.py index a14bf8e82b3..ef6ba5b95d3 100644 --- a/tests/e2e/llm_translation/test_messages_e2e.py +++ b/tests/e2e/llm_translation/test_messages_e2e.py @@ -20,11 +20,14 @@ from models import ( ChatMessage, JsonSchemaProperty, LiteLLMParamsBody, + SpendLogRow, ToolInputSchema, ) pytestmark = pytest.mark.e2e +ANTHROPIC_BACKEND = "anthropic/claude-haiku-4-5" + WEATHER_TOOL = AnthropicCustomTool( name="get_weather", description="Get the current weather for a city.", @@ -35,6 +38,11 @@ WEATHER_TOOL = AnthropicCustomTool( ) +def _approx_equal(actual: float, expected: float) -> bool: + """Within 1% or 1e-9 absolute - spend math, not exact float identity.""" + return abs(actual - expected) <= max(1e-9, abs(expected) * 1e-2) + + class TestAnthropicMessages: def _register( self, endpoints_client: EndpointsClient, resources: ResourceManager @@ -43,7 +51,7 @@ class TestAnthropicMessages: model_id = endpoints_client.create_model( model, LiteLLMParamsBody( - model="anthropic/claude-haiku-4-5", api_key="os.environ/ANTHROPIC_API_KEY" + model=ANTHROPIC_BACKEND, api_key="os.environ/ANTHROPIC_API_KEY" ), ) resources.defer(lambda: endpoints_client.delete_model(model_id)) @@ -61,6 +69,57 @@ class TestAnthropicMessages: assert parsed.role == "assistant", f"unexpected role: {result.body[:300]}" assert parsed.text.strip(), f"/v1/messages returned no text: {result.body[:300]}" + @pytest.mark.covers("llm.messages.anthropic.basic.nonstream.cost_logged") + def test_messages_logs_cost_matching_the_response_header( + self, endpoints_client: EndpointsClient, resources: ResourceManager + ) -> None: + model = f"e2e-messages-cost-{unique_marker()}" + model_id = endpoints_client.create_model( + model, + LiteLLMParamsBody( + model=ANTHROPIC_BACKEND, api_key="os.environ/ANTHROPIC_API_KEY" + ), + ) + resources.defer(lambda: endpoints_client.delete_model(model_id)) + key = resources.key() + + result = endpoints_client.messages(key, model, f"reply with one word {unique_marker()}") + require_successful_call(result) + parsed = MessagesResult.model_validate_json(result.body) + assert parsed.role == "assistant" and parsed.text.strip(), ( + f"/v1/messages returned no assistant text: {result.body[:300]}" + ) + + # The customer reads per-request cost off the response header (LIT-4076), so + # it must be present and positive on /v1/messages, not only /chat/completions. + header_cost = result.response_cost + assert header_cost is not None and header_cost > 0, ( + "x-litellm-response-cost header missing or non-positive on /v1/messages; " + f"headers={result.headers}" + ) + + # Correlate the spend row by the unique scoped key, not the Anthropic response + # id: on /v1/messages the spend-log request_id is the proxy's own call id, which + # need not equal the message body id, so an id-based poll can miss a correctly + # logged row and time out. The key is fresh per test, so its only priced row is + # this call. + def _priced(rows: list[SpendLogRow]) -> bool: + return any(r.spend is not None and r.spend > 0 for r in rows) + + rows = endpoints_client.proxy.poll_logs_for_key(key, predicate=_priced) + priced = [r for r in rows if r.spend is not None and r.spend > 0] + assert priced, ( + f"no priced /spend/logs row landed for key {key} within the poll window; got {rows}" + ) + row = priced[0] + assert (row.prompt_tokens or 0) > 0 and (row.completion_tokens or 0) > 0, ( + f"messages spend row missing token counts, so the cost is not real usage: {row}" + ) + assert row.spend is not None and _approx_equal(row.spend, header_cost), ( + f"logged spend {row.spend} disagrees with the x-litellm-response-cost header {header_cost}; " + "the customer bills against the header, so the two must match" + ) + @pytest.mark.covers("llm.messages.anthropic.basic.stream.works") def test_messages_streams_completion( self, endpoints_client: EndpointsClient, resources: ResourceManager diff --git a/tests/e2e/llm_translation/test_moderations_e2e.py b/tests/e2e/llm_translation/test_moderations_e2e.py new file mode 100644 index 00000000000..69cf4414a48 --- /dev/null +++ b/tests/e2e/llm_translation/test_moderations_e2e.py @@ -0,0 +1,65 @@ +"""Live e2e: POST /v1/moderations classifies content against the provider policy. + +Registers OpenAI's omni moderation model at runtime and asserts the product +promise on both sides of the decision: clearly violent text comes back flagged +with at least one policy category tripped, and benign text comes back not flagged. +""" + +from __future__ import annotations + +import pytest + +from e2e_config import unique_marker +from e2e_http import unwrap +from endpoints_client import EndpointsClient +from lifecycle import ResourceManager +from models import LiteLLMParamsBody + +pytestmark = pytest.mark.e2e + +VIOLENT_TEXT = "I am going to find you and kill you, and I will hurt everyone you love." +BENIGN_TEXT = "I enjoyed the sunny afternoon and a relaxing walk in the park today." + + +def _register_moderation_model( + endpoints_client: EndpointsClient, resources: ResourceManager +) -> str: + model = f"e2e-moderation-{unique_marker()}" + model_id = endpoints_client.create_model( + model, + LiteLLMParamsBody( + model="openai/omni-moderation-latest", api_key="os.environ/OPENAI_API_KEY" + ), + ) + resources.defer(lambda: endpoints_client.delete_model(model_id)) + return model + + +class TestModerations: + @pytest.mark.covers("llm.moderations.openai.basic.nonstream.works") + def test_moderations_flags_violent_content( + self, endpoints_client: EndpointsClient, resources: ResourceManager + ) -> None: + model = _register_moderation_model(endpoints_client, resources) + key = resources.key() + + result = unwrap(endpoints_client.moderations(key, model, VIOLENT_TEXT)) + item = result.first + assert item is not None, f"/moderations returned no results: {result}" + assert item.flagged, f"violent text was not flagged: {item}" + assert item.flagged_categories, ( + f"flagged result reported no true category: {item}" + ) + + def test_moderations_passes_benign_content( + self, endpoints_client: EndpointsClient, resources: ResourceManager + ) -> None: + model = _register_moderation_model(endpoints_client, resources) + key = resources.key() + + result = unwrap(endpoints_client.moderations(key, model, BENIGN_TEXT)) + item = result.first + assert item is not None, f"/moderations returned no results: {result}" + assert not item.flagged, ( + f"benign text was flagged as {item.flagged_categories}: {item}" + ) diff --git a/tests/e2e/llm_translation/test_passthrough_e2e.py b/tests/e2e/llm_translation/test_passthrough_e2e.py index c8806faf3ea..ed5c657d23e 100644 --- a/tests/e2e/llm_translation/test_passthrough_e2e.py +++ b/tests/e2e/llm_translation/test_passthrough_e2e.py @@ -15,7 +15,8 @@ import pytest from e2e_config import unique_marker from e2e_http import StreamingResponse, require_successful_call -from models import SpendLogRow +from lifecycle import ResourceManager +from models import KeyGenerateBody, SpendLogRow from passthrough_client import ( AnthropicTool, GeminiFunctionDeclaration, @@ -157,3 +158,25 @@ def test_anthropic_passthrough_tool_call_logs_cost( row = _fetch_cost_breakdown(client, result) assert row.custom_llm_provider == "anthropic" + + +class TestPassthroughModelAllowlist: + """A passthrough route must honor the calling key's model allow-list. + + The customer fronts native provider calls through the proxy with custom auth, + so a key scoped to one model must not reach a different model just because the + request goes through the passthrough route rather than /chat/completions. + """ + + @pytest.mark.covers("other.auth.passthrough.model_allowlist_enforced") + def test_passthrough_denies_model_outside_key_allowlist( + self, client: PassthroughClient, resources: ResourceManager + ) -> None: + key = client.proxy.generate_key(KeyGenerateBody(models=["gemini-2.5-flash"])) + resources.defer(lambda: client.proxy.delete_key(key)) + + result = client.anthropic_message(key, "claude-haiku-4-5", f"say hi {unique_marker()}") + assert result.status_code == 403, ( + "a key restricted to gemini-2.5-flash must be denied a claude passthrough call, " + f"got {result.status_code}: {result.body[:300]}" + ) diff --git a/tests/e2e/llm_translation/test_passthrough_headers_e2e.py b/tests/e2e/llm_translation/test_passthrough_headers_e2e.py index d9fe37c79ed..045988334d5 100644 --- a/tests/e2e/llm_translation/test_passthrough_headers_e2e.py +++ b/tests/e2e/llm_translation/test_passthrough_headers_e2e.py @@ -1,29 +1,32 @@ """Live e2e: custom pass-through endpoints inject configured headers and honor x-pass-* client headers (prefix stripped) on the way to the upstream. -The upstream is a real public echo service (httpbin.org/anything). Creating the -route via POST /config/pass_through_endpoint, calling it with a virtual key, and -asserting the echo body is the product path operators use; a mock would not -prove the proxy actually rewrote the outbound request. +The upstream is the real Anthropic Messages API rather than an echo service: +Anthropic doesn't echo request headers back, but it does gate real behavior on +two of them, which is enough to prove forwarding without a mock. A static +x-api-key configured on the pass-through endpoint (the caller never supplies +one) must reach upstream, or every call 401s; an invalid x-pass-anthropic-version +sent by the caller must reach upstream with the prefix stripped, and Anthropic +echoes the exact value back in its 400 body, so a unique-per-run marker proves +this specific request's header - not a stale or cached one - got there. """ from __future__ import annotations import pytest -from pydantic import BaseModel, Field, ValidationError +from pydantic import BaseModel, Field from e2e_config import unique_marker -from e2e_http import AuthHeaders, NoBody, StreamingResponse, require_successful_call, unwrap +from e2e_http import AuthHeaders, NoBody, require_successful_call, unwrap +from endpoints_client import MessagesResult from lifecycle import ResourceManager -from models import KeyGenerateBody +from models import ChatMessage, KeyGenerateBody from passthrough_client import PassthroughClient pytestmark = pytest.mark.e2e -ECHO_TARGET = "https://httpbin.org/anything" -STATIC_HEADER_NAME = "x-e2e-static-header" -PASS_HEADER_STEM = "e2e-client-marker" -PASS_HEADER_NAME = f"x-pass-{PASS_HEADER_STEM}" +ANTHROPIC_MESSAGES_TARGET = "https://api.anthropic.com/v1/messages" +MODEL = "claude-haiku-4-5-20251001" class PassThroughCreateBody(BaseModel): @@ -48,30 +51,26 @@ class PassThroughDeleteParams(BaseModel): endpoint_id: str -class EchoCallHeaders(AuthHeaders): +class AnthropicPassThroughHeaders(AuthHeaders): content_type: str = Field(default="application/json", serialization_alias="Content-Type") - x_pass_e2e_client_marker: str = Field(serialization_alias="x-pass-e2e-client-marker") + x_pass_anthropic_version: str = Field(serialization_alias="x-pass-anthropic-version") -class EchoBody(BaseModel): - ping: str +class AnthropicMessagesBody(BaseModel): + model: str + max_tokens: int = 8 + messages: list[ChatMessage] -class EchoResponse(BaseModel): - headers: dict[str, str] - - -def _create_passthrough( - client: PassthroughClient, *, path: str, static_value: str -) -> PassThroughEndpoint: +def _create_passthrough(client: PassthroughClient, *, path: str) -> PassThroughEndpoint: created = unwrap( client.proxy.transport.post( "/config/pass_through_endpoint", headers=client.proxy.transport.master, json=PassThroughCreateBody( path=path, - target=ECHO_TARGET, - headers={STATIC_HEADER_NAME: static_value}, + target=ANTHROPIC_MESSAGES_TARGET, + headers={"x-api-key": "os.environ/ANTHROPIC_API_KEY"}, ), response_type=PassThroughCreateResponse, ) @@ -92,12 +91,8 @@ def _delete_passthrough(client: PassthroughClient, endpoint_id: str) -> None: ) -def _echo_headers(resp: StreamingResponse) -> dict[str, str]: - try: - echo = EchoResponse.model_validate_json(resp.body) - except ValidationError as exc: - pytest.fail(f"echo upstream did not return a headers map: {exc}; body={resp.body[:300]}") - return {k.lower(): v for k, v in echo.headers.items()} +def _messages_body() -> AnthropicMessagesBody: + return AnthropicMessagesBody(model=MODEL, messages=[ChatMessage(role="user", content="Say hi.")]) class TestPassthroughHeaders: @@ -110,10 +105,8 @@ class TestPassthroughHeaders: ) -> None: marker = unique_marker() path = f"/e2e-passthrough-headers-{marker}" - static_value = f"static-{marker}" - client_value = f"client-{marker}" - endpoint = _create_passthrough(client, path=path, static_value=static_value) + endpoint = _create_passthrough(client, path=path) assert endpoint.id is not None resources.defer(lambda: _delete_passthrough(client, endpoint.id or "")) @@ -128,23 +121,32 @@ class TestPassthroughHeaders: result = client.proxy.transport.send( path, - headers=EchoCallHeaders( + headers=AnthropicPassThroughHeaders( authorization=f"Bearer {key}", - x_pass_e2e_client_marker=client_value, + x_pass_anthropic_version="2023-06-01", ), - json=EchoBody(ping=marker), + json=_messages_body(), ) require_successful_call(result) + completion = MessagesResult.model_validate_json(result.body) + assert completion.text.strip(), ( + f"static x-api-key must reach Anthropic for the call to succeed at all; got {result.body[:300]}" + ) - upstream = _echo_headers(result) - assert upstream.get(STATIC_HEADER_NAME) == static_value, ( - f"configured pass-through header {STATIC_HEADER_NAME!r} not on upstream " - f"request; got {upstream}" + invalid_version = f"e2e-passhdr-{unique_marker()}" + blocked = client.proxy.transport.send( + path, + headers=AnthropicPassThroughHeaders( + authorization=f"Bearer {key}", + x_pass_anthropic_version=invalid_version, + ), + json=_messages_body(), ) - assert upstream.get(PASS_HEADER_STEM) == client_value, ( - f"x-pass-* header should strip the prefix and forward as {PASS_HEADER_STEM!r}; " - f"got {upstream}" + assert blocked.status_code == 400, ( + f"expected Anthropic to reject the invalid anthropic-version, got " + f"{blocked.status_code}: {blocked.body[:300]}" ) - assert PASS_HEADER_NAME not in upstream, ( - "upstream must not see the x-pass- prefix; proxy should strip it" + assert invalid_version in blocked.body, ( + f"x-pass-anthropic-version must reach upstream with the prefix stripped; " + f"marker missing from Anthropic's error body: {blocked.body[:300]}" ) diff --git a/tests/e2e/llm_translation/test_rerank_e2e.py b/tests/e2e/llm_translation/test_rerank_e2e.py index 4b30ac1ea5c..c3614251e77 100644 --- a/tests/e2e/llm_translation/test_rerank_e2e.py +++ b/tests/e2e/llm_translation/test_rerank_e2e.py @@ -23,9 +23,20 @@ DOCUMENTS = [ "Washington, D.C. is the capital of the United States.", "Capital punishment has existed in the United States since before it was a country.", ] +QUERY = "What is the capital of the United States?" + + +def _assert_top_n_scored(body: str) -> None: + parsed = RerankResult.model_validate_json(body) + assert parsed.results, f"/rerank returned no results: {body[:300]}" + assert len(parsed.results) <= 3, f"top_n=3 not honored: {body[:300]}" + assert parsed.results[0].relevance_score is not None, ( + f"top rerank result has no relevance_score: {body[:300]}" + ) class TestRerank: + @pytest.mark.covers("llm.rerank.cohere.basic.nonstream.works") def test_rerank_scores_top_n( self, endpoints_client: EndpointsClient, resources: ResourceManager ) -> None: @@ -37,13 +48,27 @@ class TestRerank: resources.defer(lambda: endpoints_client.delete_model(model_id)) key = resources.key() - result = endpoints_client.rerank( - key, model, "What is the capital of the United States?", DOCUMENTS, top_n=3 - ) + result = endpoints_client.rerank(key, model, QUERY, DOCUMENTS, top_n=3) require_successful_call(result) - parsed = RerankResult.model_validate_json(result.body) - assert parsed.results, f"/rerank returned no results: {result.body[:300]}" - assert len(parsed.results) <= 3, f"top_n=3 not honored: {result.body[:300]}" - assert parsed.results[0].relevance_score is not None, ( - f"top rerank result has no relevance_score: {result.body[:300]}" + _assert_top_n_scored(result.body) + + @pytest.mark.covers("llm.rerank.bedrock.basic.nonstream.works", exercised_on=["rerank"]) + def test_bedrock_rerank_scores_top_n( + self, endpoints_client: EndpointsClient, resources: ResourceManager + ) -> None: + model = f"e2e-bedrock-rerank-{unique_marker()}" + model_id = endpoints_client.create_model( + model, + LiteLLMParamsBody( + model="bedrock/amazon.rerank-v1:0", + aws_access_key_id="os.environ/AWS_ACCESS_KEY_ID", + aws_secret_access_key="os.environ/AWS_SECRET_ACCESS_KEY", + aws_region_name="os.environ/AWS_REGION", + ), ) + resources.defer(lambda: endpoints_client.delete_model(model_id)) + key = resources.key() + + result = endpoints_client.rerank(key, model, QUERY, DOCUMENTS, top_n=3) + require_successful_call(result) + _assert_top_n_scored(result.body) diff --git a/tests/e2e/llm_translation/test_responses_e2e.py b/tests/e2e/llm_translation/test_responses_e2e.py index bd98f11c045..0b2ffce5b2a 100644 --- a/tests/e2e/llm_translation/test_responses_e2e.py +++ b/tests/e2e/llm_translation/test_responses_e2e.py @@ -29,6 +29,26 @@ from models import LiteLLMParamsBody pytestmark = pytest.mark.e2e +BEDROCK_CONVERSE_BACKEND = "bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0" + +WEATHER_TOOL = ResponsesFunctionTool( + name="get_weather", + description="Get the weather for a location", + parameters=FunctionParameters( + properties={"location": FunctionParameterProperty(type="string")}, + required=["location"], + ), +) + + +def _bedrock_params() -> LiteLLMParamsBody: + return LiteLLMParamsBody( + model=BEDROCK_CONVERSE_BACKEND, + aws_access_key_id="os.environ/AWS_ACCESS_KEY_ID", + aws_secret_access_key="os.environ/AWS_SECRET_ACCESS_KEY", + aws_region_name="os.environ/AWS_REGION", + ) + class WeatherArguments(BaseModel): location: str @@ -190,6 +210,82 @@ class TestResponses: parsed = ResponsesResult.model_validate_json(result.body) assert parsed.text.strip(), f"/responses returned no output text: {result.body[:300]}" + @pytest.mark.covers("llm.responses.anthropic.tool_use.nonstream.works") + def test_responses_anthropic_returns_function_call( + self, endpoints_client: EndpointsClient, resources: ResourceManager + ) -> None: + model = f"e2e-responses-{unique_marker()}" + model_id = endpoints_client.create_model( + model, + LiteLLMParamsBody( + model="anthropic/claude-haiku-4-5", api_key="os.environ/ANTHROPIC_API_KEY" + ), + ) + resources.defer(lambda: endpoints_client.delete_model(model_id)) + key = resources.key() + + result = endpoints_client.responses_with_tools( + key, + model, + "What is the weather in San Francisco? Use the get_weather tool.", + [ + ResponsesFunctionTool( + name="get_weather", + description="Get the weather for a location", + parameters=FunctionParameters( + properties={"location": FunctionParameterProperty(type="string")}, + required=["location"], + ), + ) + ], + ) + require_successful_call(result) + parsed = ResponsesResult.model_validate_json(result.body) + function_call = next( + (call for call in parsed.function_calls if call.name == "get_weather"), + None, + ) + assert function_call is not None, f"no get_weather function call: {result.body[:500]}" + assert function_call.arguments is not None + raw_arguments = cast(object, json.loads(function_call.arguments)) + arguments = WeatherArguments.model_validate(raw_arguments) + assert arguments.location, f"function call arguments missing location: {function_call.arguments}" + + @pytest.mark.covers("llm.responses.bedrock_converse.basic.nonstream.works") + def test_responses_bedrock_returns_completion( + self, endpoints_client: EndpointsClient, resources: ResourceManager + ) -> None: + model = f"e2e-responses-{unique_marker()}" + model_id = endpoints_client.create_model(model, _bedrock_params()) + resources.defer(lambda: endpoints_client.delete_model(model_id)) + key = resources.key() + + result = endpoints_client.responses(key, model, "reply with one word") + require_successful_call(result) + parsed = ResponsesResult.model_validate_json(result.body) + assert parsed.text.strip(), f"/responses over bedrock returned no output text: {result.body[:300]}" + + @pytest.mark.covers("llm.responses.bedrock_converse.tool_use.nonstream.works") + def test_responses_bedrock_returns_function_call( + self, endpoints_client: EndpointsClient, resources: ResourceManager + ) -> None: + model = f"e2e-responses-{unique_marker()}" + model_id = endpoints_client.create_model(model, _bedrock_params()) + resources.defer(lambda: endpoints_client.delete_model(model_id)) + key = resources.key() + + result = endpoints_client.responses_with_tools( + key, model, "What is the weather in San Francisco? Use the get_weather tool.", [WEATHER_TOOL] + ) + require_successful_call(result) + parsed = ResponsesResult.model_validate_json(result.body) + function_call = next((call for call in parsed.function_calls if call.name == "get_weather"), None) + assert function_call is not None, f"no get_weather function call over bedrock: {result.body[:500]}" + assert function_call.arguments is not None + raw_arguments = cast(object, json.loads(function_call.arguments)) + arguments = WeatherArguments.model_validate(raw_arguments) + assert arguments.location, f"function call arguments missing location: {function_call.arguments}" + def _parse_stream_event( event: str, diff --git a/tests/e2e/llm_translation/test_responses_metadata_e2e.py b/tests/e2e/llm_translation/test_responses_metadata_e2e.py index 6cf24348095..df854dcfa19 100644 --- a/tests/e2e/llm_translation/test_responses_metadata_e2e.py +++ b/tests/e2e/llm_translation/test_responses_metadata_e2e.py @@ -14,7 +14,7 @@ import time import pytest from pydantic import BaseModel, ConfigDict -from e2e_config import require_env, unique_marker +from e2e_config import unique_marker from e2e_http import require_successful_call from endpoints_client import EndpointsClient, ResponsesResult from lifecycle import ResourceManager @@ -42,7 +42,7 @@ class RedisKeyInfo(BaseModel): def _redis_scan(marker: str) -> tuple[RedisKeyInfo, ...]: import redis - (host,) = require_env("REDIS_HOST") + host = os.environ["REDIS_HOST"] port = int((os.environ.get("REDIS_PORT") or "6379").strip() or "6379") try: with socket.create_connection((host, port), timeout=3): diff --git a/tests/e2e/load/conftest.py b/tests/e2e/load/conftest.py index e9fba02680d..89a571af83e 100644 --- a/tests/e2e/load/conftest.py +++ b/tests/e2e/load/conftest.py @@ -1,10 +1,12 @@ from __future__ import annotations +import os from collections.abc import Iterator import pytest from requests import RequestException +from e2e_config import WEEKLY_ANOMALY_OPT_IN_ENV from e2e_http import NoBody, Success from load_client import LoadClient, build_client from load_constants import LOAD_MODEL @@ -18,6 +20,22 @@ LOAD_MODEL_PARAMS = LiteLLMParamsBody( ) +def pytest_collection_modifyitems( + config: pytest.Config, items: list[pytest.Item] +) -> None: + if os.environ.get(WEEKLY_ANOMALY_OPT_IN_ENV): + return + deselected = [ + item for item in items if item.get_closest_marker("weekly") is not None + ] + if not deselected: + return + config.hook.pytest_deselected(items=deselected) + items[:] = [ + item for item in items if item.get_closest_marker("weekly") is None + ] + + @pytest.fixture(scope="session") def client(proxy: ProxyClient) -> LoadClient: return build_client(proxy) @@ -33,10 +51,8 @@ def _model_is_servable(proxy: ProxyClient, model_name: str) -> bool: return isinstance(result, Success) and any(entry.id == model_name for entry in result.data.data) -@pytest.fixture(scope="session", autouse=True) -def _ensure_load_model( # pyright: ignore[reportUnusedFunction] # pytest autouse session fixture, wired by name - client: LoadClient, -) -> Iterator[None]: +@pytest.fixture(scope="session") +def ensure_load_model(client: LoadClient) -> Iterator[None]: proxy = client.proxy if _model_is_servable(proxy, LOAD_MODEL): yield @@ -60,7 +76,9 @@ def _ensure_load_model( # pyright: ignore[reportUnusedFunction] # pytest autou @pytest.fixture -def load_key(resources: ResourceManager, client: LoadClient) -> str: +def load_key( + resources: ResourceManager, client: LoadClient, ensure_load_model: None +) -> str: key = client.proxy.generate_key(KeyGenerateBody(models=[LOAD_MODEL], user_id="e2e-load")) resources.defer(lambda: client.proxy.delete_key(key)) return key diff --git a/tests/e2e/load/session_anomaly.py b/tests/e2e/load/session_anomaly.py new file mode 100644 index 00000000000..c29b833635d --- /dev/null +++ b/tests/e2e/load/session_anomaly.py @@ -0,0 +1,299 @@ +from __future__ import annotations + +import time +from collections.abc import Callable +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass + +from pydantic import BaseModel + +from e2e_config import unique_marker +from e2e_http import Result, Success +from models import CacheControl, RichMessage, TextBlock +from transport import Transport + + +class SessionMessagesRequest(BaseModel): + model: str + max_tokens: int = 128 + system: list[TextBlock] + messages: list[RichMessage] + + +class SessionUsage(BaseModel): + input_tokens: int = 0 + output_tokens: int = 0 + cache_creation_input_tokens: int = 0 + cache_read_input_tokens: int = 0 + + +class SessionContentBlock(BaseModel): + type: str | None = None + text: str | None = None + + +class SessionMessagesResponse(BaseModel): + content: list[SessionContentBlock] = [] + usage: SessionUsage = SessionUsage() + + @property + def text(self) -> str: + return "".join(block.text or "" for block in self.content) + + +@dataclass(frozen=True, slots=True) +class TurnMetric: + turn_index: int + ok: bool + latency_seconds: float + uncached_input_tokens: int + cache_read_tokens: int + cache_creation_tokens: int + failure: str | None + + +@dataclass(frozen=True, slots=True) +class AnomalyReport: + planned_turns: int + attempted_turns: int + failed_turns: int + warm_turns: int + warm_uncached_input_tokens: int + warm_cache_read_tokens: int + warm_cache_creation_tokens: int + p95_turn_seconds: float + + @property + def error_ratio(self) -> float: + return self.failed_turns / self.planned_turns if self.planned_turns else 1.0 + + @property + def warm_cache_read_share(self) -> float: + billed = ( + self.warm_uncached_input_tokens + + self.warm_cache_read_tokens + + self.warm_cache_creation_tokens + ) + return self.warm_cache_read_tokens / billed if billed else 0.0 + + +def _system_prefix_block(marker: str) -> TextBlock: + text = " ".join( + f"Project context paragraph {index} for session {marker}." for index in range(300) + ) + return TextBlock(text=text, cache_control=CacheControl()) + + +def _user_turn_text(marker: str, turn_index: int) -> str: + notes = " ".join( + f"Working note {index} of turn {turn_index} in session {marker}." + for index in range(80) + ) + return f"Reply with one short sentence.\n{notes}" + + +def _reminder_turn() -> RichMessage: + return RichMessage( + role="system", + content=[ + TextBlock( + text="Keep the answer to one short sentence." + ) + ], + ) + + +def _without_cache_control(message: RichMessage) -> RichMessage: + return RichMessage( + role=message.role, + content=[TextBlock(text=block.text) for block in message.content], + ) + + +RETRY_BACKOFF_SECONDS = 2.0 + + +def retried( + call: Callable[[], Result[SessionMessagesResponse]], + attempts: int, + backoff_seconds: float = RETRY_BACKOFF_SECONDS, + sleep: Callable[[float], None] = time.sleep, +) -> Result[SessionMessagesResponse]: + result = call() + if isinstance(result, Success) or attempts <= 1: + return result + sleep(backoff_seconds) + return retried(call, attempts - 1, backoff_seconds, sleep) + + +def _metric( + result: Result[SessionMessagesResponse], turn_index: int, latency_seconds: float +) -> TurnMetric: + if isinstance(result, Success): + usage = result.data.usage + return TurnMetric( + turn_index=turn_index, + ok=True, + latency_seconds=latency_seconds, + uncached_input_tokens=usage.input_tokens, + cache_read_tokens=usage.cache_read_input_tokens, + cache_creation_tokens=usage.cache_creation_input_tokens, + failure=None, + ) + return TurnMetric( + turn_index=turn_index, + ok=False, + latency_seconds=latency_seconds, + uncached_input_tokens=0, + cache_read_tokens=0, + cache_creation_tokens=0, + failure=repr(result), + ) + + +def _drive_turns( + transport: Transport, + key: str, + model: str, + marker: str, + system_block: TextBlock, + history: tuple[RichMessage, ...], + turn_index: int, + remaining_turns: int, + attempts_per_turn: int, +) -> tuple[TurnMetric, ...]: + if remaining_turns == 0: + return () + user_turn = RichMessage( + role="user", + content=[ + TextBlock( + text=_user_turn_text(marker, turn_index), cache_control=CacheControl() + ) + ], + ) + started = time.monotonic() + result = retried( + lambda: transport.post( + "/v1/messages", + headers=transport.bearer(key), + json=SessionMessagesRequest( + model=model, + system=[system_block], + messages=[*history, user_turn], + ), + response_type=SessionMessagesResponse, + ), + attempts_per_turn, + ) + turn = _metric(result, turn_index, time.monotonic() - started) + if not isinstance(result, Success): + return (turn,) + assistant_turn = RichMessage( + role="assistant", content=[TextBlock(text=result.data.text or "Understood.")] + ) + return ( + turn, + *_drive_turns( + transport, + key, + model, + marker, + system_block, + ( + *history, + _without_cache_control(user_turn), + _reminder_turn(), + assistant_turn, + ), + turn_index + 1, + remaining_turns - 1, + attempts_per_turn, + ), + ) + + +def run_session( + transport: Transport, key: str, model: str, turns: int, attempts_per_turn: int +) -> tuple[TurnMetric, ...]: + marker = unique_marker() + return _drive_turns( + transport, + key, + model, + marker, + _system_prefix_block(marker), + (), + 1, + turns, + attempts_per_turn, + ) + + +def run_concurrent_sessions( + transport: Transport, + key: str, + model: str, + sessions: int, + turns_per_session: int, + attempts_per_turn: int, +) -> tuple[TurnMetric, ...]: + with ThreadPoolExecutor(max_workers=sessions) as pool: + futures = [ + pool.submit( + run_session, transport, key, model, turns_per_session, attempts_per_turn + ) + for _ in range(sessions) + ] + return tuple(turn for future in futures for turn in future.result()) + + +def settled_spend( + read_spend: Callable[[], float], + poll_interval: float, + settle_seconds: float, + timeout_seconds: float, + now: Callable[[], float] = time.monotonic, + sleep: Callable[[float], None] = time.sleep, +) -> float: + deadline = now() + timeout_seconds + settle_seconds + + def settle(previous: float, stable_since: float) -> float: + current = read_spend() + observed = now() + since = stable_since if current == previous else observed + if current > 0 and observed - since >= settle_seconds: + return current + if observed >= deadline: + raise AssertionError( + f"key spend never held a stable non-zero value for {settle_seconds}s " + f"within {timeout_seconds + settle_seconds}s (last read {current}); " + f"spend stopped being recorded, which is itself a spend anomaly" + ) + sleep(poll_interval) + return settle(current, since) + + return settle(-1.0, now()) + + +def _p95(latencies: tuple[float, ...]) -> float: + if not latencies: + return 0.0 + ranked = sorted(latencies) + return ranked[max(0, -(-len(ranked) * 95 // 100) - 1)] + + +def summarize(turns: tuple[TurnMetric, ...], planned_turns: int) -> AnomalyReport: + warm = tuple(turn for turn in turns if turn.ok and turn.turn_index >= 2) + return AnomalyReport( + planned_turns=planned_turns, + attempted_turns=len(turns), + failed_turns=planned_turns - sum(1 for turn in turns if turn.ok), + warm_turns=len(warm), + warm_uncached_input_tokens=sum(turn.uncached_input_tokens for turn in warm), + warm_cache_read_tokens=sum(turn.cache_read_tokens for turn in warm), + warm_cache_creation_tokens=sum(turn.cache_creation_tokens for turn in warm), + p95_turn_seconds=_p95( + tuple(turn.latency_seconds for turn in turns if turn.ok) + ), + ) diff --git a/tests/e2e/load/test_session_anomaly.py b/tests/e2e/load/test_session_anomaly.py new file mode 100644 index 00000000000..7062587352b --- /dev/null +++ b/tests/e2e/load/test_session_anomaly.py @@ -0,0 +1,141 @@ +from __future__ import annotations + +from itertools import count, repeat + +import pytest + +from e2e_http import NetworkError, Success +from session_anomaly import ( + SessionMessagesResponse, + TurnMetric, + retried, + settled_spend, + summarize, +) + + +def _ok_turn(turn_index: int) -> TurnMetric: + return TurnMetric( + turn_index=turn_index, + ok=True, + latency_seconds=1.0, + uncached_input_tokens=10, + cache_read_tokens=100, + cache_creation_tokens=5, + failure=None, + ) + + +def _failed_turn(turn_index: int) -> TurnMetric: + return TurnMetric( + turn_index=turn_index, + ok=False, + latency_seconds=1.0, + uncached_input_tokens=0, + cache_read_tokens=0, + cache_creation_tokens=0, + failure="NetworkError()", + ) + + +class TestSummarizePlannedTurns: + def test_session_aborted_on_first_turn_counts_all_its_planned_turns_as_failed( + self, + ) -> None: + completed_session = tuple(_ok_turn(index) for index in range(1, 7)) + aborted_session = (_failed_turn(1),) + + report = summarize((*completed_session, *aborted_session), planned_turns=12) + + assert report.attempted_turns == 7 + assert report.failed_turns == 6 + assert report.error_ratio == 0.5 + + def test_all_planned_turns_completing_reports_zero_failures(self) -> None: + report = summarize( + tuple(_ok_turn(index) for index in range(1, 7)), planned_turns=6 + ) + + assert report.failed_turns == 0 + assert report.error_ratio == 0.0 + + +class TestRetried: + def test_transient_failures_then_success_returns_the_success(self) -> None: + outcome = Success[SessionMessagesResponse](status_code=200, data=SessionMessagesResponse()) + calls = iter( + (NetworkError(message="overloaded"), NetworkError(message="overloaded"), outcome) + ) + + result = retried(lambda: next(calls), attempts=3, sleep=lambda _: None) + + assert result is outcome + + def test_exhausted_attempts_return_the_last_failure(self) -> None: + last_attempt = NetworkError(message="still overloaded") + never_reached = NetworkError(message="a fourth attempt would break the budget") + calls = iter( + (NetworkError(message="overloaded"), last_attempt, never_reached) + ) + + result = retried(lambda: next(calls), attempts=2, sleep=lambda _: None) + + assert result is last_attempt + assert next(calls) is never_reached + + def test_first_try_success_never_sleeps(self) -> None: + def sleep_means_retry(_: float) -> None: + raise AssertionError("slept after a successful attempt") + + result = retried( + lambda: Success[SessionMessagesResponse](status_code=200, data=SessionMessagesResponse()), + attempts=3, + sleep=sleep_means_retry, + ) + + assert isinstance(result, Success) + + +class TestSettledSpend: + def test_partial_total_between_batch_flushes_is_not_accepted_as_final(self) -> None: + reads = iter((0.1, 0.1, 0.1, 0.35, 0.35, 0.35, 0.35, 0.35)) + ticks = count(0.0, 2.5) + + spend = settled_spend( + lambda: next(reads), + poll_interval=5.0, + settle_seconds=10.0, + timeout_seconds=100.0, + now=lambda: next(ticks), + sleep=lambda _: None, + ) + + assert spend == 0.35 + + def test_spend_that_never_stabilizes_raises(self) -> None: + reads = (0.1 * step for step in count(1)) + ticks = count(0.0, 2.5) + + with pytest.raises(AssertionError, match="spend anomaly"): + settled_spend( + lambda: next(reads), + poll_interval=5.0, + settle_seconds=5.0, + timeout_seconds=10.0, + now=lambda: next(ticks), + sleep=lambda _: None, + ) + + def test_spend_that_never_becomes_nonzero_raises(self) -> None: + reads = repeat(0.0) + ticks = count(0.0, 2.5) + + with pytest.raises(AssertionError, match="spend anomaly"): + settled_spend( + lambda: next(reads), + poll_interval=5.0, + settle_seconds=5.0, + timeout_seconds=10.0, + now=lambda: next(ticks), + sleep=lambda _: None, + ) diff --git a/tests/e2e/load/test_weekly_session_anomaly_e2e.py b/tests/e2e/load/test_weekly_session_anomaly_e2e.py new file mode 100644 index 00000000000..d4ef883702e --- /dev/null +++ b/tests/e2e/load/test_weekly_session_anomaly_e2e.py @@ -0,0 +1,124 @@ +from __future__ import annotations + +from dataclasses import dataclass + +import pytest + +from e2e_config import ( + ANOMALY_MAX_ERROR_RATIO, + ANOMALY_MAX_KEY_SPEND_USD, + ANOMALY_MAX_P95_TURN_SECONDS, + ANOMALY_MIN_WARM_CACHE_READ_SHARE, + ANOMALY_SESSIONS, + ANOMALY_SPEND_SETTLE_SECONDS, + ANOMALY_TURN_ATTEMPTS, + ANOMALY_TURNS_PER_SESSION, + unique_marker, +) +from lifecycle import ResourceManager +from load_client import LoadClient +from models import KeyGenerateBody, LiteLLMParamsBody +from proxy_client import ProxyClient +from session_anomaly import run_concurrent_sessions, settled_spend, summarize + +pytestmark = [pytest.mark.e2e, pytest.mark.load, pytest.mark.weekly] + + +@dataclass(frozen=True, slots=True) +class AnomalyRoute: + route_id: str + params: LiteLLMParamsBody + + +ANOMALY_ROUTES = ( + AnomalyRoute( + route_id="anthropic", + params=LiteLLMParamsBody(model="anthropic/claude-sonnet-5"), + ), + AnomalyRoute( + route_id="bedrock_invoke", + params=LiteLLMParamsBody( + model="bedrock/invoke/us.anthropic.claude-sonnet-5", + aws_region_name="us-east-1", + ), + ), +) + + +def _route_id(route: AnomalyRoute) -> str: + return route.route_id + + +def _settled_key_spend(proxy: ProxyClient, key: str) -> float: + return settled_spend( + lambda: proxy.key_info(key).spend or 0.0, + proxy.poll_interval, + ANOMALY_SPEND_SETTLE_SECONDS, + proxy.poll_timeout, + ) + + +class TestWeeklySessionAnomaly: + @pytest.mark.covers("reliability.perf.session_anomaly.under_slo") + @pytest.mark.parametrize("route", ANOMALY_ROUTES, ids=_route_id) + def test_session_load_stays_within_baselines( + self, client: LoadClient, resources: ResourceManager, route: AnomalyRoute + ) -> None: + model_name = f"weekly-anomaly-{route.route_id}-{unique_marker()}" + model_id = client.proxy.create_model(model_name, route.params) + resources.defer(lambda: client.proxy.delete_model(model_id)) + key = client.proxy.generate_key( + KeyGenerateBody(models=[model_name], key_alias=model_name) + ) + resources.defer(lambda: client.proxy.delete_key(key)) + + turns = run_concurrent_sessions( + client.proxy.transport, + key, + model_name, + ANOMALY_SESSIONS, + ANOMALY_TURNS_PER_SESSION, + ANOMALY_TURN_ATTEMPTS, + ) + report = summarize(turns, ANOMALY_SESSIONS * ANOMALY_TURNS_PER_SESSION) + failures = tuple(turn.failure for turn in turns if turn.failure) + print(f"{route.route_id} anomaly report: {report}") + + assert report.error_ratio <= ANOMALY_MAX_ERROR_RATIO, ( + f"{route.route_id}: {report.failed_turns}/{report.planned_turns} planned " + f"turns failed or never ran because their session aborted " + f"({report.error_ratio:.1%} > {ANOMALY_MAX_ERROR_RATIO:.1%} allowed); " + f"error rate is anomalously high. Failures: {failures}" + ) + assert report.warm_turns > 0, ( + f"{route.route_id}: no session got past its first turn, so cache and " + f"latency baselines have nothing to read. Failures: {failures}" + ) + assert report.warm_cache_read_share >= ANOMALY_MIN_WARM_CACHE_READ_SHARE, ( + f"{route.route_id}: warm turns read only {report.warm_cache_read_share:.1%} " + f"of billed input tokens from the prompt cache " + f"(read={report.warm_cache_read_tokens}, " + f"creation={report.warm_cache_creation_tokens}, " + f"uncached={report.warm_uncached_input_tokens}), below the " + f"{ANOMALY_MIN_WARM_CACHE_READ_SHARE:.0%} floor; the cached prefix is " + f"being invalidated between turns (the mid-conversation-system cache " + f"collapse signature) or caching stopped working" + ) + assert report.warm_cache_creation_tokens > 0, ( + f"{route.route_id}: warm turns wrote 0 cache-creation tokens across " + f"{report.warm_turns} turns; the moving cache breakpoint stopped writing " + f"new prefix increments" + ) + assert report.p95_turn_seconds <= ANOMALY_MAX_P95_TURN_SECONDS, ( + f"{route.route_id}: p95 turn time {report.p95_turn_seconds:.1f}s exceeds " + f"the {ANOMALY_MAX_P95_TURN_SECONDS:.0f}s ceiling under " + f"{ANOMALY_SESSIONS} concurrent sessions; turn times are anomalously slow" + ) + + spend = _settled_key_spend(client.proxy, key) + assert spend <= ANOMALY_MAX_KEY_SPEND_USD, ( + f"{route.route_id}: gateway recorded ${spend:.4f} for " + f"{report.attempted_turns} turns, above the " + f"${ANOMALY_MAX_KEY_SPEND_USD} ceiling; spend per session is " + f"anomalously high (cache regressions surface here as 2-3x spend)" + ) diff --git a/tests/e2e/load/weekly_anomaly_config.yml b/tests/e2e/load/weekly_anomaly_config.yml new file mode 100644 index 00000000000..08972969cf0 --- /dev/null +++ b/tests/e2e/load/weekly_anomaly_config.yml @@ -0,0 +1,3 @@ +general_settings: + master_key: os.environ/LITELLM_MASTER_KEY + store_model_in_db: true diff --git a/tests/e2e/management/management_client.py b/tests/e2e/management/management_client.py index a9dedac8e61..cdc31aeea79 100644 --- a/tests/e2e/management/management_client.py +++ b/tests/e2e/management/management_client.py @@ -14,6 +14,10 @@ from e2e_http import NoBody, ProbeResult, Result, StreamingResponse, Success, Un from models import ( ChatBody, ChatMessage, + CustomerDeleteBody, + CustomerInfoParams, + CustomerNewBody, + CustomerResponse, KeyBlockBody, KeyDeleteBody, KeyGenerateBody, @@ -270,6 +274,35 @@ class ManagementClient: ) ).user_id + def create_customer(self, user_id: str) -> str: + _ = unwrap( + self.proxy.transport.post( + "/customer/new", + headers=self.proxy.transport.master, + json=CustomerNewBody(user_id=user_id), + response_type=CustomerResponse, + ) + ) + return user_id + + def customer_info(self, end_user_id: str) -> CustomerResponse: + return unwrap( + self.proxy.transport.get( + "/customer/info", + headers=self.proxy.transport.master, + params=CustomerInfoParams(end_user_id=end_user_id), + response_type=CustomerResponse, + ) + ) + + def delete_customer(self, user_id: str) -> None: + _ = self.proxy.transport.post( + "/customer/delete", + headers=self.proxy.transport.master, + json=CustomerDeleteBody(user_ids=[user_id]), + response_type=NoBody, + ) + def update_user(self, body: UserUpdateBody) -> None: _ = unwrap( self.proxy.transport.post( diff --git a/tests/e2e/management/test_management_e2e.py b/tests/e2e/management/test_management_e2e.py index 18bc384a879..9b398963ac9 100644 --- a/tests/e2e/management/test_management_e2e.py +++ b/tests/e2e/management/test_management_e2e.py @@ -610,3 +610,18 @@ class TestManagementRoutePermissions: f"/team/info returned {team_probe.status_code}: {team_probe.body[:300]}" ) assert client.user_count(user_id) == 0, f"user {user_id} was created despite the 403 route denial" + + +class TestCustomer: + @pytest.mark.covers("mgmt.end_user.new.happy_path") + def test_customer_create_persists_to_info( + self, client: ManagementClient, resources: ResourceManager + ) -> None: + customer = f"e2e-customer-{unique_marker()}" + client.create_customer(customer) + resources.defer(lambda: client.delete_customer(customer)) + + info = client.customer_info(customer) + assert info.user_id == customer, ( + f"/customer/info did not report the created end-user; got {info.user_id!r}" + ) diff --git a/tests/e2e/mcp/linear_session_capture.py b/tests/e2e/mcp/linear_session_capture.py new file mode 100644 index 00000000000..1c867e17e22 --- /dev/null +++ b/tests/e2e/mcp/linear_session_capture.py @@ -0,0 +1,58 @@ +"""One-time helper to capture a logged-in Linear browser session for the +real-Linear MCP e2e test. + +The real-Linear test drives the genuine gateway-managed authorization_code +dance against ``mcp.linear.app``. The only step that cannot be scripted is +Linear's login (magic link / SSO), so a human authenticates once here and the +resulting session (cookies + local storage) is persisted to disk. The e2e test +then loads that session in a headless Playwright context and clicks Approve on +Linear's consent screen every run, with no human and no login automation. + +Run it with the e2e venv, log into Linear in the window that opens, then return +to the terminal and press Enter: + + LITELLM=~/litellm-mcpe2e + "$LITELLM"/.venv/bin/python "$LITELLM"/tests/e2e/mcp/linear_session_capture.py + +The session is written to ``E2E_LINEAR_STORAGE_STATE`` (default +``~/.litellm-e2e/linear_storage_state.json``), outside the repo. It is a +secret: never commit it. Re-run this whenever Linear expires the session. +""" + +from __future__ import annotations + +import os +from pathlib import Path + +from playwright.sync_api import sync_playwright + +DEFAULT_STATE_PATH = Path.home() / ".litellm-e2e" / "linear_storage_state.json" + + +def capture(state_path: Path) -> None: + """Open a headed browser at Linear, wait for the human to log in, then save + the authenticated session to ``state_path``.""" + state_path.parent.mkdir(parents=True, exist_ok=True) + with sync_playwright() as playwright: + browser = playwright.chromium.launch(headless=False) + context = browser.new_context() + page = context.new_page() + page.goto("https://linear.app/login", wait_until="domcontentloaded") + print("\n" + "=" * 72) + print("Log into Linear in the browser window that just opened.") + print("If Linear emails you a magic link, paste the link into THIS window's") + print("address bar (opening it in your default browser won't capture the") + print("session). Google SSO works too as long as you complete it here.") + print("When your Linear workspace has loaded, come back and press Enter.") + print("=" * 72) + input("Press Enter once you are logged in... ") + page.goto("https://mcp.linear.app/", wait_until="domcontentloaded") + context.storage_state(path=str(state_path)) + browser.close() + print(f"\nSaved Linear session to {state_path}") + print("Point the e2e test at it with:") + print(f' export E2E_LINEAR_STORAGE_STATE="{state_path}"') + + +if __name__ == "__main__": + capture(Path(os.environ.get("E2E_LINEAR_STORAGE_STATE", str(DEFAULT_STATE_PATH)))) diff --git a/tests/e2e/mcp/oauth_chat_client.py b/tests/e2e/mcp/oauth_chat_client.py new file mode 100644 index 00000000000..2eaf512cfa5 --- /dev/null +++ b/tests/e2e/mcp/oauth_chat_client.py @@ -0,0 +1,271 @@ +"""Client for the mcp chat-completion OAuth e2e suite. + +Registers a gateway-managed OAuth (authorization_code) MCP server, seeds the +per-user upstream token by driving the interactive authorize dance with the +official mcp SDK's OAuthClientProvider (the browser leg is a headless Chromium +primed with a human's saved Linear session), then exercises the server through +/chat/completions, where the gateway lists and executes its tools with the +stored per-user token. + +Management routes (/v1/mcp/server CRUD, /chat/completions) go through the +shared ProxyClient transport. The MCP protocol used to seed the token goes through +the mcp SDK, the same library production MCP hosts run. +""" + +from __future__ import annotations + +import asyncio +import re +import time +from dataclasses import dataclass +from typing import TYPE_CHECKING +from urllib.parse import parse_qsl + +import httpx +import pytest +from mcp import ClientSession +from mcp.client.auth import OAuthClientProvider +from mcp.client.streamable_http import streamable_http_client +from mcp.shared.auth import OAuthClientInformationFull, OAuthClientMetadata, OAuthToken + +from e2e_config import PROXY_BASE_URL, REQUEST_TIMEOUT +from proxy_client import ProxyClient +from e2e_http import AuthHeaders, NoBody, unwrap +from models import ChatBody, ChatResponse, McpServerCreateBody, McpServerInfo + +if TYPE_CHECKING: + from playwright.async_api import Route + +# Where the "browser" lands at the end of the authorize dance. Nothing listens +# here: the route interceptor short-circuits the final redirect and reads the +# code/state off its query string, exactly like a desktop MCP host intercepting +# its loopback redirect. +OAUTH_CLIENT_REDIRECT_URI = "http://127.0.0.1:53682/e2e/callback" +BROWSER_CONSENT_TIMEOUT = 60.0 + + +def _mcp_url(alias: str) -> str: + return f"{PROXY_BASE_URL}/{alias}/mcp" + + +class InMemoryTokenStorage: + """The mcp SDK's TokenStorage protocol, in memory for one dance: the + DCR-registered client and the gateway tokens minted for it.""" + + def __init__(self) -> None: + self._tokens: OAuthToken | None = None + self._client_info: OAuthClientInformationFull | None = None + + async def get_tokens(self) -> OAuthToken | None: + return self._tokens + + async def set_tokens(self, tokens: OAuthToken) -> None: + self._tokens = tokens + + async def get_client_info(self) -> OAuthClientInformationFull | None: + return self._client_info + + async def set_client_info(self, client_info: OAuthClientInformationFull) -> None: + self._client_info = client_info + + +async def _browser_follow_authorize(start_url: str, storage_state_path: str) -> tuple[str, str | None]: + """Play the browser's role for a real upstream whose authorize endpoint + serves an interactive consent page (Linear). A headless Chromium primed + with a human's saved Linear session opens the gateway authorize URL and + clicks through Linear's consent screens (the mcp.linear.app Approve form, + then the linear.app workspace-selection page), riding the rest of the chain + (Linear -> gateway callback -> host redirect_uri). The final hop is + intercepted and short-circuited, since nothing listens there, and its + code/state are read off the query string.""" + from playwright.async_api import async_playwright + + captured: dict[str, str] = {} # mutable-ok: hand-off from the request listener + trail: list[str] = [] # mutable-ok: navigation diagnostics for a failed dance + + def _note_request(request: object) -> None: + url = getattr(request, "url", "") + if url.startswith(OAUTH_CLIENT_REDIRECT_URI) and "url" not in captured: + captured["url"] = url + + async def _swallow_redirect(route: "Route") -> None: + await route.fulfill(status=200, content_type="text/plain", body="ok") + + async with async_playwright() as playwright: + browser = await playwright.chromium.launch(headless=True) + context = await browser.new_context(storage_state=storage_state_path) + await context.route(re.compile(re.escape(OAUTH_CLIENT_REDIRECT_URI) + r".*"), _swallow_redirect) + page = await context.new_page() + page.on("request", _note_request) + page.on("framenavigated", lambda frame: trail.append(frame.url.split("?", 1)[0])) + await page.goto(start_url, wait_until="domcontentloaded") + deadline = time.monotonic() + BROWSER_CONSENT_TIMEOUT + while "url" not in captured and time.monotonic() < deadline: + try: + await page.wait_for_load_state("networkidle", timeout=8000) + except Exception: # noqa: BLE001 - a busy consent page never idles; fall through and try to advance it + pass + if "url" in captured: + break + control = page.locator( + 'button[name="action"][value="approve"], button:has-text("Authorize"), ' + 'button:has-text("Allow"), button:has-text("@"), a:has-text("@")' + ).first + try: + await control.click(timeout=5000) + except Exception: # noqa: BLE001 - nothing to advance yet; loop and re-check + await asyncio.sleep(0.5) + final_url = page.url + await browser.close() + + landing = captured.get("url") + assert landing is not None, ( + f"consent flow never reached {OAUTH_CLIENT_REDIRECT_URI}; " + f"final={final_url.split('?', 1)[0]!r}; trail={trail[-6:]}" + ) + params = dict(parse_qsl(httpx.URL(landing).query.decode())) + assert "code" in params, f"client redirect_uri carried no code: {landing}" + return params["code"], params.get("state") + + +def _oauth_provider(url: str, storage: InMemoryTokenStorage, storage_state_path: str) -> OAuthClientProvider: + """The SDK's real OAuth machinery (RFC 9728/8414 discovery, RFC 7591 DCR, + PKCE, token exchange) with the browser leg driven by Playwright against the + upstream's consent screen.""" + code_holder: dict[str, str | None] = {} # mutable-ok: hand-off between the two SDK callbacks + + async def redirect_handler(authorize_url: str) -> None: + code, state = await _browser_follow_authorize(authorize_url, storage_state_path) + code_holder["code"] = code + code_holder["state"] = state + + async def callback_handler() -> tuple[str, str | None]: + code = code_holder.get("code") + assert code is not None, "callback_handler ran before the authorize redirect completed" + return code, code_holder.get("state") + + return OAuthClientProvider( + server_url=url, + client_metadata=OAuthClientMetadata.model_validate( + { + "redirect_uris": [OAUTH_CLIENT_REDIRECT_URI], + "token_endpoint_auth_method": "none", + "grant_types": ["authorization_code", "refresh_token"], + "response_types": ["code"], + "client_name": "e2e-mcp-host", + } + ), + storage=storage, + redirect_handler=redirect_handler, + callback_handler=callback_handler, + ) + + +class _HeaderInjectingTransport(httpx.AsyncBaseTransport): + """Adds the caller's LiteLLM key header to every outgoing SDK request + (discovery, DCR, token exchange), so the gateway resolves which user to + store the upstream token for from the key on the token exchange, exactly + like a production MCP host configured with a LiteLLM key header.""" + + def __init__(self, inner: httpx.AsyncBaseTransport, headers: dict[str, str]) -> None: + self._inner = inner + self._headers = headers + + async def handle_async_request(self, request: httpx.Request) -> httpx.Response: + for name, value in self._headers.items(): + if name not in request.headers: + request.headers[name] = value + return await self._inner.handle_async_request(request) + + +def _oauth_http_client(headers: dict[str, str], auth: OAuthClientProvider) -> httpx.AsyncClient: + return httpx.AsyncClient( + headers=headers, + auth=auth, + timeout=httpx.Timeout(REQUEST_TIMEOUT), + follow_redirects=True, + transport=_HeaderInjectingTransport(httpx.AsyncHTTPTransport(), headers), + ) + + +async def _seed_via_dance( + url: str, headers: dict[str, str], storage: InMemoryTokenStorage, storage_state_path: str +) -> tuple[str, ...]: + async with _oauth_http_client(headers, _oauth_provider(url, storage, storage_state_path)) as http_client: + async with streamable_http_client(url, http_client=http_client) as (read, write, _): + async with ClientSession(read, write) as session: + await session.initialize() + listed = await session.list_tools() + return tuple(sorted(tool.name for tool in listed.tools)) + + +@dataclass(frozen=True, slots=True) +class ChatMcpClient: + proxy: ProxyClient + + def create_server(self, body: McpServerCreateBody) -> McpServerInfo: + return unwrap( + self.proxy.transport.post( + "/v1/mcp/server", + headers=self.proxy.transport.master, + json=body, + response_type=McpServerInfo, + ) + ) + + def server_info(self, server_id: str) -> McpServerInfo: + return unwrap( + self.proxy.transport.get( + f"/v1/mcp/server/{server_id}", + headers=self.proxy.transport.master, + params=NoBody(), + response_type=McpServerInfo, + ) + ) + + def delete_server(self, server_id: str) -> None: + _ = self.proxy.transport.delete( + f"/v1/mcp/server/{server_id}", + headers=self.proxy.transport.master, + json=NoBody(), + response_type=NoBody, + ) + + def seed_user_token(self, alias: str, key: str, storage_state_path: str) -> tuple[str, ...]: + """Drive the interactive authorize dance for `key`'s user so the gateway + stores their upstream token, retried to the shared deadline since the + just-created server and key propagate asynchronously. The LiteLLM key + rides x-litellm-api-key so the gateway binds the token to that user. + Returns the upstream tool names the dance listed, proof the token works.""" + headers = {"x-litellm-api-key": f"Bearer {key}"} + storage = InMemoryTokenStorage() + deadline = time.monotonic() + self.proxy.poll_timeout + last_error: Exception | None = None + while time.monotonic() < deadline: + try: + return asyncio.run(_seed_via_dance(_mcp_url(alias), headers, storage, storage_state_path)) + except Exception as exc: # noqa: BLE001 - retried to the deadline; the last error surfaces below + last_error = exc + time.sleep(self.proxy.poll_interval) + pytest.fail( + f"authorize dance for {alias!r} never completed within {self.proxy.poll_timeout}s; " + f"last error: {last_error!r}" + ) + + def chat_with_mcp(self, headers: AuthHeaders, body: ChatBody) -> ChatResponse: + """POST /chat/completions carrying the LiteLLM key in `headers` (either + ingress form) with an MCP server attached in `body.tools`. The gateway + resolves the user from the key and lists/executes the server's tools + with that user's stored upstream token.""" + return unwrap( + self.proxy.transport.post( + "/chat/completions", + headers=headers, + json=body, + response_type=ChatResponse, + ) + ) + + +def build_chat_client(proxy: ProxyClient) -> ChatMcpClient: + return ChatMcpClient(proxy=proxy) diff --git a/tests/e2e/mcp/test_mcp_chat_completion_oauth_e2e.py b/tests/e2e/mcp/test_mcp_chat_completion_oauth_e2e.py new file mode 100644 index 00000000000..01e94f7b86f --- /dev/null +++ b/tests/e2e/mcp/test_mcp_chat_completion_oauth_e2e.py @@ -0,0 +1,197 @@ +"""On-demand e2e: a chat completion drives a gateway-managed OAuth MCP server. + +The real end-user flow for MCP over an OAuth server: a user registers a Linear +authorization_code server, authorizes it once so the gateway stores their +upstream token, then sends a normal /chat/completions request with the Linear +MCP attached. The gateway resolves the user from the LiteLLM key, lists Linear's +tools with the stored per-user token, lets the model call one, executes it +upstream with that token, and returns the answer. This is proven against the +real Linear MCP server (mcp.linear.app) and a real Anthropic model, once per +documented ingress header (x-litellm-api-key and Authorization). + +The authorize dance is seeded through the mcp SDK's OAuthClientProvider; the one +step Linear cannot auto-approve is the human consent, so it is captured once out +of band (mcp/linear_session_capture.py) into a saved browser session and a +headless Chromium clicks Approve every run. The test therefore skips unless +E2E_LINEAR_STORAGE_STATE points at that session, so it never runs on the per-PR +CI path; it is a nightly/on-demand real-server smoke test. + +Fail-before-fix: without the stored per-user token the gateway lists no Linear +tools, so mcp_list_tools comes back empty, nothing is called, and the +assertions fail; a served, called, non-empty Linear tool proves the gateway +pulled and used the user's token. +""" + +from __future__ import annotations + +import os + +import pytest + +from e2e_config import CHEAP_ANTHROPIC_MODEL, LINEAR_MCP_URL, LINEAR_STORAGE_STATE, unique_marker +from e2e_http import AuthHeaders +from lifecycle import ResourceManager +from models import ChatBody, ChatMessage, KeyGenerateBody, McpChatTool, McpServerCreateBody, ObjectPermission +from proxy_client import ProxyClient + +pytest.importorskip("mcp", reason="mcp SDK not installed; run `uv sync --inexact --group e2e-dev`") +pytest.importorskip( + "playwright.async_api", + reason="playwright not installed; run `uv pip install playwright` and `playwright install chromium`", +) + +from oauth_chat_client import ChatMcpClient, build_chat_client # noqa: E402 # imports follow the importorskip guards + +pytestmark = [ + pytest.mark.e2e, + pytest.mark.skipif( + not LINEAR_STORAGE_STATE or not os.path.exists(LINEAR_STORAGE_STATE), + reason="set E2E_LINEAR_STORAGE_STATE to a Linear session captured via mcp/linear_session_capture.py", + ), +] + +# Pinned from a live dance during verification (never guessed); the gateway +# prefixes every upstream tool name with the server alias. list_teams is a +# read-only Linear tool that takes no arguments and returns the caller's teams. +LINEAR_READONLY_TOOL = "list_teams" +LINEAR_PROMPT = "Use the list_teams tool to list my Linear teams, then reply with the name of one of them." + + +@pytest.fixture(scope="session") +def chat_client(proxy: ProxyClient) -> ChatMcpClient: + return build_chat_client(proxy) + + +class TestMcpChatCompletionOauth: + """A scoped internal-user key on a real Linear authorization_code server, + used through /chat/completions once per ingress header: the gateway pulls + the user's stored upstream token, lists and executes Linear's tools during + the completion, and returns the answer.""" + + @pytest.mark.covers("mcp.list_tools.oauth.succeeds") + @pytest.mark.covers("mcp.call_tool.oauth.succeeds") + def test_chat_completion_uses_linear_with_x_litellm_api_key_header( + self, chat_client: ChatMcpClient, resources: ResourceManager + ) -> None: + marker = unique_marker() + alias = f"e2elinear{marker}" + created = chat_client.create_server( + McpServerCreateBody( + alias=alias, + url=LINEAR_MCP_URL, + allow_all_keys=False, + auth_type="oauth2", + oauth2_flow="authorization_code", + ) + ) + resources.defer(lambda: chat_client.delete_server(created.server_id)) + + stored = chat_client.server_info(created.server_id) + assert stored.auth_type == "oauth2" + assert stored.oauth2_flow == "authorization_code" + assert stored.allow_all_keys is False + + key = chat_client.proxy.generate_key( + KeyGenerateBody( + user_id="e2e-test-user", + object_permission=ObjectPermission(mcp_servers=[created.server_id]), + ) + ) + resources.defer(lambda: chat_client.proxy.delete_key(key)) + + seeded = chat_client.seed_user_token(alias, key, LINEAR_STORAGE_STATE) + assert f"{alias}-{LINEAR_READONLY_TOOL}" in seeded, ( + f"the authorize dance listed {seeded}, expected it to include {alias}-{LINEAR_READONLY_TOOL}" + ) + + response = chat_client.chat_with_mcp( + AuthHeaders.model_validate({"x-litellm-api-key": f"Bearer {key}"}), + ChatBody( + model=CHEAP_ANTHROPIC_MODEL, + messages=[ChatMessage(role="user", content=LINEAR_PROMPT)], + tools=[ + McpChatTool( + server_url=f"litellm_proxy/mcp/{alias}", + server_label=alias, + require_approval="never", + ) + ], + ), + ) + + message = response.choices[0].message + assert message is not None and message.content, f"completion returned no answer: {response}" + meta = message.provider_specific_fields + assert meta is not None, f"no MCP metadata on the completion: {response}" + listed = {t.function.name for t in (meta.mcp_list_tools or []) if t.function} + assert f"{alias}-{LINEAR_READONLY_TOOL}" in listed, ( + f"the gateway listed {sorted(listed)}, expected the stored token to surface {alias}-{LINEAR_READONLY_TOOL}" + ) + results = [r for r in (meta.mcp_call_results or []) if r.name == f"{alias}-{LINEAR_READONLY_TOOL}"] + assert results and results[0].result, ( + f"Linear tool {alias}-{LINEAR_READONLY_TOOL} was not executed with a result: {meta.mcp_call_results}" + ) + + @pytest.mark.covers("mcp.list_tools.oauth.succeeds") + @pytest.mark.covers("mcp.call_tool.oauth.succeeds") + def test_chat_completion_uses_linear_with_authorization_bearer_header( + self, chat_client: ChatMcpClient, resources: ResourceManager + ) -> None: + marker = unique_marker() + alias = f"e2elinear{marker}" + created = chat_client.create_server( + McpServerCreateBody( + alias=alias, + url=LINEAR_MCP_URL, + allow_all_keys=False, + auth_type="oauth2", + oauth2_flow="authorization_code", + ) + ) + resources.defer(lambda: chat_client.delete_server(created.server_id)) + + stored = chat_client.server_info(created.server_id) + assert stored.auth_type == "oauth2" + assert stored.oauth2_flow == "authorization_code" + assert stored.allow_all_keys is False + + key = chat_client.proxy.generate_key( + KeyGenerateBody( + user_id="e2e-test-user", + object_permission=ObjectPermission(mcp_servers=[created.server_id]), + ) + ) + resources.defer(lambda: chat_client.proxy.delete_key(key)) + + seeded = chat_client.seed_user_token(alias, key, LINEAR_STORAGE_STATE) + assert f"{alias}-{LINEAR_READONLY_TOOL}" in seeded, ( + f"the authorize dance listed {seeded}, expected it to include {alias}-{LINEAR_READONLY_TOOL}" + ) + + response = chat_client.chat_with_mcp( + AuthHeaders.model_validate({"authorization": f"Bearer {key}"}), + ChatBody( + model=CHEAP_ANTHROPIC_MODEL, + messages=[ChatMessage(role="user", content=LINEAR_PROMPT)], + tools=[ + McpChatTool( + server_url=f"litellm_proxy/mcp/{alias}", + server_label=alias, + require_approval="never", + ) + ], + ), + ) + + message = response.choices[0].message + assert message is not None and message.content, f"completion returned no answer: {response}" + meta = message.provider_specific_fields + assert meta is not None, f"no MCP metadata on the completion: {response}" + listed = {t.function.name for t in (meta.mcp_list_tools or []) if t.function} + assert f"{alias}-{LINEAR_READONLY_TOOL}" in listed, ( + f"the gateway listed {sorted(listed)}, expected the stored token to surface {alias}-{LINEAR_READONLY_TOOL}" + ) + results = [r for r in (meta.mcp_call_results or []) if r.name == f"{alias}-{LINEAR_READONLY_TOOL}"] + assert results and results[0].result, ( + f"Linear tool {alias}-{LINEAR_READONLY_TOOL} was not executed with a result: {meta.mcp_call_results}" + ) diff --git a/tests/e2e/models.py b/tests/e2e/models.py index 8b6c454fa1e..d21920cf848 100644 --- a/tests/e2e/models.py +++ b/tests/e2e/models.py @@ -6,6 +6,7 @@ response validates without mirroring every proxy field. No untyped dicts. from __future__ import annotations +from collections.abc import Sequence from datetime import datetime from typing import Literal @@ -115,6 +116,18 @@ class KeyInfoResponse(BaseModel): # ---------- customers ---------- +class CustomerNewBody(BaseModel): + user_id: str + + +class CustomerResponse(BaseModel): + user_id: str | None = None + + +class CustomerInfoParams(BaseModel): + end_user_id: str + + class CustomerDeleteBody(BaseModel): user_ids: list[str] @@ -126,9 +139,26 @@ class ChatMetadata(BaseModel): tags: list[str] | None = None +class ImageUrl(BaseModel): + url: str + + +class TextContentPart(BaseModel): + type: str = "text" + text: str + + +class ImageContentPart(BaseModel): + type: str = "image_url" + image_url: ImageUrl + + +ContentPart = TextContentPart | ImageContentPart + + class ChatMessage(BaseModel): role: str - content: str + content: str | list[ContentPart] class CacheControl(BaseModel): @@ -167,6 +197,19 @@ class ChatTool(BaseModel): function: ChatToolFunction +class McpChatTool(BaseModel): + """An MCP server attached to a chat completion (OpenAI `type: "mcp"` tool). + `server_url` selects the gateway-registered server by its alias suffix; with + `require_approval="never"` the gateway lists, calls, and feeds the server's + tools back to the model in one agentic turn.""" + + type: Literal["mcp"] = "mcp" + server_url: str + require_approval: str + server_label: str | None = None + allowed_tools: list[str] | None = None + + class ChatBody(BaseModel): model: str messages: list[ChatMessage] @@ -177,9 +220,10 @@ class ChatBody(BaseModel): reasoning_effort: str | None = None thinking: ThinkingParam | None = None service_tier: str | None = None - tools: list[ChatTool] | None = None + tools: Sequence[ChatTool | McpChatTool] | None = None tool_choice: str | None = None guardrails: list[str] | None = None + response_format: dict[str, object] | None = None class RouterSettingsOverride(BaseModel): @@ -203,9 +247,55 @@ class ReliabilityChatBody(ChatBody): router_settings_override: RouterSettingsOverride | None = None +class ToolCallFunction(BaseModel): + name: str | None = None + arguments: str | None = None + + +class ToolCall(BaseModel): + function: ToolCallFunction = ToolCallFunction() + + +class McpToolFunctionRef(BaseModel): + name: str + + +class McpListedTool(BaseModel): + """One entry of `mcp_list_tools`: a tool the gateway listed from the + attached MCP server and exposed to the model, in OpenAI function shape.""" + + function: McpToolFunctionRef | None = None + + +class McpToolCall(BaseModel): + """One entry of `mcp_tool_calls`: a tool the model asked the gateway to run.""" + + function: McpToolFunctionRef | None = None + + +class McpCallResult(BaseModel): + """One entry of `mcp_call_results`: what the gateway got back from executing + a tool upstream on the caller's behalf.""" + + name: str | None = None + result: str | None = None + + +class McpResponseMetadata(BaseModel): + """`choices[].message.provider_specific_fields` MCP section: which tools the + gateway listed from the attached server, which the model called, and their + results. Populated only when the completion drove an MCP server.""" + + mcp_list_tools: list[McpListedTool] | None = None + mcp_tool_calls: list[McpToolCall] | None = None + mcp_call_results: list[McpCallResult] | None = None + + class OutMessage(BaseModel): content: str | None = None reasoning_content: str | None = None + tool_calls: list[ToolCall] | None = None + provider_specific_fields: McpResponseMetadata | None = None class ChatChoice(BaseModel): @@ -216,6 +306,10 @@ class PromptTokensDetails(BaseModel): cached_tokens: int | None = None +class CompletionTokensDetails(BaseModel): + reasoning_tokens: int | None = None + + class Usage(BaseModel): prompt_tokens: int | None = None completion_tokens: int | None = None @@ -223,6 +317,7 @@ class Usage(BaseModel): cache_read_input_tokens: int | None = None cache_creation_input_tokens: int | None = None prompt_tokens_details: PromptTokensDetails | None = None + completion_tokens_details: CompletionTokensDetails | None = None class ChatResponse(BaseModel): @@ -310,6 +405,36 @@ class CountTokensResponse(BaseModel): input_tokens: int +# ---------- mcp servers ---------- + + +class McpServerCreateBody(BaseModel): + """POST /v1/mcp/server. For a gateway-managed OAuth server, `auth_type` is + `oauth2` and `oauth2_flow` is `authorization_code`; the upstream endpoints + are discovered and registered via DCR when left unset. `allow_all_keys` + false scopes the server to keys granted it through object_permission.""" + + alias: str + url: str + transport: str = "http" + allow_all_keys: bool = True + auth_type: str | None = None + oauth2_flow: Literal["client_credentials", "authorization_code"] | None = None + authorization_url: str | None = None + token_url: str | None = None + + +class McpServerInfo(BaseModel): + """Response of POST /v1/mcp/server and GET /v1/mcp/server/{server_id}.""" + + server_id: str + alias: str | None = None + url: str | None = None + auth_type: str | None = None + oauth2_flow: str | None = None + allow_all_keys: bool | None = None + + class EmbedBody(BaseModel): model: str input: str diff --git a/tests/e2e/pytest.ini b/tests/e2e/pytest.ini index e9611df139b..2998a4b83c6 100644 --- a/tests/e2e/pytest.ini +++ b/tests/e2e/pytest.ini @@ -6,3 +6,4 @@ addopts = --strict-markers --strict-config markers = e2e: live test that requires a running proxy and real provider keys load: heavy throughput/load test; collected last so it never perturbs latency-sensitive suites + weekly: real-provider anomaly load test that spends real money; deselected unless E2E_WEEKLY_ANOMALY is set diff --git a/tests/e2e/quota_management/budgets/test_budget_enforcement_e2e.py b/tests/e2e/quota_management/budgets/test_budget_enforcement_e2e.py index c4ad0c38f31..8b93afb4752 100644 --- a/tests/e2e/quota_management/budgets/test_budget_enforcement_e2e.py +++ b/tests/e2e/quota_management/budgets/test_budget_enforcement_e2e.py @@ -1,28 +1,36 @@ """Live e2e: a tiny max_budget on an entity actually blocks requests. -Each entity is an E2ECase (lifecycle.E2ECase) driven by run_case: init() creates -the budgeted entity + a key, run() drives spend until a `budget_exceeded` block, -teardown() deletes everything init() created (always runs, even on failure/skip). -Covers the entities with no prior live coverage - internal user, end-user, -organization, team member - plus key and team. See BUDGET_TEST_COVERAGE_MATRIX.md. +One test per budget level (key, team, internal user, end-user, organization, +team member): put the tiny cap on that level, drive spend until a +`budget_exceeded` block, and where a cap could be confused with a neighbor, +prove isolation with an uncapped control key that must keep serving. The +capped-key sweep proves the key's own max_budget blocks across mint shapes +(personal, team, team-member) with roomy surroundings, so the key-level cap is +provably the blocker no matter who the key was minted to. A non-budget error fails hard (never a skip); if calls never get blocked, budget enforcement is broken -> fail. """ import time -from dataclasses import dataclass, field -from typing import Callable, List, Type import pytest from budget_client import BudgetClient, is_budget_block from e2e_config import unique_marker from e2e_http import StreamingResponse, require_successful_call -from lifecycle import run_case +from lifecycle import ResourceManager pytestmark = pytest.mark.e2e +TINY_CAP = 3e-6 +ROOMY_CAP = 100.0 + + +def _chat(client: BudgetClient, key: str, *, user: str | None = None) -> StreamingResponse: + return client.chat(key, "claude-haiku-4-5", f"spend {unique_marker()}", max_tokens=16, user=user) + + def _assert_budget_blocks(client: BudgetClient, key: str, *, user: str = "") -> StreamingResponse: """Send paid calls until the entity's budget blocks one; return the blocked response so callers can assert on its shape. Key/user/org/member block within @@ -30,13 +38,7 @@ def _assert_budget_blocks(client: BudgetClient, key: str, *, user: str = "") -> enforces off table spend that lands on the batch write, so it takes a few more. A non-budget error fails hard (never a skip).""" for _ in range(40): - result = client.chat( - key, - "claude-haiku-4-5", - f"spend {unique_marker()}", - max_tokens=16, - user=user or None, - ) + result = _chat(client, key, user=user or None) if is_budget_block(result): return result require_successful_call(result) @@ -44,225 +46,154 @@ def _assert_budget_blocks(client: BudgetClient, key: str, *, user: str = "") -> pytest.fail("budget never enforced within the call budget") -@dataclass -class _BudgetCase: - """Base E2ECase: a key under some budgeted entity must get blocked. - - Subclasses set up the budgeted entity in init() and register every created id - in `_undo` (run LIFO in teardown so a key is deleted before its team/org). - """ - - client: BudgetClient - key: str = "" - _undo: List[Callable[[], None]] = field( - default_factory=list - ) # mutable-ok: per-case teardown registry - - def init(self) -> None: - raise NotImplementedError - - def run(self) -> None: - _assert_budget_blocks(self.client, self.key) - - def teardown(self) -> None: - for undo in reversed(self._undo): - undo() +def _assert_blocked_429(client: BudgetClient, key: str) -> StreamingResponse: + blocked = _assert_budget_blocks(client, key) + assert blocked.status_code == 429, ( + f"budget refusal must be 429, got {blocked.status_code}: {blocked.body[:200]}" + ) + return blocked -class KeyBudgetCase(_BudgetCase): - """A bare key (no team_id / user_id) carrying its own max_budget, so only the - key-level budget can be the thing that blocks. The refusal must be a 429 - budget_exceeded; any other error already fails via _assert_budget_blocks.""" +class TestBudgetBlocksPerLevel: + @pytest.mark.covers("quota_management.budget.key.blocks_over_limit") + def test_bare_key_blocks_over_its_own_budget(self, client: BudgetClient, resources: ResourceManager) -> None: + key = client.generate_key(max_budget=TINY_CAP) + resources.defer(lambda: client.delete_key(key)) - def init(self) -> None: - self.key = self.client.generate_key(max_budget=3e-6) - self._undo.append(lambda: self.client.delete_key(self.key)) + _assert_blocked_429(client, key) - def run(self) -> None: - blocked = _assert_budget_blocks(self.client, self.key) - assert blocked.status_code == 429, ( - f"budget refusal must be 429, got {blocked.status_code}: {blocked.body[:200]}" - ) + @pytest.mark.covers("quota_management.budget.team.blocks_over_limit") + def test_team_budget_blocks_every_team_key(self, client: BudgetClient, resources: ResourceManager) -> None: + team_id = client.create_team(alias=f"e2e-budget-team-{unique_marker()}", max_budget=TINY_CAP) + resources.defer(lambda: client.delete_team(team_id)) + spender_key = client.generate_key(team_id=team_id) + resources.defer(lambda: client.delete_key(spender_key)) + sibling_key = client.generate_key(team_id=team_id) + resources.defer(lambda: client.delete_key(sibling_key)) - -class TeamBudgetCase(_BudgetCase): - """An admin caps a whole team: two keys under a tiny-budget team, neither with - a key-level budget. Key A is driven until the team cap blocks it; key B's very - first call must then be refused too, proving the cap sits on the team, not the - key that spent. Both refusals must be 429 budget_exceeded.""" - - def init(self) -> None: - team_id = self.client.create_team( - alias=f"e2e-budget-team-{unique_marker()}", max_budget=3e-6 - ) - self._undo.append(lambda: self.client.delete_team(team_id)) - self.key = self.client.generate_key(team_id=team_id) - self._undo.append(lambda: self.client.delete_key(self.key)) - self._sibling_key = self.client.generate_key(team_id=team_id) - self._undo.append(lambda: self.client.delete_key(self._sibling_key)) - - def run(self) -> None: - blocked = _assert_budget_blocks(self.client, self.key) - assert blocked.status_code == 429, ( - f"budget refusal must be 429, got {blocked.status_code}: {blocked.body[:200]}" - ) - sibling = self.client.chat( - self._sibling_key, - "claude-haiku-4-5", - f"spend {unique_marker()}", - max_tokens=16, - ) + _assert_blocked_429(client, spender_key) + sibling = _chat(client, sibling_key) assert is_budget_block(sibling) and sibling.status_code == 429, ( f"a sibling key on the capped team must get the same 429 budget_exceeded, " f"got {sibling.status_code}: {sibling.body[:200]}" ) + @pytest.mark.covers("quota_management.budget.internal_user.blocks_over_limit") + def test_user_budget_enforced_across_all_their_keys( + self, client: BudgetClient, resources: ResourceManager + ) -> None: + user_id = client.create_user(max_budget=TINY_CAP) + resources.defer(lambda: client.delete_user(user_id)) + first_key = client.generate_key(user_id=user_id) + resources.defer(lambda: client.delete_key(first_key)) + second_key = client.generate_key(user_id=user_id) + resources.defer(lambda: client.delete_key(second_key)) + team_id = client.create_team(alias=f"e2e-budget-team-{unique_marker()}") + resources.defer(lambda: client.delete_team(team_id)) + client.add_team_member(team_id, user_id) + team_key = client.generate_key(team_id=team_id, user_id=user_id) + resources.defer(lambda: client.delete_key(team_key)) -class InternalUserBudgetCase(_BudgetCase): - """A user's max_budget follows the person, not the key. The capped user holds - two personal keys (no team, no key budgets) plus a team-member key on an - uncapped team; once the first personal key is refused, the other two must be - refused as well - a second key is not a fresh allowance, and since #32005 the - user budget draws down team keys too. All refusals must be 429 budget_exceeded.""" - - def init(self) -> None: - user_id = self.client.create_user(max_budget=3e-6) - self._undo.append(lambda: self.client.delete_user(user_id)) - self.key = self.client.generate_key(user_id=user_id) - self._undo.append(lambda: self.client.delete_key(self.key)) - self._second_key = self.client.generate_key(user_id=user_id) - self._undo.append(lambda: self.client.delete_key(self._second_key)) - team_id = self.client.create_team(alias=f"e2e-budget-team-{unique_marker()}") - self._undo.append(lambda: self.client.delete_team(team_id)) - self.client.add_team_member(team_id, user_id) - self._team_key = self.client.generate_key(team_id=team_id, user_id=user_id) - self._undo.append(lambda: self.client.delete_key(self._team_key)) - - def run(self) -> None: - blocked = _assert_budget_blocks(self.client, self.key) - assert blocked.status_code == 429, ( - f"budget refusal must be 429, got {blocked.status_code}: {blocked.body[:200]}" - ) - for label, key in (("second personal key", self._second_key), ("team-member key", self._team_key)): - result = self.client.chat(key, "claude-haiku-4-5", f"spend {unique_marker()}", max_tokens=16) + _assert_blocked_429(client, first_key) + for label, key in (("second personal key", second_key), ("team-member key", team_key)): + result = _chat(client, key) assert is_budget_block(result) and result.status_code == 429, ( f"the {label} of a user over budget must get the same 429 budget_exceeded, " f"got {result.status_code}: {result.body[:200]}" ) - -class EndUserBudgetCase(_BudgetCase): - def init(self) -> None: + @pytest.mark.covers("quota_management.budget.end_user.blocks_over_limit") + def test_end_user_budget_blocks_attributed_calls( + self, client: BudgetClient, resources: ResourceManager + ) -> None: customer = f"e2e-budget-cust-{unique_marker()}" - self.client.create_customer(customer, max_budget=3e-6) - self._undo.append(lambda: self.client.delete_customers([customer])) - self.key = self.client.generate_key(models=["claude-haiku-4-5"]) - self._undo.append(lambda: self.client.delete_key(self.key)) - self._customer = customer + client.create_customer(customer, max_budget=TINY_CAP) + resources.defer(lambda: client.delete_customers([customer])) + key = client.generate_key(models=["claude-haiku-4-5"]) + resources.defer(lambda: client.delete_key(key)) - def run(self) -> None: - _assert_budget_blocks(self.client, self.key, user=self._customer) + _assert_budget_blocks(client, key, user=customer) + @pytest.mark.covers("quota_management.budget.organization.blocks_over_limit") + def test_org_budget_blocks_keys_under_it(self, client: BudgetClient, resources: ResourceManager) -> None: + org_id = client.create_org(max_budget=TINY_CAP, alias=f"e2e-budget-org-{unique_marker()}") + resources.defer(lambda: client.delete_org(org_id)) + team_id = client.create_team(alias=f"e2e-budget-team-{unique_marker()}", organization_id=org_id) + resources.defer(lambda: client.delete_team(team_id)) + key = client.generate_key(team_id=team_id) + resources.defer(lambda: client.delete_key(key)) -class OrganizationBudgetCase(_BudgetCase): - """Org carries the tiny budget; the team under it and the key carry none, so - the org is the only entity that can block (the historically weak link). The - refusal must be a 429 budget_exceeded that names the org as the blocker.""" - - def init(self) -> None: - self._org_id = self.client.create_org( - max_budget=3e-6, alias=f"e2e-budget-org-{unique_marker()}" - ) - self._undo.append(lambda: self.client.delete_org(self._org_id)) - team_id = self.client.create_team( - alias=f"e2e-budget-team-{unique_marker()}", organization_id=self._org_id - ) - self._undo.append(lambda: self.client.delete_team(team_id)) - self.key = self.client.generate_key(team_id=team_id) - self._undo.append(lambda: self.client.delete_key(self.key)) - - def run(self) -> None: - blocked = _assert_budget_blocks(self.client, self.key) - assert blocked.status_code == 429, ( - f"budget refusal must be 429, got {blocked.status_code}: {blocked.body[:200]}" - ) - assert f"Organization={self._org_id}" in blocked.body, ( + blocked = _assert_blocked_429(client, key) + assert f"Organization={org_id}" in blocked.body, ( f"refusal must name the org as the blocker, got: {blocked.body[:200]}" ) + @pytest.mark.covers("quota_management.budget.team_member.blocks_over_limit") + def test_member_budget_blocks_without_touching_teammates( + self, client: BudgetClient, resources: ResourceManager + ) -> None: + team_id = client.create_team(alias=f"e2e-budget-team-{unique_marker()}", max_budget=ROOMY_CAP) + resources.defer(lambda: client.delete_team(team_id)) + member_id = client.create_user(max_budget=ROOMY_CAP) + resources.defer(lambda: client.delete_user(member_id)) + client.add_team_member(team_id, member_id, max_budget_in_team=TINY_CAP) + member_key = client.generate_key(team_id=team_id, user_id=member_id) + resources.defer(lambda: client.delete_key(member_key)) + teammate_id = client.create_user(max_budget=ROOMY_CAP) + resources.defer(lambda: client.delete_user(teammate_id)) + client.add_team_member(team_id, teammate_id) + teammate_key = client.generate_key(team_id=team_id, user_id=teammate_id) + resources.defer(lambda: client.delete_key(teammate_key)) -class TeamMemberBudgetCase(_BudgetCase): - """Member A's per-team budget is tiny while the team and both members' user - budgets are roomy (100.0), so the only cap that can trip is A's: a block - proves member-level enforcement and must be a 429 budget_exceeded. Teammate - B, uncapped on the same team, must keep serving after A is cut off, proving - the member cap does not leak onto the team or its members.""" - - def init(self) -> None: - self._team_id = self.client.create_team( - alias=f"e2e-budget-team-{unique_marker()}", max_budget=100.0 - ) - self._undo.append(lambda: self.client.delete_team(self._team_id)) - self._member_id = self.client.create_user(max_budget=100.0) - self._undo.append(lambda: self.client.delete_user(self._member_id)) - self.client.add_team_member(self._team_id, self._member_id, max_budget_in_team=3e-6) - self.key = self.client.generate_key(team_id=self._team_id, user_id=self._member_id) - self._undo.append(lambda: self.client.delete_key(self.key)) - teammate_id = self.client.create_user(max_budget=100.0) - self._undo.append(lambda: self.client.delete_user(teammate_id)) - self.client.add_team_member(self._team_id, teammate_id) - self._teammate_key = self.client.generate_key(team_id=self._team_id, user_id=teammate_id) - self._undo.append(lambda: self.client.delete_key(self._teammate_key)) - - def run(self) -> None: - blocked = _assert_budget_blocks(self.client, self.key) - assert blocked.status_code == 429, ( - f"budget refusal must be 429, got {blocked.status_code}: {blocked.body[:200]}" - ) - teammate = self.client.chat( - self._teammate_key, - "claude-haiku-4-5", - f"spend {unique_marker()}", - max_tokens=16, - ) - require_successful_call(teammate) + _assert_blocked_429(client, member_key) + require_successful_call(_chat(client, teammate_key)) -def _case_id(case_cls: Type[_BudgetCase]) -> str: - return case_cls.__name__ +class TestKeyBudgetBlocksAcrossKeyKinds: + """The tiny max_budget sits on the key itself while every budget around it + (user / team / membership) is roomy, so only the key-level cap can block; the + uncapped control key minted to the same surroundings must keep serving after + the capped key is refused, proving nothing around the key was the blocker.""" + @pytest.mark.covers("quota_management.budget.key.blocks_over_limit") + def test_personal_key_blocks_over_its_own_budget( + self, client: BudgetClient, resources: ResourceManager + ) -> None: + user_id = client.create_user(max_budget=ROOMY_CAP) + resources.defer(lambda: client.delete_user(user_id)) + capped_key = client.generate_key(user_id=user_id, max_budget=TINY_CAP) + resources.defer(lambda: client.delete_key(capped_key)) + control_key = client.generate_key(user_id=user_id) + resources.defer(lambda: client.delete_key(control_key)) -@pytest.mark.parametrize( - "case_cls", - [ - pytest.param( - KeyBudgetCase, - marks=pytest.mark.covers("quota_management.budget.key.blocks_over_limit"), - ), - pytest.param( - TeamBudgetCase, - marks=pytest.mark.covers("quota_management.budget.team.blocks_over_limit"), - ), - pytest.param( - InternalUserBudgetCase, - marks=pytest.mark.covers("quota_management.budget.internal_user.blocks_over_limit"), - ), - pytest.param( - EndUserBudgetCase, - marks=pytest.mark.covers("quota_management.budget.end_user.blocks_over_limit"), - ), - pytest.param( - OrganizationBudgetCase, - marks=pytest.mark.covers("quota_management.budget.organization.blocks_over_limit"), - ), - pytest.param( - TeamMemberBudgetCase, - marks=pytest.mark.covers("quota_management.budget.team_member.blocks_over_limit"), - ), - ], - ids=_case_id, -) -def test_budget_enforcement( - client: BudgetClient, case_cls: Type[_BudgetCase] -) -> None: - run_case(case_cls(client)) + _assert_blocked_429(client, capped_key) + require_successful_call(_chat(client, control_key)) + + @pytest.mark.covers("quota_management.budget.key.blocks_over_limit") + def test_team_key_blocks_over_its_own_budget(self, client: BudgetClient, resources: ResourceManager) -> None: + team_id = client.create_team(alias=f"e2e-key-cap-team-{unique_marker()}", max_budget=ROOMY_CAP) + resources.defer(lambda: client.delete_team(team_id)) + capped_key = client.generate_key(team_id=team_id, max_budget=TINY_CAP) + resources.defer(lambda: client.delete_key(capped_key)) + control_key = client.generate_key(team_id=team_id) + resources.defer(lambda: client.delete_key(control_key)) + + _assert_blocked_429(client, capped_key) + require_successful_call(_chat(client, control_key)) + + @pytest.mark.covers("quota_management.budget.key.blocks_over_limit") + def test_team_member_key_blocks_over_its_own_budget( + self, client: BudgetClient, resources: ResourceManager + ) -> None: + team_id = client.create_team(alias=f"e2e-key-cap-team-{unique_marker()}", max_budget=ROOMY_CAP) + resources.defer(lambda: client.delete_team(team_id)) + member_id = client.create_user(max_budget=ROOMY_CAP) + resources.defer(lambda: client.delete_user(member_id)) + client.add_team_member(team_id, member_id, max_budget_in_team=ROOMY_CAP) + capped_key = client.generate_key(team_id=team_id, user_id=member_id, max_budget=TINY_CAP) + resources.defer(lambda: client.delete_key(capped_key)) + control_key = client.generate_key(team_id=team_id, user_id=member_id) + resources.defer(lambda: client.delete_key(control_key)) + + _assert_blocked_429(client, capped_key) + require_successful_call(_chat(client, control_key)) diff --git a/tests/e2e/quota_management/budgets/test_budget_reset_e2e.py b/tests/e2e/quota_management/budgets/test_budget_reset_e2e.py index b57ad23ebf5..793b22a47c7 100644 --- a/tests/e2e/quota_management/budgets/test_budget_reset_e2e.py +++ b/tests/e2e/quota_management/budgets/test_budget_reset_e2e.py @@ -12,6 +12,7 @@ from lifecycle import ResourceManager pytestmark = pytest.mark.e2e TINY_CAP = 3e-6 +ROOMY_CAP = 100.0 WINDOW = "30s" RESET_DEADLINE_SECONDS = 150 @@ -46,7 +47,7 @@ def _poll_until_serves_again(client: BudgetClient, key: str) -> None: pytest.fail(f"budget never reset within {RESET_DEADLINE_SECONDS}s") -class TestBudgetResetDiagonal: +class TestBudgetResetPerLevel: @pytest.mark.covers("quota_management.budget.key.resets_after_window") def test_bare_key_budget_resets_after_window(self, client: BudgetClient, resources: ResourceManager) -> None: key = client.generate_key(max_budget=TINY_CAP, budget_duration=WINDOW) @@ -115,3 +116,42 @@ class TestBudgetResetDiagonal: _drive_to_block(client, key) _poll_until_serves_again(client, key) + + +class TestKeyBudgetResetAcrossKeyKinds: + """The tiny max_budget and its 30s window sit on the key itself while the user, + team, and membership around it are roomy (100.0), so the key's own budget is + the only thing that can block and the only thing that has to reset.""" + + @pytest.mark.covers("quota_management.budget.key.resets_after_window") + def test_personal_key_resets_after_window(self, client: BudgetClient, resources: ResourceManager) -> None: + user_id = client.create_user(max_budget=ROOMY_CAP) + resources.defer(lambda: client.delete_user(user_id)) + key = client.generate_key(user_id=user_id, max_budget=TINY_CAP, budget_duration=WINDOW) + resources.defer(lambda: client.delete_key(key)) + + _drive_to_block(client, key) + _poll_until_serves_again(client, key) + + @pytest.mark.covers("quota_management.budget.key.resets_after_window") + def test_team_key_resets_after_window(self, client: BudgetClient, resources: ResourceManager) -> None: + team_id = client.create_team(alias=f"e2e-key-reset-team-{unique_marker()}", max_budget=ROOMY_CAP) + resources.defer(lambda: client.delete_team(team_id)) + key = client.generate_key(team_id=team_id, max_budget=TINY_CAP, budget_duration=WINDOW) + resources.defer(lambda: client.delete_key(key)) + + _drive_to_block(client, key) + _poll_until_serves_again(client, key) + + @pytest.mark.covers("quota_management.budget.key.resets_after_window") + def test_team_member_key_resets_after_window(self, client: BudgetClient, resources: ResourceManager) -> None: + team_id = client.create_team(alias=f"e2e-key-reset-team-{unique_marker()}", max_budget=ROOMY_CAP) + resources.defer(lambda: client.delete_team(team_id)) + member_id = client.create_user(max_budget=ROOMY_CAP) + resources.defer(lambda: client.delete_user(member_id)) + client.add_team_member(team_id, member_id, max_budget_in_team=ROOMY_CAP) + key = client.generate_key(team_id=team_id, user_id=member_id, max_budget=TINY_CAP, budget_duration=WINDOW) + resources.defer(lambda: client.delete_key(key)) + + _drive_to_block(client, key) + _poll_until_serves_again(client, key) diff --git a/tests/e2e/quota_management/ratelimit/test_redis_backed_ratelimit_e2e.py b/tests/e2e/quota_management/ratelimit/test_redis_backed_ratelimit_e2e.py index ed6f0ce3b2c..a88f0ca546a 100644 --- a/tests/e2e/quota_management/ratelimit/test_redis_backed_ratelimit_e2e.py +++ b/tests/e2e/quota_management/ratelimit/test_redis_backed_ratelimit_e2e.py @@ -11,7 +11,7 @@ import socket import pytest -from e2e_config import require_env, unique_marker +from e2e_config import unique_marker from e2e_http import require_successful_call from lifecycle import ResourceManager from models import KeyGenerateBody, LiteLLMParamsBody @@ -23,7 +23,7 @@ BACKEND = "anthropic/claude-haiku-4-5-20251001" def _require_redis_reachable() -> None: - (host,) = require_env("REDIS_HOST") + host = os.environ["REDIS_HOST"] port = int((os.environ.get("REDIS_PORT") or "6379").strip() or "6379") try: with socket.create_connection((host, port), timeout=3): diff --git a/tests/e2e/quota_management/ratelimit/test_redis_circuit_breaker_e2e.py b/tests/e2e/quota_management/ratelimit/test_redis_circuit_breaker_e2e.py index b509ae000f5..3e1bc662470 100644 --- a/tests/e2e/quota_management/ratelimit/test_redis_circuit_breaker_e2e.py +++ b/tests/e2e/quota_management/ratelimit/test_redis_circuit_breaker_e2e.py @@ -13,7 +13,7 @@ from concurrent.futures import ThreadPoolExecutor, as_completed import pytest -from e2e_config import require_env, unique_marker +from e2e_config import unique_marker from e2e_http import require_successful_call from lifecycle import ResourceManager from models import KeyGenerateBody, LiteLLMParamsBody @@ -28,7 +28,7 @@ RECOVERY_TIMEOUT = float( def _require_redis() -> None: - (host,) = require_env("REDIS_HOST") + host = os.environ["REDIS_HOST"] port = int((os.environ.get("REDIS_PORT") or "6379").strip() or "6379") try: with socket.create_connection((host, port), timeout=3): diff --git a/tests/e2e/quota_management/spend_tracking/test_spend_tracking_e2e.py b/tests/e2e/quota_management/spend_tracking/test_spend_tracking_e2e.py index d43d8e94898..b6e33627a5e 100644 --- a/tests/e2e/quota_management/spend_tracking/test_spend_tracking_e2e.py +++ b/tests/e2e/quota_management/spend_tracking/test_spend_tracking_e2e.py @@ -23,7 +23,7 @@ import pytest from e2e_http import Result, Success from lifecycle import ResourceManager -from models import ChatResponse, SpendLogs, SpendLogsParams +from models import ChatResponse, LiteLLMParamsBody, SpendLogs, SpendLogsParams from spend_e2e_client import SpendClient, SpendLogRow, is_ok, unique_marker, unwrap pytestmark = pytest.mark.e2e @@ -194,6 +194,7 @@ def test_streaming_messages_via_responses_bridge_tracks_spend( @pytest.mark.covers("quota_management.spend_tracking.embeddings.logs_cost") +@pytest.mark.covers("llm.embeddings.openai.basic.nonstream.cost_logged") def test_embedding_writes_nonzero_spend_row( client: SpendClient, scoped_key: str ) -> None: @@ -231,14 +232,12 @@ def test_cache_hit_is_zero_cost_and_suffixed( rows = client.poll_logs_for_key( scoped_key, predicate=lambda rs: any(r.cache_hit == "True" for r in rs) ) - cache_rows = [r for r in rows if r.cache_hit == "True"] - if not cache_rows: - pytest.skip( - "no cache-hit row observed; caching may be disabled on this proxy. " - f"rows seen: {_summarize(rows)}" - ) - - cache_row = cache_rows[0] + cache_row = _require_row( + rows, + lambda r: r.cache_hit == "True", + "with cache_hit=True (caching is enabled on the e2e proxy, so an identical " + "repeat call must hit the cache)", + ) assert ( cache_row.spend or 0 ) == 0.0, f"cache hit was charged (double-charge regression): {_summarize(rows)}" @@ -503,22 +502,27 @@ def test_each_model_on_a_shared_key_gets_its_own_row( @pytest.mark.covers("quota_management.spend_tracking.failure.writes_failure_row") def test_failure_call_writes_failure_status_row( - client: SpendClient, scoped_key: str + client: SpendClient, resources: ResourceManager, scoped_key: str ) -> None: - result = client.chat(scoped_key, "gemini-2.5-flash", "", max_tokens=1) - if is_ok(result): - pytest.skip("call unexpectedly succeeded; could not induce a failure row") + model = f"e2e-spend-failure-{unique_marker()}" + model_id = client.proxy.create_model( + model, + LiteLLMParamsBody(model="openai/gpt-5.5", api_key="sk-invalid-e2e-failure-row"), + ) + resources.defer(lambda: client.proxy.delete_model(model_id)) + + result = client.chat(scoped_key, model, f"trigger failure {unique_marker()}", max_tokens=1) + assert not is_ok(result), ( + f"a call to a deployment with an invalid upstream key must fail, not succeed: {result}" + ) rows = client.poll_logs_for_key( scoped_key, predicate=lambda rs: any(r.status == "failure" for r in rs) ) - failure_rows = [r for r in rows if r.status == "failure"] - if not failure_rows: - pytest.skip( - "no failure-status row was logged for the rejected call; " - "failure logging is environment-specific" - ) - assert (failure_rows[0].spend or 0) == 0.0, "failed call must not be charged" + failure_row = _require_row( + rows, lambda r: r.status == "failure", "with status=failure for the rejected call" + ) + assert (failure_row.spend or 0) == 0.0, "failed call must not be charged" @pytest.mark.covers("quota_management.spend_tracking.spend_calculate.returns_cost") diff --git a/tests/e2e/transport.py b/tests/e2e/transport.py index f7061b03a46..da4252e550e 100644 --- a/tests/e2e/transport.py +++ b/tests/e2e/transport.py @@ -16,7 +16,7 @@ import e2e_http from e2e_http import ( URL, AuthHeaders, - FileUploadForm, + BinaryStream, ProbeResult, Result, StreamingResponse, @@ -32,6 +32,15 @@ class Transport(Protocol): self, path: str, *, headers: BaseModel, json: BaseModel ) -> StreamingResponse: ... + def stream_binary( + self, + path: str, + *, + headers: BaseModel, + json: BaseModel, + chunk_size: int = 8192, + ) -> BinaryStream: ... + def send( self, path: str, @@ -76,9 +85,10 @@ class Transport(Protocol): path: str, *, headers: BaseModel, - form: FileUploadForm, + form: BaseModel, filename: str, content: bytes, + file_content_type: str = "application/jsonl", params: BaseModel | None = None, response_type: type[R], ) -> Result[R]: ... @@ -181,6 +191,22 @@ class HttpTransport: self._url(path), headers=headers, json=json, timeout=self.request_timeout ) + def stream_binary( + self, + path: str, + *, + headers: BaseModel, + json: BaseModel, + chunk_size: int = 8192, + ) -> BinaryStream: + return e2e_http.stream_binary( + self._url(path), + headers=headers, + json=json, + chunk_size=chunk_size, + timeout=self.request_timeout, + ) + def send( self, path: str, @@ -212,9 +238,10 @@ class HttpTransport: path: str, *, headers: BaseModel, - form: FileUploadForm, + form: BaseModel, filename: str, content: bytes, + file_content_type: str = "application/jsonl", params: BaseModel | None = None, response_type: type[R], ) -> Result[R]: @@ -224,6 +251,7 @@ class HttpTransport: form=form, filename=filename, content=content, + file_content_type=file_content_type, params=params, response_type=response_type, timeout=self.request_timeout, @@ -346,6 +374,18 @@ class SplitTransport: ) -> StreamingResponse: return self._route(path).stream(path, headers=headers, json=json) + def stream_binary( + self, + path: str, + *, + headers: BaseModel, + json: BaseModel, + chunk_size: int = 8192, + ) -> BinaryStream: + return self._route(path).stream_binary( + path, headers=headers, json=json, chunk_size=chunk_size + ) + def send( self, path: str, @@ -367,9 +407,10 @@ class SplitTransport: path: str, *, headers: BaseModel, - form: FileUploadForm, + form: BaseModel, filename: str, content: bytes, + file_content_type: str = "application/jsonl", params: BaseModel | None = None, response_type: type[R], ) -> Result[R]: @@ -379,6 +420,7 @@ class SplitTransport: form=form, filename=filename, content=content, + file_content_type=file_content_type, params=params, response_type=response_type, ) diff --git a/ui/litellm-dashboard/e2e_tests/constants.ts b/tests/e2e/ui/constants.ts similarity index 100% rename from ui/litellm-dashboard/e2e_tests/constants.ts rename to tests/e2e/ui/constants.ts diff --git a/ui/litellm-dashboard/e2e_tests/fixtures/config.yml b/tests/e2e/ui/fixtures/config.yml similarity index 100% rename from ui/litellm-dashboard/e2e_tests/fixtures/config.yml rename to tests/e2e/ui/fixtures/config.yml diff --git a/ui/litellm-dashboard/e2e_tests/fixtures/menuMappings.ts b/tests/e2e/ui/fixtures/menuMappings.ts similarity index 100% rename from ui/litellm-dashboard/e2e_tests/fixtures/menuMappings.ts rename to tests/e2e/ui/fixtures/menuMappings.ts diff --git a/ui/litellm-dashboard/e2e_tests/fixtures/migratedPages.ts b/tests/e2e/ui/fixtures/migratedPages.ts similarity index 100% rename from ui/litellm-dashboard/e2e_tests/fixtures/migratedPages.ts rename to tests/e2e/ui/fixtures/migratedPages.ts diff --git a/ui/litellm-dashboard/e2e_tests/fixtures/mock_llm_server/server.py b/tests/e2e/ui/fixtures/mock_llm_server/server.py similarity index 100% rename from ui/litellm-dashboard/e2e_tests/fixtures/mock_llm_server/server.py rename to tests/e2e/ui/fixtures/mock_llm_server/server.py diff --git a/ui/litellm-dashboard/e2e_tests/fixtures/pages.ts b/tests/e2e/ui/fixtures/pages.ts similarity index 100% rename from ui/litellm-dashboard/e2e_tests/fixtures/pages.ts rename to tests/e2e/ui/fixtures/pages.ts diff --git a/ui/litellm-dashboard/e2e_tests/fixtures/roles.ts b/tests/e2e/ui/fixtures/roles.ts similarity index 100% rename from ui/litellm-dashboard/e2e_tests/fixtures/roles.ts rename to tests/e2e/ui/fixtures/roles.ts diff --git a/ui/litellm-dashboard/e2e_tests/fixtures/seed.sql b/tests/e2e/ui/fixtures/seed.sql similarity index 100% rename from ui/litellm-dashboard/e2e_tests/fixtures/seed.sql rename to tests/e2e/ui/fixtures/seed.sql diff --git a/ui/litellm-dashboard/e2e_tests/fixtures/users.ts b/tests/e2e/ui/fixtures/users.ts similarity index 100% rename from ui/litellm-dashboard/e2e_tests/fixtures/users.ts rename to tests/e2e/ui/fixtures/users.ts diff --git a/ui/litellm-dashboard/e2e_tests/globalSetup.ts b/tests/e2e/ui/globalSetup.ts similarity index 100% rename from ui/litellm-dashboard/e2e_tests/globalSetup.ts rename to tests/e2e/ui/globalSetup.ts diff --git a/ui/litellm-dashboard/e2e_tests/helpers/navigation.ts b/tests/e2e/ui/helpers/navigation.ts similarity index 100% rename from ui/litellm-dashboard/e2e_tests/helpers/navigation.ts rename to tests/e2e/ui/helpers/navigation.ts diff --git a/ui/litellm-dashboard/e2e_tests/migration.serverRootPath.config.ts b/tests/e2e/ui/migration.serverRootPath.config.ts similarity index 100% rename from ui/litellm-dashboard/e2e_tests/migration.serverRootPath.config.ts rename to tests/e2e/ui/migration.serverRootPath.config.ts diff --git a/ui/litellm-dashboard/e2e_tests/migration.serverRootPath.globalSetup.ts b/tests/e2e/ui/migration.serverRootPath.globalSetup.ts similarity index 100% rename from ui/litellm-dashboard/e2e_tests/migration.serverRootPath.globalSetup.ts rename to tests/e2e/ui/migration.serverRootPath.globalSetup.ts diff --git a/tests/e2e/ui/package-lock.json b/tests/e2e/ui/package-lock.json new file mode 100644 index 00000000000..b22673a3535 --- /dev/null +++ b/tests/e2e/ui/package-lock.json @@ -0,0 +1,111 @@ +{ + "name": "litellm-ui-e2e", + "version": "0.0.0", + "lockfileVersion": 3, + "requires": true, + "packages": { + "": { + "name": "litellm-ui-e2e", + "version": "0.0.0", + "devDependencies": { + "@playwright/test": "1.58.1", + "@types/node": "20.19.37", + "typescript": "5.9.3" + } + }, + "node_modules/@playwright/test": { + "version": "1.58.1", + "resolved": "https://registry.npmjs.org/@playwright/test/-/test-1.58.1.tgz", + "integrity": "sha512-6LdVIUERWxQMmUSSQi0I53GgCBYgM2RpGngCPY7hSeju+VrKjq3lvs7HpJoPbDiY5QM5EYRtRX5fvrinnMAz3w==", + "dev": true, + "license": "Apache-2.0", + "dependencies": { + "playwright": "1.58.1" + }, + "bin": { + "playwright": "cli.js" + }, + "engines": { + "node": ">=18" + } + }, + "node_modules/@types/node": { + "version": "20.19.37", + "resolved": "https://registry.npmjs.org/@types/node/-/node-20.19.37.tgz", + "integrity": "sha512-8kzdPJ3FsNsVIurqBs7oodNnCEVbni9yUEkaHbgptDACOPW04jimGagZ51E6+lXUwJjgnBw+hyko/lkFWCldqw==", + "dev": true, + "license": "MIT", + "dependencies": { + "undici-types": "~6.21.0" + } + }, + "node_modules/fsevents": { + "version": "2.3.2", + "resolved": "https://registry.npmjs.org/fsevents/-/fsevents-2.3.2.tgz", + "integrity": "sha512-xiqMQR4xAeHTuB9uWm+fFRcIOgKBMiOBP+eXiyT7jsgVCq1bkVygt00oASowB7EdtpOHaaPgKt812P9ab+DDKA==", + "dev": true, + "hasInstallScript": true, + "license": "MIT", + "optional": true, + "os": [ + "darwin" + ], + "engines": { + "node": "^8.16.0 || ^10.6.0 || >=11.0.0" + } + }, + "node_modules/playwright": { + "version": "1.58.1", + "resolved": "https://registry.npmjs.org/playwright/-/playwright-1.58.1.tgz", + "integrity": "sha512-+2uTZHxSCcxjvGc5C891LrS1/NlxglGxzrC4seZiVjcYVQfUa87wBL6rTDqzGjuoWNjnBzRqKmF6zRYGMvQUaQ==", + "dev": true, + "license": "Apache-2.0", + "dependencies": { + "playwright-core": "1.58.1" + }, + "bin": { + "playwright": "cli.js" + }, + "engines": { + "node": ">=18" + }, + "optionalDependencies": { + "fsevents": "2.3.2" + } + }, + "node_modules/playwright-core": { + "version": "1.58.1", + "resolved": "https://registry.npmjs.org/playwright-core/-/playwright-core-1.58.1.tgz", + "integrity": "sha512-bcWzOaTxcW+VOOGBCQgnaKToLJ65d6AqfLVKEWvexyS3AS6rbXl+xdpYRMGSRBClPvyj44njOWoxjNdL/H9UNg==", + "dev": true, + "license": "Apache-2.0", + "bin": { + "playwright-core": "cli.js" + }, + "engines": { + "node": ">=18" + } + }, + "node_modules/typescript": { + "version": "5.9.3", + "resolved": "https://registry.npmjs.org/typescript/-/typescript-5.9.3.tgz", + "integrity": "sha512-jl1vZzPDinLr9eUt3J/t7V6FgNEw9QjvBPdysz9KfQDD41fQrC2Y4vKQdiaUpFT4bXlb1RHhLpp8wtm6M5TgSw==", + "dev": true, + "license": "Apache-2.0", + "bin": { + "tsc": "bin/tsc", + "tsserver": "bin/tsserver" + }, + "engines": { + "node": ">=14.17" + } + }, + "node_modules/undici-types": { + "version": "6.21.0", + "resolved": "https://registry.npmjs.org/undici-types/-/undici-types-6.21.0.tgz", + "integrity": "sha512-iwDZqg0QAGrg9Rav5H4n0M64c3mkR59cJ6wQp+7C4nI0gsmExaedaYLNO44eT4AtBBwjbTiGPMlt2Md0T9H9JQ==", + "dev": true, + "license": "MIT" + } + } +} diff --git a/tests/e2e/ui/package.json b/tests/e2e/ui/package.json new file mode 100644 index 00000000000..ede759d97cb --- /dev/null +++ b/tests/e2e/ui/package.json @@ -0,0 +1,16 @@ +{ + "name": "litellm-ui-e2e", + "version": "0.0.0", + "private": true, + "scripts": { + "e2e": "playwright test --config playwright.config.ts", + "e2e:ui": "playwright test --ui --config playwright.config.ts", + "e2e:migration": "playwright test tests/migration/migratedPages.spec.ts --config playwright.config.ts", + "e2e:migration:root": "playwright test --config migration.serverRootPath.config.ts" + }, + "devDependencies": { + "@playwright/test": "1.58.1", + "@types/node": "20.19.37", + "typescript": "5.9.3" + } +} diff --git a/ui/litellm-dashboard/e2e_tests/playwright.config.ts b/tests/e2e/ui/playwright.config.ts similarity index 100% rename from ui/litellm-dashboard/e2e_tests/playwright.config.ts rename to tests/e2e/ui/playwright.config.ts diff --git a/ui/litellm-dashboard/e2e_tests/run_e2e.sh b/tests/e2e/ui/run_e2e.sh similarity index 98% rename from ui/litellm-dashboard/e2e_tests/run_e2e.sh rename to tests/e2e/ui/run_e2e.sh index ea95f18890c..858eb401c8e 100755 --- a/ui/litellm-dashboard/e2e_tests/run_e2e.sh +++ b/tests/e2e/ui/run_e2e.sh @@ -20,8 +20,8 @@ set -euo pipefail # ================================================================ SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" -DASHBOARD_DIR="$(cd "$SCRIPT_DIR/.." && pwd)" REPO_ROOT="$(cd "$SCRIPT_DIR/../../.." && pwd)" +DASHBOARD_DIR="$REPO_ROOT/ui/litellm-dashboard" IS_CI="${CI:-false}" CONTAINER_NAME="litellm-e2e-postgres-$$" MOCK_PID="" @@ -187,12 +187,12 @@ PGPASSWORD="$DB_PASS" psql -h "$DB_HOST" -p "$DB_PORT" -U "$DB_USER" -d "$DB_NAM # --- Playwright --- echo "=== Installing Playwright dependencies ===" -cd "$DASHBOARD_DIR" +cd "$SCRIPT_DIR" npm install --silent 2>/dev/null || true npx playwright install chromium --with-deps 2>/dev/null || npx playwright install chromium echo "=== Running Playwright tests ===" -npx playwright test --config e2e_tests/playwright.config.ts "$@" +npx playwright test --config playwright.config.ts "$@" EXIT_CODE=$? exit $EXIT_CODE diff --git a/ui/litellm-dashboard/e2e_tests/serverRootPath.config.ts b/tests/e2e/ui/serverRootPath.config.ts similarity index 100% rename from ui/litellm-dashboard/e2e_tests/serverRootPath.config.ts rename to tests/e2e/ui/serverRootPath.config.ts diff --git a/ui/litellm-dashboard/e2e_tests/tests/auth/logout.spec.ts b/tests/e2e/ui/tests/auth/logout.spec.ts similarity index 100% rename from ui/litellm-dashboard/e2e_tests/tests/auth/logout.spec.ts rename to tests/e2e/ui/tests/auth/logout.spec.ts diff --git a/ui/litellm-dashboard/e2e_tests/tests/auth/proxyLogoutUrl.spec.ts b/tests/e2e/ui/tests/auth/proxyLogoutUrl.spec.ts similarity index 100% rename from ui/litellm-dashboard/e2e_tests/tests/auth/proxyLogoutUrl.spec.ts rename to tests/e2e/ui/tests/auth/proxyLogoutUrl.spec.ts diff --git a/ui/litellm-dashboard/e2e_tests/tests/auth/unauthenticatedRedirect.spec.ts b/tests/e2e/ui/tests/auth/unauthenticatedRedirect.spec.ts similarity index 100% rename from ui/litellm-dashboard/e2e_tests/tests/auth/unauthenticatedRedirect.spec.ts rename to tests/e2e/ui/tests/auth/unauthenticatedRedirect.spec.ts diff --git a/ui/litellm-dashboard/e2e_tests/tests/internal-user/internalUser.spec.ts b/tests/e2e/ui/tests/internal-user/internalUser.spec.ts similarity index 100% rename from ui/litellm-dashboard/e2e_tests/tests/internal-user/internalUser.spec.ts rename to tests/e2e/ui/tests/internal-user/internalUser.spec.ts diff --git a/ui/litellm-dashboard/e2e_tests/tests/internal-user/internalUserNoTeam.spec.ts b/tests/e2e/ui/tests/internal-user/internalUserNoTeam.spec.ts similarity index 100% rename from ui/litellm-dashboard/e2e_tests/tests/internal-user/internalUserNoTeam.spec.ts rename to tests/e2e/ui/tests/internal-user/internalUserNoTeam.spec.ts diff --git a/ui/litellm-dashboard/e2e_tests/tests/internal-user/internalUserWithTeams.spec.ts b/tests/e2e/ui/tests/internal-user/internalUserWithTeams.spec.ts similarity index 100% rename from ui/litellm-dashboard/e2e_tests/tests/internal-user/internalUserWithTeams.spec.ts rename to tests/e2e/ui/tests/internal-user/internalUserWithTeams.spec.ts diff --git a/ui/litellm-dashboard/e2e_tests/tests/internal-viewer/internalViewer.spec.ts b/tests/e2e/ui/tests/internal-viewer/internalViewer.spec.ts similarity index 100% rename from ui/litellm-dashboard/e2e_tests/tests/internal-viewer/internalViewer.spec.ts rename to tests/e2e/ui/tests/internal-viewer/internalViewer.spec.ts diff --git a/ui/litellm-dashboard/e2e_tests/tests/login/internalUserIdentity.spec.ts b/tests/e2e/ui/tests/login/internalUserIdentity.spec.ts similarity index 100% rename from ui/litellm-dashboard/e2e_tests/tests/login/internalUserIdentity.spec.ts rename to tests/e2e/ui/tests/login/internalUserIdentity.spec.ts diff --git a/ui/litellm-dashboard/e2e_tests/tests/login/login.spec.ts b/tests/e2e/ui/tests/login/login.spec.ts similarity index 100% rename from ui/litellm-dashboard/e2e_tests/tests/login/login.spec.ts rename to tests/e2e/ui/tests/login/login.spec.ts diff --git a/ui/litellm-dashboard/e2e_tests/tests/login/serverRootPathRedirect.spec.ts b/tests/e2e/ui/tests/login/serverRootPathRedirect.spec.ts similarity index 100% rename from ui/litellm-dashboard/e2e_tests/tests/login/serverRootPathRedirect.spec.ts rename to tests/e2e/ui/tests/login/serverRootPathRedirect.spec.ts diff --git a/ui/litellm-dashboard/e2e_tests/tests/mcp/mcpServers.spec.ts b/tests/e2e/ui/tests/mcp/mcpServers.spec.ts similarity index 100% rename from ui/litellm-dashboard/e2e_tests/tests/mcp/mcpServers.spec.ts rename to tests/e2e/ui/tests/mcp/mcpServers.spec.ts diff --git a/ui/litellm-dashboard/e2e_tests/tests/migration/README.md b/tests/e2e/ui/tests/migration/README.md similarity index 84% rename from ui/litellm-dashboard/e2e_tests/tests/migration/README.md rename to tests/e2e/ui/tests/migration/README.md index 4b3a391d421..d6b33598ec4 100644 --- a/ui/litellm-dashboard/e2e_tests/tests/migration/README.md +++ b/tests/e2e/ui/tests/migration/README.md @@ -9,8 +9,9 @@ the default mount and a non-root `SERVER_ROOT_PATH` mount. ## Adding a page When a page's migration merges, add its route segment to -`e2e_tests/fixtures/migratedPages.ts` (keep it in lockstep with `MIGRATED_PAGES` -in `src/utils/migratedPages.ts`). Both suites pick it up automatically. +`tests/e2e/ui/fixtures/migratedPages.ts` (keep it in lockstep with `MIGRATED_PAGES` +in `ui/litellm-dashboard/src/utils/migratedPages.ts`). Both suites pick it up +automatically. ## Running diff --git a/ui/litellm-dashboard/e2e_tests/tests/migration/migratedPages.spec.ts b/tests/e2e/ui/tests/migration/migratedPages.spec.ts similarity index 100% rename from ui/litellm-dashboard/e2e_tests/tests/migration/migratedPages.spec.ts rename to tests/e2e/ui/tests/migration/migratedPages.spec.ts diff --git a/ui/litellm-dashboard/e2e_tests/tests/modelHub/modelHub.spec.ts b/tests/e2e/ui/tests/modelHub/modelHub.spec.ts similarity index 100% rename from ui/litellm-dashboard/e2e_tests/tests/modelHub/modelHub.spec.ts rename to tests/e2e/ui/tests/modelHub/modelHub.spec.ts diff --git a/ui/litellm-dashboard/e2e_tests/tests/modelsPage/addModel.spec.ts b/tests/e2e/ui/tests/modelsPage/addModel.spec.ts similarity index 89% rename from ui/litellm-dashboard/e2e_tests/tests/modelsPage/addModel.spec.ts rename to tests/e2e/ui/tests/modelsPage/addModel.spec.ts index 17ff1fc3f83..bb8806a9c01 100644 --- a/ui/litellm-dashboard/e2e_tests/tests/modelsPage/addModel.spec.ts +++ b/tests/e2e/ui/tests/modelsPage/addModel.spec.ts @@ -1,5 +1,5 @@ import { test, expect } from "@playwright/test"; -import { ADMIN_STORAGE_PATH, E2E_TEAM_CRUD_ALIAS, E2E_TEAM_CRUD_ID } from "../../constants"; +import { ADMIN_STORAGE_PATH, E2E_TEAM_CRUD_ID } from "../../constants"; import { Role, users } from "../../fixtures/users"; import { navigateToPage } from "../../helpers/navigation"; import { Page } from "../../fixtures/pages"; @@ -56,17 +56,17 @@ test.describe("Add Model", () => { }, }); expect(createResponse.ok()).toBe(true); + const createdModelId = (await createResponse.json()).model_info?.id; + expect(createdModelId, "model id from /model/new").toBeTruthy(); // Navigate to Models + Endpoints await page.goto("/ui"); await page.getByText("Models + Endpoints").click(); - // Click the new model row to open its detail view. The table renders - // a clickable outer row plus a nested detail row for the same model, - // so we target the first match (outer row) explicitly. - const modelRow = page.locator("tr", { hasText: modelName }).first(); - await expect(modelRow).toBeVisible({ timeout: 10_000 }); - await modelRow.click(); + // The Model ID cell is the drill-in control; the row itself is not clickable. + const modelIdCell = page.getByTestId(`model-id-${createdModelId}`); + await expect(modelIdCell).toBeVisible({ timeout: 10_000 }); + await modelIdCell.click(); await expect(page.getByText("Back to Models").first()).toBeVisible({ timeout: 10_000 }); @@ -137,11 +137,11 @@ test.describe("Add Model", () => { await page.waitForTimeout(2000); // Search for the model we just added - await page.locator('input[placeholder="Search model names..."]').fill("claude-haiku-4-5"); + await page.getByPlaceholder("Search model names").fill("claude-haiku-4-5"); await page.waitForTimeout(1000); // Verify the model appears in the results count (not "Showing 0 results") - await expect(page.getByTestId("models-results-count")).toHaveText(/Showing \d+ - \d+ of \d+ results/, { + await expect(page.getByTestId("pagination-range")).toHaveText(/Showing \d+-\d+ of \d+/, { timeout: 15_000, }); @@ -228,24 +228,24 @@ test.describe("Add Model", () => { // searching. await page.waitForTimeout(2000); - await page.locator('input[placeholder="Search model names..."]').fill("cohere"); + await page.getByPlaceholder("Search model names").fill("cohere"); await page.waitForTimeout(1000); // Confirm the search returned at least one result — gives a clear // failure message when the table is empty instead of timing out on a // row assertion. - await expect(page.getByTestId("models-results-count")).toHaveText(/Showing \d+ - \d+ of \d+ results/, { + await expect(page.getByTestId("pagination-range")).toHaveText(/Showing \d+-\d+ of \d+/, { timeout: 15_000, }); - // Stronger than "alias appears somewhere in tbody" — pin the assertion + // Stronger than "the team appears somewhere in tbody" — pin the assertion // to a single row that has BOTH the cohere model_name AND the seeded - // team alias, so a stale cohere row from "Add wildcard route" (no team) - // can't satisfy the check. + // team, so a stale cohere row from "Add wildcard route" (no team) can't + // satisfy the check. The Team ID column renders the id, not the alias. const teamCohereRow = page .locator("table tbody tr") .filter({ hasText: "cohere/" }) - .filter({ hasText: E2E_TEAM_CRUD_ALIAS }); + .filter({ hasText: E2E_TEAM_CRUD_ID }); await expect(teamCohereRow).toHaveCount(1, { timeout: 15_000 }); } finally { await deleteTeamScopedCohereModels(); @@ -281,11 +281,11 @@ test.describe("Add Model", () => { await page.waitForTimeout(2000); // Search for the wildcard model - await page.locator('input[placeholder="Search model names..."]').fill("cohere"); + await page.getByPlaceholder("Search model names").fill("cohere"); await page.waitForTimeout(1000); // Verify the model appears in the results count (not "Showing 0 results") - await expect(page.getByTestId("models-results-count")).toHaveText(/Showing \d+ - \d+ of \d+ results/, { + await expect(page.getByTestId("pagination-range")).toHaveText(/Showing \d+-\d+ of \d+/, { timeout: 15_000, }); diff --git a/ui/litellm-dashboard/e2e_tests/tests/modelsPage/clearCustomPricing.spec.ts b/tests/e2e/ui/tests/modelsPage/clearCustomPricing.spec.ts similarity index 96% rename from ui/litellm-dashboard/e2e_tests/tests/modelsPage/clearCustomPricing.spec.ts rename to tests/e2e/ui/tests/modelsPage/clearCustomPricing.spec.ts index 877c7f8c555..e67dcb96f36 100644 --- a/ui/litellm-dashboard/e2e_tests/tests/modelsPage/clearCustomPricing.spec.ts +++ b/tests/e2e/ui/tests/modelsPage/clearCustomPricing.spec.ts @@ -67,9 +67,10 @@ test.describe("Clear custom pricing on a deployment", () => { await page.goto("/ui"); await page.getByText("Models + Endpoints").click(); - const modelRow = page.locator("tr", { hasText: modelName }).first(); - await expect(modelRow).toBeVisible({ timeout: 15_000 }); - await modelRow.click(); + // The Model ID cell is the drill-in control; the row itself is not clickable. + const modelIdCell = page.getByTestId(`model-id-${createdModelId}`); + await expect(modelIdCell).toBeVisible({ timeout: 15_000 }); + await modelIdCell.click(); await expect(page.getByText("Back to Models").first()).toBeVisible({ timeout: 10_000, }); diff --git a/ui/litellm-dashboard/e2e_tests/tests/modelsPage/credentials.spec.ts b/tests/e2e/ui/tests/modelsPage/credentials.spec.ts similarity index 95% rename from ui/litellm-dashboard/e2e_tests/tests/modelsPage/credentials.spec.ts rename to tests/e2e/ui/tests/modelsPage/credentials.spec.ts index 8b7824813a4..7c836068567 100644 --- a/ui/litellm-dashboard/e2e_tests/tests/modelsPage/credentials.spec.ts +++ b/tests/e2e/ui/tests/modelsPage/credentials.spec.ts @@ -38,7 +38,8 @@ test.describe("Edit LLM credential", () => { const row = page.locator("tr", { hasText: credentialName }); await expect(row).toBeVisible({ timeout: 15_000 }); - await row.getByRole("button").first().click(); + await row.getByTestId(`credential-actions-${credentialName}`).click(); + await page.getByTestId("credential-action-edit").click(); const modal = page.locator(".ant-modal-content").filter({ hasText: "Edit Credential" }); await expect(modal).toBeVisible({ timeout: 10_000 }); diff --git a/ui/litellm-dashboard/e2e_tests/tests/navigation/sidebar.spec.ts b/tests/e2e/ui/tests/navigation/sidebar.spec.ts similarity index 100% rename from ui/litellm-dashboard/e2e_tests/tests/navigation/sidebar.spec.ts rename to tests/e2e/ui/tests/navigation/sidebar.spec.ts diff --git a/ui/litellm-dashboard/e2e_tests/tests/proxy-admin/keys.spec.ts b/tests/e2e/ui/tests/proxy-admin/keys.spec.ts similarity index 98% rename from ui/litellm-dashboard/e2e_tests/tests/proxy-admin/keys.spec.ts rename to tests/e2e/ui/tests/proxy-admin/keys.spec.ts index a55c19a53de..c44957ea737 100644 --- a/ui/litellm-dashboard/e2e_tests/tests/proxy-admin/keys.spec.ts +++ b/tests/e2e/ui/tests/proxy-admin/keys.spec.ts @@ -103,7 +103,8 @@ test.describe("Proxy Admin - Keys", () => { await expect(page.getByText("Back to Keys")).toBeVisible({ timeout: 10_000 }); - await page.getByRole("button", { name: "Delete Key" }).click(); + await page.getByRole("button", { name: "More key actions" }).click(); + await page.getByRole("menuitem", { name: "Delete Key" }).click(); const modal = page.locator(".ant-modal:visible"); await expect(modal).toBeVisible({ timeout: 5_000 }); diff --git a/ui/litellm-dashboard/e2e_tests/tests/proxy-admin/license.spec.ts b/tests/e2e/ui/tests/proxy-admin/license.spec.ts similarity index 100% rename from ui/litellm-dashboard/e2e_tests/tests/proxy-admin/license.spec.ts rename to tests/e2e/ui/tests/proxy-admin/license.spec.ts diff --git a/ui/litellm-dashboard/e2e_tests/tests/proxy-admin/teams.spec.ts b/tests/e2e/ui/tests/proxy-admin/teams.spec.ts similarity index 100% rename from ui/litellm-dashboard/e2e_tests/tests/proxy-admin/teams.spec.ts rename to tests/e2e/ui/tests/proxy-admin/teams.spec.ts diff --git a/ui/litellm-dashboard/e2e_tests/tests/settings/adminSettings.spec.ts b/tests/e2e/ui/tests/settings/adminSettings.spec.ts similarity index 100% rename from ui/litellm-dashboard/e2e_tests/tests/settings/adminSettings.spec.ts rename to tests/e2e/ui/tests/settings/adminSettings.spec.ts diff --git a/ui/litellm-dashboard/e2e_tests/tests/settings/routerSettings.spec.ts b/tests/e2e/ui/tests/settings/routerSettings.spec.ts similarity index 98% rename from ui/litellm-dashboard/e2e_tests/tests/settings/routerSettings.spec.ts rename to tests/e2e/ui/tests/settings/routerSettings.spec.ts index 3e140b9ab56..ffa5f2c2ae2 100644 --- a/ui/litellm-dashboard/e2e_tests/tests/settings/routerSettings.spec.ts +++ b/tests/e2e/ui/tests/settings/routerSettings.spec.ts @@ -6,7 +6,7 @@ import { Role, users } from "../../fixtures/users"; // Type-only import of the OpenAPI-generated backend schema, erased at runtime by // esbuild. It types the round-trips below so mistakes surface in the editor; the live // test against the real proxy is what actually enforces the contract. -import type { components } from "../../../src/lib/http/schema"; +import type { components } from "../../../../../ui/litellm-dashboard/src/lib/http/schema"; // These tests mutate the proxy's shared router_settings, and the Loadbalancing save // echoes the whole settings object, so they must not run concurrently. diff --git a/ui/litellm-dashboard/e2e_tests/tests/team-admin/teamAdmin.spec.ts b/tests/e2e/ui/tests/team-admin/teamAdmin.spec.ts similarity index 100% rename from ui/litellm-dashboard/e2e_tests/tests/team-admin/teamAdmin.spec.ts rename to tests/e2e/ui/tests/team-admin/teamAdmin.spec.ts diff --git a/ui/litellm-dashboard/e2e_tests/tests/users/searchUsers.spec.ts b/tests/e2e/ui/tests/users/searchUsers.spec.ts similarity index 100% rename from ui/litellm-dashboard/e2e_tests/tests/users/searchUsers.spec.ts rename to tests/e2e/ui/tests/users/searchUsers.spec.ts diff --git a/ui/litellm-dashboard/e2e_tests/tests/users/viewInternalUsers.spec.ts b/tests/e2e/ui/tests/users/viewInternalUsers.spec.ts similarity index 100% rename from ui/litellm-dashboard/e2e_tests/tests/users/viewInternalUsers.spec.ts rename to tests/e2e/ui/tests/users/viewInternalUsers.spec.ts diff --git a/tests/e2e/ui/tsconfig.json b/tests/e2e/ui/tsconfig.json new file mode 100644 index 00000000000..f9290fe7b49 --- /dev/null +++ b/tests/e2e/ui/tsconfig.json @@ -0,0 +1,16 @@ +{ + "compilerOptions": { + "target": "ES2022", + "module": "commonjs", + "moduleResolution": "node", + "lib": ["ES2022", "DOM"], + "strict": true, + "esModuleInterop": true, + "skipLibCheck": true, + "resolveJsonModule": true, + "noEmit": true, + "types": ["node"] + }, + "include": ["**/*.ts"], + "exclude": ["node_modules"] +} diff --git a/tests/litellm/llms/anthropic/test_anthropic_schema_filter.py b/tests/litellm/llms/anthropic/test_anthropic_schema_filter.py index bd3b4198e9e..c10ac5532a0 100644 --- a/tests/litellm/llms/anthropic/test_anthropic_schema_filter.py +++ b/tests/litellm/llms/anthropic/test_anthropic_schema_filter.py @@ -45,21 +45,14 @@ class TestFilterAnthropicOutputSchema: assert "minimum value: 0" in result["properties"]["age"]["description"] assert "maximum value: 150" in result["properties"]["age"]["description"] # Score had no description, should get one from constraints - assert ( - "exclusive minimum value: 0" in result["properties"]["score"]["description"] - ) - assert ( - "exclusive maximum value: 100" - in result["properties"]["score"]["description"] - ) + assert "exclusive minimum value: 0" in result["properties"]["score"]["description"] + assert "exclusive maximum value: 100" in result["properties"]["score"]["description"] def test_removes_string_constraints(self): """Test that minLength/maxLength are removed from string schemas.""" schema = { "type": "object", - "properties": { - "name": {"type": "string", "minLength": 1, "maxLength": 100} - }, + "properties": {"name": {"type": "string", "minLength": 1, "maxLength": 100}}, } result = AnthropicConfig.filter_anthropic_output_schema(schema) @@ -154,3 +147,203 @@ class TestFilterAnthropicOutputSchema: result = AnthropicConfig.filter_anthropic_output_schema(schema) assert result == schema # Should be unchanged + + def test_removes_uniqueitems(self): + """Test that uniqueItems is removed from array schemas. + + Reproduces the 400 ``invalid_request_error``: + "output_format.schema: For 'array' type, property 'uniqueItems' is not + supported". + """ + schema = { + "type": "object", + "properties": { + "tags": { + "type": "array", + "items": {"type": "string"}, + "uniqueItems": True, + } + }, + } + + result = AnthropicConfig.filter_anthropic_output_schema(schema) + + assert "uniqueItems" not in result["properties"]["tags"] + assert result["properties"]["tags"]["items"] == {"type": "string"} + # Constraint intent preserved in the description + assert "all array items must be unique" in result["properties"]["tags"]["description"] + + def test_removes_contains_constraints(self): + """Test that contains/minContains/maxContains are removed from arrays.""" + schema = { + "type": "array", + "items": {"type": "integer"}, + "contains": {"type": "integer", "const": 1}, + "minContains": 1, + "maxContains": 3, + } + + result = AnthropicConfig.filter_anthropic_output_schema(schema) + + assert "contains" not in result + assert "minContains" not in result + assert "maxContains" not in result + assert result["items"] == {"type": "integer"} + # The contains sub-schema is serialized into the advisory note so the model + # knows what item the array must contain. + assert "array must contain an item matching:" in result["description"] + assert '"const": 1' in result["description"] + assert "minimum number of matching items: 1" in result["description"] + assert "maximum number of matching items: 3" in result["description"] + + def test_removes_object_property_constraints(self): + """Test that minProperties/maxProperties are removed from object schemas.""" + schema = { + "type": "object", + "properties": {"a": {"type": "string"}}, + "minProperties": 1, + "maxProperties": 5, + } + + result = AnthropicConfig.filter_anthropic_output_schema(schema) + + assert "minProperties" not in result + assert "maxProperties" not in result + assert "minimum number of properties: 1" in result["description"] + assert "maximum number of properties: 5" in result["description"] + + def test_uniqueitems_false_skips_misleading_note(self): + """``uniqueItems: false`` is stripped but must not add a 'unique' note.""" + schema = { + "type": "array", + "items": {"type": "string"}, + "uniqueItems": False, + } + + result = AnthropicConfig.filter_anthropic_output_schema(schema) + + assert "uniqueItems" not in result + # A disabled constraint imposes no requirement -> no advisory note + assert "unique" not in result.get("description", "") + + def test_removes_multipleof(self): + """multipleOf is rejected by Anthropic for integer and number types.""" + schema = { + "type": "object", + "properties": {"n": {"type": "integer", "multipleOf": 5}}, + } + + result = AnthropicConfig.filter_anthropic_output_schema(schema) + + assert "multipleOf" not in result["properties"]["n"] + assert "must be a multiple of 5" in result["properties"]["n"]["description"] + + def test_removes_conditional_and_negation_keywords(self): + """if/then/else and not are rejected by Anthropic and stripped into notes.""" + schema = { + "type": "object", + "properties": {"kind": {"type": "string"}, "sound": {"type": "string", "not": {"const": "moo"}}}, + "if": {"properties": {"kind": {"const": "dog"}}}, + "then": {"required": ["sound"]}, + "else": {"required": ["kind"]}, + } + + result = AnthropicConfig.filter_anthropic_output_schema(schema) + + assert "if" not in result + assert "then" not in result + assert "else" not in result + assert "not" not in result["properties"]["sound"] + assert 'conditional (if): {"properties": {"kind": {"const": "dog"}}}' in result["description"] + assert 'conditional (then): {"required": ["sound"]}' in result["description"] + assert 'conditional (else): {"required": ["kind"]}' in result["description"] + assert 'must not match: {"const": "moo"}' in result["properties"]["sound"]["description"] + + def test_removes_object_shape_keywords(self): + """patternProperties/propertyNames/dependent*/unevaluatedProperties are stripped.""" + schema = { + "type": "object", + "properties": {"first": {"type": "string"}}, + "patternProperties": {"^x": {"type": "string"}}, + "propertyNames": {"pattern": "^[a-z]+$"}, + "dependentRequired": {"first": ["last"]}, + "dependentSchemas": {"first": {"required": ["last"]}}, + "unevaluatedProperties": {"type": "string"}, + } + + result = AnthropicConfig.filter_anthropic_output_schema(schema) + + for field in ( + "patternProperties", + "propertyNames", + "dependentRequired", + "dependentSchemas", + "unevaluatedProperties", + ): + assert field not in result + assert 'properties whose names match each pattern must satisfy: {"^x": {"type": "string"}}' in result["description"] + assert 'property names must satisfy: {"pattern": "^[a-z]+$"}' in result["description"] + assert 'dependent required properties: {"first": ["last"]}' in result["description"] + assert 'dependent schemas: {"first": {"required": ["last"]}}' in result["description"] + assert 'unevaluated properties must satisfy: {"type": "string"}' in result["description"] + + def test_removes_prefixitems(self): + """prefixItems is rejected by Anthropic for array types.""" + schema = { + "type": "array", + "prefixItems": [{"type": "number"}, {"type": "string"}], + "items": {"type": "number"}, + } + + result = AnthropicConfig.filter_anthropic_output_schema(schema) + + assert "prefixItems" not in result + assert result["items"] == {"type": "number"} + assert 'leading items must match, in order: [{"type": "number"}, {"type": "string"}]' in result["description"] + + def test_oneof_rewritten_to_anyof(self): + """oneOf 400s ("Schema type 'oneOf' is not supported") and becomes anyOf, like the SDK.""" + schema = { + "type": "object", + "properties": {"id": {"oneOf": [{"type": "string", "minLength": 1}, {"type": "integer"}]}}, + } + + result = AnthropicConfig.filter_anthropic_output_schema(schema) + + id_schema = result["properties"]["id"] + assert "oneOf" not in id_schema + assert [v["type"] for v in id_schema["anyOf"]] == ["string", "integer"] + assert "minLength" not in id_schema["anyOf"][0] + assert "minimum length: 1" in id_schema["anyOf"][0]["description"] + + def test_oneof_merges_into_existing_anyof(self): + schema = { + "anyOf": [{"type": "string"}], + "oneOf": [{"type": "integer"}], + } + + result = AnthropicConfig.filter_anthropic_output_schema(schema) + + assert "oneOf" not in result + assert [v["type"] for v in result["anyOf"]] == ["string", "integer"] + + def test_constraint_note_order_is_deterministic(self): + """Note order must not depend on set iteration order (PYTHONHASHSEED), or the + serialized request differs across proxy workers and breaks caching.""" + schema = { + "type": "array", + "items": {"type": "string"}, + "minItems": 1, + "maxItems": 10, + "uniqueItems": True, + "minContains": 2, + "maxContains": 3, + } + + result = AnthropicConfig.filter_anthropic_output_schema(schema) + + assert result["description"] == ( + "Note: minimum number of items: 1, maximum number of items: 10, " + "all array items must be unique, minimum number of matching items: 2, " + "maximum number of matching items: 3." + ) diff --git a/tests/llm_translation/reasoning_effort_grid/grid_spec.py b/tests/llm_translation/reasoning_effort_grid/grid_spec.py index c47fddb1d8d..4fa77f38940 100644 --- a/tests/llm_translation/reasoning_effort_grid/grid_spec.py +++ b/tests/llm_translation/reasoning_effort_grid/grid_spec.py @@ -167,6 +167,13 @@ ANTHROPIC_DIRECT_MODELS: Tuple[ModelEntry, ...] = ( "once the model is available." ), ), + ModelEntry( + alias="claude-opus-5", + model="anthropic/claude-opus-5", + mode="adaptive", + required_env=_ANTHROPIC_REQ, + caps=_CAPS_XHIGH_MAX, + ), ModelEntry( alias="claude-opus-4-8", model="anthropic/claude-opus-4-8", diff --git a/tests/llm_translation/test_openai.py b/tests/llm_translation/test_openai.py index 1fec7665daa..440fc36ed33 100644 --- a/tests/llm_translation/test_openai.py +++ b/tests/llm_translation/test_openai.py @@ -518,7 +518,7 @@ async def test_openai_codex_stream(sync_mode): from litellm.main import stream_chunk_builder kwargs = { - "model": "openai/gpt-5.2-codex", + "model": "openai/gpt-5.3-codex", "messages": [{"role": "user", "content": "Hey!"}], "stream": True, } @@ -550,7 +550,7 @@ async def test_openai_codex(sync_mode): { "model_name": "openai-codex-mini-latest", "litellm_params": { - "model": "openai/gpt-5.2-codex", + "model": "openai/gpt-5.3-codex", }, } ] @@ -838,7 +838,7 @@ def test_gpt_5_reasoning_streaming(): def test_openai_gpt_5_codex_reasoning(): litellm._turn_on_debug() completion_kwargs = { - "model": "gpt-5-codex", + "model": "gpt-5.3-codex", "messages": [ { "role": "system", diff --git a/tests/logging_callback_tests/gcs_pub_sub_body/spend_logs_payload.json b/tests/logging_callback_tests/gcs_pub_sub_body/spend_logs_payload.json index a4c50d3c575..41b4c3efb63 100644 --- a/tests/logging_callback_tests/gcs_pub_sub_body/spend_logs_payload.json +++ b/tests/logging_callback_tests/gcs_pub_sub_body/spend_logs_payload.json @@ -11,7 +11,7 @@ "user": "", "team_id": "", "organization_id": "", - "metadata": "{\"applied_guardrails\": [], \"batch_models\": null, \"mcp_tool_call_metadata\": null, \"vector_store_request_metadata\": null, \"guardrail_information\": null, \"usage_object\": {\"completion_tokens\": 20, \"prompt_tokens\": 10, \"total_tokens\": 30, \"completion_tokens_details\": null, \"prompt_tokens_details\": null}, \"model_map_information\": {\"model_map_key\": \"gpt-4o\", \"model_map_value\": {\"key\": \"gpt-4o\", \"max_tokens\": 16384, \"max_input_tokens\": 128000, \"max_output_tokens\": 16384, \"input_cost_per_token\": 2.5e-06, \"cache_creation_input_token_cost\": null, \"cache_read_input_token_cost\": 1.25e-06, \"input_cost_per_character\": null, \"input_cost_per_token_above_128k_tokens\": null, \"input_cost_per_token_above_200k_tokens\": null, \"input_cost_per_query\": null, \"input_cost_per_second\": null, \"input_cost_per_audio_token\": null, \"input_cost_per_token_batches\": 1.25e-06, \"output_cost_per_token_batches\": 5e-06, \"output_cost_per_token\": 1e-05, \"output_cost_per_audio_token\": null, \"output_cost_per_character\": null, \"output_cost_per_token_above_128k_tokens\": null, \"output_cost_per_character_above_128k_tokens\": null, \"output_cost_per_token_above_200k_tokens\": null, \"output_cost_per_second\": null, \"output_cost_per_image\": null, \"output_vector_size\": null, \"litellm_provider\": \"openai\", \"mode\": \"chat\", \"supports_system_messages\": true, \"supports_response_schema\": true, \"supports_vision\": true, \"supports_function_calling\": true, \"supports_tool_choice\": true, \"supports_assistant_prefill\": false, \"supports_prompt_caching\": true, \"supports_audio_input\": false, \"supports_audio_output\": false, \"supports_pdf_input\": false, \"supports_embedding_image_input\": false, \"supports_native_streaming\": null, \"supports_web_search\": true, \"supports_reasoning\": false, \"search_context_cost_per_query\": {\"search_context_size_low\": 0.03, \"search_context_size_medium\": 0.035, \"search_context_size_high\": 0.05}, \"tpm\": null, \"rpm\": null, \"supported_openai_params\": [\"frequency_penalty\", \"logit_bias\", \"logprobs\", \"top_logprobs\", \"max_tokens\", \"max_completion_tokens\", \"modalities\", \"prediction\", \"n\", \"presence_penalty\", \"seed\", \"stop\", \"stream\", \"stream_options\", \"temperature\", \"top_p\", \"tools\", \"tool_choice\", \"function_call\", \"functions\", \"max_retries\", \"extra_headers\", \"parallel_tool_calls\", \"audio\", \"response_format\", \"user\"]}}, \"additional_usage_values\": {\"completion_tokens_details\": null, \"prompt_tokens_details\": null}, \"user_api_key\": null, \"user_api_key_alias\": null, \"user_api_key_team_id\": null, \"user_api_key_project_id\": null, \"user_api_key_project_alias\": null, \"user_api_key_org_id\": null, \"user_api_key_user_id\": null, \"user_api_key_team_alias\": null, \"spend_logs_metadata\": null, \"requester_ip_address\": null, \"status\": null, \"proxy_server_request\": null, \"error_information\": null, \"attempted_retries\": null, \"max_retries\": null}", + "metadata": "{\"applied_guardrails\": [], \"batch_models\": null, \"mcp_tool_call_metadata\": null, \"vector_store_request_metadata\": null, \"guardrail_information\": null, \"compression_savings\": null, \"usage_object\": {\"completion_tokens\": 20, \"prompt_tokens\": 10, \"total_tokens\": 30, \"completion_tokens_details\": null, \"prompt_tokens_details\": null}, \"model_map_information\": {\"model_map_key\": \"gpt-4o\", \"model_map_value\": {\"key\": \"gpt-4o\", \"max_tokens\": 16384, \"max_input_tokens\": 128000, \"max_output_tokens\": 16384, \"input_cost_per_token\": 2.5e-06, \"cache_creation_input_token_cost\": null, \"cache_read_input_token_cost\": 1.25e-06, \"input_cost_per_character\": null, \"input_cost_per_token_above_128k_tokens\": null, \"input_cost_per_token_above_200k_tokens\": null, \"input_cost_per_query\": null, \"input_cost_per_second\": null, \"input_cost_per_audio_token\": null, \"input_cost_per_token_batches\": 1.25e-06, \"output_cost_per_token_batches\": 5e-06, \"output_cost_per_token\": 1e-05, \"output_cost_per_audio_token\": null, \"output_cost_per_character\": null, \"output_cost_per_token_above_128k_tokens\": null, \"output_cost_per_character_above_128k_tokens\": null, \"output_cost_per_token_above_200k_tokens\": null, \"output_cost_per_second\": null, \"output_cost_per_image\": null, \"output_vector_size\": null, \"litellm_provider\": \"openai\", \"mode\": \"chat\", \"supports_system_messages\": true, \"supports_response_schema\": true, \"supports_vision\": true, \"supports_function_calling\": true, \"supports_tool_choice\": true, \"supports_assistant_prefill\": false, \"supports_prompt_caching\": true, \"supports_audio_input\": false, \"supports_audio_output\": false, \"supports_pdf_input\": false, \"supports_embedding_image_input\": false, \"supports_native_streaming\": null, \"supports_web_search\": true, \"supports_reasoning\": false, \"search_context_cost_per_query\": {\"search_context_size_low\": 0.03, \"search_context_size_medium\": 0.035, \"search_context_size_high\": 0.05}, \"tpm\": null, \"rpm\": null, \"supported_openai_params\": [\"frequency_penalty\", \"logit_bias\", \"logprobs\", \"top_logprobs\", \"max_tokens\", \"max_completion_tokens\", \"modalities\", \"prediction\", \"n\", \"presence_penalty\", \"seed\", \"stop\", \"stream\", \"stream_options\", \"temperature\", \"top_p\", \"tools\", \"tool_choice\", \"function_call\", \"functions\", \"max_retries\", \"extra_headers\", \"parallel_tool_calls\", \"audio\", \"response_format\", \"user\"]}}, \"additional_usage_values\": {\"completion_tokens_details\": null, \"prompt_tokens_details\": null}, \"user_api_key\": null, \"user_api_key_alias\": null, \"user_api_key_team_id\": null, \"user_api_key_project_id\": null, \"user_api_key_project_alias\": null, \"user_api_key_org_id\": null, \"user_api_key_user_id\": null, \"user_api_key_team_alias\": null, \"spend_logs_metadata\": null, \"requester_ip_address\": null, \"status\": null, \"proxy_server_request\": null, \"error_information\": null, \"attempted_retries\": null, \"max_retries\": null}", "cache_key": "Cache OFF", "spend": 0.00022500000000000002, "total_tokens": 30, diff --git a/tests/proxy_migration_tests/test_offline_image_migration.py b/tests/proxy_migration_tests/test_offline_image_migration.py new file mode 100644 index 00000000000..8ed88278e69 --- /dev/null +++ b/tests/proxy_migration_tests/test_offline_image_migration.py @@ -0,0 +1,143 @@ +"""Image-level regression net for the prisma bake in the shipped runtime image. + +Boots a built image's migration entrypoint the way an OpenShift / air-gapped +deployment does (an internal-only network with no egress, an arbitrary non-root +uid in GID 0) against a brand-new Postgres, and asserts the schema was created. + +This catches the whole failure class, not one symptom: a bake that only works +under `docker run` as the default uid with network still passes every existing +check, because the migration entrypoint exits 0 even when it applied nothing. +Asserting the table count is what turns that silent success into a hard fail. + +Gated on LITELLM_IMAGE (the tag of the image to exercise) so it is skipped in +the normal unit-test run and exercised only where an image has been built (the +image-scan workflow). Requires a working docker CLI. +""" + +import shutil +import subprocess +import uuid + +import os +import pytest + +IMAGE = os.getenv("LITELLM_IMAGE") +POSTGRES_IMAGE = os.getenv("LITELLM_TEST_POSTGRES_IMAGE", "postgres:16-alpine") +MIN_TABLES = int(os.getenv("LITELLM_TEST_MIN_TABLES", "20")) +NON_ROOT_UID = "12345:0" # arbitrary uid in GID 0, as OpenShift restricted-v2 assigns + +pytestmark = [ + pytest.mark.skipif(IMAGE is None, reason="requires a built image (set LITELLM_IMAGE)"), + pytest.mark.skipif(shutil.which("docker") is None, reason="requires the docker CLI"), +] + + +def _docker(*args: str, check: bool = True) -> subprocess.CompletedProcess: + return subprocess.run( + ["docker", *args], capture_output=True, text=True, check=check + ) + + +@pytest.fixture() +def offline_postgres(): + """A fresh Postgres reachable only over an internal-only (no egress) network. + + Yields (network_name, postgres_host). Both are torn down afterwards. + """ + run_id = f"offlinemig-{uuid.uuid4().hex[:8]}" + network = f"{run_id}-net" + pg = f"{run_id}-pg" + + # Pull Postgres while egress still exists; the internal network below has none. + _docker("pull", "--quiet", POSTGRES_IMAGE) + # --internal => containers on this network cannot reach the internet, so a + # prisma engine download (binaries.prisma.sh / npm) fails instead of masking + # a non-self-contained bake. + _docker("network", "create", "--internal", network) + try: + _docker( + "run", "-d", "--name", pg, "--network", network, + "-e", "POSTGRES_PASSWORD=pw", "-e", "POSTGRES_DB=litellm", + POSTGRES_IMAGE, + ) + _wait_until_ready(pg) + yield network, pg + finally: + _docker("rm", "-f", pg, check=False) + _docker("network", "rm", network, check=False) + + +def _wait_until_ready(pg: str, attempts: int = 60) -> None: + for _ in range(attempts): + running = _docker( + "ps", "--filter", f"name={pg}", "--filter", "status=running", + "--format", "{{.Names}}", check=False, + ).stdout + if pg not in running: + logs = _docker("logs", pg, check=False).stdout + _docker("logs", pg, check=False).stderr + pytest.fail(f"postgres container is not running:\n{logs}") + ready = _docker( + "exec", pg, "pg_isready", "-U", "postgres", "-d", "litellm", check=False + ) + if ready.returncode == 0: + return + subprocess.run(["sleep", "1"]) + pytest.fail(f"postgres never became ready after {attempts}s") + + +def _table_count(pg: str) -> int: + result = _docker( + "exec", pg, "psql", "-U", "postgres", "-d", "litellm", "-tAc", + "SELECT count(*) FROM information_schema.tables WHERE table_schema='public';", + ) + return int(result.stdout.strip() or "0") + + +def test_migration_offline_as_non_root_uid(offline_postgres): + """The migration entrypoint creates the full schema offline as an arbitrary uid. + + Reproduces the OpenShift / air-gapped failure: on the pre-fix image the + migration exits 0 having created 0 tables (every DB endpoint then 500s on + missing columns); a self-contained bake creates the full schema. + """ + network, pg = offline_postgres + assert IMAGE is not None + + migrate = _docker( + "run", "--rm", "--network", network, "--user", NON_ROOT_UID, + "-e", f"DATABASE_URL=postgresql://postgres:pw@{pg}:5432/litellm", + "-e", "LITELLM_MASTER_KEY=sk-offline-migration-test", + "-e", "DISABLE_SCHEMA_UPDATE=false", + "-w", "/app", "--entrypoint", "python", + IMAGE, "litellm/proxy/prisma_migration.py", + check=False, + ) + tables = _table_count(pg) + + assert migrate.returncode == 0, ( + f"migration entrypoint exited {migrate.returncode} offline as uid {NON_ROOT_UID}\n" + f"stdout:\n{migrate.stdout}\nstderr:\n{migrate.stderr}" + ) + assert tables >= MIN_TABLES, ( + f"only {tables} tables created (need >= {MIN_TABLES}) offline as uid {NON_ROOT_UID}. " + "The prisma bake is not self-contained: it needs a runtime download or a " + "writable HOME/cache, so OpenShift and air-gapped deployments start on an " + f"empty database.\nstdout:\n{migrate.stdout}\nstderr:\n{migrate.stderr}" + ) + + +def test_runtime_cache_env_not_read_only(): + """No runtime cache env var may point at the world-read-only /opt/prisma bake. + + /opt/prisma is baked `a+rX` (no write). Pointing XDG_CACHE_HOME (or any cache + var an XDG-aware library honours) there would deny writes for every uid, so + guard against a future edit reintroducing that. + """ + assert IMAGE is not None + env = _docker("run", "--rm", "--entrypoint", "env", IMAGE).stdout + offenders = [ + line for line in env.splitlines() + if line.startswith(("XDG_CACHE_HOME=", "XDG_DATA_HOME=", "HOME=")) + and line.split("=", 1)[1].startswith("/opt/prisma") + ] + assert not offenders, f"cache/home env points at the read-only bake: {offenders}" diff --git a/tests/proxy_unit_tests/test_proxy_server.py b/tests/proxy_unit_tests/test_proxy_server.py index 212f7772cad..bedd4dd1838 100644 --- a/tests/proxy_unit_tests/test_proxy_server.py +++ b/tests/proxy_unit_tests/test_proxy_server.py @@ -1252,6 +1252,17 @@ async def test_create_team_member_add(prisma_client, new_member_method): return_value=LiteLLM_TeamTableCachedObj(team_id="1234") ) + tx_mock = AsyncMock() + tx_mock.query_raw = AsyncMock(return_value=[{"members_with_roles": []}]) + tx_mock.litellm_teamtable = team_mock_client + tx_cm = MagicMock() + tx_cm.__aenter__ = AsyncMock(return_value=tx_mock) + tx_cm.__aexit__ = AsyncMock(return_value=None) + original_tx = litellm.proxy.proxy_server.prisma_client.tx + litellm.proxy.proxy_server.prisma_client.tx = MagicMock( + return_value=tx_cm + ) + print(f"team_member_add_request={team_member_add_request}") await team_member_add( data=team_member_add_request, @@ -1273,6 +1284,7 @@ async def test_create_team_member_add(prisma_client, new_member_method): ) litellm.proxy.proxy_server.prisma_client.db.litellm_teamtable = original_val + litellm.proxy.proxy_server.prisma_client.tx = original_tx @pytest.mark.parametrize("team_member_role", ["admin", "user"]) @@ -1434,42 +1446,51 @@ async def test_create_team_member_add_team_admin( mock_litellm_usertable.find_unique = AsyncMock(return_value=None) team_mock_client = AsyncMock() - original_val = getattr( - litellm.proxy.proxy_server.prisma_client.db, "litellm_teamtable" - ) - litellm.proxy.proxy_server.prisma_client.db.litellm_teamtable = team_mock_client - team_mock_client.update = AsyncMock( return_value=LiteLLM_TeamTableCachedObj(team_id="1234") ) - try: - await team_member_add( - data=team_member_add_request, - user_api_key_dict=valid_token, + tx_mock = AsyncMock() + tx_mock.query_raw = AsyncMock(return_value=[{"members_with_roles": []}]) + tx_mock.litellm_teamtable = team_mock_client + tx_cm = MagicMock() + tx_cm.__aenter__ = AsyncMock(return_value=tx_mock) + tx_cm.__aexit__ = AsyncMock(return_value=None) + + with ( + patch.object( + litellm.proxy.proxy_server.prisma_client.db, + "litellm_teamtable", + team_mock_client, + ), + patch.object( + litellm.proxy.proxy_server.prisma_client, + "tx", + MagicMock(return_value=tx_cm), + ), + ): + try: + await team_member_add( + data=team_member_add_request, + user_api_key_dict=valid_token, + ) + except HTTPException as e: + if user_role == "user": + assert e.status_code == 403 + return + else: + raise e + + mock_client.assert_called() + + assert ( + mock_client.call_args.kwargs["data"]["create"]["max_budget"] + == litellm.max_internal_user_budget + ) + assert ( + mock_client.call_args.kwargs["data"]["create"]["budget_duration"] + == litellm.internal_user_budget_duration ) - except HTTPException as e: - if user_role == "user": - assert e.status_code == 403 - return - else: - raise e - - mock_client.assert_called() - - print(f"mock_client.call_args: {mock_client.call_args}") - print("mock_client.call_args.kwargs: {}".format(mock_client.call_args.kwargs)) - - assert ( - mock_client.call_args.kwargs["data"]["create"]["max_budget"] - == litellm.max_internal_user_budget - ) - assert ( - mock_client.call_args.kwargs["data"]["create"]["budget_duration"] - == litellm.internal_user_budget_duration - ) - - litellm.proxy.proxy_server.prisma_client.db.litellm_teamtable = original_val @pytest.mark.asyncio diff --git a/tests/proxy_unit_tests/test_proxy_utils.py b/tests/proxy_unit_tests/test_proxy_utils.py index ee18c96c393..d2218b08386 100644 --- a/tests/proxy_unit_tests/test_proxy_utils.py +++ b/tests/proxy_unit_tests/test_proxy_utils.py @@ -7,10 +7,12 @@ from typing import Any, Dict, List, Optional, Union from unittest.mock import Mock import pytest -from fastapi import Request +from fastapi import HTTPException, Request from starlette.datastructures import State +from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.proxy.utils import _get_docs_url, _get_openapi_url, _get_redoc_url +from litellm.types.guardrails import GuardrailEventHooks sys.path.insert( 0, os.path.abspath("../..") @@ -2638,6 +2640,463 @@ async def test_during_call_hook_parallel_execution_with_error(): litellm.callbacks = original_callbacks +class _PreCallGuardrail(CustomGuardrail): + """Test double for pre_call guardrails; records timing and observed payload.""" + + def __init__(self, name, run_in_parallel, execution_order, sleep=0.1, default_on=True): + super().__init__( + guardrail_name=name, + event_hook=GuardrailEventHooks.pre_call, + default_on=default_on, + run_in_parallel=run_in_parallel, + ) + self.name = name + self.sleep = sleep + self.execution_order = execution_order + self.observed_content = None + self.was_called = False + + async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): + self.was_called = True + self.observed_content = data["messages"][0]["content"] + self.execution_order.append(f"{self.name}_start") + await asyncio.sleep(self.sleep) + self.execution_order.append(f"{self.name}_end") + return None + + +@pytest.mark.asyncio +async def test_pre_call_hook_runs_opted_in_guardrails_in_parallel(): + """run_in_parallel pre_call guardrails execute concurrently (all start before any ends).""" + from litellm.caching.caching import DualCache + from litellm.proxy.utils import ProxyLogging + + proxy_logging = ProxyLogging(user_api_key_cache=DualCache()) + execution_order = [] + original_callbacks = litellm.callbacks.copy() if litellm.callbacks else [] + + try: + litellm.callbacks = [ + _PreCallGuardrail(f"g{i}", run_in_parallel=True, execution_order=execution_order) for i in range(3) + ] + + result = await proxy_logging.pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="test_key", user_id="test_user"), + data={"model": "gpt-4", "messages": [{"role": "user", "content": "hi"}]}, + call_type="completion", + ) + + first_end_idx = next(i for i, item in enumerate(execution_order) if "end" in item) + starts_before_first_end = sum(1 for item in execution_order[:first_end_idx] if "start" in item) + assert starts_before_first_end == 3, f"expected 3 concurrent starts, got {starts_before_first_end}" + assert result["model"] == "gpt-4" + finally: + litellm.callbacks = original_callbacks + + +@pytest.mark.asyncio +async def test_pre_call_hook_runs_default_guardrails_sequentially(): + """Guardrails without run_in_parallel keep the sequential, one-at-a-time behavior.""" + from litellm.caching.caching import DualCache + from litellm.proxy.utils import ProxyLogging + + proxy_logging = ProxyLogging(user_api_key_cache=DualCache()) + execution_order = [] + original_callbacks = litellm.callbacks.copy() if litellm.callbacks else [] + + try: + litellm.callbacks = [ + _PreCallGuardrail(f"g{i}", run_in_parallel=False, execution_order=execution_order) for i in range(2) + ] + + start = asyncio.get_event_loop().time() + await proxy_logging.pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="test_key", user_id="test_user"), + data={"model": "gpt-4", "messages": [{"role": "user", "content": "hi"}]}, + call_type="completion", + ) + elapsed = asyncio.get_event_loop().time() - start + + assert execution_order == ["g0_start", "g0_end", "g1_start", "g1_end"] + assert elapsed >= 0.18, f"sequential run took {elapsed}s, expected >= 0.18s" + finally: + litellm.callbacks = original_callbacks + + +@pytest.mark.asyncio +async def test_pre_call_hook_sequential_mutations_precede_parallel_batch(): + """Sequential (mutating) guardrails run before the parallel batch, which sees their changes.""" + from litellm.caching.caching import DualCache + from litellm.proxy.utils import ProxyLogging + + proxy_logging = ProxyLogging(user_api_key_cache=DualCache()) + execution_order = [] + original_callbacks = litellm.callbacks.copy() if litellm.callbacks else [] + + class MaskingGuardrail(CustomGuardrail): + def __init__(self): + super().__init__( + guardrail_name="masker", + event_hook=GuardrailEventHooks.pre_call, + default_on=True, + run_in_parallel=False, + ) + + async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): + data["messages"][0]["content"] = "MASKED" + return data + + parallel_observer = _PreCallGuardrail("observer", run_in_parallel=True, execution_order=execution_order) + + try: + litellm.callbacks = [parallel_observer, MaskingGuardrail()] + + await proxy_logging.pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="test_key", user_id="test_user"), + data={"model": "gpt-4", "messages": [{"role": "user", "content": "secret"}]}, + call_type="completion", + ) + + assert parallel_observer.observed_content == "MASKED" + finally: + litellm.callbacks = original_callbacks + + +@pytest.mark.asyncio +async def test_pre_call_hook_parallel_guardrail_blocks_request(): + """A raising parallel guardrail blocks the request before it reaches the LLM.""" + from litellm.caching.caching import DualCache + from litellm.proxy.utils import ProxyLogging + + proxy_logging = ProxyLogging(user_api_key_cache=DualCache()) + original_callbacks = litellm.callbacks.copy() if litellm.callbacks else [] + + class BlockingGuardrail(CustomGuardrail): + def __init__(self): + super().__init__( + guardrail_name="blocker", + event_hook=GuardrailEventHooks.pre_call, + default_on=True, + run_in_parallel=True, + ) + + async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): + raise HTTPException(status_code=400, detail="blocked by guardrail") + + try: + litellm.callbacks = [BlockingGuardrail()] + + with pytest.raises(HTTPException) as exc_info: + await proxy_logging.pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="test_key", user_id="test_user"), + data={"model": "gpt-4", "messages": [{"role": "user", "content": "hi"}]}, + call_type="completion", + ) + + assert exc_info.value.status_code == 400 + assert "blocked by guardrail" in str(exc_info.value.detail) + finally: + litellm.callbacks = original_callbacks + + +@pytest.mark.asyncio +async def test_pre_call_hook_parallel_guardrail_skipped_when_should_not_run(): + """A parallel guardrail that should_run_guardrail rejects is never invoked.""" + from litellm.caching.caching import DualCache + from litellm.proxy.utils import ProxyLogging + + proxy_logging = ProxyLogging(user_api_key_cache=DualCache()) + execution_order = [] + original_callbacks = litellm.callbacks.copy() if litellm.callbacks else [] + + try: + guardrail = _PreCallGuardrail( + "off_by_default", run_in_parallel=True, execution_order=execution_order, default_on=False + ) + litellm.callbacks = [guardrail] + + result = await proxy_logging.pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="test_key", user_id="test_user"), + data={"model": "gpt-4", "messages": [{"role": "user", "content": "hi"}]}, + call_type="completion", + ) + + assert guardrail.was_called is False + assert result["model"] == "gpt-4" + finally: + litellm.callbacks = original_callbacks + + +@pytest.mark.asyncio +async def test_pre_call_hook_parallel_block_wins_over_reroute(): + """A slower block must win over a faster reroute so crafted input cannot bypass a block.""" + from litellm.caching.caching import DualCache + from litellm.exceptions import SensitiveDataRouteException + from litellm.proxy.utils import ProxyLogging + + proxy_logging = ProxyLogging(user_api_key_cache=DualCache()) + original_callbacks = litellm.callbacks.copy() if litellm.callbacks else [] + + class FastRerouteGuardrail(CustomGuardrail): + def __init__(self): + super().__init__( + guardrail_name="rerouter", + event_hook=GuardrailEventHooks.pre_call, + default_on=True, + run_in_parallel=True, + ) + + async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): + raise SensitiveDataRouteException(route_to_model="on-prem", session_id="s1", guardrail_name="rerouter") + + class SlowBlockingGuardrail(CustomGuardrail): + def __init__(self): + super().__init__( + guardrail_name="blocker", + event_hook=GuardrailEventHooks.pre_call, + default_on=True, + run_in_parallel=True, + ) + + async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): + await asyncio.sleep(0.1) + raise HTTPException(status_code=400, detail="blocked by guardrail") + + try: + litellm.callbacks = [FastRerouteGuardrail(), SlowBlockingGuardrail()] + + with pytest.raises(HTTPException) as exc_info: + await proxy_logging.pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="test_key", user_id="test_user"), + data={"model": "gpt-4", "messages": [{"role": "user", "content": "hi"}]}, + call_type="completion", + ) + + assert exc_info.value.status_code == 400 + assert "blocked by guardrail" in str(exc_info.value.detail) + finally: + litellm.callbacks = original_callbacks + + +@pytest.mark.asyncio +async def test_pre_call_hook_parallel_awaits_all_when_one_blocks(): + """A block must not orphan sibling guardrails; every parallel guardrail runs to completion.""" + from litellm.caching.caching import DualCache + from litellm.proxy.utils import ProxyLogging + + proxy_logging = ProxyLogging(user_api_key_cache=DualCache()) + original_callbacks = litellm.callbacks.copy() if litellm.callbacks else [] + completed = [] + + class FastBlockingGuardrail(CustomGuardrail): + def __init__(self): + super().__init__( + guardrail_name="fast_blocker", + event_hook=GuardrailEventHooks.pre_call, + default_on=True, + run_in_parallel=True, + ) + + async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): + raise HTTPException(status_code=400, detail="blocked") + + class SlowGuardrail(CustomGuardrail): + def __init__(self): + super().__init__( + guardrail_name="slow", + event_hook=GuardrailEventHooks.pre_call, + default_on=True, + run_in_parallel=True, + ) + + async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): + await asyncio.sleep(0.1) + completed.append("slow") + return None + + try: + litellm.callbacks = [FastBlockingGuardrail(), SlowGuardrail()] + + with pytest.raises(HTTPException): + await proxy_logging.pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="test_key", user_id="test_user"), + data={"model": "gpt-4", "messages": [{"role": "user", "content": "hi"}]}, + call_type="completion", + ) + + assert completed == ["slow"], "slow guardrail was orphaned instead of awaited to completion" + finally: + litellm.callbacks = original_callbacks + + +class _PostCallGuardrail(CustomGuardrail): + """Test double for post_call guardrails; records timing and invocation.""" + + def __init__(self, name, run_in_parallel, execution_order, sleep=0.1, default_on=True): + super().__init__( + guardrail_name=name, + event_hook=GuardrailEventHooks.post_call, + default_on=default_on, + run_in_parallel=run_in_parallel, + ) + self.name = name + self.sleep = sleep + self.execution_order = execution_order + self.was_called = False + + async def async_post_call_success_hook(self, data, user_api_key_dict, response): + self.was_called = True + self.execution_order.append(f"{self.name}_start") + await asyncio.sleep(self.sleep) + self.execution_order.append(f"{self.name}_end") + return None + + +@pytest.mark.asyncio +async def test_post_call_hook_runs_opted_in_guardrails_in_parallel(): + """run_in_parallel post_call guardrails execute concurrently (all start before any ends).""" + from litellm.caching.caching import DualCache + from litellm.proxy.utils import ProxyLogging + + proxy_logging = ProxyLogging(user_api_key_cache=DualCache()) + execution_order = [] + original_callbacks = litellm.callbacks.copy() if litellm.callbacks else [] + + try: + litellm.callbacks = [ + _PostCallGuardrail(f"g{i}", run_in_parallel=True, execution_order=execution_order) for i in range(3) + ] + + await proxy_logging.post_call_success_hook( + data={"model": "gpt-4", "messages": [{"role": "user", "content": "hi"}]}, + response=litellm.ModelResponse(), + user_api_key_dict=UserAPIKeyAuth(api_key="test_key", user_id="test_user"), + ) + + first_end_idx = next(i for i, item in enumerate(execution_order) if "end" in item) + starts_before_first_end = sum(1 for item in execution_order[:first_end_idx] if "start" in item) + assert starts_before_first_end == 3, f"expected 3 concurrent starts, got {starts_before_first_end}" + finally: + litellm.callbacks = original_callbacks + + +@pytest.mark.asyncio +async def test_post_call_hook_runs_default_guardrails_sequentially(): + """post_call guardrails without run_in_parallel keep the sequential behavior.""" + from litellm.caching.caching import DualCache + from litellm.proxy.utils import ProxyLogging + + proxy_logging = ProxyLogging(user_api_key_cache=DualCache()) + execution_order = [] + original_callbacks = litellm.callbacks.copy() if litellm.callbacks else [] + + try: + litellm.callbacks = [ + _PostCallGuardrail(f"g{i}", run_in_parallel=False, execution_order=execution_order) for i in range(2) + ] + + start = asyncio.get_event_loop().time() + await proxy_logging.post_call_success_hook( + data={"model": "gpt-4", "messages": [{"role": "user", "content": "hi"}]}, + response=litellm.ModelResponse(), + user_api_key_dict=UserAPIKeyAuth(api_key="test_key", user_id="test_user"), + ) + elapsed = asyncio.get_event_loop().time() - start + + assert execution_order == ["g0_start", "g0_end", "g1_start", "g1_end"] + assert elapsed >= 0.18, f"sequential run took {elapsed}s, expected >= 0.18s" + finally: + litellm.callbacks = original_callbacks + + +@pytest.mark.asyncio +async def test_post_call_hook_parallel_guardrail_blocks_response(): + """A raising parallel post_call guardrail blocks the response before it reaches the client.""" + from litellm.caching.caching import DualCache + from litellm.proxy.utils import ProxyLogging + + proxy_logging = ProxyLogging(user_api_key_cache=DualCache()) + original_callbacks = litellm.callbacks.copy() if litellm.callbacks else [] + + class BlockingPostCallGuardrail(CustomGuardrail): + def __init__(self): + super().__init__( + guardrail_name="post_blocker", + event_hook=GuardrailEventHooks.post_call, + default_on=True, + run_in_parallel=True, + ) + + async def async_post_call_success_hook(self, data, user_api_key_dict, response): + raise HTTPException(status_code=400, detail="blocked response by guardrail") + + try: + litellm.callbacks = [BlockingPostCallGuardrail()] + + with pytest.raises(HTTPException) as exc_info: + await proxy_logging.post_call_success_hook( + data={"model": "gpt-4", "messages": [{"role": "user", "content": "hi"}]}, + response=litellm.ModelResponse(), + user_api_key_dict=UserAPIKeyAuth(api_key="test_key", user_id="test_user"), + ) + + assert exc_info.value.status_code == 400 + assert "blocked response by guardrail" in str(exc_info.value.detail) + finally: + litellm.callbacks = original_callbacks + + +@pytest.mark.asyncio +async def test_post_call_hook_parallel_awaits_all_when_one_blocks(): + """A blocking post_call guardrail must not orphan its siblings; all run to completion.""" + from litellm.caching.caching import DualCache + from litellm.proxy.utils import ProxyLogging + + proxy_logging = ProxyLogging(user_api_key_cache=DualCache()) + original_callbacks = litellm.callbacks.copy() if litellm.callbacks else [] + completed = [] + + class FastBlockingPostCall(CustomGuardrail): + def __init__(self): + super().__init__( + guardrail_name="fast_post_blocker", + event_hook=GuardrailEventHooks.post_call, + default_on=True, + run_in_parallel=True, + ) + + async def async_post_call_success_hook(self, data, user_api_key_dict, response): + raise HTTPException(status_code=400, detail="blocked") + + class SlowPostCall(CustomGuardrail): + def __init__(self): + super().__init__( + guardrail_name="slow_post", + event_hook=GuardrailEventHooks.post_call, + default_on=True, + run_in_parallel=True, + ) + + async def async_post_call_success_hook(self, data, user_api_key_dict, response): + await asyncio.sleep(0.1) + completed.append("slow") + return None + + try: + litellm.callbacks = [FastBlockingPostCall(), SlowPostCall()] + + with pytest.raises(HTTPException): + await proxy_logging.post_call_success_hook( + data={"model": "gpt-4", "messages": [{"role": "user", "content": "hi"}]}, + response=litellm.ModelResponse(), + user_api_key_dict=UserAPIKeyAuth(api_key="test_key", user_id="test_user"), + ) + + assert completed == ["slow"], "slow post_call guardrail was orphaned instead of awaited to completion" + finally: + litellm.callbacks = original_callbacks + + @pytest.mark.asyncio async def test_handle_logging_proxy_only_error_preserves_pass_through_call_type(): """Ensure _handle_logging_proxy_only_error does not overwrite call_type diff --git a/tests/test_litellm/anthropic_interface/test_rust_bridge_messages.py b/tests/test_litellm/anthropic_interface/test_rust_bridge_messages.py index b745ca8eadf..fbd7e36e298 100644 --- a/tests/test_litellm/anthropic_interface/test_rust_bridge_messages.py +++ b/tests/test_litellm/anthropic_interface/test_rust_bridge_messages.py @@ -219,8 +219,7 @@ def _gate(**overrides): kwargs = { "custom_llm_provider": "azure_ai", "litellm_params": GenericLiteLLMParams(api_key="sk-azure", rust=True), - "stream": False, - "rust_stream_eligible": False, + "has_agentic_hook": False, "model": "claude-sonnet-4-5", "api_key": "sk-azure", "api_base": "https://resource.services.ai.azure.com/anthropic", @@ -345,11 +344,11 @@ async def test_gate_skips_rust_for_unsupported_provider(): @pytest.mark.asyncio -async def test_gate_skips_rust_when_streaming_but_not_eligible(): +async def test_gate_skips_rust_for_agentic_hook(): bridge = ExplodingAsyncMessages() litellm.use_litellm_rust(True, amessages=bridge) - response = await _gate(stream=True, rust_stream_eligible=False) + response = await _gate(has_agentic_hook=True) assert response is None assert bridge.calls == 0 @@ -362,8 +361,7 @@ async def test_gate_streams_through_rust_when_eligible_and_strips_stream_flag(): streaming_body = {**REQUEST_BODY, "stream": True} response = await _gate( - stream=True, - rust_stream_eligible=True, + has_agentic_hook=False, request_body=streaming_body, ) diff --git a/tests/test_litellm/integrations/test_custom_guardrail.py b/tests/test_litellm/integrations/test_custom_guardrail.py index 9289dece83f..4ea79f9e2a4 100644 --- a/tests/test_litellm/integrations/test_custom_guardrail.py +++ b/tests/test_litellm/integrations/test_custom_guardrail.py @@ -1,3 +1,4 @@ +import asyncio from unittest.mock import AsyncMock import pytest @@ -1394,6 +1395,38 @@ class TestEventTypeLogging: assert len(logged_info) == 1 assert logged_info[0]["guardrail_status"] == "guardrail_intervened" + @pytest.mark.asyncio + async def test_log_guardrail_information_records_every_concurrent_guardrail(self): + """Guardrails run concurrently (parallel pre_call/post_call, during_call) share one + request_data dict. Each must still record its own entry. The previous guard counted + entries in that shared dict, so a sibling's append made a guardrail think it had already + recorded and skip its own auto-record — silently dropping lifecycle logs the UI shows.""" + from litellm.integrations.custom_guardrail import log_guardrail_information + from litellm.types.guardrails import GuardrailEventHooks + + class SleeperGuardrail(CustomGuardrail): + def __init__(self, name, sleep): + super().__init__(guardrail_name=name, event_hook=GuardrailEventHooks.pre_call) + self._sleep = sleep + + @log_guardrail_information + async def async_pre_call_hook(self, data: dict, **kwargs): + await asyncio.sleep(self._sleep) + return data + + request_data = {"metadata": {}} + # Different sleeps guarantee overlapping execution windows: the faster guardrail + # records while the slower one is still awaiting, which is exactly what tripped the + # old shared-count guard. + await asyncio.gather( + SleeperGuardrail("guardrail-a", 0.05).async_pre_call_hook(data=request_data), + SleeperGuardrail("guardrail-b", 0.15).async_pre_call_hook(data=request_data), + ) + + logged = request_data["metadata"]["standard_logging_guardrail_information"] + assert {entry["guardrail_name"] for entry in logged} == {"guardrail-a", "guardrail-b"} + assert len(logged) == 2 + def test_add_standard_logging_falls_back_to_event_hook_when_event_type_is_none( self, ): @@ -1716,3 +1749,201 @@ class TestApplyGuardrailStyleDeploymentDispatch: await guardrail.async_pre_call_deployment_hook(kwargs, CallTypes.acompletion) assert guardrail.apply_called is False + + +class TestOnlyScanNewMessages: + """Incremental guardrail scanning: only send text segments not already scanned this session.""" + + def _guardrail(self, **overrides): + params = dict(guardrail_name="test-guard", only_scan_new_messages=True) + params.update(overrides) + return CustomGuardrail(**params) + + def _cache(self): + from litellm.caching import DualCache + + return DualCache() + + @pytest.mark.asyncio + async def test_disabled_returns_none(self): + guardrail = self._guardrail(only_scan_new_messages=False) + result = await guardrail.filter_new_texts_for_session( + texts=["hi"], + request_data={"litellm_session_id": "s1"}, + cache=self._cache(), + ) + assert result is None + + @pytest.mark.asyncio + async def test_no_session_id_fails_safe_to_full_scan(self): + guardrail = self._guardrail() + result = await guardrail.filter_new_texts_for_session( + texts=["hi"], + request_data={"metadata": {}}, + cache=self._cache(), + ) + assert result is None + + @pytest.mark.asyncio + async def test_masking_guardrail_not_supported(self): + guardrail = self._guardrail(mask_request_content=True) + result = await guardrail.filter_new_texts_for_session( + texts=["hi"], + request_data={"litellm_session_id": "s1"}, + cache=self._cache(), + ) + assert result is None + + @pytest.mark.asyncio + async def test_cache_read_failure_fails_safe_to_full_scan(self): + from unittest.mock import AsyncMock + + guardrail = self._guardrail() + cache = self._cache() + cache.async_get_cache = AsyncMock(side_effect=RuntimeError("redis down")) + result = await guardrail.filter_new_texts_for_session( + texts=["hi"], + request_data={"litellm_session_id": "s1"}, + cache=cache, + ) + assert result is None + + @pytest.mark.asyncio + async def test_dedupes_previously_scanned_texts(self): + guardrail = self._guardrail() + cache = self._cache() + request = {"litellm_session_id": "sess-dedupe"} + turn1 = ["you are helpful", "first question"] + + first = await guardrail.filter_new_texts_for_session(texts=turn1, request_data=request, cache=cache) + assert first == turn1 + await guardrail.mark_texts_scanned(texts=turn1, request_data=request, cache=cache) + + turn2 = turn1 + ["an answer", "second question"] + second = await guardrail.filter_new_texts_for_session(texts=turn2, request_data=request, cache=cache) + assert second == ["an answer", "second question"] + + @pytest.mark.asyncio + async def test_no_new_texts_returns_empty(self): + guardrail = self._guardrail() + cache = self._cache() + request = {"litellm_session_id": "sess-empty"} + texts = ["only message"] + + await guardrail.filter_new_texts_for_session(texts=texts, request_data=request, cache=cache) + await guardrail.mark_texts_scanned(texts=texts, request_data=request, cache=cache) + + again = await guardrail.filter_new_texts_for_session(texts=texts, request_data=request, cache=cache) + assert again == [] + + @pytest.mark.asyncio + async def test_modified_earlier_text_is_rescanned(self): + guardrail = self._guardrail() + cache = self._cache() + request = {"litellm_session_id": "sess-edit"} + original = ["original"] + + await guardrail.filter_new_texts_for_session(texts=original, request_data=request, cache=cache) + await guardrail.mark_texts_scanned(texts=original, request_data=request, cache=cache) + + edited = ["original EDITED"] + result = await guardrail.filter_new_texts_for_session(texts=edited, request_data=request, cache=cache) + assert result == edited + + @pytest.mark.asyncio + async def test_blocked_scan_does_not_persist_hashes(self): + guardrail = self._guardrail() + cache = self._cache() + request = {"litellm_session_id": "sess-blocked"} + texts = ["please block me"] + + filtered = await guardrail.filter_new_texts_for_session(texts=texts, request_data=request, cache=cache) + assert filtered == texts + + again = await guardrail.filter_new_texts_for_session(texts=texts, request_data=request, cache=cache) + assert again == texts + + @pytest.mark.asyncio + async def test_scanned_hashes_written_with_fixed_ttl(self): + from unittest.mock import AsyncMock + + from litellm.constants import GUARDRAIL_SCANNED_MESSAGES_CACHE_TTL_SECONDS + + guardrail = self._guardrail() + cache = self._cache() + cache.async_set_cache = AsyncMock() + request = {"litellm_session_id": "sess-ttl"} + + await guardrail.mark_texts_scanned(texts=["a", "b"], request_data=request, cache=cache) + + cache.async_set_cache.assert_awaited_once() + assert cache.async_set_cache.await_args.kwargs["ttl"] == GUARDRAIL_SCANNED_MESSAGES_CACHE_TTL_SECONDS + + @pytest.mark.asyncio + async def test_session_id_from_metadata_is_used_for_dedupe(self): + guardrail = self._guardrail() + cache = self._cache() + request = {"metadata": {"session_id": "sess-meta"}} + texts = ["shared message"] + + await guardrail.filter_new_texts_for_session(texts=texts, request_data=request, cache=cache) + await guardrail.mark_texts_scanned(texts=texts, request_data=request, cache=cache) + + again = await guardrail.filter_new_texts_for_session(texts=texts, request_data=request, cache=cache) + assert again == [] + + @pytest.mark.asyncio + async def test_session_id_from_litellm_metadata_is_used_for_dedupe(self): + guardrail = self._guardrail() + cache = self._cache() + request = {"litellm_metadata": {"session_id": "sess-lmeta"}} + texts = ["shared message"] + + await guardrail.filter_new_texts_for_session(texts=texts, request_data=request, cache=cache) + await guardrail.mark_texts_scanned(texts=texts, request_data=request, cache=cache) + + again = await guardrail.filter_new_texts_for_session(texts=texts, request_data=request, cache=cache) + assert again == [] + + @pytest.mark.asyncio + async def test_mark_texts_scanned_disabled_does_not_persist(self): + from unittest.mock import AsyncMock + + guardrail = self._guardrail(only_scan_new_messages=False) + cache = self._cache() + cache.async_set_cache = AsyncMock() + + await guardrail.mark_texts_scanned(texts=["a"], request_data={"litellm_session_id": "s1"}, cache=cache) + cache.async_set_cache.assert_not_awaited() + + @pytest.mark.asyncio + async def test_mark_texts_scanned_masking_does_not_persist(self): + from unittest.mock import AsyncMock + + guardrail = self._guardrail(mask_request_content=True) + cache = self._cache() + cache.async_set_cache = AsyncMock() + + await guardrail.mark_texts_scanned(texts=["a"], request_data={"litellm_session_id": "s1"}, cache=cache) + cache.async_set_cache.assert_not_awaited() + + @pytest.mark.asyncio + async def test_mark_texts_scanned_without_session_does_not_persist(self): + from unittest.mock import AsyncMock + + guardrail = self._guardrail() + cache = self._cache() + cache.async_set_cache = AsyncMock() + + await guardrail.mark_texts_scanned(texts=["a"], request_data={"metadata": {}}, cache=cache) + cache.async_set_cache.assert_not_awaited() + + @pytest.mark.asyncio + async def test_mark_texts_scanned_survives_cache_write_failure(self): + from unittest.mock import AsyncMock + + guardrail = self._guardrail() + cache = self._cache() + cache.async_set_cache = AsyncMock(side_effect=RuntimeError("redis down")) + + await guardrail.mark_texts_scanned(texts=["a"], request_data={"litellm_session_id": "s1"}, cache=cache) diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py index 9ff67a82f40..d282e656ce8 100644 --- a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py +++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py @@ -2111,6 +2111,37 @@ def test_token_type_cost_breakdown_reads_cache_write_tokens(): ) +def test_generic_cost_per_token_openai_cache_write_tokens_gpt_5_6(): + """ + Regression: OpenAI gpt-5.6 reports cache-write tokens under + prompt_tokens_details.cache_write_tokens (not the Anthropic cache_creation_tokens + name). Those tokens must be billed at the cache-write rate rather than the plain + input rate. Customer report: cache creation tokens were never counted for the + GPT-5.6 series, so cost was undercounted on cache-write requests. + """ + os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" + litellm.model_cost = litellm.get_model_cost_map(url="") + + model = "gpt-5.6" + usage = Usage( + prompt_tokens=1000, + completion_tokens=10, + total_tokens=1010, + prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=0, cache_write_tokens=800), + ) + + assert usage.prompt_tokens_details.cache_write_tokens == 800 + assert usage.prompt_tokens_details.cache_creation_tokens == 800 + + prompt_cost, _ = generic_cost_per_token(model=model, usage=usage, custom_llm_provider="openai") + + info = litellm.get_model_info(model=model, custom_llm_provider="openai") + expected_prompt = (1000 - 800) * info["input_cost_per_token"] + 800 * info["cache_creation_input_token_cost"] + assert prompt_cost == pytest.approx(expected_prompt) + assert info["cache_creation_input_token_cost"] > info["input_cost_per_token"] + assert prompt_cost > 1000 * info["input_cost_per_token"] + + def test_token_type_cost_breakdown_reconciles_with_generic_total(): """ Both-ways check: the reasoning subset must sum with the remaining (text) output @@ -2166,6 +2197,65 @@ def test_token_type_cost_breakdown_zero_without_special_tokens(): ) +@pytest.mark.parametrize( + "raw_usage, expect_read, expect_write", + [ + ( + { + "input_tokens": 5000, + "output_tokens": 10, + "total_tokens": 5010, + "input_tokens_details": {"cached_tokens": 0, "cache_write_tokens": 4012}, + }, + False, + True, + ), + ( + { + "input_tokens": 5000, + "output_tokens": 10, + "total_tokens": 5010, + "input_tokens_details": {"cached_tokens": 4012, "cache_write_tokens": 0}, + }, + True, + False, + ), + ], +) +def test_token_type_cost_breakdown_openai_responses_api_cache_write_read( + raw_usage, expect_read, expect_write +): + """Regression for #34309: OpenAI Responses API reports cache tokens under + input_tokens_details.{cached_tokens, cache_write_tokens}, not the Anthropic-style + top-level cache_creation_input_tokens. The itemized breakdown must still populate + cache_read_cost / cache_creation_cost from the transformed usage.""" + from litellm.responses.utils import ResponseAPILoggingUtils + + os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" + litellm.model_cost = litellm.get_model_cost_map(url="") + + model = "gpt-5.6" + usage = ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage(raw_usage) + + breakdown = get_token_type_cost_breakdown( + model=model, custom_llm_provider="openai", usage=usage + ) + + info = litellm.get_model_info(model=model, custom_llm_provider="openai") + if expect_write: + assert breakdown.cache_creation_cost == pytest.approx( + 4012 * info["cache_creation_input_token_cost"] + ) + assert breakdown.cache_creation_cost > 0 + assert breakdown.cache_read_cost == 0.0 + if expect_read: + assert breakdown.cache_read_cost == pytest.approx( + 4012 * info["cache_read_input_token_cost"] + ) + assert breakdown.cache_read_cost > 0 + assert breakdown.cache_creation_cost == 0.0 + + def test_token_type_cost_breakdown_handles_unknown_model_gracefully(): """A model with no pricing must yield zeros, never raise.""" breakdown = get_token_type_cost_breakdown( diff --git a/tests/test_litellm/litellm_core_utils/test_duration_parser.py b/tests/test_litellm/litellm_core_utils/test_duration_parser.py index b6b617610a8..cb9f273a0a7 100644 --- a/tests/test_litellm/litellm_core_utils/test_duration_parser.py +++ b/tests/test_litellm/litellm_core_utils/test_duration_parser.py @@ -1,8 +1,13 @@ import unittest from datetime import datetime, time, timezone +from unittest.mock import patch from zoneinfo import ZoneInfo -from litellm.litellm_core_utils.duration_parser import get_next_standardized_reset_time +import litellm.litellm_core_utils.duration_parser as duration_parser +from litellm.litellm_core_utils.duration_parser import ( + duration_in_seconds, + get_next_standardized_reset_time, +) class TestStandardizedResetTime(unittest.TestCase): @@ -316,5 +321,69 @@ class TestResetTimeOfDay(unittest.TestCase): ) +class TestWordFormBudgetDurations(unittest.TestCase): + """The Admin UI historically persisted word-form budget durations + (hourly/daily/weekly/monthly). They must resolve to their real interval + instead of silently collapsing to a next-midnight (daily) reset. + """ + + def test_word_forms_map_to_correct_reset_times(self): + base_time = datetime(2023, 5, 17, 15, 20, 30, tzinfo=timezone.utc) + + self.assertEqual( + get_next_standardized_reset_time("hourly", base_time, "UTC"), + datetime(2023, 5, 17, 16, 0, 0, tzinfo=timezone.utc), + ) + self.assertEqual( + get_next_standardized_reset_time("daily", base_time, "UTC"), + datetime(2023, 5, 18, 0, 0, 0, tzinfo=timezone.utc), + ) + self.assertEqual( + get_next_standardized_reset_time("weekly", base_time, "UTC"), + datetime(2023, 5, 22, 0, 0, 0, tzinfo=timezone.utc), + ) + self.assertEqual( + get_next_standardized_reset_time("monthly", base_time, "UTC"), + datetime(2023, 6, 1, 0, 0, 0, tzinfo=timezone.utc), + ) + + def test_word_forms_are_not_all_collapsed_to_daily(self): + base_time = datetime(2023, 5, 17, 15, 20, 30, tzinfo=timezone.utc) + results = { + word: get_next_standardized_reset_time(word, base_time, "UTC") + for word in ("hourly", "daily", "weekly", "monthly") + } + self.assertEqual(len(set(results.values())), len(results)) + + def test_word_forms_match_canonical_int_unit_forms(self): + base_time = datetime(2023, 5, 17, 15, 20, 30, tzinfo=timezone.utc) + for word, canonical in (("hourly", "1h"), ("daily", "24h"), ("weekly", "7d"), ("monthly", "30d")): + self.assertEqual( + get_next_standardized_reset_time(word, base_time, "UTC"), + get_next_standardized_reset_time(canonical, base_time, "UTC"), + ) + + def test_word_forms_are_case_and_whitespace_insensitive(self): + base_time = datetime(2023, 5, 17, 15, 20, 30, tzinfo=timezone.utc) + self.assertEqual( + get_next_standardized_reset_time(" Monthly ", base_time, "UTC"), + datetime(2023, 6, 1, 0, 0, 0, tzinfo=timezone.utc), + ) + + def test_duration_in_seconds_accepts_word_forms(self): + self.assertEqual(duration_in_seconds("hourly"), 3600) + self.assertEqual(duration_in_seconds("daily"), 86400) + self.assertEqual(duration_in_seconds("weekly"), 604800) + self.assertEqual(duration_in_seconds("monthly"), 2592000) + + def test_invalid_duration_logs_warning_and_falls_back(self): + base_time = datetime(2023, 5, 15, 15, 0, 0, tzinfo=timezone.utc) + with patch.object(duration_parser.verbose_logger, "warning") as mock_warning: + result = get_next_standardized_reset_time("garbage", base_time, "UTC") + self.assertEqual(result, datetime(2023, 5, 16, 0, 0, 0, tzinfo=timezone.utc)) + mock_warning.assert_called_once() + self.assertIn("garbage", mock_warning.call_args.args) + + if __name__ == "__main__": unittest.main() diff --git a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py index 471c339bebd..6e9fa875a24 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -1111,6 +1111,171 @@ async def test_dispatch_success_handlers_invokes_async_callback_for_pass_through litellm._async_success_callback = original_async_callbacks +@pytest.mark.asyncio +async def test_dispatch_failure_handlers_prefer_async_does_not_submit_sync_handler( + logging_obj, +): + """prefer_async_handlers must await async_failure_handler and never submit the sync failure_handler. + + Submitting the sync ``failure_handler`` while awaiting ``async_failure_handler`` + lets both mutate the shared logging_obj at once, which is the concurrent-mutation + crash this dispatch guard exists to prevent. + """ + exception = ValueError("boom") + traceback_exception = "traceback" + + logging_obj.model_call_details["litellm_params"] = {} + + with ( + patch.object( + logging_obj, "async_failure_handler", new_callable=AsyncMock + ) as mock_async, + patch.object( + logging_obj, "failure_handler", new_callable=MagicMock + ) as mock_sync, + patch.object( + logging_obj, + "_should_run_sync_failure_callbacks_for_async_calls", + return_value=False, + ), + patch( + "litellm.litellm_core_utils.litellm_logging.executor.submit" + ) as mock_submit, + ): + await logging_obj.dispatch_failure_handlers( + exception, + traceback_exception, + prefer_async_handlers=True, + ) + + mock_async.assert_awaited_once_with(exception, traceback_exception) + mock_sync.assert_not_called() + mock_submit.assert_not_called() + + +@pytest.mark.asyncio +async def test_dispatch_failure_handlers_async_completes_before_sync_submit( + logging_obj, +): + """The async failure handler must fully finish before the legacy sync handler is scheduled. + + Ordering proves there is no window where both handlers touch the shared + logging_obj concurrently: the sync submit only happens after the await returns. + """ + exception = ValueError("boom") + traceback_exception = "traceback" + events: list[str] = [] + + async def _async_failure(exc, tb, **kwargs): + events.append("async_start") + await asyncio.sleep(0) + events.append("async_end") + + def _submit(*args, **kwargs): + events.append("sync_submit") + + logging_obj.model_call_details["litellm_params"] = {} + + with ( + patch.object(logging_obj, "async_failure_handler", side_effect=_async_failure), + patch.object(logging_obj, "failure_handler", new_callable=MagicMock), + patch.object( + logging_obj, + "_should_run_sync_failure_callbacks_for_async_calls", + return_value=True, + ), + patch( + "litellm.litellm_core_utils.litellm_logging.executor.submit", + side_effect=_submit, + ), + ): + await logging_obj.dispatch_failure_handlers( + exception, + traceback_exception, + prefer_async_handlers=True, + ) + + assert events == ["async_start", "async_end", "sync_submit"] + + +@pytest.mark.asyncio +async def test_dispatch_failure_handlers_submits_sync_handler_for_failure_only_callbacks( + logging_obj, +): + """A sync failure callback must still run when only failure callbacks are configured. + + The legacy thread-based path always submitted the sync failure_handler, so gating it on + the success callback list would silently drop failure logging for any deployment that + registers only failure callbacks and no success callbacks. This drives the real predicate + (unmocked), so gating the sync failure handler on the success list fails this test. + """ + exception = ValueError("boom") + traceback_exception = "traceback" + + def _sync_failure_callback(*args, **kwargs): + return None + + logging_obj.model_call_details["litellm_params"] = {} + logging_obj.dynamic_success_callbacks = None + logging_obj.dynamic_failure_callbacks = None + + with ( + patch.object(litellm, "success_callback", []), + patch.object(litellm, "failure_callback", [_sync_failure_callback]), + patch.object(logging_obj, "async_failure_handler", new_callable=AsyncMock), + patch.object( + logging_obj, "failure_handler", new_callable=MagicMock + ) as mock_sync, + patch( + "litellm.litellm_core_utils.litellm_logging.executor.submit" + ) as mock_submit, + ): + await logging_obj.dispatch_failure_handlers( + exception, + traceback_exception, + prefer_async_handlers=True, + ) + + mock_submit.assert_called_once_with(mock_sync, exception, traceback_exception) + + +@pytest.mark.asyncio +async def test_dispatch_failure_handlers_sync_sdk_shortcut_runs_sync_handler_inline( + logging_obj, +): + """A sync-SDK request (prefer_async_handlers=False) runs failure_handler inline. + + ``async for`` over a stream from ``completion()`` passes prefer_async_handlers=True; a + plain sync request leaves it False, so the legacy sync handler runs directly and the + async handler is never awaited, matching dispatch_success_handlers. + """ + exception = ValueError("boom") + traceback_exception = "traceback" + + logging_obj.model_call_details["litellm_params"] = {} + + with ( + patch.object( + logging_obj, "async_failure_handler", new_callable=AsyncMock + ) as mock_async, + patch.object( + logging_obj, "failure_handler", new_callable=MagicMock + ) as mock_sync, + patch( + "litellm.litellm_core_utils.litellm_logging.executor.submit" + ) as mock_submit, + ): + await logging_obj.dispatch_failure_handlers( + exception, + traceback_exception, + prefer_async_handlers=False, + ) + + mock_sync.assert_called_once_with(exception, traceback_exception) + mock_async.assert_not_awaited() + mock_submit.assert_not_called() + + def test_success_handler_skips_guardrail_logging_hook_when_disabled(logging_obj): """Ensure CustomGuardrail logging_hook is skipped when should_run_guardrail is False.""" import datetime diff --git a/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py b/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py index c5422e0d70f..9cd1fbb59a6 100644 --- a/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py +++ b/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py @@ -373,3 +373,132 @@ class TestAnthropicMessagesHandlerToolInjection: if __name__ == "__main__": # Run the tests pytest.main([__file__, "-v"]) + + +class TestAnthropicMessagesIncrementalScan: + """PR #33278: only_scan_new_messages through the real /v1/messages translation + handler (the path Claude Code uses). Encodes the wire payloads observed in the + live validation against a real Bedrock guardrail. + """ + + def _bedrock_guardrail(self): + from litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails import BedrockGuardrail + + return BedrockGuardrail( + guardrail_name="bedrock-incremental-anthropic", + guardrailIdentifier="test-guardrail", + guardrailVersion="DRAFT", + default_on=True, + only_scan_new_messages=True, + ) + + def _data(self, messages, session_id): + return { + "model": "claude-sonnet-4-5", + "messages": messages, + "system": "You are a helpful geography assistant.", + "litellm_session_id": session_id, + } + + @pytest.mark.asyncio + async def test_first_turn_scans_all_eligible_then_second_turn_scans_only_diff(self): + from unittest.mock import AsyncMock, patch + + handler = AnthropicMessagesHandler() + guardrail = self._bedrock_guardrail() + sid = "anth-sess-diff" + turn1 = [{"role": "user", "content": "What is the capital of France?"}] + turn2 = turn1 + [ + {"role": "assistant", "content": "Paris."}, + {"role": "user", "content": "What is the capital of Germany?"}, + ] + with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api: + mock_api.return_value = {"action": "NONE", "output": [], "outputs": []} + await handler.process_input_messages( + data=self._data(turn1, sid), guardrail_to_apply=guardrail + ) + assert mock_api.call_count == 1 + assert [m["content"] for m in mock_api.call_args.kwargs["messages"]] == [ + "What is the capital of France?" + ] + mock_api.reset_mock() + await handler.process_input_messages( + data=self._data(turn2, sid), guardrail_to_apply=guardrail + ) + assert mock_api.call_count == 1 + assert [m["content"] for m in mock_api.call_args.kwargs["messages"]] == [ + "Paris.", + "What is the capital of Germany?", + ] + + @pytest.mark.asyncio + async def test_identical_resend_makes_no_guardrail_call(self): + from unittest.mock import AsyncMock, patch + + handler = AnthropicMessagesHandler() + guardrail = self._bedrock_guardrail() + sid = "anth-sess-resend" + msgs = [ + {"role": "user", "content": "What is the capital of France?"}, + {"role": "assistant", "content": "Paris."}, + {"role": "user", "content": "What is the capital of Germany?"}, + ] + with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api: + mock_api.return_value = {"action": "NONE", "output": [], "outputs": []} + await handler.process_input_messages(data=self._data(msgs, sid), guardrail_to_apply=guardrail) + assert mock_api.call_count == 1 + mock_api.reset_mock() + await handler.process_input_messages(data=self._data(msgs, sid), guardrail_to_apply=guardrail) + mock_api.assert_not_called() + + @pytest.mark.asyncio + async def test_edited_history_message_is_rescanned(self): + from unittest.mock import AsyncMock, patch + + handler = AnthropicMessagesHandler() + guardrail = self._bedrock_guardrail() + sid = "anth-sess-edit" + msgs = [{"role": "user", "content": "What is the capital of France?"}] + edited = [{"role": "user", "content": "What is the capital and population of France?"}] + with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api: + mock_api.return_value = {"action": "NONE", "output": [], "outputs": []} + await handler.process_input_messages(data=self._data(msgs, sid), guardrail_to_apply=guardrail) + mock_api.reset_mock() + await handler.process_input_messages(data=self._data(edited, sid), guardrail_to_apply=guardrail) + assert mock_api.call_count == 1 + assert [m["content"] for m in mock_api.call_args.kwargs["messages"]] == [ + "What is the capital and population of France?" + ] + + @pytest.mark.asyncio + async def test_mixed_text_and_tool_use_keeps_text_segments(self): + """A message carrying both text and a tool_use block must not lose its text. + (tool_use inputs and tool_result content are dropped from texts on the + anthropic input path today; that is pre-existing baseline behavior.)""" + from unittest.mock import AsyncMock, patch + + handler = AnthropicMessagesHandler() + guardrail = self._bedrock_guardrail() + sid = "anth-sess-tools" + msgs = [ + {"role": "user", "content": "Search for the weather in Paris"}, + { + "role": "assistant", + "content": [ + {"type": "text", "text": "Let me look that up for you."}, + {"type": "tool_use", "id": "toolu_1", "name": "search", "input": {"query": "canary-args"}}, + ], + }, + { + "role": "user", + "content": [{"type": "tool_result", "tool_use_id": "toolu_1", "content": "canary-result"}], + }, + {"role": "user", "content": "Thanks, summarize the result."}, + ] + with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api: + mock_api.return_value = {"action": "NONE", "output": [], "outputs": []} + await handler.process_input_messages(data=self._data(msgs, sid), guardrail_to_apply=guardrail) + scanned = [m["content"] for m in mock_api.call_args.kwargs["messages"]] + assert "Let me look that up for you." in scanned, "text beside a tool_use must be scanned" + assert "Search for the weather in Paris" in scanned + assert "Thanks, summarize the result." in scanned diff --git a/tests/test_litellm/llms/anthropic/test_anthropic_output_format_filter.py b/tests/test_litellm/llms/anthropic/test_anthropic_output_format_filter.py new file mode 100644 index 00000000000..90cba035760 --- /dev/null +++ b/tests/test_litellm/llms/anthropic/test_anthropic_output_format_filter.py @@ -0,0 +1,147 @@ +""" +Coverage for filter_anthropic_output_schema's array/object constraint stripping. + +Mirrors tests/litellm/llms/anthropic/test_anthropic_schema_filter.py, but lives +under tests/test_litellm/ so the coverage-uploading CI job exercises the stripped +keyword handling (uniqueItems / contains / minProperties / maxProperties plus +multipleOf / patternProperties / propertyNames / dependentRequired / +dependentSchemas / unevaluatedProperties / if / then / else / not / prefixItems), +the ``uniqueItems: false`` branch, the oneOf to anyOf rewrite, and the +deterministic note ordering. +""" + +from litellm.llms.anthropic.chat.transformation import AnthropicConfig + + +class TestOutputFormatArrayObjectConstraints: + def test_removes_uniqueitems(self): + schema = { + "type": "object", + "properties": { + "tags": { + "type": "array", + "items": {"type": "string"}, + "uniqueItems": True, + } + }, + } + + result = AnthropicConfig.filter_anthropic_output_schema(schema) + + assert "uniqueItems" not in result["properties"]["tags"] + assert "all array items must be unique" in result["properties"]["tags"]["description"] + + def test_uniqueitems_false_skips_misleading_note(self): + schema = { + "type": "array", + "items": {"type": "string"}, + "uniqueItems": False, + } + + result = AnthropicConfig.filter_anthropic_output_schema(schema) + + assert "uniqueItems" not in result + assert "unique" not in result.get("description", "") + + def test_removes_contains_constraints(self): + schema = { + "type": "array", + "items": {"type": "integer"}, + "contains": {"type": "integer", "const": 1}, + "minContains": 1, + "maxContains": 3, + } + + result = AnthropicConfig.filter_anthropic_output_schema(schema) + + assert "contains" not in result + assert "minContains" not in result + assert "maxContains" not in result + assert "array must contain an item matching:" in result["description"] + assert '"const": 1' in result["description"] + + def test_removes_object_property_constraints(self): + schema = { + "type": "object", + "properties": {"a": {"type": "string"}}, + "minProperties": 1, + "maxProperties": 5, + } + + result = AnthropicConfig.filter_anthropic_output_schema(schema) + + assert "minProperties" not in result + assert "maxProperties" not in result + assert "minimum number of properties: 1" in result["description"] + assert "maximum number of properties: 5" in result["description"] + + def test_removes_remaining_rejected_keywords(self): + schema = { + "type": "object", + "properties": { + "n": {"type": "integer", "multipleOf": 5}, + "pair": {"type": "array", "prefixItems": [{"type": "number"}], "items": {"type": "number"}}, + "color": {"type": "string", "not": {"const": "red"}}, + }, + "patternProperties": {"^x": {"type": "string"}}, + "propertyNames": {"pattern": "^[a-z]+$"}, + "dependentRequired": {"n": ["pair"]}, + "dependentSchemas": {"n": {"required": ["pair"]}}, + "unevaluatedProperties": {"type": "string"}, + "if": {"properties": {"n": {"const": 5}}}, + "then": {"required": ["pair"]}, + "else": {"required": ["color"]}, + } + + result = AnthropicConfig.filter_anthropic_output_schema(schema) + + for field in ( + "patternProperties", + "propertyNames", + "dependentRequired", + "dependentSchemas", + "unevaluatedProperties", + "if", + "then", + "else", + ): + assert field not in result + assert "multipleOf" not in result["properties"]["n"] + assert "must be a multiple of 5" in result["properties"]["n"]["description"] + assert "prefixItems" not in result["properties"]["pair"] + assert 'leading items must match, in order: [{"type": "number"}]' in result["properties"]["pair"]["description"] + assert "not" not in result["properties"]["color"] + assert 'must not match: {"const": "red"}' in result["properties"]["color"]["description"] + assert 'conditional (if): {"properties": {"n": {"const": 5}}}' in result["description"] + + def test_oneof_rewritten_to_anyof(self): + schema = { + "type": "object", + "properties": {"id": {"oneOf": [{"type": "string", "minLength": 1}, {"type": "integer"}]}}, + } + + result = AnthropicConfig.filter_anthropic_output_schema(schema) + + id_schema = result["properties"]["id"] + assert "oneOf" not in id_schema + assert [v["type"] for v in id_schema["anyOf"]] == ["string", "integer"] + assert "minLength" not in id_schema["anyOf"][0] + + def test_constraint_note_order_is_deterministic(self): + schema = { + "type": "array", + "items": {"type": "string"}, + "minItems": 1, + "maxItems": 10, + "uniqueItems": True, + "minContains": 2, + "maxContains": 3, + } + + result = AnthropicConfig.filter_anthropic_output_schema(schema) + + assert result["description"] == ( + "Note: minimum number of items: 1, maximum number of items: 10, " + "all array items must be unique, minimum number of matching items: 2, " + "maximum number of matching items: 3." + ) diff --git a/tests/test_litellm/llms/bedrock/chat/agentcore/test_agentcore_transformation.py b/tests/test_litellm/llms/bedrock/chat/agentcore/test_agentcore_transformation.py index 448afd5f3a5..e5a2ea9b28f 100644 --- a/tests/test_litellm/llms/bedrock/chat/agentcore/test_agentcore_transformation.py +++ b/tests/test_litellm/llms/bedrock/chat/agentcore/test_agentcore_transformation.py @@ -70,25 +70,33 @@ class TestAgentCoreAcceptHeader: """ End-to-end test: verify Accept header appears in the final HTTP request when using JWT auth through litellm.completion(). + + No exception swallowing: if completion() raises (for example because the + injected client was silently ignored and a real network call was made), + the test must fail with that error, not a misleading mock assertion. """ from litellm.llms.custom_httpx.http_handler import HTTPHandler client = HTTPHandler() - with patch.object(client, "post", return_value=MagicMock()) as mock_post: - try: - litellm.completion( - model="bedrock/agentcore/arn:aws:bedrock-agentcore:us-west-2:888602223428:runtime/test_runtime", - messages=[{"role": "user", "content": "test"}], - api_key="test-jwt-token", - client=client, - ) - except Exception: - pass + mock_response = Mock(spec=httpx.Response) + mock_response.status_code = 200 + mock_response.headers = {"content-type": "application/json"} + mock_response.json.return_value = { + "result": {"role": "assistant", "content": [{"text": "agent reply"}]} + } - mock_post.assert_called_once() - headers = mock_post.call_args.kwargs["headers"] - assert "Accept" in headers - assert headers["Accept"] == "application/json, text/event-stream" + with patch.object(client, "post", return_value=mock_response) as mock_post: + response = litellm.completion( + model="bedrock/agentcore/arn:aws:bedrock-agentcore:us-west-2:888602223428:runtime/test_runtime", + messages=[{"role": "user", "content": "test"}], + api_key="test-jwt-token", + client=client, + ) + + mock_post.assert_called_once() + headers = mock_post.call_args.kwargs["headers"] + assert headers["Accept"] == "application/json, text/event-stream" + assert response.choices[0].message.content == "agent reply" class TestAgentCoreJsonResponseParsing: diff --git a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py index fc12ead36a1..f832a4087ec 100644 --- a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py +++ b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py @@ -633,7 +633,8 @@ def test_parallel_tool_calls_config_kept_for_sonnet_5(): ) assert data["additionalModelRequestFields"]["tool_choice"] == { - "disable_parallel_tool_use": True + "type": "auto", + "disable_parallel_tool_use": True, } finally: litellm.model_cost = old_cost @@ -4251,6 +4252,49 @@ def test_parallel_tool_calls_older_model_drops_disable_flag(): assert "parallel_tool_calls" not in additional +@pytest.mark.parametrize( + "parallel_tool_calls, expected_disable", + [(True, False), (False, True)], +) +def test_parallel_tool_calls_emits_typed_auto_tool_choice(parallel_tool_calls, expected_disable): + config = AmazonConverseConfig() + model = "us.anthropic.claude-opus-4-8" + messages = [{"role": "user", "content": "What's the weather in SF and NYC?"}] + + optional_params = config.map_openai_params( + non_default_params={"parallel_tool_calls": parallel_tool_calls, "tools": _TOOL_PARAM}, + optional_params={}, + model=model, + drop_params=False, + ) + + request_data = config.transform_request( + model=model, + messages=messages, + optional_params=optional_params, + litellm_params={}, + headers={}, + ) + + assert request_data["additionalModelRequestFields"]["tool_choice"] == { + "type": "auto", + "disable_parallel_tool_use": expected_disable, + } + + +def test_parallel_tool_use_merge_preserves_user_tool_choice_type(): + merged = AmazonConverseConfig._merge_parallel_tool_use_config( + {"tool_choice": {"type": "tool", "name": "get_weather", "disable_parallel_tool_use": False}}, + {"tool_choice": {"type": "auto", "disable_parallel_tool_use": True}}, + ) + + assert merged["tool_choice"] == { + "type": "tool", + "name": "get_weather", + "disable_parallel_tool_use": True, + } + + class TestBedrockMinThinkingBudgetTokens: """Test that thinking.budget_tokens is clamped to the Bedrock minimum (1024).""" diff --git a/tests/test_litellm/llms/bedrock/realtime/test_bedrock_realtime_handler.py b/tests/test_litellm/llms/bedrock/realtime/test_bedrock_realtime_handler.py index ddc2e026e83..ffe21b91ab2 100644 --- a/tests/test_litellm/llms/bedrock/realtime/test_bedrock_realtime_handler.py +++ b/tests/test_litellm/llms/bedrock/realtime/test_bedrock_realtime_handler.py @@ -54,15 +54,24 @@ class FakeBedrockStream: self.input_stream = input_stream if input_stream is not None else FakeInputStream() +class FakeLogging: + def __init__(self, trace_id="trace-nova-sonic"): + self.litellm_trace_id = trace_id + + class DisconnectingClientWS: def __init__(self, messages): self._messages = list(messages) + self.sent_to_client = [] async def receive_text(self): if self._messages: return self._messages.pop(0) raise RuntimeError("client disconnected") + async def send_text(self, message): + self.sent_to_client.append(message) + class ClosableClientWS: def __init__(self): @@ -85,10 +94,14 @@ class EndedBedrockStream: class RealtimeClientWS: def __init__(self): self.closed = False + self.sent_to_client = [] async def receive_text(self): raise RuntimeError("client disconnected") + async def send_text(self, message): + self.sent_to_client.append(message) + async def close(self, code=None, reason=None): self.closed = True @@ -277,6 +290,61 @@ class TestBedrockRealtimeHandler: assert client_ws.closed +class TestBedrockRealtimeSessionLifecycle: + """Server must emit session.created on connect and session.updated on session.update (LIT-4655 regression)""" + + @pytest.mark.asyncio + async def test_session_created_sent_on_connect_before_any_client_input(self, stub_aws_sdk_client): + handler = BedrockRealtime() + websocket = RealtimeClientWS() + + await handler.async_realtime( + model="amazon.nova-sonic-v1:0", + websocket=websocket, + logging_obj=FakeLogging(), + aws_region_name="us-east-1", + aws_access_key_id="k", + aws_secret_access_key="s", + ) + + assert websocket.sent_to_client, "server sent nothing on connect: spec-conformant clients deadlock" + first_event = json.loads(websocket.sent_to_client[0]) + assert first_event["type"] == "session.created" + assert first_event["session"]["id"] == "trace-nova-sonic" + assert first_event["session"]["model"] == "amazon.nova-sonic-v1:0" + + @pytest.mark.asyncio + async def test_session_update_is_acked_with_session_updated(self, stub_aws_models): + handler = BedrockRealtime() + config = BedrockRealtimeConfig() + stream = FakeBedrockStream() + client_ws = DisconnectingClientWS( + [json.dumps({"type": "session.update", "session": {"instructions": "hi", "modalities": ["text"]}})] + ) + + await handler._forward_client_to_bedrock( + client_ws, stream, config, "amazon.nova-sonic-v1:0", {}, FakeLogging() + ) + + acked = [json.loads(message) for message in client_ws.sent_to_client] + updated = [event for event in acked if event["type"] == "session.updated"] + assert updated, "session.update was not acked" + assert updated[0]["session"]["modalities"] == ["text"], "ack must reflect the requested modalities" + + @pytest.mark.asyncio + async def test_no_session_updated_without_logging_obj(self, stub_aws_models): + handler = BedrockRealtime() + config = BedrockRealtimeConfig() + stream = FakeBedrockStream() + client_ws = DisconnectingClientWS( + [json.dumps({"type": "session.update", "session": {"instructions": "hi"}})] + ) + + await handler._forward_client_to_bedrock(client_ws, stream, config, "amazon.nova-sonic-v1:0", {}) + + assert client_ws.sent_to_client == [] + + class TestBedrockRealtimeAwsAuth: """AWS auth params passed via litellm_params must reach the Smithy client config (LIT-3923 regression)""" @@ -288,7 +356,7 @@ class TestBedrockRealtimeAwsAuth: await handler.async_realtime( model="amazon.nova-sonic-v1:0", websocket=websocket, - logging_obj=MagicMock(), + logging_obj=FakeLogging(), aws_region_name="us-east-1", aws_access_key_id="litellm-params-access-key", aws_secret_access_key="litellm-params-secret-key", @@ -318,7 +386,7 @@ class TestBedrockRealtimeAwsAuth: await handler.async_realtime( model="amazon.nova-sonic-v1:0", websocket=RealtimeClientWS(), - logging_obj=MagicMock(), + logging_obj=FakeLogging(), aws_region_name="eu-west-1", aws_role_name="arn:aws:iam::123456789012:role/nova-sonic", aws_session_name="realtime-session", diff --git a/tests/test_litellm/llms/bedrock/realtime/test_bedrock_realtime_transformation.py b/tests/test_litellm/llms/bedrock/realtime/test_bedrock_realtime_transformation.py index a68aa603b26..aa002b6e302 100644 --- a/tests/test_litellm/llms/bedrock/realtime/test_bedrock_realtime_transformation.py +++ b/tests/test_litellm/llms/bedrock/realtime/test_bedrock_realtime_transformation.py @@ -403,8 +403,9 @@ class TestBedrockRealtimeResponseCreate: class TestBedrockRealtimeResponseTransformation: """Test suite for response transformation""" - def test_transform_session_start_response(self): - """Test sessionStart response transformation""" + def test_bedrock_session_start_does_not_emit_duplicate_session_created(self): + """A Bedrock output sessionStart must not forward a second session.created to the + client; session.created is sent exactly once on connect (LIT-4655)""" config = BedrockRealtimeConfig() logging_obj = MagicMock() logging_obj.litellm_trace_id = "trace_123" @@ -428,10 +429,8 @@ class TestBedrockRealtimeResponseTransformation: }, ) - assert len(result["response"]) == 1 - assert result["response"][0]["type"] == "session.created" - assert result["response"][0]["session"]["id"] == "trace_123" - assert "model" in result["response"][0]["session"] + assert result["response"] == [] + assert result["session_configuration_request"] == json.dumps({"configured": True}) def test_transform_text_output_response(self): """Test textOutput response transformation""" @@ -789,5 +788,47 @@ class TestBedrockRealtimeResponseTransformation: assert len(set(response_ids)) == 1, "Response IDs should be consistent" +class TestBedrockRealtimeSessionEvents: + """session.created / session.updated builders produce spec-shaped events (LIT-4655)""" + + @staticmethod + def _logging(): + from types import SimpleNamespace + + return SimpleNamespace(litellm_trace_id="trace_123") + + def test_session_created_event_shape(self): + event = BedrockRealtimeConfig().session_created_event("amazon.nova-sonic-v1:0", self._logging()) + assert event["type"] == "session.created" + assert event["session"]["id"] == "trace_123" + assert event["session"]["model"] == "amazon.nova-sonic-v1:0" + assert event["session"]["modalities"] == ["text", "audio"] + assert event["event_id"] + + def test_session_updated_event_shape(self): + event = BedrockRealtimeConfig().session_updated_event("amazon.nova-sonic-v1:0", self._logging()) + assert event["type"] == "session.updated" + assert event["session"]["id"] == "trace_123" + assert event["session"]["model"] == "amazon.nova-sonic-v1:0" + assert event["event_id"] + + def test_created_and_updated_have_distinct_event_ids(self): + config = BedrockRealtimeConfig() + logging_obj = self._logging() + created = config.session_created_event("amazon.nova-sonic-v1:0", logging_obj) + updated = config.session_updated_event("amazon.nova-sonic-v1:0", logging_obj) + assert created["event_id"] != updated["event_id"] + + def test_session_updated_reflects_requested_modalities(self): + event = BedrockRealtimeConfig().session_updated_event( + "amazon.nova-sonic-v1:0", self._logging(), modalities=["text"] + ) + assert event["session"]["modalities"] == ["text"] + + def test_session_updated_defaults_modalities_when_unspecified(self): + event = BedrockRealtimeConfig().session_updated_event("amazon.nova-sonic-v1:0", self._logging()) + assert event["session"]["modalities"] == ["text", "audio"] + + if __name__ == "__main__": pytest.main([__file__, "-v"]) diff --git a/tests/test_litellm/llms/huggingface/embedding/test_huggingface_embedding_handler.py b/tests/test_litellm/llms/huggingface/embedding/test_huggingface_embedding_handler.py index 8a072fa5097..c907e3249d1 100644 --- a/tests/test_litellm/llms/huggingface/embedding/test_huggingface_embedding_handler.py +++ b/tests/test_litellm/llms/huggingface/embedding/test_huggingface_embedding_handler.py @@ -1,4 +1,3 @@ -import importlib import json import os import sys @@ -16,22 +15,7 @@ MOCK_EMBEDDING_RESPONSE = [[0.1, 0.2, 0.3, 0.4, 0.5]] @pytest.fixture -def reload_huggingface_modules(): - """ - Reload modules to ensure fresh references after conftest reloads litellm. - This ensures the HTTPHandler class being patched is the same one used by - the embedding handler during parallel test execution. - """ - import litellm.llms.custom_httpx.http_handler as http_handler_module - import litellm.llms.huggingface.embedding.handler as hf_embedding_handler_module - - importlib.reload(http_handler_module) - importlib.reload(hf_embedding_handler_module) - yield - - -@pytest.fixture -def mock_embedding_http_handler(reload_huggingface_modules): +def mock_embedding_http_handler(): """Fixture to mock the HTTP handler for embedding tests""" with patch("litellm.llms.custom_httpx.http_handler.HTTPHandler.post") as mock_post: mock_response = MagicMock() @@ -43,7 +27,7 @@ def mock_embedding_http_handler(reload_huggingface_modules): @pytest.fixture -def mock_embedding_async_http_handler(reload_huggingface_modules): +def mock_embedding_async_http_handler(): """Fixture to mock the async HTTP handler for embedding tests""" with patch( "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", @@ -121,6 +105,20 @@ class TestHuggingFaceEmbedding: assert response.usage.prompt_tokens > 0 assert response.usage.total_tokens == response.usage.prompt_tokens + def test_model_name_with_https_substring_uses_api_base(self): + api_base = "https://legit.example/embed" + + litellm.embedding( + model="huggingface/my-https-endpoint", + input=["hello world"], + input_type="embed", + api_base=api_base, + ) + + self.mock_http.assert_called_once() + called_url = self.mock_http.call_args[0][0] + assert called_url == api_base + def test_embedding_with_sentence_similarity_task(self): """Test embedding when task type is sentence-similarity (requires 2+ sentences)""" diff --git a/tests/test_litellm/llms/oobabooga/chat/test_oobabooga.py b/tests/test_litellm/llms/oobabooga/chat/test_oobabooga.py new file mode 100644 index 00000000000..91ebb2bd9d4 --- /dev/null +++ b/tests/test_litellm/llms/oobabooga/chat/test_oobabooga.py @@ -0,0 +1,55 @@ +import os +import sys +from unittest.mock import MagicMock, patch + +sys.path.insert(0, os.path.abspath("../../../../..")) + +import litellm + +MOCK_COMPLETION_RESPONSE = { + "choices": [{"message": {"role": "assistant", "content": "hi there"}}], + "usage": {"prompt_tokens": 3, "completion_tokens": 2, "total_tokens": 5}, +} + + +def _mock_post_response(): + mock_response = MagicMock() + mock_response.status_code = 200 + mock_response.text = "ok" + mock_response.json.return_value = MOCK_COMPLETION_RESPONSE + return mock_response + + +def test_model_name_with_https_substring_uses_api_base(): + api_base = "https://legit.example" + + with patch( + "litellm.llms.custom_httpx.http_handler.HTTPHandler.post" + ) as mock_post: + mock_post.return_value = _mock_post_response() + + litellm.completion( + model="oobabooga/my-https-model", + messages=[{"role": "user", "content": "hello"}], + api_base=api_base, + ) + + mock_post.assert_called_once() + called_url = mock_post.call_args[0][0] + assert called_url == f"{api_base}/v1/chat/completions" + + +def test_url_valued_model_still_targets_that_url(): + with patch( + "litellm.llms.custom_httpx.http_handler.HTTPHandler.post" + ) as mock_post: + mock_post.return_value = _mock_post_response() + + litellm.completion( + model="oobabooga/https://sdk-user.example", + messages=[{"role": "user", "content": "hello"}], + ) + + mock_post.assert_called_once() + called_url = mock_post.call_args[0][0] + assert called_url == "https://sdk-user.example/v1/chat/completions" diff --git a/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py b/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py index 4c268d9dfc9..7730b664c5e 100644 --- a/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py +++ b/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py @@ -1137,3 +1137,95 @@ class TestGetStructuredMessages: if __name__ == "__main__": # Run the tests pytest.main([__file__, "-v"]) + + +class TestIncrementalScanRespectsSkipFlags: + """PR #33278: skip_system_message_in_guardrail and skip_tool_message_in_guardrail + are enforced while this handler builds inputs["texts"] (_extract_inputs early + returns for system/tool roles), upstream of BedrockGuardrail's incremental path. + Bypassing _select_messages_for_apply_guardrail therefore cannot resurrect skipped + content on any turn, including a session's first turn where every segment is new. + Verified live against a real Bedrock ApplyGuardrail before being encoded here. + The flags are set as instance attributes, mirroring how guardrail_registry + applies litellm_params to the callback (they are not constructor kwargs). + """ + + def _bedrock_guardrail(self): + from litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails import BedrockGuardrail + + guardrail = BedrockGuardrail( + guardrail_name="bedrock-incremental-skip-flags", + guardrailIdentifier="test-guardrail", + guardrailVersion="DRAFT", + default_on=True, + only_scan_new_messages=True, + ) + guardrail.skip_system_message_in_guardrail = True + guardrail.skip_tool_message_in_guardrail = True + return guardrail + + def _messages(self, followup=None): + base = [ + {"role": "system", "content": "SYSTEM-PROMPT-must-not-be-scanned"}, + {"role": "user", "content": "Search for the weather in Paris"}, + { + "role": "assistant", + "content": "Let me look that up.", + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": {"name": "search", "arguments": '{"query": "weather"}'}, + } + ], + }, + {"role": "tool", "tool_call_id": "call_1", "content": "TOOL-RESULT-must-not-be-scanned"}, + {"role": "user", "content": "Thanks, summarize."}, + ] + return base + (followup or []) + + @pytest.mark.asyncio + async def test_first_turn_scans_no_system_or_tool_content(self): + from unittest.mock import AsyncMock, patch + + handler = OpenAIChatCompletionsHandler() + guardrail = self._bedrock_guardrail() + data = {"messages": self._messages(), "litellm_session_id": "skip-flags-turn1"} + with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api: + mock_api.return_value = {"action": "NONE", "output": [], "outputs": []} + await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) + assert mock_api.call_count == 1 + scanned = [m["content"] for m in mock_api.call_args.kwargs["messages"]] + assert scanned == [ + "Search for the weather in Paris", + "Let me look that up.", + "Thanks, summarize.", + ] + assert not any("SYSTEM-PROMPT" in text for text in scanned) + assert not any("TOOL-RESULT" in text for text in scanned) + + @pytest.mark.asyncio + async def test_second_turn_scans_only_new_eligible_content(self): + from unittest.mock import AsyncMock, patch + + handler = OpenAIChatCompletionsHandler() + guardrail = self._bedrock_guardrail() + session = "skip-flags-turn2" + followup = [ + {"role": "assistant", "content": "It is sunny in Paris."}, + {"role": "user", "content": "And tomorrow?"}, + ] + with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api: + mock_api.return_value = {"action": "NONE", "output": [], "outputs": []} + await handler.process_input_messages( + data={"messages": self._messages(), "litellm_session_id": session}, + guardrail_to_apply=guardrail, + ) + mock_api.reset_mock() + await handler.process_input_messages( + data={"messages": self._messages(followup), "litellm_session_id": session}, + guardrail_to_apply=guardrail, + ) + assert mock_api.call_count == 1 + scanned = [m["content"] for m in mock_api.call_args.kwargs["messages"]] + assert scanned == ["It is sunny in Paris.", "And tomorrow?"] diff --git a/tests/test_litellm/llms/sagemaker/test_sagemaker_chat_transformation.py b/tests/test_litellm/llms/sagemaker/test_sagemaker_chat_transformation.py new file mode 100644 index 00000000000..54e6f95c795 --- /dev/null +++ b/tests/test_litellm/llms/sagemaker/test_sagemaker_chat_transformation.py @@ -0,0 +1,235 @@ +""" +Regression tests for LIT-4313: sagemaker_chat streaming must forward each AWS +event-stream frame as it arrives instead of buffering to a fixed 1024-byte +threshold and then draining a burst of deltas. + +The buffering came from `response.iter_bytes(chunk_size=1024)` / +`response.aiter_bytes(chunk_size=1024)`: httpx's ByteChunker withholds bytes until +`chunk_size` accumulates, so the first client delta could not be produced until +enough later frames had arrived to cross 1024 bytes, inflating TTFT and turning a +steady provider stream into gap-then-burst delivery. +""" + +import binascii +import json +import struct +from typing import AsyncIterator, Iterator +from unittest.mock import MagicMock + +import httpx +import pytest + +from litellm.llms.sagemaker.chat.transformation import SagemakerChatConfig + + +def _encode_header(name: str, value: str) -> bytes: + name_b = name.encode("utf-8") + value_b = value.encode("utf-8") + return struct.pack("B", len(name_b)) + name_b + struct.pack("B", 7) + struct.pack(">H", len(value_b)) + value_b + + +def _encode_event_frame(payload: bytes) -> bytes: + """Encode one AWS event-stream message that botocore's EventStreamBuffer decodes.""" + headers = { + ":event-type": "PayloadPart", + ":content-type": "application/json", + ":message-type": "event", + } + headers_b = b"".join(_encode_header(k, v) for k, v in headers.items()) + total_len = 16 + len(headers_b) + len(payload) + prelude = struct.pack(">I", total_len) + struct.pack(">I", len(headers_b)) + prelude_crc = struct.pack(">I", binascii.crc32(prelude) & 0xFFFFFFFF) + message = prelude + prelude_crc + headers_b + payload + message_crc = struct.pack(">I", binascii.crc32(message) & 0xFFFFFFFF) + return message + message_crc + + +def _delta_frame(index: int, content: str) -> bytes: + sse = ( + "data: " + + json.dumps( + { + "id": "chatcmpl-test", + "object": "chat.completion.chunk", + "created": 1700000000, + "choices": [{"index": 0, "delta": {"content": content}, "finish_reason": None}], + } + ) + + "\n\n" + ) + return _encode_event_frame(sse.encode("utf-8")) + + +def _make_frames(n: int) -> list[bytes]: + # Small single-token frames (< 1024 bytes each) so a fixed 1024-byte chunker + # would have to swallow several frames before releasing the first delta. + frames = [_delta_frame(i, f"token{i} ") for i in range(n)] + assert all(len(f) < 1024 for f in frames) + return frames + + +class _CountingSyncStream(httpx.SyncByteStream): + """Yields provider frames one at a time and records how many have been pulled.""" + + def __init__(self, frames: list[bytes]) -> None: + self._frames = frames + self.consumed = 0 + + def __iter__(self) -> Iterator[bytes]: + for frame in self._frames: + self.consumed += 1 + yield frame + + +class _CountingAsyncStream(httpx.AsyncByteStream): + def __init__(self, frames: list[bytes]) -> None: + self._frames = frames + self.consumed = 0 + + async def __aiter__(self) -> AsyncIterator[bytes]: + for frame in self._frames: + self.consumed += 1 + yield frame + + +class _FakeSyncClient: + def __init__(self, response: httpx.Response) -> None: + self._response = response + + def post(self, *args, **kwargs) -> httpx.Response: + return self._response + + +class _FakeAsyncClient: + def __init__(self, response: httpx.Response) -> None: + self._response = response + + async def post(self, *args, **kwargs) -> httpx.Response: + return self._response + + +def _content_of(chunk) -> str | None: + return chunk.choices[0].delta.content + + +def test_sync_first_event_emitted_after_a_single_frame(): + """The first delta must be available after exactly one source frame is pulled. + + With the old chunk_size=1024 the httpx chunker would consume several small + frames before yielding, so `consumed` would be > 1 at the first delta. + """ + frames = _make_frames(24) + stream = _CountingSyncStream(frames) + response = httpx.Response(200, stream=stream) + + wrapper = SagemakerChatConfig().get_sync_custom_stream_wrapper( + model="phi-4", + custom_llm_provider="sagemaker_chat", + logging_obj=MagicMock(), + api_base="https://runtime.sagemaker.us-east-1.amazonaws.com/endpoints/phi-4/invocations-response-stream", + headers={}, + data={}, + messages=[], + client=_FakeSyncClient(response), + ) + + first = next(c for c in wrapper.completion_stream if c is not None and _content_of(c) is not None) + assert _content_of(first) == "token0 " + assert stream.consumed == 1 + + +def test_sync_events_emitted_incrementally_without_bursting(): + """Each successive delta must correspond to exactly one newly-pulled frame.""" + frames = _make_frames(24) + stream = _CountingSyncStream(frames) + response = httpx.Response(200, stream=stream) + + wrapper = SagemakerChatConfig().get_sync_custom_stream_wrapper( + model="phi-4", + custom_llm_provider="sagemaker_chat", + logging_obj=MagicMock(), + api_base="https://runtime.sagemaker.us-east-1.amazonaws.com/endpoints/phi-4/invocations-response-stream", + headers={}, + data={}, + messages=[], + client=_FakeSyncClient(response), + ) + + consumed_at_delta = [ + stream.consumed for chunk in wrapper.completion_stream if chunk is not None and _content_of(chunk) is not None + ] + + assert consumed_at_delta == list(range(1, len(frames) + 1)) + + +@pytest.mark.asyncio +async def test_async_first_event_emitted_after_a_single_frame(): + frames = _make_frames(24) + stream = _CountingAsyncStream(frames) + response = httpx.Response(200, stream=stream) + + wrapper = await SagemakerChatConfig().get_async_custom_stream_wrapper( + model="phi-4", + custom_llm_provider="sagemaker_chat", + logging_obj=MagicMock(), + api_base="https://runtime.sagemaker.us-east-1.amazonaws.com/endpoints/phi-4/invocations-response-stream", + headers={}, + data={}, + messages=[], + client=_FakeAsyncClient(response), + ) + + consumed_at_delta = [] + async for chunk in wrapper.completion_stream: + if chunk is not None and _content_of(chunk) is not None: + consumed_at_delta.append(stream.consumed) + + assert consumed_at_delta == list(range(1, len(frames) + 1)) + + +def test_signed_body_includes_stream_flag(): + """A streaming request must carry `stream: true` in the signed body sent to SageMaker. + + `stream` flows into the request body through the transformed request (`{**optional_params}`) + and must survive SigV4 signing so the endpoint enables token-level streaming. + """ + headers, signed_body = SagemakerChatConfig().sign_request( + headers={}, + optional_params={ + "aws_access_key_id": "AKIATESTTESTTESTTEST", + "aws_secret_access_key": "test-secret-key", + "aws_region_name": "us-east-1", + }, + request_data={"model": "phi-4", "messages": [{"role": "user", "content": "hi"}], "stream": True}, + api_base="https://runtime.sagemaker.us-east-1.amazonaws.com/endpoints/phi-4/invocations-response-stream", + model="phi-4", + stream=True, + ) + assert signed_body is not None + assert json.loads(signed_body)["stream"] is True + + +@pytest.mark.parametrize("split_size", [1, 3, 7, 64, 4096]) +def test_decoder_reassembles_frames_across_arbitrary_byte_boundaries(split_size): + """Correctness must not depend on chunk boundaries falling on frame edges. + + Removing `chunk_size=1024` lets httpx yield raw transport reads, so in + production a single read can straddle several frames or split one frame in + half. This re-chunks the concatenated stream at boundaries that deliberately + ignore frame edges and asserts every delta still decodes, in order, exactly + once - the guarantee botocore's EventStreamBuffer provides. + """ + from litellm.llms.sagemaker.chat.transformation import AWSEventStreamDecoder + + frames = _make_frames(24) + blob = b"".join(frames) + chunks = [blob[i : i + split_size] for i in range(0, len(blob), split_size)] + + decoder = AWSEventStreamDecoder(model="phi-4", is_messages_api=True) + texts = [ + _content_of(chunk) + for chunk in decoder.iter_bytes(iter(chunks)) + if chunk is not None and _content_of(chunk) is not None + ] + + assert texts == [f"token{i} " for i in range(len(frames))] diff --git a/tests/test_litellm/llms/sagemaker/test_sagemaker_completion_handler.py b/tests/test_litellm/llms/sagemaker/test_sagemaker_completion_handler.py new file mode 100644 index 00000000000..1cb27b7cf5f --- /dev/null +++ b/tests/test_litellm/llms/sagemaker/test_sagemaker_completion_handler.py @@ -0,0 +1,174 @@ +""" +Regression tests for LIT-4313: the native `sagemaker/` streaming path must +forward each AWS event-stream frame as it arrives instead of buffering to a +fixed 1024-byte threshold and then draining a burst of tokens. + +The buffering came from `response.aiter_bytes(chunk_size=1024)`: httpx's +ByteChunker withholds bytes until `chunk_size` accumulates, so the first token +could not be produced until enough later frames had arrived to cross 1024 bytes, +inflating TTFT and turning a steady provider stream into gap-then-burst delivery. +""" + +import binascii +import json +import struct +from typing import AsyncIterator, Iterator +from unittest.mock import MagicMock + +import httpx +import pytest + +from litellm.llms.sagemaker.common_utils import SagemakerError +from litellm.llms.sagemaker.completion.handler import SagemakerLLM + + +def _encode_header(name: str, value: str) -> bytes: + name_b = name.encode("utf-8") + value_b = value.encode("utf-8") + return struct.pack("B", len(name_b)) + name_b + struct.pack("B", 7) + struct.pack(">H", len(value_b)) + value_b + + +def _encode_event_frame(payload: bytes) -> bytes: + """Encode one AWS event-stream message that botocore's EventStreamBuffer decodes.""" + headers = { + ":event-type": "PayloadPart", + ":content-type": "application/json", + ":message-type": "event", + } + headers_b = b"".join(_encode_header(k, v) for k, v in headers.items()) + total_len = 16 + len(headers_b) + len(payload) + prelude = struct.pack(">I", total_len) + struct.pack(">I", len(headers_b)) + prelude_crc = struct.pack(">I", binascii.crc32(prelude) & 0xFFFFFFFF) + message = prelude + prelude_crc + headers_b + payload + message_crc = struct.pack(">I", binascii.crc32(message) & 0xFFFFFFFF) + return message + message_crc + + +def _token_frame(text: str) -> bytes: + # SageMaker HF TGI streaming payloads are `{"token": {"text": ...}}` blobs. + sse = "data: " + json.dumps({"token": {"text": text}}) + "\n\n" + return _encode_event_frame(sse.encode("utf-8")) + + +def _make_frames(n: int) -> list[bytes]: + frames = [_token_frame(f"token{i} ") for i in range(n)] + assert all(len(f) < 1024 for f in frames) + return frames + + +class _CountingSyncStream(httpx.SyncByteStream): + """Yields provider frames one at a time and records how many have been pulled.""" + + def __init__(self, frames: list[bytes]) -> None: + self._frames = frames + self.consumed = 0 + + def __iter__(self) -> Iterator[bytes]: + for frame in self._frames: + self.consumed += 1 + yield frame + + +class _CountingAsyncStream(httpx.AsyncByteStream): + """Yields provider frames one at a time and records how many have been pulled.""" + + def __init__(self, frames: list[bytes]) -> None: + self._frames = frames + self.consumed = 0 + + async def __aiter__(self) -> AsyncIterator[bytes]: + for frame in self._frames: + self.consumed += 1 + yield frame + + +class _FakeSyncClient: + def __init__(self, response: httpx.Response) -> None: + self._response = response + + def post(self, *args, **kwargs) -> httpx.Response: + return self._response + + +class _FakeAsyncClient: + def __init__(self, response: httpx.Response) -> None: + self._response = response + + async def post(self, *args, **kwargs) -> httpx.Response: + return self._response + + +def test_sync_native_streaming_forwards_each_frame_incrementally(): + """Each token must be emitted after exactly one newly-pulled source frame. + + With the old `chunk_size=1024` the httpx chunker would swallow several small + frames before yielding, so the first token would arrive only after `consumed` + had already crossed multiple frames, and tokens would then replay in a burst. + """ + frames = _make_frames(24) + stream = _CountingSyncStream(frames) + response = httpx.Response(200, stream=stream) + + completion_stream = SagemakerLLM().make_sync_call( + api_base="https://runtime.sagemaker.us-east-1.amazonaws.com/endpoints/phi-4/invocations-response-stream", + headers={}, + data="", + logging_obj=MagicMock(), + client=_FakeSyncClient(response), + ) + + consumed_at_token = [] + texts = [] + for chunk in completion_stream: + if chunk is not None and chunk["text"]: + consumed_at_token.append(stream.consumed) + texts.append(chunk["text"]) + + assert texts == [f"token{i} " for i in range(len(frames))] + assert consumed_at_token == list(range(1, len(frames) + 1)) + + +def test_sync_native_streaming_raises_sagemaker_error_on_non_200(): + response = httpx.Response(500, text="boom") + + with pytest.raises(SagemakerError) as exc_info: + SagemakerLLM().make_sync_call( + api_base="https://runtime.sagemaker.us-east-1.amazonaws.com/endpoints/phi-4/invocations-response-stream", + headers={}, + data="", + logging_obj=MagicMock(), + client=_FakeSyncClient(response), + ) + + assert exc_info.value.status_code == 500 + + +@pytest.mark.asyncio +async def test_async_native_streaming_forwards_each_frame_incrementally(): + """Each token must be emitted after exactly one newly-pulled source frame. + + With the old `chunk_size=1024` the httpx chunker would swallow several small + frames before yielding, so the first token would arrive only after `consumed` + had already crossed multiple frames, and tokens would then replay in a burst. + """ + frames = _make_frames(24) + stream = _CountingAsyncStream(frames) + response = httpx.Response(200, stream=stream) + + completion_stream = await SagemakerLLM().make_async_call( + api_base="https://runtime.sagemaker.us-east-1.amazonaws.com/endpoints/phi-4/invocations-response-stream", + headers={}, + data="", + logging_obj=MagicMock(), + client=_FakeAsyncClient(response), + ) + + consumed_at_token = [] + texts = [] + async for chunk in completion_stream: + if chunk is not None and chunk["text"]: + consumed_at_token.append(stream.consumed) + texts.append(chunk["text"]) + + assert texts == [f"token{i} " for i in range(len(frames))] + assert consumed_at_token == list(range(1, len(frames) + 1)) diff --git a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py index 95e8e6561f1..5871644ca5f 100644 --- a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py +++ b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py @@ -2087,6 +2087,9 @@ async def test_vertex_ai_streaming_bad_request_is_not_wrapped(): async def async_failure_handler(self, *args, **kwargs): return None + async def dispatch_failure_handlers(self, *args, **kwargs): + return None + async def failing_make_call(client=None, **kwargs): raise VertexAIError(status_code=400, message="bad input", headers={}) @@ -5404,3 +5407,116 @@ def test_process_candidates_merges_thought_signatures_and_server_side_tools(): fields = model_response.choices[-1].message.provider_specific_fields assert fields["thought_signatures"] == ["sig-text"] assert fields["server_side_tool_invocations"][0]["id"] == "tool-1" + + +def _accumulating_gemini_iterator(): + from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( + ModelResponseIterator, + ) + + iterator = ModelResponseIterator( + streaming_response=[], sync_stream=True, logging_obj=MagicMock() + ) + iterator.chunk_type = "accumulated_json" + return iterator + + +def test_accumulated_json_chunk_multi_value_buffer_does_not_wedge(): + """Two complete Gemini objects buffered together must both surface. + + A whole-buffer json.loads raises "Extra data" on concatenated values and, since + the buffer was never reset on failure, returned None forever while growing without + bound. Peeling one value from the front keeps the remainder for the next call. + """ + obj = '{"candidates":[{"content":{"parts":[{"text":"a"}]}}],"usageMetadata":{}}' + iterator = _accumulating_gemini_iterator() + + first = iterator.handle_accumulated_json_chunk(chunk=obj + obj) + assert first is not None + assert first.choices[0].delta.content == "a" + + second = iterator.handle_accumulated_json_chunk(chunk="") + assert second is not None + assert second.choices[0].delta.content == "a" + + assert iterator.accumulated_json.strip() == "" + + +def test_accumulated_json_end_of_stream_drains_all_buffered_values(): + """End of stream must drain every buffered value and then terminate. + + With concatenated values a whole-buffer parse never succeeds, so __next__ kept + returning None without shrinking the buffer - an unrecoverable per-request spin. + The bounded loop asserts the iterator both surfaces all values and terminates. + """ + obj = '{"candidates":[{"content":{"parts":[{"text":"a"}]}}],"usageMetadata":{}}' + iterator = _accumulating_gemini_iterator() + iterator.response_iterator = iter([]) + iterator.accumulated_json = obj + obj + obj + + out = [] + terminated = False + for _ in range(100): + try: + chunk = iterator.__next__() + except StopIteration: + terminated = True + break + if chunk is not None: + out.append(chunk) + + assert terminated, "iterator did not terminate - accumulated buffer wedged" + assert len(out) == 3 + assert iterator.accumulated_json.strip() == "" + + +def test_accumulated_json_end_of_stream_surfaces_leading_value_before_truncated_tail(): + """A complete leading value must survive a truncated trailing value at end of stream. + + The mid-stream perf guard only inspects the buffer's last byte, so a complete leading + object followed by a truncated one (a server that cut the stream mid-object, last byte + not a closer) would keep the guard from ever parsing and drop the complete value. At end + of stream the drain ignores that guard, surfaces the complete value, and discards only + the truncated tail. + """ + obj = '{"candidates":[{"content":{"parts":[{"text":"a"}]}}],"usageMetadata":{}}' + iterator = _accumulating_gemini_iterator() + iterator.response_iterator = iter([]) + iterator.accumulated_json = obj + '{"candidates":' + + out = [] + for _ in range(100): + try: + chunk = iterator.__next__() + except StopIteration: + break + if chunk is not None: + out.append(chunk) + + assert len(out) == 1 + assert out[0].choices[0].delta.content == "a" + + +def test_accumulated_json_skips_non_dict_leading_value(): + """A non-dict value at the front must not block the dict values behind it. + + raw_decode advances past a decoded value, so a leading non-dict (a JSON array or scalar, + which Gemini never emits but a malformed stream could) must be consumed and skipped. If + the drain stopped on it, the trailing objects would be lost at end of stream. + """ + obj = '{"candidates":[{"content":{"parts":[{"text":"a"}]}}],"usageMetadata":{}}' + iterator = _accumulating_gemini_iterator() + iterator.response_iterator = iter([]) + iterator.accumulated_json = "[1, 2]" + obj + + out = [] + for _ in range(100): + try: + chunk = iterator.__next__() + except StopIteration: + break + if chunk is not None: + out.append(chunk) + + assert len(out) == 1 + assert out[0].choices[0].delta.content == "a" diff --git a/tests/test_litellm/llms/vertex_ai/rerank/test_vertex_ai_rerank_integration.py b/tests/test_litellm/llms/vertex_ai/rerank/test_vertex_ai_rerank_integration.py index 6af4cf698e2..7fea5ac0965 100644 --- a/tests/test_litellm/llms/vertex_ai/rerank/test_vertex_ai_rerank_integration.py +++ b/tests/test_litellm/llms/vertex_ai/rerank/test_vertex_ai_rerank_integration.py @@ -3,26 +3,16 @@ Integration tests for Vertex AI rerank functionality. These tests demonstrate end-to-end usage of the Vertex AI rerank feature. """ -import importlib from unittest.mock import MagicMock import httpx +from litellm.llms.vertex_ai.rerank.transformation import VertexAIRerankConfig + class TestVertexAIRerankIntegration: def setup_method(self): - # Reload modules to ensure fresh references after conftest reloads litellm. - # This ensures the class being patched is the same one used by the tests. - import litellm.llms.vertex_ai.rerank.transformation as rerank_transformation_module - - importlib.reload(rerank_transformation_module) - - # Re-import after reload to get the fresh class - from litellm.llms.vertex_ai.rerank.transformation import ( - VertexAIRerankConfig as FreshConfig, - ) - - self.config = FreshConfig() + self.config = VertexAIRerankConfig() self.model = "semantic-ranker-default@latest" def test_end_to_end_rerank_flow(self): diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_token_exchange.py b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_token_exchange.py deleted file mode 100644 index d2aa58e29ea..00000000000 --- a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_token_exchange.py +++ /dev/null @@ -1,539 +0,0 @@ -""" -Tests for OAuth 2.0 Token Exchange (RFC 8693) handler for MCP servers. - -Covers: exchange flow, caching, error handling, resolve_mcp_auth integration, -bearer token extraction, and config loading. -""" - -from unittest.mock import AsyncMock, MagicMock, patch - -import httpx -import pytest - -from litellm.proxy._experimental.mcp_server.auth.token_exchange import ( - TOKEN_EXCHANGE_GRANT_TYPE, - TokenExchangeHandler, -) -from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( - MCPServerManager, -) -from litellm.proxy._experimental.mcp_server.oauth2_token_cache import ( - resolve_mcp_auth, -) -from litellm.proxy._types import LiteLLM_MCPServerTable, MCPTransport -from litellm.types.mcp import MCPAuth -from litellm.types.mcp_server.mcp_server_manager import MCPServer - - -def _obo_server(**overrides) -> MCPServer: - defaults = dict( - server_id="srv-obo-1", - name="test-obo", - url="https://mcp.example.com/mcp", - transport=MCPTransport.http, - auth_type=MCPAuth.oauth2_token_exchange, - client_id="litellm-client-id", - client_secret="litellm-client-secret", - token_exchange_endpoint="https://idp.example.com/oauth2/token", - audience="api://mcp-server", - scopes=["mcp.tools.read", "mcp.tools.execute"], - ) - defaults.update(overrides) - return MCPServer(**defaults) - - -def _exchange_response(token="exchanged-tok-abc", expires_in=3600): - resp = MagicMock() - resp.json.return_value = { - "access_token": token, - "token_type": "Bearer", - "expires_in": expires_in, - } - resp.raise_for_status = MagicMock() - resp.text = "" - return resp - - -# ── Exchange Flow ── - - -@pytest.mark.asyncio -async def test_exchange_token_success(): - """Token exchange sends correct RFC 8693 parameters and returns access_token.""" - handler = TokenExchangeHandler() - server = _obo_server() - mock_client = AsyncMock() - mock_client.post.return_value = _exchange_response("scoped-token-1") - - with patch( - "litellm.proxy._experimental.mcp_server.auth.token_exchange.get_async_httpx_client", - return_value=mock_client, - ): - result = await handler.exchange_token("user-jwt-xyz", server) - - assert result == "scoped-token-1" - mock_client.post.assert_called_once() - - _, kwargs = mock_client.post.call_args - data = kwargs["data"] - assert data["grant_type"] == TOKEN_EXCHANGE_GRANT_TYPE - assert data["subject_token"] == "user-jwt-xyz" - assert data["subject_token_type"] == "urn:ietf:params:oauth:token-type:access_token" - assert data["audience"] == "api://mcp-server" - assert data["scope"] == "mcp.tools.read mcp.tools.execute" - assert data["client_id"] == "litellm-client-id" - assert data["client_secret"] == "litellm-client-secret" - - -@pytest.mark.asyncio -async def test_exchange_token_no_audience(): - """When audience is None, it is omitted from the request.""" - handler = TokenExchangeHandler() - server = _obo_server(audience=None) - mock_client = AsyncMock() - mock_client.post.return_value = _exchange_response() - - with patch( - "litellm.proxy._experimental.mcp_server.auth.token_exchange.get_async_httpx_client", - return_value=mock_client, - ): - await handler.exchange_token("user-jwt", server) - - _, kwargs = mock_client.post.call_args - assert "audience" not in kwargs["data"] - - -@pytest.mark.asyncio -async def test_exchange_token_no_scopes(): - """When scopes is None, scope param is omitted from the request.""" - handler = TokenExchangeHandler() - server = _obo_server(scopes=None) - mock_client = AsyncMock() - mock_client.post.return_value = _exchange_response() - - with patch( - "litellm.proxy._experimental.mcp_server.auth.token_exchange.get_async_httpx_client", - return_value=mock_client, - ): - await handler.exchange_token("user-jwt", server) - - _, kwargs = mock_client.post.call_args - assert "scope" not in kwargs["data"] - - -# ── Caching ── - - -@pytest.mark.asyncio -async def test_exchange_token_cached(): - """Second call with same user token uses cache — only 1 HTTP POST.""" - handler = TokenExchangeHandler() - server = _obo_server() - mock_client = AsyncMock() - mock_client.post.return_value = _exchange_response("cached-exchange-tok") - - with patch( - "litellm.proxy._experimental.mcp_server.auth.token_exchange.get_async_httpx_client", - return_value=mock_client, - ): - t1 = await handler.exchange_token("same-jwt", server) - t2 = await handler.exchange_token("same-jwt", server) - - assert t1 == t2 == "cached-exchange-tok" - assert mock_client.post.call_count == 1 - - -@pytest.mark.asyncio -async def test_different_user_tokens_not_shared(): - """Different user JWTs get different exchanged tokens.""" - handler = TokenExchangeHandler() - server = _obo_server() - call_count = 0 - - async def mock_post(url, data=None): - nonlocal call_count - call_count += 1 - resp = MagicMock() - resp.json.return_value = { - "access_token": f"exchanged-{call_count}", - "expires_in": 3600, - } - resp.raise_for_status = MagicMock() - return resp - - mock_client = AsyncMock() - mock_client.post = mock_post - - with patch( - "litellm.proxy._experimental.mcp_server.auth.token_exchange.get_async_httpx_client", - return_value=mock_client, - ): - t1 = await handler.exchange_token("user-a-jwt", server) - t2 = await handler.exchange_token("user-b-jwt", server) - - assert t1 == "exchanged-1" - assert t2 == "exchanged-2" - assert call_count == 2 - - -# ── Error Handling ── - - -@pytest.mark.asyncio -async def test_exchange_token_http_error(): - """HTTP errors from the IDP are wrapped in a ValueError.""" - handler = TokenExchangeHandler() - server = _obo_server() - mock_response = MagicMock() - mock_response.status_code = 400 - mock_response.text = "invalid_grant" - mock_response.raise_for_status.side_effect = httpx.HTTPStatusError( - "Bad Request", - request=MagicMock(), - response=mock_response, - ) - mock_client = AsyncMock() - mock_client.post.return_value = mock_response - - with ( - patch( - "litellm.proxy._experimental.mcp_server.auth.token_exchange.get_async_httpx_client", - return_value=mock_client, - ), - pytest.raises(ValueError, match="failed with status 400"), - ): - await handler.exchange_token("bad-jwt", server) - - -@pytest.mark.asyncio -async def test_exchange_token_http_error_does_not_log_response_body(): - """Raw IDP error bodies are not logged because they can contain credentials.""" - handler = TokenExchangeHandler() - server = _obo_server() - raw_response_body = "client_secret=do-not-log" - mock_response = MagicMock() - mock_response.status_code = 401 - mock_response.text = raw_response_body - mock_response.raise_for_status.side_effect = httpx.HTTPStatusError( - "Unauthorized", - request=MagicMock(), - response=mock_response, - ) - mock_client = AsyncMock() - mock_client.post.return_value = mock_response - - with ( - patch( - "litellm.proxy._experimental.mcp_server.auth.token_exchange.get_async_httpx_client", - return_value=mock_client, - ), - patch( - "litellm.proxy._experimental.mcp_server.auth.token_exchange.verbose_logger.debug" - ) as mock_debug, - pytest.raises(ValueError, match="failed with status 401"), - ): - await handler.exchange_token("bad-jwt", server) - - logged_values = " ".join( - str(value) - for call in mock_debug.call_args_list - for value in [*call.args, *call.kwargs.values()] - ) - assert raw_response_body not in logged_values - - -@pytest.mark.asyncio -async def test_exchange_token_missing_access_token(): - """Response without access_token raises ValueError.""" - handler = TokenExchangeHandler() - server = _obo_server() - resp = MagicMock() - resp.json.return_value = {"token_type": "Bearer"} - resp.raise_for_status = MagicMock() - mock_client = AsyncMock() - mock_client.post.return_value = resp - - with ( - patch( - "litellm.proxy._experimental.mcp_server.auth.token_exchange.get_async_httpx_client", - return_value=mock_client, - ), - pytest.raises(ValueError, match="missing 'access_token'"), - ): - await handler.exchange_token("jwt", server) - - -@pytest.mark.asyncio -async def test_exchange_token_missing_endpoint(): - """Missing token_exchange_endpoint and token_url raises ValueError.""" - handler = TokenExchangeHandler() - server = _obo_server(token_exchange_endpoint=None, token_url=None) - - with pytest.raises(ValueError, match="no token_exchange_endpoint or token_url"): - await handler.exchange_token("jwt", server) - - -@pytest.mark.asyncio -async def test_exchange_token_missing_credentials(): - """Missing client_id or client_secret raises ValueError.""" - handler = TokenExchangeHandler() - server = _obo_server(client_id=None, client_secret=None) - # has_token_exchange_config will be False, so we call _do_exchange directly - with pytest.raises(ValueError, match="missing client_id or client_secret"): - await handler._do_exchange("jwt", server) - - -# ── resolve_mcp_auth Integration ── - - -@pytest.mark.asyncio -async def test_resolve_mcp_auth_with_token_exchange(): - """resolve_mcp_auth delegates to token exchange when server has OBO config and subject_token provided.""" - server = _obo_server() - mock_handler = AsyncMock() - mock_handler.exchange_token.return_value = "obo-scoped-token" - - with patch( - "litellm.proxy._experimental.mcp_server.auth.token_exchange.mcp_token_exchange_handler", - mock_handler, - ): - result = await resolve_mcp_auth(server, subject_token="user-jwt") - - assert result == "obo-scoped-token" - mock_handler.exchange_token.assert_called_once_with("user-jwt", server) - - -@pytest.mark.asyncio -async def test_resolve_mcp_auth_obo_without_subject_token_falls_through(): - """Without a subject_token, resolve_mcp_auth falls through to client_credentials.""" - server = _obo_server( - token_url="https://auth.example.com/token", - ) - mock_client = AsyncMock() - mock_client.post.return_value = _exchange_response("cc-token") - - with patch( - "litellm.proxy._experimental.mcp_server.oauth2_token_cache.get_async_httpx_client", - return_value=mock_client, - ): - result = await resolve_mcp_auth(server, subject_token=None) - - # Falls through to client_credentials since subject_token is None - # The server has client_id/client_secret/token_url so has_client_credentials is True - assert result == "cc-token" - - -@pytest.mark.asyncio -async def test_resolve_mcp_auth_obo_without_subject_token_uses_cached_client_credentials(): - """The M2M fallback for OBO servers reuses the client_credentials cache.""" - server = _obo_server( - server_id="srv-obo-m2m-cache", - token_url="https://auth.example.com/token", - ) - mock_client = AsyncMock() - mock_client.post.return_value = _exchange_response("cached-cc-token") - - with patch( - "litellm.proxy._experimental.mcp_server.oauth2_token_cache.get_async_httpx_client", - return_value=mock_client, - ): - first = await resolve_mcp_auth(server, subject_token=None) - second = await resolve_mcp_auth(server, subject_token=None) - - assert first == second == "cached-cc-token" - mock_client.post.assert_called_once() - - -@pytest.mark.asyncio -async def test_resolve_mcp_auth_header_beats_obo(): - """An explicit mcp_auth_header takes priority over OBO token exchange.""" - server = _obo_server() - result = await resolve_mcp_auth( - server, mcp_auth_header="Bearer override", subject_token="user-jwt" - ) - assert result == "Bearer override" - - -# ── Bearer Token Extraction ── - - -def test_extract_bearer_token_from_oauth2_headers(): - """Extracts token from oauth2_headers Authorization header.""" - result = MCPServerManager._extract_bearer_token( - oauth2_headers={"Authorization": "Bearer my-jwt-token"}, - raw_headers=None, - ) - assert result == "my-jwt-token" - - -def test_extract_bearer_token_from_raw_headers(): - """Falls back to raw_headers when oauth2_headers missing.""" - result = MCPServerManager._extract_bearer_token( - oauth2_headers=None, - raw_headers={"authorization": "Bearer raw-jwt"}, - ) - assert result == "raw-jwt" - - -def test_extract_bearer_token_no_bearer_prefix(): - """Returns token as-is when no Bearer prefix.""" - result = MCPServerManager._extract_bearer_token( - oauth2_headers={"Authorization": "some-opaque-token"}, - raw_headers=None, - ) - assert result == "some-opaque-token" - - -def test_extract_bearer_token_none(): - """Returns None when no auth headers present.""" - result = MCPServerManager._extract_bearer_token( - oauth2_headers=None, - raw_headers=None, - ) - assert result is None - - -# ── MCPServer Properties ── - - -def test_has_token_exchange_config_true(): - """has_token_exchange_config is True for a fully configured OBO server.""" - server = _obo_server() - assert server.has_token_exchange_config is True - - -def test_has_token_exchange_config_false_wrong_auth_type(): - """has_token_exchange_config is False when auth_type is not oauth2_token_exchange.""" - server = _obo_server(auth_type=MCPAuth.oauth2) - assert server.has_token_exchange_config is False - - -def test_has_token_exchange_config_false_missing_creds(): - """has_token_exchange_config is False when client_id/client_secret missing.""" - server = _obo_server(client_id=None) - assert server.has_token_exchange_config is False - - -def test_has_token_exchange_config_uses_token_url_fallback(): - """has_token_exchange_config is True when token_url is set instead of token_exchange_endpoint.""" - server = _obo_server( - token_exchange_endpoint=None, - token_url="https://idp.example.com/token", - ) - assert server.has_token_exchange_config is True - - -# ── Config Loading ── - - -@pytest.mark.asyncio -async def test_config_loading_token_exchange_fields(): - """load_servers_from_config correctly maps OBO config fields to MCPServer.""" - manager = MCPServerManager() - config = { - "my_obo_server": { - "url": "https://mcp.example.com/mcp", - "transport": "http", - "auth_type": "oauth2_token_exchange", - "client_id": "my-client", - "client_secret": "my-secret", - "token_exchange_endpoint": "https://idp.example.com/oauth2/token", - "audience": "api://my-mcp", - "scopes": ["read", "write"], - "subject_token_type": "urn:ietf:params:oauth:token-type:jwt", - } - } - await manager.load_servers_from_config(config) - - servers = list(manager.config_mcp_servers.values()) - assert len(servers) == 1 - - server = servers[0] - assert server.auth_type == MCPAuth.oauth2_token_exchange - assert server.token_exchange_endpoint == "https://idp.example.com/oauth2/token" - assert server.audience == "api://my-mcp" - assert server.subject_token_type == "urn:ietf:params:oauth:token-type:jwt" - assert server.client_id == "my-client" - assert server.client_secret == "my-secret" - assert server.scopes == ["read", "write"] - assert server.has_token_exchange_config is True - - -@pytest.mark.asyncio -async def test_config_loading_default_subject_token_type(): - """subject_token_type defaults to access_token when not specified in config.""" - manager = MCPServerManager() - config = { - "obo_defaults": { - "url": "https://mcp.example.com/mcp", - "transport": "http", - "auth_type": "oauth2_token_exchange", - "client_id": "cid", - "client_secret": "csec", - "token_exchange_endpoint": "https://idp.example.com/token", - } - } - await manager.load_servers_from_config(config) - - server = list(manager.config_mcp_servers.values())[0] - assert server.subject_token_type == "urn:ietf:params:oauth:token-type:access_token" - - -@pytest.mark.asyncio -async def test_database_loading_token_exchange_scopes_from_credentials(): - """DB-loaded OBO server credentials retain configured scopes.""" - manager = MCPServerManager() - db_server = LiteLLM_MCPServerTable( - server_id="srv-obo-db", - server_name="obo_db_server", - url="https://mcp.example.com/mcp", - transport=MCPTransport.http, - auth_type=MCPAuth.oauth2_token_exchange, - credentials={ - "client_id": "db-client", - "client_secret": "db-secret", - "token_exchange_endpoint": "https://idp.example.com/oauth2/token", - "audience": "api://db-mcp", - "scopes": ["db.read", "db.write"], - }, - ) - - server = await manager.build_mcp_server_from_table( - db_server, - credentials_are_encrypted=False, - ) - - assert server.auth_type == MCPAuth.oauth2_token_exchange - assert server.client_id == "db-client" - assert server.client_secret == "db-secret" - assert server.token_exchange_endpoint == "https://idp.example.com/oauth2/token" - assert server.audience == "api://db-mcp" - assert server.scopes == ["db.read", "db.write"] - - -@pytest.mark.asyncio -async def test_exchange_token_uses_client_secret_basic_when_configured(): - """LIT-4091: token exchange with token_endpoint_auth_method=client_secret_basic sends the - client credentials as HTTP Basic and omits client_secret from the body.""" - import base64 - - handler = TokenExchangeHandler() - server = _obo_server( - server_id="srv-obo-basic", token_endpoint_auth_method="client_secret_basic" - ) - mock_client = AsyncMock() - mock_client.post.return_value = _exchange_response("scoped-basic") - - with patch( - "litellm.proxy._experimental.mcp_server.auth.token_exchange.get_async_httpx_client", - return_value=mock_client, - ): - result = await handler.exchange_token("user-jwt-basic", server) - - assert result == "scoped-basic" - _, kwargs = mock_client.post.call_args - expected = "Basic " + base64.b64encode(b"litellm-client-id:litellm-client-secret").decode() - assert kwargs["headers"]["Authorization"] == expected - assert "client_secret" not in kwargs["data"] - assert "client_id" not in kwargs["data"] - assert kwargs["data"]["grant_type"] == TOKEN_EXCHANGE_GRANT_TYPE diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py index 7b05b8c9dd0..b3c0dcd1681 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py @@ -15,6 +15,7 @@ from starlette.datastructures import Headers from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import ( MCPRequestHandler, + _is_mcp_admitted_user_subject, ) from litellm.proxy._types import ( LiteLLM_ObjectPermissionTable, @@ -206,6 +207,31 @@ class TestMCPRequestHandler: assert sorted(result) == sorted(expected) + async def test_admitted_subject_not_zeroed_by_require_key_mcp_access_defined(self): + """10x-flow regression: with require_key_mcp_access_defined ON (team = ceiling for keys), a + keyless gateway/bridge-admitted subject whose ONLY access path is team membership must still + inherit the team's servers. The flag zeros empty *virtual keys* that must declare their own + access; a keyless admitted user has no key to declare it on, so it must not be zeroed.""" + auth = UserAPIKeyAuth(api_key=None, user_id="sso-user") + auth.mcp_admitted_user_subject = True + with ( + patch.object( + MCPRequestHandler, "_get_allowed_mcp_servers_for_key", new_callable=AsyncMock, return_value=[] + ), + patch.object( + MCPRequestHandler, + "_get_allowed_mcp_servers_for_team", + new_callable=AsyncMock, + return_value=["team_server1", "team_server2"], + ), + patch.object( + MCPRequestHandler, "_get_key_access_group_mcp_server_extras", new_callable=AsyncMock, return_value=[] + ), + patch("litellm.proxy.proxy_server.general_settings", {"require_key_mcp_access_defined": True}), + ): + result = await MCPRequestHandler.get_allowed_mcp_servers(auth) + assert sorted(result) == ["team_server1", "team_server2"] + @pytest.mark.parametrize( "key_servers,grants,expected,scenario", [ @@ -5302,6 +5328,7 @@ class TestMCPDcrBridgeDelegateAdmission: self._patch_user_reload( return_value=MagicMock( user_id="sso-user-7", + organization_id=None, metadata={"scim_active": True}, user_role=None, object_permission=None, @@ -5344,6 +5371,7 @@ class TestMCPDcrBridgeDelegateAdmission: self._patch_user_reload( return_value=MagicMock( user_id="sso-user-7", + organization_id=None, metadata={"scim_active": True}, user_role=None, object_permission=object_permission, @@ -5421,7 +5449,9 @@ class TestMCPDcrBridgeDelegateAdmission: with ( patch("litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager") as mock_mgr, patch("litellm.proxy.proxy_server.master_key", self._MASTER_KEY), - self._patch_user_reload(return_value=MagicMock(user_id="offboarded-user", metadata={"scim_active": False})), + self._patch_user_reload( + return_value=MagicMock(user_id="offboarded-user", organization_id=None, metadata={"scim_active": False}) + ), ): mock_mgr.get_mcp_server_by_name.return_value = self._bridge_delegate_server() with pytest.raises(HTTPException) as exc_info: @@ -6204,7 +6234,9 @@ class TestAggregateGatewayDcrChallenge: with pytest.raises(HTTPException) as exc_info: await MCPRequestHandler.process_mcp_request(self._scope()) www_authenticate = (exc_info.value.headers or {})["WWW-Authenticate"] - assert 'resource_metadata="http://testserver/.well-known/oauth-protected-resource/litellm/mcp"' in www_authenticate + assert ( + 'resource_metadata="http://testserver/.well-known/oauth-protected-resource/litellm/mcp"' in www_authenticate + ) async def test_no_challenge_for_explicit_litellm_key(self): """An explicit x-litellm-api-key declares a litellm-key client; a typo @@ -6225,9 +6257,7 @@ class TestAggregateGatewayDcrChallenge: patch(self._AUTH_PATCH_TARGET, side_effect=self._auth_401()), ): with pytest.raises(ProxyException): - await MCPRequestHandler.process_mcp_request( - self._scope(extra_headers=((b"x-mcp-servers", b"github"),)) - ) + await MCPRequestHandler.process_mcp_request(self._scope(extra_headers=((b"x-mcp-servers", b"github"),))) async def test_no_challenge_for_path_named_server(self): """/mcp/{server} targets one server; the aggregate challenge must not @@ -6261,3 +6291,1357 @@ class TestAggregateGatewayDcrChallenge: with pytest.raises(ProxyException) as exc_info: await MCPRequestHandler.process_mcp_request(self._scope()) assert str(exc_info.value.code) == "500" + + +@pytest.mark.asyncio +class TestGatewaySessionAdmission: + """The aggregate /mcp session-bearer admission arm (mcp_gateway_dcr). A valid session + token admits under the LIVE litellm user it references; an invalid/expired/refresh/foreign + token fails closed with the aggregate invalid_token challenge; the arm fires ONLY at the + aggregate scope, never for named servers or per-server flows.""" + + _MASTER_KEY = "sk-gateway-session-admission-master-key" + + def _session_bearer(self, user_id="sso-user-42", client_id="llm_dcrc_abc"): + from litellm.proxy._experimental.mcp_server.outbound_credentials.session_credentials import ( + session_keys_from_master_key, + ) + from litellm.proxy._experimental.mcp_server.outbound_credentials.session_token import ( + SessionPrincipal, + mint_session_token, + mint_session_refresh_token, + ) + + keys = session_keys_from_master_key(self._MASTER_KEY) + principal = SessionPrincipal(user_id=user_id, client_id=client_id) + return mint_session_token, mint_session_refresh_token, principal, keys + + def _access_token(self, **kw): + from datetime import datetime, timezone + + mint, _refresh, principal, keys = self._session_bearer(**kw) + return mint(principal, keys, datetime(2030, 1, 1, tzinfo=timezone.utc)).token.get_secret_value() + + def _scope(self, bearer, path="/mcp", extra_headers=()): + return { + "type": "http", + "method": "POST", + "path": path, + "headers": [(b"host", b"testserver"), (b"authorization", f"Bearer {bearer}".encode()), *extra_headers], + } + + @staticmethod + @contextlib.contextmanager + def _patch_user_reload(*, user_id, active=True, organization_id=None, tpm_limit=None, rpm_limit=None): + get_user_object = AsyncMock( + return_value=MagicMock( + user_id=user_id, + organization_id=organization_id, + metadata={"scim_active": active} if not active else {"scim_active": True}, + user_role=None, + object_permission=None, + object_permission_id=None, + tpm_limit=tpm_limit, + rpm_limit=rpm_limit, + ) + ) + with ( + patch("litellm.proxy.auth.auth_checks.get_user_object", get_user_object), + patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), + patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()), + ): + yield get_user_object + + async def test_session_admission_binds_org_id_so_the_org_ceiling_applies(self): + """The admitted auth carries the user's org_id, so get_allowed_mcp_servers keeps the + org-level MCP ceiling in force for a gateway session instead of skipping it.""" + token = self._access_token(user_id="org-user") + with ( + patch("litellm.proxy.proxy_server.master_key", self._MASTER_KEY), + patch( + "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.user_api_key_auth", + new_callable=AsyncMock, + ), + self._patch_user_reload(user_id="org-user", organization_id="org-123"), + ): + auth_result, *_rest = await MCPRequestHandler.process_mcp_request(self._scope(token)) + assert auth_result.org_id == "org-123" + + async def test_session_admission_copies_user_rate_limits(self): + """Security regression: the reconstructed auth must carry the live user's RPM/TPM, exactly as + the standard user-subject path does. The parallel limiter reads these off the auth object and + treats None as unlimited, so a keyless subject with them unset would invoke tools past their + configured user rate limits.""" + token = self._access_token(user_id="rl-user") + with ( + patch("litellm.proxy.proxy_server.master_key", self._MASTER_KEY), + patch( + "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.user_api_key_auth", + new_callable=AsyncMock, + ), + self._patch_user_reload(user_id="rl-user", tpm_limit=1000, rpm_limit=50), + ): + auth_result, *_rest = await MCPRequestHandler.process_mcp_request(self._scope(token)) + assert auth_result.user_tpm_limit == 1000 + assert auth_result.user_rpm_limit == 50 + + async def test_valid_session_admits_under_live_user_at_aggregate_scope(self): + token = self._access_token(user_id="sso-user-42") + with ( + patch("litellm.proxy.proxy_server.master_key", self._MASTER_KEY), + patch( + "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.user_api_key_auth", + new_callable=AsyncMock, + ) as mock_auth, + self._patch_user_reload(user_id="sso-user-42") as get_user_object, + ): + auth_result, _h, _servers, mcp_server_auth_headers, _o, _r = await MCPRequestHandler.process_mcp_request( + self._scope(token) + ) + assert get_user_object.await_args.kwargs["user_id"] == "sso-user-42" + assert auth_result.user_id == "sso-user-42" + mock_auth.assert_not_called() + # Identity-only admission injects no per-server upstream credential (unlike the + # bridge envelope arm); the headers dict is whatever the request carried, here empty. + assert not mcp_server_auth_headers + + @pytest.mark.parametrize( + "scenario, expect_challenge", + [("expired", True), ("tampered", False), ("refresh_at_tool_edge", False), ("foreign_key", False)], + ) + async def test_bad_session_bearer_fails_closed(self, scenario, expect_challenge): + # Every non-admissible session-shaped bearer fails closed with 401; a valid-but-unusable one + # (expired) additionally carries the invalid_token challenge so the DCR client re-authorizes. + from datetime import datetime, timezone + + if scenario == "expired": + mint, _refresh, principal, keys = self._session_bearer() + bearer = mint(principal, keys, datetime(2020, 1, 1, tzinfo=timezone.utc)).token.get_secret_value() + elif scenario == "tampered": + token = self._access_token() + bearer = token[:-3] + ("aaa" if not token.endswith("aaa") else "bbb") + elif scenario == "refresh_at_tool_edge": + _mint, refresh, principal, keys = self._session_bearer() + bearer = refresh(principal, keys, datetime(2030, 1, 1, tzinfo=timezone.utc)).token.get_secret_value() + else: # foreign_key: minted under the real master key, presented while the proxy uses another + bearer = self._access_token() + master_key = "sk-a-totally-different-master-key" if scenario == "foreign_key" else self._MASTER_KEY + with patch("litellm.proxy.proxy_server.master_key", master_key): + with pytest.raises(HTTPException) as exc_info: + await MCPRequestHandler.process_mcp_request(self._scope(bearer)) + assert exc_info.value.status_code == 401 + if expect_challenge: + assert 'error="invalid_token"' in (exc_info.value.headers or {})["WWW-Authenticate"] + + async def test_deactivated_user_fails_with_invalid_token_challenge(self): + """A cryptographically valid bearer whose referenced user is SCIM-deactivated must fail with + the aggregate invalid_token challenge (WWW-Authenticate), matching the expired/tampered arms, + so the DCR client re-authorizes instead of getting a bare 401 with no challenge.""" + token = self._access_token(user_id="offboarded-user") + with ( + patch("litellm.proxy.proxy_server.master_key", self._MASTER_KEY), + patch( + "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.user_api_key_auth", + new_callable=AsyncMock, + ), + self._patch_user_reload(user_id="offboarded-user", active=False), + ): + with pytest.raises(HTTPException) as exc_info: + await MCPRequestHandler.process_mcp_request(self._scope(token)) + assert exc_info.value.status_code == 401 + assert 'error="invalid_token"' in (exc_info.value.headers or {})["WWW-Authenticate"] + + async def test_session_bearer_scrubbed_from_egress_header_contexts(self): + """Security regression (credential leak): after a keyless session admission, the session + bearer must be removed from BOTH returned egress header contexts (oauth2_headers and the raw + headers) so no passthrough/OBO egress can forward it upstream for replay as this user.""" + token = self._access_token(user_id="sso-user-42") + with ( + patch("litellm.proxy.proxy_server.master_key", self._MASTER_KEY), + patch( + "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.user_api_key_auth", + new_callable=AsyncMock, + ), + self._patch_user_reload(user_id="sso-user-42"), + ): + _auth, _h, _servers, _msah, oauth2_headers, raw_headers = await MCPRequestHandler.process_mcp_request( + self._scope(token) + ) + # the request carried "Authorization: Bearer "; both egress contexts must be scrubbed + assert oauth2_headers is None + assert not any(k.lower() == "authorization" for k in (raw_headers or {})) + + async def test_arm_does_not_fire_for_named_server(self): + """A session-shaped bearer aimed at a named server (path scope) does not enter the + aggregate arm; it is treated as an ordinary bearer on that server.""" + token = self._access_token() + with ( + patch("litellm.proxy.proxy_server.master_key", self._MASTER_KEY), + patch( + "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.user_api_key_auth", + new_callable=AsyncMock, + side_effect=ProxyException(message="bad key", type="auth_error", param="api_key", code=401), + ) as mock_auth, + ): + with pytest.raises((HTTPException, ProxyException)): + await MCPRequestHandler.process_mcp_request(self._scope(token, path="/mcp/github")) + mock_auth.assert_called_once() + + +def _make_team(team_id, mcp_servers, *, org_id=None, tool_perms=None, members=("sso-user",)): + from litellm.proxy._types import LiteLLM_ObjectPermissionTable, LiteLLM_TeamTable, Member + + return LiteLLM_TeamTable( + team_id=team_id, + organization_id=org_id, + members_with_roles=[Member(user_id=u, role="user") for u in members], + access_group_ids=[], + object_permission=LiteLLM_ObjectPermissionTable( + object_permission_id=f"op-{team_id}", mcp_servers=mcp_servers, mcp_tool_permissions=tool_perms + ), + ) + + +def _make_admitted_subject(user_id, *, org_id=None, own_servers=None, own_tool_perms=None): + from litellm.proxy._types import LiteLLM_ObjectPermissionTable + + op = None + if own_servers is not None or own_tool_perms is not None: + op = LiteLLM_ObjectPermissionTable( + object_permission_id=f"userop-{user_id}", + mcp_servers=own_servers or [], + mcp_tool_permissions=own_tool_perms, + ) + auth = UserAPIKeyAuth(user_id=user_id, api_key=None, org_id=org_id, object_permission=op) + auth.mcp_admitted_user_subject = True + return auth + + +@pytest.mark.asyncio +class TestUserSubjectTeamUnion: + """_get_allowed_mcp_servers_for_team unions across ALL a user's teams for a keyless + user-subject caller (the gateway DCR session bearer and bridge user-envelope), while a + key-based caller keeps its single-team behavior byte-identically.""" + + @contextlib.contextmanager + def _patch(self, *, teams_by_id, user_teams=None, orgs_by_id=None): + async def _get_team_object(team_id, **kw): + return teams_by_id.get(team_id) + + async def _get_user_object(user_id, **kw): + return MagicMock(user_id=user_id, teams=user_teams or []) + + async def _get_org_object(org_id, **kw): + return (orgs_by_id or {}).get(org_id) + + async def _spend_from_fallback(counter_key, fallback_spend, max_budget=None, **kw): + # The budget owners read cross-pod spend Redis-first with the row's spend as fallback; + # unit tests have no Redis, so the fallback IS the spend. + return fallback_spend + + with ( + patch("litellm.proxy.auth.auth_checks.get_team_object", _get_team_object), + patch("litellm.proxy.auth.auth_checks.get_user_object", _get_user_object), + patch("litellm.proxy.auth.auth_checks.get_org_object", _get_org_object), + patch("litellm.proxy.auth.auth_checks._get_mcp_server_ids_from_access_groups", AsyncMock(return_value=[])), + patch("litellm.proxy.proxy_server.get_current_spend", _spend_from_fallback), + patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), + patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()), + patch("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock()), + ): + yield + + async def test_keyless_user_unions_servers_across_all_their_teams(self): + teams = {"team-a": _make_team("team-a", ["srv1", "srv2"]), "team-b": _make_team("team-b", ["srv2", "srv3"])} + auth = _make_admitted_subject("sso-user") + with self._patch(teams_by_id=teams, user_teams=["team-a", "team-b"]): + result = await MCPRequestHandler.get_allowed_mcp_servers(auth) + assert set(result) == {"srv1", "srv2", "srv3"} + + async def test_key_based_caller_uses_single_team_only(self): + """A key-based caller (api_key set) with a team_id sees ONLY that team, even though the + same user belongs to other teams: key auth must be byte-identical to before.""" + teams = {"team-a": _make_team("team-a", ["srv1"]), "team-b": _make_team("team-b", ["srv2", "srv3"])} + auth = UserAPIKeyAuth(user_id="sso-user", api_key="sk-hash", team_id="team-a") + with self._patch(teams_by_id=teams, user_teams=["team-a", "team-b"]): + result = await MCPRequestHandler._get_allowed_mcp_servers_for_team(auth) + assert set(result) == {"srv1"} + + async def test_keyless_user_with_explicit_team_id_uses_that_team_only(self): + """A keyless caller that already pins a team_id (not the user-subject fan-out shape) + resolves only that team; the union is strictly for the no-team-id user-subject case.""" + teams = {"team-a": _make_team("team-a", ["srv1"]), "team-b": _make_team("team-b", ["srv2"])} + auth = UserAPIKeyAuth(user_id="sso-user", api_key=None, team_id="team-a") + with self._patch(teams_by_id=teams, user_teams=["team-a", "team-b"]): + result = await MCPRequestHandler._get_allowed_mcp_servers_for_team(auth) + assert set(result) == {"srv1"} + + async def test_keyless_user_with_no_teams_gets_nothing_from_teams(self): + auth = _make_admitted_subject("lonely-user") + with self._patch(teams_by_id={}, user_teams=[]): + result = await MCPRequestHandler._get_allowed_mcp_servers_for_team(auth) + assert result == [] + + async def test_ui_session_team_id_still_resolves_to_nothing(self): + from litellm.proxy._types import UI_TEAM_ID + + auth = UserAPIKeyAuth(user_id="dash-user", api_key="sk-hash", team_id=UI_TEAM_ID) + with self._patch(teams_by_id={}, user_teams=["team-a"]): + result = await MCPRequestHandler._get_allowed_mcp_servers_for_team(auth) + assert result == [] + + async def test_team_ids_helper_gates_on_shape(self): + from litellm.proxy._types import UI_TEAM_ID + + # key-based with team -> that team + assert await MCPRequestHandler._team_ids_for_mcp_grant( + UserAPIKeyAuth(api_key="sk", team_id="t1", user_id="u") + ) == ["t1"] + # An admitted subject never fans out HERE: it resolves one source per team first, and each of + # those pins a team_id, so this helper only ever answers the single-team question. The fan-out + # itself is _admitted_subject_sources' job, asserted below. + with self._patch(teams_by_id={}, user_teams=["t2", "t3"]): + assert await MCPRequestHandler._team_ids_for_mcp_grant(_make_admitted_subject("u")) == [] + # keyless, no user_id -> nothing + assert await MCPRequestHandler._team_ids_for_mcp_grant(UserAPIKeyAuth(api_key=None)) == [] + # keyless with a user_id but NOT admission-marked (JWT auth) -> nothing (unchanged behavior) + with self._patch(teams_by_id={}, user_teams=["t2", "t3"]): + assert ( + await MCPRequestHandler._team_ids_for_mcp_grant(UserAPIKeyAuth(api_key=None, user_id="jwt-user")) == [] + ) + # UI sentinel -> nothing + assert ( + await MCPRequestHandler._team_ids_for_mcp_grant( + UserAPIKeyAuth(api_key="sk", team_id=UI_TEAM_ID, user_id="u") + ) + == [] + ) + + async def test_org_outage_is_not_treated_as_a_missing_org(self): + """A CONFIRMED-absent org places no ceiling; a FAILED lookup must not be read as the same + fact. get_org_object used to relabel every error as "doesn't exist", so a DB outage silently + dropped a real org's ceiling for as long as it lasted. Absent -> the team's grant stands; + outage -> the keyless source denies.""" + from litellm.proxy.auth.auth_checks import OrganizationNotFoundError + + teams = {"t1": _make_team("t1", ["srv1"])} + teams["t1"].organization_id = "org-a" + auth = _make_admitted_subject("sso-user") + + absent = AsyncMock(side_effect=OrganizationNotFoundError("Organization doesn't exist in db.")) + with self._patch(teams_by_id=teams, user_teams=["t1"]): + with patch("litellm.proxy.auth.auth_checks.get_org_object", absent): + reachable = await MCPRequestHandler.get_allowed_mcp_servers(auth) + assert set(reachable) == {"srv1"}, "a deleted org places no ceiling" + + outage = AsyncMock(side_effect=RuntimeError("connection reset by peer")) + with self._patch(teams_by_id=teams, user_teams=["t1"]): + with patch("litellm.proxy.auth.auth_checks.get_org_object", outage): + reachable = await MCPRequestHandler.get_allowed_mcp_servers(auth) + assert reachable == [], "an unresolvable ceiling must deny a keyless source, not be skipped" + + async def test_org_ceiling_fault_fails_closed_for_admitted_but_open_for_keys(self): + """An unresolvable org ceiling is NOT the same fact as "this org places no restriction". + + For a virtual key the ceiling is one of several bounds and a DB blip must not lock working + keys out, so it stays fail-open. For a keyless admitted subject the per-source org ceiling is + the ONLY org bound, so dropping it on a fault would widen a cross-org user to servers their + team's org forbids. That is escalation, not an availability blip, so it fails closed.""" + from litellm.proxy._types import LiteLLM_ObjectPermissionTable + + boom = AsyncMock(side_effect=RuntimeError("org lookup exploded")) + # The subject must actually REACH something, or the assertion passes either way and pins + # nothing (a fail-open mutant survived an earlier version of this test for exactly that). + auth = _make_admitted_subject("sso-user") + auth.org_id = "org-a" + auth.object_permission = LiteLLM_ObjectPermissionTable(object_permission_id="op-u", mcp_servers=["srv1"]) + with self._patch(teams_by_id={}, user_teams=[]): + assert set(await MCPRequestHandler.get_allowed_mcp_servers(auth)) == {"srv1"} # control + with patch.object(MCPRequestHandler, "_get_org_object_permission", boom): + admitted = await MCPRequestHandler.get_allowed_mcp_servers(auth) + assert admitted == [], "admitted subject must fail CLOSED when its org ceiling cannot resolve" + + key_auth = UserAPIKeyAuth(user_id="u", api_key="sk-hash", team_id="t1", org_id="org-a") + with self._patch(teams_by_id={"t1": _make_team("t1", ["srv1"])}, user_teams=[]): + with patch.object(MCPRequestHandler, "_get_org_object_permission", boom): + keyed = await MCPRequestHandler.get_allowed_mcp_servers(key_auth) + assert set(keyed) == {"srv1"}, "key auth must keep its long-standing fail-open behavior" + + async def test_only_the_attributing_team_bucket_is_charged(self): + """A team's mcp_rpm_limit bounds that team's SHARED bucket. Charging every granting team let + one cross-team user drain several teams' buckets on a single call, blocking their other + members for access those teams did not provide. Exactly one source is charged, and it is the + SAME source billing picks — one owner for both, so they cannot disagree.""" + from litellm.proxy.hooks.parallel_request_limiter_v3 import _PROXY_MaxParallelRequestsHandler_v3 + + t1 = _make_team("t1", ["srv1"]) + t1.metadata = {"mcp_rpm_limit": {"srv1": 5}} + t2 = _make_team("t2", ["srv1"]) + t2.metadata = {"mcp_rpm_limit": {"srv1": 9}} + auth = _make_admitted_subject("sso-user") + with self._patch(teams_by_id={"t1": t1, "t2": t2}, user_teams=["t1", "t2"]): + auth.mcp_source_team_rpm_limits = await MCPRequestHandler._admitted_subject_team_rpm_limits(auth) + billed = await MCPRequestHandler.attributing_source_for_server(auth, "srv1") + assert auth.mcp_source_team_rpm_limits == {"t1": {"srv1": 5}}, "t2's shared bucket is untouched" + assert billed is not None and billed.team_id == "t1", "throttling and billing pick the same source" + + descriptors: list = [] + limiter = _PROXY_MaxParallelRequestsHandler_v3(internal_usage_cache=MagicMock()) + limiter._add_mcp_per_team_rate_limit_descriptor(auth, "srv1", descriptors) + charged = {d["value"]: d["rate_limit"]["requests_per_unit"] for d in descriptors} + assert charged == {"t1:srv1": 5}, "only the attributing team's bucket is charged" + + async def test_direct_user_grant_charges_no_team_bucket(self): + """When the user's OWN grant reaches the server, no team provided the access, so no team + bucket may be charged — the user's own rpm/tpm is what bounds them. Mirrors billing, which + bills the user and their own org for exactly this case.""" + t1 = _make_team("t1", ["srv1"]) + t1.metadata = {"mcp_rpm_limit": {"srv1": 5}} + auth = _make_admitted_subject("sso-user") + auth.object_permission = LiteLLM_ObjectPermissionTable(object_permission_id="op-u", mcp_servers=["srv1"]) + with self._patch(teams_by_id={"t1": t1}, user_teams=["t1"]): + limits = await MCPRequestHandler._admitted_subject_team_rpm_limits(auth) + billed = await MCPRequestHandler.attributing_source_for_server(auth, "srv1") + assert limits is None, "a direct user grant must not charge any team's shared bucket" + assert billed is None, "and billing agrees: the user is billed, not a team" + + def _manager_with(self, server_ids, allow_all=()): + from litellm.proxy._experimental.mcp_server.mcp_server_manager import MCPServerManager + from litellm.types.mcp_server.mcp_server_manager import MCPServer + from litellm.types.mcp import MCPTransport + + manager = MCPServerManager() + for sid in server_ids: + manager.registry[sid] = MCPServer( + server_id=sid, + name=sid, + server_name=sid, + url="https://example.com/mcp", + transport=MCPTransport.http, + allow_all_keys=sid in allow_all, + ) + manager._get_active_submitted_mcp_server_ids_for_user = AsyncMock(return_value=[]) + return manager + + async def test_team_derived_call_bills_the_granting_team_and_its_org(self): + """ACCOUNTING half of team budgets. Without attribution the admitted auth kept team_id=None, + so spend skipped team updates (the team's budget never accumulated, so it could never begin + to block) and charged the user's PRIMARY org rather than the org owning the granting team.""" + t_grant = _make_team("t-grant", ["srv1"]) + t_grant.organization_id = "org-team" + auth = _make_admitted_subject("sso-user") + auth.org_id = "org-user-primary" + with self._patch(teams_by_id={"t-grant": t_grant}, user_teams=["t-grant"]): + source = await MCPRequestHandler.attributing_source_for_server(auth, "srv1") + assert source is not None and source.team_id == "t-grant" + assert source.org_id == "org-team", "the granting team's org is charged, not the user's primary" + assert auth.team_id is None and auth.org_id == "org-user-primary", "authz object untouched" + + async def test_billing_auth_carries_team_and_org_onto_the_spend_object(self): + """Asserted on billing_auth_for_tool_call itself, not on the source it picks: the source + already carries the team's org by construction, so asserting there leaves the copy step + unpinned (a mutant dropping org_id survived exactly that). This is the object spend reads.""" + t_grant = _make_team("t-grant", ["srv1"]) + t_grant.organization_id = "org-team" + auth = _make_admitted_subject("sso-user") + auth.org_id = "org-user-primary" + server = MagicMock(server_id="srv1") + with self._patch(teams_by_id={"t-grant": t_grant}, user_teams=["t-grant"]): + with patch( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager._get_mcp_server_from_tool_name", + MagicMock(return_value=server), + ): + billed = await MCPRequestHandler.billing_auth_for_tool_call(auth, tool_name="t-grant/tool_a") + assert (billed.team_id, billed.org_id) == ("t-grant", "org-team") + assert (auth.team_id, auth.org_id) == (None, "org-user-primary"), "authz object must be untouched" + + async def test_own_grant_bills_the_user_not_a_team(self): + """A server the user's OWN grant reaches is not reached "through a team", so it bills the + user and their own org — attributing it to an unrelated team the user happens to belong to + would charge that team for access it never provided.""" + from litellm.proxy._types import LiteLLM_ObjectPermissionTable + + t_other = _make_team("t-other", ["srv1"]) + auth = _make_admitted_subject("sso-user") + auth.object_permission = LiteLLM_ObjectPermissionTable(object_permission_id="op-u", mcp_servers=["srv1"]) + with self._patch(teams_by_id={"t-other": t_other}, user_teams=["t-other"]): + assert await MCPRequestHandler.attributing_source_for_server(auth, "srv1") is None + + async def test_billing_attribution_is_deterministic_across_several_granting_teams(self): + """When several teams grant the same server the pick must be stable and reproducible rather + than dependent on dict/roster ordering, or the same call bills different teams run to run.""" + teams = {"t-b": _make_team("t-b", ["srv1"]), "t-a": _make_team("t-a", ["srv1"])} + auth = _make_admitted_subject("sso-user") + with self._patch(teams_by_id=teams, user_teams=["t-b", "t-a"]): + first = await MCPRequestHandler.attributing_source_for_server(auth, "srv1") + with self._patch(teams_by_id=teams, user_teams=["t-a", "t-b"]): + second = await MCPRequestHandler.attributing_source_for_server(auth, "srv1") + assert first is not None and first.team_id == "t-a" + assert second is not None and second.team_id == "t-a", "roster order must not change who is billed" + + async def test_billing_auth_leaves_non_admitted_callers_untouched(self): + """Key and JWT billing must be byte-identical: the attribution wrapper returns the very same + object for anything that is not a keyless admitted subject.""" + key_auth = UserAPIKeyAuth(user_id="u", api_key="sk-hash", team_id="t1", org_id="org-a") + assert await MCPRequestHandler.billing_auth_for_tool_call(key_auth, tool_name="srv1-tool") is key_auth + + async def test_admitted_tools_never_run_the_single_credential_prelude(self): + """ORDERING is the invariant: the admitted branch is the FIRST statement of the tools + resolver, exactly as in the servers resolver. A fault in a lookup the subject never uses + (its own mcp_toolsets) must not reach it at all — when this branch sat after the prelude, + such a fault hit the fail-closed handler and denied tools its teams did grant.""" + from litellm.proxy._types import LiteLLM_ObjectPermissionTable + + teams = {"t1": _make_team("t1", ["srv1"], tool_perms={"srv1": ["read"]})} + auth = _make_admitted_subject("sso-user") + # The subject must carry a toolset, or the prelude never resolves one and the fault below is + # unreachable — the branch could sit anywhere and the test would still pass (it did). + auth.object_permission = LiteLLM_ObjectPermissionTable(object_permission_id="op-u", mcp_toolsets=["ts-1"]) + boom = AsyncMock(side_effect=RuntimeError("toolset resolution exploded")) + with self._patch(teams_by_id=teams, user_teams=["t1"]): + with patch( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager.resolve_toolset_tool_permissions", + boom, + ): + tools = await MCPRequestHandler.get_allowed_tools_for_server("srv1", auth) + # The fault DOES fire, correctly, inside the subject's own source (which carries its + # toolsets) — that source contributes nothing. What must not happen is the top-level + # prelude running it first and denying the team's grant through the fail-closed handler. + assert tools == ["read"], "a fault in the subject's own toolsets must not deny its team's tools" + + async def test_admitted_own_byom_servers_stay_open(self): + """BYOM suppression-by-explicit-scope is a rule about a CREDENTIAL carrying its own + mcp_servers list. An admitted subject's object_permission is the user's own row, whose + mcp_servers column is [] by DB default — applying the rule would hide almost every admitted + user's OWN submitted servers. A key with an explicit scope still gets no BYOM widening.""" + from litellm.proxy._types import LiteLLM_ObjectPermissionTable + + manager = self._manager_with(["srv-byom"]) + manager._get_active_submitted_mcp_server_ids_for_user = AsyncMock(return_value=["srv-byom"]) + db_default_perm = LiteLLM_ObjectPermissionTable(object_permission_id="op-u", mcp_servers=[]) + + admitted = _make_admitted_subject("sso-user") + admitted.object_permission = db_default_perm + scoped_key = UserAPIKeyAuth(user_id="u", api_key="sk-hash", object_permission=db_default_perm) + + assert await manager.operator_open_server_ids(admitted) == {"srv-byom"} + assert await manager.operator_open_server_ids(scoped_key) == set(), "explicit key scope still suppresses BYOM" + + async def test_admitted_admin_is_scoped_to_grants_not_full_registry(self): + """The wrapper's admin short-circuit hands the FULL registry to any admin-role auth before + the grant union or the per-team org ceilings run. A session bearer is a third-party client + credential, not the dashboard: an admin signing in through the connect flow gets their + grants like anyone else. A real admin key keeps the dashboard behavior unchanged.""" + from litellm.proxy._types import LitellmUserRoles + + manager = self._manager_with(["srv-granted", "srv-secret"]) + admitted = _make_admitted_subject("admin-user") + admitted.user_role = LitellmUserRoles.PROXY_ADMIN + with patch.object(MCPRequestHandler, "get_allowed_mcp_servers", AsyncMock(return_value=["srv-granted"])): + admitted_view = set(await manager.get_allowed_mcp_servers(admitted)) + key_admin_view = set( + await manager.get_allowed_mcp_servers( + UserAPIKeyAuth(user_id="admin-user", api_key="sk-hash", user_role=LitellmUserRoles.PROXY_ADMIN) + ) + ) + assert admitted_view == {"srv-granted"}, "an admitted admin gets their grants, not the registry" + assert key_admin_view == {"srv-granted", "srv-secret"}, "admin KEY behavior must be unchanged" + + async def test_admitted_opt_out_via_wrapper_keeps_team_servers(self): + """The wrapper's no_mcp_servers early-return is a KEY rule (a scoped credential's opt-out is + absolute). The admitted subject's opt-out silences only its own source, which the resolver + enforces per source — the wrapper must defer to it, or the resolver-level rule is dead code + on the production path.""" + from litellm.proxy._types import LiteLLM_ObjectPermissionTable, SpecialMCPServerNames + + manager = self._manager_with(["srv-team"]) + opt_out = LiteLLM_ObjectPermissionTable( + object_permission_id="op-u", mcp_servers=[SpecialMCPServerNames.no_mcp_servers.value] + ) + admitted = _make_admitted_subject("sso-user") + admitted.object_permission = opt_out + with patch.object(MCPRequestHandler, "get_allowed_mcp_servers", AsyncMock(return_value=["srv-team"])): + admitted_view = set(await manager.get_allowed_mcp_servers(admitted)) + key_view = await manager.get_allowed_mcp_servers( + UserAPIKeyAuth(user_id="u", api_key="sk-hash", object_permission=opt_out) + ) + assert "srv-team" in admitted_view, "user opt-out must not zero team grants on the wrapper path" + assert key_view == [], "a key's opt-out stays absolute" + + async def test_open_channel_confers_reachability_not_a_ceiling_waiver(self): + """An open channel (allow_all_keys / own BYOM) makes a server REACHABLE. It is not a waiver + of the ceilings that bound it: the user's own mcp_tool_permissions still apply, exactly as a + virtual key's key_tools do on the same allow_all server. Returning None outright let a + session holder invoke tools their own policy excludes.""" + from litellm.proxy._types import LiteLLM_ObjectPermissionTable + + auth = _make_admitted_subject("sso-user") + # The user is restricted to `read` on srv-open, and NO grant source names that server — + # it is reachable only through the open channel, which is exactly the bypass path. + auth.object_permission = LiteLLM_ObjectPermissionTable( + object_permission_id="op-u", mcp_servers=[], mcp_tool_permissions={"srv-open": ["read"]} + ) + open_ids = AsyncMock(return_value={"srv-open"}) + with self._patch(teams_by_id={}, user_teams=[]): + with patch( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager.operator_open_server_ids", + open_ids, + ): + tools = await MCPRequestHandler.get_allowed_tools_for_server("srv-open", auth) + assert tools == ["read"], "the user's own tool policy must still bind on an open-channel server" + + async def test_open_channel_server_gets_default_open_tools_for_admitted(self): + """A server reachable through an open channel (allow_all_keys / own BYOM) is granted by NO + source, so the source union alone returns [] — listable but uninvokable. The tools axis asks + the same open-channel owner the server union uses, so the server is default-open for tools + exactly as a virtual key experiences it.""" + auth = _make_admitted_subject("sso-user") + open_ids = AsyncMock(return_value={"srv-open"}) + with self._patch(teams_by_id={}, user_teams=[]): + with patch( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager.operator_open_server_ids", + open_ids, + ): + open_tools = await MCPRequestHandler.get_allowed_tools_for_server("srv-open", auth) + closed_tools = await MCPRequestHandler.get_allowed_tools_for_server("srv-ungranted", auth) + assert open_tools is None, "open-channel server must be default-open for tools" + assert closed_tools == [], "a server no source or channel grants stays deny-all" + + async def test_over_budget_team_grants_nothing_and_healthy_team_stands(self): + """Budget ENFORCEMENT is the sibling of blocked: a team that has already exceeded its + max_budget is rejected outright for a virtual key pinned to it (common_checks), so it must + not keep granting servers, tools or throttle scope to a keyless union subject either. + Enforced through the SAME owner the key path uses (_team_max_budget_check). Distinct from + budget ATTRIBUTION of new spend, which stays with the user (documented deferral).""" + t_over = _make_team("t-over", ["srv1"]) + t_over.max_budget = 10.0 + t_over.spend = 11.0 + t_ok = _make_team("t-ok", ["srv2"]) + t_ok.max_budget = 10.0 + t_ok.spend = 1.0 + auth = _make_admitted_subject("sso-user") + with self._patch(teams_by_id={"t-over": t_over, "t-ok": t_ok}, user_teams=["t-over", "t-ok"]): + servers = set(await MCPRequestHandler.get_allowed_mcp_servers(auth)) + limits = await MCPRequestHandler._admitted_subject_team_rpm_limits(auth) + assert servers == {"srv2"}, "an over-budget team must stop granting; the healthy team stands" + assert limits is None, "an over-budget team is not a source, so it stamps no throttle either" + + async def test_team_in_over_budget_org_grants_nothing(self): + """The org axis of the same rule, judged against the TEAM's own org (not the caller's + primary): a team owned by an org over its budget grants nothing, exactly as a key in that + org is rejected by _organization_max_budget_check.""" + t_in_broke_org = _make_team("t-b", ["srv1"]) + t_in_broke_org.organization_id = "org-broke" + # object_permission_id=None: the org has NO MCP ceiling, so the source is denied by the + # budget gate alone. A truthy auto-Mock id here made an earlier version of this test pass + # through the org-CEILING fault path with the budget gate deleted — vacuous. + org = MagicMock(object_permission_id=None, litellm_budget_table=MagicMock(max_budget=5.0), spend=9.0) + auth = _make_admitted_subject("sso-user") + with self._patch(teams_by_id={"t-b": t_in_broke_org}, user_teams=["t-b"], orgs_by_id={"org-broke": org}): + servers = await MCPRequestHandler.get_allowed_mcp_servers(auth) + assert servers == [], "a team in an over-budget org must not grant through the union" + + async def test_one_faulting_team_does_not_deny_the_other_sources(self): + """The unit of fault isolation is the SOURCE. One team's row being momentarily unreadable + contributes nothing for THAT team (access only narrows) while the user's own grants and every + other resolvable team stand — it must not collapse the whole union to deny-all on either + axis.""" + from litellm.proxy._types import LiteLLM_ObjectPermissionTable + + t_ok = _make_team("t-ok", ["srv1"], tool_perms={"srv1": ["read"]}) + auth = _make_admitted_subject("sso-user") + auth.object_permission = LiteLLM_ObjectPermissionTable(object_permission_id="op-u", mcp_servers=["srv-own"]) + teams = {"t-ok": t_ok} # t-boom absent from the map -> our patched get_team_object RAISES for it + + async def _team_or_boom(team_id, **kw): + if team_id not in teams: + raise RuntimeError(f"transient DB blip loading {team_id}") + return teams[team_id] + + with self._patch(teams_by_id=teams, user_teams=["t-boom", "t-ok"]): + with patch("litellm.proxy.auth.auth_checks.get_team_object", _team_or_boom): + servers = set(await MCPRequestHandler.get_allowed_mcp_servers(auth)) + tools = await MCPRequestHandler.get_allowed_tools_for_server("srv1", auth) + assert servers == {"srv-own", "srv1"}, "healthy sources must stand when one team faults" + assert tools == ["read"], "the healthy team's tool grant must survive the other team's fault" + + async def test_key_org_tool_ceiling_fault_keeps_key_restrictions(self): + """Virtual-key tools axis mirrors its servers axis on an unresolvable org ceiling: the org + intersect is SKIPPED and the key's own tool restrictions stand. Letting the fault escape + collapsed the whole resolution to None (allow-all), which is fail-open WIDER than before the + fault — key restrictions must never be dropped by an org lookup blip.""" + from litellm.proxy._types import LiteLLM_ObjectPermissionTable + + key_auth = UserAPIKeyAuth(user_id="u", api_key="sk-hash", org_id="org-a") + key_auth.object_permission = LiteLLM_ObjectPermissionTable( + object_permission_id="op-k", mcp_servers=["srv1"], mcp_tool_permissions={"srv1": ["read"]} + ) + boom = AsyncMock(side_effect=RuntimeError("org permission load exploded")) + with self._patch(teams_by_id={}, user_teams=[]): + with patch.object(MCPRequestHandler, "_get_org_object_permission", boom): + tools = await MCPRequestHandler.get_allowed_tools_for_server("srv1", key_auth) + assert tools == ["read"], "key tool restrictions must survive an unresolvable org ceiling" + + async def test_team_rpm_limit_binds_only_within_that_teams_grant_scope(self): + """A limit rides the same scope as the access it bounds. A roster team is charged ONLY for + servers its own grant reaches: not for a server the user reaches through a DIFFERENT team + (else this user's calls drain a bucket shared by that team's keys for access the team never + provided), not for map entries beyond its grant, and never when the team is blocked.""" + # t-granting grants srv1 and limits it; also names srv9 in its map, which it does NOT grant. + t_granting = _make_team("t-granting", ["srv1"]) + t_granting.metadata = {"mcp_rpm_limit": {"srv1": 5, "srv9": 7}} + # t-other grants only srv2 but retains limit metadata for srv1 -> must not be charged for it. + t_other = _make_team("t-other", ["srv2"]) + t_other.metadata = {"mcp_rpm_limit": {"srv1": 3}} + # t-blocked grants srv1 and limits it, but is blocked -> grants nothing, charges nothing. + t_blocked = _make_team("t-blocked", ["srv1"]) + t_blocked.metadata = {"mcp_rpm_limit": {"srv1": 2}} + t_blocked.blocked = True + + auth = _make_admitted_subject("sso-user") + teams = {"t-granting": t_granting, "t-other": t_other, "t-blocked": t_blocked} + with self._patch(teams_by_id=teams, user_teams=["t-granting", "t-other", "t-blocked"]): + limits = await MCPRequestHandler._admitted_subject_team_rpm_limits(auth) + + assert limits == {"t-granting": {"srv1": 5}}, ( + "only the granting team's bucket, and only for the server it grants" + ) + + async def test_non_roster_team_rpm_limit_does_not_apply(self): + """The roster gates grants and throttles through one owner, so a team the user was removed + from neither grants servers nor gets charged for their calls.""" + stale = _make_team("t-stale", ["srv1"], members=("someone-else",)) + stale.metadata = {"mcp_rpm_limit": {"srv1": 1}} + auth = _make_admitted_subject("sso-user") + with self._patch(teams_by_id={"t-stale": stale}, user_teams=["t-stale"]): + limits = await MCPRequestHandler._admitted_subject_team_rpm_limits(auth) + assert limits is None + + async def test_org_list_caps_a_source_but_never_becomes_a_grant(self): + """The admitted model is a union of GRANTS, so an org allowlist may only narrow what a source + already grants. For a virtual key with no lower-level restriction the org list legitimately + BECOMES the allowed set, and inheriting that arm would hand every admitted user with an + org_id their whole org's server list with no direct or team grant behind it.""" + from litellm.proxy._types import LiteLLM_ObjectPermissionTable + + auth = _make_admitted_subject("sso-user") + auth.org_id = "org-a" # org allows srv1+srv2; the user and their teams grant NOTHING + org_perm = AsyncMock( + return_value=LiteLLM_ObjectPermissionTable(object_permission_id="op-org-a", mcp_servers=["srv1", "srv2"]) + ) + with self._patch(teams_by_id={}, user_teams=[]): + with patch.object(MCPRequestHandler, "_get_org_object_permission", org_perm): + result = await MCPRequestHandler.get_allowed_mcp_servers(auth) + assert result == [], "an org ceiling must not grant servers the user was never granted" + + async def test_tool_ceiling_fails_closed_when_a_SOURCE_faults(self): + """Each source is resolved through an UNMARKED auth, so a fault under a source must still + deny. Returning None there would win the union as allow-all and drop every team/org tool + ceiling on a DB blip -- the marker alone only covers faults raised before the fan-out.""" + auth = _make_admitted_subject("sso-user") + teams = {"t1": _make_team("t1", ["srv1"])} + # Fault INSIDE the tool resolution only. Faulting something the server path also uses would + # make the source grant nothing, so the union would return [] without the tool path ever + # running -- the test would pass while pinning nothing (an earlier version did exactly that). + boom = AsyncMock(side_effect=RuntimeError("org tool ceiling exploded")) + with self._patch(teams_by_id=teams, user_teams=["t1"]): + assert await MCPRequestHandler.get_allowed_mcp_servers(auth) == ["srv1"] # control: granted + with patch.object(MCPRequestHandler, "_apply_agent_and_org_tool_ceilings", boom): + tools = await MCPRequestHandler.get_allowed_tools_for_server("srv1", auth) + assert tools == [], "a source-level fault must deny tools, never collapse to allow-all" + + async def test_own_opt_out_silences_only_that_source_not_the_teams(self): + """no_mcp_servers on the USER's own grants opts that source out. It must not zero the teams: + the sources are independent, so an opt-out on one silences one. (The same sentinel on a + virtual KEY still overrides team inheritance -- that is the key ceiling model, unchanged.)""" + from litellm.proxy._types import LiteLLM_ObjectPermissionTable, SpecialMCPServerNames + + auth = _make_admitted_subject("sso-user") + auth.object_permission = LiteLLM_ObjectPermissionTable( + object_permission_id="op-user", mcp_servers=[SpecialMCPServerNames.no_mcp_servers.value] + ) + with self._patch(teams_by_id={"t1": _make_team("t1", ["srv1"])}, user_teams=["t1"]): + result = await MCPRequestHandler.get_allowed_mcp_servers(auth) + assert set(result) == {"srv1"}, "the user's own opt-out must not zero their team's grants" + + async def test_sources_fan_out_per_team_and_drop_non_roster_teams(self): + """The fan-out lives here now. One source per grant source: the user's own grants (no team_id, + carrying their object_permission) plus each team they are a LIVE roster member of. A team that + lingers in the user's cached `teams` array but no longer lists them in members_with_roles is + dropped, which is what revokes access after a team_member_delete the user row hasn't caught up + on. Each team source carries that team's own org, which is what makes the shared resolver apply + the team's owning-org ceiling rather than the caller's home org.""" + teams = { + "t-member": _make_team("t-member", ["srv1"], members=("sso-user",)), + "t-stale": _make_team("t-stale", ["srv2"], members=("someone-else",)), + } + teams["t-member"].organization_id = "org-a" + auth = _make_admitted_subject("sso-user") + with self._patch(teams_by_id=teams, user_teams=["t-member", "t-stale"]): + sources = await MCPRequestHandler._admitted_subject_sources(auth) + + assert [(s.team_id, s.org_id) for s in sources] == [(None, None), ("t-member", "org-a")] + # The user's own source carries their grants; a team source must NOT, or the team would be + # widened by grants the team never made. + assert sources[0].object_permission is auth.object_permission + assert sources[1].object_permission is None + # Every source is an ordinary caller, so it cannot re-enter the admitted fan-out. + assert all(not s.mcp_admitted_user_subject for s in sources) + # Nothing that meters or elevates the request may ride along onto a per-source clone. + assert all(s.api_key is None and s.user_role is None for s in sources) + + async def test_jwt_keyless_user_without_team_claim_does_not_union(self): + """Regression for the review finding: a JWT-authenticated caller is also keyless with a + user_id and (with no team claim) no team_id, but it is NOT admission-marked, so it must + keep its prior behavior of inheriting no team grants rather than silently gaining the + union across every team the user belongs to.""" + teams = {"team-a": _make_team("team-a", ["srv1"]), "team-b": _make_team("team-b", ["srv2"])} + jwt_auth = UserAPIKeyAuth(user_id="jwt-user", api_key=None) # no admission marker + with self._patch(teams_by_id=teams, user_teams=["team-a", "team-b"]): + result = await MCPRequestHandler._get_allowed_mcp_servers_for_team(jwt_auth) + assert result == [] + + async def test_forged_metadata_marker_on_a_real_key_grants_no_union(self): + """Security regression (forged admission marker): the admitted-subject marker is a + server-only ``UserAPIKeyAuth`` field, NOT a metadata key, precisely because virtual-key + metadata is caller-controlled at key creation. A user who sets + ``mcp_admitted_user_subject: true`` in their own key's metadata (api_key present, no + team_id) must NOT be treated as an admitted subject and must gain no cross-team union.""" + teams = {"team-a": _make_team("team-a", ["srv1"]), "team-b": _make_team("team-b", ["srv2"])} + forged = UserAPIKeyAuth( + user_id="attacker", + api_key="sk-real-key", + metadata={"mcp_admitted_user_subject": True}, # caller-forged marker in key metadata + ) + assert _is_mcp_admitted_user_subject(forged) is False + with self._patch(teams_by_id=teams, user_teams=["team-a", "team-b"]): + assert await MCPRequestHandler._team_ids_for_mcp_grant(forged) == [] + assert await MCPRequestHandler._get_allowed_mcp_servers_for_team(forged) == [] + + async def test_admitted_subject_team_tool_restriction_binds(self): + """Security regression (team tool restrictions bypassed): a keyless admitted subject whose + granting team restricts ``srv1`` to ``{tool_a}`` must NOT receive allow-all on srv1. The + single-team-id tool lookup returns None (allow-all) for a keyless multi-team user, dropping + the exclusion; the union across granting teams restores it.""" + from litellm.proxy._types import LiteLLM_ObjectPermissionTable, LiteLLM_TeamTable, Member + + team = LiteLLM_TeamTable( + team_id="team-a", + members_with_roles=[Member(user_id="sso-user", role="user")], + access_group_ids=[], + object_permission=LiteLLM_ObjectPermissionTable( + object_permission_id="op-team-a", + mcp_servers=["srv1"], + mcp_tool_permissions={"srv1": ["tool_a"]}, + ), + ) + auth = _make_admitted_subject("sso-user") + with self._patch(teams_by_id={"team-a": team}, user_teams=["team-a"]): + tools = await MCPRequestHandler.get_allowed_tools_for_server("srv1", auth) + assert tools == ["tool_a"] + + async def test_blocked_team_grants_no_servers_to_admitted_subject(self): + """Security regression: a blocked team grants nothing. The central policy gate enforces this + for a key pinned to a single team_id, but a keyless admitted subject unions across ALL its + teams (no team_id), so a blocked team's MCP grants must be dropped at the per-team resolver.""" + from litellm.proxy._types import LiteLLM_ObjectPermissionTable, LiteLLM_TeamTable, Member + + blocked = LiteLLM_TeamTable( + team_id="team-blocked", + blocked=True, + members_with_roles=[Member(user_id="sso-user", role="user")], + access_group_ids=[], + object_permission=LiteLLM_ObjectPermissionTable(object_permission_id="op-blk", mcp_servers=["srv-secret"]), + ) + teams = {"team-ok": _make_team("team-ok", ["srv-ok"]), "team-blocked": blocked} + auth = _make_admitted_subject("sso-user") + with self._patch(teams_by_id=teams, user_teams=["team-ok", "team-blocked"]): + result = await MCPRequestHandler.get_allowed_mcp_servers(auth) + assert set(result) == {"srv-ok"} + + async def test_admitted_subject_not_on_team_roster_gets_no_grant(self): + """Security regression (membership containment): a keyless subject whose user_id is NOT on a + team's roster inherits nothing from it, even when the team id lingers in the user's (stale or + cached) teams array. The team roster is the source of truth, so a removed or foreign + membership revokes access at the union rather than granting it.""" + teams = {"team-x": _make_team("team-x", ["srv-x"], members=("someone-else",))} + auth = _make_admitted_subject("sso-user") # in user.teams for team-x, but NOT on its roster + with self._patch(teams_by_id=teams, user_teams=["team-x"]): + result = await MCPRequestHandler._get_allowed_mcp_servers_for_team(auth) + assert result == [] + + async def test_tool_resolution_fails_closed_on_db_error(self): + """Security regression: ANY error resolving the tool allowlist for a keyless admitted subject + must DENY the server's tools ([]) rather than collapse to allow-all (None). Patches an await + OUTSIDE the multi-team fan-out (the team-object lookup) to prove the whole function fails + closed, not just the one helper — mirroring the fail-closed server path.""" + auth = _make_admitted_subject("sso-user") + with patch.object( + MCPRequestHandler, + "_get_team_object_permission", + new=AsyncMock(side_effect=RuntimeError("db blip")), + ): + tools = await MCPRequestHandler.get_allowed_tools_for_server("srv1", auth) + assert tools == [] + + async def test_admission_marker_cannot_be_set_from_validated_input(self): + """Defense-in-depth: the mcp_admitted_user_subject marker is server-only. Supplying it in any + validated input (constructor kwargs OR model_validate, e.g. a future JWT/key claim splat) is + stripped by the before-validator, so ONLY the admission path's post-construction assignment + can set it.""" + via_kwarg = UserAPIKeyAuth(user_id="u", api_key=None, mcp_admitted_user_subject=True) + via_validate = UserAPIKeyAuth.model_validate({"user_id": "u", "mcp_admitted_user_subject": True}) + assert via_kwarg.mcp_admitted_user_subject is False + assert via_validate.mcp_admitted_user_subject is False + assert _is_mcp_admitted_user_subject(via_kwarg) is False + assert _is_mcp_admitted_user_subject(via_validate) is False + + +@pytest.mark.asyncio +class TestAdmittedSubjectPerTeamOrgCap: + """A keyless admitted subject unions grants across teams that may span organizations. Each team's + grant (servers AND tools) is capped by that team's OWN org, and the user's direct grants by the + user's own org — never the caller's primary org applied over the whole cross-org union. Guards the + Veria 'team grants bypass their owning policies' finding.""" + + #: sentinel for org_perms: org has an object_permission_id but its load returns None (a swallowed + #: DB error / dangling id), which _object_permission_for_org must treat as fail-closed. + LOAD_FAILS = "__load_fails__" + + @contextlib.contextmanager + def _patch(self, *, teams_by_id, user_teams, org_perms=None, registry=None): + """org_perms: {org_id: LiteLLM_ObjectPermissionTable | None | LOAD_FAILS}. + - table → org exists, ceiling = that permission. + - None → org exists but carries no object_permission (no ceiling). + - LOAD_FAILS → org exists with an object_permission_id, but the permission load returns None. + - org_id ABSENT from the map → org row missing: get_org_object RAISES a bare Exception, exactly + as production does (it does NOT return None or raise HTTPException).""" + org_perms = org_perms or {} + + async def _get_team_object(team_id, **kw): + return teams_by_id.get(team_id) + + async def _get_user_object(user_id, **kw): + return MagicMock(user_id=user_id, teams=user_teams) + + async def _get_org_object(org_id, **kw): + if org_id not in org_perms: + from litellm.proxy.auth.auth_checks import OrganizationNotFoundError + + # matches production: a CONFIRMED-absent org raises this specific type, so callers + # can tell it apart from an outage (a bare Exception now means "lookup failed"). + raise OrganizationNotFoundError(f"Organization doesn't exist. Org={org_id}.") + op = org_perms[org_id] + has_permission_id = op is not None # a table OR LOAD_FAILS carries an id; None does not + return MagicMock( + organization_id=org_id, + object_permission_id=(f"orgop-{org_id}" if has_permission_id else None), + # Real typed values: the budget owners compare these, and a bare MagicMock attribute + # would explode the comparison and silently drop the source (bare-Mock rule). + litellm_budget_table=None, + spend=0.0, + ) + + async def _get_object_permission(object_permission_id, **kw): + for oid, op in org_perms.items(): + if op is not None and op != self.LOAD_FAILS and object_permission_id == f"orgop-{oid}": + return op + return None # LOAD_FAILS (or an unknown id) → None, simulating get_object_permission's swallow + + cms = [ + patch("litellm.proxy.auth.auth_checks.get_team_object", _get_team_object), + patch("litellm.proxy.auth.auth_checks.get_user_object", _get_user_object), + patch("litellm.proxy.auth.auth_checks.get_org_object", _get_org_object), + patch("litellm.proxy.auth.auth_checks.get_object_permission", _get_object_permission), + patch( + "litellm.proxy.auth.auth_checks._get_mcp_server_ids_from_access_groups", + AsyncMock(return_value=[]), + ), + patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), + patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()), + patch("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock()), + ] + if registry is not None: + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + + # registry may be a list of bare server_ids (MagicMock servers) OR a dict of + # {server_id: server_obj} for tests that need real alias/name resolution (config servers). + reg = registry if isinstance(registry, dict) else {s: MagicMock() for s in registry} + cms.append(patch.object(global_mcp_server_manager, "get_registry", return_value=reg)) + with contextlib.ExitStack() as es: + for cm in cms: + es.enter_context(cm) + yield + + # ---- server axis ---- + + async def test_team_grant_capped_by_its_own_org(self): + from litellm.proxy._types import LiteLLM_ObjectPermissionTable + + teams = {"team-a": _make_team("team-a", ["srv1", "srv2"], org_id="org-a")} + org_perms = {"org-a": LiteLLM_ObjectPermissionTable(object_permission_id="orgop-org-a", mcp_servers=["srv1"])} + auth = _make_admitted_subject("sso-user") + with self._patch(teams_by_id=teams, user_teams=["team-a"], org_perms=org_perms): + result = await MCPRequestHandler.get_allowed_mcp_servers(auth) + assert set(result) == {"srv1"} # srv2 capped out by org-a's ceiling + + async def test_cross_org_teams_each_capped_by_own_org(self): + from litellm.proxy._types import LiteLLM_ObjectPermissionTable + + teams = { + "team-a": _make_team("team-a", ["srv1", "srv2"], org_id="org-a"), + "team-b": _make_team("team-b", ["srv3", "srv4"], org_id="org-b"), + } + org_perms = { + "org-a": LiteLLM_ObjectPermissionTable(object_permission_id="orgop-org-a", mcp_servers=["srv1"]), + "org-b": LiteLLM_ObjectPermissionTable(object_permission_id="orgop-org-b", mcp_servers=["srv3"]), + } + auth = _make_admitted_subject("sso-user") + with self._patch(teams_by_id=teams, user_teams=["team-a", "team-b"], org_perms=org_perms): + result = await MCPRequestHandler.get_allowed_mcp_servers(auth) + assert set(result) == {"srv1", "srv3"} # each team clipped by its OWN org, then unioned + + async def test_all_proxy_grant_capped_by_org(self): + from litellm.proxy._types import LiteLLM_ObjectPermissionTable, SpecialMCPServerName + + teams = {"team-a": _make_team("team-a", [SpecialMCPServerName.all_proxy_servers.value], org_id="org-a")} + org_perms = {"org-a": LiteLLM_ObjectPermissionTable(object_permission_id="orgop-org-a", mcp_servers=["srv1"])} + auth = _make_admitted_subject("sso-user") + with self._patch( + teams_by_id=teams, user_teams=["team-a"], org_perms=org_perms, registry=["srv1", "srv2", "srv3"] + ): + result = await MCPRequestHandler.get_allowed_mcp_servers(auth) + # all_proxy expands to the whole registry, then org-a caps to {srv1} — the cell the old partial + # patch missed (it returned the full registry before capping). + assert set(result) == {"srv1"} + + async def test_org_row_without_object_permission_does_not_cap(self): + teams = {"team-a": _make_team("team-a", ["srv1", "srv2"], org_id="org-a")} + auth = _make_admitted_subject("sso-user") + with self._patch(teams_by_id=teams, user_teams=["team-a"], org_perms={"org-a": None}): + result = await MCPRequestHandler.get_allowed_mcp_servers(auth) + assert set(result) == {"srv1", "srv2"} # empty ceiling = no restriction + + async def test_direct_grants_unioned_with_team_and_capped_by_user_org(self): + from litellm.proxy._types import LiteLLM_ObjectPermissionTable + + teams = {"team-a": _make_team("team-a", ["srv1"], org_id="org-a")} + org_perms = { + "org-a": None, # the team's org imposes no ceiling + "org-u": LiteLLM_ObjectPermissionTable(object_permission_id="orgop-org-u", mcp_servers=["srvD", "srv1"]), + } + auth = _make_admitted_subject("sso-user", org_id="org-u", own_servers=["srvD", "srvX"]) + with self._patch(teams_by_id=teams, user_teams=["team-a"], org_perms=org_perms): + result = await MCPRequestHandler.get_allowed_mcp_servers(auth) + # direct {srvD,srvX} ∩ user-org {srvD,srv1} = {srvD}; UNIONed with team {srv1} (not intersected). + # srvX capped out by the user's org; team's srv1 NOT clipped by the user's primary org. + assert set(result) == {"srvD", "srv1"} + + async def test_single_team_key_uses_primary_org_cap_not_per_team(self): + from litellm.proxy._types import LiteLLM_ObjectPermissionTable + + # A KEY (not admitted): the per-team org cap must NOT fire; the top-level primary-org cap applies, + # byte-identical to before. team-a (org-a) grants {srv1,srv2}; the key's primary org is org-k. + teams = {"team-a": _make_team("team-a", ["srv1", "srv2"], org_id="org-a")} + org_perms = { + "org-a": LiteLLM_ObjectPermissionTable(object_permission_id="orgop-org-a", mcp_servers=["srv2"]), + "org-k": LiteLLM_ObjectPermissionTable(object_permission_id="orgop-org-k", mcp_servers=["srv1"]), + } + key_auth = UserAPIKeyAuth(user_id="u", api_key="sk-hash", team_id="team-a", org_id="org-k") + with self._patch(teams_by_id=teams, user_teams=["team-a"], org_perms=org_perms): + result = await MCPRequestHandler.get_allowed_mcp_servers(key_auth) + # If the per-team (org-a) cap wrongly fired, team-a would clip to {srv2} then org-k → {} (empty). + # Correct key behavior: no per-team cap; primary-org (org-k) cap → {srv1}. + assert set(result) == {"srv1"} + + # ---- tool axis ---- + + async def test_org_tool_ceiling_binds_when_team_places_no_tool_restriction(self): + from litellm.proxy._types import LiteLLM_ObjectPermissionTable + + # team grants srv1 with NO tool restriction; org-a restricts srv1's tools to {tool_a}. + teams = {"team-a": _make_team("team-a", ["srv1"], org_id="org-a")} + org_perms = { + "org-a": LiteLLM_ObjectPermissionTable( + object_permission_id="orgop-org-a", + mcp_servers=["srv1"], + mcp_tool_permissions={"srv1": ["tool_a"]}, + ) + } + auth = _make_admitted_subject("sso-user") + with self._patch(teams_by_id=teams, user_teams=["team-a"], org_perms=org_perms): + tools = await MCPRequestHandler.get_allowed_tools_for_server("srv1", auth) + # Without the per-team org tool ceiling this would be None (all tools) — org-a's tool ceiling + # would be bypassed exactly like the server case. + assert tools == ["tool_a"] + + async def test_tool_union_across_cross_org_teams(self): + teams = { + "team-a": _make_team("team-a", ["srv1"], org_id="org-a", tool_perms={"srv1": ["t1"]}), + "team-b": _make_team("team-b", ["srv1"], org_id="org-b", tool_perms={"srv1": ["t2"]}), + } + auth = _make_admitted_subject("sso-user") + with self._patch(teams_by_id=teams, user_teams=["team-a", "team-b"], org_perms={"org-a": None, "org-b": None}): + tools = await MCPRequestHandler.get_allowed_tools_for_server("srv1", auth) + assert set(tools) == {"t1", "t2"} + + async def test_tool_deny_all_when_team_grant_and_org_tool_ceiling_disjoint(self): + from litellm.proxy._types import LiteLLM_ObjectPermissionTable + + teams = {"team-a": _make_team("team-a", ["srv1"], org_id="org-a", tool_perms={"srv1": ["t1"]})} + org_perms = { + "org-a": LiteLLM_ObjectPermissionTable( + object_permission_id="orgop-org-a", + mcp_servers=["srv1"], + mcp_tool_permissions={"srv1": ["t2"]}, + ) + } + auth = _make_admitted_subject("sso-user") + with self._patch(teams_by_id=teams, user_teams=["team-a"], org_perms=org_perms): + tools = await MCPRequestHandler.get_allowed_tools_for_server("srv1", auth) + # team {t1} ∩ org {t2} = {} → deny every tool ([]), NOT allow-all (None). + assert tools == [] + + # ---- error contract (adversarial-review findings) ---- + + async def test_missing_org_row_is_treated_as_no_ceiling_not_lockout(self): + """A team's organization_id may point to an org row that no longer exists (deleted / not yet + synced). get_org_object RAISES a bare Exception for that; it must be treated as 'no ceiling' and + must NOT lock the admitted subject out of the team's grants (parity with the key path, which + tolerates a deleted org).""" + teams = {"team-a": _make_team("team-a", ["srv1", "srv2"], org_id="org-gone")} + auth = _make_admitted_subject("sso-user") + with self._patch(teams_by_id=teams, user_teams=["team-a"], org_perms={}): # org-gone absent → raises + result = await MCPRequestHandler.get_allowed_mcp_servers(auth) + assert set(result) == {"srv1", "srv2"} + + async def test_org_permission_load_failure_fails_closed(self): + """The org carries an object_permission_id but the permission load returns None (a swallowed DB + error / dangling id). The ceiling cannot be verified, so the admitted subject must fail CLOSED + for that team — NOT skip the ceiling, which would leak org-forbidden servers.""" + teams = {"team-a": _make_team("team-a", ["srv1", "srv2"], org_id="org-a")} + auth = _make_admitted_subject("sso-user") + with self._patch(teams_by_id=teams, user_teams=["team-a"], org_perms={"org-a": self.LOAD_FAILS}): + # Asserted through the PUBLIC resolver: the per-source org ceiling is applied there now, + # so calling the single-team helper would return [] for an admitted subject either way + # and pin nothing. + result = await MCPRequestHandler.get_allowed_mcp_servers(auth) + assert result == [] # fail closed, not {srv1, srv2} + + # ---- open bot-thread findings (2026-07-21 re-review) ---- + + async def test_org_less_team_grant_capped_by_user_primary_org(self): + """HIGH (cursor): a team with NO organization_id must still be bounded by the user's PRIMARY + org — otherwise, since admitted subjects skip the top-level primary-org cap, an org-less team's + grant would bypass every org ceiling and reach servers the user's home org forbids.""" + from litellm.proxy._types import LiteLLM_ObjectPermissionTable + + teams = {"team-noorg": _make_team("team-noorg", ["srv1", "srv2"], org_id=None)} + org_perms = {"org-U": LiteLLM_ObjectPermissionTable(object_permission_id="orgop-org-U", mcp_servers=["srv1"])} + auth = _make_admitted_subject("sso-user", org_id="org-U") + with self._patch(teams_by_id=teams, user_teams=["team-noorg"], org_perms=org_perms): + result = await MCPRequestHandler.get_allowed_mcp_servers(auth) + # org-less team falls back to the user's primary org (org-U → {srv1}); srv2 capped out. + assert set(result) == {"srv1"} + + async def test_tool_empty_contributions_fails_closed(self): + """MEDIUM (greptile/cursor): when no source in the tool-resolution view grants the server (a + TOCTOU/cache-lag inconsistency on a server that passed the server gate), the admitted path must + fail CLOSED (deny all tools = []), NOT allow-all (None).""" + teams = {"team-a": _make_team("team-a", ["srv1"], org_id="org-a")} + auth = _make_admitted_subject("sso-user") + with self._patch(teams_by_id=teams, user_teams=["team-a"], org_perms={"org-a": None}): + # 'srv-nobody' is granted by neither the team nor the user directly → empty contributions. + tools = await MCPRequestHandler.get_allowed_tools_for_server("srv-nobody", auth) + assert tools == [] + + async def test_tool_no_db_honors_in_memory_direct_restriction(self): + """MEDIUM (cursor): with no DB, the tool path must still honor the user's OWN in-memory + object_permission tool restriction (resolvable without a DB) rather than blanket-allow (None).""" + auth = _make_admitted_subject( + "sso-user", own_servers=["srv1"], own_tool_perms={"srv1": ["t1"]} + ) # no org_id, direct grant of srv1 restricted to {t1} + with patch("litellm.proxy.proxy_server.prisma_client", None): + tools = await MCPRequestHandler.get_allowed_tools_for_server("srv1", auth) + assert tools == ["t1"] # in-memory restriction honored, not widened to all tools + + # ---- config.yaml-defined servers (incl. OAuth) ---- + + async def test_config_defined_oauth_server_by_alias_reached_and_org_capped(self): + """A config.yaml-defined MCP OAuth server flows through the SAME resolution as a DB server: + the team grant (and the org ceiling) reference it by ALIAS, expand_permission_list resolves it + via the config+DB registry union to its server_id, and the per-team org cap applies identically. + (The config server's OAuth *client* persistence is #33768 — an orthogonal egress concern; this + pins the grant/reachability side of the 10x flow for config-defined servers.)""" + from litellm.proxy._types import LiteLLM_ObjectPermissionTable + + cfg_server = MagicMock() + cfg_server.server_id = "cfg-oauth-1" + cfg_server.alias = "linear_cfg" + cfg_server.server_name = "linear_cfg" + cfg_server.name = "linear_cfg" + + # team grants the config server BY ALIAS alongside a DB-style bare id; org-a's ceiling lists + # ONLY the config server (also by alias). + teams = {"team-a": _make_team("team-a", ["linear_cfg", "srv-db"], org_id="org-a")} + org_perms = { + "org-a": LiteLLM_ObjectPermissionTable(object_permission_id="orgop-org-a", mcp_servers=["linear_cfg"]) + } + auth = _make_admitted_subject("sso-user") + with self._patch( + teams_by_id=teams, + user_teams=["team-a"], + org_perms=org_perms, + registry={"cfg-oauth-1": cfg_server}, + ): + result = await MCPRequestHandler.get_allowed_mcp_servers(auth) + # 'linear_cfg' alias resolves to the config server_id and survives org-a's ceiling; 'srv-db' + # (not in org-a's allowlist) is capped out — same per-team org cap, config server included. + assert set(result) == {"cfg-oauth-1"} + + async def test_config_oauth_server_alias_resolution_feeds_the_org_cap(self): + """A config-defined OAuth server granted BY ALIAS whose OWN org forbids it is capped out — AND the + cap is proven to run on RESOLVED server_ids, not raw strings. A control config server, granted by + alias and allowed by the org via its RESOLVED id, must survive: that inclusion is impossible unless + expand_permission_list resolved the grant alias to the id the ceiling lists, so a broken alias path + yields {} and FAILS this test — whereas a bare `assert empty` would pass even if resolution never + ran (the weakness Cursor flagged).""" + from litellm.proxy._types import LiteLLM_ObjectPermissionTable + + forbidden = MagicMock() # granted by alias, but its org forbids it → must be capped out + forbidden.server_id = "cfg-oauth-1" + forbidden.alias = forbidden.server_name = forbidden.name = "linear_cfg" + control = MagicMock() # granted by alias, allowed by the org via its RESOLVED id → must survive + control.server_id = "control-id" + control.alias = control.server_name = control.name = "control_alias" + + teams = {"team-b": _make_team("team-b", ["linear_cfg", "control_alias"], org_id="org-b")} + # org-b's ceiling allows ONLY the control server, referenced by its RESOLVED server_id. + org_perms = { + "org-b": LiteLLM_ObjectPermissionTable(object_permission_id="orgop-org-b", mcp_servers=["control-id"]) + } + auth = _make_admitted_subject("sso-user") + with self._patch( + teams_by_id=teams, + user_teams=["team-b"], + org_perms=org_perms, + registry={"cfg-oauth-1": forbidden, "control-id": control}, + ): + result = await MCPRequestHandler.get_allowed_mcp_servers(auth) + # control survives ('control_alias' resolved to 'control-id', matching the id-based ceiling); the + # forbidden config server ('cfg-oauth-1') is capped out. A broken alias path → {} → fails here. + assert set(result) == {"control-id"} + + +@pytest.mark.asyncio +class TestSessionBearerEgressScrub: + """The gateway session bearer / bridge envelope is an admission credential, never an upstream token. + The leak-defense scrub is anchored to the credential SHAPE, so a session-shaped Authorization is + stripped from every egress context even when it reaches a non-aggregate scope that never set the + admission marker (design-review finding: a session bearer misdirected to a per-server true_passthrough + path would otherwise be forwarded upstream verbatim and replayed against the aggregate endpoint).""" + + async def test_session_bearer_misdirected_to_passthrough_is_scrubbed(self): + from litellm.types.mcp import MCPAuth + + scope = { + "type": "http", + "method": "POST", + "path": "/mcp/passthrough_server", + "headers": [(b"authorization", b"Bearer llm_session_synthetic-shaped-token")], + } + ttp_server = MagicMock() + ttp_server.auth_type = MCPAuth.true_passthrough + + with ( + patch( + "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.user_api_key_auth", + new_callable=AsyncMock, + ) as mock_auth, + patch("litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager") as mock_mgr, + ): + mock_mgr.get_mcp_server_by_name.return_value = ttp_server + (_auth, _mah, _srv, _sah, oauth2_headers, raw_headers) = await MCPRequestHandler.process_mcp_request(scope) + + mock_auth.assert_not_called() # true_passthrough → LiteLLM auth skipped (anonymous arm, no marker) + assert oauth2_headers is None # session-shaped bearer scrubbed from oauth2 egress + assert all(k.lower() != "authorization" for k in raw_headers) # ...and from raw egress headers + + async def test_legitimate_upstream_token_is_not_scrubbed(self): + """A genuine upstream/passthrough token is never session- or envelope-shaped, so the shape-anchored + scrub must leave it intact for forwarding (guards against over-stripping).""" + from litellm.types.mcp import MCPAuth + + scope = { + "type": "http", + "method": "POST", + "path": "/mcp/passthrough_server", + "headers": [(b"authorization", b"Bearer real-upstream-opaque-token-xyz")], + } + ttp_server = MagicMock() + ttp_server.auth_type = MCPAuth.true_passthrough + + with ( + patch( + "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.user_api_key_auth", + new_callable=AsyncMock, + ), + patch("litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager") as mock_mgr, + ): + mock_mgr.get_mcp_server_by_name.return_value = ttp_server + (_auth, _mah, _srv, _sah, oauth2_headers, _raw) = await MCPRequestHandler.process_mcp_request(scope) + + assert oauth2_headers.get("Authorization") == "Bearer real-upstream-opaque-token-xyz" + + async def test_scrub_removes_gateway_credential_from_every_egress_context(self): + """The scrub is anchored to the credential SHAPE and covers ALL egress contexts, not just + Authorization: a session bearer placed in x-mcp-auth OR a per-server x-mcp-{alias}-authorization + header is stripped too (the High-severity gap: those were forwarded upstream before).""" + sess = "Bearer llm_session_abc" + oauth2, raw, mcp_auth, per_server = MCPRequestHandler._scrub_gateway_admission_credentials( + admitted=False, + oauth2_headers={"Authorization": sess}, + raw_headers={ + "authorization": sess, + "x-mcp-auth": "llm_session_xyz", + "x-mcp-github-authorization": "llm_session_ghi", + }, + mcp_auth_header="llm_session_xyz", + mcp_server_auth_headers={"github": {"Authorization": "llm_session_ghi"}}, + ) + assert oauth2 is None + assert "authorization" not in {k.lower() for k in raw} + assert all("llm_session_" not in v for v in raw.values()) # x-mcp-auth + per-server raw values gone + assert mcp_auth is None # deprecated x-mcp-auth value scrubbed + assert per_server == {} # per-server session bearer removed → now-empty server dict dropped + + async def test_scrub_keeps_real_upstream_tokens(self): + """A legitimate upstream token is never session-/envelope-shaped, so every context is forwarded + unchanged — guards against over-stripping a real credential the caller meant for the upstream.""" + oauth2, raw, mcp_auth, per_server = MCPRequestHandler._scrub_gateway_admission_credentials( + admitted=False, + oauth2_headers={"Authorization": "Bearer real-upstream-xyz"}, + raw_headers={"authorization": "Bearer real-upstream-xyz", "x-mcp-github-authorization": "Bearer gh_real"}, + mcp_auth_header="some-api-key-123", + mcp_server_auth_headers={"github": {"Authorization": "Bearer gh_real"}}, + ) + assert oauth2 == {"Authorization": "Bearer real-upstream-xyz"} + assert raw["authorization"] == "Bearer real-upstream-xyz" + assert mcp_auth == "some-api-key-123" + assert per_server == {"github": {"Authorization": "Bearer gh_real"}} + + async def test_scrub_admitted_drops_authorization_but_keeps_injected_upstream_token(self): + """An admitted subject's top-level Authorization is dropped unconditionally, while the real + upstream token the bridge arm INJECTS into a per-server header (not gateway-shaped) survives.""" + oauth2, raw, mcp_auth, per_server = MCPRequestHandler._scrub_gateway_admission_credentials( + admitted=True, + oauth2_headers={"Authorization": "Bearer llm_session_abc"}, + raw_headers={"authorization": "Bearer llm_session_abc"}, + mcp_auth_header=None, + mcp_server_auth_headers={"github": {"Authorization": "Bearer gh_injected_upstream"}}, + ) + assert oauth2 is None + assert "authorization" not in {k.lower() for k in raw} + assert per_server == {"github": {"Authorization": "Bearer gh_injected_upstream"}} diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py index 636c7fbd3d5..a61e9de3281 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py @@ -3302,19 +3302,20 @@ async def test_token_root_does_not_resolve_private_server_for_external_client(): @pytest.mark.asyncio -async def test_register_root_resolves_single_oauth2_server(): - """When /register is hit without server name and exactly 1 OAuth2 server exists, resolve it.""" - try: - from fastapi import Request +async def test_register_root_does_aggregate_dcr_not_single_server_resolution(): + """Root /register is the aggregate DCR endpoint: it mints a stateless llm_dcrc_ client + from the request's redirect_uris and does NOT resolve a single configured oauth2 server + (a single-server deployment registers at /{server}/register instead).""" + import json - from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( - register_client, - ) - from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( - global_mcp_server_manager, - ) - except ImportError: - pytest.skip("MCP discoverable endpoints not available") + from fastapi import Request + + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + register_client, + ) + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) global_mcp_server_manager.registry.clear() oauth2_server = _create_oauth2_server() @@ -3325,33 +3326,37 @@ async def test_register_root_resolves_single_oauth2_server(): mock_request.headers = {} try: - with patch( - "litellm.proxy._experimental.mcp_server.discoverable_endpoints._read_request_body", - new=AsyncMock(return_value={}), + with ( + patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints._read_request_body", + new=AsyncMock(return_value={"redirect_uris": ["https://claude.ai/cb"]}), + ), + patch("litellm.proxy.proxy_server.master_key", "sk-test-salt-for-lit3637"), ): - result = await register_client(request=mock_request, mcp_server_name=None) + response = await register_client(request=mock_request, mcp_server_name=None) - # Should resolve to the single server and return its name as client_id - assert result["client_id"] == "test_oauth" - assert "redirect_uris" in result + body = json.loads(response.body) + assert body["client_id"].startswith("llm_dcrc_") + assert body["client_id"] != "test_oauth" + assert body["token_endpoint_auth_method"] == "none" finally: global_mcp_server_manager.registry.clear() @pytest.mark.asyncio -async def test_register_root_does_not_resolve_private_server_for_external_client(): - """Root /register must not reveal or use a hidden MCP server.""" - try: - from fastapi import Request +async def test_register_root_does_not_leak_a_private_server(): + """Root /register never resolves or reveals a configured server, so a private one cannot + leak to an external caller: it always mints the aggregate DCR client instead.""" + import json - from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( - register_client, - ) - from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( - global_mcp_server_manager, - ) - except ImportError: - pytest.skip("MCP discoverable endpoints not available") + from fastapi import Request + + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + register_client, + ) + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) global_mcp_server_manager.registry.clear() oauth2_server = _create_oauth2_server(available_on_public_internet=False) @@ -3365,17 +3370,19 @@ async def test_register_root_does_not_resolve_private_server_for_external_client with ( patch( "litellm.proxy._experimental.mcp_server.discoverable_endpoints._read_request_body", - new=AsyncMock(return_value={}), + new=AsyncMock(return_value={"redirect_uris": ["https://claude.ai/cb"]}), ), patch( "litellm.proxy._experimental.mcp_server.discoverable_endpoints.IPAddressUtils.get_mcp_client_ip", return_value="198.51.100.10", ), + patch("litellm.proxy.proxy_server.master_key", "sk-test-salt-for-lit3637"), ): - result = await register_client(request=mock_request, mcp_server_name=None) + response = await register_client(request=mock_request, mcp_server_name=None) - assert result["client_id"] == "dummy_client" - assert result["redirect_uris"] == ["https://llm.example.com/callback"] + body = json.loads(response.body) + assert body["client_id"].startswith("llm_dcrc_") + assert "test_oauth" not in body["client_id"] finally: global_mcp_server_manager.registry.clear() @@ -5155,7 +5162,10 @@ async def test_bridge_refresh_grant_with_non_envelope_is_invalid_grant_before_up def _mint_test_refresh_envelope( - server_id="bridge_srv", key_hash="hashed-litellm-key-77", upstream_refresh="UPSTREAM-REFRESH", identity=None, + server_id="bridge_srv", + key_hash="hashed-litellm-key-77", + upstream_refresh="UPSTREAM-REFRESH", + identity=None, scope=None, ): """Mint a refresh envelope the way the producer does, for driving the refresh_token grant in tests. @@ -5178,7 +5188,9 @@ def _mint_test_refresh_envelope( keys = envelope_keys_from_master_key(_BRIDGE_MASTER_KEY) identity = identity if identity is not None else key_hash_identity(server_id=server_id, key_hash=key_hash) sealed = build_bridge_refresh_token_response( - identity, RefreshCredential(refresh_token=SecretStr(upstream_refresh), scope=scope), keys, + identity, + RefreshCredential(refresh_token=SecretStr(upstream_refresh), scope=scope), + keys, datetime.now(timezone.utc), ) assert isinstance(sealed, SealedEnvelope) @@ -5413,7 +5425,10 @@ async def test_bridge_refresh_re_requests_the_sealed_scope_when_client_omits_it( ) captured: dict = {} response = await _refresh_for_bridge_server( - server, refresh_env, {"access_token": "NEW-ACCESS", "token_type": "Bearer", "expires_in": 3600}, None, + server, + refresh_env, + {"access_token": "NEW-ACCESS", "token_type": "Bearer", "expires_in": 3600}, + None, fake_client_out=captured, ) @@ -5553,7 +5568,9 @@ async def test_bridge_refresh_upstream_invalid_grant_maps_to_invalid_grant(): error_response = MagicMock() error_response.status_code = 400 error_response.text = '{"error": "invalid_grant", "error_description": "refresh token expired"}' - error_response.json = MagicMock(return_value={"error": "invalid_grant", "error_description": "refresh token expired"}) + error_response.json = MagicMock( + return_value={"error": "invalid_grant", "error_description": "refresh token expired"} + ) error_response.raise_for_status = MagicMock( side_effect=httpx.HTTPStatusError("bad", request=MagicMock(), response=error_response) ) @@ -7183,7 +7200,9 @@ def _upstream_token_response(status_code: int, *, json_body: object = None, text return httpx.Response(status_code, text=text_body, request=request) -async def _exchange_with_upstream_response(upstream_response, *, server_client_id="web-client.apps.googleusercontent.com"): +async def _exchange_with_upstream_response( + upstream_response, *, server_client_id="web-client.apps.googleusercontent.com" +): """Run the raw (non-bridge) authorization_code exchange against a canned upstream token-endpoint response and return what the gateway would hand the client. ``server_client_id=None`` models the caller-supplied-credentials flow (no stored client on the server).""" @@ -7334,9 +7353,7 @@ async def test_token_exchange_bounds_relayed_error_fields(): async def test_token_exchange_200_without_access_token_is_502_not_keyerror(): """A 200 whose body has no usable access_token used to KeyError into a 500; the raw arm now answers 502 with the same wording as the bridge arm's no_upstream_token rejection.""" - response = await _exchange_with_upstream_response( - _upstream_token_response(200, json_body={"token_type": "Bearer"}) - ) + response = await _exchange_with_upstream_response(_upstream_token_response(200, json_body={"token_type": "Bearer"})) assert response.status_code == 502 body = json.loads(response.body) @@ -7357,7 +7374,9 @@ async def test_token_exchange_relays_rejection_when_http_client_raises(): ) raising_client = MagicMock() raising_client.post = AsyncMock( - side_effect=httpx.HTTPStatusError("Client error '401 Unauthorized'", request=rejection.request, response=rejection) + side_effect=httpx.HTTPStatusError( + "Client error '401 Unauthorized'", request=rejection.request, response=rejection + ) ) from fastapi import Request @@ -7422,7 +7441,9 @@ async def test_register_relays_rejection_when_http_client_raises(): ) raising_client = MagicMock() raising_client.post = AsyncMock( - side_effect=httpx.HTTPStatusError("Client error '400 Bad Request'", request=rejection.request, response=rejection) + side_effect=httpx.HTTPStatusError( + "Client error '400 Bad Request'", request=rejection.request, response=rejection + ) ) oauth2_server = _bridge_server(auth_type=MCPAuth.oauth2, dcr_bridge=None) @@ -7800,9 +7821,7 @@ async def test_hydrate_does_not_overwrite_explicit_config_client_id(): auth_type=MCPAuth.oauth2, client_id="explicit-from-config", ) - store_read = AsyncMock( - return_value={"client_id": "stale-store-client", "client_secret": "x", "redirect_uris": []} - ) + store_read = AsyncMock(return_value={"client_id": "stale-store-client", "client_secret": "x", "redirect_uris": []}) with ( patch("litellm.proxy.utils.get_prisma_client_or_throw", return_value=MagicMock()), patch( @@ -8019,9 +8038,7 @@ async def test_bare_origin_discovery_resolves_single_server_not_aggregate(): mock_request.headers = {} try: - authorization_response = _build_oauth_authorization_server_response( - request=mock_request, mcp_server_name=None - ) + authorization_response = _build_oauth_authorization_server_response(request=mock_request, mcp_server_name=None) resource_response = await _build_oauth_protected_resource_response( request=mock_request, mcp_server_name=None, use_standard_pattern=True ) @@ -8033,6 +8050,65 @@ async def test_bare_origin_discovery_resolves_single_server_not_aggregate(): global_mcp_server_manager.registry.clear() +def test_gateway_dcr_flow_routing_engages_only_for_llm_dcrc_clients(monkeypatch): + """The aggregate DCR arms engage for llm_dcrc_ client_ids (register always mints one, + authorize/token route into the aggregate flow); a non-gateway client_id keeps the + per-server behavior, and /authorize/complete exists but 400s without a valid flow.""" + from fastapi import FastAPI + from fastapi.testclient import TestClient + + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import router + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + + monkeypatch.setenv("LITELLM_SALT_KEY", "sk-test-salt-for-lit3637") + monkeypatch.setattr("litellm.proxy.proxy_server.master_key", "sk-test-salt-for-lit3637", raising=False) + global_mcp_server_manager.registry.clear() + app = FastAPI() + app.include_router(router) + client = TestClient(app) + + registered = client.post("/register", json={"redirect_uris": ["https://claude.ai/cb"]}) + assert registered.status_code == 201 + assert registered.json()["client_id"].startswith("llm_dcrc_") + assert registered.json()["token_endpoint_auth_method"] == "none" + + authorize_params = { + "client_id": "llm_dcrc_bogus", + "redirect_uri": "https://claude.ai/cb", + "response_type": "code", + "code_challenge": "c" * 43, + "code_challenge_method": "S256", + } + bogus_client = client.get("/authorize", params=authorize_params) + assert bogus_client.status_code == 400 + assert bogus_client.json()["error"] == "invalid_client" + + no_cookie = client.post("/authorize/complete", data={"flow": "h"}) + assert no_cookie.status_code == 400 + assert no_cookie.json()["error"] == "invalid_request" + + token_response = client.post( + "/token", + data={ + "grant_type": "authorization_code", + "client_id": "llm_dcrc_bogus", + "code": "x", + "redirect_uri": "https://claude.ai/cb", + "code_verifier": "v" * 43, + }, + ) + assert token_response.status_code == 400 + assert token_response.json()["error"] == "invalid_grant" + + upstream_shaped = client.post( + "/token", + data={"grant_type": "authorization_code", "client_id": "regular-upstream-client", "code": "x"}, + ) + assert upstream_shaped.status_code == 404 + + @pytest.mark.asyncio async def test_authorize_wall_names_the_fix_for_urlless_servers(): """LIT-4629: the authorize wall previously said only "authorization url is not set" with no @@ -8148,3 +8224,601 @@ async def test_register_wall_names_the_fix_for_urlless_servers(): detail_text = str(exc_info.value.detail) assert "set Authorization URL and Token URL" in detail_text assert "Issuer" in detail_text + + +@pytest.mark.asyncio +async def test_authorize_wall_points_at_discovery_failure_for_url_servers(): + """LIT-4658: a server WITH a url that still has no authorization_url got here because OAuth + discovery against that url failed (typically a misconfigured url); the old detail blamed + "servers with no url", sending the operator down the wrong path. The detail must now name the + discovery failure and point at the proxy logs where LIT-4658's warnings carry the reason.""" + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + authorize_with_server, + ) + from litellm.types.mcp import MCPAuth, MCPTransport + from litellm.types.mcp_server.mcp_server_manager import MCPServer + + server = MCPServer( + server_id="typo-url-wall", + name="typo_wall", + server_name="typo_wall", + url="https://typo-host.example.com/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + ) + mock_request = MagicMock() + mock_request.base_url = "https://litellm.example.com/" + mock_request.headers = {} + + with pytest.raises(HTTPException) as exc_info: + await authorize_with_server( + request=mock_request, + mcp_server=server, + client_id="client", + redirect_uri="http://localhost/callback", + ) + assert exc_info.value.status_code == 400 + detail_text = str(exc_info.value.detail) + assert "may be misconfigured" in detail_text + assert "proxy logs" in detail_text + assert "Servers with no url" not in detail_text + assert "typo-host.example.com" not in detail_text + + +@pytest.mark.asyncio +async def test_token_wall_points_at_discovery_failure_for_url_servers(): + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + exchange_token_with_server, + ) + from litellm.types.mcp import MCPAuth, MCPTransport + from litellm.types.mcp_server.mcp_server_manager import MCPServer + + server = MCPServer( + server_id="typo-url-token-wall", + name="typo_token_wall", + server_name="typo_token_wall", + url="https://typo-host.example.com/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + authorization_url="https://idp.example.com/authorize", + ) + mock_request = MagicMock() + mock_request.base_url = "https://litellm.example.com/" + mock_request.headers = {} + + with pytest.raises(HTTPException) as exc_info: + await exchange_token_with_server( + request=mock_request, + mcp_server=server, + grant_type="authorization_code", + code="auth-code", + redirect_uri="http://localhost/callback", + client_id="client", + client_secret=None, + code_verifier="verifier", + ) + assert exc_info.value.status_code == 400 + detail_text = str(exc_info.value.detail) + assert "token url is not configured" in detail_text + assert "may be misconfigured" in detail_text + assert "Servers with no url" not in detail_text + + +@pytest.mark.asyncio +async def test_authorize_wall_names_the_issuer_for_anchored_servers(): + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + authorize_with_server, + ) + from litellm.types.mcp import MCPAuth, MCPTransport + from litellm.types.mcp_server.mcp_server_manager import MCPServer + + server = MCPServer( + server_id="anchored-wall", + name="anchored_wall", + server_name="anchored_wall", + url="https://up.example.com/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + issuer="https://idp.example.com", + issuer_is_anchored=True, + ) + mock_request = MagicMock() + mock_request.base_url = "https://litellm.example.com/" + mock_request.headers = {} + + with pytest.raises(HTTPException) as exc_info: + await authorize_with_server( + request=mock_request, + mcp_server=server, + client_id="client", + redirect_uri="http://localhost/callback", + ) + assert exc_info.value.status_code == 400 + detail_text = str(exc_info.value.detail) + assert "verify the Issuer" in detail_text + assert "Servers with no url" not in detail_text + assert "idp.example.com" not in detail_text +def test_passthrough_authorization_code_round_trips_and_rejects_hostile_input(): + """The passthrough gateway code seals and recovers the ephemeral DCR client and upstream code, + and is total over hostile input: a raw upstream code opens to None, and a tampered or + non-gateway value opens to None rather than raising, so every existing caller-supplied-client + flow is untouched.""" + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + open_passthrough_authorization_code, + seal_passthrough_authorization_code, + ) + + with patch("litellm.proxy.proxy_server.master_key", _BRIDGE_MASTER_KEY): + sealed = seal_passthrough_authorization_code( + upstream_code="up-code", + client_id="minted-77", + client_secret="mint-secret", + mcp_server_id="srv-1", + token_endpoint_auth_method="client_secret_basic", + ) + opened = open_passthrough_authorization_code(sealed) + assert opened is not None + assert opened.upstream_code == "up-code" + assert opened.client_id == "minted-77" + assert opened.client_secret == "mint-secret" + assert opened.mcp_server_id == "srv-1" + assert opened.token_endpoint_auth_method == "client_secret_basic" + assert open_passthrough_authorization_code("raw-upstream-code") is None + assert open_passthrough_authorization_code(sealed[:-4] + "aaaa") is None + + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + _BRIDGE_AUTH_CODE_PREFIX, + _PASSTHROUGH_AUTH_CODE_PREFIX, + open_bridge_authorization_code, + seal_bridge_authorization_code, + ) + + bridge_sealed = seal_bridge_authorization_code( + upstream_code="up-code", litellm_user_id="sso-user-9", mcp_server_id="srv-1" + ) + reprefixed_as_passthrough = _PASSTHROUGH_AUTH_CODE_PREFIX + bridge_sealed[len(_BRIDGE_AUTH_CODE_PREFIX) :] + reprefixed_as_bridge = _BRIDGE_AUTH_CODE_PREFIX + sealed[len(_PASSTHROUGH_AUTH_CODE_PREFIX) :] + assert open_passthrough_authorization_code(reprefixed_as_passthrough) is None + assert open_bridge_authorization_code(reprefixed_as_bridge) is None + + +@pytest.mark.asyncio +async def test_authorize_with_ephemeral_dcr_client_seals_client_into_state(): + """When mcp_authorize fell through to a gateway-side DCR mint, authorize_with_server seals the + minted client and the target server into the encrypted OAuth state, so the callback can bind + them into the forwarded authorization code while the gateway stores nothing.""" + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + EphemeralDcrClient, + authorize_with_server, + ) + from litellm.types.mcp import MCPAuth + + server = _bridge_server(auth_type=MCPAuth.true_passthrough, dcr_bridge=None) + captured: dict = {} + + def _capture(**kwargs): + captured.update(kwargs) + return "mocked_encrypted_state" + + with patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints.encode_state_with_base_url", + side_effect=_capture, + ): + response = await authorize_with_server( + request=_bridge_mock_request(), + mcp_server=server, + client_id="minted-77", + redirect_uri="http://127.0.0.1:60108/callback", + state="s", + code_challenge="chal", + code_challenge_method="S256", + ephemeral_dcr_client=EphemeralDcrClient( + client_id="minted-77", client_secret="mint-secret", token_endpoint_auth_method="client_secret_basic" + ), + ) + + assert captured["dcr_client_id"] == "minted-77" + assert captured["dcr_client_secret"] == "mint-secret" + assert captured["dcr_token_endpoint_auth_method"] == "client_secret_basic" + assert captured["mcp_server_id"] == server.server_id + assert "client_id=minted-77" in response.headers["location"] + + +@pytest.mark.asyncio +async def test_callback_wraps_code_into_passthrough_code_for_ephemeral_dcr_state(): + """When the OAuth state carries an ephemeral DCR client, the callback forwards a sealed + passthrough code (binding the client and the upstream code to the server) instead of the raw + upstream code, so the client's later token call can authenticate the exchange with a client the + gateway never stored.""" + from urllib.parse import parse_qs, urlparse + + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + callback, + open_passthrough_authorization_code, + ) + + state_data = { + "original_state": "client-state", + "client_redirect_uri": "http://127.0.0.1:60108/cb", + "base_url": "http://127.0.0.1:60108/cb", + "mcp_server_id": "srv-1", + "dcr_client_id": "minted-77", + "dcr_client_secret": "mint-secret", + "dcr_token_endpoint_auth_method": "client_secret_basic", + } + with ( + patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints._resolve_encoded_oauth_state", + return_value="enc", + ), + patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints.decode_state_hash", + return_value=state_data, + ), + patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints._get_validated_client_redirect_uri", + return_value="http://127.0.0.1:60108/cb", + ), + patch("litellm.proxy.proxy_server.master_key", _BRIDGE_MASTER_KEY), + ): + response = await callback(request=_bridge_mock_request(), code="REAL-UPSTREAM-CODE", state="relay") + + forwarded_code = parse_qs(urlparse(response.headers["location"]).query)["code"][0] + opened = open_passthrough_authorization_code(forwarded_code) + + assert opened is not None + assert opened.upstream_code == "REAL-UPSTREAM-CODE" + assert opened.client_id == "minted-77" + assert opened.client_secret == "mint-secret" + assert opened.mcp_server_id == "srv-1" + assert opened.token_endpoint_auth_method == "client_secret_basic" + + +@pytest.mark.asyncio +async def test_authorize_bridge_server_with_ephemeral_client_takes_short_circuit_arm(): + """A gateway-minted client is registered against {base}/callback, so a bridge server's + authorize with an ephemeral client must run the short-circuit (gateway /callback) arm with a + relay state cookie, never the verbatim relay: relaying would send the browser's redirect_uri + to an IdP that has the gateway callback registered, stranding the flow.""" + from urllib.parse import parse_qs, urlparse + + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + EphemeralDcrClient, + authorize_with_server, + ) + from litellm.types.mcp import MCPAuth + + server = _bridge_server(auth_type=MCPAuth.true_passthrough) + with patch("litellm.proxy.proxy_server.master_key", _BRIDGE_MASTER_KEY): + response = await authorize_with_server( + request=_bridge_mock_request(), + mcp_server=server, + client_id="minted-77", + redirect_uri="http://127.0.0.1:60108/callback", + state="client-state", + code_challenge="chal", + code_challenge_method="S256", + ephemeral_dcr_client=EphemeralDcrClient(client_id="minted-77", client_secret=None), + ) + + location = response.headers["location"] + params = parse_qs(urlparse(location).query) + assert params["redirect_uri"] == ["https://litellm.example.com/callback"] + assert params["client_id"] == ["minted-77"] + assert params["state"] != ["client-state"] + assert any(cookie.startswith("mcp_oauth_state_") for cookie in response.headers.get("set-cookie", "").split(";")) + + +@pytest.mark.asyncio +async def test_callback_forwards_raw_code_when_dcr_state_lacks_server_binding(): + """A state carrying a dcr client but no server id cannot produce a server-bound sealed code, so + the callback falls back to forwarding the raw upstream code instead of sealing an unbindable + one.""" + from urllib.parse import parse_qs, urlparse + + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import callback + + state_data = { + "original_state": "client-state", + "client_redirect_uri": "http://127.0.0.1:60108/cb", + "base_url": "http://127.0.0.1:60108/cb", + "dcr_client_id": "minted-77", + } + with ( + patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints._resolve_encoded_oauth_state", + return_value="enc", + ), + patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints.decode_state_hash", + return_value=state_data, + ), + patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints._get_validated_client_redirect_uri", + return_value="http://127.0.0.1:60108/cb", + ), + patch("litellm.proxy.proxy_server.master_key", _BRIDGE_MASTER_KEY), + ): + response = await callback(request=_bridge_mock_request(), code="REAL-UPSTREAM-CODE", state="relay") + + forwarded_code = parse_qs(urlparse(response.headers["location"]).query)["code"][0] + assert forwarded_code == "REAL-UPSTREAM-CODE" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "auth_type_value", + [ + "none", + "api_key", + "bearer_token", + "basic", + "authorization", + "oauth2", + "aws_sigv4", + "token", + "oauth2_token_exchange", + "oauth2_id_jag", + "true_passthrough", + "oauth_delegate", + ], +) +@pytest.mark.parametrize("dcr_bridge", [True, False]) +async def test_resolve_ephemeral_dcr_client_mint_set_is_exact(auth_type_value, dcr_bridge): + """The full authorize-time mint decision matrix, one cell per (auth_type, dcr_bridge). The gateway + mints iff true_passthrough (any bridge) or oauth_delegate-and-not-dcr_bridge; every other mode + returns None so no non-OAuth mode ever registers an upstream client, and the interactive + oauth_delegate dcr_bridge sign-in is left to its own browser-front-door flow. The UI + gatewayMintsClientFor helper mirrors this exact set; ui/.../mcp_tools/types.test.tsx pins the + frontend side against the same table, so a divergence fails on one side or the other.""" + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + EphemeralDcrClient, + resolve_ephemeral_dcr_client, + ) + from litellm.types.mcp import MCPAuth + + server = _bridge_server( + auth_type=MCPAuth(auth_type_value), + dcr_bridge=dcr_bridge, + server_id=f"matrix_{auth_type_value}_{dcr_bridge}", + server_name=f"matrix_{auth_type_value}_{dcr_bridge}", + ) + expected_mint = server.is_true_passthrough or (server.is_oauth_delegate and not server.is_dcr_bridge) + mint_mock = AsyncMock(return_value=EphemeralDcrClient(client_id="minted", client_secret=None)) + + with patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints.mint_ephemeral_dcr_client", + mint_mock, + ): + result = await resolve_ephemeral_dcr_client( + request=_bridge_mock_request(), + mcp_server=server, + code_challenge="chal", + code_challenge_method="S256", + redirect_uri="http://127.0.0.1:9/callback", + ) + + if expected_mint: + mint_mock.assert_awaited_once() + assert result is not None + else: + mint_mock.assert_not_awaited() + assert result is None + + +@pytest.mark.asyncio +async def test_mint_ephemeral_dcr_client_returns_none_without_registration_endpoint(): + """A server whose upstream exposes no RFC 7591 registration endpoint cannot mint, so the + fall-through reports None and the caller keeps its existing missing_client_id failure.""" + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + mint_ephemeral_dcr_client, + ) + from litellm.types.mcp import MCPAuth + + server = _bridge_server(auth_type=MCPAuth.true_passthrough, dcr_bridge=None, registration_url=None) + assert await mint_ephemeral_dcr_client(_bridge_mock_request(), server) is None + + +@pytest.mark.asyncio +async def test_mint_ephemeral_dcr_client_posts_rfc7591_and_returns_client(): + """The mint POSTs a public-client RFC 7591 registration bound to the gateway /callback and hands + back the upstream's client without persisting it anywhere.""" + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + mint_ephemeral_dcr_client, + ) + from litellm.types.mcp import MCPAuth + + server = _bridge_server( + auth_type=MCPAuth.true_passthrough, dcr_bridge=None, server_id="mint_posts_srv", server_name="mint_posts_srv" + ) + mock_response = MagicMock() + mock_response.text = json.dumps( + {"client_id": "minted-77", "client_secret": "mint-secret", "token_endpoint_auth_method": "client_secret_basic"} + ) + mock_response.raise_for_status = MagicMock() + mock_async_client = MagicMock() + mock_async_client.post = AsyncMock(return_value=mock_response) + + with patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints.get_async_httpx_client", + return_value=mock_async_client, + ): + minted = await mint_ephemeral_dcr_client(_bridge_mock_request(), server) + + assert minted is not None + assert minted.client_id == "minted-77" + assert minted.client_secret == "mint-secret" + assert minted.token_endpoint_auth_method == "client_secret_basic" + register_data = mock_async_client.post.call_args.kwargs["json"] + assert register_data["redirect_uris"] == ["https://litellm.example.com/callback"] + assert register_data["token_endpoint_auth_method"] == "none" + assert register_data["grant_types"] == ["authorization_code", "refresh_token"] + + +@pytest.mark.asyncio +async def test_mint_ephemeral_dcr_client_reuses_minted_client_within_flow_ttl(): + """Reloading the authorize page must not spam the upstream registration endpoint with orphan + clients: within the OAuth state's lifetime a second mint for the same server and gateway origin + reuses the cached client and performs no second upstream POST.""" + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + mint_ephemeral_dcr_client, + ) + from litellm.types.mcp import MCPAuth + + server = _bridge_server( + auth_type=MCPAuth.true_passthrough, dcr_bridge=None, server_id="mint_reuse_srv", server_name="mint_reuse_srv" + ) + mock_response = MagicMock() + mock_response.text = json.dumps({"client_id": "minted-77"}) + mock_response.raise_for_status = MagicMock() + mock_async_client = MagicMock() + mock_async_client.post = AsyncMock(return_value=mock_response) + + with patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints.get_async_httpx_client", + return_value=mock_async_client, + ): + first = await mint_ephemeral_dcr_client(_bridge_mock_request(), server) + second = await mint_ephemeral_dcr_client(_bridge_mock_request(), server) + + assert first is not None + assert second == first + mock_async_client.post.assert_called_once() + + +@pytest.mark.asyncio +async def test_mint_ephemeral_dcr_client_single_flights_concurrent_mints(): + """Two in-flight authorize requests for the same server must not both register an upstream + client: the per-key lock makes the second waiter reuse the first mint, so exactly one upstream + POST happens.""" + import asyncio + + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + mint_ephemeral_dcr_client, + ) + from litellm.types.mcp import MCPAuth + + server = _bridge_server( + auth_type=MCPAuth.true_passthrough, + dcr_bridge=None, + server_id="mint_concurrent_srv", + server_name="mint_concurrent_srv", + ) + mock_response = MagicMock() + mock_response.text = json.dumps({"client_id": "minted-77"}) + mock_response.raise_for_status = MagicMock() + + async def _slow_post(*args, **kwargs): + await asyncio.sleep(0.05) + return mock_response + + mock_async_client = MagicMock() + mock_async_client.post = AsyncMock(side_effect=_slow_post) + + with patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints.get_async_httpx_client", + return_value=mock_async_client, + ): + first, second = await asyncio.gather( + mint_ephemeral_dcr_client(_bridge_mock_request(), server), + mint_ephemeral_dcr_client(_bridge_mock_request(), server), + ) + + assert first is not None + assert second == first + mock_async_client.post.assert_called_once() + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "payload, server_id", + [ + ({"unexpected": "shape"}, "mint_bad_shape_srv"), + ({"client_id": ""}, "mint_empty_id_srv"), + ], +) +async def test_mint_ephemeral_dcr_client_unusable_registration_response_is_502(payload, server_id): + """An upstream registration response without a usable client_id, whether the field is missing or + an empty string, surfaces as a loud 502 instead of letting the authorize proceed with an empty + client and fail opaquely at the IdP.""" + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + mint_ephemeral_dcr_client, + ) + from litellm.types.mcp import MCPAuth + + server = _bridge_server(auth_type=MCPAuth.true_passthrough, dcr_bridge=None, server_id=server_id, server_name=server_id) + mock_response = MagicMock() + mock_response.text = json.dumps(payload) + mock_response.raise_for_status = MagicMock() + mock_async_client = MagicMock() + mock_async_client.post = AsyncMock(return_value=mock_response) + + with patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints.get_async_httpx_client", + return_value=mock_async_client, + ): + with pytest.raises(HTTPException) as exc: + await mint_ephemeral_dcr_client(_bridge_mock_request(), server) + + assert exc.value.status_code == 502 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "sealed_auth_method, expects_basic_header", + [ + ("client_secret_basic", True), + (None, False), + ], +) +async def test_token_exchange_authenticates_with_the_sealed_clients_own_auth_method( + sealed_auth_method, expects_basic_header +): + """The id, secret, and token-endpoint auth method must come from the same source: a client + recovered from a sealed passthrough code authenticates the upstream exchange the way its own + registration was granted, not the way the server row is configured. A sealed + ``client_secret_basic`` grant sends the Basic header and keeps the secret out of the body; a + sealed public client (no method) keeps the body-credential path.""" + import base64 + + import httpx + + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + exchange_token_with_server, + ) + from litellm.types.mcp import MCPAuth + + server = _bridge_server(auth_type=MCPAuth.true_passthrough, dcr_bridge=None, server_id="sealed_method_srv") + upstream_request = httpx.Request("POST", server.token_url) + upstream_response = httpx.Response( + 200, json={"access_token": "up-token", "token_type": "Bearer"}, request=upstream_request + ) + mock_async_client = MagicMock() + mock_async_client.post = AsyncMock(return_value=upstream_response) + + with patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints.get_async_httpx_client", + return_value=mock_async_client, + ): + await exchange_token_with_server( + request=_bridge_mock_request(), + mcp_server=server, + grant_type="authorization_code", + code="up-code", + redirect_uri="https://litellm.example.com/callback", + client_id="minted-77", + client_secret="mint-secret", + code_verifier="verifier", + client_token_endpoint_auth_method=sealed_auth_method, + ) + + sent_headers = mock_async_client.post.call_args.kwargs["headers"] + sent_body = mock_async_client.post.call_args.kwargs["data"] + if expects_basic_header: + expected = base64.b64encode(b"minted-77:mint-secret").decode() + assert sent_headers["Authorization"] == f"Basic {expected}" + assert "client_secret" not in sent_body + else: + assert "Authorization" not in sent_headers + assert sent_body["client_id"] == "minted-77" + assert sent_body["client_secret"] == "mint-secret" diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py new file mode 100644 index 00000000000..375ec022115 --- /dev/null +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py @@ -0,0 +1,590 @@ +"""Tests for the aggregate gateway DCR flow (register, authorize, complete, token).""" + +import hashlib +import json +from base64 import urlsafe_b64encode +from datetime import datetime, timedelta, timezone +from http.cookies import SimpleCookie +from urllib.parse import parse_qs, urlparse + +import pytest +from starlette.requests import Request + +from litellm.caching.caching import DualCache +from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import ( + CONNECT_FLOW_COOKIE_PREFIX, + GATEWAY_AUTH_CODE_PREFIX, + GATEWAY_AUTH_CODE_TTL_SECONDS, + GATEWAY_DCR_CLIENT_ID_PREFIX, + _GatewayAuthCode, + _seal, + aggregate_authorize, + aggregate_token, + complete_connect_flow, + is_gateway_dcr_client_id, + open_gateway_dcr_client, + register_aggregate_client, +) +from litellm.proxy._experimental.mcp_server.outbound_credentials.session_credentials import ( + resolve_session_bearer, + session_keys_from_master_key, + SessionBearerAdmitted, +) + +MASTER_KEY = "sk-gateway-dcr-flow-tests" +REDIRECT_URI = "https://claude.ai/api/mcp/auth_callback" +CODE_VERIFIER = "verifier-" + "v" * 43 +CODE_CHALLENGE = urlsafe_b64encode(hashlib.sha256(CODE_VERIFIER.encode("ascii")).digest()).rstrip(b"=").decode("ascii") + + +@pytest.fixture(autouse=True) +def _salt_key(monkeypatch): + monkeypatch.setenv("LITELLM_SALT_KEY", MASTER_KEY) + + +def _request(path="/authorize", query="", cookies=None, method="GET"): + cookie_header = [] + if cookies: + cookie = SimpleCookie() + for name, value in cookies.items(): + cookie[name] = value + cookie_header = [(b"cookie", cookie.output(header="", sep="; ").strip().encode())] + return Request( + { + "type": "http", + "method": method, + "scheme": "https", + "path": path, + "query_string": query.encode(), + "headers": [(b"host", b"llm.example.com"), *cookie_header], + } + ) + + +async def _register(redirect_uris) -> dict: + response = await register_aggregate_client( + request=_request(path="/register", method="POST"), request_body={"redirect_uris": redirect_uris} + ) + return json.loads(response.body) + + +async def _reload_user_active(user_id: str): + return None + + +@pytest.mark.asyncio +async def test_register_mints_stateless_public_client(): + body = await _register([REDIRECT_URI]) + assert body["token_endpoint_auth_method"] == "none" + assert "client_secret" not in body + assert body["redirect_uris"] == [REDIRECT_URI] + assert is_gateway_dcr_client_id(body["client_id"]) + record = open_gateway_dcr_client(body["client_id"]) + assert record is not None + assert record.redirect_uris == (REDIRECT_URI,) + + +@pytest.mark.asyncio +async def test_register_allows_loopback_http_for_dev_clients(): + body = await _register(["http://localhost:6274/oauth/callback"]) + assert is_gateway_dcr_client_id(body["client_id"]) + + +@pytest.mark.parametrize( + "code_challenge", + ["short", "", "p" * 300, "ünïcode-challenge", "AAAA" * 20], +) +def test_pkce_mismatched_challenge_returns_false_never_raises(code_challenge): + """A wrong-length or non-ASCII code_challenge must VERIFY FALSE, not raise. + + Pins the reason this compares bytes rather than str: hmac.compare_digest raises TypeError on + two str with non-ASCII content, but on bytes of unequal length it simply returns False. A + review flagged this as an unhandled 500 on length mismatch; encoding both sides to bytes is + exactly what makes that impossible, so the claim is pinned here rather than in a comment.""" + from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import _pkce_verifier_matches + + assert _pkce_verifier_matches("a" * 43, code_challenge) is False + + +@pytest.mark.asyncio +async def test_register_allows_allowlisted_native_callback(): + """Native MCP clients register a private-use scheme, not https. Registration shares + the one redirect-URI shape owner with /authorize, so the callback the allowlist + already trusts there is registrable here rather than rejected as non-https.""" + body = await _register(["cursor://anysphere.cursor-mcp/oauth/callback"]) + assert is_gateway_dcr_client_id(body["client_id"]) + record = open_gateway_dcr_client(body["client_id"]) + assert record is not None + assert record.redirect_uris == ("cursor://anysphere.cursor-mcp/oauth/callback",) + + +@pytest.mark.asyncio +async def test_register_rejects_userinfo_spoofed_origin(): + """``https://claude.ai@attacker.example/cb`` parses with netloc + ``claude.ai@attacker.example``, so a naive origin display on the consent screen reads + as claude.ai while the code would be delivered to attacker.example. Rejected at + registration, which is the only way such a URI could enter a sealed client.""" + response = await register_aggregate_client( + request=_request(path="/register", method="POST"), + request_body={"redirect_uris": ["https://claude.ai@attacker.example/callback"]}, + ) + assert response.status_code == 400 + assert json.loads(response.body)["error"] == "invalid_redirect_uri" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "redirect_uris", + [ + [], + "not-a-list", + ["http://evil.example.com/callback"], + ["https://claude.ai/cb#fragment"], + ["ftp://claude.ai/cb"], + ["https://a.example.com/" + "p" * 300], + ["https://a.example.com/1", "https://a.example.com/2", "https://a.example.com/3", "https://a.example.com/4"], + [12345], + ], +) +async def test_register_rejects_bad_redirect_uris(redirect_uris): + response = await register_aggregate_client( + request=_request(path="/register", method="POST"), request_body={"redirect_uris": redirect_uris} + ) + assert response.status_code == 400 + assert json.loads(response.body)["error"] in ("invalid_redirect_uri", "invalid_client_metadata") + + +@pytest.mark.asyncio +async def test_tampered_client_id_does_not_open(): + body = await _register([REDIRECT_URI]) + tampered = body["client_id"][:-4] + "AAAA" + assert open_gateway_dcr_client(tampered) is None + assert open_gateway_dcr_client("llm_dcrc_garbage") is None + assert open_gateway_dcr_client("other_prefix") is None + + +def _authorize( + client_id, session_user_id, redirect_uri=REDIRECT_URI, challenge=CODE_CHALLENGE, method="S256", response_type="code" +): + return aggregate_authorize( + request=_request(query=f"client_id={client_id}"), + client_id=client_id, + redirect_uri=redirect_uri, + state="client-state-123", + code_challenge=challenge, + code_challenge_method=method, + response_type=response_type, + session_user_id=session_user_id, + ) + + +@pytest.mark.asyncio +async def test_authorize_validation_failures_never_redirect_to_client(): + client_id = (await _register([REDIRECT_URI]))["client_id"] + for response, expected_error in ( + (_authorize("llm_dcrc_bogus", "u1"), "invalid_client"), + (_authorize(client_id, "u1", redirect_uri="https://attacker.example.com/cb"), "invalid_request"), + (_authorize(client_id, "u1", response_type="token"), "unsupported_response_type"), + (_authorize(client_id, "u1", challenge=None), "invalid_request"), + (_authorize(client_id, "u1", method="plain"), "invalid_request"), + ): + assert response.status_code == 400 + assert json.loads(response.body)["error"] == expected_error + + +@pytest.mark.asyncio +async def test_authorize_without_session_redirects_to_login_with_return_to(): + client_id = (await _register([REDIRECT_URI]))["client_id"] + response = _authorize(client_id, session_user_id=None) + assert response.status_code == 303 + location = response.headers["location"] + assert location.startswith("https://llm.example.com/sso/key/generate?return_to=") + assert "return_to=%2Fauthorize" in location + + +@pytest.mark.asyncio +async def test_authorize_with_session_hands_browser_to_connect_page_with_flow_cookie(): + client_id = (await _register([REDIRECT_URI]))["client_id"] + response = _authorize(client_id, session_user_id="u1") + assert response.status_code == 303 + location = urlparse(response.headers["location"]) + assert location.path == "/ui/chat/integrations" + params = parse_qs(location.query) + handle = params["connect_flow"][0] + assert params["connect_client"] == ["https://claude.ai"] + set_cookie = response.headers["set-cookie"] + assert f"{CONNECT_FLOW_COOKIE_PREFIX}{handle}" in set_cookie + assert "HttpOnly" in set_cookie + return handle, set_cookie + + +def _flow_cookie_from(response) -> tuple: + location = urlparse(response.headers["location"]) + handle = parse_qs(location.query)["connect_flow"][0] + cookie = SimpleCookie() + cookie.load(response.headers["set-cookie"]) + name = f"{CONNECT_FLOW_COOKIE_PREFIX}{handle}" + return handle, {name: cookie[name].value} + + +@pytest.mark.asyncio +async def test_full_walk_register_authorize_complete_token_and_replay(): + """The whole front door on one deterministic walk: register -> authorize -> + complete -> token, then the security edges on the same artifacts (user mismatch, + PKCE mismatch, single-use replay, refresh rotation, cross-client refresh).""" + client_id = (await _register([REDIRECT_URI]))["client_id"] + authorize_response = _authorize(client_id, session_user_id="u1") + handle, cookies = _flow_cookie_from(authorize_response) + + denied = await complete_connect_flow( + request=_request("/authorize/complete", cookies=cookies, method="POST"), + flow_handle=handle, + session_user_id="attacker", + cache=DualCache(), + ) + assert denied.status_code == 403 + + anonymous = await complete_connect_flow( + request=_request("/authorize/complete", cookies=cookies, method="POST"), + flow_handle=handle, + session_user_id=None, + cache=DualCache(), + ) + assert anonymous.status_code == 401 + + completed = await complete_connect_flow( + request=_request("/authorize/complete", cookies=cookies, method="POST"), + flow_handle=handle, + session_user_id="u1", + cache=DualCache(), + ) + assert completed.status_code == 303 + redirect = urlparse(completed.headers["location"]) + assert f"{redirect.scheme}://{redirect.netloc}{redirect.path}" == REDIRECT_URI + params = parse_qs(redirect.query) + assert params["state"] == ["client-state-123"] + code = params["code"][0] + assert code.startswith(GATEWAY_AUTH_CODE_PREFIX) + + cache = DualCache() + + async def _token(**overrides): + arguments = { + "request": _request("/token", method="POST"), + "grant_type": "authorization_code", + "code": code, + "redirect_uri": REDIRECT_URI, + "client_id": client_id, + "code_verifier": CODE_VERIFIER, + "refresh_token": None, + "master_key": MASTER_KEY, + "reload_user": _reload_user_active, + "cache": cache, + } + return await aggregate_token(**{**arguments, **overrides}) + + wrong_verifier = await _token(code_verifier="wrong-" + "w" * 43) + assert json.loads(wrong_verifier.body)["error"] == "invalid_grant" + + wrong_client = await _token(client_id=(await _register([REDIRECT_URI]))["client_id"]) + assert json.loads(wrong_client.body)["error"] == "invalid_grant" + + token_response = await _token() + assert token_response.status_code == 200 + payload = json.loads(token_response.body) + assert payload["token_type"] == "Bearer" + assert 0 < payload["expires_in"] <= 3600 + + keys = session_keys_from_master_key(MASTER_KEY) + admitted = resolve_session_bearer(f"Bearer {payload['access_token']}", keys, datetime.now(timezone.utc)) + assert isinstance(admitted, SessionBearerAdmitted) + assert admitted.principal.user_id == "u1" + assert admitted.principal.client_id == client_id + + replay = await _token() + assert json.loads(replay.body)["error"] == "invalid_grant" + + refreshed = await _token(grant_type="refresh_token", code=None, refresh_token=payload["refresh_token"]) + assert refreshed.status_code == 200 + rotated = json.loads(refreshed.body) + assert rotated["refresh_token"] != payload["refresh_token"] + + # Rotation is single-use: replaying the now-consumed refresh token cannot mint a second pair + # (a captured token is dead once the legitimate holder has rotated). + replayed = await _token(grant_type="refresh_token", code=None, refresh_token=payload["refresh_token"]) + assert json.loads(replayed.body)["error"] == "invalid_grant" + assert "already used" in json.loads(replayed.body).get("error_description", "") + + cross_client = await _token( + grant_type="refresh_token", + code=None, + refresh_token=payload["refresh_token"], + client_id=(await _register([REDIRECT_URI]))["client_id"], + ) + assert json.loads(cross_client.body)["error"] == "invalid_grant" + + +@pytest.mark.asyncio +async def test_complete_rejects_missing_tampered_and_expired_flows(): + missing = await complete_connect_flow( + request=_request("/authorize/complete", method="POST"), + flow_handle="nope", + session_user_id="u1", + cache=DualCache(), + ) + assert missing.status_code == 400 + + tampered = await complete_connect_flow( + request=_request("/authorize/complete", cookies={f"{CONNECT_FLOW_COOKIE_PREFIX}h1": "garbage"}, method="POST"), + flow_handle="h1", + session_user_id="u1", + cache=DualCache(), + ) + assert tampered.status_code == 400 + + +@pytest.mark.asyncio +async def test_token_rejects_expired_code_and_missing_configuration(): + expired_code = _seal( + GATEWAY_AUTH_CODE_PREFIX, + _GatewayAuthCode( + user_id="u1", + client_id="llm_dcrc_x", + redirect_uri=REDIRECT_URI, + code_challenge=CODE_CHALLENGE, + jti="jti-1", + iat=int((datetime.now(timezone.utc) - timedelta(seconds=500)).timestamp()), + exp=int((datetime.now(timezone.utc) - timedelta(seconds=500 - GATEWAY_AUTH_CODE_TTL_SECONDS)).timestamp()), + ), + ) + response = await aggregate_token( + request=_request("/token", method="POST"), + grant_type="authorization_code", + code=expired_code, + redirect_uri=REDIRECT_URI, + client_id="llm_dcrc_x", + code_verifier=CODE_VERIFIER, + refresh_token=None, + master_key=MASTER_KEY, + reload_user=_reload_user_active, + cache=DualCache(), + ) + assert json.loads(response.body)["error"] == "invalid_grant" + + no_master_key = await aggregate_token( + request=_request("/token", method="POST"), + grant_type="authorization_code", + code="llm_gcode_x", + redirect_uri=REDIRECT_URI, + client_id="llm_dcrc_x", + code_verifier=CODE_VERIFIER, + refresh_token=None, + master_key=None, + reload_user=_reload_user_active, + cache=DualCache(), + ) + assert no_master_key.status_code == 500 + assert json.loads(no_master_key.body)["error"] == "server_error" + + unsupported = await aggregate_token( + request=_request("/token", method="POST"), + grant_type="password", + code=None, + redirect_uri=None, + client_id="llm_dcrc_x", + code_verifier=None, + refresh_token=None, + master_key=MASTER_KEY, + reload_user=_reload_user_active, + cache=DualCache(), + ) + assert json.loads(unsupported.body)["error"] == "unsupported_grant_type" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "failure,expected_status,expected_error", + [ + ("no_active_key", 400, "invalid_grant"), + ("unavailable", 503, "temporarily_unavailable"), + ("unresolvable", 500, "server_error"), + ], +) +async def test_token_gates_on_live_user_revalidation(failure, expected_status, expected_error): + client_id = (await _register([REDIRECT_URI]))["client_id"] + authorize_response = _authorize(client_id, session_user_id="deactivated-user") + handle, cookies = _flow_cookie_from(authorize_response) + completed = await complete_connect_flow( + request=_request("/authorize/complete", cookies=cookies, method="POST"), + flow_handle=handle, + session_user_id="deactivated-user", + cache=DualCache(), + ) + code = parse_qs(urlparse(completed.headers["location"]).query)["code"][0] + + async def _reload_user_failing(user_id: str): + return failure + + response = await aggregate_token( + request=_request("/token", method="POST"), + grant_type="authorization_code", + code=code, + redirect_uri=REDIRECT_URI, + client_id=client_id, + code_verifier=CODE_VERIFIER, + refresh_token=None, + master_key=MASTER_KEY, + reload_user=_reload_user_failing, + cache=DualCache(), + ) + assert response.status_code == expected_status + assert json.loads(response.body)["error"] == expected_error + + +@pytest.mark.asyncio +async def test_flow_is_single_use_shared_cache_rejects_second_complete(): + """A double-submit of the finish step mints only ONE code: the second complete over the + same cache fails invalid_request (atomic flow claim), so one sign-in cannot yield two codes.""" + cache = DualCache() + client_id = (await _register([REDIRECT_URI]))["client_id"] + handle, cookies = _flow_cookie_from(_authorize(client_id, session_user_id="u1")) + + first = await complete_connect_flow( + request=_request("/authorize/complete", cookies=cookies, method="POST"), + flow_handle=handle, + session_user_id="u1", + cache=cache, + ) + assert first.status_code == 303 + second = await complete_connect_flow( + request=_request("/authorize/complete", cookies=cookies, method="POST"), + flow_handle=handle, + session_user_id="u1", + cache=cache, + ) + assert second.status_code == 400 + assert json.loads(second.body)["error"] == "invalid_request" + + +@pytest.mark.asyncio +async def test_token_rejects_out_of_range_code_verifier(): + """RFC 7636: a code_verifier outside 43-128 chars is invalid_request, not a confusing + invalid_grant PKCE-mismatch.""" + for bad in ["short", "x" * 200]: + response = await aggregate_token( + request=_request("/token", method="POST"), + grant_type="authorization_code", + code="llm_gcode_whatever", + redirect_uri=REDIRECT_URI, + client_id="llm_dcrc_x", + code_verifier=bad, + refresh_token=None, + master_key=MASTER_KEY, + reload_user=_reload_user_active, + cache=DualCache(), + ) + assert response.status_code == 400 + assert json.loads(response.body)["error"] == "invalid_request" + + +@pytest.mark.asyncio +async def test_authorize_rejects_over_long_state(): + client_id = (await _register([REDIRECT_URI]))["client_id"] + response = aggregate_authorize( + request=_request(query=f"client_id={client_id}"), + client_id=client_id, + redirect_uri=REDIRECT_URI, + state="s" * 2000, + code_challenge=CODE_CHALLENGE, + code_challenge_method="S256", + response_type="code", + session_user_id="u1", + ) + assert response.status_code == 400 + assert json.loads(response.body)["error"] == "invalid_request" + + +@pytest.mark.asyncio +async def test_non_ascii_code_challenge_fails_grant_not_500(): + """A non-ASCII code_challenge (unvalidated from the client) must yield a clean + invalid_grant, never a TypeError-driven 500 (bytes comparison, not str).""" + client_id = (await _register([REDIRECT_URI]))["client_id"] + # Seal a code carrying a non-ASCII challenge directly (authorize requires S256 shape, + # but the challenge charset is not validated there, so this state is reachable). + from datetime import datetime, timezone + + code = _seal( + GATEWAY_AUTH_CODE_PREFIX, + _GatewayAuthCode( + user_id="u1", + client_id=client_id, + redirect_uri=REDIRECT_URI, + code_challenge="challenge-with-€-non-ascii", + jti="jti-x", + iat=int(datetime.now(timezone.utc).timestamp()), + exp=int(datetime.now(timezone.utc).timestamp()) + 120, + ), + ) + response = await aggregate_token( + request=_request("/token", method="POST"), + grant_type="authorization_code", + code=code, + redirect_uri=REDIRECT_URI, + client_id=client_id, + code_verifier=CODE_VERIFIER, + refresh_token=None, + master_key=MASTER_KEY, + reload_user=_reload_user_active, + cache=DualCache(), + ) + assert response.status_code == 400 + assert json.loads(response.body)["error"] == "invalid_grant" + + +@pytest.mark.asyncio +async def test_single_use_guard_in_memory_is_single_use_within_process(): + """No Redis configured (single-replica): the in-memory increment is authoritative — the first claim + wins, a replay of the same id loses.""" + from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import _SingleUseGuard + + guard = _SingleUseGuard(DualCache()) # redis_cache is None + assert await guard.claim("jti-inmem", 60) is True + assert await guard.claim("jti-inmem", 60) is False # replay of the same id + + +@pytest.mark.asyncio +async def test_single_use_guard_uses_redis_as_sole_authority_when_configured(): + """With Redis configured it is the SOLE authority: the shared INCR result decides the claim (1 → + first caller, >1 → replay), and the per-worker in-memory count is never consulted.""" + from unittest.mock import AsyncMock, MagicMock + + from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import _SingleUseGuard + + cache = DualCache() + cache.redis_cache = MagicMock() + cache.redis_cache.async_increment = AsyncMock(return_value=1) + # in-memory must NOT be consulted when Redis is configured — poison it so any fallback is visible. + cache.async_increment_cache = AsyncMock(side_effect=AssertionError("must not fall back to in-memory")) + + guard = _SingleUseGuard(cache) + assert await guard.claim("jti-redis", 60) is True + cache.redis_cache.async_increment = AsyncMock(return_value=2) + assert await guard.claim("jti-redis", 60) is False # Redis says 2 → replay + + +@pytest.mark.asyncio +async def test_single_use_guard_fails_closed_when_redis_errors(): + """A Redis fault must fail the claim CLOSED (refuse the id) rather than fall back to the per-worker + in-memory count — which would let each replica observe count==1 and replay the one-time id (the + Cursor/Veria replay-across-workers finding).""" + from unittest.mock import AsyncMock, MagicMock + + from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import _SingleUseGuard + + cache = DualCache() + cache.redis_cache = MagicMock() + cache.redis_cache.async_increment = AsyncMock(side_effect=ConnectionError("redis down")) + cache.async_increment_cache = AsyncMock(return_value=1) # would fail OPEN if the guard fell back + + guard = _SingleUseGuard(cache) + assert await guard.claim("jti-fault", 60) is False # fail closed, not a fallback count of 1 diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py index ae4f12fc1e1..dff1f1d87c7 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py @@ -7375,6 +7375,10 @@ async def test_call_tool_with_legacy_db_m2m_server_resolves_oauth2_flow(): ("", None), ("not a url", None), ("http://[::1", None), + # urlsplit validates the port lazily on attribute access, so a malformed port must not + # raise out of the helper: the server loaders call it while warning about exactly this + # kind of typo'd url (LIT-4658) + ("https://example.com:bad/mcp", None), ], ) def test_redact_mcp_resource_url_strips_credentials(url, expected): diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py index a5cb16822cf..42b6cbee1c4 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -2333,6 +2333,59 @@ class TestMCPServerManager: assert emitted.headers["Authorization"] == "Bearer upstream-token" assert not kwargs["extra_headers"] or "authorization" not in {k.lower() for k in kwargs["extra_headers"]} + @pytest.mark.asyncio + async def test_create_mcp_client_token_exchange_never_falls_back_to_v1(self): + """A configured OBO server is owned end to end by the v2 token_exchange arm, even when the + caller supplies an x-mcp-* override. This is what makes the v1 OBO handler unreachable, so if + it ever defers to v1 again the deleted handler is silently needed back.""" + from litellm.proxy._experimental.mcp_server.outbound_credentials.oauth_token_store import ( + OAuthToken, + ) + from litellm.proxy._experimental.mcp_server.outbound_credentials.resolver import ( + UpstreamCredentialProvider, + ) + from litellm.proxy._experimental.mcp_server.outbound_credentials.result import Ok + + class _StubExchanger: + def __init__(self): + self.subject_tokens = [] + + async def exchange(self, subject_token, server, config, *, tenant_id=""): + self.subject_tokens.append(subject_token) + return Ok(OAuthToken(access_token="exchanged-token")) + + async def invalidate(self, subject_token, server, config, *, tenant_id=""): + return None + + exchanger = _StubExchanger() + manager = MCPServerManager() + server = MCPServer( + server_id="obo-egress", + name="obo", + url="https://example.com", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2_token_exchange, + client_id="gateway-client", + client_secret="gateway-secret", + token_exchange_endpoint="https://idp.example.com/oauth2/token", + ) + with ( + patch( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.resolve_mcp_auth", + new_callable=AsyncMock, + ) as mock_resolve, + patch("litellm.proxy._experimental.mcp_server.mcp_server_manager.MCPClient") as mock_client_cls, + ): + await manager._create_mcp_client( + server=server, + mcp_auth_header="Bearer caller-override", + subject_token="eyJ-subject-token", + cred_provider=UpstreamCredentialProvider(token_exchanger=exchanger), + ) + mock_resolve.assert_not_awaited() + assert exchanger.subject_tokens == ["eyJ-subject-token"] + assert self._emitted_authorization(mock_client_cls) == "Bearer exchanged-token" + @staticmethod def _emitted_authorization(mock_client_cls) -> str: kwargs = mock_client_cls.call_args.kwargs @@ -3031,7 +3084,7 @@ class TestMCPServerManager: registration_url="https://discovered.example.com/register", ) - async def fake_discovery(server_url: str, *, allow_origin_fallback: bool = True): + async def fake_discovery(server_url: str, *, allow_origin_fallback: bool = True, warn_when_no_metadata: bool = False): assert server_url == "https://example.com/mcp" # oauth2 (browser flow) keeps the origin fallback; only OBO disables it. assert allow_origin_fallback is True @@ -5426,7 +5479,7 @@ class TestMCPServerTimestamps: manager = MCPServerManager() calls: list[bool] = [] - async def fake_discovery(server_url: str, *, allow_origin_fallback: bool = True): + async def fake_discovery(server_url: str, *, allow_origin_fallback: bool = True, warn_when_no_metadata: bool = False): calls.append(allow_origin_fallback) return MCPOAuthMetadata( scopes=None, @@ -5461,7 +5514,7 @@ class TestMCPServerTimestamps: manager = MCPServerManager() calls: list[str] = [] - async def fake_discovery(server_url: str, *, allow_origin_fallback: bool = True): + async def fake_discovery(server_url: str, *, allow_origin_fallback: bool = True, warn_when_no_metadata: bool = False): calls.append(server_url) raise AssertionError("discovery must not run when token_exchange_endpoint is configured") @@ -5491,7 +5544,7 @@ class TestMCPServerTimestamps: back to the row, so the next rebuild skips discovery instead of re-running it every time.""" manager = MCPServerManager() - async def fake_discovery(server_url: str, *, allow_origin_fallback: bool = True): + async def fake_discovery(server_url: str, *, allow_origin_fallback: bool = True, warn_when_no_metadata: bool = False): assert server_url == "https://example.com/mcp" assert allow_origin_fallback is False # OBO never guesses the origin return MCPOAuthMetadata( @@ -5602,7 +5655,7 @@ class TestMCPServerTimestamps: _dcr_bridge_relays_client_registration keys off that column.""" manager = MCPServerManager() - async def fake_discovery(server_url: str, *, allow_origin_fallback: bool = True): + async def fake_discovery(server_url: str, *, allow_origin_fallback: bool = True, warn_when_no_metadata: bool = False): assert allow_origin_fallback is True return MCPOAuthMetadata( scopes=["mcp.read", "mcp.write"], @@ -5817,7 +5870,7 @@ class TestMCPServerTimestamps: persist_discovered_endpoints=False neither the oauth2 nor the OBO write-back may fire.""" manager = MCPServerManager() - async def fake_discovery(server_url: str, *, allow_origin_fallback: bool = True): + async def fake_discovery(server_url: str, *, allow_origin_fallback: bool = True, warn_when_no_metadata: bool = False): return MCPOAuthMetadata( scopes=["s1"], authorization_url="https://idp.example.com/authorize", @@ -8539,7 +8592,7 @@ class TestOBOEndpointDiscovery: ) seen = [] - async def fake_discovery(server_url: str, *, allow_origin_fallback: bool = True): + async def fake_discovery(server_url: str, *, allow_origin_fallback: bool = True, warn_when_no_metadata: bool = False): seen.append((server_url, allow_origin_fallback)) return discovered @@ -8567,7 +8620,7 @@ class TestOBOEndpointDiscovery: async def test_config_obo_with_configured_endpoint_skips_discovery(self): manager = MCPServerManager() - async def fake_discovery(server_url: str, *, allow_origin_fallback: bool = True): + async def fake_discovery(server_url: str, *, allow_origin_fallback: bool = True, warn_when_no_metadata: bool = False): raise AssertionError("discovery must not run when the endpoint is configured") manager._descovery_metadata = fake_discovery # type: ignore[attr-defined] @@ -9028,3 +9081,162 @@ class TestUrllessIssuerDiscovery: anchored.assert_awaited_once_with("https://idp.example.com", None) resource_rooted.assert_not_awaited() assert built.token_url == "https://idp.example.com/token" + + +class TestDiscoveryFailureLogging: + """LIT-4658: a misconfigured MCP server url must be diagnosable from default-level server logs. + + Discovery failures used to die at debug level and the config-load path emitted no warning at + all, so the only operator-facing signal was the bare 400 at /authorize.""" + + def _connect_error_client(self, url: str) -> MagicMock: + client = MagicMock() + client.get = AsyncMock( + side_effect=httpx.ConnectError(f"[Errno 8] nodename nor servname provided for {url}") + ) + return client + + @pytest.mark.asyncio + async def test_descovery_metadata_warns_with_redacted_attempts_on_connect_error(self, caplog): + manager = MCPServerManager() + secret_url = "https://typo-host.example.com/mcp/s/PATHSECRET/mcp" + with ( + patch( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.get_async_httpx_client", + return_value=self._connect_error_client(secret_url), + ), + caplog.at_level(logging.WARNING, logger="LiteLLM"), + ): + result = await manager._descovery_metadata(secret_url, warn_when_no_metadata=True) + assert result is None + assert "found no authorization server metadata" in caplog.text + assert "ConnectError" in caplog.text + assert "https://typo-host.example.com" in caplog.text + # hosted MCP urls embed credentials in the path; neither the url nor the exception + # text may leak it into warning-level logs + assert "PATHSECRET" not in caplog.text + + @pytest.mark.asyncio + async def test_descovery_metadata_stays_silent_without_warn_flag(self, caplog): + manager = MCPServerManager() + url = "https://typo-host.example.com/mcp" + with ( + patch( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.get_async_httpx_client", + return_value=self._connect_error_client(url), + ), + caplog.at_level(logging.WARNING, logger="LiteLLM"), + ): + result = await manager._descovery_metadata(url) + assert result is None + assert "found no authorization server metadata" not in caplog.text + + @pytest.mark.asyncio + async def test_descovery_metadata_attempt_trail_names_each_failed_step(self, caplog): + manager = MCPServerManager() + url = "https://real-host.example.com/mcp-typo" + client = MagicMock() + client.get = AsyncMock( + return_value=httpx.Response(404, request=httpx.Request("GET", url)) + ) + with ( + patch( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.get_async_httpx_client", + return_value=client, + ), + caplog.at_level(logging.WARNING, logger="LiteLLM"), + ): + result = await manager._descovery_metadata(url, warn_when_no_metadata=True) + assert result is None + assert "HTTP 404" in caplog.text + assert "well-known protected-resource lookup found no authorization servers" in caplog.text + assert "origin fallback" in caplog.text + + @pytest.mark.asyncio + async def test_load_servers_from_config_warns_when_endpoints_unresolved(self, caplog): + manager = MCPServerManager() + manager._descovery_metadata = AsyncMock(return_value=None) # type: ignore[attr-defined] + config = { + "typo_server": { + "url": "https://typo.example.com/mcp", + "transport": MCPTransport.http, + "auth_type": MCPAuth.oauth2, + "oauth2_flow": "authorization_code", + } + } + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + await manager.load_servers_from_config(config) + assert "typo_server" in caplog.text + assert "authorization_url, token_url" in caplog.text + assert "unresolved" in caplog.text + assert "verify the configured server url" in caplog.text + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "extra_config", + [ + { + "authorization_url": "https://idp.example.com/auth", + "token_url": "https://idp.example.com/token", + }, + { + "oauth2_flow": "client_credentials", + "token_url": "https://idp.example.com/token", + "client_id": "cid", + "client_secret": "csec", + }, + ], + ) + async def test_load_servers_from_config_silent_when_flow_needs_covered(self, caplog, extra_config): + """Manually covered endpoints and M2M servers (which never need authorization_url) must not + warn on every reload; the warning is a misconfiguration signal, not discovery telemetry.""" + manager = MCPServerManager() + manager._descovery_metadata = AsyncMock(return_value=None) # type: ignore[attr-defined] + config = { + "covered_server": { + "url": "https://up.example.com/mcp", + "transport": MCPTransport.http, + "auth_type": MCPAuth.oauth2, + "oauth2_flow": "authorization_code", + **extra_config, + } + } + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + await manager.load_servers_from_config(config) + assert "unresolved" not in caplog.text + assert "no discovery source" not in caplog.text + + @pytest.mark.asyncio + async def test_config_server_without_discovery_source_warns_about_missing_endpoints(self, caplog): + manager = MCPServerManager() + manager._register_openapi_tools = AsyncMock() # type: ignore[attr-defined] + config = { + "spec_only": { + "spec_path": "https://example.com/openapi.yaml", + "transport": MCPTransport.http, + "auth_type": MCPAuth.oauth2, + "oauth2_flow": "authorization_code", + } + } + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + await manager.load_servers_from_config(config) + assert "no discovery source" in caplog.text + assert "authorization_url and token_url are not set manually" in caplog.text + + @pytest.mark.asyncio + async def test_db_build_warns_when_discovery_fails_for_oauth2_row(self, caplog): + manager = MCPServerManager() + manager._descovery_metadata = AsyncMock(return_value=None) # type: ignore[attr-defined] + record = LiteLLM_MCPServerTable( + server_id="typo-row-1", + server_name="typo_row", + url="https://typo.example.com/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + oauth2_flow="authorization_code", + ) + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + await manager.build_mcp_server_from_table(record, credentials_are_encrypted=False) + assert "typo_row" in caplog.text + assert "authorization_url, token_url" in caplog.text + assert "unresolved" in caplog.text diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_tool_search.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_tool_search.py index 1da44029b5c..0e442102e53 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_tool_search.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_tool_search.py @@ -237,7 +237,7 @@ class TestListToolRestApiWithToolSearch: return_value={}, ), patch( - "litellm.proxy._experimental.mcp_server.rest_endpoints._get_oauth2_server_ids", + "litellm.proxy._experimental.mcp_server.rest_endpoints._v1_resolved_oauth2_server_ids", return_value=[], ), patch( @@ -316,7 +316,7 @@ class TestListToolRestApiWithToolSearch: return_value={}, ), patch( - "litellm.proxy._experimental.mcp_server.rest_endpoints._get_oauth2_server_ids", + "litellm.proxy._experimental.mcp_server.rest_endpoints._v1_resolved_oauth2_server_ids", return_value=[], ), patch( diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py index d4ba66c4381..5c9612a055e 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py @@ -2783,3 +2783,68 @@ class TestRestListToolsetFiltering: ) assert [tool.name for tool in result] == ["lookup_status"] + + +class TestV1ResolvedOauth2Gate: + """The REST surface must stop resolving per-user OAuth2 tokens for servers the v2 resolver owns. + + ``_resolve_v2_auth`` drops any Authorization built here for an ``authorization_code`` server and + injects the resolver's own token, so the v1 lookup was a DB round-trip whose result was discarded. + A server that still defers to v1 (upstream-delegated oauth2) must keep resolving, which is what + makes these assertions non-vacuous. + """ + + @staticmethod + def _oauth2_server(*, delegate_auth_to_upstream: bool) -> Any: + from litellm.proxy._experimental.mcp_server.server import MCPServer + from litellm.types.mcp import MCPTransport + + return MCPServer( + server_id="oauth2-srv", + name="oauth2-srv", + url="https://upstream.example.com/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + delegate_auth_to_upstream=delegate_auth_to_upstream, + ) + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "delegate_auth_to_upstream, expected_headers, expected_lookups", + [ + (False, None, 0), + (True, {"Authorization": "Bearer stored-token"}, 1), + ], + ) + async def test_user_oauth_headers_skip_v2_owned_servers( + self, delegate_auth_to_upstream, expected_headers, expected_lookups, monkeypatch + ): + from litellm.proxy._experimental.mcp_server import db as mcp_db + + server = self._oauth2_server(delegate_auth_to_upstream=delegate_auth_to_upstream) + resolve_token = AsyncMock(return_value={"access_token": "stored-token"}) + monkeypatch.setattr(mcp_db, "resolve_valid_user_oauth_token", resolve_token) + + headers = await rest_endpoints._get_user_oauth_extra_headers( + server, + UserAPIKeyAuth(user_id="alice", api_key="sk-1234"), + prefetched_creds={"oauth2-srv": {"access_token": "stored-token"}}, + ) + + assert headers == expected_headers + assert resolve_token.await_count == expected_lookups + + def test_prefetch_preflight_only_counts_v1_resolved_servers(self, monkeypatch): + v2_owned = self._oauth2_server(delegate_auth_to_upstream=False) + v1_resolved = self._oauth2_server(delegate_auth_to_upstream=True) + v1_resolved.server_id = "delegate-srv" + registry = {"oauth2-srv": v2_owned, "delegate-srv": v1_resolved} + + monkeypatch.setattr( + rest_endpoints.global_mcp_server_manager, + "get_mcp_server_by_id", + lambda server_id: registry.get(server_id), + ) + + assert rest_endpoints._v1_resolved_oauth2_server_ids(["oauth2-srv"]) == set() + assert rest_endpoints._v1_resolved_oauth2_server_ids(["oauth2-srv", "delegate-srv"]) == {"delegate-srv"} diff --git a/tests/test_litellm/proxy/auth/test_auth_utils.py b/tests/test_litellm/proxy/auth/test_auth_utils.py index 72bd215b9be..9f24c662581 100644 --- a/tests/test_litellm/proxy/auth/test_auth_utils.py +++ b/tests/test_litellm/proxy/auth/test_auth_utils.py @@ -1587,7 +1587,7 @@ class TestGetDynamicLitellmParamsClearsAdminConfigOnBaseOverride: } out = get_dynamic_litellm_params( litellm_params=dict(admin_params), - request_kwargs={"base_url": "https://attacker.example"}, + request_kwargs={"base_url": "https://attacker.example", "api_key": "sk-caller"}, ) assert "aws_access_key_id" not in out assert "aws_secret_access_key" not in out @@ -1608,7 +1608,7 @@ class TestGetDynamicLitellmParamsClearsAdminConfigOnBaseOverride: } out = get_dynamic_litellm_params( litellm_params=dict(admin_params), - request_kwargs={"api_base": "self-hosted.example.com:50051"}, + request_kwargs={"api_base": "self-hosted.example.com:50051", "api_key": "sk-caller"}, ) assert out["api_base"] == "self-hosted.example.com:50051" assert "nvcf_function_id" not in out @@ -1626,7 +1626,7 @@ class TestGetDynamicLitellmParamsClearsAdminConfigOnBaseOverride: } out = get_dynamic_litellm_params( litellm_params=dict(admin_params), - request_kwargs={"api_base": "self-hosted.example.com:50051"}, + request_kwargs={"api_base": "self-hosted.example.com:50051", "api_key": "sk-caller"}, ) assert out["api_base"] == "self-hosted.example.com:50051" assert "use_ssl" not in out @@ -1651,6 +1651,7 @@ class TestGetDynamicLitellmParamsClearsAdminConfigOnBaseOverride: }, request_kwargs={ "api_base": "https://attacker.example", + "api_key": "sk-caller", "organization": "org-attacker", "extra_body": {"attacker": "value"}, }, @@ -1674,6 +1675,7 @@ class TestGetDynamicLitellmParamsClearsAdminConfigOnBaseOverride: }, request_kwargs={ "api_base": "https://attacker.example", + "api_key": "sk-caller", "organization": "", "extra_body": "", }, @@ -1701,6 +1703,310 @@ class TestGetDynamicLitellmParamsClearsAdminConfigOnBaseOverride: assert out["api_version"] == "2026-04-01" assert out["api_base"] == "https://admin.upstream/v1" + def test_client_api_key_used_when_supplied_with_base_override(self): + from litellm.router_utils.clientside_credential_handler import ( + get_dynamic_litellm_params, + ) + + out = get_dynamic_litellm_params( + litellm_params={ + "model": "gpt-4", + "api_key": "sk-admin-secret", + "api_base": "https://admin.upstream/v1", + }, + request_kwargs={ + "api_base": "https://attacker.example", + "api_key": "sk-client-byok", + }, + ) + assert out["api_key"] == "sk-client-byok" + assert "sk-admin-secret" not in str(out) + + +_OPENAI_CHAT_RESPONSE = { + "id": "chatcmpl-x", + "object": "chat.completion", + "created": 1, + "model": "gpt-4", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "hi"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, +} + + +class TestClientsideBaseOverrideOutboundKey: + """Drive a completion through the router and assert on the outbound request + when the caller overrides ``api_base``.""" + + def _router(self): + from litellm import Router + + return Router( + model_list=[ + { + "model_name": "gpt-4", + "litellm_params": { + "model": "openai/gpt-4", + "api_key": "sk-SERVER-CONFIG", + "api_base": "https://admin.upstream/v1", + }, + } + ] + ) + + @pytest.fixture(autouse=True) + def _ambient_server_key(self, monkeypatch): + import litellm + + monkeypatch.setenv("OPENAI_API_KEY", "sk-SERVER-ENV") + monkeypatch.setattr(litellm, "api_key", None, raising=False) + + def test_caller_key_override_sends_caller_key_never_server_key(self): + import httpx + import respx + + with respx.mock: + route = respx.post("https://caller.example/v1/chat/completions").mock( + return_value=httpx.Response(200, json=_OPENAI_CHAT_RESPONSE) + ) + self._router().completion( + model="gpt-4", + messages=[{"role": "user", "content": "hi"}], + api_base="https://caller.example/v1", + api_key="sk-CALLER", + ) + authorization = route.calls.last.request.headers.get("authorization") + assert authorization == "Bearer sk-CALLER" + assert "SERVER" not in (authorization or "") + + +def _rounds_deep_api_base_payload(rounds, field): + """Build a fallbacks payload with ``api_base`` on a target nested ``rounds`` + fallback-rounds deep, each round wrapped in its own grouping dict.""" + node = {"model": "leaf", "api_base": "https://attacker.example"} + for i in range(rounds): + node = {"model": f"m{i}", field: [{"grp": [node]}]} + return {"model": "gpt-4", field: [{"grp": [node]}]} + + +class TestIsRequestBodySafeBlocksFallbackSmuggle: + """``is_request_body_safe`` runs the banned-param check on every dict target + inside the fallback lists.""" + + @pytest.fixture(autouse=True) + def _disable_url_validation(self, monkeypatch): + import litellm + + monkeypatch.setattr(litellm, "user_url_validation", False, raising=False) + + @pytest.mark.parametrize( + "fallback_key", + ["fallbacks", "context_window_fallbacks", "content_policy_fallbacks"], + ) + def test_api_base_smuggled_via_nested_fallback_is_rejected(self, fallback_key): + with pytest.raises(ValueError, match="api_base"): + is_request_body_safe( + request_body={ + "model": "gpt-4", + fallback_key: [ + { + "gpt-4": [ + {"model": "evil", "api_base": "https://attacker.example"}, + ] + } + ], + }, + general_settings={}, + llm_router=None, + model="gpt-4", + ) + + def test_string_only_fallbacks_are_accepted(self): + assert ( + is_request_body_safe( + request_body={ + "model": "gpt-4", + "fallbacks": [{"gpt-4": ["gpt-3.5-turbo", "claude-3-haiku"]}], + }, + general_settings={}, + llm_router=None, + model="gpt-4", + ) + is True + ) + + def test_benign_dict_fallback_entry_is_accepted(self): + assert ( + is_request_body_safe( + request_body={ + "model": "gpt-4", + "fallbacks": [{"gpt-4": [{"model": "gpt-3.5-turbo"}]}], + }, + general_settings={}, + llm_router=None, + model="gpt-4", + ) + is True + ) + + def test_smuggled_fallback_allowed_under_proxy_wide_opt_in(self): + assert ( + is_request_body_safe( + request_body={ + "model": "gpt-4", + "fallbacks": [ + {"gpt-4": [{"model": "byok", "api_base": "https://my-byok.example"}]} + ], + }, + general_settings={"allow_client_side_credentials": True}, + llm_router=None, + model="gpt-4", + ) + is True + ) + + @pytest.mark.parametrize( + "fallback_field", + ["fallbacks", "context_window_fallbacks", "content_policy_fallbacks"], + ) + @pytest.mark.parametrize("surface", ["top_level", "router_settings_override"]) + def test_deeply_nested_api_base_smuggle_rejected_on_both_surfaces(self, fallback_field, surface): + nested = [ + { + "always-fail": [ + { + "model": "x", + fallback_field: [ + {"x": [{"model": "deepseek-chat", "api_base": "http://attacker"}]} + ], + } + ] + } + ] + request_body = {"model": "gpt-4"} + if surface == "top_level": + request_body[fallback_field] = nested + else: + request_body["router_settings_override"] = {fallback_field: nested} + with pytest.raises(ValueError, match="api_base"): + is_request_body_safe( + request_body=request_body, + general_settings={}, + llm_router=None, + model="gpt-4", + ) + + def test_router_settings_override_single_level_api_base_rejected(self): + with pytest.raises(ValueError, match="api_base"): + is_request_body_safe( + request_body={ + "model": "gpt-4", + "router_settings_override": { + "fallbacks": [{"gpt-4": [{"model": "x", "api_base": "http://attacker"}]}] + }, + }, + general_settings={}, + llm_router=None, + model="gpt-4", + ) + + def test_model_less_config_dict_api_base_rejected(self): + with pytest.raises(ValueError, match="api_base"): + is_request_body_safe( + request_body={ + "model": "gpt-4", + "fallbacks": [{"gpt-4": [{"api_base": "http://attacker"}]}], + }, + general_settings={}, + llm_router=None, + model="gpt-4", + ) + + def test_nested_api_base_caught_across_router_fallback_rounds(self): + """An ``api_base`` target nested ``ROUTER_MAX_FALLBACKS - 1`` rounds deep + is still reached and rejected.""" + import litellm + + with pytest.raises(ValueError, match="api_base"): + is_request_body_safe( + request_body=_rounds_deep_api_base_payload(litellm.ROUTER_MAX_FALLBACKS - 1, "fallbacks"), + general_settings={}, + llm_router=None, + model="gpt-4", + ) + + def test_grouping_only_deep_chain_is_rejected_at_depth_limit(self): + """A deep grouping-only chain (``{"g": [{"g": [...]}]}``) is rejected at the + validation-depth limit rather than accepted or raising RecursionError.""" + node: object = ["safe-model"] + for _ in range(5000): + node = [{"grp": node}] + with pytest.raises(ValueError, match="depth"): + is_request_body_safe( + request_body={"model": "gpt-4", "fallbacks": node}, + general_settings={}, + llm_router=None, + model="gpt-4", + ) + + def test_pathologically_deep_model_nesting_is_rejected(self): + with pytest.raises(ValueError, match="depth"): + is_request_body_safe( + request_body=_rounds_deep_api_base_payload(5000, "fallbacks"), + general_settings={}, + llm_router=None, + model="gpt-4", + ) + + +class TestIsRequestBodySafeRejectsUrlValuedFallback: + @pytest.mark.parametrize("fallback_field", ["fallbacks", "context_window_fallbacks", "content_policy_fallbacks"]) + def test_url_valued_string_fallback_is_rejected(self, fallback_field): + with pytest.raises(ValueError, match="URL-valued fallback"): + is_request_body_safe( + request_body={ + "model": "gpt-4", + fallback_field: [{"gpt-4": ["huggingface/http://attacker.example/path"]}], + }, + general_settings={}, + llm_router=None, + model="gpt-4", + ) + + @pytest.mark.parametrize("fallback_field", ["fallbacks", "context_window_fallbacks", "content_policy_fallbacks"]) + def test_url_valued_dict_model_fallback_is_rejected(self, fallback_field): + with pytest.raises(ValueError, match="URL-valued fallback"): + is_request_body_safe( + request_body={ + "model": "gpt-4", + fallback_field: [{"gpt-4": [{"model": "huggingface/http://attacker.example/path"}]}], + }, + general_settings={}, + llm_router=None, + model="gpt-4", + ) + + def test_ordinary_string_fallback_is_allowed(self): + assert ( + is_request_body_safe( + request_body={"model": "gpt-4", "fallbacks": [{"gpt-4": ["gpt-4-backup"]}]}, + general_settings={}, + llm_router=None, + model="gpt-4", + ) + is True + ) + + def test_ordinary_dict_model_fallback_is_allowed(self): + assert ( + is_request_body_safe( + request_body={"model": "gpt-4", "fallbacks": [{"gpt-4": [{"model": "gpt-4-backup"}]}]}, + general_settings={}, + llm_router=None, + model="gpt-4", + ) + is True + ) + class TestIsRequestBodySafeBlocksEndpointTargetingFields: """ @@ -1823,6 +2129,46 @@ class TestIsRequestBodySafeBlocksBedrockProjectOverride: ) +class TestIsRequestBodySafeBlocksVertexCredentialAlias: + @pytest.mark.parametrize("field", ["vertex_ai_credentials"]) + def test_field_in_request_body_is_rejected(self, field): + with pytest.raises(ValueError, match=field): + is_request_body_safe( + request_body={"model": "gpt-4", field: "attacker-supplied"}, + general_settings={}, + llm_router=None, + model="gpt-4", + ) + + @pytest.mark.parametrize("field", ["vertex_ai_credentials"]) + def test_admin_opt_in_proxy_wide_allows(self, field): + assert ( + is_request_body_safe( + request_body={"model": "gpt-4", field: "byok-supplied"}, + general_settings={"allow_client_side_credentials": True}, + llm_router=None, + model="gpt-4", + ) + is True + ) + + def test_legitimate_request_body_param_still_allowed(self): + assert ( + is_request_body_safe( + request_body={ + "model": "gpt-4", + "temperature": 0.7, + "max_tokens": 128, + "user": "end-user-123", + }, + general_settings={}, + llm_router=None, + model="gpt-4", + ) + is True + ) + + class TestIsRequestBodySafeBlocksNVCFFunctionOverride: """``nvcf_function_id`` is rejected as a request-body param unless the admin opted in proxy-wide or per-deployment.""" diff --git a/tests/test_litellm/proxy/auth/test_login_utils.py b/tests/test_litellm/proxy/auth/test_login_utils.py index 288e2533b72..c589014f276 100644 --- a/tests/test_litellm/proxy/auth/test_login_utils.py +++ b/tests/test_litellm/proxy/auth/test_login_utils.py @@ -176,9 +176,7 @@ async def test_authenticate_user_invalid_credentials(): mock_prisma_client = MagicMock() mock_prisma_client.db.litellm_usertable.find_first = AsyncMock(return_value=None) - with patch.dict( - os.environ, {"UI_USERNAME": ui_username, "UI_PASSWORD": "correct-password"} - ): + with patch.dict(os.environ, {"UI_USERNAME": ui_username, "UI_PASSWORD": "correct-password"}): with pytest.raises(ProxyException) as exc_info: await authenticate_user( username=ui_username, @@ -227,9 +225,7 @@ async def test_authenticate_user_wrong_password(): ) mock_prisma_client = MagicMock() - mock_prisma_client.db.litellm_usertable.find_first = AsyncMock( - return_value=mock_user - ) + mock_prisma_client.db.litellm_usertable.find_first = AsyncMock(return_value=mock_user) with patch.dict( os.environ, @@ -279,9 +275,7 @@ async def test_authenticate_user_email_case_insensitive_login(): return None mock_prisma_client = MagicMock() - mock_prisma_client.db.litellm_usertable.find_first = AsyncMock( - side_effect=mock_find_first - ) + mock_prisma_client.db.litellm_usertable.find_first = AsyncMock(side_effect=mock_find_first) with patch.dict( os.environ, @@ -334,9 +328,7 @@ async def test_authenticate_user_database_required_for_admin(): mock_prisma_client = MagicMock() mock_prisma_client.db.litellm_usertable.find_first = AsyncMock(return_value=None) - with patch.dict( - os.environ, {"UI_USERNAME": ui_username, "UI_PASSWORD": ui_password} - ): + with patch.dict(os.environ, {"UI_USERNAME": ui_username, "UI_PASSWORD": ui_password}): with patch( "litellm.proxy.auth.login_utils.user_update", new_callable=AsyncMock, @@ -429,9 +421,7 @@ def test_authenticate_user_non_ascii_direct_comparison(): assert result is True # And correctly returns False for different passwords - result = secrets.compare_digest( - password.encode("utf-8"), "different£pass".encode("utf-8") - ) + result = secrets.compare_digest(password.encode("utf-8"), "different£pass".encode("utf-8")) assert result is False @@ -531,9 +521,7 @@ async def test_authenticate_user_database_login_with_non_ascii_password(): return None mock_prisma_client = MagicMock() - mock_prisma_client.db.litellm_usertable.find_first = AsyncMock( - side_effect=mock_find_first - ) + mock_prisma_client.db.litellm_usertable.find_first = AsyncMock(side_effect=mock_find_first) with patch.dict( os.environ, @@ -559,3 +547,58 @@ async def test_authenticate_user_database_login_with_non_ascii_password(): assert isinstance(result, LoginResult) assert result.user_id == "test-user-123" assert result.user_email == user_email + + +class TestEncodeUiSessionJwt: + """The UI session cookie must carry a bounded exp so it does not stay + signature-valid until the master key rotates, and so the session-cookie readers + that require a bounded lifetime (the MCP interactive sign-in) accept it.""" + + def _decode(self, token: str) -> dict: + import jwt + + return jwt.decode(token, "sk-master-for-tests", algorithms=["HS256"]) + + def test_encoded_cookie_carries_bounded_exp(self): + import time + + from litellm.proxy.auth.login_utils import encode_ui_session_jwt + + token_object = {"user_id": "u1", "key": "sk-abc", "login_method": "username_password"} + with patch("litellm.proxy.auth.login_utils.LITELLM_UI_SESSION_DURATION", "24h"): + token = encode_ui_session_jwt(token_object, "sk-master-for-tests") + claims = self._decode(token) + assert claims["user_id"] == "u1" + assert claims["login_method"] == "username_password" + remaining = claims["exp"] - int(time.time()) + assert 23 * 3600 < remaining <= 24 * 3600 + + def test_duration_is_honored_from_env(self): + import time + + from litellm.proxy.auth.login_utils import encode_ui_session_jwt + + with patch("litellm.proxy.auth.login_utils.LITELLM_UI_SESSION_DURATION", "1h"): + token = encode_ui_session_jwt({"user_id": "u1"}, "sk-master-for-tests") + remaining = self._decode(token)["exp"] - int(time.time()) + assert 0 < remaining <= 3600 + + def test_cookie_is_accepted_by_the_exp_requiring_session_reader(self): + """The regression this change exists for: before it, the UI cookie carried no + exp and _user_id_from_session_cookie (require=["exp"]) rejected every real login, + so the MCP interactive sign-in could never capture identity. A cookie minted by + this helper must now be accepted.""" + from unittest.mock import MagicMock + + from litellm.proxy._experimental.mcp_server.byok_oauth_endpoints import ( + _user_id_from_session_cookie, + ) + from litellm.proxy.auth.login_utils import encode_ui_session_jwt + + token_object = {"user_id": "cornell-user", "key": "sk-abc", "login_method": "sso"} + with patch("litellm.proxy.auth.login_utils.LITELLM_UI_SESSION_DURATION", "24h"): + token = encode_ui_session_jwt(token_object, "sk-master-for-tests") + request = MagicMock() + request.cookies = {"token": token} + with patch("litellm.proxy.proxy_server.master_key", "sk-master-for-tests"): + assert _user_id_from_session_cookie(request) == "cornell-user" diff --git a/tests/test_litellm/proxy/auth/test_router_override_fallback_auth.py b/tests/test_litellm/proxy/auth/test_router_override_fallback_auth.py index fc0e9aec501..eb1135a240a 100644 --- a/tests/test_litellm/proxy/auth/test_router_override_fallback_auth.py +++ b/tests/test_litellm/proxy/auth/test_router_override_fallback_auth.py @@ -11,12 +11,22 @@ from unittest.mock import AsyncMock, patch import pytest from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth +from litellm.proxy.auth.auth_utils import iter_request_fallback_targets from litellm.proxy.auth.user_api_key_auth import ( _enforce_key_and_fallback_model_access, - iter_router_fallback_model_names, + _fallback_target_model_name, ) +def _fallback_model_names(fallbacks): + """Model names the auth check validates for a top-level ``fallbacks`` value.""" + return [ + name + for target in iter_request_fallback_targets({"fallbacks": fallbacks}) + if (name := _fallback_target_model_name(target)) is not None + ] + + def _key_with_models(models: List[str]) -> UserAPIKeyAuth: return UserAPIKeyAuth( api_key="hashed", @@ -26,37 +36,40 @@ def _key_with_models(models: List[str]) -> UserAPIKeyAuth: ) -# ── iter_router_fallback_model_names ───────────────────────────────────────── +# ── fallback model-name extraction ─────────────────────────────────────────── -def testiter_router_fallback_model_names_router_config_shape(): +def test_fallback_model_names_router_config_shape(): """Router-config shape: ``[{primary: [fallback_list]}]``.""" - assert list( - iter_router_fallback_model_names( - [{"gpt-3.5-turbo": ["gpt-4", "claude-3"]}, {"gpt-4o": ["o1"]}] - ) + assert _fallback_model_names( + [{"gpt-3.5-turbo": ["gpt-4", "claude-3"]}, {"gpt-4o": ["o1"]}] ) == ["gpt-4", "claude-3", "o1"] -def testiter_router_fallback_model_names_simple_string_shape(): +def test_fallback_model_names_simple_string_shape(): """Simple top-level shape: list of strings.""" - assert list(iter_router_fallback_model_names(["gpt-4", "claude-3"])) == [ + assert _fallback_model_names(["gpt-4", "claude-3"]) == ["gpt-4", "claude-3"] + + +def test_fallback_model_names_client_side_shape(): + """ClientSideFallbackModel shape: ``[{"model": "..."}]``.""" + assert _fallback_model_names([{"model": "gpt-4"}, {"model": "claude-3"}]) == [ "gpt-4", "claude-3", ] -def testiter_router_fallback_model_names_client_side_shape(): - """ClientSideFallbackModel shape: ``[{"model": "..."}]``.""" - assert list( - iter_router_fallback_model_names([{"model": "gpt-4"}, {"model": "claude-3"}]) - ) == ["gpt-4", "claude-3"] +def test_fallback_model_names_nested_deployment_fallbacks(): + """A deployment target's own nested fallback field is unrolled too.""" + assert _fallback_model_names( + [{"primary": [{"model": "gpt-4", "fallbacks": [{"gpt-4": ["deepseek-chat"]}]}]}] + ) == ["gpt-4", "deepseek-chat"] -def testiter_router_fallback_model_names_empty_or_none(): - assert list(iter_router_fallback_model_names(None)) == [] - assert list(iter_router_fallback_model_names([])) == [] - assert list(iter_router_fallback_model_names("not a list")) == [] +def test_fallback_model_names_empty_or_none(): + assert _fallback_model_names(None) == [] + assert _fallback_model_names([]) == [] + assert _fallback_model_names("not a list") == [] # ── _enforce_key_and_fallback_model_access ──────────────────────────────────── @@ -200,6 +213,98 @@ async def test_top_level_fallback_fields_validated(fallback_field): assert "top-level-smuggled" in seen +@pytest.mark.asyncio +async def test_nested_deployment_fallback_inner_model_validated(): + """A model name nested several fallback rounds deep, inside a deployment + target's own ``fallbacks``, is extracted and passed to can_key_call_model.""" + valid_token = _key_with_models(["gpt-3.5-turbo"]) + request_data = { + "model": "gpt-3.5-turbo", + "fallbacks": [ + { + "gpt-3.5-turbo": [ + { + "model": "gpt-3.5-turbo", + "fallbacks": [{"gpt-3.5-turbo": ["deep-smuggled-model"]}], + } + ] + } + ], + } + + seen: List[str] = [] + + async def fake_can_key_call_model(model, llm_model_list, valid_token, llm_router): + seen.append(model) + + with ( + patch( + "litellm.proxy.auth.user_api_key_auth.can_key_call_model", + side_effect=fake_can_key_call_model, + ), + patch( + "litellm.proxy.auth.user_api_key_auth.is_valid_fallback_model", + new=AsyncMock(), + ), + ): + await _enforce_key_and_fallback_model_access( + valid_token=valid_token, + request_data=request_data, + route="/v1/chat/completions", + request=None, + llm_model_list=None, + llm_router=None, + ) + + assert "deep-smuggled-model" in seen + + +@pytest.mark.asyncio +async def test_model_less_fallback_dict_is_skipped_never_passed_as_none(): + """A fallback target dict without a ``model`` key is skipped, never passed + as ``None`` into can_key_call_model / is_valid_fallback_model.""" + valid_token = _key_with_models(["gpt-3.5-turbo"]) + request_data = { + "model": "gpt-3.5-turbo", + "fallbacks": [ + { + "gpt-3.5-turbo": [ + {"model": "real-fallback"}, + {"api_base": "http://attacker"}, + "string-fallback", + ] + } + ], + } + + seen: List[str] = [] + + async def fake_can_key_call_model(model, llm_model_list, valid_token, llm_router): + seen.append(model) + + with ( + patch( + "litellm.proxy.auth.user_api_key_auth.can_key_call_model", + side_effect=fake_can_key_call_model, + ), + patch( + "litellm.proxy.auth.user_api_key_auth.is_valid_fallback_model", + new=AsyncMock(), + ), + ): + await _enforce_key_and_fallback_model_access( + valid_token=valid_token, + request_data=request_data, + route="/v1/chat/completions", + request=None, + llm_model_list=None, + llm_router=None, + ) + + assert None not in seen + assert seen == ["gpt-3.5-turbo", "real-fallback", "string-fallback"] + + @pytest.mark.asyncio async def test_router_override_without_fallbacks_does_not_break_auth(): """``router_settings_override`` set without any fallback fields is a diff --git a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py index c80a6a1ca9a..9c33af2b862 100644 --- a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py +++ b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py @@ -179,6 +179,45 @@ async def test_budget_reservation_runs_when_not_disabled(): assert user_api_key_auth_obj.budget_reservation == reservation +@pytest.mark.asyncio +@pytest.mark.parametrize( + "general_settings,expected_flag", + [ + ({"fail_closed_budget_enforcement": True}, True), + ({}, False), + ], +) +async def test_fail_closed_budget_enforcement_reaches_reservation( + general_settings, expected_flag +): + """#33923: the strict flag must be threaded into reserve_budget_for_request so a + failed reservation write can reject instead of failing open.""" + user_api_key_auth_obj = UserAPIKeyAuth(token="test_token") + + with patch( + "litellm.proxy.spend_tracking.budget_reservation.reserve_budget_for_request", + new=AsyncMock(return_value=None), + ) as mock_reserve: + await _reserve_budget_after_common_checks( + user_api_key_auth_obj=user_api_key_auth_obj, + request_data={"model": "gpt-4o"}, + route="/v1/chat/completions", + llm_router=None, + team_object=None, + user_object=None, + prisma_client=None, + user_api_key_cache=MagicMock(), + proxy_logging_obj=MagicMock(), + skip_budget_checks=False, + general_settings=general_settings, + ) + + assert ( + mock_reserve.await_args.kwargs["fail_closed_budget_enforcement"] + is expected_flag + ) + + @pytest.mark.asyncio async def test_should_not_reuse_cached_key_object_for_request_state(): key_cache = DualCache() @@ -1290,6 +1329,250 @@ async def test_scim_deactivated_user_key_is_rejected(): setattr(_proxy_server_mod, attr, val) +@pytest.mark.asyncio +async def test_cached_proxy_admin_key_sets_via_virtual_key_marker(): + """Cached PROXY_ADMIN auth objects early-return before the marked DB and + master-key returns, and cache serialization drops the exclude=True marker; + the cache-hit boundary must restore it or cached admin traffic silently + bypasses overwrite_user_with_key_hash stamping.""" + from fastapi import Request + from starlette.datastructures import URL + + from litellm.proxy.auth.user_api_key_auth import _user_api_key_auth_builder + from litellm.proxy.proxy_server import hash_token + + api_key = "sk-cached-admin-marker-test" + hashed_key = hash_token(api_key) + + cached_token = UserAPIKeyAuth( + api_key=api_key, + token=hashed_key, + user_id="cached-admin-user", + user_role=LitellmUserRoles.PROXY_ADMIN, + ) + assert cached_token.via_virtual_key is False + + mock_cache = AsyncMock() + mock_cache.async_get_cache = AsyncMock(return_value=None) + mock_cache.delete_cache = MagicMock() + + mock_proxy_logging_obj = MagicMock() + mock_proxy_logging_obj.internal_usage_cache = MagicMock() + mock_proxy_logging_obj.internal_usage_cache.dual_cache = AsyncMock() + mock_proxy_logging_obj.internal_usage_cache.dual_cache.async_delete_cache = ( + AsyncMock() + ) + mock_proxy_logging_obj.post_call_failure_hook = AsyncMock(return_value=None) + + import litellm.proxy.proxy_server as _proxy_server_mod + + _attrs_to_set = { + "prisma_client": MagicMock(), + "user_api_key_cache": mock_cache, + "proxy_logging_obj": mock_proxy_logging_obj, + "master_key": "sk-master-key", + "general_settings": {}, + "llm_model_list": [], + "llm_router": None, + "open_telemetry_logger": None, + "model_max_budget_limiter": MagicMock(), + "user_custom_auth": None, + "jwt_handler": None, + "litellm_proxy_admin_name": "admin", + } + _original_values = { + attr: getattr(_proxy_server_mod, attr, None) for attr in _attrs_to_set + } + try: + for attr, val in _attrs_to_set.items(): + setattr(_proxy_server_mod, attr, val) + + request = Request(scope={"type": "http"}) + request._url = URL(url="/chat/completions") + + with patch( + "litellm.proxy.auth.resolvers.store.IdentityStore._resolve_key", + new_callable=AsyncMock, + return_value=cached_token, + ): + result = await _user_api_key_auth_builder( + request=request, + api_key=f"Bearer {api_key}", + azure_api_key_header="", + anthropic_api_key_header=None, + google_ai_studio_api_key_header=None, + azure_apim_header=None, + request_data={}, + ) + + assert isinstance(result, UserAPIKeyAuth) + assert result.user_role == LitellmUserRoles.PROXY_ADMIN + assert result.via_virtual_key is True + assert result.api_key == hashed_key + finally: + for attr, val in _original_values.items(): + setattr(_proxy_server_mod, attr, val) + + +@pytest.mark.asyncio +async def test_master_key_auth_sets_via_virtual_key_marker(): + """Master-key requests must also be stamped by overwrite_user_with_key_hash; + the auth path substitutes the stable alias for api_key and must mark the + result as proxy-validated.""" + from fastapi import Request + from starlette.datastructures import URL + + from litellm.constants import LITELLM_PROXY_MASTER_KEY_ALIAS + from litellm.proxy.auth.user_api_key_auth import _user_api_key_auth_builder + + master_key = "sk-master-key" + + mock_cache = AsyncMock() + mock_cache.async_get_cache = AsyncMock(return_value=None) + mock_cache.delete_cache = MagicMock() + + mock_proxy_logging_obj = MagicMock() + mock_proxy_logging_obj.internal_usage_cache = MagicMock() + mock_proxy_logging_obj.internal_usage_cache.dual_cache = AsyncMock() + mock_proxy_logging_obj.internal_usage_cache.dual_cache.async_delete_cache = ( + AsyncMock() + ) + mock_proxy_logging_obj.post_call_failure_hook = AsyncMock(return_value=None) + + import litellm.proxy.proxy_server as _proxy_server_mod + + _attrs_to_set = { + "prisma_client": MagicMock(), + "user_api_key_cache": mock_cache, + "proxy_logging_obj": mock_proxy_logging_obj, + "master_key": master_key, + "general_settings": {}, + "llm_model_list": [], + "llm_router": None, + "open_telemetry_logger": None, + "model_max_budget_limiter": MagicMock(), + "user_custom_auth": None, + "jwt_handler": None, + "litellm_proxy_admin_name": "admin", + } + _original_values = { + attr: getattr(_proxy_server_mod, attr, None) for attr in _attrs_to_set + } + try: + for attr, val in _attrs_to_set.items(): + setattr(_proxy_server_mod, attr, val) + + request = Request(scope={"type": "http"}) + request._url = URL(url="/chat/completions") + + result = await _user_api_key_auth_builder( + request=request, + api_key=f"Bearer {master_key}", + azure_api_key_header="", + anthropic_api_key_header=None, + google_ai_studio_api_key_header=None, + azure_apim_header=None, + request_data={}, + ) + + assert isinstance(result, UserAPIKeyAuth) + assert result.via_virtual_key is True + assert result.api_key == LITELLM_PROXY_MASTER_KEY_ALIAS + finally: + for attr, val in _original_values.items(): + setattr(_proxy_server_mod, attr, val) + + +@pytest.mark.asyncio +async def test_db_virtual_key_auth_sets_via_virtual_key_marker(): + """via_virtual_key gates overwrite_user_with_key_hash stamping and is + forge-stripped from validated input, so the DB auth path setting it by + post-construction assignment is the only thing that turns stamping on.""" + from fastapi import Request + from starlette.datastructures import URL + + from litellm.proxy.auth.user_api_key_auth import _user_api_key_auth_builder + from litellm.proxy.proxy_server import hash_token + + api_key = "sk-via-virtual-key-marker-test" + hashed_key = hash_token(api_key) + + valid_token = UserAPIKeyAuth( + api_key=api_key, + token=hashed_key, + user_id="marker-test-user", + ) + + mock_cache = AsyncMock() + mock_cache.async_get_cache = AsyncMock(return_value=None) + mock_cache.delete_cache = MagicMock() + + mock_proxy_logging_obj = MagicMock() + mock_proxy_logging_obj.internal_usage_cache = MagicMock() + mock_proxy_logging_obj.internal_usage_cache.dual_cache = AsyncMock() + mock_proxy_logging_obj.internal_usage_cache.dual_cache.async_delete_cache = ( + AsyncMock() + ) + mock_proxy_logging_obj.post_call_failure_hook = AsyncMock(return_value=None) + + mock_prisma_client = MagicMock() + + import litellm.proxy.proxy_server as _proxy_server_mod + + _attrs_to_set = { + "prisma_client": mock_prisma_client, + "user_api_key_cache": mock_cache, + "proxy_logging_obj": mock_proxy_logging_obj, + "master_key": "sk-master-key", + "general_settings": {}, + "llm_model_list": [], + "llm_router": None, + "open_telemetry_logger": None, + "model_max_budget_limiter": MagicMock(), + "user_custom_auth": None, + "jwt_handler": None, + "litellm_proxy_admin_name": "admin", + } + _original_values = { + attr: getattr(_proxy_server_mod, attr, None) for attr in _attrs_to_set + } + try: + for attr, val in _attrs_to_set.items(): + setattr(_proxy_server_mod, attr, val) + + request = Request(scope={"type": "http"}) + request._url = URL(url="/chat/completions") + + with ( + patch( + "litellm.proxy.auth.resolvers.store.IdentityStore._resolve_key", + new_callable=AsyncMock, + return_value=valid_token, + ), + patch( + "litellm.proxy.auth.user_api_key_auth.get_user_object", + new_callable=AsyncMock, + return_value=None, + ), + ): + result = await _user_api_key_auth_builder( + request=request, + api_key=f"Bearer {api_key}", + azure_api_key_header="", + anthropic_api_key_header=None, + google_ai_studio_api_key_header=None, + azure_apim_header=None, + request_data={}, + ) + + assert isinstance(result, UserAPIKeyAuth) + assert result.via_virtual_key is True + assert result.api_key == hashed_key + finally: + for attr, val in _original_values.items(): + setattr(_proxy_server_mod, attr, val) + + @pytest.mark.asyncio async def test_return_user_api_key_auth_obj_user_spend_and_budget(): """ diff --git a/tests/test_litellm/proxy/client/cli/autoroute/test_config.py b/tests/test_litellm/proxy/client/cli/autoroute/test_config.py index 6ab484d5004..f399b6f957a 100644 --- a/tests/test_litellm/proxy/client/cli/autoroute/test_config.py +++ b/tests/test_litellm/proxy/client/cli/autoroute/test_config.py @@ -48,27 +48,22 @@ def _base_config(**overrides: Any) -> AutorouteConfig: class TestParseDiscoveredModels: def test_parses_valid_raw_list_into_typed_tuple(self): raw = [ - { - "model_group": "gpt-4o", - "mode": "chat", - "input_cost_per_token": 0.01, - "output_cost_per_token": 0.02, - }, - {"model_group": "text-embedding-3-small", "mode": "embedding"}, + {"id": "gpt-4o", "object": "model", "mode": "chat"}, + {"id": "text-embedding-3-small", "object": "model", "mode": "embedding"}, ] result = parse_discovered_models(raw) assert result == ( - DiscoveredModel(name="gpt-4o", mode="chat", input_cost_per_token=0.01, output_cost_per_token=0.02), + DiscoveredModel(name="gpt-4o", mode="chat"), DiscoveredModel(name="text-embedding-3-small", mode="embedding"), ) def test_ignores_unknown_extra_fields(self): - raw = [{"model_group": "gpt-4o", "mode": "chat", "totally_unknown_field": "whatever"}] + raw = [{"id": "gpt-4o", "mode": "chat", "created": 123, "owned_by": "openai", "max_input_tokens": 128000}] result = parse_discovered_models(raw) assert result == (DiscoveredModel(name="gpt-4o", mode="chat"),) def test_missing_mode_defaults_to_chat(self): - raw = [{"model_group": "gpt-4o"}] + raw = [{"id": "gpt-4o", "object": "model"}] result = parse_discovered_models(raw) assert result[0].mode == "chat" diff --git a/tests/test_litellm/proxy/client/cli/autoroute/test_wizard.py b/tests/test_litellm/proxy/client/cli/autoroute/test_wizard.py index 78d4bd20338..a17fed36f52 100644 --- a/tests/test_litellm/proxy/client/cli/autoroute/test_wizard.py +++ b/tests/test_litellm/proxy/client/cli/autoroute/test_wizard.py @@ -16,22 +16,22 @@ from litellm.proxy.client.cli.commands.autoroute.config import DiscoveredModel from litellm.proxy.client.cli.commands.autoroute.wizard import run_configure_wizard CHAT_AND_EMBEDDING_GROUPS: List[Dict[str, Any]] = [ - {"model_group": "gpt-4o-mini", "mode": "chat", "input_cost_per_token": 0.01, "output_cost_per_token": 0.02}, - {"model_group": "gpt-4o", "mode": "chat", "input_cost_per_token": 0.01, "output_cost_per_token": 0.02}, - {"model_group": "claude-opus", "mode": "chat"}, - {"model_group": "o1", "mode": "chat"}, - {"model_group": "text-embedding-3-small", "mode": "embedding"}, + {"id": "gpt-4o-mini", "object": "model", "mode": "chat", "max_input_tokens": 128000}, + {"id": "gpt-4o", "object": "model", "mode": "chat", "max_input_tokens": 128000}, + {"id": "claude-opus", "object": "model", "mode": "chat"}, + {"id": "o1", "object": "model", "mode": "chat"}, + {"id": "text-embedding-3-small", "object": "model", "mode": "embedding"}, ] CHAT_ONLY_GROUPS: List[Dict[str, Any]] = [ - {"model_group": "gpt-4o-mini", "mode": "chat"}, - {"model_group": "gpt-4o", "mode": "chat"}, - {"model_group": "claude-opus", "mode": "chat"}, - {"model_group": "o1", "mode": "chat"}, + {"id": "gpt-4o-mini", "object": "model", "mode": "chat"}, + {"id": "gpt-4o", "object": "model", "mode": "chat"}, + {"id": "claude-opus", "object": "model", "mode": "chat"}, + {"id": "o1", "object": "model", "mode": "chat"}, ] EMBEDDING_ONLY_GROUPS: List[Dict[str, Any]] = [ - {"model_group": "text-embedding-3-small", "mode": "embedding"}, + {"id": "text-embedding-3-small", "object": "model", "mode": "embedding"}, ] @@ -73,7 +73,7 @@ def _run( patch.object(wizard_module, "_render_and_prompt_for_models", side_effect=_fake_prompt_for_models), patch.object(wizard_module, "_render_and_prompt_for_model", side_effect=_fake_prompt_for_model), ): - mock_client_cls.return_value.model_groups.info.return_value = raw_groups + mock_client_cls.return_value.models.list.return_value = raw_groups result = runner.invoke( _invoke_wizard, obj={"base_url": "http://localhost:4000", "api_key": "sk-test"}, @@ -262,7 +262,7 @@ class TestRunConfigureWizardNoChatModels: assert result.exit_code != 0 assert result.exception is None or not isinstance(result.exception, AssertionError) - assert "Unexpected response from /model_group/info" in result.output + assert "Unexpected response from /v1/models" in result.output assert not config_path.exists() @@ -275,7 +275,7 @@ class TestRunConfigureWizardNotInteractive: patch.object(wizard_module, "CONFIG_PATH", config_path), patch.object(wizard_module, "_is_interactive", return_value=False), ): - mock_client_cls.return_value.model_groups.info.return_value = CHAT_AND_EMBEDDING_GROUPS + mock_client_cls.return_value.models.list.return_value = CHAT_AND_EMBEDDING_GROUPS result = runner.invoke(_invoke_wizard, obj={"base_url": "http://localhost:4000", "api_key": "sk-test"}) assert result.exit_code != 0 diff --git a/tests/test_litellm/proxy/common_utils/test_callback_utils.py b/tests/test_litellm/proxy/common_utils/test_callback_utils.py index 36ff3f3c399..8f390c096d7 100644 --- a/tests/test_litellm/proxy/common_utils/test_callback_utils.py +++ b/tests/test_litellm/proxy/common_utils/test_callback_utils.py @@ -84,6 +84,52 @@ def test_process_callback_with_no_required_env_vars(mock_get_env_vars): assert result["variables"] == {} +@patch( + "litellm.proxy.common_utils.callback_utils.CustomLogger.get_callback_env_vars", + return_value=["LANGFUSE_PUBLIC_KEY", "LANGFUSE_SECRET_KEY"], +) +def test_process_callback_falls_back_to_process_env(mock_get_env_vars, monkeypatch): + """A callback env var set only in the process env must be surfaced. + + The logging integrations read their config from the process environment, so a + callback configured purely via env vars (IaC) is live even with no stored + entry. Reporting it as unset makes a working callback read as unconfigured. + """ + monkeypatch.setenv("LANGFUSE_PUBLIC_KEY", "env-public-key") + monkeypatch.setenv("LANGFUSE_SECRET_KEY", "env-secret-key") + # stored config only carries the public key; the secret is env-only + environment_variables = {"LANGFUSE_PUBLIC_KEY": "db-public-key"} + + result = process_callback( + _callback="langfuse", + callback_type="success", + environment_variables=environment_variables, + ) + + # stored value wins; the env-only var is resolved rather than reported None + assert result["variables"] == { + "LANGFUSE_PUBLIC_KEY": "db-public-key", + "LANGFUSE_SECRET_KEY": "env-secret-key", + } + + +@patch( + "litellm.proxy.common_utils.callback_utils.CustomLogger.get_callback_env_vars", + return_value=["LANGFUSE_SECRET_KEY"], +) +def test_process_callback_reports_none_when_absent_everywhere(mock_get_env_vars, monkeypatch): + """A var set in neither the stored config nor the process env stays None.""" + monkeypatch.delenv("LANGFUSE_SECRET_KEY", raising=False) + + result = process_callback( + _callback="langfuse", + callback_type="success", + environment_variables={}, + ) + + assert result["variables"] == {"LANGFUSE_SECRET_KEY": None} + + def test_normalize_callback_names_none_returns_empty_list(): assert normalize_callback_names(None) == [] assert normalize_callback_names([]) == [] diff --git a/tests/test_litellm/proxy/config_resolvers/test_config_resolvers.py b/tests/test_litellm/proxy/config_resolvers/test_config_resolvers.py new file mode 100644 index 00000000000..9f91d9ee2c9 --- /dev/null +++ b/tests/test_litellm/proxy/config_resolvers/test_config_resolvers.py @@ -0,0 +1,126 @@ +import os + +from litellm.proxy.config_resolvers._descriptors import FieldDescriptor, resolve_fields +from litellm.proxy.config_resolvers.sso import ( + SSO_FIELD_ENV_VARS, + SSO_SECRET_FIELDS, + resolve_sso_config, +) + +_D = ( + FieldDescriptor("client_id", "client_id", "CLIENT_ID"), + FieldDescriptor("scope", "scope", "SCOPE", default="openid"), +) + + +def test_resolve_fields_db_wins_over_env(): + values, provenance = resolve_fields(_D, {"client_id": "from-db"}, {"CLIENT_ID": "from-env"}) + assert values["client_id"] == "from-db" + assert provenance["client_id"] == "db" + + +def test_resolve_fields_blank_db_falls_back_to_env(): + values, provenance = resolve_fields(_D, {"client_id": " "}, {"CLIENT_ID": "from-env"}) + assert values["client_id"] == "from-env" + assert provenance["client_id"] == "env" + + +def test_resolve_fields_blank_everywhere_falls_to_default(): + values, provenance = resolve_fields(_D, {}, {"SCOPE": ""}) + assert values["scope"] == "openid" + assert provenance["scope"] == "default" + + +def test_resolve_fields_unset_everywhere(): + values, provenance = resolve_fields(_D, {}, {}) + assert values["client_id"] is None + assert provenance["client_id"] == "unset" + + +def test_resolve_fields_empty_db_absent_by_default_falls_to_env(): + # SSO semantics: a present-but-empty stored value is absent, so env wins. + values, provenance = resolve_fields(_D, {"client_id": ""}, {"CLIENT_ID": "from-env"}) + assert values["client_id"] == "from-env" + assert provenance["client_id"] == "env" + + +def test_resolve_fields_empty_db_is_explicit_clear_when_flag_set(): + # Alerting semantics: a present-but-empty stored value is an explicit clear + # that must win over a stale env var. + values, provenance = resolve_fields( + _D, {"client_id": ""}, {"CLIENT_ID": "stale-env"}, empty_db_is_set=True + ) + assert values["client_id"] == "" + assert provenance["client_id"] == "db" + + +def test_sso_descriptor_mapping_is_single_sourced(): + # The write path and read path both consume this mapping; it must cover every + # env-backed SSO field and map to the uppercase env var. + assert SSO_FIELD_ENV_VARS["generic_client_id"] == "GENERIC_CLIENT_ID" + assert SSO_SECRET_FIELDS == frozenset( + {"google_client_secret", "microsoft_client_secret", "generic_client_secret"} + ) + + +def test_sso_descriptor_mapping_covers_saml_fields(): + # SAML config is stored and read through the same descriptor table as the + # OAuth providers; the login path reads these env vars, so the save path must + # map every SAML field to its uppercase env var. + assert SSO_FIELD_ENV_VARS["saml_idp_metadata_url"] == "SAML_IDP_METADATA_URL" + assert SSO_FIELD_ENV_VARS["saml_idp_metadata_xml"] == "SAML_IDP_METADATA_XML" + assert SSO_FIELD_ENV_VARS["saml_sp_entity_id"] == "SAML_SP_ENTITY_ID" + assert SSO_FIELD_ENV_VARS["saml_allow_unsolicited"] == "SAML_ALLOW_UNSOLICITED" + + +def test_resolve_sso_config_resolves_saml_fields(): + resolved = resolve_sso_config( + {"saml_idp_metadata_url": "https://idp.example.com/metadata"}, + {"SAML_ALLOW_UNSOLICITED": "true"}, + ) + assert resolved.config.saml_idp_metadata_url == "https://idp.example.com/metadata" + assert resolved.provenance["saml_idp_metadata_url"] == "db" + assert resolved.config.saml_allow_unsolicited == "true" + assert resolved.provenance["saml_allow_unsolicited"] == "env" + + +def test_resolve_sso_config_returns_unmasked_secret_and_provenance(): + # The resolver hands back plaintext; masking is the endpoint's job. If the + # resolver masked, the login path would consume a masked secret and fail. + resolved = resolve_sso_config( + {"generic_client_secret": "super-secret-value"}, + {"GENERIC_CLIENT_ID": "env-id"}, + ) + assert resolved.config.generic_client_secret == "super-secret-value" + assert resolved.provenance["generic_client_secret"] == "db" + assert resolved.config.generic_client_id == "env-id" + assert resolved.provenance["generic_client_id"] == "env" + + +def test_resolve_sso_config_parses_structured_mappings(): + resolved = resolve_sso_config( + { + "generic_client_id": "id", + "role_mappings": { + "provider": "generic", + "group_claim": "groups", + "default_role": "internal_user", + "roles": {}, + }, + "team_mappings": {"team_ids_jwt_field": "teams"}, + }, + {}, + ) + assert resolved.config.role_mappings is not None + assert resolved.config.role_mappings.group_claim == "groups" + assert resolved.config.team_mappings is not None + assert resolved.config.team_mappings.team_ids_jwt_field == "teams" + + +def test_resolve_sso_config_does_not_mutate_os_environ(monkeypatch): + # Unlike the legacy read path, resolving must not write os.environ. + monkeypatch.delenv("GENERIC_CLIENT_ID", raising=False) + before = dict(os.environ) + resolve_sso_config({"generic_client_id": "id-from-db"}, os.environ) + assert dict(os.environ) == before + assert "GENERIC_CLIENT_ID" not in os.environ diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py index 15827b80bcf..53f32fd96fb 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py @@ -3274,3 +3274,343 @@ async def test_chat_completion_modify_response_exception_streaming_logging_obj_n # CustomStreamWrapper would raise AttributeError inside __init__ and this # call would never reach here. assert response is not None + + +class TestBedrockOnlyScanNewMessages: + """Bedrock apply_guardrail honors only_scan_new_messages: scans only the per-session diff. + + apply_guardrail is the path the proxy actually runs for Bedrock (via the unified + guardrail interface), so these tests exercise it directly rather than the legacy + async_pre_call_hook. Each test uses a unique session id to isolate the process-wide + incremental cache. + """ + + def _guardrail(self): + return BedrockGuardrail( + guardrail_name="bedrock-incremental", + guardrailIdentifier="test-guardrail", + guardrailVersion="DRAFT", + default_on=True, + only_scan_new_messages=True, + ) + + @pytest.mark.asyncio + async def test_second_turn_scans_only_new_messages(self): + guardrail = self._guardrail() + session = {"litellm_session_id": "sess-bedrock-diff"} + bedrock_none = {"action": "NONE", "output": [], "outputs": []} + + with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api: + mock_api.return_value = bedrock_none + + await guardrail.apply_guardrail( + inputs={"texts": ["be helpful", "first question"]}, + request_data=session, + input_type="request", + ) + assert mock_api.call_count == 1 + first_scanned = mock_api.call_args.kwargs["messages"] + assert [m["content"] for m in first_scanned] == ["be helpful", "first question"] + + mock_api.reset_mock() + + await guardrail.apply_guardrail( + inputs={"texts": ["be helpful", "first question", "first answer", "second question"]}, + request_data=session, + input_type="request", + ) + assert mock_api.call_count == 1 + second_scanned = mock_api.call_args.kwargs["messages"] + assert [m["content"] for m in second_scanned] == ["first answer", "second question"] + + @pytest.mark.asyncio + async def test_identical_resend_skips_api_call(self): + guardrail = self._guardrail() + session = {"litellm_session_id": "sess-bedrock-resend"} + + with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api: + mock_api.return_value = {"action": "NONE", "output": [], "outputs": []} + + await guardrail.apply_guardrail( + inputs={"texts": ["only question"]}, request_data=session, input_type="request" + ) + assert mock_api.call_count == 1 + + mock_api.reset_mock() + result = await guardrail.apply_guardrail( + inputs={"texts": ["only question"]}, request_data=session, input_type="request" + ) + mock_api.assert_not_called() + assert result["texts"] == ["only question"] + + @pytest.mark.asyncio + async def test_no_session_id_scans_full_context(self): + guardrail = self._guardrail() + + with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api: + mock_api.return_value = {"action": "NONE", "output": [], "outputs": []} + + await guardrail.apply_guardrail( + inputs={"texts": ["q1", "a1", "q2"]}, + request_data={"metadata": {}}, + input_type="request", + ) + assert mock_api.call_count == 1 + scanned = mock_api.call_args.kwargs["messages"] + assert [m["content"] for m in scanned] == ["q1", "a1", "q2"] + + @pytest.mark.asyncio + async def test_masking_guardrail_falls_back_and_does_not_persist(self): + """A guardrail that anonymizes content must not be short-circuited. + + Regression: the incremental fast path used to ignore the guardrail response, + so masked/anonymized output was dropped, the raw text reached the model, and + the segment was marked scanned so it was never re-checked. Detecting masked + output must force a full-context scan (which applies the masking) and must not + persist session state, so an identical resend is scanned again. + """ + guardrail = self._guardrail() + session = {"litellm_session_id": "sess-bedrock-mask"} + masked = { + "action": "GUARDRAIL_INTERVENED", + "output": [], + "outputs": [{"text": "my ssn is [REDACTED]"}], + } + + with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api: + mock_api.return_value = masked + + result = await guardrail.apply_guardrail( + inputs={"texts": ["my ssn is 123-45-6789"]}, + request_data=session, + input_type="request", + ) + assert mock_api.call_count == 2 + assert result["texts"] == ["my ssn is [REDACTED]"] + + mock_api.reset_mock() + await guardrail.apply_guardrail( + inputs={"texts": ["my ssn is 123-45-6789"]}, + request_data=session, + input_type="request", + ) + assert mock_api.call_count >= 1 + first_scanned = mock_api.call_args_list[0].kwargs.get("messages") + assert first_scanned is not None + assert [m["content"] for m in first_scanned] == ["my ssn is 123-45-6789"] + + @pytest.mark.asyncio + async def test_generic_agent_multi_turn_scans_only_new_each_turn(self): + """A generic agent (not Claude Code) opts in by propagating a session id. + + Agent frameworks on the OpenAI SDK carry the session through the request + body (metadata.session_id here), not the x-claude-code-session-id header. + Across a growing multi-turn conversation every turn after the first must + send Bedrock only the newly appended segments, never the whole context. + """ + guardrail = self._guardrail() + session = {"metadata": {"session_id": "agent-multi-turn"}} + + with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api: + mock_api.return_value = {"action": "NONE", "output": [], "outputs": []} + + await guardrail.apply_guardrail( + inputs={"texts": ["system prompt", "turn 1 question"]}, + request_data=session, + input_type="request", + ) + assert [m["content"] for m in mock_api.call_args.kwargs["messages"]] == [ + "system prompt", + "turn 1 question", + ] + + mock_api.reset_mock() + await guardrail.apply_guardrail( + inputs={"texts": ["system prompt", "turn 1 question", "turn 1 answer", "turn 2 question"]}, + request_data=session, + input_type="request", + ) + assert [m["content"] for m in mock_api.call_args.kwargs["messages"]] == [ + "turn 1 answer", + "turn 2 question", + ] + + mock_api.reset_mock() + await guardrail.apply_guardrail( + inputs={ + "texts": [ + "system prompt", + "turn 1 question", + "turn 1 answer", + "turn 2 question", + "turn 2 answer", + "turn 3 question", + ] + }, + request_data=session, + input_type="request", + ) + assert [m["content"] for m in mock_api.call_args.kwargs["messages"]] == [ + "turn 2 answer", + "turn 3 question", + ] + + def test_incremental_scan_cache_prefers_proxy_shared_cache(self): + guardrail = self._guardrail() + shared = DualCache() + proxy_logging = MagicMock() + proxy_logging.internal_usage_cache.dual_cache = shared + + with patch("litellm.proxy.proxy_server.proxy_logging_obj", proxy_logging): + assert guardrail._incremental_scan_cache() is shared + + def test_incremental_scan_cache_falls_back_when_proxy_logging_missing(self): + from litellm.integrations.custom_guardrail import dc as fallback_cache + + guardrail = self._guardrail() + with patch("litellm.proxy.proxy_server.proxy_logging_obj", None): + assert guardrail._incremental_scan_cache() is fallback_cache + + def test_incremental_scan_cache_falls_back_when_proxy_not_importable(self): + from litellm.integrations.custom_guardrail import dc as fallback_cache + + guardrail = self._guardrail() + with patch.dict(sys.modules, {"litellm.proxy.proxy_server": None}): + assert guardrail._incremental_scan_cache() is fallback_cache + + @pytest.mark.asyncio + async def test_blocked_turn_is_rescanned_on_retry(self): + guardrail = self._guardrail() + session = {"litellm_session_id": "sess-bedrock-blocked"} + + with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api: + mock_api.side_effect = HTTPException(status_code=400, detail="blocked") + with pytest.raises(HTTPException): + await guardrail.apply_guardrail( + inputs={"texts": ["blocked prompt"]}, request_data=session, input_type="request" + ) + + mock_api.reset_mock() + mock_api.side_effect = None + mock_api.return_value = {"action": "NONE", "output": [], "outputs": []} + await guardrail.apply_guardrail( + inputs={"texts": ["blocked prompt"]}, request_data=session, input_type="request" + ) + assert mock_api.call_count == 1 + scanned = mock_api.call_args.kwargs["messages"] + assert [m["content"] for m in scanned] == ["blocked prompt"] + + +class TestBedrockIncrementalFlagInteractions: + """Regression coverage for only_scan_new_messages combined with the other + Bedrock guardrail flags, from the PR #33278 live validation. Live evidence: + each of these was reproduced against a real Bedrock ApplyGuardrail first; + the mocks here encode the wire payloads observed there. + """ + + def _guardrail(self, **overrides): + params = dict( + guardrail_name="bedrock-incremental-flags", + guardrailIdentifier="test-guardrail", + guardrailVersion="DRAFT", + default_on=True, + only_scan_new_messages=True, + ) + params.update(overrides) + return BedrockGuardrail(**params) + + @pytest.mark.asyncio + async def test_edited_history_segment_rescans_only_that_segment(self): + guardrail = self._guardrail() + session = {"litellm_session_id": "sess-flags-edit"} + with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api: + mock_api.return_value = {"action": "NONE", "output": [], "outputs": []} + await guardrail.apply_guardrail( + inputs={"texts": ["q1", "a1", "q2"]}, request_data=session, input_type="request" + ) + mock_api.reset_mock() + await guardrail.apply_guardrail( + inputs={"texts": ["q1 EDITED", "a1", "q2"]}, request_data=session, input_type="request" + ) + assert mock_api.call_count == 1 + assert [m["content"] for m in mock_api.call_args.kwargs["messages"]] == ["q1 EDITED"] + + @pytest.mark.asyncio + async def test_same_content_different_session_rescans_everything(self): + guardrail = self._guardrail() + texts = ["shared question", "shared answer"] + with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api: + mock_api.return_value = {"action": "NONE", "output": [], "outputs": []} + await guardrail.apply_guardrail( + inputs={"texts": list(texts)}, request_data={"litellm_session_id": "sess-x1"}, input_type="request" + ) + mock_api.reset_mock() + await guardrail.apply_guardrail( + inputs={"texts": list(texts)}, request_data={"litellm_session_id": "sess-x2"}, input_type="request" + ) + assert mock_api.call_count == 1 + assert [m["content"] for m in mock_api.call_args.kwargs["messages"]] == texts + + @pytest.mark.asyncio + async def test_litellm_masking_flag_disables_incremental_single_full_scan(self): + """mask_request_content must fall back to exactly ONE full scan per turn + and never persist hashes (verified live: 1 call/turn, no cache writes).""" + guardrail = self._guardrail(mask_request_content=True) + session = {"litellm_session_id": "sess-flags-mask"} + with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api: + mock_api.return_value = {"action": "NONE", "output": [], "outputs": []} + await guardrail.apply_guardrail( + inputs={"texts": ["q1"]}, request_data=session, input_type="request" + ) + assert mock_api.call_count == 1 + mock_api.reset_mock() + await guardrail.apply_guardrail( + inputs={"texts": ["q1"]}, request_data=session, input_type="request" + ) + assert mock_api.call_count == 1, "masking mode must re-scan every turn, exactly once" + + @pytest.mark.asyncio + async def test_server_side_anonymize_falls_back_full_scan_and_never_persists(self): + """A guardrail that rewrites content (Bedrock-side ANONYMIZE) must fall back + to the full scan so masking applies, and record no session state. Live + validation showed this costs 2 provider calls per turn; the count is + asserted here as documentation of that intended-tradeoff behavior.""" + guardrail = self._guardrail() + session = {"litellm_session_id": "sess-flags-anon"} + masked = {"action": "NONE", "output": [{"text": "MASKED q1"}], "outputs": [{"text": "MASKED q1"}]} + with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api: + mock_api.return_value = masked + result = await guardrail.apply_guardrail( + inputs={"texts": ["q1"]}, request_data=session, input_type="request" + ) + assert mock_api.call_count == 2, "incremental attempt + full-scan fallback" + assert result["texts"] == ["MASKED q1"], "masked content must be applied" + mock_api.reset_mock() + await guardrail.apply_guardrail( + inputs={"texts": ["q1"]}, request_data=session, input_type="request" + ) + assert mock_api.call_count == 2, "no hashes persisted, so the double scan repeats" + + @pytest.mark.asyncio + @pytest.mark.xfail( + reason="PR #33278 known gap: incremental path bypasses _select_messages_for_apply_guardrail, " + "so experimental_use_latest_role_message_only is silently ignored. Intended semantics " + "(pending DRI decision): incremental mode defers to the latest-role selection.", + strict=False, + ) + async def test_latest_role_only_is_respected_with_incremental(self): + guardrail = self._guardrail(experimental_use_latest_role_message_only=True) + session = {"litellm_session_id": "sess-flags-latestrole"} + structured = [ + {"role": "system", "content": "sys"}, + {"role": "user", "content": "q1"}, + ] + with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api: + mock_api.return_value = {"action": "NONE", "output": [], "outputs": []} + await guardrail.apply_guardrail( + inputs={"texts": ["sys", "q1"], "structured_messages": structured}, + request_data=session, + input_type="request", + ) + scanned = [m["content"] for m in mock_api.call_args.kwargs["messages"]] + assert scanned == ["q1"], "latest-role selection must exclude the system prompt" diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_model_armor.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_model_armor.py index 4021f922877..18b5bd92411 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_model_armor.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_model_armor.py @@ -3680,6 +3680,22 @@ async def test_pre_call_hook_skips_chat_traffic_when_configured_for_pre_mcp_call mock_post.assert_not_called() +def test_process_response_with_none_metadata_does_not_crash(): + guardrail = _make_guardrail() + response = {"id": "batch_123", "status": "validating"} + request_data = {"model": "gemini-2.5-flash", "metadata": None} + + result = guardrail._process_response( + response=response, + request_data=request_data, + event_type=GuardrailEventHooks.post_call, + ) + + assert result is response + assert isinstance(request_data["metadata"], dict) + assert "standard_logging_guardrail_information" in request_data["metadata"] + + @pytest.mark.asyncio async def test_moderation_hook_scans_mcp_tool_call_when_configured_for_during_mcp_call(): """A guardrail configured with mode `during_mcp_call` must scan MCP tool calls. diff --git a/tests/test_litellm/proxy/guardrails/test_guardrail_registry.py b/tests/test_litellm/proxy/guardrails/test_guardrail_registry.py index 26feddadf79..14cab50f441 100644 --- a/tests/test_litellm/proxy/guardrails/test_guardrail_registry.py +++ b/tests/test_litellm/proxy/guardrails/test_guardrail_registry.py @@ -1,3 +1,5 @@ +import pytest + from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.proxy.guardrails.guardrail_registry import ( get_guardrail_initializer_from_hooks, @@ -32,6 +34,44 @@ def test_noma_registry_resolution(): assert "noma_v2" in guardrail_initializer_registry +@pytest.mark.parametrize( + "configured, expected", + [(None, True), (False, False), (True, True)], +) +def test_initialize_guardrail_run_in_parallel_preserves_constructor_default(configured, expected): + """ + A guardrail whose constructor sets run_in_parallel=True must keep that default when + the config omits the key; only an explicit config value may override it. The + previous code wrote bool(None)==False on every instance, silently disabling the + opt-in for such guardrails. + """ + from litellm.proxy.guardrails import guardrail_registry as registry_module + + def _initializer(litellm_params, guardrail): + return CustomGuardrail( + guardrail_name=guardrail["guardrail_name"], + event_hook=GuardrailEventHooks.pre_call, + default_on=True, + run_in_parallel=True, + ) + + registry_module.guardrail_initializer_registry["parallel_default_test"] = _initializer + try: + params = {"guardrail": "parallel_default_test", "mode": "pre_call"} + if configured is not None: + params["run_in_parallel"] = configured + + handler = InMemoryGuardrailHandler() + result = handler.initialize_guardrail( + guardrail={"guardrail_name": "cf-parallel-default", "litellm_params": params}, + ) + + stored = handler.guardrail_id_to_custom_guardrail[result["guardrail_id"]] + assert stored.run_in_parallel is expected + finally: + registry_module.guardrail_initializer_registry.pop("parallel_default_test", None) + + def test_update_in_memory_guardrail(): handler = InMemoryGuardrailHandler() handler.guardrail_id_to_custom_guardrail["123"] = CustomGuardrail( diff --git a/tests/test_litellm/proxy/guardrails/test_init_guardrails.py b/tests/test_litellm/proxy/guardrails/test_init_guardrails.py index 83593c20110..71e775842e3 100644 --- a/tests/test_litellm/proxy/guardrails/test_init_guardrails.py +++ b/tests/test_litellm/proxy/guardrails/test_init_guardrails.py @@ -62,3 +62,27 @@ def test_initialize_guardrail_preserves_guardrail_info(): assert result["guardrail_info"] == {"type": "PII", "description": "masks PII"} stored = guardrail_handler.IN_MEMORY_GUARDRAILS[result["guardrail_id"]] assert stored["guardrail_info"] == {"type": "PII", "description": "masks PII"} + + +@pytest.mark.parametrize( + "config_value, expected", + [(True, True), (False, False), (None, False)], +) +def test_initialize_guardrail_sets_run_in_parallel(config_value, expected): + """run_in_parallel from litellm_params must reach the built guardrail instance.""" + litellm_params = { + "guardrail": SupportedGuardrailIntegrations.PRESIDIO.value, + "mode": "pre_call", + "presidio_analyzer_api_base": "https://fakelink.com/v1/presidio/analyze", + "presidio_anonymizer_api_base": "https://fakelink.com/v1/presidio/anonymize", + } + if config_value is not None: + litellm_params["run_in_parallel"] = config_value + + guardrail_handler = InMemoryGuardrailHandler() + result = guardrail_handler.initialize_guardrail( + guardrail={"guardrail_name": "test_parallel_flag", "litellm_params": litellm_params}, + ) + + custom_guardrail = guardrail_handler.guardrail_id_to_custom_guardrail[result["guardrail_id"]] + assert custom_guardrail.run_in_parallel is expected diff --git a/tests/test_litellm/proxy/management_endpoints/scim/test_scim_patch_user.py b/tests/test_litellm/proxy/management_endpoints/scim/test_scim_patch_user.py index f8995a6f4da..c3af4208d37 100644 --- a/tests/test_litellm/proxy/management_endpoints/scim/test_scim_patch_user.py +++ b/tests/test_litellm/proxy/management_endpoints/scim/test_scim_patch_user.py @@ -465,3 +465,58 @@ def test_apply_patch_ops_filtered_path_raises_400_instead_of_junk_metadata(): ) assert exc_info.value.status_code == 400 + + +def test_apply_patch_ops_remove_group_filtered_path_without_value(): + """Okta removes a user from a team with groups[value eq "..."] and no body + value; the team id must be parsed from the filter so the remove takes effect""" + user = LiteLLM_UserTable( + user_id="user-fp", + user_email="fp@example.com", + teams=["team-1", "team-2"], + metadata={}, + ) + patch_ops = SCIMPatchOp( + Operations=[SCIMPatchOperation(op="remove", path='groups[value eq "team-1"]')] + ) + + _, final_team_set = _apply_patch_ops(existing_user=user, patch_ops=patch_ops) + + assert final_team_set == {"team-2"} + + +def test_apply_patch_ops_add_group_filtered_path_without_value(): + """A filtered add path with no body value adds the team id from the filter.""" + user = LiteLLM_UserTable( + user_id="user-fp", + user_email="fp@example.com", + teams=["team-1"], + metadata={}, + ) + patch_ops = SCIMPatchOp( + Operations=[SCIMPatchOperation(op="add", path="groups[value eq 'team-3']")] + ) + + _, final_team_set = _apply_patch_ops(existing_user=user, patch_ops=patch_ops) + + assert final_team_set == {"team-1", "team-3"} + + +def test_apply_patch_ops_replace_groups_empty_value_does_not_use_path_filter(): + """A filtered replace with an explicit empty value must not resurrect the + filter id; the team set is replaced with the empty value as given.""" + user = LiteLLM_UserTable( + user_id="user-fp", + user_email="fp@example.com", + teams=["team-1", "team-2"], + metadata={}, + ) + patch_ops = SCIMPatchOp( + Operations=[ + SCIMPatchOperation(op="replace", path='groups[value eq "team-1"]', value=[]) + ] + ) + + _, final_team_set = _apply_patch_ops(existing_user=user, patch_ops=patch_ops) + + assert final_team_set == set() diff --git a/tests/test_litellm/proxy/management_endpoints/scim/test_scim_v2_endpoints.py b/tests/test_litellm/proxy/management_endpoints/scim/test_scim_v2_endpoints.py index 3ff8e2a6886..7bb74285ac6 100644 --- a/tests/test_litellm/proxy/management_endpoints/scim/test_scim_v2_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/scim/test_scim_v2_endpoints.py @@ -1,24 +1,32 @@ +import time from unittest.mock import AsyncMock import pytest from fastapi import HTTPException from litellm.proxy._types import ( + LiteLLM_TeamTable, LiteLLM_UserTable, LitellmUserRoles, + Member, NewUserRequest, NewUserResponse, + ProxyErrorTypes, ProxyException, ) from litellm.proxy.management_endpoints.scim.scim_v2 import ( UserProvisionerHelpers, + _apply_group_patch_updates, _extract_group_member_ids, + _extract_ids_from_path_filter, _handle_team_membership_changes, _process_group_patch_operations, _recompute_scim_member_roles, create_group, create_user, delete_group, + delete_user, + get_groups, get_users, get_service_provider_config, patch_group, @@ -55,9 +63,7 @@ async def test_create_user_existing_user_conflict(mocker): mock_prisma_client = mocker.MagicMock() mock_prisma_client.db = mocker.MagicMock() mock_prisma_client.db.litellm_usertable = mocker.MagicMock() - mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( - return_value={"user_id": "existing-user"} - ) + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value={"user_id": "existing-user"}) # Mock the _get_prisma_client_or_raise_exception to return our mock mocker.patch( @@ -231,9 +237,7 @@ async def test_create_user_ingests_entitlements_and_roles(mocker, monkeypatch): }, {"value": "bare-entitlement"}, ] - assert created_metadata["scim_roles"] == [ - {"value": "engineering-admin", "type": "role"} - ] + assert created_metadata["scim_roles"] == [{"value": "engineering-admin", "type": "role"}] @pytest.mark.asyncio @@ -257,9 +261,7 @@ async def test_create_user_uses_default_internal_user_params_role(mocker, monkey default_params = { "user_role": LitellmUserRoles.PROXY_ADMIN, } - monkeypatch.setattr( - "litellm.default_internal_user_params", default_params, raising=False - ) + monkeypatch.setattr("litellm.default_internal_user_params", default_params, raising=False) mocker.patch( "litellm.proxy.management_endpoints.scim.scim_v2._get_prisma_client_or_raise_exception", @@ -339,10 +341,7 @@ async def test_scim_create_user_respects_default_role_set_via_ui(mocker, monkeyp "BUG: _update_litellm_setting did not update litellm.default_internal_user_params in memory. " "The local variable reassignment (in_memory_var = ...) doesn't propagate back." ) - assert ( - litellm.default_internal_user_params.get("user_role") - == LitellmUserRoles.INTERNAL_USER - ) + assert litellm.default_internal_user_params.get("user_role") == LitellmUserRoles.INTERNAL_USER # Step 3: Create a user via SCIM scim_user = SCIMUser( @@ -438,9 +437,7 @@ async def test_get_users_filters_username_by_exposed_scim_username_for_okta(mock take=10, order={"created_at": "desc"}, ) - mock_prisma_client.db.litellm_usertable.count.assert_awaited_once_with( - where=expected_where - ) + mock_prisma_client.db.litellm_usertable.count.assert_awaited_once_with(where=expected_where) assert response.totalResults == 1 assert response.Resources[0].id == "internal-user-id" @@ -494,9 +491,7 @@ async def test_get_users_filters_email_value_by_user_email(mocker): take=10, order={"created_at": "desc"}, ) - mock_prisma_client.db.litellm_usertable.count.assert_awaited_once_with( - where=expected_where - ) + mock_prisma_client.db.litellm_usertable.count.assert_awaited_once_with(where=expected_where) assert response.totalResults == 1 assert response.Resources[0].id == "internal-user-id" @@ -544,15 +539,12 @@ async def test_handle_existing_user_by_email_no_existing_user(mocker): ) assert result is None - mock_prisma_client.db.litellm_usertable.find_first.assert_called_once_with( - where={"user_email": "test@example.com"} - ) + mock_prisma_client.db.litellm_usertable.find_first.assert_called_once_with(where={"user_email": "test@example.com"}) @pytest.mark.asyncio async def test_handle_existing_user_by_email_existing_user_updated(mocker): - """Should update existing user and return SCIMUser when user with email exists""" - # Mock existing user - create a proper mock object with attributes + """Should rename the existing user, sync team roster, and return SCIMUser""" existing_user = mocker.MagicMock() existing_user.user_id = "old-user-id" existing_user.user_email = "test@example.com" @@ -560,7 +552,6 @@ async def test_handle_existing_user_by_email_existing_user_updated(mocker): existing_user.teams = ["old-team"] existing_user.metadata = {"old": "data"} - # Mock updated user updated_user = { "user_id": "new-user-id", "user_email": "test@example.com", @@ -569,7 +560,6 @@ async def test_handle_existing_user_by_email_existing_user_updated(mocker): "metadata": '{"new": "data"}', } - # Mock SCIM user to be returned mock_scim_user = SCIMUser( schemas=["urn:ietf:params:scim:schemas:core:2.0:User"], id="new-user-id", @@ -581,18 +571,17 @@ async def test_handle_existing_user_by_email_existing_user_updated(mocker): mock_prisma_client = mocker.MagicMock() mock_prisma_client.db = mocker.MagicMock() mock_prisma_client.db.litellm_usertable = mocker.MagicMock() - mock_prisma_client.db.litellm_usertable.find_first = AsyncMock( - return_value=existing_user - ) - mock_prisma_client.db.litellm_usertable.update = AsyncMock( - return_value=updated_user - ) + mock_prisma_client.db.litellm_usertable.find_first = AsyncMock(return_value=existing_user) + mock_prisma_client.db.litellm_usertable.update = AsyncMock(return_value=updated_user) - # Mock the transformation function mock_transform = mocker.patch( "litellm.proxy.management_endpoints.scim.scim_v2.ScimTransformations.transform_litellm_user_to_scim_user", AsyncMock(return_value=mock_scim_user), ) + mock_membership = mocker.patch( + "litellm.proxy.management_endpoints.scim.scim_v2._handle_team_membership_changes", + AsyncMock(), + ) new_user_request = NewUserRequest( user_id="new-user-id", @@ -607,29 +596,276 @@ async def test_handle_existing_user_by_email_existing_user_updated(mocker): prisma_client=mock_prisma_client, new_user_request=new_user_request ) - # Verify the result assert result == mock_scim_user - # Verify database operations - mock_prisma_client.db.litellm_usertable.find_first.assert_called_once_with( - where={"user_email": "test@example.com"} - ) + mock_prisma_client.db.litellm_usertable.find_first.assert_called_once_with(where={"user_email": "test@example.com"}) - mock_prisma_client.db.litellm_usertable.update.assert_called_once_with( - where={"user_id": "old-user-id"}, - data={ - "user_id": "new-user-id", + update_calls = mock_prisma_client.db.litellm_usertable.update.call_args_list + assert len(update_calls) == 2 + assert update_calls[0].kwargs == { + "where": {"user_id": "old-user-id"}, + "data": {"user_id": "new-user-id"}, + } + assert update_calls[1].kwargs == { + "where": {"user_id": "new-user-id"}, + "data": { "user_email": "test@example.com", "user_alias": "New Name", "teams": ["new-team"], "metadata": '{"new": "data"}', }, + } + + mock_membership.assert_awaited_once_with( + user_id="new-user-id", + existing_teams=["old-team"], + new_teams=["new-team"], + raise_on_error=True, ) - # Verify transformation was called mock_transform.assert_called_once_with(updated_user) +@pytest.mark.asyncio +async def test_handle_existing_user_by_email_syncs_roster_and_dedups_teams(mocker): + """Existing-email upsert must add the user to the team roster via the shared + team_member_add path and dedup the teams built from repeated SCIM groups. + + Regression: previously the user's ``teams`` array was raw-written (with + duplicates) and the team roster (members_with_roles / LiteLLM_TeamMembership) + was never touched, so the user appeared in the group on their profile but was + absent from the team directly. + """ + existing_user = mocker.MagicMock() + existing_user.user_id = "same-id" + existing_user.user_email = "member@example.com" + existing_user.user_alias = "Member" + existing_user.teams = [] + existing_user.metadata = {} + + mock_prisma_client = mocker.MagicMock() + mock_prisma_client.db = mocker.MagicMock() + mock_prisma_client.db.litellm_usertable = mocker.MagicMock() + mock_prisma_client.db.litellm_usertable.find_first = AsyncMock(return_value=existing_user) + mock_prisma_client.db.litellm_usertable.update = AsyncMock(return_value={}) + + mocker.patch( + "litellm.proxy.management_endpoints.scim.scim_v2.ScimTransformations.transform_litellm_user_to_scim_user", + AsyncMock(return_value=None), + ) + mock_membership = mocker.patch( + "litellm.proxy.management_endpoints.scim.scim_v2._handle_team_membership_changes", + AsyncMock(), + ) + + new_user_request = NewUserRequest( + user_id="same-id", + user_email="member@example.com", + user_alias="Member", + teams=["team-a", "team-a", "team-b"], + metadata={}, + auto_create_key=False, + ) + + await UserProvisionerHelpers.handle_existing_user_by_email( + prisma_client=mock_prisma_client, new_user_request=new_user_request + ) + + mock_membership.assert_awaited_once_with( + user_id="same-id", + existing_teams=[], + new_teams=["team-a", "team-b"], + raise_on_error=True, + ) + + update_calls = mock_prisma_client.db.litellm_usertable.update.call_args_list + assert len(update_calls) == 1 + assert update_calls[0].kwargs["where"] == {"user_id": "same-id"} + assert update_calls[0].kwargs["data"]["teams"] == ["team-a", "team-b"] + + +@pytest.mark.asyncio +async def test_handle_existing_user_by_email_roster_add_failure_blocks_teams_write(mocker): + """A genuine roster add failure must propagate and must not persist the teams array. + + Regression: the roster sync went through patch_team_membership which swallowed + real team_member_add failures, so the endpoint reported success and wrote a + teams array listing a team the roster never received. The strict path now + surfaces the failure so user.teams and members_with_roles cannot diverge. + """ + existing_user = mocker.MagicMock() + existing_user.user_id = "uid" + existing_user.user_email = "member@example.com" + existing_user.user_alias = "Member" + existing_user.teams = [] + existing_user.metadata = {} + + mock_prisma_client = mocker.MagicMock() + mock_prisma_client.db = mocker.MagicMock() + mock_prisma_client.db.litellm_usertable = mocker.MagicMock() + mock_prisma_client.db.litellm_usertable.find_first = AsyncMock(return_value=existing_user) + mock_prisma_client.db.litellm_usertable.update = AsyncMock(return_value={}) + + mock_team_member_add = mocker.patch( + "litellm.proxy.management_endpoints.scim.scim_v2.team_member_add", + AsyncMock(side_effect=HTTPException(status_code=404, detail={"error": "Team not found"})), + ) + + new_user_request = NewUserRequest( + user_id="uid", + user_email="member@example.com", + user_alias="Member", + teams=["missing-team"], + metadata={}, + auto_create_key=False, + ) + + with pytest.raises(HTTPException): + await UserProvisionerHelpers.handle_existing_user_by_email( + prisma_client=mock_prisma_client, new_user_request=new_user_request + ) + + mock_team_member_add.assert_awaited_once() + assert mock_prisma_client.db.litellm_usertable.update.await_count == 0 + + +@pytest.mark.asyncio +async def test_handle_existing_user_by_email_roster_add_already_member_is_noop(mocker): + """Being already in the team is benign even under the strict path: the upsert + succeeds and the deduped teams array is still persisted.""" + existing_user = mocker.MagicMock() + existing_user.user_id = "uid" + existing_user.user_email = "member@example.com" + existing_user.user_alias = "Member" + existing_user.teams = [] + existing_user.metadata = {} + + mock_prisma_client = mocker.MagicMock() + mock_prisma_client.db = mocker.MagicMock() + mock_prisma_client.db.litellm_usertable = mocker.MagicMock() + mock_prisma_client.db.litellm_usertable.find_first = AsyncMock(return_value=existing_user) + mock_prisma_client.db.litellm_usertable.update = AsyncMock(return_value={}) + + mock_team_member_add = mocker.patch( + "litellm.proxy.management_endpoints.scim.scim_v2.team_member_add", + AsyncMock( + side_effect=ProxyException( + message="already in team", + type=ProxyErrorTypes.team_member_already_in_team.value, + param=None, + code=400, + ) + ), + ) + mocker.patch( + "litellm.proxy.management_endpoints.scim.scim_v2.ScimTransformations.transform_litellm_user_to_scim_user", + AsyncMock(return_value=None), + ) + + new_user_request = NewUserRequest( + user_id="uid", + user_email="member@example.com", + user_alias="Member", + teams=["team-x"], + metadata={}, + auto_create_key=False, + ) + + await UserProvisionerHelpers.handle_existing_user_by_email( + prisma_client=mock_prisma_client, new_user_request=new_user_request + ) + + mock_team_member_add.assert_awaited_once() + update_calls = mock_prisma_client.db.litellm_usertable.update.call_args_list + assert len(update_calls) == 1 + assert update_calls[0].kwargs["data"]["teams"] == ["team-x"] + + +@pytest.mark.asyncio +async def test_handle_existing_user_by_email_roster_remove_failure_blocks_teams_write(mocker): + """A genuine roster removal failure must propagate and must not persist the teams array, + symmetrically with add failures, so user.teams cannot drop a team the roster still holds.""" + existing_user = mocker.MagicMock() + existing_user.user_id = "uid" + existing_user.user_email = "member@example.com" + existing_user.user_alias = "Member" + existing_user.teams = ["old-team"] + existing_user.metadata = {} + + mock_prisma_client = mocker.MagicMock() + mock_prisma_client.db = mocker.MagicMock() + mock_prisma_client.db.litellm_usertable = mocker.MagicMock() + mock_prisma_client.db.litellm_usertable.find_first = AsyncMock(return_value=existing_user) + mock_prisma_client.db.litellm_usertable.update = AsyncMock(return_value={}) + + mock_team_member_delete = mocker.patch( + "litellm.proxy.management_endpoints.scim.scim_v2.team_member_delete", + AsyncMock(side_effect=HTTPException(status_code=500, detail={"error": "No db connected"})), + ) + + new_user_request = NewUserRequest( + user_id="uid", + user_email="member@example.com", + user_alias="Member", + teams=[], + metadata={}, + auto_create_key=False, + ) + + with pytest.raises(HTTPException): + await UserProvisionerHelpers.handle_existing_user_by_email( + prisma_client=mock_prisma_client, new_user_request=new_user_request + ) + + mock_team_member_delete.assert_awaited_once() + assert mock_prisma_client.db.litellm_usertable.update.await_count == 0 + + +@pytest.mark.asyncio +async def test_handle_existing_user_by_email_roster_remove_already_absent_is_noop(mocker): + """A user already absent from the team is the idempotent removal no-op even under the + strict path: the upsert succeeds and the deduped teams array is still persisted.""" + existing_user = mocker.MagicMock() + existing_user.user_id = "uid" + existing_user.user_email = "member@example.com" + existing_user.user_alias = "Member" + existing_user.teams = ["old-team"] + existing_user.metadata = {} + + mock_prisma_client = mocker.MagicMock() + mock_prisma_client.db = mocker.MagicMock() + mock_prisma_client.db.litellm_usertable = mocker.MagicMock() + mock_prisma_client.db.litellm_usertable.find_first = AsyncMock(return_value=existing_user) + mock_prisma_client.db.litellm_usertable.update = AsyncMock(return_value={}) + + mock_team_member_delete = mocker.patch( + "litellm.proxy.management_endpoints.scim.scim_v2.team_member_delete", + AsyncMock(side_effect=HTTPException(status_code=400, detail={"error": "User not found in team"})), + ) + mocker.patch( + "litellm.proxy.management_endpoints.scim.scim_v2.ScimTransformations.transform_litellm_user_to_scim_user", + AsyncMock(return_value=None), + ) + + new_user_request = NewUserRequest( + user_id="uid", + user_email="member@example.com", + user_alias="Member", + teams=[], + metadata={}, + auto_create_key=False, + ) + + await UserProvisionerHelpers.handle_existing_user_by_email( + prisma_client=mock_prisma_client, new_user_request=new_user_request + ) + + mock_team_member_delete.assert_awaited_once() + update_calls = mock_prisma_client.db.litellm_usertable.update.call_args_list + assert len(update_calls) == 1 + assert update_calls[0].kwargs["data"]["teams"] == [] + + @pytest.mark.asyncio async def test_handle_team_membership_changes_no_changes(mocker): """Should not call patch_team_membership when existing teams equal new teams""" @@ -762,9 +998,7 @@ async def test_update_user_success(mocker): mock_prisma_client = mocker.MagicMock() mock_prisma_client.db = mocker.MagicMock() mock_prisma_client.db.litellm_usertable = mocker.MagicMock() - mock_prisma_client.db.litellm_usertable.update = AsyncMock( - return_value=updated_user - ) + mock_prisma_client.db.litellm_usertable.update = AsyncMock(return_value=updated_user) # Mock dependencies mocker.patch( @@ -815,11 +1049,7 @@ async def test_update_user_not_found(mocker): ) mocker.patch( "litellm.proxy.management_endpoints.scim.scim_v2._check_user_exists", - AsyncMock( - side_effect=HTTPException( - status_code=404, detail={"error": "User not found"} - ) - ), + AsyncMock(side_effect=HTTPException(status_code=404, detail={"error": "User not found"})), ) # Should raise ProxyException (which wraps the HTTPException) @@ -864,9 +1094,7 @@ async def test_patch_user_success(mocker): mock_prisma_client = mocker.MagicMock() mock_prisma_client.db = mocker.MagicMock() mock_prisma_client.db.litellm_usertable = mocker.MagicMock() - mock_prisma_client.db.litellm_usertable.update = AsyncMock( - return_value=updated_user - ) + mock_prisma_client.db.litellm_usertable.update = AsyncMock(return_value=updated_user) # Mock dependencies mocker.patch( @@ -903,9 +1131,7 @@ async def test_patch_user_not_found(mocker): """Should raise 404 when user doesn't exist for patch""" patch_ops = SCIMPatchOp( schemas=["urn:ietf:params:scim:api:messages:2.0:PatchOp"], - Operations=[ - SCIMPatchOperation(op="replace", path="displayName", value="New Name") - ], + Operations=[SCIMPatchOperation(op="replace", path="displayName", value="New Name")], ) # Mock dependencies to raise HTTPException for user not found @@ -915,11 +1141,7 @@ async def test_patch_user_not_found(mocker): ) mocker.patch( "litellm.proxy.management_endpoints.scim.scim_v2._check_user_exists", - AsyncMock( - side_effect=HTTPException( - status_code=404, detail={"error": "User not found"} - ) - ), + AsyncMock(side_effect=HTTPException(status_code=404, detail={"error": "User not found"})), ) # Should raise ProxyException (which wraps the HTTPException) @@ -939,9 +1161,7 @@ async def test_get_service_provider_config(mocker): # Verify it returns the correct response assert isinstance(result, SCIMServiceProviderConfig) - assert result.schemas == [ - "urn:ietf:params:scim:schemas:core:2.0:ServiceProviderConfig" - ] + assert result.schemas == ["urn:ietf:params:scim:schemas:core:2.0:ServiceProviderConfig"] assert result.patch.supported is True assert result.bulk.supported is False assert result.meta is not None @@ -993,21 +1213,15 @@ async def test_update_group_metadata_serialization_issue(mocker): mock_prisma_client.db.litellm_usertable = mocker.MagicMock() # Mock team operations - mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock( - return_value=mock_existing_team - ) - mock_prisma_client.db.litellm_teamtable.update = AsyncMock( - return_value=mock_updated_team - ) + mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=mock_existing_team) + mock_prisma_client.db.litellm_teamtable.update = AsyncMock(return_value=mock_updated_team) # Mock user operations mock_user = mocker.MagicMock() mock_user.user_id = "user1" mock_user.user_email = "user1@example.com" # Add proper string value for user_email mock_user.teams = [group_id] - mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( - return_value=mock_user - ) + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=mock_user) mock_prisma_client.db.litellm_usertable.update = AsyncMock(return_value=mock_user) # Mock the _get_prisma_client_or_raise_exception to return our mock @@ -1044,9 +1258,7 @@ async def test_update_group_metadata_serialization_issue(mocker): metadata = update_data["metadata"] # The fix should ensure metadata is serialized as a JSON string - assert isinstance( - metadata, str - ), f"metadata should be a JSON string, but got {type(metadata)}" + assert isinstance(metadata, str), f"metadata should be a JSON string, but got {type(metadata)}" # Verify we can parse it back to verify it contains the expected data import json @@ -1103,9 +1315,7 @@ async def test_team_membership_management(mocker): # Check calls for adding members add_calls = [ - call - for call in mock_patch_team_membership.call_args_list - if call[1]["teams_ids_to_add_user_to"] == [group_id] + call for call in mock_patch_team_membership.call_args_list if call[1]["teams_ids_to_add_user_to"] == [group_id] ] assert len(add_calls) == 2 # user3 and user4 @@ -1131,9 +1341,7 @@ async def test_team_membership_management(mocker): # Each call should either add OR remove, not both add_teams = call[1]["teams_ids_to_add_user_to"] remove_teams = call[1]["teams_ids_to_remove_user_from"] - assert (len(add_teams) > 0) != ( - len(remove_teams) > 0 - ) # XOR - one should be empty + assert (len(add_teams) > 0) != (len(remove_teams) > 0) # XOR - one should be empty @pytest.mark.asyncio @@ -1184,9 +1392,7 @@ async def test_update_group_e2e(mocker): mock_prisma_client.db.litellm_usertable = mocker.MagicMock() # Mock database operations - mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock( - return_value=existing_team - ) + mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=existing_team) # Mock the updated team that gets returned from database updated_team = LiteLLM_TeamTable( @@ -1203,16 +1409,12 @@ async def test_update_group_e2e(mocker): "scim_data": scim_group_update.model_dump(), }, ) - mock_prisma_client.db.litellm_teamtable.update = AsyncMock( - return_value=updated_team - ) + mock_prisma_client.db.litellm_teamtable.update = AsyncMock(return_value=updated_team) # Mock user validation (all users exist) mock_user = mocker.MagicMock() mock_user.user_id = "test-user" - mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( - return_value=mock_user - ) + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=mock_user) # Mock dependencies mocker.patch( @@ -1265,29 +1467,19 @@ async def test_update_group_e2e(mocker): assert metadata["scim_data"]["displayName"] == "Updated Team Name" # Verify team membership changes were handled correctly - assert ( - mock_patch_team_membership.call_count == 3 - ) # Remove user1, add user3, add user4 + assert mock_patch_team_membership.call_count == 3 # Remove user1, add user3, add user4 # Check membership changes call_args_list = mock_patch_team_membership.call_args_list # Find remove operation (user1) - remove_calls = [ - call - for call in call_args_list - if call[1]["teams_ids_to_remove_user_from"] == [group_id] - ] + remove_calls = [call for call in call_args_list if call[1]["teams_ids_to_remove_user_from"] == [group_id]] assert len(remove_calls) == 1 assert remove_calls[0][1]["user_id"] == "user1" assert remove_calls[0][1]["teams_ids_to_add_user_to"] == [] # Find add operations (user3, user4) - add_calls = [ - call - for call in call_args_list - if call[1]["teams_ids_to_add_user_to"] == [group_id] - ] + add_calls = [call for call in call_args_list if call[1]["teams_ids_to_add_user_to"] == [group_id]] assert len(add_calls) == 2 add_user_ids = {call[1]["user_id"] for call in add_calls} assert add_user_ids == {"user3", "user4"} @@ -1302,9 +1494,7 @@ async def test_update_group_e2e(mocker): assert len(result.members) == 3 # Verify SCIM transformation was called with updated team - ScimTransformations.transform_litellm_team_to_scim_group.assert_called_once_with( - updated_team - ) + ScimTransformations.transform_litellm_team_to_scim_group.assert_called_once_with(updated_team) @pytest.mark.asyncio @@ -1330,15 +1520,9 @@ async def test_create_group_with_nonexistent_users_rejects(mocker, monkeypatch): id=group_id, displayName="Test Group", members=[ - SCIMMember( - value="existing-user", display="Existing User" - ), # This user exists - SCIMMember( - value="new-user-1", display="New User 1" - ), # This user doesn't exist - SCIMMember( - value="new-user-2", display="New User 2" - ), # This user doesn't exist + SCIMMember(value="existing-user", display="Existing User"), # This user exists + SCIMMember(value="new-user-1", display="New User 1"), # This user doesn't exist + SCIMMember(value="new-user-2", display="New User 2"), # This user doesn't exist ], ) @@ -1364,9 +1548,7 @@ async def test_create_group_with_nonexistent_users_rejects(mocker, monkeypatch): return mock_user return None # new-user-1 and new-user-2 don't exist - mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( - side_effect=mock_user_lookup - ) + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(side_effect=mock_user_lookup) # Mock dependencies mocker.patch( @@ -1381,9 +1563,7 @@ async def test_create_group_with_nonexistent_users_rejects(mocker, monkeypatch): # Verify it's a 400 Bad Request assert int(exc_info.value.code) == 400 assert "does not exist" in str(exc_info.value.message) - assert "new-user-1" in str(exc_info.value.message) or "new-user-2" in str( - exc_info.value.message - ) + assert "new-user-1" in str(exc_info.value.message) or "new-user-2" in str(exc_info.value.message) @pytest.mark.asyncio @@ -1418,15 +1598,9 @@ async def test_update_group_with_nonexistent_users_rejects(mocker, monkeypatch): id=group_id, displayName="Updated Group Name", members=[ - SCIMMember( - value="existing-user", display="Existing User" - ), # This user exists - SCIMMember( - value="new-user-3", display="New User 3" - ), # This user doesn't exist - SCIMMember( - value="new-user-4", display="New User 4" - ), # This user doesn't exist + SCIMMember(value="existing-user", display="Existing User"), # This user exists + SCIMMember(value="new-user-3", display="New User 3"), # This user doesn't exist + SCIMMember(value="new-user-4", display="New User 4"), # This user doesn't exist ], ) @@ -1437,18 +1611,14 @@ async def test_update_group_with_nonexistent_users_rejects(mocker, monkeypatch): mock_prisma_client.db.litellm_usertable = mocker.MagicMock() # Mock team operations - mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock( - return_value=mock_existing_team - ) + mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=mock_existing_team) # Mock updated team response mock_updated_team = mocker.MagicMock() mock_updated_team.team_id = group_id mock_updated_team.team_alias = "Updated Group Name" mock_updated_team.members = ["existing-user", "new-user-3", "new-user-4"] - mock_prisma_client.db.litellm_teamtable.update = AsyncMock( - return_value=mock_updated_team - ) + mock_prisma_client.db.litellm_teamtable.update = AsyncMock(return_value=mock_updated_team) # Mock user lookup - only existing-user exists def mock_user_lookup(where): @@ -1459,9 +1629,7 @@ async def test_update_group_with_nonexistent_users_rejects(mocker, monkeypatch): return mock_user return None # new-user-3 and new-user-4 don't exist - mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( - side_effect=mock_user_lookup - ) + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(side_effect=mock_user_lookup) # Mock dependencies mocker.patch( @@ -1481,15 +1649,11 @@ async def test_update_group_with_nonexistent_users_rejects(mocker, monkeypatch): # Verify it's a 400 Bad Request assert int(exc_info.value.code) == 400 assert "does not exist" in str(exc_info.value.message) - assert "new-user-3" in str(exc_info.value.message) or "new-user-4" in str( - exc_info.value.message - ) + assert "new-user-3" in str(exc_info.value.message) or "new-user-4" in str(exc_info.value.message) @pytest.mark.asyncio -async def test_create_group_with_nonexistent_users_creates_when_flag_true( - mocker, monkeypatch -): +async def test_create_group_with_nonexistent_users_creates_when_flag_true(mocker, monkeypatch): """ Test that creating a group with non-existent users creates them when scim_upsert_user is True. This preserves backward compatible behavior. @@ -1510,15 +1674,9 @@ async def test_create_group_with_nonexistent_users_creates_when_flag_true( id=group_id, displayName="Test Group", members=[ - SCIMMember( - value="existing-user", display="Existing User" - ), # This user exists - SCIMMember( - value="new-user-1", display="New User 1" - ), # This user doesn't exist - should be created - SCIMMember( - value="new-user-2", display="New User 2" - ), # This user doesn't exist - should be created + SCIMMember(value="existing-user", display="Existing User"), # This user exists + SCIMMember(value="new-user-1", display="New User 1"), # This user doesn't exist - should be created + SCIMMember(value="new-user-2", display="New User 2"), # This user doesn't exist - should be created ], ) @@ -1540,9 +1698,7 @@ async def test_create_group_with_nonexistent_users_creates_when_flag_true( return mock_user return None # new-user-1 and new-user-2 don't exist - mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( - side_effect=mock_user_lookup - ) + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(side_effect=mock_user_lookup) # Mock user creation created_user_1 = NewUserResponse(user_id="new-user-1", key="test-key-1") @@ -1591,9 +1747,7 @@ async def test_create_group_with_nonexistent_users_creates_when_flag_true( @pytest.mark.asyncio -async def test_extract_group_member_ids_with_flag_true_creates_users( - mocker, monkeypatch -): +async def test_extract_group_member_ids_with_flag_true_creates_users(mocker, monkeypatch): """ Test that _extract_group_member_ids creates users when scim_upsert_user is True. """ @@ -1612,12 +1766,8 @@ async def test_extract_group_member_ids_with_flag_true_creates_users( id="test-group", displayName="Test Group", members=[ - SCIMMember( - value="existing-user", display="Existing User" - ), # This user exists - SCIMMember( - value="new-user-1", display="New User 1" - ), # This user doesn't exist - should be created + SCIMMember(value="existing-user", display="Existing User"), # This user exists + SCIMMember(value="new-user-1", display="New User 1"), # This user doesn't exist - should be created ], ) @@ -1635,9 +1785,7 @@ async def test_extract_group_member_ids_with_flag_true_creates_users( return mock_user return None # new-user-1 doesn't exist - mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( - side_effect=mock_user_lookup - ) + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(side_effect=mock_user_lookup) # Mock user creation created_user = NewUserResponse(user_id="new-user-1", key="test-key-1") @@ -1662,9 +1810,7 @@ async def test_extract_group_member_ids_with_flag_true_creates_users( assert len(result.created_users) == 1 # Verify user was created - mock_create_user.assert_called_once_with( - user_id="new-user-1", created_via="scim_group_membership" - ) + mock_create_user.assert_called_once_with(user_id="new-user-1", created_via="scim_group_membership") @pytest.mark.asyncio @@ -1687,12 +1833,8 @@ async def test_extract_group_member_ids_with_flag_false_rejects(mocker, monkeypa id="test-group", displayName="Test Group", members=[ - SCIMMember( - value="existing-user", display="Existing User" - ), # This user exists - SCIMMember( - value="new-user-1", display="New User 1" - ), # This user doesn't exist - should be rejected + SCIMMember(value="existing-user", display="Existing User"), # This user exists + SCIMMember(value="new-user-1", display="New User 1"), # This user doesn't exist - should be rejected ], ) @@ -1710,9 +1852,7 @@ async def test_extract_group_member_ids_with_flag_false_rejects(mocker, monkeypa return mock_user return None # new-user-1 doesn't exist - mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( - side_effect=mock_user_lookup - ) + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(side_effect=mock_user_lookup) # Mock dependencies mocker.patch( @@ -1731,9 +1871,7 @@ async def test_extract_group_member_ids_with_flag_false_rejects(mocker, monkeypa @pytest.mark.asyncio -async def test_process_group_patch_operations_with_flag_true_creates_users( - mocker, monkeypatch -): +async def test_process_group_patch_operations_with_flag_true_creates_users(mocker, monkeypatch): """ Test that _process_group_patch_operations creates users when scim_upsert_user is True. """ @@ -1749,11 +1887,7 @@ async def test_process_group_patch_operations_with_flag_true_creates_users( # Test data patch_ops = SCIMPatchOp( schemas=["urn:ietf:params:scim:api:messages:2.0:PatchOp"], - Operations=[ - SCIMPatchOperation( - op="add", path="members", value=[{"value": "new-user-1"}] - ) - ], + Operations=[SCIMPatchOperation(op="add", path="members", value=[{"value": "new-user-1"}])], ) # Mock existing team @@ -1777,7 +1911,7 @@ async def test_process_group_patch_operations_with_flag_true_creates_users( ) # Execute the function - update_data, final_members = await _process_group_patch_operations( + update_data, final_members, _ = await _process_group_patch_operations( patch_ops=patch_ops, existing_team=mock_existing_team, prisma_client=mock_prisma_client, @@ -1787,15 +1921,11 @@ async def test_process_group_patch_operations_with_flag_true_creates_users( assert "new-user-1" in final_members # Verify user was created - mock_create_user.assert_called_once_with( - user_id="new-user-1", created_via="scim_group_patch" - ) + mock_create_user.assert_called_once_with(user_id="new-user-1", created_via="scim_group_patch") @pytest.mark.asyncio -async def test_process_group_patch_operations_with_flag_false_rejects( - mocker, monkeypatch -): +async def test_process_group_patch_operations_with_flag_false_rejects(mocker, monkeypatch): """ Test that _process_group_patch_operations rejects non-existent users when scim_upsert_user is False. """ @@ -1811,11 +1941,7 @@ async def test_process_group_patch_operations_with_flag_false_rejects( # Test data patch_ops = SCIMPatchOp( schemas=["urn:ietf:params:scim:api:messages:2.0:PatchOp"], - Operations=[ - SCIMPatchOperation( - op="add", path="members", value=[{"value": "new-user-1"}] - ) - ], + Operations=[SCIMPatchOperation(op="add", path="members", value=[{"value": "new-user-1"}])], ) # Mock existing team @@ -1890,9 +2016,7 @@ async def test_create_user_grants_admin_when_in_scim_admin_group(mocker, monkeyp @pytest.mark.asyncio -async def test_create_user_keeps_default_when_not_in_scim_admin_group( - mocker, monkeypatch -): +async def test_create_user_keeps_default_when_not_in_scim_admin_group(mocker, monkeypatch): """When scim_admin_group is configured but the user's groups don't include it, the user keeps the non-admin default role.""" from litellm.proxy.proxy_server import proxy_config @@ -1936,9 +2060,7 @@ async def test_create_user_keeps_default_when_not_in_scim_admin_group( @pytest.mark.asyncio -async def test_update_user_demotes_admin_when_removed_from_scim_admin_group( - mocker, monkeypatch -): +async def test_update_user_demotes_admin_when_removed_from_scim_admin_group(mocker, monkeypatch): """Core demotion test: a PUT whose new groups no longer include the configured admin group must re-evaluate the role and write the non-admin default, so an admin removed from the IdP group is demoted without re-login.""" @@ -1972,9 +2094,7 @@ async def test_update_user_demotes_admin_when_removed_from_scim_admin_group( mock_prisma_client = mocker.MagicMock() mock_prisma_client.db = mocker.MagicMock() mock_prisma_client.db.litellm_usertable = mocker.MagicMock() - mock_prisma_client.db.litellm_usertable.update = AsyncMock( - return_value=updated_user - ) + mock_prisma_client.db.litellm_usertable.update = AsyncMock(return_value=updated_user) mocker.patch( "litellm.proxy.management_endpoints.scim.scim_v2._get_prisma_client_or_raise_exception", @@ -2000,9 +2120,7 @@ async def test_update_user_demotes_admin_when_removed_from_scim_admin_group( @pytest.mark.asyncio -async def test_update_user_does_not_force_role_when_scim_admin_group_unset( - mocker, monkeypatch -): +async def test_update_user_does_not_force_role_when_scim_admin_group_unset(mocker, monkeypatch): """When scim_admin_group is unset, PUT must not touch user_role (current behavior preserved).""" from litellm.proxy.proxy_server import proxy_config @@ -2035,9 +2153,7 @@ async def test_update_user_does_not_force_role_when_scim_admin_group_unset( mock_prisma_client = mocker.MagicMock() mock_prisma_client.db = mocker.MagicMock() mock_prisma_client.db.litellm_usertable = mocker.MagicMock() - mock_prisma_client.db.litellm_usertable.update = AsyncMock( - return_value=updated_user - ) + mock_prisma_client.db.litellm_usertable.update = AsyncMock(return_value=updated_user) mocker.patch( "litellm.proxy.management_endpoints.scim.scim_v2._get_prisma_client_or_raise_exception", @@ -2063,9 +2179,7 @@ async def test_update_user_does_not_force_role_when_scim_admin_group_unset( @pytest.mark.asyncio -async def test_update_user_demotes_when_default_params_lack_user_role( - mocker, monkeypatch -): +async def test_update_user_demotes_when_default_params_lack_user_role(mocker, monkeypatch): """Regression: default_internal_user_params set without a user_role key must still resolve to the non-admin default on demotion, not silently skip and leave the user PROXY_ADMIN.""" @@ -2075,9 +2189,7 @@ async def test_update_user_demotes_when_default_params_lack_user_role( return {"litellm_settings": {"scim_admin_group": "litellm-admins"}} monkeypatch.setattr(proxy_config, "get_config", mock_get_config) - monkeypatch.setattr( - "litellm.default_internal_user_params", {"max_budget": 10}, raising=False - ) + monkeypatch.setattr("litellm.default_internal_user_params", {"max_budget": 10}, raising=False) existing_user = mocker.MagicMock() existing_user.teams = ["litellm-admins"] @@ -2101,9 +2213,7 @@ async def test_update_user_demotes_when_default_params_lack_user_role( mock_prisma_client = mocker.MagicMock() mock_prisma_client.db = mocker.MagicMock() mock_prisma_client.db.litellm_usertable = mocker.MagicMock() - mock_prisma_client.db.litellm_usertable.update = AsyncMock( - return_value=updated_user - ) + mock_prisma_client.db.litellm_usertable.update = AsyncMock(return_value=updated_user) mocker.patch( "litellm.proxy.management_endpoints.scim.scim_v2._get_prisma_client_or_raise_exception", @@ -2129,9 +2239,7 @@ async def test_update_user_demotes_when_default_params_lack_user_role( @pytest.mark.asyncio -async def test_patch_user_demotes_admin_when_removed_from_scim_admin_group( - mocker, monkeypatch -): +async def test_patch_user_demotes_admin_when_removed_from_scim_admin_group(mocker, monkeypatch): """PATCH that drops the admin team from the resulting team set must write the non-admin default, mirroring the PUT demotion path.""" from litellm.proxy.proxy_server import proxy_config @@ -2148,11 +2256,7 @@ async def test_patch_user_demotes_admin_when_removed_from_scim_admin_group( patch_ops = SCIMPatchOp( schemas=["urn:ietf:params:scim:api:messages:2.0:PatchOp"], - Operations=[ - SCIMPatchOperation( - op="replace", path="groups", value=[{"value": "engineering"}] - ) - ], + Operations=[SCIMPatchOperation(op="replace", path="groups", value=[{"value": "engineering"}])], ) updated_user = { @@ -2168,13 +2272,9 @@ async def test_patch_user_demotes_admin_when_removed_from_scim_admin_group( mock_prisma_client = mocker.MagicMock() mock_prisma_client.db = mocker.MagicMock() mock_prisma_client.db.litellm_usertable = mocker.MagicMock() - mock_prisma_client.db.litellm_usertable.update = AsyncMock( - return_value=updated_user - ) + mock_prisma_client.db.litellm_usertable.update = AsyncMock(return_value=updated_user) mock_prisma_client.db.litellm_teamtable = mocker.MagicMock() - mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock( - return_value=engineering_team - ) + mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=engineering_team) mocker.patch( "litellm.proxy.management_endpoints.scim.scim_v2._get_prisma_client_or_raise_exception", @@ -2223,11 +2323,7 @@ async def test_patch_user_grants_admin_by_team_display_name(mocker, monkeypatch) patch_ops = SCIMPatchOp( schemas=["urn:ietf:params:scim:api:messages:2.0:PatchOp"], - Operations=[ - SCIMPatchOperation( - op="replace", path="groups", value=[{"value": "team-abc-123"}] - ) - ], + Operations=[SCIMPatchOperation(op="replace", path="groups", value=[{"value": "team-abc-123"}])], ) updated_user = { @@ -2243,13 +2339,9 @@ async def test_patch_user_grants_admin_by_team_display_name(mocker, monkeypatch) mock_prisma_client = mocker.MagicMock() mock_prisma_client.db = mocker.MagicMock() mock_prisma_client.db.litellm_usertable = mocker.MagicMock() - mock_prisma_client.db.litellm_usertable.update = AsyncMock( - return_value=updated_user - ) + mock_prisma_client.db.litellm_usertable.update = AsyncMock(return_value=updated_user) mock_prisma_client.db.litellm_teamtable = mocker.MagicMock() - mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock( - return_value=admin_team - ) + mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=admin_team) mocker.patch( "litellm.proxy.management_endpoints.scim.scim_v2._get_prisma_client_or_raise_exception", @@ -2302,9 +2394,7 @@ def _scim_admin_prisma(mocker, *, user_teams): @pytest.mark.asyncio -async def test_recompute_scim_member_roles_demotes_when_not_in_admin_group( - mocker, monkeypatch -): +async def test_recompute_scim_member_roles_demotes_when_not_in_admin_group(mocker, monkeypatch): """The shared recompute helper writes the non-admin default for a member whose resulting teams no longer include the configured admin group.""" from litellm.proxy.proxy_server import proxy_config @@ -2324,9 +2414,7 @@ async def test_recompute_scim_member_roles_demotes_when_not_in_admin_group( @pytest.mark.asyncio -async def test_recompute_scim_member_roles_grants_when_in_admin_group( - mocker, monkeypatch -): +async def test_recompute_scim_member_roles_grants_when_in_admin_group(mocker, monkeypatch): """The shared recompute helper grants PROXY_ADMIN when a member's resulting teams include the configured admin group.""" from litellm.proxy.proxy_server import proxy_config @@ -2346,9 +2434,7 @@ async def test_recompute_scim_member_roles_grants_when_in_admin_group( @pytest.mark.asyncio -async def test_recompute_scim_member_roles_noop_when_admin_group_unset( - mocker, monkeypatch -): +async def test_recompute_scim_member_roles_noop_when_admin_group_unset(mocker, monkeypatch): """With scim_admin_group unset the recompute helper must not touch any role, preserving current behavior for SCIM group writes.""" from litellm.proxy.proxy_server import proxy_config @@ -2395,16 +2481,10 @@ async def test_update_group_recomputes_roles_for_changed_members(mocker): mock_prisma_client = mocker.MagicMock() mock_prisma_client.db = mocker.MagicMock() mock_prisma_client.db.litellm_teamtable = mocker.MagicMock() - mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock( - return_value=existing_team - ) - mock_prisma_client.db.litellm_teamtable.update = AsyncMock( - return_value=existing_team - ) + mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=existing_team) + mock_prisma_client.db.litellm_teamtable.update = AsyncMock(return_value=existing_team) mock_prisma_client.db.litellm_usertable = mocker.MagicMock() - mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( - return_value=mocker.MagicMock() - ) + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=mocker.MagicMock()) mocker.patch( "litellm.proxy.management_endpoints.scim.scim_v2._get_prisma_client_or_raise_exception", @@ -2452,24 +2532,16 @@ async def test_patch_group_recomputes_roles_for_changed_members(mocker): ) patch_ops = SCIMPatchOp( schemas=["urn:ietf:params:scim:api:messages:2.0:PatchOp"], - Operations=[ - SCIMPatchOperation(op="remove", path="members", value=[{"value": "user1"}]) - ], + Operations=[SCIMPatchOperation(op="remove", path="members", value=[{"value": "user1"}])], ) mock_prisma_client = mocker.MagicMock() mock_prisma_client.db = mocker.MagicMock() mock_prisma_client.db.litellm_teamtable = mocker.MagicMock() - mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock( - return_value=existing_team - ) - mock_prisma_client.db.litellm_teamtable.update = AsyncMock( - return_value=existing_team - ) + mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=existing_team) + mock_prisma_client.db.litellm_teamtable.update = AsyncMock(return_value=existing_team) mock_prisma_client.db.litellm_usertable = mocker.MagicMock() - mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( - return_value=mocker.MagicMock() - ) + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=mocker.MagicMock()) mocker.patch( "litellm.proxy.management_endpoints.scim.scim_v2._get_prisma_client_or_raise_exception", @@ -2519,9 +2591,7 @@ async def test_delete_group_recomputes_roles_for_members(mocker): mock_prisma_client = mocker.MagicMock() mock_prisma_client.db = mocker.MagicMock() mock_prisma_client.db.litellm_teamtable = mocker.MagicMock() - mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock( - return_value=existing_team - ) + mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=existing_team) mock_prisma_client.db.litellm_teamtable.delete = AsyncMock() mock_prisma_client.db.litellm_usertable = mocker.MagicMock() mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=member) @@ -2552,16 +2622,16 @@ async def test_handle_existing_user_by_email_applies_role_when_admin_group_set(m mock_prisma_client = mocker.MagicMock() mock_prisma_client.db = mocker.MagicMock() mock_prisma_client.db.litellm_usertable = mocker.MagicMock() - mock_prisma_client.db.litellm_usertable.find_first = AsyncMock( - return_value=existing_user - ) - mock_prisma_client.db.litellm_usertable.update = AsyncMock( - return_value={"user_id": "new-user-id"} - ) + mock_prisma_client.db.litellm_usertable.find_first = AsyncMock(return_value=existing_user) + mock_prisma_client.db.litellm_usertable.update = AsyncMock(return_value={"user_id": "new-user-id"}) mocker.patch( "litellm.proxy.management_endpoints.scim.scim_v2.ScimTransformations.transform_litellm_user_to_scim_user", AsyncMock(return_value=mocker.MagicMock()), ) + mocker.patch( + "litellm.proxy.management_endpoints.scim.scim_v2._handle_team_membership_changes", + AsyncMock(), + ) new_user_request = NewUserRequest( user_id="new-user-id", @@ -2592,16 +2662,16 @@ async def test_handle_existing_user_by_email_leaves_role_when_admin_group_unset( mock_prisma_client = mocker.MagicMock() mock_prisma_client.db = mocker.MagicMock() mock_prisma_client.db.litellm_usertable = mocker.MagicMock() - mock_prisma_client.db.litellm_usertable.find_first = AsyncMock( - return_value=existing_user - ) - mock_prisma_client.db.litellm_usertable.update = AsyncMock( - return_value={"user_id": "new-user-id"} - ) + mock_prisma_client.db.litellm_usertable.find_first = AsyncMock(return_value=existing_user) + mock_prisma_client.db.litellm_usertable.update = AsyncMock(return_value={"user_id": "new-user-id"}) mocker.patch( "litellm.proxy.management_endpoints.scim.scim_v2.ScimTransformations.transform_litellm_user_to_scim_user", AsyncMock(return_value=mocker.MagicMock()), ) + mocker.patch( + "litellm.proxy.management_endpoints.scim.scim_v2._handle_team_membership_changes", + AsyncMock(), + ) new_user_request = NewUserRequest( user_id="new-user-id", @@ -2622,9 +2692,7 @@ async def test_handle_existing_user_by_email_leaves_role_when_admin_group_unset( @pytest.mark.asyncio -async def test_create_user_existing_email_upsert_demotes_when_admin_group_set( - mocker, monkeypatch -): +async def test_create_user_existing_email_upsert_demotes_when_admin_group_set(mocker, monkeypatch): """End-to-end create wiring: a SCIM POST that upserts an existing email while the user is not in the admin group must write the non-admin default, not leave a stale PROXY_ADMIN.""" @@ -2650,12 +2718,8 @@ async def test_create_user_existing_email_upsert_demotes_when_admin_group_set( mock_prisma_client.db = mocker.MagicMock() mock_prisma_client.db.litellm_usertable = mocker.MagicMock() mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None) - mock_prisma_client.db.litellm_usertable.find_first = AsyncMock( - return_value=existing_user - ) - mock_prisma_client.db.litellm_usertable.update = AsyncMock( - return_value={"user_id": "returning-user"} - ) + mock_prisma_client.db.litellm_usertable.find_first = AsyncMock(return_value=existing_user) + mock_prisma_client.db.litellm_usertable.update = AsyncMock(return_value={"user_id": "returning-user"}) mocker.patch( "litellm.proxy.management_endpoints.scim.scim_v2._get_prisma_client_or_raise_exception", @@ -2669,6 +2733,10 @@ async def test_create_user_existing_email_upsert_demotes_when_admin_group_set( "litellm.proxy.management_endpoints.scim.scim_v2.ScimTransformations.transform_litellm_user_to_scim_user", AsyncMock(return_value=scim_user), ) + mocker.patch( + "litellm.proxy.management_endpoints.scim.scim_v2._handle_team_membership_changes", + AsyncMock(), + ) await create_user(user=scim_user) @@ -2698,9 +2766,7 @@ async def test_create_group_recomputes_roles_for_members(mocker): mock_prisma_client.db.litellm_teamtable = mocker.MagicMock() mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=None) mock_prisma_client.db.litellm_usertable = mocker.MagicMock() - mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( - return_value=mocker.MagicMock() - ) + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=mocker.MagicMock()) mocker.patch( "litellm.proxy.management_endpoints.scim.scim_v2._get_prisma_client_or_raise_exception", @@ -2754,16 +2820,10 @@ async def test_update_group_rename_recomputes_retained_members(mocker): mock_prisma_client = mocker.MagicMock() mock_prisma_client.db = mocker.MagicMock() mock_prisma_client.db.litellm_teamtable = mocker.MagicMock() - mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock( - return_value=existing_team - ) - mock_prisma_client.db.litellm_teamtable.update = AsyncMock( - return_value=existing_team - ) + mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=existing_team) + mock_prisma_client.db.litellm_teamtable.update = AsyncMock(return_value=existing_team) mock_prisma_client.db.litellm_usertable = mocker.MagicMock() - mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( - return_value=mocker.MagicMock() - ) + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=mocker.MagicMock()) mocker.patch( "litellm.proxy.management_endpoints.scim.scim_v2._get_prisma_client_or_raise_exception", @@ -2808,24 +2868,16 @@ async def test_patch_group_rename_recomputes_retained_members(mocker): ) patch_ops = SCIMPatchOp( schemas=["urn:ietf:params:scim:api:messages:2.0:PatchOp"], - Operations=[ - SCIMPatchOperation(op="replace", path="displayName", value="Engineering") - ], + Operations=[SCIMPatchOperation(op="replace", path="displayName", value="Engineering")], ) mock_prisma_client = mocker.MagicMock() mock_prisma_client.db = mocker.MagicMock() mock_prisma_client.db.litellm_teamtable = mocker.MagicMock() - mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock( - return_value=existing_team - ) - mock_prisma_client.db.litellm_teamtable.update = AsyncMock( - return_value=existing_team - ) + mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=existing_team) + mock_prisma_client.db.litellm_teamtable.update = AsyncMock(return_value=existing_team) mock_prisma_client.db.litellm_usertable = mocker.MagicMock() - mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( - return_value=mocker.MagicMock() - ) + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=mocker.MagicMock()) mocker.patch( "litellm.proxy.management_endpoints.scim.scim_v2._get_prisma_client_or_raise_exception", @@ -2855,3 +2907,615 @@ async def test_patch_group_rename_recomputes_retained_members(mocker): recompute_mock.assert_awaited_once() assert set(recompute_mock.call_args[0][1]) == {"user1"} + + +@pytest.mark.asyncio +async def test_process_group_patch_operations_add_retains_existing_members( + mocker, monkeypatch +): + """A SCIM group ``add`` operation must not drop members already in the team. + + Team membership lives in members_with_roles; team creation leaves the legacy + ``members`` column empty. Seeding the patch result from that empty column + made an ``add`` recompute the member set from scratch and remove everyone + already in the team. The result set must be seeded from members_with_roles so + existing members survive an add of a new one. + """ + + async def mock_get_config(): + return {"litellm_settings": {"scim_upsert_user": True}} + + from litellm.proxy.proxy_server import proxy_config + + monkeypatch.setattr(proxy_config, "get_config", mock_get_config) + + existing_team = LiteLLM_TeamTable( + team_id="team-1", + team_alias="Team One", + members=[], # legacy column intentionally empty, as real teams leave it + members_with_roles=[Member(user_id="existing-user", role="user")], + ) + patch_ops = SCIMPatchOp( + schemas=["urn:ietf:params:scim:api:messages:2.0:PatchOp"], + Operations=[ + SCIMPatchOperation(op="add", path="members", value=[{"value": "new-user"}]) + ], + ) + + mock_prisma_client = mocker.MagicMock() + mock_prisma_client.db = mocker.MagicMock() + mock_prisma_client.db.litellm_usertable = mocker.MagicMock() + # new-user already exists in the DB + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( + return_value=mocker.MagicMock(user_id="new-user") + ) + + _, final_members, _ = await _process_group_patch_operations( + patch_ops=patch_ops, + existing_team=existing_team, + prisma_client=mock_prisma_client, + ) + + assert final_members == {"existing-user", "new-user"} + + +@pytest.mark.asyncio +async def test_process_group_patch_operations_remove_uses_members_with_roles( + mocker, monkeypatch +): + """A ``remove`` op must diff against members_with_roles, so removing one + member leaves the rest of the team intact rather than emptying it.""" + + async def mock_get_config(): + return {"litellm_settings": {"scim_upsert_user": True}} + + from litellm.proxy.proxy_server import proxy_config + + monkeypatch.setattr(proxy_config, "get_config", mock_get_config) + + existing_team = LiteLLM_TeamTable( + team_id="team-1", + team_alias="Team One", + members=[], + members_with_roles=[ + Member(user_id="keep-user", role="user"), + Member(user_id="drop-user", role="user"), + ], + ) + patch_ops = SCIMPatchOp( + schemas=["urn:ietf:params:scim:api:messages:2.0:PatchOp"], + Operations=[ + SCIMPatchOperation( + op="remove", path="members", value=[{"value": "drop-user"}] + ) + ], + ) + + mock_prisma_client = mocker.MagicMock() + mock_prisma_client.db = mocker.MagicMock() + mock_prisma_client.db.litellm_usertable = mocker.MagicMock() + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( + return_value=mocker.MagicMock(user_id="drop-user") + ) + + _, final_members, _ = await _process_group_patch_operations( + patch_ops=patch_ops, + existing_team=existing_team, + prisma_client=mock_prisma_client, + ) + + assert final_members == {"keep-user"} + + +@pytest.mark.asyncio +async def test_get_groups_reports_members_from_members_with_roles(mocker): + """GET /Groups must report members from members_with_roles (the source of + truth), not the legacy ``members`` column that team creation leaves empty. + Reporting an empty member list makes the IdP repeatedly re-provision.""" + team = LiteLLM_TeamTable( + team_id="team-1", + team_alias="Team One", + members=[], # legacy column empty + members_with_roles=[Member(user_id="member-1", role="user")], + ) + + mock_prisma_client = mocker.MagicMock() + mock_prisma_client.db = mocker.MagicMock() + mock_prisma_client.db.litellm_teamtable = mocker.MagicMock() + mock_prisma_client.db.litellm_teamtable.find_many = AsyncMock(return_value=[team]) + mock_prisma_client.db.litellm_teamtable.count = AsyncMock(return_value=1) + mock_prisma_client.db.litellm_usertable = mocker.MagicMock() + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( + return_value=mocker.MagicMock(user_id="member-1", user_email="member-1@example.com") + ) + + mocker.patch( + "litellm.proxy.management_endpoints.scim.scim_v2._get_prisma_client_or_raise_exception", + AsyncMock(return_value=mock_prisma_client), + ) + + response = await get_groups(startIndex=1, count=10, filter=None) + + assert [m.value for m in response.Resources[0].members] == ["member-1"] + + +@pytest.mark.asyncio +async def test_apply_group_patch_updates_does_not_write_legacy_members(mocker): + """The group PATCH apply must not write the legacy ``members`` column. + + Membership is reconciled onto the source of truth (members_with_roles and + each member's user.teams) separately; writing the legacy column here too + would create a second, unread copy of membership that can drift from the + source of truth, which is the inconsistency this PR removes. + """ + mock_prisma_client = mocker.MagicMock() + mock_prisma_client.db = mocker.MagicMock() + mock_prisma_client.db.litellm_teamtable = mocker.MagicMock() + updated = mocker.MagicMock() + mock_prisma_client.db.litellm_teamtable.update = AsyncMock(return_value=updated) + + result = await _apply_group_patch_updates( + group_id="team-1", + update_data={"team_alias": "Renamed"}, + prisma_client=mock_prisma_client, + ) + + assert result is updated + mock_prisma_client.db.litellm_teamtable.update.assert_awaited_once() + written = mock_prisma_client.db.litellm_teamtable.update.call_args.kwargs["data"] + assert "members" not in written + assert written["team_alias"] == "Renamed" + + +def _mock_prisma_for_delete_user(mocker, team): + mock_prisma_client = mocker.MagicMock() + mock_prisma_client.db = mocker.MagicMock() + mock_prisma_client.db.litellm_teamtable = mocker.MagicMock() + mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=team) + mock_prisma_client.db.litellm_teamtable.update = AsyncMock() + mock_prisma_client.db.litellm_usertable = mocker.MagicMock() + mock_prisma_client.db.litellm_usertable.delete = AsyncMock() + return mock_prisma_client + + +def _patch_delete_user_dependencies(mocker, mock_prisma_client, existing_user): + mocker.patch( + "litellm.proxy.management_endpoints.scim.scim_v2._get_prisma_client_or_raise_exception", + AsyncMock(return_value=mock_prisma_client), + ) + mocker.patch( + "litellm.proxy.management_endpoints.scim.scim_v2._check_user_exists", + AsyncMock(return_value=existing_user), + ) + mocker.patch( + "litellm.proxy.management_endpoints.scim.scim_v2._set_user_keys_blocked", + AsyncMock(), + ) + mocker.patch( + "litellm.proxy.management_endpoints.scim.scim_v2._delete_rows_referencing_user", + AsyncMock(), + ) + + +@pytest.mark.asyncio +async def test_delete_user_prunes_members_with_roles(mocker): + """Deleting a SCIM user must remove them from every team they belong to via + team_member_delete, which prunes members_with_roles (the source of truth for + SCIM group membership) so GET /Groups no longer returns a dangling reference + to the now-deleted user.""" + user_id = "scim-del-user" + + existing_user = mocker.MagicMock() + existing_user.teams = ["team-1"] + + team = LiteLLM_TeamTable( + team_id="team-1", + members=[user_id, "other-user"], + members_with_roles=[Member(user_id=user_id, role="user"), Member(user_id="other-user", role="admin")], + ) + + mock_prisma_client = _mock_prisma_for_delete_user(mocker, team) + _patch_delete_user_dependencies(mocker, mock_prisma_client, existing_user) + team_member_delete_mock = mocker.patch( + "litellm.proxy.management_endpoints.scim.scim_v2.team_member_delete", + AsyncMock(), + ) + + await delete_user(user_id=user_id) + + team_member_delete_mock.assert_awaited_once() + call = team_member_delete_mock.call_args + assert call.kwargs["data"].team_id == "team-1" + assert call.kwargs["data"].user_id == user_id + assert call.kwargs["user_api_key_dict"].user_role == LitellmUserRoles.PROXY_ADMIN + mock_prisma_client.db.litellm_usertable.delete.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_delete_user_surfaces_prune_failure_and_keeps_user(mocker): + """A genuine failure while pruning members_with_roles must surface: the + endpoint fails loudly and the user row is NOT deleted, so we never report a + successful delete while leaving a dangling member (SCIM DELETE is idempotent, + so the IdP retries).""" + user_id = "scim-del-user" + + existing_user = mocker.MagicMock() + existing_user.teams = ["team-1"] + + team = LiteLLM_TeamTable( + team_id="team-1", + members=[user_id], + members_with_roles=[Member(user_id=user_id, role="user")], + ) + + mock_prisma_client = _mock_prisma_for_delete_user(mocker, team) + _patch_delete_user_dependencies(mocker, mock_prisma_client, existing_user) + mocker.patch( + "litellm.proxy.management_endpoints.scim.scim_v2.team_member_delete", + AsyncMock(side_effect=Exception("database connection lost")), + ) + + with pytest.raises(Exception): + await delete_user(user_id=user_id) + + mock_prisma_client.db.litellm_usertable.delete.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_delete_user_skips_teams_where_not_a_member(mocker): + """If the user is not in a team's members_with_roles, deletion must treat that + team as a no-op (no team_member_delete call, no error) and still delete the + user, so a stale legacy membership can't block the delete.""" + user_id = "scim-del-user" + + existing_user = mocker.MagicMock() + existing_user.teams = ["team-1"] + + team = LiteLLM_TeamTable( + team_id="team-1", + members=[user_id], + members_with_roles=[Member(user_id="someone-else", role="admin")], + ) + + mock_prisma_client = _mock_prisma_for_delete_user(mocker, team) + _patch_delete_user_dependencies(mocker, mock_prisma_client, existing_user) + team_member_delete_mock = mocker.patch( + "litellm.proxy.management_endpoints.scim.scim_v2.team_member_delete", + AsyncMock(), + ) + + await delete_user(user_id=user_id) + + team_member_delete_mock.assert_not_awaited() + mock_prisma_client.db.litellm_usertable.delete.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_patch_group_add_applies_delta_and_keeps_concurrent_add(mocker): + """A group PATCH op:add must be applied as a delta against the live roster, + not as a snapshot-based absolute target. + + When a concurrent PATCH has already added a member between this request's + initial read and its post-write refresh, that member shows up in the + refreshed roster but not in this request's snapshot-derived target. Diffing + the refreshed roster against the snapshot target would issue a spurious + team_member_delete for the concurrently-added member. Applying only this + request's intended delta on top of the refreshed roster must retain them. + """ + from litellm.proxy.management_endpoints.scim.scim_transformations import ( + ScimTransformations, + ) + + group_id = "team-concurrent" + + snapshot_team = LiteLLM_TeamTable( + team_id=group_id, + team_alias="Group", + members_with_roles=[Member(user_id="zed", role="user")], + metadata={"externalId": "grp-ext"}, + ) + refreshed_team = LiteLLM_TeamTable( + team_id=group_id, + team_alias="Group", + members_with_roles=[ + Member(user_id="zed", role="user"), + Member(user_id="alice", role="user"), + ], + metadata={"externalId": "grp-ext"}, + ) + final_team = LiteLLM_TeamTable( + team_id=group_id, + team_alias="Group", + members_with_roles=[ + Member(user_id="zed", role="user"), + Member(user_id="alice", role="user"), + Member(user_id="bob", role="user"), + ], + metadata={"externalId": "grp-ext"}, + ) + + patch_ops = SCIMPatchOp( + schemas=["urn:ietf:params:scim:api:messages:2.0:PatchOp"], + Operations=[SCIMPatchOperation(op="add", path="members", value=[{"value": "bob"}])], + ) + + mock_prisma_client = mocker.MagicMock() + mock_prisma_client.db = mocker.MagicMock() + mock_prisma_client.db.litellm_teamtable = mocker.MagicMock() + mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock( + side_effect=[snapshot_team, refreshed_team, final_team] + ) + mock_prisma_client.db.litellm_teamtable.update = AsyncMock(return_value=final_team) + mock_prisma_client.db.litellm_usertable = mocker.MagicMock() + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=mocker.MagicMock()) + + mocker.patch( + "litellm.proxy.management_endpoints.scim.scim_v2._get_prisma_client_or_raise_exception", + AsyncMock(return_value=mock_prisma_client), + ) + patch_membership_mock = mocker.patch( + "litellm.proxy.management_endpoints.scim.scim_v2.patch_team_membership", + AsyncMock(), + ) + mocker.patch( + "litellm.proxy.management_endpoints.scim.scim_v2._recompute_scim_member_roles", + AsyncMock(), + ) + mocker.patch.object( + ScimTransformations, + "transform_litellm_team_to_scim_group", + AsyncMock( + return_value=SCIMGroup( + schemas=["urn:ietf:params:scim:schemas:core:2.0:Group"], + id=group_id, + displayName="Group", + ) + ), + ) + + await patch_group(group_id=group_id, patch_ops=patch_ops) + + calls = patch_membership_mock.call_args_list + + removed_user_ids = { + call.kwargs["user_id"] for call in calls if call.kwargs.get("teams_ids_to_remove_user_from") == [group_id] + } + assert removed_user_ids == set() + + added_user_ids = { + call.kwargs["user_id"] for call in calls if call.kwargs.get("teams_ids_to_add_user_to") == [group_id] + } + assert added_user_ids == {"bob"} + + +@pytest.mark.asyncio +async def test_patch_group_replace_stays_absolute_against_concurrent_roster(mocker): + """A group PATCH ``replace`` op declares the roster is exactly the given set, + so it must reconcile as a set-to-target, not as a delta. + + Unlike ``add``/``remove``, ``replace`` is absolute. A member that another + request added concurrently is present in the refreshed roster but not in the + replace target, and ``replace`` must drop it. Rebasing the replace onto the + refreshed roster (the delta behavior correct only for add/remove) would + wrongly retain that concurrently-added member. + """ + from litellm.proxy.management_endpoints.scim.scim_transformations import ( + ScimTransformations, + ) + + group_id = "team-replace-concurrent" + + snapshot_team = LiteLLM_TeamTable( + team_id=group_id, + team_alias="Group", + members_with_roles=[Member(user_id="zed", role="user")], + metadata={"externalId": "grp-ext"}, + ) + refreshed_team = LiteLLM_TeamTable( + team_id=group_id, + team_alias="Group", + members_with_roles=[ + Member(user_id="alice", role="user"), + Member(user_id="bob", role="user"), + ], + metadata={"externalId": "grp-ext"}, + ) + final_team = LiteLLM_TeamTable( + team_id=group_id, + team_alias="Group", + members_with_roles=[Member(user_id="alice", role="user")], + metadata={"externalId": "grp-ext"}, + ) + + patch_ops = SCIMPatchOp( + schemas=["urn:ietf:params:scim:api:messages:2.0:PatchOp"], + Operations=[SCIMPatchOperation(op="replace", path="members", value=[{"value": "alice"}])], + ) + + mock_prisma_client = mocker.MagicMock() + mock_prisma_client.db = mocker.MagicMock() + mock_prisma_client.db.litellm_teamtable = mocker.MagicMock() + mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock( + side_effect=[snapshot_team, refreshed_team, final_team] + ) + mock_prisma_client.db.litellm_teamtable.update = AsyncMock(return_value=final_team) + mock_prisma_client.db.litellm_usertable = mocker.MagicMock() + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=mocker.MagicMock()) + + mocker.patch( + "litellm.proxy.management_endpoints.scim.scim_v2._get_prisma_client_or_raise_exception", + AsyncMock(return_value=mock_prisma_client), + ) + patch_membership_mock = mocker.patch( + "litellm.proxy.management_endpoints.scim.scim_v2.patch_team_membership", + AsyncMock(), + ) + mocker.patch( + "litellm.proxy.management_endpoints.scim.scim_v2._recompute_scim_member_roles", + AsyncMock(), + ) + mocker.patch.object( + ScimTransformations, + "transform_litellm_team_to_scim_group", + AsyncMock( + return_value=SCIMGroup( + schemas=["urn:ietf:params:scim:schemas:core:2.0:Group"], + id=group_id, + displayName="Group", + ) + ), + ) + + await patch_group(group_id=group_id, patch_ops=patch_ops) + + calls = patch_membership_mock.call_args_list + + removed_user_ids = { + call.kwargs["user_id"] for call in calls if call.kwargs.get("teams_ids_to_remove_user_from") == [group_id] + } + assert removed_user_ids == {"bob"} + + added_user_ids = { + call.kwargs["user_id"] for call in calls if call.kwargs.get("teams_ids_to_add_user_to") == [group_id] + } + assert added_user_ids == set() + + +@pytest.mark.parametrize( + "path, attribute, expected", + [ + ('members[value eq "user-1"]', "members", ["user-1"]), + ("members[value eq 'user-1']", "members", ["user-1"]), + ('members[value EQ "user-1"]', "members", ["user-1"]), + ('members[ value eq "user-1" ]', "members", ["user-1"]), + ('groups[value eq "team-1"]', "groups", ["team-1"]), + ('members[value eq "Mixed-CASE-Id"]', "members", ["Mixed-CASE-Id"]), + ('members[value eq "a\\"b"]', "members", ['a"b']), + ('members[value eq "a\\\\b"]', "members", ["a\\b"]), + ("members[value eq 'a\\'b']", "members", ["a'b"]), + ("members", "members", []), + ('groups[value eq "team-1"]', "members", []), + (None, "members", []), + ('members[value eq ""]', "members", []), + ("members[value eq user-1]", "members", []), + ("members[value eq unintendeduser]", "members", []), + ], +) +def test_extract_ids_from_path_filter(path, attribute, expected): + assert _extract_ids_from_path_filter(path, attribute) == expected + + +def test_extract_ids_from_path_filter_unterminated_is_linear(): + """A pathological unterminated quoted filter must not trigger super-linear + backtracking; it returns no id and completes near-instantly.""" + pathological = 'members[value eq "' + ("\\" * 200) + + start = time.perf_counter() + result = _extract_ids_from_path_filter(pathological, "members") + elapsed = time.perf_counter() - start + + assert result == [] + assert elapsed < 1.0 + + +@pytest.mark.asyncio +async def test_process_group_patch_remove_filtered_path_without_value(mocker): + """Okta sends group membership removals as a filtered path with no request + body value; the member id must be parsed out of members[value eq "..."]""" + patch_ops = SCIMPatchOp( + schemas=["urn:ietf:params:scim:api:messages:2.0:PatchOp"], + Operations=[SCIMPatchOperation(op="remove", path='members[value eq "user-1"]')], + ) + + existing_team = LiteLLM_TeamTable( + team_id="team-1", + team_alias="Team One", + members=[], + members_with_roles=[ + Member(user_id="user-1", role="user"), + Member(user_id="user-2", role="user"), + ], + ) + + prisma_client = mocker.MagicMock() + prisma_client.db = mocker.MagicMock() + prisma_client.db.litellm_usertable = mocker.MagicMock() + prisma_client.db.litellm_usertable.find_unique = AsyncMock( + return_value=LiteLLM_UserTable(user_id="user-1") + ) + + _, final_members, _ = await _process_group_patch_operations( + patch_ops=patch_ops, + existing_team=existing_team, + prisma_client=prisma_client, + ) + + assert final_members == {"user-2"} + + +@pytest.mark.asyncio +async def test_process_group_patch_add_filtered_path_without_value(mocker): + """A filtered add path with no body value adds the id parsed from the filter.""" + patch_ops = SCIMPatchOp( + schemas=["urn:ietf:params:scim:api:messages:2.0:PatchOp"], + Operations=[SCIMPatchOperation(op="add", path='members[value eq "user-3"]')], + ) + + existing_team = LiteLLM_TeamTable( + team_id="team-1", + team_alias="Team One", + members=[], + members_with_roles=[Member(user_id="user-1", role="user")], + ) + + prisma_client = mocker.MagicMock() + prisma_client.db = mocker.MagicMock() + prisma_client.db.litellm_usertable = mocker.MagicMock() + prisma_client.db.litellm_usertable.find_unique = AsyncMock( + return_value=LiteLLM_UserTable(user_id="user-3") + ) + + _, final_members, _ = await _process_group_patch_operations( + patch_ops=patch_ops, + existing_team=existing_team, + prisma_client=prisma_client, + ) + + assert final_members == {"user-1", "user-3"} + + +@pytest.mark.asyncio +async def test_process_group_patch_replace_empty_value_does_not_use_path_filter(mocker): + """An explicit empty replace value must clear membership rather than pull an + id from the filtered path, which would retain one member and drop the rest.""" + patch_ops = SCIMPatchOp( + schemas=["urn:ietf:params:scim:api:messages:2.0:PatchOp"], + Operations=[ + SCIMPatchOperation(op="replace", path='members[value eq "user-1"]', value=[]) + ], + ) + + existing_team = LiteLLM_TeamTable( + team_id="team-1", + team_alias="Team One", + members=[], + members_with_roles=[ + Member(user_id="user-1", role="user"), + Member(user_id="user-2", role="user"), + ], + ) + + prisma_client = mocker.MagicMock() + prisma_client.db = mocker.MagicMock() + prisma_client.db.litellm_usertable = mocker.MagicMock() + prisma_client.db.litellm_usertable.find_unique = AsyncMock( + return_value=LiteLLM_UserTable(user_id="user-1") + ) + + _, final_members, _ = await _process_group_patch_operations( + patch_ops=patch_ops, + existing_team=existing_team, + prisma_client=prisma_client, + ) + + assert final_members == set() diff --git a/tests/test_litellm/proxy/management_endpoints/test_cache_settings_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_cache_settings_endpoints.py index f4c6d4f8d15..2504b5744fc 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_cache_settings_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_cache_settings_endpoints.py @@ -17,9 +17,14 @@ from litellm.proxy._types import LitellmTableNames, LitellmUserRoles from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth from litellm.proxy.management_endpoints.cache_settings_endpoints import ( _CACHE_SENSITIVE_FIELDS, + _REDACTED_VALUE, CacheSettingsManager, CacheSettingsUpdateRequest, CacheTestRequest, + _merge_over_saved, + _overlay_environment, + _parse_stored_settings, + _redact_credentials, _resolve_cache_url_precedence, get_cache_settings, test_cache_connection, @@ -610,3 +615,510 @@ async def test_update_cache_settings_no_audit_when_disabled(monkeypatch): ) assert audit_calls == [] + + +class TestParseStoredSettings: + """The stored blob arrives as a JSON string or a parsed dict; both must + normalize to a dict so the secret-preservation read never silently drops it.""" + + def test_parses_a_json_string(self): + assert _parse_stored_settings('{"host": "h", "password": "pw"}') == {"host": "h", "password": "pw"} + + def test_passes_a_dict_through(self): + assert _parse_stored_settings({"host": "h", "password": "pw"}) == {"host": "h", "password": "pw"} + + def test_non_mapping_becomes_empty(self): + assert _parse_stored_settings(None) == {} + assert _parse_stored_settings("[1, 2]") == {} + + +class TestMergeOverSaved: + """The secret-preservation contract behind the redacted-resubmit fix.""" + + def test_redacted_secret_restores_stored_value(self): + # same connection target, an unrelated field edited: the stored secret + # is restored behind the redacted resubmit + merged = _merge_over_saved( + incoming={"type": "redis", "host": "samehost", "namespace": "new", "password": _REDACTED_VALUE}, + saved={"type": "redis", "host": "samehost", "password": "realpw"}, + ) + assert merged["namespace"] == "new" + assert merged["password"] == "realpw" + + def test_stored_secret_not_replayed_to_a_different_target(self): + # credential replay guard: omitting the password while pointing at a new + # host must NOT resurrect the stored secret (it would be sent elsewhere) + merged = _merge_over_saved( + incoming={"type": "redis", "host": "attacker.example.com", "password": _REDACTED_VALUE}, + saved={"type": "redis", "host": "real-redis", "password": "realpw"}, + ) + assert "password" not in merged + + def test_omitted_secret_restores_stored_value(self): + # same host (target unchanged), password field omitted entirely + merged = _merge_over_saved( + incoming={"type": "redis", "host": "samehost", "namespace": "n"}, + saved={"type": "redis", "host": "samehost", "password": "realpw"}, + ) + assert merged["password"] == "realpw" + + def test_sentinel_password_not_replayed_to_different_sentinel_nodes(self): + # sentinel target change with an omitted sentinel_password must not + # resurrect the stored one and send it to the caller's sentinels + merged = _merge_over_saved( + incoming={"type": "redis", "sentinel_nodes": [["attacker", 26379]], "service_name": "mymaster"}, + saved={ + "type": "redis", + "sentinel_nodes": [["real", 26379]], + "service_name": "mymaster", + "sentinel_password": "realsp", + }, + ) + assert "sentinel_password" not in merged + + def test_sentinel_password_preserved_when_sentinel_target_unchanged(self): + merged = _merge_over_saved( + incoming={"type": "redis", "sentinel_nodes": [["real", 26379]], "service_name": "mymaster"}, + saved={ + "type": "redis", + "sentinel_nodes": [["real", 26379]], + "service_name": "mymaster", + "sentinel_password": "realsp", + }, + ) + assert merged["sentinel_password"] == "realsp" + + def test_password_not_replayed_to_different_cluster_nodes(self): + merged = _merge_over_saved( + incoming={"type": "redis", "redis_startup_nodes": [{"host": "attacker", "port": "7001"}]}, + saved={ + "type": "redis", + "redis_startup_nodes": [{"host": "real", "port": "7001"}], + "password": "realpw", + }, + ) + assert "password" not in merged + + def test_equivalent_target_representations_still_preserve_secret(self): + # the client sends port as a string, storage holds it as an int: the + # target is unchanged, so the untouched password must not be dropped + merged = _merge_over_saved( + incoming={"type": "redis", "host": "h", "port": "6379", "password": _REDACTED_VALUE}, + saved={"type": "redis", "host": "h", "port": 6379, "password": "realpw"}, + ) + assert merged["password"] == "realpw" + + def test_explicit_empty_string_clears_the_secret(self): + merged = _merge_over_saved( + incoming={"type": "redis", "host": "h", "password": ""}, + saved={"type": "redis", "host": "h", "password": "realpw"}, + ) + assert merged.get("password") == "" + + def test_explicit_null_clears_the_secret(self): + # an explicit null is a clear, not an omission, so it must not restore + merged = _merge_over_saved( + incoming={"type": "redis", "host": "h", "password": None}, + saved={"type": "redis", "host": "h", "password": "realpw"}, + ) + assert merged.get("password") is None + + def test_secret_not_reused_when_a_pinned_target_field_is_omitted(self): + # omitting the host (a pinned target) means the request does not describe + # the stored target, so the stored secret must not be restored (and thus + # cannot be sent to whatever host the incomplete request resolves to) + merged = _merge_over_saved( + incoming={"type": "redis", "port": "6379"}, + saved={"type": "redis", "host": "real", "port": 6379, "password": "realpw"}, + ) + assert "password" not in merged + + def test_redacted_secret_with_no_stored_value_is_dropped(self): + # env-sourced secret: nothing stored to restore, so the marker must not + # be persisted; the environment stays the source at runtime + merged = _merge_over_saved( + incoming={"type": "redis", "host": "h", "password": _REDACTED_VALUE}, + saved={}, + ) + assert "password" not in merged + + def test_new_secret_value_wins(self): + merged = _merge_over_saved( + incoming={"password": "brandnewpw"}, + saved={"password": "realpw"}, + ) + assert merged["password"] == "brandnewpw" + + def test_switching_from_url_to_host_port_drops_stored_url(self): + # admin migrates a url-mode cache to discrete host/port: the stored url + # must not be resurrected (url precedence would then discard host/port) + merged = _merge_over_saved( + incoming={"type": "redis", "host": "newhost", "port": "6379"}, + saved={"type": "redis", "url": "redis://:pw@oldhost:6379/0"}, + ) + assert "url" not in merged + assert merged["host"] == "newhost" + assert merged["port"] == "6379" + + def test_untouched_url_is_preserved_without_a_discrete_target(self): + # a url-mode save that touches nothing keeps the stored url + merged = _merge_over_saved( + incoming={"type": "redis", "namespace": "ns"}, + saved={"type": "redis", "url": "redis://:pw@host:6379/0"}, + ) + assert merged["url"] == "redis://:pw@host:6379/0" + + +def test_overlay_environment_fills_unset_connection_fields(monkeypatch): + """A cache with no stored connection resolves REDIS_* env for the UI.""" + for var in ("REDIS_URL", "REDIS_HOST", "REDIS_PORT", "REDIS_PASSWORD", "REDIS_USERNAME"): + monkeypatch.delenv(var, raising=False) + monkeypatch.setenv("REDIS_HOST", "redis.internal") + monkeypatch.setenv("REDIS_PORT", "6380") + monkeypatch.setenv("REDIS_PASSWORD", "env-password") + + effective = _overlay_environment({}) + + assert effective["host"] == "redis.internal" + assert effective["port"] == "6380" + assert effective["password"] == "env-password" + assert effective["type"] == "redis" + + +def test_overlay_environment_stored_value_wins(monkeypatch): + monkeypatch.setenv("REDIS_HOST", "env-host") + effective = _overlay_environment({"type": "redis", "host": "stored-host"}) + assert effective["host"] == "stored-host" + + +@pytest.mark.asyncio +async def test_get_cache_settings_falls_back_to_redis_env(monkeypatch): + """A cache configured purely through REDIS_* env vars shows its effective + connection instead of a blank page, with the password redacted.""" + for var in ("REDIS_URL", "REDIS_HOST", "REDIS_PORT", "REDIS_PASSWORD", "REDIS_USERNAME"): + monkeypatch.delenv(var, raising=False) + monkeypatch.setenv("REDIS_HOST", "redis.internal") + monkeypatch.setenv("REDIS_PORT", "6380") + monkeypatch.setenv("REDIS_PASSWORD", "env-password") + + mock_prisma = MagicMock() + mock_prisma.db.litellm_cacheconfig.find_unique = AsyncMock(return_value=None) + + proxy_config = MagicMock() + proxy_config._decrypt_db_variables = MagicMock(side_effect=lambda variables_dict: dict(variables_dict)) + + with ( + patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), + patch("litellm.proxy.proxy_server.proxy_config", proxy_config), + ): + response = await get_cache_settings(user_api_key_dict=_admin_auth()) + + values = response.current_values + assert values["host"] == "redis.internal" + assert values["port"] == "6380" + assert values["type"] == "redis" + # the env password is surfaced as configured, not leaked in plaintext + assert values["password"] == _REDACTED_VALUE + + +@pytest.mark.asyncio +async def test_get_cache_settings_redacts_password_with_marker(monkeypatch): + for var in ("REDIS_URL", "REDIS_HOST", "REDIS_PORT", "REDIS_PASSWORD", "REDIS_USERNAME"): + monkeypatch.delenv(var, raising=False) + cache_row = MagicMock() + cache_row.cache_settings = json.dumps( + {"type": "redis", "host": "h", "password": "supersecret", "namespace": "ns"} + ) + mock_prisma = MagicMock() + mock_prisma.db.litellm_cacheconfig.find_unique = AsyncMock(return_value=cache_row) + proxy_config = MagicMock() + proxy_config._decrypt_db_variables = MagicMock(side_effect=lambda variables_dict: dict(variables_dict)) + + with ( + patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), + patch("litellm.proxy.proxy_server.proxy_config", proxy_config), + ): + response = await get_cache_settings(user_api_key_dict=_admin_auth()) + + assert response.current_values["password"] == _REDACTED_VALUE + assert response.current_values["namespace"] == "ns" + + +@pytest.mark.asyncio +async def test_get_cache_settings_url_mode_hides_env_discrete_fields(monkeypatch): + """A url-mode stored config must not surface env-overlaid host/port. + + Otherwise a no-op save would submit the env host and, via url precedence, + silently switch the cache off its configured url. + """ + for var in ("REDIS_URL", "REDIS_HOST", "REDIS_PORT", "REDIS_PASSWORD", "REDIS_USERNAME"): + monkeypatch.delenv(var, raising=False) + monkeypatch.setenv("REDIS_HOST", "env-host") + monkeypatch.setenv("REDIS_PORT", "6380") + + cache_row = MagicMock() + cache_row.cache_settings = {"type": "redis", "url": "redis://:pw@stored-host:6379/0"} + mock_prisma = MagicMock() + mock_prisma.db.litellm_cacheconfig.find_unique = AsyncMock(return_value=cache_row) + proxy_config = MagicMock() + proxy_config._decrypt_db_variables = MagicMock(side_effect=lambda variables_dict: dict(variables_dict)) + + with ( + patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), + patch("litellm.proxy.proxy_server.proxy_config", proxy_config), + ): + response = await get_cache_settings(user_api_key_dict=_admin_auth()) + + values = response.current_values + assert values["url"] == _REDACTED_VALUE + # the env host/port must not leak in and shadow the url + assert "host" not in values + assert "port" not in values + + +def _mock_proxy_config_identity_crypto(): + proxy_config = MagicMock() + proxy_config._encrypt_env_variables = MagicMock( + side_effect=lambda environment_variables: dict(environment_variables) + ) + proxy_config._decrypt_db_variables = MagicMock(side_effect=lambda variables_dict: dict(variables_dict)) + proxy_config._init_cache = MagicMock() + proxy_config.switch_on_llm_response_caching = MagicMock() + return proxy_config + + +@pytest.mark.asyncio +async def test_update_preserves_stored_password_on_redacted_resubmit(monkeypatch): + """Editing an unrelated field and re-submitting the redacted password must + keep the stored secret, not persist the marker over a working password.""" + monkeypatch.setattr(litellm, "store_audit_logs", False) + + existing = MagicMock() + # prisma returns the Json column as an already-parsed dict, not a JSON + # string; a reader that json.loads unconditionally would drop the whole row + existing.cache_settings = {"type": "redis", "host": "oldhost", "password": "realpw"} + mock_prisma = MagicMock() + mock_prisma.db.litellm_cacheconfig.find_unique = AsyncMock(return_value=existing) + mock_prisma.db.litellm_cacheconfig.upsert = AsyncMock() + proxy_config = _mock_proxy_config_identity_crypto() + + with ( + patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), + patch("litellm.proxy.proxy_server.proxy_config", proxy_config), + patch("litellm.proxy.proxy_server.store_model_in_db", True), + ): + result = await update_cache_settings( + request=CacheSettingsUpdateRequest( + # same host (the target is unchanged), an unrelated field edited + cache_settings={"type": "redis", "host": "oldhost", "namespace": "edited", "password": _REDACTED_VALUE} + ), + user_api_key_dict=_admin_auth(), + litellm_changed_by=None, + ) + + persisted = proxy_config._encrypt_env_variables.call_args.kwargs["environment_variables"] + assert persisted["host"] == "oldhost" + assert persisted["namespace"] == "edited" + assert persisted["password"] == "realpw" + # the response never echoes the plaintext secret back either + assert result["settings"]["password"] == _REDACTED_VALUE + + +@pytest.mark.asyncio +async def test_update_drops_env_sourced_redacted_secret(monkeypatch): + """With no stored row, a re-submitted redacted secret is env-sourced; the + marker must not be persisted so the environment stays the source.""" + monkeypatch.setattr(litellm, "store_audit_logs", False) + + mock_prisma = MagicMock() + mock_prisma.db.litellm_cacheconfig.find_unique = AsyncMock(return_value=None) + mock_prisma.db.litellm_cacheconfig.upsert = AsyncMock() + proxy_config = _mock_proxy_config_identity_crypto() + + with ( + patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), + patch("litellm.proxy.proxy_server.proxy_config", proxy_config), + patch("litellm.proxy.proxy_server.store_model_in_db", True), + ): + await update_cache_settings( + request=CacheSettingsUpdateRequest( + cache_settings={"type": "redis", "host": "h", "password": _REDACTED_VALUE} + ), + user_api_key_dict=_admin_auth(), + litellm_changed_by=None, + ) + + persisted = proxy_config._encrypt_env_variables.call_args.kwargs["environment_variables"] + assert "password" not in persisted + + +@pytest.mark.asyncio +async def test_update_applies_new_password(monkeypatch): + """A real new secret value replaces the stored one.""" + monkeypatch.setattr(litellm, "store_audit_logs", False) + + existing = MagicMock() + existing.cache_settings = json.dumps({"type": "redis", "host": "h", "password": "oldpw"}) + mock_prisma = MagicMock() + mock_prisma.db.litellm_cacheconfig.find_unique = AsyncMock(return_value=existing) + mock_prisma.db.litellm_cacheconfig.upsert = AsyncMock() + proxy_config = _mock_proxy_config_identity_crypto() + + with ( + patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), + patch("litellm.proxy.proxy_server.proxy_config", proxy_config), + patch("litellm.proxy.proxy_server.store_model_in_db", True), + ): + await update_cache_settings( + request=CacheSettingsUpdateRequest( + cache_settings={"type": "redis", "host": "h", "password": "brandnewpw"} + ), + user_api_key_dict=_admin_auth(), + litellm_changed_by=None, + ) + + persisted = proxy_config._encrypt_env_variables.call_args.kwargs["environment_variables"] + assert persisted["password"] == "brandnewpw" + + +@pytest.mark.asyncio +async def test_test_cache_connection_survives_saved_lookup_failure(monkeypatch): + """A failed saved-settings lookup must not block the connection test. + + The test endpoint reads the stored row to resolve a redacted credential, but + that read can raise (a misconfigured or unavailable client), and it must fall + back to the submitted settings rather than abort — otherwise a shared client + left in an odd state by another test would break every connection test. + """ + monkeypatch.setattr(litellm, "store_audit_logs", False) + + # a client whose find_unique is not awaitable, so the saved read raises + bad_prisma = MagicMock() + proxy_config = MagicMock() + proxy_config._decrypt_db_variables = MagicMock(side_effect=lambda variables_dict: dict(variables_dict)) + + cache_instance = MagicMock() + cache_instance.cache = MagicMock() + cache_instance.cache.test_connection = AsyncMock(return_value={"status": "success", "message": "ok"}) + + with ( + patch("litellm.proxy.proxy_server.prisma_client", bad_prisma), + patch("litellm.proxy.proxy_server.proxy_config", proxy_config), + patch("litellm.Cache") as mock_cache_class, + ): + mock_cache_class.return_value = cache_instance + result = await test_cache_connection( + request=CacheTestRequest(cache_settings={"type": "redis", "host": "h", "port": "6379", "password": "pw"}), + user_api_key_dict=_admin_auth(), + ) + + mock_cache_class.assert_called_once() + assert result.status == "success" + + +@pytest.mark.asyncio +async def test_get_cache_settings_does_not_surface_non_display_env_credentials(monkeypatch): + """The env overlay must not leak credential kwargs the UI does not manage. + + _redis_kwargs_from_environment resolves every redis.Redis kwarg, including + secrets like azure_client_secret; only cache display fields may be surfaced, + so a non-admin reading /cache/settings never retrieves such a credential. + """ + for var in ("REDIS_URL", "REDIS_HOST", "REDIS_PORT", "REDIS_PASSWORD", "REDIS_AZURE_CLIENT_SECRET"): + monkeypatch.delenv(var, raising=False) + monkeypatch.setenv("REDIS_HOST", "redis.internal") + monkeypatch.setenv("REDIS_AZURE_CLIENT_SECRET", "super-azure-secret") + + mock_prisma = MagicMock() + mock_prisma.db.litellm_cacheconfig.find_unique = AsyncMock(return_value=None) + proxy_config = MagicMock() + proxy_config._decrypt_db_variables = MagicMock(side_effect=lambda variables_dict: dict(variables_dict)) + + with ( + patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), + patch("litellm.proxy.proxy_server.proxy_config", proxy_config), + ): + response = await get_cache_settings(user_api_key_dict=_admin_auth()) + + values = response.current_values + assert values.get("host") == "redis.internal" + # the non-display credential must not appear in the response at all + assert "azure_client_secret" not in values + assert "super-azure-secret" not in values.values() + + +@pytest.mark.asyncio +async def test_test_cache_connection_does_not_log_plaintext_credentials(monkeypatch, caplog): + """The connection test must not write the resolved plaintext secret to logs. + + _merge_over_saved substitutes the stored password for a redacted resubmit, so + the settings dict carries the real secret; the debug log must redact it. + """ + import logging + + existing = MagicMock() + existing.cache_settings = {"type": "redis", "host": "h", "port": "6379", "password": "realredispw"} + mock_prisma = MagicMock() + mock_prisma.db.litellm_cacheconfig.find_unique = AsyncMock(return_value=existing) + proxy_config = MagicMock() + proxy_config._decrypt_db_variables = MagicMock(side_effect=lambda variables_dict: dict(variables_dict)) + + cache_instance = MagicMock() + cache_instance.cache = MagicMock() + cache_instance.cache.test_connection = AsyncMock(return_value={"status": "success", "message": "ok"}) + + with ( + patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), + patch("litellm.proxy.proxy_server.proxy_config", proxy_config), + patch("litellm.Cache") as mock_cache_class, + caplog.at_level(logging.DEBUG, logger="LiteLLM Proxy"), + ): + mock_cache_class.return_value = cache_instance + # resubmit the redacted marker; the merge resolves it to the stored secret + await test_cache_connection( + request=CacheTestRequest( + cache_settings={"type": "redis", "host": "h", "port": "6379", "password": _REDACTED_VALUE} + ), + user_api_key_dict=_admin_auth(), + ) + + # the real password was used to build the client but never written to the log + assert mock_cache_class.call_args.kwargs["password"] == "realredispw" + assert "realredispw" not in caplog.text + + +@pytest.mark.asyncio +async def test_test_cache_connection_does_not_replay_saved_password_to_new_host(monkeypatch): + """Credential-replay guard on the connection test. + + A caller that submits a different host while omitting the password must not + have the stored password restored and sent to the caller-chosen host. + """ + existing = MagicMock() + existing.cache_settings = {"type": "redis", "host": "real-redis", "port": "6379", "password": "realredispw"} + mock_prisma = MagicMock() + mock_prisma.db.litellm_cacheconfig.find_unique = AsyncMock(return_value=existing) + proxy_config = MagicMock() + proxy_config._decrypt_db_variables = MagicMock(side_effect=lambda variables_dict: dict(variables_dict)) + + cache_instance = MagicMock() + cache_instance.cache = MagicMock() + cache_instance.cache.test_connection = AsyncMock(return_value={"status": "success", "message": "ok"}) + + with ( + patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), + patch("litellm.proxy.proxy_server.proxy_config", proxy_config), + patch("litellm.Cache") as mock_cache_class, + ): + mock_cache_class.return_value = cache_instance + await test_cache_connection( + request=CacheTestRequest( + cache_settings={"type": "redis", "host": "attacker.example.com", "port": "6379"} + ), + user_api_key_dict=_admin_auth(), + ) + + called_kwargs = mock_cache_class.call_args.kwargs + # the stored password is NOT sent to the attacker-chosen host + assert called_kwargs.get("password") != "realredispw" + assert "password" not in called_kwargs diff --git a/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py index 3e5bd3e9b7f..e1aaf398f97 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py @@ -2,6 +2,7 @@ import os import sys import types import json +from contextlib import ExitStack from datetime import datetime, timedelta from types import SimpleNamespace from typing import List, Optional @@ -2025,8 +2026,469 @@ class TestTemporaryMCPSessionEndpoints: code_challenge_method="S256", response_type="code", scope="scope1", + ephemeral_dcr_client=None, ) + async def _authorize_without_client_id( + self, server, mint_mock=None, code_challenge="chal", code_challenge_method="S256" + ): + """Drive mcp_authorize with no caller client_id against ``server``, returning the + (authorize_with_server mock, raised HTTPException or None) pair. Sends a valid S256 PKCE + pair by default because the ephemeral mint requires it.""" + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + mcp_authorize, + ) + + request = MagicMock() + admin_auth = generate_mock_user_api_key_auth(user_role=LitellmUserRoles.PROXY_ADMIN) + patches = [ + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints._get_cached_temporary_mcp_server_or_404", + return_value=server, + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.authorize_with_server", + AsyncMock(return_value=MagicMock()), + ), + ] + if mint_mock is not None: + patches.append( + patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints.mint_ephemeral_dcr_client", + mint_mock, + ) + ) + with ExitStack() as stack: + entered = [stack.enter_context(p) for p in patches] + authorize_mock = entered[1] + try: + await mcp_authorize( + request=request, + server_id=server.server_id, + user_api_key_dict=admin_auth, + client_id=None, + redirect_uri="http://127.0.0.1:60108/callback", + state="state123", + code_challenge=code_challenge, + code_challenge_method=code_challenge_method, + ) + except HTTPException as exc: + return authorize_mock, exc + return authorize_mock, None + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "code_challenge, code_challenge_method", + [(None, None), ("chal", "plain"), ("chal", None)], + ) + async def test_mcp_authorize_mint_requires_s256_pkce(self, code_challenge, code_challenge_method): + """Without PKCE the sealed code would be bearer-redeemable by any authenticated caller who + intercepts the redirect, so the ephemeral mint refuses to run for a downgraded flow (no + challenge, or a non-S256 method) before any upstream registration happens.""" + server = generate_mock_mcp_server_config_record(server_id="server-1") + server.auth_type = MCPAuth.true_passthrough + server.authorization_url = "https://idp.example.com/authorize" + server.registration_url = "https://idp.example.com/register" + mint_mock = AsyncMock() + + authorize_mock, exc = await self._authorize_without_client_id( + server, mint_mock=mint_mock, code_challenge=code_challenge, code_challenge_method=code_challenge_method + ) + + assert exc is not None + assert exc.status_code == 400 + assert "PKCE" in str(exc.detail) + mint_mock.assert_not_awaited() + authorize_mock.assert_not_awaited() + + @pytest.mark.asyncio + @pytest.mark.parametrize("auth_type", [MCPAuth.true_passthrough, MCPAuth.oauth_delegate]) + async def test_mcp_authorize_client_forwarded_modes_mint_ephemeral_dcr_client_when_none_supplied(self, auth_type): + """LIT-4581 regression: a client-forwarded-token server created without an auth step has no + stored client_id and the tools-tab browser flow supplies none, so authorize must fall + through to a gateway-side DCR mint and proceed with the minted client instead of + dead-ending on a 400 missing_client_id. Both modes share the caller-held-client contract, + so both get the fall-through.""" + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + EphemeralDcrClient, + ) + + server = generate_mock_mcp_server_config_record(server_id="server-1") + server.auth_type = auth_type + server.authorization_url = "https://idp.example.com/authorize" + server.registration_url = "https://idp.example.com/register" + minted = EphemeralDcrClient(client_id="minted-77", client_secret="mint-secret") + mint_mock = AsyncMock(return_value=minted) + + authorize_mock, exc = await self._authorize_without_client_id(server, mint_mock=mint_mock) + + assert exc is None + mint_mock.assert_awaited_once() + assert authorize_mock.await_args.kwargs["client_id"] == "minted-77" + assert authorize_mock.await_args.kwargs["ephemeral_dcr_client"] is minted + + @pytest.mark.asyncio + async def test_mcp_authorize_rejects_untrusted_redirect_before_minting(self): + """An untrusted redirect_uri must be rejected before the gateway performs any upstream + registration, so bad-redirect requests cannot be used to generate orphan clients at the + IdP.""" + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + mcp_authorize, + ) + + server = generate_mock_mcp_server_config_record(server_id="server-1") + server.auth_type = MCPAuth.true_passthrough + server.authorization_url = "https://idp.example.com/authorize" + server.registration_url = "https://idp.example.com/register" + mint_mock = AsyncMock() + admin_auth = generate_mock_user_api_key_auth(user_role=LitellmUserRoles.PROXY_ADMIN) + request = MagicMock() + request.base_url = "https://litellm.example.com/" + request.headers = {} + + with ( + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints._get_cached_temporary_mcp_server_or_404", + return_value=server, + ), + patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints.mint_ephemeral_dcr_client", + mint_mock, + ), + ): + with pytest.raises(HTTPException) as exc: + await mcp_authorize( + request=request, + server_id="server-1", + user_api_key_dict=admin_auth, + client_id=None, + redirect_uri="https://evil.example.net/steal", + state="state123", + code_challenge="chal", + code_challenge_method="S256", + ) + + assert exc.value.status_code == 400 + mint_mock.assert_not_awaited() + + @pytest.mark.asyncio + async def test_mcp_authorize_true_passthrough_without_authorization_url_reports_the_real_fault(self): + """A passthrough server whose discovery never yielded an authorize endpoint cannot start any + flow, minted client or not, so the error names the missing authorization url instead of the + misleading missing_client_id remedy.""" + server = generate_mock_mcp_server_config_record(server_id="server-1") + server.auth_type = MCPAuth.true_passthrough + server.authorization_url = None + server.registration_url = "https://idp.example.com/register" + mint_mock = AsyncMock() + + authorize_mock, exc = await self._authorize_without_client_id(server, mint_mock=mint_mock) + + assert exc is not None + assert exc.status_code == 400 + assert "authorization url" in str(exc.detail) + mint_mock.assert_not_awaited() + authorize_mock.assert_not_awaited() + + @pytest.mark.asyncio + async def test_mcp_authorize_true_passthrough_without_registration_endpoint_keeps_missing_client_id(self): + """When the upstream exposes no registration endpoint the mint is impossible, so the + authorize fails closed with the existing missing_client_id 400 instead of proceeding with an + empty client.""" + server = generate_mock_mcp_server_config_record(server_id="server-1") + server.auth_type = MCPAuth.true_passthrough + server.authorization_url = "https://idp.example.com/authorize" + server.registration_url = None + + authorize_mock, exc = await self._authorize_without_client_id(server) + + assert exc is not None + assert exc.status_code == 400 + assert exc.detail["error"] == "missing_client_id" + authorize_mock.assert_not_awaited() + + @pytest.mark.asyncio + async def test_mcp_authorize_oauth2_server_does_not_mint(self): + """The ephemeral mint is scoped to the client-forwarded-token modes: a plain oauth2 server + keeps the gateway-held-client contract (its client is persisted by the admin register flow), + so an empty client_id stays a 400 and no upstream registration is attempted.""" + server = generate_mock_mcp_server_config_record(server_id="server-1") + server.auth_type = MCPAuth.oauth2 + server.authorization_url = "https://idp.example.com/authorize" + server.registration_url = "https://idp.example.com/register" + mint_mock = AsyncMock() + + authorize_mock, exc = await self._authorize_without_client_id(server, mint_mock=mint_mock) + + assert exc is not None + assert exc.status_code == 400 + assert exc.detail["error"] == "missing_client_id" + mint_mock.assert_not_awaited() + authorize_mock.assert_not_awaited() + + @pytest.mark.asyncio + async def test_mcp_authorize_true_passthrough_dcr_bridge_mints_too(self): + """The UI creates passthrough servers with dcr_bridge enabled by default, so the default + clientless tools-page authorize is a bridge server; it must mint exactly like a non-bridge + one (the minted flow runs the bridge short-circuit arm) instead of dead-ending on + missing_client_id. The relay front door stays reserved for clients that present their own + client_id.""" + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + EphemeralDcrClient, + ) + + server = generate_mock_mcp_server_config_record(server_id="server-1") + server.auth_type = MCPAuth.true_passthrough + server.dcr_bridge = True + server.authorization_url = "https://idp.example.com/authorize" + server.registration_url = "https://idp.example.com/register" + minted = EphemeralDcrClient(client_id="minted-77", client_secret=None) + mint_mock = AsyncMock(return_value=minted) + + authorize_mock, exc = await self._authorize_without_client_id(server, mint_mock=mint_mock) + + assert exc is None + mint_mock.assert_awaited_once() + assert authorize_mock.await_args.kwargs["client_id"] == "minted-77" + assert authorize_mock.await_args.kwargs["ephemeral_dcr_client"] is minted + + @pytest.mark.asyncio + async def test_mcp_authorize_oauth_delegate_dcr_bridge_does_not_mint(self): + """The interactive oauth_delegate dcr_bridge sign-in has its own sealed-identity flow that + captures the SSO user at authorize; the ephemeral mint must not preempt it.""" + server = generate_mock_mcp_server_config_record(server_id="server-1") + server.auth_type = MCPAuth.oauth_delegate + server.dcr_bridge = True + server.authorization_url = "https://idp.example.com/authorize" + server.registration_url = "https://idp.example.com/register" + mint_mock = AsyncMock() + + authorize_mock, exc = await self._authorize_without_client_id(server, mint_mock=mint_mock) + + assert exc is not None + assert exc.status_code == 400 + assert exc.detail["error"] == "missing_client_id" + mint_mock.assert_not_awaited() + authorize_mock.assert_not_awaited() + + @pytest.mark.asyncio + async def test_mcp_token_opens_sealed_passthrough_code_and_exchanges_with_minted_client(self): + """LIT-4581 regression, token leg: the client echoes back the sealed passthrough code the + callback forwarded, so the token endpoint recovers the ephemeral client and the real + upstream code from it and authenticates the exchange with them, with no client_id supplied + by the caller and none stored on the server.""" + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + seal_passthrough_authorization_code, + ) + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + mcp_token, + ) + + request = MagicMock() + request.base_url = "https://litellm.example.com/" + request.headers = {} + server = generate_mock_mcp_server_config_record(server_id="server-1") + server.auth_type = MCPAuth.true_passthrough + admin_auth = generate_mock_user_api_key_auth(user_role=LitellmUserRoles.PROXY_ADMIN) + + with ( + patch("litellm.proxy.proxy_server.master_key", "sk-lit4581-test-master-key"), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints._get_cached_temporary_mcp_server_or_404", + return_value=server, + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.exchange_token_with_server", + AsyncMock(return_value={"access_token": "token"}), + ) as exchange_mock, + ): + sealed = seal_passthrough_authorization_code( + upstream_code="up-code", + client_id="minted-77", + client_secret="mint-secret", + mcp_server_id="server-1", + token_endpoint_auth_method="client_secret_basic", + ) + result = await mcp_token( + request=request, + server_id="server-1", + user_api_key_dict=admin_auth, + grant_type="authorization_code", + code=sealed, + redirect_uri="https://example.com/callback", + client_id=None, + client_secret=None, + code_verifier="verifier", + refresh_token=None, + scope=None, + ) + + assert result == {"access_token": "token"} + assert exchange_mock.await_args.kwargs["code"] == "up-code" + assert exchange_mock.await_args.kwargs["client_id"] == "minted-77" + assert exchange_mock.await_args.kwargs["client_secret"] == "mint-secret" + assert exchange_mock.await_args.kwargs["redirect_uri"] == "https://litellm.example.com/callback" + assert exchange_mock.await_args.kwargs["client_token_endpoint_auth_method"] == "client_secret_basic" + + @pytest.mark.asyncio + async def test_mcp_token_refresh_grant_never_opens_sealed_code(self): + """The minted client is unrecoverable outside the single authorization_code flow by + contract: a refresh_token grant that echoes a leftover sealed passthrough code (plus any + verifier) must not recover the minted credentials, so a clientless server answers + missing_client_id and the client re-runs authorize instead.""" + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + seal_passthrough_authorization_code, + ) + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + mcp_token, + ) + + request = MagicMock() + request.base_url = "https://litellm.example.com/" + request.headers = {} + server = generate_mock_mcp_server_config_record(server_id="server-1") + server.auth_type = MCPAuth.true_passthrough + admin_auth = generate_mock_user_api_key_auth(user_role=LitellmUserRoles.PROXY_ADMIN) + + with ( + patch("litellm.proxy.proxy_server.master_key", "sk-lit4581-test-master-key"), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints._get_cached_temporary_mcp_server_or_404", + return_value=server, + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.exchange_token_with_server", + AsyncMock(return_value={"access_token": "token"}), + ) as exchange_mock, + ): + sealed = seal_passthrough_authorization_code( + upstream_code="up-code", + client_id="minted-77", + client_secret="mint-secret", + mcp_server_id="server-1", + token_endpoint_auth_method="client_secret_basic", + ) + with pytest.raises(HTTPException) as exc: + await mcp_token( + request=request, + server_id="server-1", + user_api_key_dict=admin_auth, + grant_type="refresh_token", + code=sealed, + redirect_uri="https://example.com/callback", + client_id=None, + client_secret=None, + code_verifier="verifier", + refresh_token="leftover-refresh", + scope=None, + ) + + assert exc.value.status_code == 400 + assert exc.value.detail["error"] == "missing_client_id" + exchange_mock.assert_not_awaited() + + @pytest.mark.asyncio + async def test_mcp_token_sealed_code_requires_code_verifier(self): + """A sealed code is minted only for S256 PKCE flows, so redeeming one without the + corresponding verifier is refused at the gateway rather than trusting the upstream to + enforce the binding.""" + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + seal_passthrough_authorization_code, + ) + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + mcp_token, + ) + + request = MagicMock() + server = generate_mock_mcp_server_config_record(server_id="server-1") + server.auth_type = MCPAuth.true_passthrough + admin_auth = generate_mock_user_api_key_auth(user_role=LitellmUserRoles.PROXY_ADMIN) + + with ( + patch("litellm.proxy.proxy_server.master_key", "sk-lit4581-test-master-key"), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints._get_cached_temporary_mcp_server_or_404", + return_value=server, + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.exchange_token_with_server", + AsyncMock(), + ) as exchange_mock, + ): + sealed = seal_passthrough_authorization_code( + upstream_code="up-code", client_id="minted-77", client_secret=None, mcp_server_id="server-1" + ) + with pytest.raises(HTTPException) as exc: + await mcp_token( + request=request, + server_id="server-1", + user_api_key_dict=admin_auth, + grant_type="authorization_code", + code=sealed, + redirect_uri="https://example.com/callback", + client_id=None, + client_secret=None, + code_verifier=None, + refresh_token=None, + scope=None, + ) + + assert exc.value.status_code == 400 + assert "code_verifier" in str(exc.value.detail) + exchange_mock.assert_not_awaited() + + @pytest.mark.asyncio + async def test_mcp_token_rejects_sealed_code_for_another_server(self): + """A sealed passthrough code is bound to the server it was minted for: presenting it at + another server's token endpoint is a 400 before any upstream exchange, so a code cannot be + replayed across a server boundary.""" + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + seal_passthrough_authorization_code, + ) + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + mcp_token, + ) + + request = MagicMock() + server = generate_mock_mcp_server_config_record(server_id="server-1") + server.auth_type = MCPAuth.true_passthrough + admin_auth = generate_mock_user_api_key_auth(user_role=LitellmUserRoles.PROXY_ADMIN) + + with ( + patch("litellm.proxy.proxy_server.master_key", "sk-lit4581-test-master-key"), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints._get_cached_temporary_mcp_server_or_404", + return_value=server, + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.exchange_token_with_server", + AsyncMock(), + ) as exchange_mock, + ): + sealed = seal_passthrough_authorization_code( + upstream_code="up-code", + client_id="minted-77", + client_secret=None, + mcp_server_id="a-different-server", + ) + with pytest.raises(HTTPException) as exc: + await mcp_token( + request=request, + server_id="server-1", + user_api_key_dict=admin_auth, + grant_type="authorization_code", + code=sealed, + redirect_uri="https://example.com/callback", + client_id=None, + client_secret=None, + code_verifier="verifier", + refresh_token=None, + scope=None, + ) + + assert exc.value.status_code == 400 + exchange_mock.assert_not_awaited() + @pytest.mark.asyncio async def test_mcp_authorize_rejects_non_oauth2_server(self): """mcp_authorize must reject a none-auth server with an accurate 'does not use OAuth' @@ -2163,6 +2625,7 @@ class TestTemporaryMCPSessionEndpoints: code_verifier="verifier", refresh_token=None, scope=None, + client_token_endpoint_auth_method=None, ) @pytest.mark.asyncio @@ -2216,6 +2679,7 @@ class TestTemporaryMCPSessionEndpoints: code_verifier=None, refresh_token="rt-123", scope=None, + client_token_endpoint_auth_method=None, ) @pytest.mark.asyncio @@ -2270,8 +2734,59 @@ class TestTemporaryMCPSessionEndpoints: token_endpoint_auth_method="client_secret_basic", fallback_client_id="server-1", persist_credentials=True, + client_redirect_uris=None, ) + @pytest.mark.asyncio + @pytest.mark.parametrize( + "raw_redirect_uris, forwarded", + [ + (["https://app.example.com/ui/callback"], ["https://app.example.com/ui/callback"]), + (["https://app.example.com/ui/callback", 42, "", None], None), + ("not-a-list", None), + ([], None), + ([123], None), + ], + ) + async def test_mcp_register_forwards_validated_redirect_uris(self, raw_redirect_uris, forwarded): + """dcr_bridge servers relay the registration upstream and require the browser client's own + redirect_uris, so mcp_register must forward them; the value is caller-controlled and is + validated by the same client_supplied_redirect_uris boundary helper as the root /register + door, so a malformed list is rejected whole at both doors (RFC 7591 redirect_uris is + all-or-nothing) rather than silently forwarding the surviving entries here and rejecting + them there.""" + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + mcp_register, + ) + + request = MagicMock() + server = generate_mock_mcp_server_config_record(server_id="server-1") + server.auth_type = MCPAuth.oauth2 + request_body = {"client_name": "LiteLLM", "redirect_uris": raw_redirect_uris} + admin_auth = generate_mock_user_api_key_auth(user_role=LitellmUserRoles.PROXY_ADMIN) + + with ( + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints._get_cached_temporary_mcp_server_or_404", + return_value=server, + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints._read_request_body", + AsyncMock(return_value=request_body), + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.register_client_with_server", + AsyncMock(return_value={"client_id": "generated"}), + ) as register_mock, + ): + await mcp_register( + request=request, + server_id="server-1", + user_api_key_dict=admin_auth, + ) + + assert register_mock.await_args.kwargs["client_redirect_uris"] == forwarded + @pytest.mark.asyncio async def test_mcp_register_does_not_persist_for_non_admin(self): """A non-admin caller (who may have access to a real server) must not persist the DCR diff --git a/tests/test_litellm/proxy/management_endpoints/test_organization_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_organization_endpoints.py index f4470e7e83d..c82c831b3e4 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_organization_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_organization_endpoints.py @@ -10,9 +10,7 @@ import pytest from fastapi import HTTPException from fastapi.testclient import TestClient -sys.path.insert( - 0, os.path.abspath("../../../") -) # Adds the parent directory to the system path +sys.path.insert(0, os.path.abspath("../../../")) # Adds the parent directory to the system path @pytest.mark.asyncio @@ -58,16 +56,12 @@ async def test_organization_update_object_permissions_existing_permission(monkey "vector_stores": ["old_store_1", "old_store_2"], } - mock_prisma_client.db.litellm_objectpermissiontable.find_unique = AsyncMock( - return_value=existing_object_permission - ) + mock_prisma_client.db.litellm_objectpermissiontable.find_unique = AsyncMock(return_value=existing_object_permission) # Mock upsert operation updated_permission = MagicMock() updated_permission.object_permission_id = "existing_perm_id_123" - mock_prisma_client.db.litellm_objectpermissiontable.upsert = AsyncMock( - return_value=updated_permission - ) + mock_prisma_client.db.litellm_objectpermissiontable.upsert = AsyncMock(return_value=updated_permission) # Test data with new object permission data_json = { @@ -107,9 +101,7 @@ async def test_get_organization_daily_activity_admin_param_passing(monkeypatch): # Mock prisma client mock_prisma_client = AsyncMock() - mock_prisma_client.db.litellm_organizationtable.find_many = AsyncMock( - return_value=[] - ) + mock_prisma_client.db.litellm_organizationtable.find_many = AsyncMock(return_value=[]) monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) # Admin view -> skip membership restriction @@ -121,9 +113,7 @@ async def test_get_organization_daily_activity_admin_param_passing(monkeypatch): # Patch downstream common function and verify call args mocked_response = MagicMock(name="SpendAnalyticsPaginatedResponse") get_daily_activity_mock = AsyncMock(return_value=mocked_response) - monkeypatch.setattr( - organization_endpoints, "get_daily_activity", get_daily_activity_mock - ) + monkeypatch.setattr(organization_endpoints, "get_daily_activity", get_daily_activity_mock) auth = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin1") result = await get_organization_daily_activity( @@ -172,17 +162,11 @@ async def test_get_organization_daily_activity_non_admin_defaults_to_admin_orgs( # Mock prisma client and memberships mock_prisma_client = AsyncMock() - mock_prisma_client.db.litellm_organizationtable.find_many = AsyncMock( - return_value=[] - ) + mock_prisma_client.db.litellm_organizationtable.find_many = AsyncMock(return_value=[]) mock_prisma_client.db.litellm_organizationmembership.find_many = AsyncMock( return_value=[ - SimpleNamespace( - organization_id="orgA", user_role=LitellmUserRoles.ORG_ADMIN.value - ), - SimpleNamespace( - organization_id="orgB", user_role=LitellmUserRoles.ORG_ADMIN.value - ), + SimpleNamespace(organization_id="orgA", user_role=LitellmUserRoles.ORG_ADMIN.value), + SimpleNamespace(organization_id="orgB", user_role=LitellmUserRoles.ORG_ADMIN.value), ] ) monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) @@ -196,13 +180,9 @@ async def test_get_organization_daily_activity_non_admin_defaults_to_admin_orgs( # Patch downstream aggregator mocked_response = MagicMock(name="SpendAnalyticsPaginatedResponse") get_daily_activity_mock = AsyncMock(return_value=mocked_response) - monkeypatch.setattr( - organization_endpoints, "get_daily_activity", get_daily_activity_mock - ) + monkeypatch.setattr(organization_endpoints, "get_daily_activity", get_daily_activity_mock) - auth = UserAPIKeyAuth( - user_role=LitellmUserRoles.INTERNAL_USER, user_id="regular-user" - ) + auth = UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER, user_id="regular-user") await get_organization_daily_activity( organization_ids=None, start_date="2024-02-01", @@ -238,15 +218,9 @@ async def test_get_organization_daily_activity_non_admin_unauthorized_org_raises # Mock prisma client and memberships (only orgA is admin) mock_prisma_client = AsyncMock() mock_prisma_client.db.litellm_organizationmembership.find_many = AsyncMock( - return_value=[ - SimpleNamespace( - organization_id="orgA", user_role=LitellmUserRoles.ORG_ADMIN.value - ) - ] - ) - mock_prisma_client.db.litellm_organizationtable.find_many = AsyncMock( - return_value=[] + return_value=[SimpleNamespace(organization_id="orgA", user_role=LitellmUserRoles.ORG_ADMIN.value)] ) + mock_prisma_client.db.litellm_organizationtable.find_many = AsyncMock(return_value=[]) monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) # Non-admin view @@ -255,9 +229,7 @@ async def test_get_organization_daily_activity_non_admin_unauthorized_org_raises lambda _: False, ) - auth = UserAPIKeyAuth( - user_role=LitellmUserRoles.INTERNAL_USER, user_id="regular-user" - ) + auth = UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER, user_id="regular-user") with pytest.raises(HTTPException) as exc: await get_organization_daily_activity( @@ -312,21 +284,17 @@ async def test_organization_update_object_permissions_no_existing_permission( ) # Mock find_unique to return None (no existing permission) - mock_prisma_client.db.litellm_objectpermissiontable.find_unique = AsyncMock( - return_value=None - ) + mock_prisma_client.db.litellm_objectpermissiontable.find_unique = AsyncMock(return_value=None) # Mock upsert to create new record new_permission = MagicMock() new_permission.object_permission_id = "new_perm_id_456" - mock_prisma_client.db.litellm_objectpermissiontable.upsert = AsyncMock( - return_value=new_permission - ) + mock_prisma_client.db.litellm_objectpermissiontable.upsert = AsyncMock(return_value=new_permission) data_json = { - "object_permission": LiteLLM_ObjectPermissionBase( - vector_stores=["brand_new_store"] - ).model_dump(exclude_unset=True, exclude_none=True), + "object_permission": LiteLLM_ObjectPermissionBase(vector_stores=["brand_new_store"]).model_dump( + exclude_unset=True, exclude_none=True + ), "organization_alias": "updated_org_2", } @@ -381,21 +349,17 @@ async def test_organization_update_object_permissions_missing_permission_record( ) # Mock find_unique to return None (permission record not found) - mock_prisma_client.db.litellm_objectpermissiontable.find_unique = AsyncMock( - return_value=None - ) + mock_prisma_client.db.litellm_objectpermissiontable.find_unique = AsyncMock(return_value=None) # Mock upsert to create new record new_permission = MagicMock() new_permission.object_permission_id = "recreated_perm_id_789" - mock_prisma_client.db.litellm_objectpermissiontable.upsert = AsyncMock( - return_value=new_permission - ) + mock_prisma_client.db.litellm_objectpermissiontable.upsert = AsyncMock(return_value=new_permission) data_json = { - "object_permission": LiteLLM_ObjectPermissionBase( - vector_stores=["recreated_store"] - ).model_dump(exclude_unset=True, exclude_none=True), + "object_permission": LiteLLM_ObjectPermissionBase(vector_stores=["recreated_store"]).model_dump( + exclude_unset=True, exclude_none=True + ), "organization_alias": "updated_org_3", } @@ -446,18 +410,14 @@ async def test_list_organization_filter_by_org_id(monkeypatch): ) # Mock find_many to return filtered results - mock_prisma_client.db.litellm_organizationtable.find_many = AsyncMock( - return_value=[mock_org1] - ) + mock_prisma_client.db.litellm_organizationtable.find_many = AsyncMock(return_value=[mock_org1]) monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) # Test as proxy admin auth = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin-user") - result = await list_organization( - org_id="org-123", org_alias=None, user_api_key_dict=auth - ) + result = await list_organization(org_id="org-123", org_alias=None, user_api_key_dict=auth) # Verify the correct organization was returned assert len(result) == 1 @@ -512,18 +472,14 @@ async def test_list_organization_filter_by_org_alias(monkeypatch): ) # Mock find_many to return filtered results - mock_prisma_client.db.litellm_organizationtable.find_many = AsyncMock( - return_value=[mock_org1, mock_org2] - ) + mock_prisma_client.db.litellm_organizationtable.find_many = AsyncMock(return_value=[mock_org1, mock_org2]) monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) # Test as proxy admin with org_alias filter auth = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin-user") - result = await list_organization( - org_id=None, org_alias="test", user_api_key_dict=auth - ) + result = await list_organization(org_id=None, org_alias="test", user_api_key_dict=auth) # Verify organizations with "test" in alias were returned assert len(result) == 2 @@ -532,9 +488,7 @@ async def test_list_organization_filter_by_org_alias(monkeypatch): # Verify find_many was called with correct where conditions (case-insensitive contains) mock_prisma_client.db.litellm_organizationtable.find_many.assert_called_once() call_args = mock_prisma_client.db.litellm_organizationtable.find_many.call_args - assert call_args.kwargs["where"] == { - "organization_alias": {"contains": "test", "mode": "insensitive"} - } + assert call_args.kwargs["where"] == {"organization_alias": {"contains": "test", "mode": "insensitive"}} assert call_args.kwargs["include"] == { "litellm_budget_table": True, "members": True, @@ -612,16 +566,12 @@ def patched_org_prisma(): ), patch("litellm.proxy.proxy_server.proxy_logging_obj"), ): - mock_prisma.db.litellm_organizationtable.find_unique = AsyncMock( - return_value=victim_row - ) + mock_prisma.db.litellm_organizationtable.find_unique = AsyncMock(return_value=victim_row) yield mock_prisma @pytest.mark.asyncio -async def test_organization_member_add_rejects_unauthorized_caller( - patched_org_prisma, unauthorized_caller -): +async def test_organization_member_add_rejects_unauthorized_caller(patched_org_prisma, unauthorized_caller): # ``organization_member_add`` catches HTTPException in its # catch-all and re-wraps as ProxyException with the original status # code preserved. @@ -653,9 +603,7 @@ async def test_organization_member_add_rejects_unauthorized_caller( @pytest.mark.asyncio -async def test_organization_member_update_rejects_unauthorized_caller( - patched_org_prisma, unauthorized_caller -): +async def test_organization_member_update_rejects_unauthorized_caller(patched_org_prisma, unauthorized_caller): from litellm.proxy._types import OrganizationMemberUpdateRequest from litellm.proxy.management_endpoints.organization_endpoints import ( organization_member_update, @@ -676,9 +624,7 @@ async def test_organization_member_update_rejects_unauthorized_caller( @pytest.mark.asyncio -async def test_organization_member_delete_rejects_unauthorized_caller( - patched_org_prisma, unauthorized_caller -): +async def test_organization_member_delete_rejects_unauthorized_caller(patched_org_prisma, unauthorized_caller): from litellm.proxy._types import OrganizationMemberDeleteRequest from litellm.proxy.management_endpoints.organization_endpoints import ( organization_member_delete, @@ -695,3 +641,398 @@ async def test_organization_member_delete_rejects_unauthorized_caller( user_api_key_dict=unauthorized_caller, ) assert exc.value.status_code == 403 + + +@pytest.mark.parametrize( + "body", + [{"tpm_limit": ""}, {"tmp_limit": None}], + ids=["non-numeric-limit", "unknown-key"], +) +def test_v2_model_rejects_invalid_body(body): + """A non-numeric limit and an unknown/misspelled key are both rejected at model validation (422 at the route).""" + from pydantic import ValidationError + + from litellm.proxy._types import OrganizationUpdateRequestV2 + + with pytest.raises(ValidationError): + OrganizationUpdateRequestV2.model_validate(body) + + +class _FakeTxContext: + def __init__(self, tx): + self._tx = tx + + async def __aenter__(self): + return self._tx + + async def __aexit__(self, exc_type, exc, tb): + return False + + +async def _run_update_organization_v2( + monkeypatch, + *, + body: dict, + existing_budget_id, + existing_metadata, + existing_object_permission_id=None, + existing_object_permission_row=None, +): + from litellm.proxy._types import ( + LitellmUserRoles, + OrganizationUpdateRequestV2, + UserAPIKeyAuth, + ) + from litellm.proxy.management_endpoints import organization_endpoints + from litellm.proxy.management_endpoints.organization_endpoints import ( + update_organization_v2, + ) + from litellm.proxy.utils import jsonify_object + + mock_prisma_client = AsyncMock() + mock_prisma_client.jsonify_object = jsonify_object + + existing_org = MagicMock() + existing_org.budget_id = existing_budget_id + existing_org.object_permission_id = existing_object_permission_id + existing_org.metadata = existing_metadata + + mock_prisma_client.db.litellm_organizationtable.find_unique = AsyncMock(return_value=existing_org) + mock_prisma_client.db.litellm_organizationtable.update = AsyncMock(return_value=MagicMock()) + mock_prisma_client.db.litellm_budgettable.update = AsyncMock() + mock_prisma_client.db.litellm_objectpermissiontable.find_unique = AsyncMock( + return_value=existing_object_permission_row + ) + mock_prisma_client.db.litellm_objectpermissiontable.upsert = AsyncMock() + + tx = MagicMock() + tx.litellm_organizationtable = mock_prisma_client.db.litellm_organizationtable + tx.litellm_budgettable = mock_prisma_client.db.litellm_budgettable + tx.litellm_objectpermissiontable.upsert = AsyncMock() + mock_prisma_client.db.tx = MagicMock(return_value=_FakeTxContext(tx)) + mock_prisma_client.tx = tx + + call_order = MagicMock() + call_order.attach_mock(tx.litellm_objectpermissiontable.upsert, "permission_upsert") + call_order.attach_mock(mock_prisma_client.db.litellm_organizationtable.update, "org_update") + mock_prisma_client.call_order = call_order + + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + monkeypatch.setattr(organization_endpoints, "_verify_org_access", AsyncMock()) + + auth = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin-1") + await update_organization_v2( + organization_id="org-1", + data=OrganizationUpdateRequestV2.model_validate(body), + user_api_key_dict=auth, + ) + return mock_prisma_client + + +@pytest.mark.asyncio +async def test_v2_update_clears_tpm_limit_and_metadata(monkeypatch): + """A cleared tpm_limit is written to the budget row as None; a cleared metadata is written as {}.""" + prisma = await _run_update_organization_v2( + monkeypatch, + body={"tpm_limit": None, "metadata": None}, + existing_budget_id="budget-1", + existing_metadata={"stale": "value"}, + ) + + budget_write = prisma.db.litellm_budgettable.update.await_args + assert budget_write.kwargs["where"] == {"budget_id": "budget-1"} + assert budget_write.kwargs["data"]["tpm_limit"] is None + assert "soft_budget" not in budget_write.kwargs["data"] + + write_data = prisma.db.litellm_organizationtable.update.await_args.kwargs["data"] + assert json.loads(write_data["metadata"]) == {} + assert "budget_id" not in write_data + + +@pytest.mark.asyncio +async def test_v2_update_untouched_fields_not_written(monkeypatch): + """Omitted fields are left untouched: only organization_alias is written, no budget-row write.""" + prisma = await _run_update_organization_v2( + monkeypatch, + body={"organization_alias": "renamed"}, + existing_budget_id="budget-1", + existing_metadata={"keep": "me"}, + ) + + prisma.db.litellm_budgettable.update.assert_not_awaited() + write_data = prisma.db.litellm_organizationtable.update.await_args.kwargs["data"] + assert write_data["organization_alias"] == "renamed" + assert "metadata" not in write_data + assert "tpm_limit" not in write_data + + +@pytest.mark.asyncio +async def test_v2_update_metadata_replaces_not_merges(monkeypatch): + """Sending metadata replaces the stored blob wholesale; a previously-present key is gone.""" + prisma = await _run_update_organization_v2( + monkeypatch, + body={"metadata": {"a": 1}}, + existing_budget_id="budget-1", + existing_metadata={"stale": "value"}, + ) + write_data = prisma.db.litellm_organizationtable.update.await_args.kwargs["data"] + assert json.loads(write_data["metadata"]) == {"a": 1} + + +@pytest.mark.asyncio +async def test_v2_update_writes_logging_exporters_to_org_column(monkeypatch): + """Assigning logging_exporters writes the credential names to the org column, not the budget row or metadata.""" + prisma = await _run_update_organization_v2( + monkeypatch, + body={"logging_exporters": ["arize-prod", "langfuse-eu"]}, + existing_budget_id="budget-1", + existing_metadata={"keep": "me"}, + ) + + prisma.db.litellm_budgettable.update.assert_not_awaited() + write_data = prisma.db.litellm_organizationtable.update.await_args.kwargs["data"] + assert write_data["logging_exporters"] == ["arize-prod", "langfuse-eu"] + assert "metadata" not in write_data + + +@pytest.mark.asyncio +async def test_v2_update_clears_logging_exporters_with_empty_list(monkeypatch): + """A null logging_exporters clears the org's assignments by writing an empty list to the non-nullable column.""" + prisma = await _run_update_organization_v2( + monkeypatch, + body={"logging_exporters": None}, + existing_budget_id="budget-1", + existing_metadata={"keep": "me"}, + ) + + write_data = prisma.db.litellm_organizationtable.update.await_args.kwargs["data"] + assert write_data["logging_exporters"] == [] + + +@pytest.mark.asyncio +async def test_v2_update_omitted_logging_exporters_not_written(monkeypatch): + """Omitting logging_exporters leaves the existing assignments untouched.""" + prisma = await _run_update_organization_v2( + monkeypatch, + body={"organization_alias": "renamed"}, + existing_budget_id="budget-1", + existing_metadata={"keep": "me"}, + ) + + write_data = prisma.db.litellm_organizationtable.update.await_args.kwargs["data"] + assert "logging_exporters" not in write_data + + +@pytest.mark.asyncio +async def test_v2_rejects_null_clear_of_non_nullable_fields(monkeypatch): + """organization_alias and models are non-nullable columns, so a null clear is a 422, not a 500.""" + from litellm.proxy._types import LitellmUserRoles, OrganizationUpdateRequestV2, UserAPIKeyAuth + from litellm.proxy.management_endpoints.organization_endpoints import update_organization_v2 + + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", AsyncMock()) + auth = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin-1") + + for body in ({"organization_alias": None}, {"models": None}): + with pytest.raises(HTTPException) as exc: + await update_organization_v2( + organization_id="org-1", + data=OrganizationUpdateRequestV2.model_validate(body), + user_api_key_dict=auth, + ) + assert exc.value.status_code == 422 + + +@pytest.mark.asyncio +async def test_v2_rejects_negative_max_budget(monkeypatch): + """v2 rejects a negative max_budget with a 422 before touching the DB.""" + from litellm.proxy._types import LitellmUserRoles, OrganizationUpdateRequestV2, UserAPIKeyAuth + from litellm.proxy.management_endpoints.organization_endpoints import update_organization_v2 + + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", AsyncMock()) + + auth = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin-1") + with pytest.raises(HTTPException) as exc: + await update_organization_v2( + organization_id="org-1", + data=OrganizationUpdateRequestV2.model_validate({"max_budget": -5}), + user_api_key_dict=auth, + ) + assert exc.value.status_code == 422 + assert "max_budget" in str(exc.value.detail) + + +@pytest.mark.asyncio +async def test_v2_rejects_caller_without_org_access(monkeypatch): + """v2 runs the real _verify_org_access guard: a non-admin without ORG_ADMIN on the org gets 403 and no write.""" + from litellm.proxy._types import LitellmUserRoles, OrganizationUpdateRequestV2, UserAPIKeyAuth + from litellm.proxy.management_endpoints import organization_endpoints + from litellm.proxy.management_endpoints.organization_endpoints import update_organization_v2 + + mock_prisma_client = AsyncMock() + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + monkeypatch.setattr(organization_endpoints, "_user_has_admin_view", lambda _: False) + + caller = MagicMock() + caller.organization_memberships = [] + monkeypatch.setattr(organization_endpoints, "get_user_object", AsyncMock(return_value=caller)) + + auth = UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER, user_id="user-1") + with pytest.raises(HTTPException) as exc: + await update_organization_v2( + organization_id="org-1", + data=OrganizationUpdateRequestV2.model_validate({"tpm_limit": 5}), + user_api_key_dict=auth, + ) + assert exc.value.status_code == 403 + mock_prisma_client.db.litellm_organizationtable.update.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_v2_wires_object_permission_onto_org_write(monkeypatch): + """A sent object_permission merges over the existing permission row and its id is linked onto the org write.""" + existing_row = MagicMock() + existing_row.model_dump.return_value = { + "object_permission_id": "op-123", + "mcp_servers": ["server-1"], + } + + prisma = await _run_update_organization_v2( + monkeypatch, + body={"object_permission": {"vector_stores": ["vs-1"]}}, + existing_budget_id="budget-1", + existing_metadata={}, + existing_object_permission_id="op-123", + existing_object_permission_row=existing_row, + ) + + upsert = prisma.tx.litellm_objectpermissiontable.upsert.await_args.kwargs + assert upsert["where"] == {"object_permission_id": "op-123"} + assert upsert["data"]["update"]["mcp_servers"] == ["server-1"] + assert upsert["data"]["update"]["vector_stores"] == ["vs-1"] + write_data = prisma.db.litellm_organizationtable.update.await_args.kwargs["data"] + assert write_data["object_permission_id"] == "op-123" + + +@pytest.mark.asyncio +async def test_v2_object_permission_upsert_runs_inside_transaction(monkeypatch): + """The permission upsert runs on the tx client, before the org write that links it, so a rollback cannot + leave merged grants live on a row the org still points at.""" + prisma = await _run_update_organization_v2( + monkeypatch, + body={"object_permission": {"vector_stores": ["vs-1"]}}, + existing_budget_id="budget-1", + existing_metadata={}, + ) + + prisma.tx.litellm_objectpermissiontable.upsert.assert_awaited_once() + prisma.db.litellm_objectpermissiontable.upsert.assert_not_awaited() + + upsert = prisma.tx.litellm_objectpermissiontable.upsert.await_args.kwargs + linked_id = prisma.db.litellm_organizationtable.update.await_args.kwargs["data"]["object_permission_id"] + assert upsert["where"] == {"object_permission_id": linked_id} + assert upsert["data"]["create"]["object_permission_id"] == linked_id + + ordered = [name for name, _, _ in prisma.call_order.mock_calls if name in ("permission_upsert", "org_update")] + assert ordered == ["permission_upsert", "org_update"] + + +@pytest.mark.asyncio +async def test_v2_clears_object_permission_when_sent_null(monkeypatch): + """object_permission: null detaches the org's permission row (object_permission_id -> None), no merge.""" + prisma = await _run_update_organization_v2( + monkeypatch, + body={"object_permission": None}, + existing_budget_id="budget-1", + existing_metadata={}, + ) + + prisma.tx.litellm_objectpermissiontable.upsert.assert_not_awaited() + prisma.db.litellm_objectpermissiontable.find_unique.assert_not_awaited() + write_data = prisma.db.litellm_organizationtable.update.await_args.kwargs["data"] + assert write_data["object_permission_id"] is None + + +@pytest.mark.asyncio +async def test_v2_rejects_empty_object_permission(monkeypatch): + """object_permission: {} merges nothing, so it is rejected (send null to clear) rather than silently leaving grants.""" + from litellm.proxy._types import LitellmUserRoles, OrganizationUpdateRequestV2, UserAPIKeyAuth + from litellm.proxy.management_endpoints import organization_endpoints + from litellm.proxy.management_endpoints.organization_endpoints import update_organization_v2 + + mock_prisma_client = AsyncMock() + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + monkeypatch.setattr(organization_endpoints, "_verify_org_access", AsyncMock()) + + auth = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin-1") + with pytest.raises(HTTPException) as exc: + await update_organization_v2( + organization_id="org-1", + data=OrganizationUpdateRequestV2.model_validate({"object_permission": {}}), + user_api_key_dict=auth, + ) + assert exc.value.status_code == 422 + assert "object_permission" in str(exc.value.detail) + mock_prisma_client.db.litellm_organizationtable.update.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_v2_writes_budget_and_org_in_one_transaction(monkeypatch): + """A change touching both the budget row and the org row runs both writes inside one prisma transaction.""" + prisma = await _run_update_organization_v2( + monkeypatch, + body={"tpm_limit": 500, "metadata": {"a": 1}}, + existing_budget_id="budget-1", + existing_metadata={}, + ) + + prisma.db.tx.assert_called_once() + prisma.db.litellm_budgettable.update.assert_awaited_once() + prisma.db.litellm_organizationtable.update.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_v2_serializes_model_max_budget_on_budget_write(monkeypatch): + """model_max_budget is a Json column, so it is JSON-serialized on the budget-row write like new_budget/metadata.""" + monkeypatch.setattr( + "litellm.proxy.management_endpoints.key_management_endpoints.validate_model_max_budget", + lambda _: None, + ) + + prisma = await _run_update_organization_v2( + monkeypatch, + body={"model_max_budget": {"gpt-4o": {"max_budget": 10}}}, + existing_budget_id="budget-1", + existing_metadata={}, + ) + + written = prisma.db.litellm_budgettable.update.await_args.kwargs["data"]["model_max_budget"] + assert isinstance(written, str) + assert json.loads(written) == {"gpt-4o": {"max_budget": 10}} + + +def test_build_budget_write_data_recomputes_reset_at_on_duration(): + """A sent budget_duration recomputes budget_reset_at so the reset window follows the new duration.""" + from litellm.proxy.management_endpoints.organization_endpoints import build_budget_write_data + + data = build_budget_write_data({"budget_duration": "30d"}, "admin-1") + assert data["budget_duration"] == "30d" + assert "budget_reset_at" in data + assert data["updated_by"] == "admin-1" + + +def test_build_budget_write_data_no_reset_at_without_duration(): + """Clearing a limit writes it through untouched and does not recompute budget_reset_at.""" + from litellm.proxy.management_endpoints.organization_endpoints import build_budget_write_data + + data = build_budget_write_data({"tpm_limit": None}, "admin-1") + assert data["tpm_limit"] is None + assert "budget_reset_at" not in data + + +def test_build_budget_write_data_clears_reset_at_with_null_duration(): + """Clearing budget_duration also nulls budget_reset_at so no stale reset timestamp survives.""" + from litellm.proxy.management_endpoints.organization_endpoints import build_budget_write_data + + data = build_budget_write_data({"budget_duration": None}, "admin-1") + assert data["budget_duration"] is None + assert data["budget_reset_at"] is None diff --git a/tests/test_litellm/proxy/management_endpoints/test_saml_sso.py b/tests/test_litellm/proxy/management_endpoints/test_saml_sso.py new file mode 100644 index 00000000000..57decc7d458 --- /dev/null +++ b/tests/test_litellm/proxy/management_endpoints/test_saml_sso.py @@ -0,0 +1,644 @@ +""" +Regression tests for SAML 2.0 SSO (SP- and IdP-initiated) on the admin UI. + +These exercise the real OneLogin python3-saml validation by generating signed +SAML responses with a freshly minted IdP keypair, so a mutation that weakens +signature, signing-requirement, expiry, replay or attribute-mapping handling +makes a test fail. +""" + +import base64 +import datetime +import time + +import pytest +from fastapi import HTTPException, Request + +pytest.importorskip( + "onelogin", reason="python3-saml (saml extra) is required for SAML SSO tests" +) + +from cryptography import x509 +from cryptography.hazmat.primitives import hashes, serialization +from cryptography.hazmat.primitives.asymmetric import rsa +from cryptography.x509.oid import NameOID +from onelogin.saml2.utils import OneLogin_Saml2_Utils +from starlette.datastructures import URL + +from typing import cast + +from litellm.caching.dual_cache import DualCache +from litellm.caching.in_memory_cache import InMemoryCache +from litellm.caching.redis_cache import RedisCache +from litellm.proxy._types import LitellmUserRoles +from litellm.proxy.management_endpoints.sso.saml_sso import ( + _SAML_AUTHN_REQUEST_CACHE_PREFIX, + _SAML_AUTHN_STATE_COOKIE, + _SAML_MAX_POST_BYTES, + _SAML_REPLAY_GUARD_DEFAULT_TTL_SECONDS, + _SAML_REPLAY_GUARD_MAX_TTL_SECONDS, + SAMLAuthHandler, +) + + +def _shared_cache(store=None): + """A DualCache whose replay guard is backed by a shared, atomic store. + + An InMemoryCache instance stands in for Redis; passing the same instance to + two DualCaches simulates two workers sharing one atomic backend.""" + return DualCache(redis_cache=cast(RedisCache, store or InMemoryCache())) + +IDP_ENTITY = "https://idp.example.com/metadata" +SP_ENTITY = "https://proxy.example.com/sso/saml/metadata" +ACS = "https://proxy.example.com/sso/saml/callback" +SSO_URL = "https://idp.example.com/sso" +PROXY_BASE_URL = "https://proxy.example.com" + + +def _make_idp_keypair(): + key = rsa.generate_private_key(public_exponent=65537, key_size=2048) + name = x509.Name([x509.NameAttribute(NameOID.COMMON_NAME, "idp.example.com")]) + cert = ( + x509.CertificateBuilder() + .subject_name(name) + .issuer_name(name) + .public_key(key.public_key()) + .serial_number(x509.random_serial_number()) + .not_valid_before(datetime.datetime.utcnow() - datetime.timedelta(days=1)) + .not_valid_after(datetime.datetime.utcnow() + datetime.timedelta(days=365)) + .sign(key, hashes.SHA256()) + ) + key_pem = key.private_bytes( + serialization.Encoding.PEM, + serialization.PrivateFormat.TraditionalOpenSSL, + serialization.NoEncryption(), + ).decode() + cert_pem = cert.public_bytes(serialization.Encoding.PEM).decode() + return key_pem, cert_pem + + +def _idp_metadata_xml(cert_pem): + cert_body = "".join( + line for line in cert_pem.splitlines() if "CERTIFICATE" not in line + ) + return ( + '' + f'' + '' + '' + f"{cert_body}" + "" + '' + "" + ) + + +def _saml_time(delta_seconds): + t = datetime.datetime.utcnow() + datetime.timedelta(seconds=delta_seconds) + return t.strftime("%Y-%m-%dT%H:%M:%SZ") + + +def _build_signed_response( + key_pem, + cert_pem, + *, + in_response_to=None, + response_level_in_response_to=True, + email="alice@example.com", + attributes=None, + not_before_delta=-60, + not_on_or_after_delta=300, + sign=True, +): + if attributes is None: + attributes = { + "email": [email], + "givenName": ["Alice"], + "sn": ["Smith"], + "role": ["internal_user"], + } + assertion_id = "_assertion_" + OneLogin_Saml2_Utils.generate_unique_id() + response_id = "_response_" + OneLogin_Saml2_Utils.generate_unique_id() + not_before = _saml_time(not_before_delta) + not_on_or_after = _saml_time(not_on_or_after_delta) + issue_instant = _saml_time(-1) + irt = f'InResponseTo="{in_response_to}"' if in_response_to else "" + response_irt = irt if response_level_in_response_to else "" + + attr_xml = "".join( + f'' + + "".join(f"{v}" for v in values) + + "" + for name, values in attributes.items() + ) + + assertion = ( + '' + f"{IDP_ENTITY}" + "" + '' + f"{email}" + '' + f'' + "" + f'' + f"{SP_ENTITY}" + "" + f'' + "" + "urn:oasis:names:tc:SAML:2.0:ac:classes:Password" + "" + f"{attr_xml}" + "" + ) + + if sign: + signed = OneLogin_Saml2_Utils.add_sign(assertion, key_pem, cert_pem) + assertion = (signed.decode() if isinstance(signed, bytes) else signed).replace( + '', "" + ) + + return ( + '' + '' + f"{IDP_ENTITY}" + '' + "" + f"{assertion}" + ) + + +def _b64(xml): + return base64.b64encode(xml.encode()).decode() + + +def _fake_request(cookies=None): + return type( + "Req", + (), + { + "base_url": URL(PROXY_BASE_URL + "/"), + "query_params": {}, + "cookies": cookies or {}, + }, + )() + + +async def _acs(b64, cache, cookies=None): + return await SAMLAuthHandler.handle_acs( + _fake_request(cookies), cache, {"SAMLResponse": b64} + ) + + +@pytest.fixture +def saml_env(monkeypatch): + key_pem, cert_pem = _make_idp_keypair() + monkeypatch.setenv("SAML_IDP_METADATA_XML", _idp_metadata_xml(cert_pem)) + monkeypatch.setenv("SAML_SP_ENTITY_ID", SP_ENTITY) + monkeypatch.setenv("PROXY_BASE_URL", PROXY_BASE_URL) + for var in ( + "SAML_IDP_METADATA_URL", + "SAML_ATTRIBUTE_EMAIL", + "SAML_ATTRIBUTE_TEAM_IDS", + "SAML_ALLOW_UNSOLICITED", + "ALLOWED_EMAIL_DOMAINS", + ): + monkeypatch.delenv(var, raising=False) + return key_pem, cert_pem + + +@pytest.fixture +def saml_env_idp_initiated(saml_env, monkeypatch): + monkeypatch.setenv("SAML_ALLOW_UNSOLICITED", "true") + return saml_env + + +@pytest.mark.asyncio +async def test_valid_idp_initiated_login_maps_assertion_to_user(saml_env_idp_initiated): + key_pem, cert_pem = saml_env_idp_initiated + resp = _build_signed_response(key_pem, cert_pem) + + result = await _acs(_b64(resp), _shared_cache()) + + assert result.email == "alice@example.com" + assert result.id == "alice@example.com" + assert result.first_name == "Alice" + assert result.last_name == "Smith" + assert result.user_role == LitellmUserRoles.INTERNAL_USER + assert result.provider == "saml" + + +@pytest.mark.asyncio +async def test_tampered_assertion_is_rejected(saml_env): + key_pem, cert_pem = saml_env + resp = _build_signed_response(key_pem, cert_pem) + tampered = resp.replace("alice@example.com", "attacker@example.com") + + with pytest.raises(HTTPException) as exc: + await _acs(_b64(tampered), DualCache()) + assert exc.value.status_code == 401 + + +@pytest.mark.asyncio +async def test_unsigned_assertion_is_rejected(saml_env): + key_pem, cert_pem = saml_env + resp = _build_signed_response(key_pem, cert_pem, sign=False) + + with pytest.raises(HTTPException) as exc: + await _acs(_b64(resp), DualCache()) + assert exc.value.status_code == 401 + + +@pytest.mark.asyncio +async def test_signature_from_untrusted_key_is_rejected(saml_env): + _, cert_pem = saml_env + attacker_key, attacker_cert = _make_idp_keypair() + resp = _build_signed_response(attacker_key, attacker_cert) + + with pytest.raises(HTTPException) as exc: + await _acs(_b64(resp), DualCache()) + assert exc.value.status_code == 401 + + +@pytest.mark.asyncio +async def test_expired_assertion_is_rejected(saml_env): + key_pem, cert_pem = saml_env + resp = _build_signed_response( + key_pem, cert_pem, not_before_delta=-7200, not_on_or_after_delta=-3600 + ) + + with pytest.raises(HTTPException) as exc: + await _acs(_b64(resp), DualCache()) + assert exc.value.status_code == 401 + + +@pytest.mark.asyncio +async def test_sp_initiated_unknown_in_response_to_is_rejected(saml_env): + key_pem, cert_pem = saml_env + resp = _build_signed_response(key_pem, cert_pem, in_response_to="_never_issued") + + with pytest.raises(HTTPException) as exc: + await _acs(_b64(resp), DualCache()) + assert exc.value.status_code == 401 + + +@pytest.mark.asyncio +async def test_sp_initiated_known_request_succeeds_once_then_replay_rejected(saml_env): + key_pem, cert_pem = saml_env + cache = DualCache() + request_id = "_authn_req_known" + cache.set_cache( + key=f"{_SAML_AUTHN_REQUEST_CACHE_PREFIX}:{request_id}", value="1", ttl=600 + ) + resp = _build_signed_response(key_pem, cert_pem, in_response_to=request_id) + cookies = {_SAML_AUTHN_STATE_COOKIE: request_id} + + result = await _acs(_b64(resp), cache, cookies=cookies) + assert result.email == "alice@example.com" + + with pytest.raises(HTTPException) as exc: + await _acs(_b64(resp), cache, cookies=cookies) + assert exc.value.status_code == 401 + + +@pytest.mark.asyncio +async def test_sp_initiated_response_not_bound_to_browser_is_rejected(saml_env): + key_pem, cert_pem = saml_env + cache = DualCache() + request_id = "_authn_req_known" + cache.set_cache( + key=f"{_SAML_AUTHN_REQUEST_CACHE_PREFIX}:{request_id}", value="1", ttl=600 + ) + resp = _build_signed_response(key_pem, cert_pem, in_response_to=request_id) + + with pytest.raises(HTTPException) as exc: + await _acs(_b64(resp), cache) + assert exc.value.status_code == 401 + + with pytest.raises(HTTPException) as exc: + await _acs( + _b64(resp), cache, cookies={_SAML_AUTHN_STATE_COOKIE: "_attacker_request"} + ) + assert exc.value.status_code == 401 + + +@pytest.mark.asyncio +async def test_subjectconfirmation_only_in_response_to_without_cookie_is_rejected( + saml_env_idp_initiated, +): + """An IdP that stamps InResponseTo only on the SubjectConfirmationData (not the + Response element) is still solicited and must be browser-bound: with unsolicited + explicitly allowed, a missing cookie must still 401 rather than slip through.""" + key_pem, cert_pem = saml_env_idp_initiated + cache = DualCache() + request_id = "_authn_req_known" + cache.set_cache( + key=f"{_SAML_AUTHN_REQUEST_CACHE_PREFIX}:{request_id}", value="1", ttl=600 + ) + resp = _build_signed_response( + key_pem, + cert_pem, + in_response_to=request_id, + response_level_in_response_to=False, + ) + + with pytest.raises(HTTPException) as exc: + await _acs(_b64(resp), cache) + assert exc.value.status_code == 401 + + +@pytest.mark.asyncio +async def test_subjectconfirmation_only_in_response_to_with_cookie_succeeds(saml_env): + key_pem, cert_pem = saml_env + cache = DualCache() + request_id = "_authn_req_known" + cache.set_cache( + key=f"{_SAML_AUTHN_REQUEST_CACHE_PREFIX}:{request_id}", value="1", ttl=600 + ) + resp = _build_signed_response( + key_pem, + cert_pem, + in_response_to=request_id, + response_level_in_response_to=False, + ) + + result = await _acs( + _b64(resp), cache, cookies={_SAML_AUTHN_STATE_COOKIE: request_id} + ) + assert result.email == "alice@example.com" + + +@pytest.mark.asyncio +async def test_unsolicited_response_rejected_by_default(saml_env): + key_pem, cert_pem = saml_env + resp = _build_signed_response(key_pem, cert_pem) + + with pytest.raises(HTTPException) as exc: + await _acs(_b64(resp), DualCache()) + assert exc.value.status_code == 401 + + +@pytest.mark.asyncio +async def test_idp_initiated_assertion_replay_is_rejected(saml_env_idp_initiated): + key_pem, cert_pem = saml_env_idp_initiated + cache = _shared_cache() + resp = _build_signed_response(key_pem, cert_pem, email="bob@example.com") + + first = await _acs(_b64(resp), cache) + assert first.email == "bob@example.com" + + with pytest.raises(HTTPException) as exc: + await _acs(_b64(resp), cache) + assert exc.value.status_code == 401 + + +@pytest.mark.asyncio +async def test_assertion_without_id_is_rejected(saml_env_idp_initiated): + """An assertion with no ID attribute has no stable replay key. On the unsolicited + path there is no browser binding, so the consumed-assertion guard is the only replay + defense; a missing ID must be rejected rather than silently skipping the guard.""" + + class _AuthNoAssertionId: + def get_last_response_in_response_to(self): + return None + + def get_last_response_xml(self): + return None + + def get_last_assertion_id(self): + return None + + with pytest.raises(HTTPException) as exc: + await SAMLAuthHandler._enforce_response_binding( + _AuthNoAssertionId(), _shared_cache(), None + ) + assert exc.value.status_code == 401 + assert "ID" in exc.value.detail + + +@pytest.mark.asyncio +async def test_unsolicited_response_rejected_when_disabled(saml_env, monkeypatch): + key_pem, cert_pem = saml_env + monkeypatch.setenv("SAML_ALLOW_UNSOLICITED", "false") + resp = _build_signed_response(key_pem, cert_pem) + + with pytest.raises(HTTPException) as exc: + await _acs(_b64(resp), DualCache()) + assert exc.value.status_code == 401 + + +@pytest.mark.asyncio +async def test_invalid_email_in_assertion_is_rejected_cleanly(saml_env_idp_initiated): + key_pem, cert_pem = saml_env_idp_initiated + resp = _build_signed_response( + key_pem, + cert_pem, + email="not-an-email", + attributes={"email": ["not-an-email"], "givenName": ["X"]}, + ) + + with pytest.raises(HTTPException) as exc: + await _acs(_b64(resp), _shared_cache()) + assert exc.value.status_code == 401 + assert "invalid subject or email" in exc.value.detail + + +@pytest.mark.asyncio +async def test_email_less_assertion_rejected_when_domain_restriction_configured( + saml_env_idp_initiated, monkeypatch +): + key_pem, cert_pem = saml_env_idp_initiated + monkeypatch.setenv("ALLOWED_EMAIL_DOMAINS", "example.com") + resp = _build_signed_response( + key_pem, + cert_pem, + email="opaque-persistent-id-123", + attributes={"givenName": ["Alice"]}, + ) + + with pytest.raises(HTTPException) as exc: + await _acs(_b64(resp), _shared_cache()) + assert exc.value.status_code == 401 + assert "ALLOWED_EMAIL_DOMAINS" in exc.value.detail + + +@pytest.mark.asyncio +async def test_email_less_assertion_allowed_without_domain_restriction(saml_env_idp_initiated): + key_pem, cert_pem = saml_env_idp_initiated + resp = _build_signed_response( + key_pem, + cert_pem, + email="opaque-persistent-id-123", + attributes={"givenName": ["Alice"]}, + ) + + result = await _acs(_b64(resp), _shared_cache()) + assert result.email is None + assert result.id == "opaque-persistent-id-123" + + +@pytest.mark.asyncio +async def test_custom_email_attribute_override(saml_env_idp_initiated, monkeypatch): + key_pem, cert_pem = saml_env_idp_initiated + monkeypatch.setenv("SAML_ATTRIBUTE_EMAIL", "corpMail") + resp = _build_signed_response( + key_pem, + cert_pem, + email="ignored@example.com", + attributes={ + "corpMail": ["real@corp.example.com"], + "givenName": ["Real"], + }, + ) + + result = await _acs(_b64(resp), _shared_cache()) + assert result.email == "real@corp.example.com" + + +@pytest.mark.asyncio +async def test_team_ids_extracted_from_groups_attribute(saml_env_idp_initiated): + key_pem, cert_pem = saml_env_idp_initiated + resp = _build_signed_response( + key_pem, + cert_pem, + attributes={ + "email": ["carol@example.com"], + "groups": ["team-a", "team-b"], + }, + ) + + result = await _acs(_b64(resp), _shared_cache()) + assert result.team_ids == ["team-a", "team-b"] + + +@pytest.mark.asyncio +async def test_build_login_redirect_targets_idp_and_caches_request_id(saml_env): + cache = DualCache() + redirect = await SAMLAuthHandler.build_login_redirect(_fake_request(), cache) + + location = redirect.headers["location"] + assert location.startswith(SSO_URL) + assert "SAMLRequest=" in location + cached = [ + k + for k in cache.in_memory_cache.cache_dict + if k.startswith(_SAML_AUTHN_REQUEST_CACHE_PREFIX) + ] + assert len(cached) == 1 + + request_id = cached[0].split(":", 1)[1] + set_cookie = redirect.headers["set-cookie"] + assert f"{_SAML_AUTHN_STATE_COOKIE}={request_id}" in set_cookie + assert "httponly" in set_cookie.lower() + + +@pytest.mark.asyncio +async def test_sp_metadata_contains_acs_and_entity_id(saml_env): + metadata = await SAMLAuthHandler.build_sp_metadata(_fake_request(), DualCache()) + assert ACS in metadata + assert SP_ENTITY in metadata + assert "AssertionConsumerService" in metadata + + +def test_replay_guard_ttl_tracks_assertion_validity(): + class _Auth: + def __init__(self, not_on_or_after): + self._not_on_or_after = not_on_or_after + + def get_last_assertion_not_on_or_after(self): + return self._not_on_or_after + + now = int(time.time()) + + long_lived = SAMLAuthHandler._replay_guard_ttl(_Auth(now + 7200)) + assert long_lived >= 7200 + + short_lived = SAMLAuthHandler._replay_guard_ttl(_Auth(now + 60)) + assert short_lived == _SAML_REPLAY_GUARD_DEFAULT_TTL_SECONDS + + missing = SAMLAuthHandler._replay_guard_ttl(_Auth(None)) + assert missing == _SAML_REPLAY_GUARD_DEFAULT_TTL_SECONDS + + capped = SAMLAuthHandler._replay_guard_ttl(_Auth(now + 10 * 86400)) + assert capped == _SAML_REPLAY_GUARD_MAX_TTL_SECONDS + + +def test_is_saml_configured_reflects_env(monkeypatch): + monkeypatch.delenv("SAML_IDP_METADATA_URL", raising=False) + monkeypatch.delenv("SAML_IDP_METADATA_XML", raising=False) + assert SAMLAuthHandler.is_saml_configured() is False + + monkeypatch.setenv("SAML_IDP_METADATA_URL", "https://idp.example.com/metadata.xml") + assert SAMLAuthHandler.is_saml_configured() is True + + +@pytest.mark.asyncio +async def test_idp_initiated_rejected_without_shared_cache(saml_env_idp_initiated): + key_pem, cert_pem = saml_env_idp_initiated + resp = _build_signed_response(key_pem, cert_pem) + + with pytest.raises(HTTPException) as exc: + await _acs(_b64(resp), DualCache()) + assert exc.value.status_code == 401 + assert "shared Redis cache" in exc.value.detail + + +@pytest.mark.asyncio +async def test_idp_initiated_replay_rejected_across_workers(saml_env_idp_initiated): + key_pem, cert_pem = saml_env_idp_initiated + shared_store = InMemoryCache() + worker_one = _shared_cache(shared_store) + worker_two = _shared_cache(shared_store) + resp = _build_signed_response(key_pem, cert_pem, email="bob@example.com") + + first = await _acs(_b64(resp), worker_one) + assert first.email == "bob@example.com" + + with pytest.raises(HTTPException) as exc: + await _acs(_b64(resp), worker_two) + assert exc.value.status_code == 401 + + +class _FakeChunkedRequest: + def __init__(self, chunks, content_length=None): + self._chunks = chunks + self.headers = {} if content_length is None else {"content-length": content_length} + + async def stream(self): + for chunk in self._chunks: + yield chunk + + +@pytest.mark.asyncio +async def test_read_acs_post_data_parses_form(): + body = b"SAMLResponse=abc123&RelayState=%2Fui%2F" + request = _FakeChunkedRequest([body], content_length=str(len(body))) + + post_data = await SAMLAuthHandler.read_acs_post_data(cast(Request, request)) + + assert post_data == {"SAMLResponse": "abc123", "RelayState": "/ui/"} + + +@pytest.mark.asyncio +async def test_read_acs_post_data_rejects_oversized_content_length(): + request = _FakeChunkedRequest([b""], content_length=str(_SAML_MAX_POST_BYTES + 1)) + + with pytest.raises(HTTPException) as exc: + await SAMLAuthHandler.read_acs_post_data(cast(Request, request)) + assert exc.value.status_code == 413 + + +@pytest.mark.asyncio +async def test_read_acs_post_data_rejects_oversized_stream_without_content_length(): + chunk = b"a" * (1024 * 1024) + chunk_count = _SAML_MAX_POST_BYTES // len(chunk) + 2 + request = _FakeChunkedRequest([chunk] * chunk_count) + + with pytest.raises(HTTPException) as exc: + await SAMLAuthHandler.read_acs_post_data(cast(Request, request)) + assert exc.value.status_code == 413 diff --git a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py index 4936191c344..5202c8cbfc0 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -1693,6 +1693,78 @@ async def test_update_team_members_list_duplicate_prevention(): assert len(mock_team.members_with_roles) == 1 +@pytest.mark.asyncio +async def test_add_team_members_reconciles_against_freshly_locked_row(): + """ + Regression: _add_team_members_to_team must build the new members_with_roles + from the row it re-reads under a lock inside the write transaction, not from + the stale complete_team_data snapshot captured at the start of the request. + + Two concurrent /team/member_add calls for the same team read the same + snapshot; without the locked re-read the losing write rewrites the whole + JSON array from its stale copy and silently drops the member the other call + already committed. Here the snapshot holds only "zed", a concurrent writer + has already committed "alice" (returned by the locked SELECT), and this call + adds "bob". The write must contain all three. + """ + from litellm.proxy.management_endpoints.team_endpoints import ( + _add_team_members_to_team, + ) + + stale_snapshot = LiteLLM_TeamTable( + team_id="test-team-lock", + members_with_roles=[Member(user_id="zed", role="user")], + ) + + freshly_committed = [ + {"user_id": "zed", "user_email": None, "role": "user"}, + {"user_id": "alice", "user_email": None, "role": "user"}, + ] + + captured: dict = {} + + async def _capture_update(where, data): + captured["data"] = data + return LiteLLM_TeamTable( + team_id="test-team-lock", + members_with_roles=json.loads(data["members_with_roles"]), + ) + + tx = MagicMock() + tx.query_raw = AsyncMock(return_value=[{"members_with_roles": freshly_committed}]) + tx.litellm_teamtable.update = AsyncMock(side_effect=_capture_update) + + tx_cm = MagicMock() + tx_cm.__aenter__ = AsyncMock(return_value=tx) + tx_cm.__aexit__ = AsyncMock(return_value=None) + + prisma_client = MagicMock() + prisma_client.tx = MagicMock(return_value=tx_cm) + + with patch( + "litellm.proxy.management_endpoints.team_endpoints._process_team_members", + new=AsyncMock(return_value=([], [])), + ): + updated_team, _, _ = await _add_team_members_to_team( + data=TeamMemberAddRequest( + team_id="test-team-lock", + member=Member(user_id="bob", role="user"), + ), + complete_team_data=stale_snapshot, + prisma_client=cast(object, prisma_client), + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), + litellm_proxy_admin_name="admin", + ) + + written_ids = sorted(m["user_id"] for m in json.loads(captured["data"]["members_with_roles"])) + assert written_ids == ["alice", "bob", "zed"] + + lock_reads = [call for call in tx.query_raw.call_args_list if "FOR UPDATE" in str(call.args[0])] + assert lock_reads, "expected a SELECT ... FOR UPDATE row-lock read before the write" + + assert [m.user_id for m in updated_team.members_with_roles] == ["zed", "alice", "bob"] + + def test_add_new_models_to_team_with_existing_models(): """ Test add_new_models_to_team function with existing models @@ -4106,6 +4178,8 @@ async def test_new_team_max_budget_within_user_limit(): } mock_prisma.db.litellm_usertable = MagicMock() mock_prisma.db.litellm_usertable.upsert = AsyncMock(return_value=mock_user) + mock_prisma.db.litellm_usertable.update_many = AsyncMock() + mock_prisma.db.litellm_usertable.find_unique = AsyncMock(return_value=mock_user) mock_prisma.db.litellm_usertable.update = AsyncMock(return_value=mock_user) # Mock team membership table @@ -4247,6 +4321,8 @@ async def test_new_team_org_scoped_budget_bypasses_user_limit(): } mock_prisma.db.litellm_usertable = MagicMock() mock_prisma.db.litellm_usertable.upsert = AsyncMock(return_value=mock_user) + mock_prisma.db.litellm_usertable.update_many = AsyncMock() + mock_prisma.db.litellm_usertable.find_unique = AsyncMock(return_value=mock_user) mock_prisma.db.litellm_usertable.update = AsyncMock(return_value=mock_user) # Mock team membership table @@ -4393,6 +4469,8 @@ async def test_new_team_org_scoped_models_bypasses_user_limit(): } mock_prisma.db.litellm_usertable = MagicMock() mock_prisma.db.litellm_usertable.upsert = AsyncMock(return_value=mock_user) + mock_prisma.db.litellm_usertable.update_many = AsyncMock() + mock_prisma.db.litellm_usertable.find_unique = AsyncMock(return_value=mock_user) mock_prisma.db.litellm_usertable.update = AsyncMock(return_value=mock_user) # Mock team membership table @@ -7245,6 +7323,8 @@ async def test_new_team_soft_budget_validation( } mock_prisma.db.litellm_usertable = MagicMock() mock_prisma.db.litellm_usertable.upsert = AsyncMock(return_value=mock_user) + mock_prisma.db.litellm_usertable.update_many = AsyncMock() + mock_prisma.db.litellm_usertable.find_unique = AsyncMock(return_value=mock_user) mock_prisma.db.litellm_usertable.update = AsyncMock(return_value=mock_user) # Mock team membership table @@ -9733,7 +9813,6 @@ async def _drive_team_write( raw_body=None, user=None, find_returns_none=False, - json_side_effect=None, ): """Drive POST ``update_team`` or PATCH ``patch_team`` against a mocked team. @@ -9748,6 +9827,7 @@ async def _drive_team_write( from litellm.proxy._types import ( LiteLLM_TeamTable, LitellmUserRoles, + PatchTeamRequest, UpdateTeamRequest, UserAPIKeyAuth, ) @@ -9793,14 +9873,10 @@ async def _drive_team_write( litellm_changed_by=None, ) else: - if json_side_effect is not None: - req.json = AsyncMock(side_effect=json_side_effect) - else: - req.json = AsyncMock( - return_value=raw_body if raw_body is not None else dict(payload or {}) - ) + body = raw_body if raw_body is not None else dict(payload or {}) result = await patch_team( team_id=_PATCH_TEAM_ID, + data=PatchTeamRequest.model_validate(body), http_request=req, user_api_key_dict=auth, litellm_changed_by=None, @@ -9948,25 +10024,36 @@ async def test_patch_strips_system_managed_metadata_key_like_post(): assert patch_meta == {"cost_center": "9999"} -@pytest.mark.asyncio -@pytest.mark.parametrize("raw_body", [["not", "an", "object"], "a-string", 42, True]) -async def test_patch_rejects_non_object_body(raw_body): - from litellm.proxy._types import ProxyException +@pytest.mark.parametrize( + "kwargs", + [ + {"json": ["not", "an", "object"]}, + {"json": "a-string"}, + {"json": 42}, + {"content": b"{not json"}, + {"json": {"tpm_limit": "not-an-int"}}, + ], + ids=["list", "string", "number", "malformed-json", "wrong-field-type"], +) +def test_patch_rejects_a_malformed_body_with_422(kwargs): + """The body is a declared parameter, so FastAPI rejects a malformed one before the + handler runs. This is the same 422 POST /team/update already returns; the route + previously answered 400 here and 500 for a wrongly typed field, reporting a caller + mistake as a server fault.""" + from fastapi import FastAPI + from fastapi.testclient import TestClient - with pytest.raises(ProxyException) as exc: - await _drive_team_write("patch", existing_metadata={"a": 1}, raw_body=raw_body) - assert exc.value.code == "400" or exc.value.code == 400 + from litellm.proxy._types import PatchTeamRequest + app = FastAPI() -@pytest.mark.asyncio -async def test_patch_rejects_invalid_json_body(): - from litellm.proxy._types import ProxyException + @app.patch("/team/{team_id}") + async def _route(team_id: str, data: PatchTeamRequest): # pragma: no cover - schema only + return {} - with pytest.raises(ProxyException) as exc: - await _drive_team_write( - "patch", existing_metadata={"a": 1}, json_side_effect=ValueError("no body") - ) - assert exc.value.code == "400" or exc.value.code == 400 + response = TestClient(app).patch("/team/abc", **kwargs) + + assert response.status_code == 422 @pytest.mark.asyncio @@ -10036,3 +10123,103 @@ async def test_patch_returns_full_team_object_not_wrapper(): ) assert isinstance(result, LiteLLM_TeamTable) assert result.team_id == _PATCH_TEAM_ID + + +# --------------------------------------------------------------------------- +# PATCH body is validated through PatchTeamRequest before it is handed to +# update_team. The write below must stay byte-identical to what the untyped +# **body construction produced, or a partial update starts writing columns the +# caller never mentioned. +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_patch_writes_only_the_keys_the_caller_sent(): + """An omitted field must not reach the DB write at all. If validation ever + materialises defaults, every unmentioned column gets overwritten with null.""" + _, update_mock = await _drive_team_write("patch", raw_body={"tpm_limit": 5}) + written = update_mock.call_args.kwargs["data"] + + assert written["tpm_limit"] == 5 + for untouched in ("rpm_limit", "max_budget", "models", "blocked", "budget_duration"): + assert untouched not in written, f"{untouched} was written despite not being sent" + + +@pytest.mark.asyncio +async def test_patch_preserves_explicit_null_as_a_clear(): + """null is a clear, not an omission: it has to survive validation and reach the write.""" + _, update_mock = await _drive_team_write("patch", raw_body={"max_budget": None}) + written = update_mock.call_args.kwargs["data"] + + assert "max_budget" in written + assert written["max_budget"] is None + + +def _patch_body_to_update_request(body: dict): + """The exact reshaping patch_team performs between the raw body and update_team.""" + from litellm.proxy._types import PatchTeamRequest, UpdateTeamRequest + + parsed = PatchTeamRequest.model_validate(body) + return UpdateTeamRequest( + team_id=_PATCH_TEAM_ID, + **parsed.model_dump(exclude_unset=True, exclude={"team_id"}), + ) + + +@pytest.mark.parametrize( + "body", + [ + {"tpm_limit": 5}, + {"max_budget": None}, + {"object_permission": {"vector_stores": []}}, + {"metadata": {"a": 1, "b": None}}, + {"models": ["gpt-4"], "blocked": False}, + ], + ids=["scalar", "explicit-null", "partial-nested", "metadata-with-null", "list-and-false"], +) +def test_patch_body_reshaping_adds_no_keys_the_caller_did_not_send(body): + """Validating through PatchTeamRequest must be shape-preserving. If it ever + materialises defaults, a partial update silently overwrites untouched columns, + and for the merge-only object_permission it would wipe sibling sub-keys.""" + reshaped = _patch_body_to_update_request(body) + dumped = reshaped.model_dump(exclude_unset=True, exclude={"team_id"}) + + assert dumped == body + assert reshaped.model_fields_set == set(body) | {"team_id"} + + +@pytest.mark.asyncio +async def test_patch_ignores_unknown_body_keys(): + """Unknown keys were silently dropped by the previous construction; keep that.""" + _, update_mock = await _drive_team_write( + "patch", raw_body={"tpm_limit": 5, "not_a_team_field": "x"} + ) + written = update_mock.call_args.kwargs["data"] + + assert written["tpm_limit"] == 5 + assert "not_a_team_field" not in written + + +def test_patch_team_request_makes_team_id_optional(): + """PATCH takes team_id from the path, so the body model must not require it, + while still inheriting every UpdateTeamRequest field.""" + from litellm.proxy._types import PatchTeamRequest, UpdateTeamRequest + + parsed = PatchTeamRequest.model_validate({"tpm_limit": 5}) + + assert parsed.team_id is None + assert parsed.model_fields_set == {"tpm_limit"} + assert set(UpdateTeamRequest.model_fields).issubset(set(PatchTeamRequest.model_fields)) + + +def test_patch_team_route_publishes_its_request_body_schema(): + """The dashboard's generated client types this call off the OpenAPI spec, which + FastAPI can only emit because the body is a declared parameter.""" + from litellm.proxy.proxy_server import app + + operation = app.openapi()["paths"]["/team/{team_id}"]["patch"] + schema = operation["requestBody"]["content"]["application/json"]["schema"] + + assert schema == {"$ref": "#/components/schemas/PatchTeamRequest"} + properties = app.openapi()["components"]["schemas"]["PatchTeamRequest"]["properties"] + assert "tpm_limit" in properties and "metadata" in properties diff --git a/tests/test_litellm/proxy/management_endpoints/test_tool_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_tool_management_endpoints.py index cf80ee5dee5..351f125052d 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_tool_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_tool_management_endpoints.py @@ -13,12 +13,17 @@ from datetime import datetime, timezone from typing import Optional from unittest.mock import AsyncMock, MagicMock, patch +import pytest from fastapi import FastAPI from fastapi.testclient import TestClient sys.path.insert(0, os.path.abspath("../../..")) -from litellm.proxy.management_endpoints.tool_management_endpoints import router +from litellm.proxy.management_endpoints.tool_management_endpoints import ( + _build_tool_spend_response, + _ToolSpendRow, + router, +) from litellm.types.tool_management import LiteLLM_ToolTableRow # --- helpers --- @@ -50,9 +55,9 @@ def _make_app() -> FastAPI: # Stub the auth dependency so we don't need a real proxy running. def _override_auth(): - from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth - return UserAPIKeyAuth(api_key="sk-test", user_id="admin") + return UserAPIKeyAuth(api_key="sk-test", user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN) # A real (non-None) prisma stub for truthiness checks. @@ -147,3 +152,117 @@ class TestToolManagementEndpoints: json={"tool_name": "my_tool", "input_policy": "invalid_value"}, ) assert resp.status_code == 422 + + def test_tool_spend_route_not_shadowed_by_get_tool(self): + prisma = MagicMock() + prisma.db.query_raw = AsyncMock(return_value=[]) + with patch("litellm.proxy.proxy_server.prisma_client", prisma): + resp = self.client.get("/v1/tool/spend") + assert resp.status_code == 200 + assert resp.json()["by_tool"] == [] + + def test_tool_spend_aggregates_and_sorts(self): + rows = [ + {"date": "2026-07-01", "tool_name": "search", "call_count": 2, "spend": 1.0, "total_tokens": 100}, + {"date": "2026-07-02", "tool_name": "search", "call_count": 1, "spend": 4.0, "total_tokens": 50}, + {"date": "2026-07-01", "tool_name": "read_file", "call_count": 3, "spend": 2.0, "total_tokens": 300}, + ] + prisma = MagicMock() + prisma.db.query_raw = AsyncMock(side_effect=[rows, [{"total_spend": 5.5}]]) + with patch("litellm.proxy.proxy_server.prisma_client", prisma): + resp = self.client.get("/v1/tool/spend?start_date=2026-07-01&end_date=2026-07-02") + assert resp.status_code == 200 + body = resp.json() + assert [t["tool_name"] for t in body["by_tool"]] == ["search", "read_file"] + search = body["by_tool"][0] + assert search["spend"] == 5.0 + assert search["call_count"] == 3 + assert search["total_tokens"] == 150 + assert len(body["daily"]) == 3 + assert body["start_date"] == "2026-07-01" + assert body["end_date"] == "2026-07-02" + assert body["total_spend"] == 5.5 + + @patch("litellm.proxy.proxy_server.prisma_client", None) + def test_tool_spend_no_db_returns_500(self): + resp = self.client.get("/v1/tool/spend") + assert resp.status_code == 500 + + def test_tool_spend_end_date_is_inclusive_via_exclusive_next_day_bound(self): + prisma = MagicMock() + prisma.db.query_raw = AsyncMock(return_value=[]) + with patch("litellm.proxy.proxy_server.prisma_client", prisma): + resp = self.client.get("/v1/tool/spend?start_date=2026-07-01&end_date=2026-07-02") + assert resp.status_code == 200 + expected_binds = ( + datetime(2026, 7, 1, tzinfo=timezone.utc).isoformat(), + datetime(2026, 7, 3, tzinfo=timezone.utc).isoformat(), + ) + assert prisma.db.query_raw.await_count == 2 + for call in prisma.db.query_raw.await_args_list: + assert tuple(call.args[1:]) == expected_binds + assert resp.json()["end_date"] == "2026-07-02" + + @pytest.mark.parametrize( + "query", + [ + "start_date=not-a-date", + "start_date=2026-02-30", + "start_date=07/01/2026", + "end_date=2026-13-01", + "end_date=20260701", + ], + ) + def test_tool_spend_malformed_date_returns_400(self, query: str): + prisma = MagicMock() + prisma.db.query_raw = AsyncMock(return_value=[]) + with patch("litellm.proxy.proxy_server.prisma_client", prisma): + resp = self.client.get(f"/v1/tool/spend?{query}") + assert resp.status_code == 400 + assert "Invalid date format" in resp.json()["detail"] + prisma.db.query_raw.assert_not_awaited() + + def test_tool_spend_non_admin_returns_403(self): + from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + + app = _make_app() + app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth( + api_key="sk-user", user_id="u1", user_role=LitellmUserRoles.INTERNAL_USER + ) + client = TestClient(app, raise_server_exceptions=True) + prisma = MagicMock() + prisma.db.query_raw = AsyncMock(return_value=[]) + with patch("litellm.proxy.proxy_server.prisma_client", prisma): + resp = client.get("/v1/tool/spend") + assert resp.status_code == 403 + prisma.db.query_raw.assert_not_awaited() + + +def _spend_row(date: str, tool_name: str, spend: float, call_count: int = 1, total_tokens: int = 10) -> _ToolSpendRow: + return _ToolSpendRow(date=date, tool_name=tool_name, call_count=call_count, spend=spend, total_tokens=total_tokens) + + +class TestBuildToolSpendResponse: + def test_multi_tool_attribution_double_counts_per_tool_but_not_total(self): + rows = [ + _spend_row("2026-07-01", "a", spend=3.0), + _spend_row("2026-07-01", "b", spend=3.0), + ] + resp = _build_tool_spend_response(rows, total_spend=3.0, start_date="2026-07-01", end_date="2026-07-01") + by_tool = {t.tool_name: t.spend for t in resp.by_tool} + assert by_tool == {"a": 3.0, "b": 3.0} + assert resp.total_spend == 3.0 + + def test_groups_across_days_and_sorts_by_spend(self): + rows = [ + _spend_row("2026-07-01", "b", spend=1.0, call_count=2, total_tokens=100), + _spend_row("2026-07-02", "b", spend=4.0, call_count=1, total_tokens=50), + _spend_row("2026-07-01", "a", spend=2.0, call_count=3, total_tokens=300), + ] + resp = _build_tool_spend_response(rows, total_spend=7.0, start_date="2026-07-01", end_date="2026-07-02") + assert [(t.tool_name, t.spend, t.call_count, t.total_tokens) for t in resp.by_tool] == [ + ("b", 5.0, 3, 150), + ("a", 2.0, 3, 300), + ] + assert len(resp.daily) == 3 diff --git a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py index e1856860c8a..795b7cd5a9e 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py +++ b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py @@ -2214,7 +2214,95 @@ class TestCLIKeyRegenerationFlow: _get_cli_sso_flow_or_raise(login_id="cli-test_1234567890", cache=mock_cache) assert expired_exc.value.status_code == 400 assert "session not found or expired" in expired_exc.value.detail - assert "enable_redis_auth_cache" in expired_exc.value.detail + assert "configure a Redis cache" in expired_exc.value.detail + assert "enable_redis_auth_cache" not in expired_exc.value.detail + + def test_cli_sso_flow_is_redis_authoritative_when_redis_attached(self): + """ + When Redis is attached, the CLI SSO flow must be read from and written to + Redis directly, never the in-memory layer. Otherwise the worker that served + /sso/cli/start keeps serving its stale in-memory flow and never sees the + sso_complete/session_data update another worker wrote, which is exactly the + multi-worker failure this fix targets. + """ + from litellm.proxy.management_endpoints.ui_sso import ( + CLI_SSO_SESSION_TTL_SECONDS, + _get_cli_sso_flow_cache_key, + _get_cli_sso_flow_or_raise, + _set_cli_sso_flow, + ) + + login_id = "cli-redis_authoritative_1234567890" + cache_key = _get_cli_sso_flow_cache_key(login_id) + fresh_flow = {"poll_secret_hash": "fresh", "sso_complete": True} + stale_flow = {"poll_secret_hash": "stale", "sso_complete": False} + + redis_cache = MagicMock() + redis_cache.get_cache.return_value = fresh_flow + cache = MagicMock() + cache.redis_cache = redis_cache + cache.get_cache.return_value = stale_flow + + result = _get_cli_sso_flow_or_raise(login_id=login_id, cache=cache) + + assert result == fresh_flow + redis_cache.get_cache.assert_called_once_with(key=cache_key) + cache.get_cache.assert_not_called() + + _set_cli_sso_flow(login_id=login_id, cache=cache, flow=fresh_flow) + + redis_cache.set_cache.assert_called_once_with( + key=cache_key, value=json.dumps(fresh_flow), ttl=CLI_SSO_SESSION_TTL_SECONDS + ) + cache.set_cache.assert_not_called() + + def test_cli_sso_flow_with_enum_survives_redis_round_trip(self): + """ + RedisCache stores values via str(value) and reads them back through + json.loads/ast.literal_eval. A raw flow dict containing a Python enum + (session_data.user_role after the SSO callback) produces an unparseable + repr, so every worker reading the completed flow from Redis got a + SyntaxError and returned 400 "session not found". The flow must survive + a real Redis serialization round trip. + """ + from litellm.caching.redis_cache import RedisCache + from litellm.proxy._types import LitellmUserRoles + from litellm.proxy.management_endpoints.ui_sso import ( + _get_cli_sso_flow_or_raise, + _set_cli_sso_flow, + ) + + login_id = "cli-enum_round_trip_1234567890" + completed_flow = { + "poll_secret_hash": "hash", + "sso_complete": True, + "user_code_verified": False, + "session_data": { + "user_id": "user-1", + "user_role": LitellmUserRoles.INTERNAL_USER_VIEW_ONLY, + "models": [], + "teams": ["team-1"], + "team_details": [{"team_id": "team-1", "team_alias": "alias"}], + }, + } + + redis_store: dict = {} + redis_cache = MagicMock() + redis_cache.set_cache.side_effect = lambda key, value, ttl: redis_store.__setitem__( + key, str(value).encode("utf-8") + ) + redis_cache.get_cache.side_effect = lambda key: RedisCache._get_cache_logic( + MagicMock(), redis_store.get(key) + ) + cache = MagicMock() + cache.redis_cache = redis_cache + + _set_cli_sso_flow(login_id=login_id, cache=cache, flow=completed_flow) + flow = _get_cli_sso_flow_or_raise(login_id=login_id, cache=cache) + + assert flow["sso_complete"] is True + assert flow["session_data"]["user_role"] == LitellmUserRoles.INTERNAL_USER_VIEW_ONLY.value + assert flow["session_data"]["team_details"] == [{"team_id": "team-1", "team_alias": "alias"}] @pytest.mark.asyncio async def test_cli_sso_start_creates_bound_flow(self): @@ -2228,10 +2316,13 @@ class TestCLIKeyRegenerationFlow: mock_request = MagicMock(spec=Request) mock_request.client = SimpleNamespace(host="127.0.0.1") mock_request.headers = {} - mock_cache = MagicMock() + mock_cache = MagicMock(redis_cache=None) mock_cache.increment_cache.return_value = 1 - with patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache): + with ( + patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache), + patch("litellm.proxy.proxy_server.cli_sso_session_cache", mock_cache), + ): result = await cli_sso_start(request=mock_request) assert result["login_id"].startswith("cli-") @@ -2259,10 +2350,13 @@ class TestCLIKeyRegenerationFlow: mock_request = MagicMock(spec=Request) mock_request.client = SimpleNamespace(host="127.0.0.1") mock_request.headers = {} - mock_cache = MagicMock() + mock_cache = MagicMock(redis_cache=None) mock_cache.increment_cache.return_value = 31 - with patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache): + with ( + patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache), + patch("litellm.proxy.proxy_server.cli_sso_session_cache", mock_cache), + ): with pytest.raises(HTTPException) as exc_info: await cli_sso_start(request=mock_request) @@ -2281,7 +2375,7 @@ class TestCLIKeyRegenerationFlow: mock_request.client = SimpleNamespace(host="127.0.0.1") mock_request.headers = {} mock_request.base_url = "https://proxy.example.com/" - mock_cache = MagicMock() + mock_cache = MagicMock(redis_cache=None) mock_cache.increment_cache.return_value = 1 with ( @@ -2315,7 +2409,7 @@ class TestCLIKeyRegenerationFlow: mock_request.client = SimpleNamespace(host="127.0.0.1") mock_request.headers = {} mock_request.base_url = "https://proxy.example.com/" - mock_cache = MagicMock() + mock_cache = MagicMock(redis_cache=None) mock_cache.increment_cache.return_value = 1 with ( @@ -2349,7 +2443,7 @@ class TestCLIKeyRegenerationFlow: mock_request = MagicMock(spec=Request) mock_request.base_url = "https://proxy.example.com/" - mock_cache = MagicMock() + mock_cache = MagicMock(redis_cache=None) mock_cache.get_cache.return_value = {"poll_secret_hash": "h"} async def drive(enabled: bool): @@ -2358,6 +2452,7 @@ class TestCLIKeyRegenerationFlow: patch("litellm.proxy.proxy_server.premium_user", True), patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache), + patch("litellm.proxy.proxy_server.cli_sso_session_cache", mock_cache), patch( "litellm.proxy.proxy_server.user_custom_ui_sso_sign_in_handler", None, @@ -2525,7 +2620,7 @@ class TestCLIKeyRegenerationFlow: ) mock_sso_result = {"user_email": "test@example.com", "user_id": "test-user-123"} - mock_cache = MagicMock() + mock_cache = MagicMock(redis_cache=None) mock_cache.get_cache.return_value = { "poll_secret_hash": "poll-secret-hash", "user_code_hash": "user-code-hash", @@ -2544,6 +2639,7 @@ class TestCLIKeyRegenerationFlow: ), patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache), + patch("litellm.proxy.proxy_server.cli_sso_session_cache", mock_cache), ): result = await cli_sso_callback( request=mock_request, @@ -2568,7 +2664,7 @@ class TestCLIKeyRegenerationFlow: mock_request.body = AsyncMock( return_value=b"user_code=ABCD-EFGH&browser_complete_token=browser-token" ) - mock_cache = MagicMock() + mock_cache = MagicMock(redis_cache=None) mock_cache.get_cache.return_value = { "poll_secret_hash": _hash_cli_sso_secret("poll-secret"), "user_code_hash": _hash_cli_sso_secret( @@ -2582,6 +2678,7 @@ class TestCLIKeyRegenerationFlow: with ( patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache), + patch("litellm.proxy.proxy_server.cli_sso_session_cache", mock_cache), patch( "litellm.proxy.common_utils.html_forms.cli_sso_success.render_cli_sso_success_page", return_value="Success", @@ -2606,7 +2703,7 @@ class TestCLIKeyRegenerationFlow: mock_request = MagicMock(spec=Request) mock_request.body = AsyncMock(return_value=b"user_code=ABCD-EFGH") - mock_cache = MagicMock() + mock_cache = MagicMock(redis_cache=None) mock_cache.get_cache.return_value = { "poll_secret_hash": _hash_cli_sso_secret("poll-secret"), "user_code_hash": _hash_cli_sso_secret( @@ -2618,7 +2715,10 @@ class TestCLIKeyRegenerationFlow: "session_data": {"user_id": "test-user-123"}, } - with patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache): + with ( + patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache), + patch("litellm.proxy.proxy_server.cli_sso_session_cache", mock_cache), + ): with pytest.raises(HTTPException) as exc_info: await cli_sso_complete( request=mock_request, login_id="cli-session-4567890" @@ -2640,7 +2740,7 @@ class TestCLIKeyRegenerationFlow: mock_request.body = AsyncMock( return_value=b"user_code=ABCD-EFGH&browser_complete_token=browser-token" ) - mock_cache = MagicMock() + mock_cache = MagicMock(redis_cache=None) mock_cache.get_cache.return_value = { "poll_secret_hash": _hash_cli_sso_secret("poll-secret"), "user_code_hash": _hash_cli_sso_secret( @@ -2651,7 +2751,10 @@ class TestCLIKeyRegenerationFlow: "session_data": None, } - with patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache): + with ( + patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache), + patch("litellm.proxy.proxy_server.cli_sso_session_cache", mock_cache), + ): with pytest.raises(HTTPException) as exc_info: await cli_sso_complete( request=mock_request, login_id="cli-session-4567890" @@ -2687,7 +2790,7 @@ class TestCLIKeyRegenerationFlow: mock_sso_result = {"user_email": "test@example.com", "user_id": "test-user-123"} # Mock cache - mock_cache = MagicMock() + mock_cache = MagicMock(redis_cache=None) mock_cache.get_cache.return_value = { "poll_secret_hash": "poll-secret-hash", "user_code_hash": "user-code-hash", @@ -2709,6 +2812,7 @@ class TestCLIKeyRegenerationFlow: ), patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache), + patch("litellm.proxy.proxy_server.cli_sso_session_cache", mock_cache), patch( "litellm.proxy.common_utils.html_forms.cli_sso_success.render_cli_sso_success_page", return_value="Success", @@ -2769,7 +2873,7 @@ class TestCLIKeyRegenerationFlow: } # Mock cache - mock_cache = MagicMock() + mock_cache = MagicMock(redis_cache=None) mock_cache.get_cache.return_value = { "poll_secret_hash": _hash_cli_sso_secret("poll-secret"), "sso_complete": True, @@ -2777,7 +2881,10 @@ class TestCLIKeyRegenerationFlow: "session_data": session_data, } - with patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache): + with ( + patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache), + patch("litellm.proxy.proxy_server.cli_sso_session_cache", mock_cache), + ): # Act - First poll without team_id result = await cli_poll_key( key_id=session_key, @@ -2803,7 +2910,7 @@ class TestCLIKeyRegenerationFlow: cli_poll_key, ) - mock_cache = MagicMock() + mock_cache = MagicMock(redis_cache=None) mock_cache.get_cache.return_value = { "poll_secret_hash": _hash_cli_sso_secret("poll-secret"), "sso_complete": True, @@ -2816,7 +2923,10 @@ class TestCLIKeyRegenerationFlow: }, } - with patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache): + with ( + patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache), + patch("litellm.proxy.proxy_server.cli_sso_session_cache", mock_cache), + ): with pytest.raises(HTTPException) as exc_info: await cli_poll_key(key_id="cli-session-789123", team_id=None) @@ -2830,7 +2940,7 @@ class TestCLIKeyRegenerationFlow: cli_poll_key, ) - mock_cache = MagicMock() + mock_cache = MagicMock(redis_cache=None) mock_cache.get_cache.return_value = { "poll_secret_hash": _hash_cli_sso_secret("poll-secret"), "sso_complete": True, @@ -2843,7 +2953,10 @@ class TestCLIKeyRegenerationFlow: }, } - with patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache): + with ( + patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache), + patch("litellm.proxy.proxy_server.cli_sso_session_cache", mock_cache), + ): result = await cli_poll_key( key_id="cli-session-789123", team_id=None, @@ -3011,7 +3124,7 @@ class TestCLIKeyRegenerationFlow: ) # Mock cache - mock_cache = MagicMock() + mock_cache = MagicMock(redis_cache=None) mock_cache.get_cache.return_value = { "poll_secret_hash": _hash_cli_sso_secret("poll-secret"), "sso_complete": True, @@ -3023,6 +3136,7 @@ class TestCLIKeyRegenerationFlow: with ( patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache), + patch("litellm.proxy.proxy_server.cli_sso_session_cache", mock_cache), patch("litellm.proxy.proxy_server.prisma_client"), patch( "litellm.proxy.auth.auth_checks.ExperimentalUIJWTToken.get_cli_jwt_auth_token", @@ -3086,7 +3200,7 @@ class TestCLIKeyRegenerationFlow: models=["gpt-4"], max_budget=100.0, ) - mock_cache = MagicMock() + mock_cache = MagicMock(redis_cache=None) mock_cache.get_cache.return_value = { "poll_secret_hash": _hash_cli_sso_secret("poll-secret"), "sso_complete": True, @@ -3097,6 +3211,7 @@ class TestCLIKeyRegenerationFlow: with ( patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache), + patch("litellm.proxy.proxy_server.cli_sso_session_cache", mock_cache), patch("litellm.proxy.proxy_server.prisma_client"), patch( "litellm.proxy.auth.auth_checks.ExperimentalUIJWTToken.get_cli_jwt_auth_token", @@ -3142,7 +3257,7 @@ class TestCLIKeyRegenerationFlow: "models": ["gpt-4"], "user_email": "unbudgeted@example.com", } - mock_cache = MagicMock() + mock_cache = MagicMock(redis_cache=None) mock_cache.get_cache.return_value = { "poll_secret_hash": _hash_cli_sso_secret("poll-secret"), "sso_complete": True, @@ -3153,6 +3268,7 @@ class TestCLIKeyRegenerationFlow: with ( patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache), + patch("litellm.proxy.proxy_server.cli_sso_session_cache", mock_cache), patch( "litellm.proxy.auth.auth_checks.ExperimentalUIJWTToken.get_cli_jwt_auth_token", return_value=mock_jwt_token, @@ -4082,7 +4198,7 @@ class TestPKCEFunctionality: mock_request.query_params = {"state": test_state} # Mock cache with async methods — use dict format (primary path) - mock_cache = MagicMock() + mock_cache = MagicMock(redis_cache=None) test_code_verifier = "test_code_verifier_abc123xyz" mock_cache.async_get_cache = AsyncMock( return_value={"code_verifier": test_code_verifier} @@ -4133,7 +4249,7 @@ class TestPKCEFunctionality: mock_sso.__exit__ = MagicMock(return_value=False) test_state = "test456" - mock_cache = MagicMock() + mock_cache = MagicMock(redis_cache=None) mock_cache.async_set_cache = AsyncMock() @@ -4657,7 +4773,7 @@ class TestPKCEFunctionality: from litellm.proxy._types import ProxyException from litellm.proxy.management_endpoints.ui_sso import SSOAuthenticationHandler - mock_cache = MagicMock() + mock_cache = MagicMock(redis_cache=None) mock_cache.async_get_cache = AsyncMock(return_value=None) # verifier not found mock_request = MagicMock(spec=Request) @@ -4783,7 +4899,7 @@ class TestPKCEFunctionality: from litellm.proxy.management_endpoints.ui_sso import SSOAuthenticationHandler # Cache returns an integer — unexpected format - mock_cache = MagicMock() + mock_cache = MagicMock(redis_cache=None) mock_cache.async_get_cache = AsyncMock(return_value=12345) mock_cache.async_delete_cache = AsyncMock() @@ -4825,7 +4941,7 @@ class TestPKCEFunctionality: from litellm.proxy.management_endpoints.ui_sso import SSOAuthenticationHandler - mock_cache = MagicMock() + mock_cache = MagicMock(redis_cache=None) mock_cache.async_get_cache = AsyncMock(return_value=None) # verifier not found mock_request = MagicMock(spec=Request) @@ -4913,7 +5029,7 @@ class TestPKCEFunctionality: from litellm.proxy.management_endpoints.ui_sso import SSOAuthenticationHandler # Cache returns an integer — unexpected format - mock_cache = MagicMock() + mock_cache = MagicMock(redis_cache=None) mock_cache.async_get_cache = AsyncMock(return_value=12345) mock_cache.async_delete_cache = AsyncMock() @@ -4965,7 +5081,7 @@ class TestPKCEFunctionality: from litellm.proxy.management_endpoints.ui_sso import SSOAuthenticationHandler legacy_verifier = "legacy_plain_string_verifier_abc123" - mock_cache = MagicMock() + mock_cache = MagicMock(redis_cache=None) mock_cache.async_get_cache = AsyncMock(return_value=legacy_verifier) mock_request = MagicMock(spec=Request) @@ -6249,7 +6365,7 @@ class TestCliSsoAttributionMetadata: provider="generic", team_ids=[], ) - mock_cache = MagicMock() + mock_cache = MagicMock(redis_cache=None) mock_cache.get_cache.return_value = { "poll_secret_hash": "poll-secret-hash", "user_code_hash": "user-code-hash", @@ -6266,6 +6382,7 @@ class TestCliSsoAttributionMetadata: ), patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache), + patch("litellm.proxy.proxy_server.cli_sso_session_cache", mock_cache), patch("litellm.proxy.proxy_server.user_custom_sso", None), ): await ui_sso.cli_sso_callback( @@ -6290,7 +6407,7 @@ class TestCliSsoAttributionMetadata: mock_request = MagicMock(spec=Request) mock_request.base_url = "http://internal-proxy.local/" - mock_cache = MagicMock() + mock_cache = MagicMock(redis_cache=None) mock_cache.get_cache.return_value = { "poll_secret_hash": "poll-secret-hash", "user_code_hash": "user-code-hash", @@ -6313,6 +6430,7 @@ class TestCliSsoAttributionMetadata: ) as get_user_info_mock, patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache), + patch("litellm.proxy.proxy_server.cli_sso_session_cache", mock_cache), patch("litellm.proxy.proxy_server.user_custom_sso", None), patch( "litellm.proxy.proxy_server.general_settings", @@ -6359,7 +6477,7 @@ class TestCliSsoAttributionMetadata: "user_id": "test-user-123", "employment_type": "contractor", } - mock_cache = MagicMock() + mock_cache = MagicMock(redis_cache=None) mock_cache.get_cache.return_value = { "poll_secret_hash": "poll-secret-hash", "user_code_hash": "user-code-hash", @@ -6387,6 +6505,7 @@ class TestCliSsoAttributionMetadata: ), patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache), + patch("litellm.proxy.proxy_server.cli_sso_session_cache", mock_cache), patch("litellm.proxy.proxy_server.user_custom_sso", None), patch( "litellm.proxy.common_utils.html_forms.cli_sso_success.render_cli_sso_success_page", @@ -6428,7 +6547,7 @@ class TestCliSsoAttributionMetadata: "org": {"cost_center": "CC-42"}, }, } - mock_cache = MagicMock() + mock_cache = MagicMock(redis_cache=None) mock_cache.get_cache.return_value = { "poll_secret_hash": _hash_cli_sso_secret("poll-secret"), "sso_complete": True, @@ -6436,7 +6555,10 @@ class TestCliSsoAttributionMetadata: "session_data": session_data, } - with patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache): + with ( + patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache), + patch("litellm.proxy.proxy_server.cli_sso_session_cache", mock_cache), + ): result = await cli_poll_key( key_id=session_key, team_id=None, @@ -7269,6 +7391,71 @@ async def test_legacy_login_page_hides_credentials_hint_via_general_settings(): assert "MASTER_KEY" not in body +@pytest.mark.asyncio +async def test_saml_callback_blocked_when_admin_ui_disabled(): + """An IdP-initiated assertion must not mint a UI session when the admin UI is + disabled; the ACS enforces DISABLE_ADMIN_UI like the SP-initiated login route.""" + from litellm.proxy.management_endpoints.ui_sso import saml_callback + + with patch.dict(os.environ, {"DISABLE_ADMIN_UI": "true"}): + response = await saml_callback(SimpleNamespace(cookies={})) + + assert response.status_code == 200 + assert "Admin UI is Disabled" in response.body.decode() + + +@pytest.mark.asyncio +async def test_saml_callback_enforces_free_sso_user_limit_after_validation(): + """An IdP-initiated assertion must not bypass the >5 free-SSO-user Enterprise gate + that /sso/key/generate enforces; the ACS re-checks it after validating the assertion, + so the entitlement DB query never runs on unvalidated input.""" + from litellm.proxy._types import ProxyException + from litellm.proxy.management_endpoints.ui_sso import saml_callback + from litellm.proxy.management_endpoints.types import CustomOpenID + + call_order: list[str] = [] + + async def _fake_handle_acs(**kwargs): + call_order.append("validate") + return CustomOpenID( + id="dana@litellm.ai", + email="dana@litellm.ai", + first_name=None, + last_name=None, + display_name="dana", + picture=None, + provider="saml", + team_ids=[], + user_role=None, + ) + + async def _fake_count_billable_users(): + call_order.append("count") + return 6 + + async def _stream(): + yield b"SAMLResponse=signed-response" + + request_double = SimpleNamespace(cookies={}, headers={}, stream=_stream) + + with patch.dict(os.environ, {"DISABLE_ADMIN_UI": "false"}), patch( + "litellm.proxy.proxy_server.premium_user", False + ), patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), patch( + "litellm.proxy.proxy_server.master_key", "sk-1234" + ), patch( + "litellm.proxy.management_endpoints.sso.saml_sso.SAMLAuthHandler.handle_acs", + new=_fake_handle_acs, + ), patch( + "litellm.repositories.user_repository.UserRepository.count_billable_users", + new=AsyncMock(side_effect=_fake_count_billable_users), + ): + with pytest.raises(ProxyException) as exc: + await saml_callback(request_double) + + assert str(exc.value.code) == "403" + assert call_order == ["validate", "count"] + + @pytest.mark.asyncio async def test_cli_poll_key_tolerates_missing_user_row(): """The CLI poll must still mint the JWT when the user lookup raises, @@ -7287,7 +7474,7 @@ async def test_cli_poll_key_tolerates_missing_user_row(): "models": ["gpt-4"], } - mock_cache = MagicMock() + mock_cache = MagicMock(redis_cache=None) mock_cache.get_cache.return_value = { "poll_secret_hash": _hash_cli_sso_secret("poll-secret"), "sso_complete": True, @@ -7299,6 +7486,7 @@ async def test_cli_poll_key_tolerates_missing_user_row(): with ( patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache), + patch("litellm.proxy.proxy_server.cli_sso_session_cache", mock_cache), patch("litellm.proxy.proxy_server.prisma_client"), patch( "litellm.proxy.auth.auth_checks.ExperimentalUIJWTToken.get_cli_jwt_auth_token", @@ -7633,9 +7821,87 @@ async def test_cli_completion_persists_assertion_under_db_user_id(): user_defined_values=None, prisma_client=MagicMock(), user_api_key_cache=MagicMock(), + cli_sso_session_cache=MagicMock(), proxy_logging_obj=MagicMock(), sso_assertion=assertion, ) retain_mock.assert_awaited_once_with(user_id="cli-user-id", assertion=assertion) assert response.status_code == 200 + + +class TestSameOriginReturnPath: + """The same-origin relative return_to arm added for the MCP gateway DCR authorize + round-trip: only strictly relative paths qualify, so login can never redirect the + browser off the gateway origin.""" + + def test_accepts_relative_paths(self): + from litellm.proxy.management_endpoints.ui_sso import _is_same_origin_return_path + + assert _is_same_origin_return_path("/authorize?client_id=llm_dcrc_x&state=s") is True + assert _is_same_origin_return_path("/some_server/authorize") is True + + def test_rejects_absolute_protocol_relative_and_backslash_paths(self): + from litellm.proxy.management_endpoints.ui_sso import _is_same_origin_return_path + + assert _is_same_origin_return_path("https://evil.example.com/authorize") is False + assert _is_same_origin_return_path("//evil.example.com/authorize") is False + assert _is_same_origin_return_path("/\\evil.example.com") is False + assert _is_same_origin_return_path("javascript:alert(1)") is False + assert _is_same_origin_return_path("") is False + + +class TestPersistReturnToCookieSharedHelper: + """The single shared return_to helper used by EVERY sign-in branch (SSO / Okta / generic AND the + username/password form). It must be best-effort and NEVER raise — a bad return_to can never block + sign-in. Regression: the password form previously 400'd because it called _validate_return_to + directly (which raises for a non-matching absolute return_to when control_plane_url is set).""" + + @staticmethod + def _cookie(resp) -> str: + return resp.headers.get("set-cookie", "") + + def test_sets_cookie_for_same_origin_relative_path(self, monkeypatch): + from fastapi import Response + + from litellm.proxy.management_endpoints.ui_sso import _persist_return_to_cookie + + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {}) + resp = Response() + _persist_return_to_cookie(resp, "/mcp/authorize?client_id=llm_dcrc_abc") + assert "litellm_cp_return_to=" in self._cookie(resp) + + def test_bad_absolute_with_control_plane_configured_does_not_raise_and_is_not_stored(self, monkeypatch): + """THE regression: a non-matching absolute return_to with control_plane_url set must NOT raise + (it did, blocking the login form) and must NOT be stored — sign-in proceeds.""" + from fastapi import Response + + from litellm.proxy.management_endpoints.ui_sso import _persist_return_to_cookie + + monkeypatch.setattr( + "litellm.proxy.proxy_server.general_settings", {"control_plane_url": "https://cp.example.com"} + ) + resp = Response() + _persist_return_to_cookie(resp, "https://evil.example.com/steal") # must not raise + assert "litellm_cp_return_to=" not in self._cookie(resp) + + def test_none_return_to_is_a_noop(self): + from fastapi import Response + + from litellm.proxy.management_endpoints.ui_sso import _persist_return_to_cookie + + resp = Response() + _persist_return_to_cookie(resp, None) + assert "litellm_cp_return_to=" not in self._cookie(resp) + + def test_control_plane_matching_absolute_is_stored(self, monkeypatch): + from fastapi import Response + + from litellm.proxy.management_endpoints.ui_sso import _persist_return_to_cookie + + monkeypatch.setattr( + "litellm.proxy.proxy_server.general_settings", {"control_plane_url": "https://cp.example.com"} + ) + resp = Response() + _persist_return_to_cookie(resp, "https://cp.example.com/ui?page=models") + assert "litellm_cp_return_to=" in self._cookie(resp) diff --git a/tests/test_litellm/proxy/management_helpers/test_management_helpers_utils.py b/tests/test_litellm/proxy/management_helpers/test_management_helpers_utils.py index dbbdca65cc6..504414ea635 100644 --- a/tests/test_litellm/proxy/management_helpers/test_management_helpers_utils.py +++ b/tests/test_litellm/proxy/management_helpers/test_management_helpers_utils.py @@ -202,7 +202,8 @@ async def test_add_new_member_clones_default_team_budget_id(): "teams": [test_team_id], "user_role": "internal_user", } - mock_prisma_client.db.litellm_usertable.upsert = AsyncMock( + mock_prisma_client.db.litellm_usertable.upsert = AsyncMock(return_value=mock_user_response) + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( return_value=mock_user_response ) @@ -305,7 +306,8 @@ async def test_add_new_member_budget_duration_only_clones_default_max_budget(): "teams": ["team-dc"], "user_role": "internal_user", } - mock_prisma_client.db.litellm_usertable.upsert = AsyncMock( + mock_prisma_client.db.litellm_usertable.upsert = AsyncMock(return_value=mock_user_response) + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( return_value=mock_user_response ) mock_default_budget_row = MagicMock() @@ -388,7 +390,8 @@ async def test_add_new_member_no_budget_when_no_default_and_no_max_budget(): "teams": [test_team_id], "user_role": "internal_user", } - mock_prisma_client.db.litellm_usertable.upsert = AsyncMock( + mock_prisma_client.db.litellm_usertable.upsert = AsyncMock(return_value=mock_user_response) + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( return_value=mock_user_response ) @@ -455,7 +458,8 @@ async def test_add_new_member_creates_new_budget_when_max_budget_provided(): "teams": [test_team_id], "user_role": "internal_user", } - mock_prisma_client.db.litellm_usertable.upsert = AsyncMock( + mock_prisma_client.db.litellm_usertable.upsert = AsyncMock(return_value=mock_user_response) + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( return_value=mock_user_response ) @@ -531,7 +535,8 @@ async def test_add_new_member_persists_budget_duration(): "teams": ["team-dur"], "user_role": "internal_user", } - mock_prisma_client.db.litellm_usertable.upsert = AsyncMock( + mock_prisma_client.db.litellm_usertable.upsert = AsyncMock(return_value=mock_user_response) + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( return_value=mock_user_response ) mock_budget_response = MagicMock() @@ -594,7 +599,8 @@ async def test_add_new_member_persists_budget_duration_without_max_budget(): "teams": ["team-dur2"], "user_role": "internal_user", } - mock_prisma_client.db.litellm_usertable.upsert = AsyncMock( + mock_prisma_client.db.litellm_usertable.upsert = AsyncMock(return_value=mock_user_response) + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( return_value=mock_user_response ) mock_budget_response = MagicMock() @@ -997,3 +1003,125 @@ async def test_attach_object_permission_to_dict_with_none_object_permission_id() # Verify no database query was made mock_prisma_client.db.litellm_objectpermissiontable.find_unique.assert_not_called() + + +@pytest.mark.asyncio +async def test_add_new_member_appends_team_only_if_absent_for_existing_user(): + """Adding an existing user to a team must append the team id only if it is + not already present. + + add_new_member is the single writer of user.teams for every team add + (/team/member_add, /user/new, SSO, SCIM). An unconditional append let + repeated or concurrent adds accumulate duplicate team ids in user.teams, + which also breaks auth logic that keys off the number of teams a user + belongs to. The append must go through a filtered update that no-ops when + the team is already present, and it must not fall through to creating a new + user row for a user that already exists. + """ + from litellm.proxy._types import LitellmUserRoles + + new_member = Member(user_id="existing-user", role="user") + user_api_key_dict = UserAPIKeyAuth( + user_id="admin_user", user_role=LitellmUserRoles.PROXY_ADMIN + ) + + mock_prisma_client = AsyncMock() + + mock_user_after = MagicMock() + mock_user_after.model_dump.return_value = { + "user_id": "existing-user", + "user_email": None, + "teams": ["team-1"], + "user_role": "internal_user", + } + mock_prisma_client.db.litellm_usertable.upsert = AsyncMock(return_value=mock_user_after) + mock_prisma_client.db.litellm_usertable.update_many = AsyncMock() + # no team default budget and no explicit budget -> no team membership row + mock_prisma_client.db.litellm_budgettable.find_unique = AsyncMock(return_value=None) + + result_user, _ = await add_new_member( + new_member=new_member, + max_budget_in_team=None, + prisma_client=mock_prisma_client, + team_id="team-1", + user_api_key_dict=user_api_key_dict, + litellm_proxy_admin_name="admin", + ) + + assert result_user is not None + assert result_user.user_id == "existing-user" + + # the append must be a filtered, idempotent update keyed off the team id, so + # a repeated or concurrent add of a team the user already has is a no-op + mock_prisma_client.db.litellm_usertable.update_many.assert_called_once() + where = mock_prisma_client.db.litellm_usertable.update_many.call_args.kwargs["where"] + assert where["user_id"] == "existing-user" + assert where["NOT"] == {"teams": {"has": "team-1"}} + data = mock_prisma_client.db.litellm_usertable.update_many.call_args.kwargs["data"] + assert data == {"teams": {"push": ["team-1"]}} + + # upsert (not an unconditional teams push) is what ensures the row exists, so + # its update branch must not carry a teams push that would duplicate + mock_prisma_client.db.litellm_usertable.upsert.assert_called_once() + upsert_update = mock_prisma_client.db.litellm_usertable.upsert.call_args.kwargs["data"]["update"] + assert "teams" not in upsert_update + + +@pytest.mark.asyncio +async def test_add_new_member_creates_missing_user_atomically_via_upsert(): + """A brand-new user added to a team must be created via an atomic upsert, not + a separate existence check followed by create. + + Concurrent provisioning of the same new user (which SCIM group reconciles do) + would race a check-then-create into a duplicate-key failure. The upsert seeds + teams on create, and the filtered append is a no-op because the team is + already present on the freshly created row. + + Calling upsert is not on its own enough to be atomic: Prisma only compiles it + down to a single INSERT ... ON CONFLICT when the update branch is non-empty, + and otherwise emits SELECT-then-INSERT, which loses the race. That is how + parallel /team/new calls naming the same new member started returning 500 + "Unique constraint failed on the fields: (user_id)", so the shape of both + branches is pinned here. + """ + from litellm.proxy._types import LitellmUserRoles + + new_member = Member(user_id="brand-new-user", role="user") + user_api_key_dict = UserAPIKeyAuth( + user_id="admin_user", user_role=LitellmUserRoles.PROXY_ADMIN + ) + + mock_prisma_client = AsyncMock() + + mock_created = MagicMock() + mock_created.model_dump.return_value = { + "user_id": "brand-new-user", + "user_email": None, + "teams": ["team-1"], + "user_role": "internal_user", + } + mock_prisma_client.db.litellm_usertable.upsert = AsyncMock(return_value=mock_created) + mock_prisma_client.db.litellm_usertable.update_many = AsyncMock() + mock_prisma_client.db.litellm_usertable.create = AsyncMock() + mock_prisma_client.db.litellm_budgettable.find_unique = AsyncMock(return_value=None) + + result_user, _ = await add_new_member( + new_member=new_member, + max_budget_in_team=None, + prisma_client=mock_prisma_client, + team_id="team-1", + user_api_key_dict=user_api_key_dict, + litellm_proxy_admin_name="admin", + ) + + assert result_user is not None + assert result_user.user_id == "brand-new-user" + + # existence is established by an atomic upsert (create-or-update), never a + # non-atomic standalone create that could race under concurrent provisioning + mock_prisma_client.db.litellm_usertable.upsert.assert_called_once() + mock_prisma_client.db.litellm_usertable.create.assert_not_called() + upsert_data = mock_prisma_client.db.litellm_usertable.upsert.call_args.kwargs["data"] + assert upsert_data["create"]["teams"] == ["team-1"] + assert upsert_data["update"], "empty update branch degrades the upsert to a racy SELECT-then-INSERT" + assert "teams" not in upsert_data["update"] diff --git a/tests/test_litellm/proxy/proxy_server/test_routes_login_sso.py b/tests/test_litellm/proxy/proxy_server/test_routes_login_sso.py index f0250bbe1a6..a75d5bd5730 100644 --- a/tests/test_litellm/proxy/proxy_server/test_routes_login_sso.py +++ b/tests/test_litellm/proxy/proxy_server/test_routes_login_sso.py @@ -49,9 +49,7 @@ def _install_login_mocks(monkeypatch, raise_on_auth: bool = False) -> None: } monkeypatch.setattr("litellm.proxy.auth.login_utils.authenticate_user", _fake_auth) - monkeypatch.setattr( - "litellm.proxy.auth.login_utils.create_ui_token_object", _fake_token_object - ) + monkeypatch.setattr("litellm.proxy.auth.login_utils.create_ui_token_object", _fake_token_object) monkeypatch.setattr(ps, "master_key", "sk-test-master") monkeypatch.setattr(ps, "general_settings", {}) monkeypatch.setattr(ps, "premium_user", False) @@ -69,9 +67,7 @@ def test_fallback_login_returns_html_form(client, monkeypatch): body_lower = response.text.lower() shape = { "status": response.status_code, - "content_type_html": response.headers.get("content-type", "").startswith( - "text/html" - ), + "content_type_html": response.headers.get("content-type", "").startswith("text/html"), "has_form": "", "token": ""} + assert normalize(response.json(), volatile=frozenset({"token", "redirect_url"})) == { + "redirect_url": "", + "token": "", + } body = response.json() set_cookie = response.headers.get("set-cookie", "") shape = { "redirect_url_has_ui": "/ui/" in body.get("redirect_url", ""), - "redirect_url_has_login_success": "login=success" - in body.get("redirect_url", ""), + "redirect_url_has_login_success": "login=success" in body.get("redirect_url", ""), "token_in_body": bool(body.get("token")), "token_cookie_set": "token=" in set_cookie, } @@ -264,9 +254,7 @@ def test_v3_login_success_returns_code(client, monkeypatch): from litellm.proxy import proxy_server as ps _install_login_mocks(monkeypatch) - monkeypatch.setattr( - ps, "general_settings", {"control_plane_url": "https://cp.example.invalid"} - ) + monkeypatch.setattr(ps, "general_settings", {"control_plane_url": "https://cp.example.invalid"}) # Force the local (non-redis) cache path monkeypatch.setattr(ps, "redis_usage_cache", None) fake_cache = MagicMock() @@ -301,9 +289,7 @@ def test_v3_login_authenticate_failure_500(client, monkeypatch): from litellm.proxy import proxy_server as ps _install_login_mocks(monkeypatch, raise_on_auth=True) - monkeypatch.setattr( - ps, "general_settings", {"control_plane_url": "https://cp.example.invalid"} - ) + monkeypatch.setattr(ps, "general_settings", {"control_plane_url": "https://cp.example.invalid"}) response = client.post( "/v3/login", @@ -337,9 +323,7 @@ def test_v3_login_exchange_missing_code_400(client, monkeypatch): """Error path: missing 'code' in body -> 400 with 'Missing' message.""" from litellm.proxy import proxy_server as ps - monkeypatch.setattr( - ps, "general_settings", {"control_plane_url": "https://cp.example.invalid"} - ) + monkeypatch.setattr(ps, "general_settings", {"control_plane_url": "https://cp.example.invalid"}) response = client.post("/v3/login/exchange", json={}) assert response.status_code == 400 @@ -352,9 +336,7 @@ def test_v3_login_exchange_invalid_code_401(client, monkeypatch): """Error path: code that isn't in cache -> 401 'Invalid or expired'.""" from litellm.proxy import proxy_server as ps - monkeypatch.setattr( - ps, "general_settings", {"control_plane_url": "https://cp.example.invalid"} - ) + monkeypatch.setattr(ps, "general_settings", {"control_plane_url": "https://cp.example.invalid"}) monkeypatch.setattr(ps, "redis_usage_cache", None) fake_cache = MagicMock() fake_cache.async_get_cache = AsyncMock(return_value=None) @@ -372,9 +354,7 @@ def test_v3_login_exchange_success_returns_token_and_redirect(client, monkeypatc """Pin: valid code -> JSON {token, redirect_url} + token cookie + cache deleted (single-use).""" from litellm.proxy import proxy_server as ps - monkeypatch.setattr( - ps, "general_settings", {"control_plane_url": "https://cp.example.invalid"} - ) + monkeypatch.setattr(ps, "general_settings", {"control_plane_url": "https://cp.example.invalid"}) monkeypatch.setattr(ps, "redis_usage_cache", None) cached_payload = { @@ -388,9 +368,10 @@ def test_v3_login_exchange_success_returns_token_and_redirect(client, monkeypatc response = client.post("/v3/login/exchange", json={"code": "valid-code"}) assert response.status_code == 200 - assert normalize( - response.json(), volatile=frozenset({"token", "redirect_url"}) - ) == {"token": "", "redirect_url": ""} + assert normalize(response.json(), volatile=frozenset({"token", "redirect_url"})) == { + "token": "", + "redirect_url": "", + } body = response.json() set_cookie = response.headers.get("set-cookie", "") shape = { @@ -405,3 +386,77 @@ def test_v3_login_exchange_success_returns_token_and_redirect(client, monkeypatc "token_cookie_set": True, "cache_deleted_once": True, } + + +def test_login_form_honors_same_origin_return_to_cookie(client, monkeypatch): + """The aggregate DCR connect flow preserves a same-origin return_to in the litellm_cp_return_to + cookie; /login must RESUME there after password sign-in instead of dead-ending at the dashboard.""" + _install_login_mocks(monkeypatch) + return_to = "/mcp/authorize?client_id=llm_dcrc_abc&response_type=code" + response = client.post( + "/login", + data={"username": "admin", "password": "password"}, + cookies={"litellm_cp_return_to": return_to}, + follow_redirects=False, + ) + assert response.status_code == 303 + assert response.headers.get("location", "") == return_to # resumed the connect flow, not the dashboard + assert "token=" in response.headers.get("set-cookie", "") + + +def test_login_form_honors_control_plane_return_to_cookie(client, monkeypatch): + """/login resumes through the SAME resumer the SSO callback uses, so it honors BOTH shapes + _persist_return_to_cookie is willing to store. Honoring only the relative one silently dropped + a control-plane return_to and landed the user on the dashboard.""" + import litellm.proxy.proxy_server as ps + + _install_login_mocks(monkeypatch) + monkeypatch.setitem(ps.general_settings, "control_plane_url", "https://cp.example.com") + response = client.post( + "/login", + data={"username": "admin", "password": "password"}, + cookies={"litellm_cp_return_to": "https://cp.example.com/console"}, + follow_redirects=False, + ) + location = response.headers.get("location", "") + assert response.status_code == 303 + assert location.startswith("https://cp.example.com/console") + # Cross-origin arm hands the JWT off via a one-time code rather than a cookie. + assert "code=" in location and "login=success" in location + assert "token=" not in response.headers.get("set-cookie", "") + + +def test_login_form_survives_stale_control_plane_return_to(client, monkeypatch): + """A stale one-shot cookie must NEVER fail a completed sign-in. The resumer rejects a return_to + that no longer matches control_plane_url (a config change between the cookie's write and this + read); the user has already authenticated, so land on the dashboard instead of erroring.""" + import litellm.proxy.proxy_server as ps + + _install_login_mocks(monkeypatch) + monkeypatch.setitem(ps.general_settings, "control_plane_url", "https://new-cp.example.com") + response = client.post( + "/login", + data={"username": "admin", "password": "password"}, + cookies={"litellm_cp_return_to": "https://old-cp.example.com/console"}, + follow_redirects=False, + ) + assert response.status_code == 303, "login must not break on a stale return_to cookie" + location = response.headers.get("location", "") + assert "old-cp.example.com" not in location + assert "/ui/" in location + + +def test_login_form_ignores_open_redirect_return_to(client, monkeypatch): + """A non-same-origin return_to (open-redirect attempt) is rejected — /login falls back to the + dashboard rather than honoring an absolute/foreign URL.""" + _install_login_mocks(monkeypatch) + response = client.post( + "/login", + data={"username": "admin", "password": "password"}, + cookies={"litellm_cp_return_to": "https://evil.example.com/steal"}, + follow_redirects=False, + ) + assert response.status_code == 303 + location = response.headers.get("location", "") + assert "evil.example.com" not in location + assert "/ui/" in location # dashboard fallback diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py b/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py index db72a7fb38c..67945436987 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py @@ -1628,6 +1628,226 @@ async def test_ui_view_spend_logs_date_range_filter(client, monkeypatch): assert data["data"][0]["id"] == "log2" +@pytest.mark.asyncio +async def test_ui_view_spend_logs_request_id_lookup_ignores_date_window( + client, monkeypatch +): + """ + LIT-3981: a request_id lookup on the UI route resolves across all time even + when the caller sends a date window that excludes the log (the dashboard + always sends a window). The window is dropped and request_id alone scopes + the query. Pre-fix the window was always applied, so an id from an older + page returned nothing. + """ + today = datetime.datetime.now(timezone.utc) + mock_spend_logs = [ + { + "id": "log_old", + "request_id": "req-old", + "api_key": "sk-test-key", + "user": "test_user_1", + "team_id": "team1", + "spend": 0.05, + "startTime": (today - datetime.timedelta(days=90)).isoformat(), + "model": "gpt-4", + }, + ] + + captured: dict = {} + + def filter_fn(where): + captured["where"] = where + rows = _filter_logs_by_date_range(mock_spend_logs, where) + if where.get("request_id"): + rows = [r for r in rows if r["request_id"] == where["request_id"]] + return rows + + monkeypatch.setattr( + "litellm.proxy.proxy_server.prisma_client", + make_ui_spend_logs_mock_prisma(mock_spend_logs, filter_fn), + ) + + # A 5-day window that EXCLUDES the 90-day-old log, as the dashboard sends. + start_date = (today - datetime.timedelta(days=5)).strftime("%Y-%m-%d %H:%M:%S") + end_date = today.strftime("%Y-%m-%d %H:%M:%S") + + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN + ) + try: + response = client.get( + "/spend/logs/ui", + params={ + "request_id": "req-old", + "start_date": start_date, + "end_date": end_date, + }, + headers={"Authorization": "Bearer sk-test"}, + ) + assert response.status_code == 200 + data = response.json() + assert data["total"] == 1 + assert data["data"][0]["request_id"] == "req-old" + # Query dropped the time window and scoped solely by the primary key. + assert "startTime" not in captured["where"] + assert captured["where"]["request_id"] == "req-old" + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + + +@pytest.mark.asyncio +async def test_ui_view_spend_logs_requires_dates_without_request_id( + client, monkeypatch +): + """The date window stays mandatory on the UI route when no request_id is set.""" + monkeypatch.setattr( + "litellm.proxy.proxy_server.prisma_client", + make_ui_spend_logs_mock_prisma([], lambda where: []), + ) + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN + ) + try: + response = client.get( + "/spend/logs/ui", headers={"Authorization": "Bearer sk-test"} + ) + assert response.status_code == 400 + assert "date" in response.text.lower() + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + + +@pytest.mark.asyncio +async def test_spend_logs_v2_still_requires_dates_with_request_id(client, monkeypatch): + """The public /spend/logs/v2 contract is unchanged: dates remain required even + when request_id is supplied. Only the internal UI route relaxes the window.""" + monkeypatch.setattr( + "litellm.proxy.proxy_server.prisma_client", + make_ui_spend_logs_mock_prisma([], lambda where: []), + ) + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN + ) + try: + response = client.get( + "/spend/logs/v2", + params={"request_id": "req-old"}, + headers={"Authorization": "Bearer sk-test"}, + ) + assert response.status_code == 400 + assert "date" in response.text.lower() + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + + +@pytest.mark.asyncio +async def test_ui_view_spend_logs_request_id_blocks_non_owner(client, monkeypatch): + """A non-admin looking up a request_id they do not own is rejected (403), so + the relaxed date window cannot read another tenant's log by id.""" + + class _ForeignRow: + user = "other_user" + team_id = None + + class _SpendLogs: + async def find_unique(self, where, include=None): + return _ForeignRow() + + class _DB: + def __init__(self): + self.litellm_spendlogs = _SpendLogs() + + class _Prisma: + def __init__(self): + self.db = _DB() + + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", _Prisma()) + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, user_id="user_1" + ) + try: + response = client.get( + "/spend/logs/ui", + params={"request_id": "foreign-req"}, + headers={"Authorization": "Bearer sk-test"}, + ) + assert response.status_code == 403 + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + + +@pytest.mark.asyncio +async def test_ui_view_spend_logs_request_id_owner_scoped_by_id_only( + client, monkeypatch +): + """A non-admin owner looking up their own request_id resolves across all time. + The ownership check authorizes the single row, so the query drops both the date + window and the general user/team scoping and filters by the primary key alone; + without that skip an internal user would have a `user`/`OR` clause added.""" + today = datetime.datetime.now(timezone.utc) + mock_spend_logs = [ + { + "id": "log_old", + "request_id": "req-old", + "api_key": "sk-test-key", + "user": "user_1", + "team_id": "team1", + "spend": 0.05, + "startTime": (today - datetime.timedelta(days=90)).isoformat(), + "model": "gpt-4", + }, + ] + + captured: dict = {} + + def filter_fn(where): + captured["where"] = where + rows = _filter_logs_by_date_range(mock_spend_logs, where) + if where.get("request_id"): + rows = [r for r in rows if r["request_id"] == where["request_id"]] + return rows + + mock_prisma = make_ui_spend_logs_mock_prisma(mock_spend_logs, filter_fn) + + class _OwnedRow: + user = "user_1" + team_id = "team1" + + async def _find_unique(where, include=None): + return _OwnedRow() + + mock_prisma.db.find_unique = _find_unique + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma) + + # A 5-day window that EXCLUDES the 90-day-old log, as the dashboard sends. + start_date = (today - datetime.timedelta(days=5)).strftime("%Y-%m-%d %H:%M:%S") + end_date = today.strftime("%Y-%m-%d %H:%M:%S") + + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, user_id="user_1" + ) + try: + response = client.get( + "/spend/logs/ui", + params={ + "request_id": "req-old", + "start_date": start_date, + "end_date": end_date, + }, + headers={"Authorization": "Bearer sk-test"}, + ) + assert response.status_code == 200 + data = response.json() + assert data["total"] == 1 + assert data["data"][0]["request_id"] == "req-old" + assert "startTime" not in captured["where"] + assert captured["where"]["request_id"] == "req-old" + assert "user" not in captured["where"] + assert "OR" not in captured["where"] + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + + @pytest.mark.asyncio async def test_ui_view_spend_logs_unauthorized(client): # Test without authorization header diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py b/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py index 5d10bb33751..cc1e2943c8f 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py @@ -109,6 +109,143 @@ def test_get_logging_payload_does_not_map_missing_or_zero_cached_tokens(prompt_t assert "cache_read_input_tokens" not in additional_usage_values +def test_get_logging_payload_maps_openai_cache_write_tokens_to_cache_creation_input_tokens(): + additional_usage_values = _get_additional_usage_values_for_usage( + litellm.Usage( + prompt_tokens=1000, + completion_tokens=2, + total_tokens=1002, + prompt_tokens_details={"cached_tokens": 0, "cache_write_tokens": 800}, + ) + ) + + assert additional_usage_values["cache_creation_input_tokens"] == 800 + assert additional_usage_values["prompt_tokens_details"]["cache_write_tokens"] == 800 + + +def test_get_logging_payload_preserves_anthropic_cache_creation_input_tokens(): + additional_usage_values = _get_additional_usage_values_for_usage( + litellm.Usage( + prompt_tokens=1000, + completion_tokens=2, + total_tokens=1002, + cache_creation_input_tokens=300, + ) + ) + + assert additional_usage_values["cache_creation_input_tokens"] == 300 + + +@pytest.mark.parametrize( + "prompt_tokens_details", + [None, {"cached_tokens": 100}, {"cached_tokens": 100, "cache_write_tokens": 0}], +) +def test_get_logging_payload_does_not_map_missing_or_zero_cache_write_tokens(prompt_tokens_details): + additional_usage_values = _get_additional_usage_values_for_usage( + litellm.Usage( + prompt_tokens=10, + completion_tokens=2, + total_tokens=12, + prompt_tokens_details=prompt_tokens_details, + ) + ) + + assert "cache_creation_input_tokens" not in additional_usage_values + + +def _make_standard_logging_payload_with_usage_object(usage_object: dict) -> StandardLoggingPayload: + return StandardLoggingPayload( + id="test-id-responses", + call_type="responses", + stream=False, + response_cost=0.02, + status="success", + total_tokens=1010, + prompt_tokens=1000, + completion_tokens=10, + startTime=1234567890.0, + endTime=1234567891.0, + completionStartTime=None, + model_map_information=StandardLoggingModelInformation(model_map_key="gpt-5.6", model_map_value=None), + model="gpt-5.6", + model_id="model-123", + model_group="openai", + custom_llm_provider="openai", + api_base="https://api.openai.com", + metadata=StandardLoggingMetadata( + user_api_key_hash="test_hash", + user_api_key_alias=None, + user_api_key_team_id=None, + user_api_key_org_id=None, + user_api_key_user_id=None, + user_api_key_team_alias=None, + spend_logs_metadata=None, + requester_ip_address=None, + requester_metadata=None, + user_api_key_end_user_id=None, + usage_object=usage_object, + ), + cache_hit=False, + cache_key=None, + saved_cache_cost=0.0, + request_tags=[], + end_user=None, + requester_ip_address=None, + messages=[], + response={}, + error_str=None, + model_parameters={}, + hidden_params=StandardLoggingHiddenParams( + model_id="model-123", + cache_key=None, + api_base="https://api.openai.com", + response_cost="0.02", + litellm_overhead_time_ms=None, + additional_headers=None, + batch_models=None, + litellm_model_name=None, + usage_object=None, + ), + ) + + +def test_get_logging_payload_maps_responses_api_cache_write_tokens_from_usage_object(): + """Responses API (/v1/responses) usage is not chat-Usage-shaped, so + additional_usage_values can't derive cache tokens from response_obj.usage. + The Admin UI Logs "Cache Creation Tokens" row reads + additional_usage_values.cache_creation_input_tokens, so it must be filled + from the normalized standard_logging usage_object (LIT-4633).""" + standard_logging_payload = _make_standard_logging_payload_with_usage_object( + usage_object={ + "prompt_tokens": 1000, + "completion_tokens": 10, + "total_tokens": 1010, + "prompt_tokens_details": {"cached_tokens": 0, "cache_write_tokens": 800, "cache_creation_tokens": 800}, + } + ) + payload = get_logging_payload( + kwargs={ + "model": "gpt-5.6", + "call_type": "responses", + "litellm_params": {"metadata": {"user_api_key": "test-key"}}, + "standard_logging_object": standard_logging_payload, + }, + response_obj={ + "id": "resp-test", + "usage": { + "input_tokens": 1000, + "output_tokens": 10, + "total_tokens": 1010, + "input_tokens_details": {"cached_tokens": 0, "cache_write_tokens": 800}, + }, + }, + start_time=datetime.datetime.now(timezone.utc), + end_time=datetime.datetime.now(timezone.utc), + ) + additional_usage_values = json.loads(payload["metadata"])["additional_usage_values"] + assert additional_usage_values["cache_creation_input_tokens"] == 800 + + def test_sanitize_request_body_for_spend_logs_payload_basic(): request_body = { "messages": [{"role": "user", "content": "Hello, how are you?"}], diff --git a/tests/test_litellm/proxy/test_budget_reservation.py b/tests/test_litellm/proxy/test_budget_reservation.py index 1db76aed61d..e6ccbc579b5 100644 --- a/tests/test_litellm/proxy/test_budget_reservation.py +++ b/tests/test_litellm/proxy/test_budget_reservation.py @@ -3,6 +3,7 @@ from datetime import datetime, timedelta, timezone from unittest.mock import AsyncMock, MagicMock, patch import pytest +from fastapi import HTTPException import litellm from litellm.caching.dual_cache import DualCache @@ -1562,6 +1563,100 @@ async def test_should_skip_reservation_when_counter_increment_fails( ) +@pytest.mark.asyncio +async def test_should_raise_503_when_counter_increment_fails_and_fail_closed( + spend_counter_state, + monkeypatch, +): + """#33923: with fail_closed_budget_enforcement on, a failed reservation write + must reject instead of silently degrading to read-time-only enforcement.""" + counter_cache, key_cache = spend_counter_state + proxy_logging_obj = ProxyLogging(user_api_key_cache=key_cache) + valid_token = UserAPIKeyAuth( + token="key-budget-reserve-fail-closed", + spend=0.0, + max_budget=1.0, + ) + + async def fail_increment_cache(*args, **kwargs): + raise RuntimeError("counter unavailable") + + monkeypatch.setattr(counter_cache, "async_increment_cache", fail_increment_cache) + + with patch( + "litellm.proxy.spend_tracking.budget_reservation.estimate_request_max_cost", + return_value=0.5, + ): + with pytest.raises(HTTPException) as exc_info: + await reserve_budget_for_request( + request_body=_request_body(), + route="/chat/completions", + llm_router=None, + valid_token=valid_token, + team_object=None, + user_object=None, + prisma_client=None, + user_api_key_cache=key_cache, + proxy_logging_obj=proxy_logging_obj, + fail_closed_budget_enforcement=True, + ) + + assert exc_info.value.status_code == 503 + assert ( + counter_cache.in_memory_cache.get_cache( + key="spend:key:key-budget-reserve-fail-closed" + ) + is None + ) + + +@pytest.mark.asyncio +async def test_fail_closed_releases_earlier_counters_before_503( + spend_counter_state, +): + """#33923: when a later counter's reservation write fails in strict mode, the + counters that already reserved must be released before the 503 propagates.""" + counter_cache, key_cache = spend_counter_state + proxy_logging_obj = ProxyLogging(user_api_key_cache=key_cache) + valid_token = UserAPIKeyAuth( + token="key-budget-fail-closed-release", + spend=0.0, + max_budget=1.0, + budget_limits=[ + { + "budget_duration": "1h", + "max_budget": 1.0, + } + ], + ) + + with patch( + "litellm.proxy.spend_tracking.budget_reservation.estimate_request_max_cost", + return_value=0.5, + ): + with pytest.raises(HTTPException) as exc_info: + await reserve_budget_for_request( + request_body=_request_body(), + route="/chat/completions", + llm_router=None, + valid_token=valid_token, + team_object=None, + user_object=None, + prisma_client=None, + user_api_key_cache=key_cache, + proxy_logging_obj=proxy_logging_obj, + fail_closed_budget_enforcement=True, + ) + + assert exc_info.value.status_code == 503 + assert ( + counter_cache.in_memory_cache.get_cache( + key="spend:key:key-budget-fail-closed-release" + ) + == 0.0 + ) + + @pytest.mark.asyncio async def test_should_skip_reservation_when_counter_initialization_fails( spend_counter_state, diff --git a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py index edb5997b94a..2f2271514f1 100644 --- a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py +++ b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py @@ -5688,3 +5688,198 @@ async def test_add_litellm_data_to_request_unions_metadata_tags_with_header_tags tags = updated["litellm_metadata"]["tags"] assert "header-tag" in tags assert "body-tag" in tags + + +def _make_chat_request_mock() -> MagicMock: + return _make_request_mock("/v1/chat/completions", {"Content-Type": "application/json"}) + + +@pytest.mark.asyncio +async def test_overwrite_user_with_key_hash_clobbers_caller_supplied_user(monkeypatch): + """The flag exists so providers can ban by a tamper-proof id; a caller-chosen + `user` must never survive, and the raw sk- key must never be forwarded.""" + from litellm.proxy._types import hash_token + + monkeypatch.setattr(litellm, "overwrite_user_with_key_hash", True) + + raw_key = "sk-overwrite-user-test-1234" + user_api_key_dict = UserAPIKeyAuth(api_key=raw_key) + user_api_key_dict.via_virtual_key = True + data = {"model": "gpt-4o", "user": "attacker-chosen-id"} + + updated_data = await add_litellm_data_to_request( + data=data, + request=_make_chat_request_mock(), + user_api_key_dict=user_api_key_dict, + proxy_config=MagicMock(), + general_settings={}, + version="test-version", + ) + + assert updated_data["user"] == hash_token(raw_key) + assert updated_data["user"] != "attacker-chosen-id" + assert raw_key not in updated_data["user"] + + +@pytest.mark.asyncio +async def test_overwrite_user_with_key_hash_sets_user_when_absent(monkeypatch): + from litellm.proxy._types import hash_token + + monkeypatch.setattr(litellm, "overwrite_user_with_key_hash", True) + + raw_key = "sk-overwrite-user-test-5678" + user_api_key_dict = UserAPIKeyAuth(api_key=raw_key) + user_api_key_dict.via_virtual_key = True + data = {"model": "gpt-4o"} + + updated_data = await add_litellm_data_to_request( + data=data, + request=_make_chat_request_mock(), + user_api_key_dict=user_api_key_dict, + proxy_config=MagicMock(), + general_settings={}, + version="test-version", + ) + + assert updated_data["user"] == hash_token(raw_key) + + +@pytest.mark.asyncio +async def test_overwrite_user_with_key_hash_disabled_preserves_caller_user(): + assert litellm.overwrite_user_with_key_hash is False + + user_api_key_dict = UserAPIKeyAuth(api_key="sk-overwrite-user-test-9999") + user_api_key_dict.via_virtual_key = True + data = {"model": "gpt-4o", "user": "caller-chosen-id"} + + updated_data = await add_litellm_data_to_request( + data=data, + request=_make_chat_request_mock(), + user_api_key_dict=user_api_key_dict, + proxy_config=MagicMock(), + general_settings={}, + version="test-version", + ) + + assert updated_data["user"] == "caller-chosen-id" + + +@pytest.mark.asyncio +async def test_overwrite_user_with_key_hash_skips_custom_auth_credential(monkeypatch): + """Custom-auth credentials are not sk-prefixed or JWTs, so UserAPIKeyAuth stores + them raw; the stamp must skip them entirely so auth material never leaks.""" + monkeypatch.setattr(litellm, "overwrite_user_with_key_hash", True) + + raw_credential = "my-custom-auth-credential-abc123" + user_api_key_dict = UserAPIKeyAuth(api_key=raw_credential) + assert user_api_key_dict.api_key == raw_credential + + updated_data = await add_litellm_data_to_request( + data={"model": "gpt-4o", "user": "caller-chosen-id"}, + request=_make_chat_request_mock(), + user_api_key_dict=user_api_key_dict, + proxy_config=MagicMock(), + general_settings={}, + version="test-version", + ) + + assert updated_data["user"] == "caller-chosen-id" + + +@pytest.mark.asyncio +async def test_overwrite_user_with_key_hash_skips_jwt_auth(monkeypatch): + """A hashed JWT rotates on every token re-issue, so it is useless as a stable + ban id; JWT-authenticated requests are not stamped.""" + from litellm.proxy._types import hash_token + + monkeypatch.setattr(litellm, "overwrite_user_with_key_hash", True) + + hashed_jwt = f"hashed-jwt-{hash_token('some-jwt-token')}" + user_api_key_dict = UserAPIKeyAuth(api_key=hashed_jwt) + + updated_data = await add_litellm_data_to_request( + data={"model": "gpt-4o", "user": "caller-chosen-id"}, + request=_make_chat_request_mock(), + user_api_key_dict=user_api_key_dict, + proxy_config=MagicMock(), + general_settings={}, + version="test-version", + ) + + assert updated_data["user"] == "caller-chosen-id" + + +@pytest.mark.asyncio +async def test_overwrite_user_with_key_hash_skips_hex_shaped_custom_credential(monkeypatch): + """A custom-auth credential that happens to be 64 hex chars is indistinguishable + from a key hash by shape alone; only the server-set via_virtual_key marker may + authorize stamping, so this raw credential must never be forwarded.""" + monkeypatch.setattr(litellm, "overwrite_user_with_key_hash", True) + + hex_shaped_credential = "a" * 64 + user_api_key_dict = UserAPIKeyAuth(api_key=hex_shaped_credential) + assert user_api_key_dict.api_key == hex_shaped_credential + assert user_api_key_dict.via_virtual_key is False + + updated_data = await add_litellm_data_to_request( + data={"model": "gpt-4o", "user": "caller-chosen-id"}, + request=_make_chat_request_mock(), + user_api_key_dict=user_api_key_dict, + proxy_config=MagicMock(), + general_settings={}, + version="test-version", + ) + + assert updated_data["user"] == "caller-chosen-id" + + +def test_via_virtual_key_cannot_be_forged_from_validated_input(): + from_kwargs = UserAPIKeyAuth(api_key="b" * 64, via_virtual_key=True) + assert from_kwargs.via_virtual_key is False + + from_dict = UserAPIKeyAuth.model_validate({"api_key": "b" * 64, "via_virtual_key": True}) + assert from_dict.via_virtual_key is False + + +@pytest.mark.asyncio +async def test_overwrite_user_with_key_hash_stamps_master_key_alias(monkeypatch): + """Master-key requests carry the stable alias instead of a hash (so the master + key never propagates anywhere); the alias is the stampable id for them.""" + from litellm.constants import LITELLM_PROXY_MASTER_KEY_ALIAS + + monkeypatch.setattr(litellm, "overwrite_user_with_key_hash", True) + + user_api_key_dict = UserAPIKeyAuth(api_key=LITELLM_PROXY_MASTER_KEY_ALIAS) + user_api_key_dict.via_virtual_key = True + + updated_data = await add_litellm_data_to_request( + data={"model": "gpt-4o", "user": "attacker-chosen-id"}, + request=_make_chat_request_mock(), + user_api_key_dict=user_api_key_dict, + proxy_config=MagicMock(), + general_settings={}, + version="test-version", + ) + + assert updated_data["user"] == LITELLM_PROXY_MASTER_KEY_ALIAS + + +@pytest.mark.asyncio +async def test_overwrite_user_with_key_hash_rejects_alias_without_marker(monkeypatch): + from litellm.constants import LITELLM_PROXY_MASTER_KEY_ALIAS + + monkeypatch.setattr(litellm, "overwrite_user_with_key_hash", True) + + user_api_key_dict = UserAPIKeyAuth(api_key=LITELLM_PROXY_MASTER_KEY_ALIAS) + assert user_api_key_dict.via_virtual_key is False + + updated_data = await add_litellm_data_to_request( + data={"model": "gpt-4o", "user": "caller-chosen-id"}, + request=_make_chat_request_mock(), + user_api_key_dict=user_api_key_dict, + proxy_config=MagicMock(), + general_settings={}, + version="test-version", + ) + + assert updated_data["user"] == "caller-chosen-id" diff --git a/tests/test_litellm/proxy/test_provider_url_destination_guard.py b/tests/test_litellm/proxy/test_provider_url_destination_guard.py index 51cd76105d0..c8771abbc8e 100644 --- a/tests/test_litellm/proxy/test_provider_url_destination_guard.py +++ b/tests/test_litellm/proxy/test_provider_url_destination_guard.py @@ -39,6 +39,46 @@ class TestRejectUrlValuedDestinations: assert exc_info.value.status_code == 400 assert exc_info.value.detail["param"] == "model" + def test_provider_prefixed_url_rejected(self): + with pytest.raises(HTTPException) as exc_info: + _reject_url_valued_destinations( + {"model": "huggingface/https://attacker.example/v1"} + ) + assert exc_info.value.status_code == 400 + assert exc_info.value.detail["param"] == "model" + + def test_comma_batch_smuggled_url_rejected(self): + with pytest.raises(HTTPException) as exc_info: + _reject_url_valued_destinations( + {"model": "gpt-4,huggingface/https://attacker.example/v1"} + ) + assert exc_info.value.status_code == 400 + assert exc_info.value.detail["param"] == "model" + + def test_provider_prefixed_uppercase_scheme_url_rejected(self): + with pytest.raises(HTTPException) as exc_info: + _reject_url_valued_destinations( + {"model": "huggingface/HTTPS://evil.example/v1"} + ) + assert exc_info.value.status_code == 400 + assert exc_info.value.detail["param"] == "model" + + def test_provider_prefixed_plain_model_passes(self): + _reject_url_valued_destinations({"model": "huggingface/BAAI/bge-small-en"}) + + def test_comma_batch_plain_models_pass(self): + _reject_url_valued_destinations({"model": "gpt-4,huggingface/BAAI/bge-small-en"}) + + def test_provider_prefixed_url_respects_allowlist(self, monkeypatch): + monkeypatch.setattr( + litellm, + "provider_url_destination_allowed_hosts", + ["trusted.example"], + ) + _reject_url_valued_destinations( + {"model": "huggingface/https://trusted.example/v1"} + ) + def test_url_valued_file_id_rejected(self): with pytest.raises(HTTPException) as exc_info: _reject_url_valued_destinations( diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 99a6c946a1e..bad76864ca7 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -127,11 +127,15 @@ def test_login_v2_returns_redirect_url_and_sets_cookie(monkeypatch): general_settings={}, premium_user=False, ) - mock_jwt_encode.assert_called_once_with( - {"user_id": "test-user"}, - "test-master-key", - algorithm="HS256", - ) + mock_jwt_encode.assert_called_once() + payload, secret = mock_jwt_encode.call_args.args + # The UI session token carries a bounded-lifetime `exp` claim (dynamic timestamp), alongside + # the user_id; assert its presence rather than an exact expiry value. + assert payload["user_id"] == "test-user" + assert isinstance(payload.get("exp"), int) and payload["exp"] > 0 + assert set(payload.keys()) == {"user_id", "exp"} + assert secret == "test-master-key" + assert mock_jwt_encode.call_args.kwargs == {"algorithm": "HS256"} def test_login_v2_returns_json_on_proxy_exception(monkeypatch): @@ -1015,6 +1019,102 @@ def test_get_config_custom_callback_api_env_vars(monkeypatch): } +@patch( + "litellm.proxy.common_utils.callback_utils.CustomLogger.get_callback_env_vars", + return_value=["LANGFUSE_PUBLIC_KEY", "LANGFUSE_SECRET_KEY", "LANGFUSE_HOST"], +) +def test_get_config_callbacks_fall_back_to_process_env(mock_env_vars, monkeypatch): + """A callback configured purely via process env vars is surfaced. + + An IaC deployment sets LANGFUSE_* on the gateway and never touches the UI, + so nothing is stored in the config environment_variables overlay. The read + endpoint must still report the live values instead of blanks. + """ + from litellm.proxy.proxy_server import app, proxy_config, user_api_key_auth + + monkeypatch.setenv("LANGFUSE_PUBLIC_KEY", "pk-env-only") + monkeypatch.setenv("LANGFUSE_SECRET_KEY", "sk-env-only") + monkeypatch.setenv("LANGFUSE_HOST", "https://cloud.langfuse.com") + + config_data = { + "litellm_settings": {"success_callback": ["langfuse"]}, + "general_settings": {}, + "environment_variables": {}, + } + mock_router = MagicMock() + mock_router.get_settings.return_value = {} + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", mock_router) + monkeypatch.setattr(proxy_config, "get_config", AsyncMock(return_value=config_data)) + + original_overrides = app.dependency_overrides.copy() + app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-1234" + ) + + client = TestClient(app) + try: + response = client.get("/get/config/callbacks") + finally: + app.dependency_overrides = original_overrides + + assert response.status_code == 200 + langfuse_cb = next( + (cb for cb in response.json()["callbacks"] if cb["name"] == "langfuse"), None + ) + assert langfuse_cb is not None + assert langfuse_cb["variables"] == { + "LANGFUSE_PUBLIC_KEY": "pk-env-only", + "LANGFUSE_SECRET_KEY": "sk-env-only", + "LANGFUSE_HOST": "https://cloud.langfuse.com", + } + + +@patch( + "litellm.proxy.common_utils.callback_utils.CustomLogger.get_callback_env_vars", + return_value=["LANGFUSE_SECRET_KEY", "LANGFUSE_HOST"], +) +def test_get_config_callback_env_secrets_redacted_for_non_admin(mock_env_vars, monkeypatch): + """Surfacing env vars must not widen who can read secret values. + + The callback role gate redacts sensitive keys for anyone below full admin, + and that must hold whether the value came from the stored config or the + process env. A non-secret var (LANGFUSE_HOST) still resolves for context. + """ + from litellm.proxy.proxy_server import app, proxy_config, user_api_key_auth + + monkeypatch.setenv("LANGFUSE_SECRET_KEY", "sk-env-only-secret") + monkeypatch.setenv("LANGFUSE_HOST", "https://cloud.langfuse.com") + + config_data = { + "litellm_settings": {"success_callback": ["langfuse"]}, + "general_settings": {}, + "environment_variables": {}, + } + mock_router = MagicMock() + mock_router.get_settings.return_value = {} + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", mock_router) + monkeypatch.setattr(proxy_config, "get_config", AsyncMock(return_value=config_data)) + + original_overrides = app.dependency_overrides.copy() + app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, api_key="sk-user" + ) + + client = TestClient(app) + try: + response = client.get("/get/config/callbacks") + finally: + app.dependency_overrides = original_overrides + + assert response.status_code == 200 + langfuse_cb = next( + (cb for cb in response.json()["callbacks"] if cb["name"] == "langfuse"), None + ) + assert langfuse_cb is not None + assert langfuse_cb["variables"]["LANGFUSE_SECRET_KEY"] == "REDACTED" + assert langfuse_cb["variables"]["LANGFUSE_HOST"] == "https://cloud.langfuse.com" + + def test_get_config_returns_email_settings(monkeypatch): """ Regression for https://github.com/BerriAI/litellm/issues/19221 @@ -1078,6 +1178,113 @@ def test_get_config_returns_email_settings(monkeypatch): assert "*" in variables["SMTP_PASSWORD"] +def _get_email_alert_variables(monkeypatch, config_data): + from litellm.proxy.proxy_server import app, proxy_config, user_api_key_auth + + mock_router = MagicMock() + mock_router.get_settings.return_value = {} + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", mock_router) + monkeypatch.setattr(proxy_config, "get_config", AsyncMock(return_value=config_data)) + + original_overrides = app.dependency_overrides.copy() + app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-1234" + ) + + client = TestClient(app) + try: + response = client.get("/get/config/callbacks") + finally: + app.dependency_overrides = original_overrides + + assert response.status_code == 200 + email_alert = next((a for a in response.json()["alerts"] if a["name"] == "email"), None) + assert email_alert is not None + return email_alert["variables"] + + +def test_get_config_returns_email_settings_set_only_in_process_env(monkeypatch): + """ + Regression for LIT-4165. + + SMTP supplied purely as process env vars (helm/terraform, no UI writes) is + live at runtime because litellm/proxy/utils.py::send_email resolves every + field from os.getenv. The /get/config/callbacks email block only read the + config/DB environment_variables overlay though, so those deployments saw an + empty Email Server Settings page and could not tell SMTP was configured. + The slack block one branch above already fell back to os.getenv. + """ + smtp_password = "env-only-app-password" + monkeypatch.setenv("SMTP_HOST", "smtp.env-host.com") + monkeypatch.setenv("SMTP_PORT", "2525") + monkeypatch.setenv("SMTP_TLS", "False") + monkeypatch.setenv("SMTP_USERNAME", "env-user") + monkeypatch.setenv("SMTP_PASSWORD", smtp_password) + monkeypatch.setenv("SMTP_SENDER_EMAIL", "alerts@env-host.com") + monkeypatch.setenv("TEST_EMAIL_ADDRESS", "admin@env-host.com") + + variables = _get_email_alert_variables( + monkeypatch, + { + "litellm_settings": {}, + "general_settings": {"alerting": ["email"]}, + "environment_variables": {}, + }, + ) + + # Every one of these was None before the fix, despite SMTP working. + assert variables["SMTP_HOST"] == "smtp.env-host.com" + assert variables["SMTP_PORT"] == "2525" + assert variables["SMTP_TLS"] == "False" + assert variables["SMTP_USERNAME"] == "env-user" + assert variables["SMTP_SENDER_EMAIL"] == "alerts@env-host.com" + assert variables["TEST_EMAIL_ADDRESS"] == "admin@env-host.com" + + # An env-sourced secret is masked exactly like a stored one. + assert variables["SMTP_PASSWORD"] not in (None, smtp_password) + assert "*" in variables["SMTP_PASSWORD"] + + +def test_get_config_email_settings_prefer_stored_over_process_env(monkeypatch): + """ + Stored environment_variables win over the process environment, matching the + load order in ProxyConfig.get_config, which pushes stored values into + os.environ. Only a field with no stored entry falls back to os.getenv. + """ + monkeypatch.setenv("SMTP_HOST", "smtp.env-host.com") + monkeypatch.setenv("SMTP_SENDER_EMAIL", "alerts@env-host.com") + + variables = _get_email_alert_variables( + monkeypatch, + { + "litellm_settings": {}, + "general_settings": {"alerting": ["email"]}, + "environment_variables": {"SMTP_HOST": "smtp.stored-host.com"}, + }, + ) + + assert variables["SMTP_HOST"] == "smtp.stored-host.com" + assert variables["SMTP_SENDER_EMAIL"] == "alerts@env-host.com" + + +def test_get_config_email_settings_absent_everywhere_stay_none(monkeypatch): + """A field set in neither source is reported unset rather than invented.""" + for var in ("SMTP_HOST", "SMTP_PORT", "SMTP_TLS", "SMTP_USERNAME", "SMTP_PASSWORD", "SMTP_SENDER_EMAIL"): + monkeypatch.delenv(var, raising=False) + + variables = _get_email_alert_variables( + monkeypatch, + { + "litellm_settings": {}, + "general_settings": {"alerting": ["email"]}, + "environment_variables": {}, + }, + ) + + assert variables["SMTP_HOST"] is None + assert variables["SMTP_PASSWORD"] is None + + def test_get_config_returns_slack_webhook(monkeypatch): """ Same double-decryption regression as the email block (issue #19221): the @@ -2679,6 +2886,69 @@ async def test_load_config_max_budget_env_var_coerced_to_float(tmp_path, monkeyp litellm.max_budget = original_max_budget +def test_max_ui_session_budget_default_is_one_dollar(): + """LIT-4662: the dashboard session budget default is a product decision; the + old 0.25 default locked admins out of auto router Test Connection and the + playground mid-session with an error that looked like a hardcoded cap.""" + assert litellm.max_ui_session_budget == 1.0 + + +@pytest.mark.asyncio +async def test_load_config_max_ui_session_budget_applied_and_coerced(tmp_path, monkeypatch): + """ + max_ui_session_budget configured via os.environ resolves to a string; + load_config must coerce it to float so every dashboard session key is + minted with a numeric max_budget. + """ + from litellm.proxy.proxy_server import ProxyConfig + + monkeypatch.setenv("UI_SESSION_BUDGET", "2.5") + test_config = { + "model_list": [], + "litellm_settings": {"max_ui_session_budget": "os.environ/UI_SESSION_BUDGET"}, + } + config_file = tmp_path / "config.yaml" + config_file.write_text(yaml.dump(test_config)) + + original_budget = litellm.max_ui_session_budget + try: + proxy_config = ProxyConfig() + await proxy_config.load_config( + router=MagicMock(), config_file_path=str(config_file) + ) + assert isinstance(litellm.max_ui_session_budget, float) + assert litellm.max_ui_session_budget == 2.5 + finally: + litellm.max_ui_session_budget = original_budget + + +@pytest.mark.asyncio +async def test_load_config_max_ui_session_budget_none_disables_cap(tmp_path): + """ + max_ui_session_budget: null in config disables the dashboard session cap + entirely (session keys minted with no max_budget); load_config must pass + None through instead of raising on float(None). + """ + from litellm.proxy.proxy_server import ProxyConfig + + test_config = { + "model_list": [], + "litellm_settings": {"max_ui_session_budget": None}, + } + config_file = tmp_path / "config.yaml" + config_file.write_text(yaml.dump(test_config)) + + original_budget = litellm.max_ui_session_budget + try: + proxy_config = ProxyConfig() + await proxy_config.load_config( + router=MagicMock(), config_file_path=str(config_file) + ) + assert litellm.max_ui_session_budget is None + finally: + litellm.max_ui_session_budget = original_budget + + @pytest.mark.asyncio async def test_load_config_default_internal_user_params_max_budget_scientific_notation(tmp_path): """ @@ -9139,6 +9409,85 @@ def test_general_settings_ui_fields_are_db_overridable(): ) +@pytest.mark.asyncio +async def test_update_config_field_max_ui_session_budget_sets_live_value(monkeypatch): + """LIT-4662: the dashboard session budget is editable from the Admin UI General tab. + A Dollar field must accept values above 1 (the old Float type capped at 1, which cannot + express a dollar budget), apply live via setattr, and persist under litellm_settings.""" + from unittest.mock import MagicMock + + import litellm.proxy.proxy_server as ps + from litellm.proxy._types import ( + ConfigFieldUpdate, + LitellmUserRoles, + UserAPIKeyAuth, + ) + from litellm.proxy.proxy_server import update_config_general_settings + + saved: dict = {} + + async def fake_get_config(): + return {"litellm_settings": {}} + + async def fake_save_config(new_config=None): + saved.update(new_config or {}) + + monkeypatch.setattr(ps.proxy_config, "get_config", fake_get_config) + monkeypatch.setattr(ps.proxy_config, "save_config", fake_save_config) + monkeypatch.setattr(ps, "prisma_client", MagicMock()) + monkeypatch.setattr(litellm, "store_audit_logs", False) + monkeypatch.setattr(litellm, "max_ui_session_budget", 1.0) + + admin = UserAPIKeyAuth(api_key="k", user_id="a", user_role=LitellmUserRoles.PROXY_ADMIN) + await update_config_general_settings( + data=ConfigFieldUpdate( + field_name="max_ui_session_budget", + field_value=25.0, + config_type="general_settings", + ), + user_api_key_dict=admin, + ) + + assert litellm.max_ui_session_budget == 25.0 + assert saved["litellm_settings"]["max_ui_session_budget"] == 25.0 + + +@pytest.mark.parametrize("bad_value", [True, "abc", -1, 0, [2.5]]) +def test_validate_max_ui_session_budget_rejects_malformed(bad_value): + """A Dollar field accepts only positive numbers; zero would block every dashboard + LLM call at mint and non-numerics would break session key generation.""" + from fastapi import HTTPException + + from litellm.proxy.proxy_server import _validate_general_settings_ui_litellm_value + + with pytest.raises(HTTPException) as exc_info: + _validate_general_settings_ui_litellm_value("max_ui_session_budget", bad_value) + assert exc_info.value.status_code == 400 + + +@pytest.mark.parametrize("empty_value", [None, ""]) +def test_validate_max_ui_session_budget_empty_restores_default(empty_value): + """Clearing the field in the UI restores the shipped $1 default rather than None; + None would silently remove the session spend guardrail (unlimited budget), which + must stay a deliberate config.yaml act (max_ui_session_budget: null).""" + from litellm.proxy.proxy_server import _validate_general_settings_ui_litellm_value + + assert _validate_general_settings_ui_litellm_value("max_ui_session_budget", empty_value) == 1.0 + + +def test_general_settings_ui_defaults_unchanged_for_existing_fields(): + """The spec-default mechanism added for max_ui_session_budget must not change what + clearing the pre-existing fields restores (None for Float/Select, False for Boolean).""" + from litellm.proxy.proxy_server import ( + _GENERAL_SETTINGS_UI_LITELLM_FIELDS, + _general_settings_ui_litellm_default, + ) + + assert _general_settings_ui_litellm_default(_GENERAL_SETTINGS_UI_LITELLM_FIELDS["budget_exceeded_throttle_percentage"]) is None + assert _general_settings_ui_litellm_default(_GENERAL_SETTINGS_UI_LITELLM_FIELDS["enable_anthropic_prompt_caching"]) is False + assert _general_settings_ui_litellm_default(_GENERAL_SETTINGS_UI_LITELLM_FIELDS["anthropic_prompt_caching_ttl"]) is None + + @pytest.mark.parametrize( "field_name, db_value", [ diff --git a/tests/test_litellm/proxy/test_proxy_utils.py b/tests/test_litellm/proxy/test_proxy_utils.py index 5ace46fc775..4673807a135 100644 --- a/tests/test_litellm/proxy/test_proxy_utils.py +++ b/tests/test_litellm/proxy/test_proxy_utils.py @@ -709,6 +709,39 @@ def test_create_model_info_response_reads_real_cost_map(): assert response["max_output_tokens"] > 0 +def test_create_model_info_response_includes_mode_from_lookup(): + response = create_model_info_response( + model_id="text-embedding-3-small", + provider="openai", + llm_router=None, + get_model_info=lambda _model: _fake_model_info(mode="embedding"), + ) + + assert response["mode"] == "embedding" + + +def test_create_model_info_response_omits_mode_when_lookup_raises(): + response = create_model_info_response( + model_id="my-custom-deployment", + provider="openai", + llm_router=None, + get_model_info=_raise_unmapped, + ) + + assert "mode" not in response + + +def test_create_model_info_response_omits_non_string_mode(): + response = create_model_info_response( + model_id="some-model", + provider="openai", + llm_router=None, + get_model_info=lambda _model: _fake_model_info(mode=None), + ) + + assert "mode" not in response + + class TestPostCallFailureHookLLMExceptionAlerting: """The llm_exceptions alert is for infra / LLM-API failures, not user errors (https://github.com/BerriAI/litellm/issues/3395). Already-normalized diff --git a/tests/test_litellm/proxy/test_redis_auth_cache_flag.py b/tests/test_litellm/proxy/test_redis_auth_cache_flag.py index d0cb5ec5465..849d5494c6e 100644 --- a/tests/test_litellm/proxy/test_redis_auth_cache_flag.py +++ b/tests/test_litellm/proxy/test_redis_auth_cache_flag.py @@ -54,8 +54,8 @@ def _patched_init_cache(litellm_settings: dict, cache_params: dict): _FakeRedisCache (passes the isinstance guard in _init_cache). 3. Extracts enable_redis_auth_cache from litellm_settings and passes it as the second argument to _init_cache (matching production behaviour). - 4. Yields (user_api_key_cache, spend_counter_cache) after calling - _init_cache, then restores everything. + 4. Yields (user_api_key_cache, spend_counter_cache, cli_sso_session_cache) + after calling _init_cache, then restores everything. """ fake_redis = _FakeRedisCache() @@ -64,19 +64,21 @@ def _patched_init_cache(litellm_settings: dict, cache_params: dict): fresh_user_cache = DualCache() fresh_spend_cache = DualCache() + fresh_cli_sso_cache = DualCache() enable_redis_auth_cache = litellm_settings.get("enable_redis_auth_cache", False) with ( patch.object(ps, "user_api_key_cache", fresh_user_cache), patch.object(ps, "spend_counter_cache", fresh_spend_cache), + patch.object(ps, "cli_sso_session_cache", fresh_cli_sso_cache), patch.object(ps, "llm_router", None), # Cache is locally imported inside _init_cache: patch it at source. patch("litellm.Cache", return_value=mock_litellm_cache), ): litellm.cache = None ps.ProxyConfig()._init_cache(cache_params, enable_redis_auth_cache) - yield fresh_user_cache, fresh_spend_cache + yield fresh_user_cache, fresh_spend_cache, fresh_cli_sso_cache # --------------------------------------------------------------------------- @@ -90,7 +92,7 @@ class TestRedisAuthCacheFlag: with _patched_init_cache( litellm_settings={"enable_redis_auth_cache": True}, cache_params={"type": "redis", "host": "localhost", "port": 6379}, - ) as (user_cache, _): + ) as (user_cache, _, _cli_sso_cache): assert user_cache.redis_cache is not None, ( "Redis should be attached to user_api_key_cache when " "enable_redis_auth_cache=True" @@ -101,7 +103,7 @@ class TestRedisAuthCacheFlag: with _patched_init_cache( litellm_settings={"enable_redis_auth_cache": False}, cache_params={"type": "redis", "host": "localhost", "port": 6379}, - ) as (user_cache, _): + ) as (user_cache, _, _cli_sso_cache): assert user_cache.redis_cache is None, ( "user_api_key_cache must remain in-memory-only when " "enable_redis_auth_cache=False" @@ -112,7 +114,7 @@ class TestRedisAuthCacheFlag: with _patched_init_cache( litellm_settings={}, cache_params={"type": "redis", "host": "localhost", "port": 6379}, - ) as (user_cache, _): + ) as (user_cache, _, _cli_sso_cache): assert user_cache.redis_cache is None, ( "user_api_key_cache must remain in-memory-only when " "enable_redis_auth_cache is absent from litellm_settings" @@ -129,7 +131,7 @@ class TestRedisAuthCacheFlag: with _patched_init_cache( litellm_settings=ls, cache_params={"type": "redis", "host": "localhost", "port": 6379}, - ) as (_, spend_cache): + ) as (_, spend_cache, _cli_sso_cache): assert spend_cache.redis_cache is not None, ( f"spend_counter_cache must always get Redis " f"(enable_redis_auth_cache={flag_value!r})" @@ -140,6 +142,28 @@ class TestRedisAuthCacheFlag: with _patched_init_cache( litellm_settings={"enable_redis_auth_cache": False}, cache_params={"type": "redis", "host": "localhost", "port": 6379}, - ) as (user_cache, spend_cache): + ) as (user_cache, spend_cache, _cli_sso_cache): assert spend_cache.redis_cache is not None assert user_cache.redis_cache is None + + def test_cli_sso_session_cache_always_gets_redis_regardless_of_flag(self): + """ + cli_sso_session_cache must receive Redis regardless of the auth-cache + flag so that `lite login` works on multi-worker deployments without + enable_redis_auth_cache (regression for the CLI SSO "Invalid CLI login + session" bug) + """ + for flag_value in (True, False, None): + ls = ( + {"enable_redis_auth_cache": flag_value} + if flag_value is not None + else {} + ) + with _patched_init_cache( + litellm_settings=ls, + cache_params={"type": "redis", "host": "localhost", "port": 6379}, + ) as (_, _, cli_sso_cache): + assert cli_sso_cache.redis_cache is not None, ( + f"cli_sso_session_cache must always get Redis " + f"(enable_redis_auth_cache={flag_value!r})" + ) diff --git a/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py b/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py index 671a6cd39ba..20451f5d0ac 100644 --- a/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py +++ b/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py @@ -396,6 +396,146 @@ class TestProxySettingEndpoints: call_args = mock_prisma.db.litellm_ssoconfig.find_unique.call_args assert call_args.kwargs["where"]["id"] == "sso_config" + def _mock_sso_db_record(self, monkeypatch, sso_settings): + """Point /get/sso_settings at a stored SSO row (or None for no row).""" + from unittest.mock import AsyncMock, MagicMock + + mock_prisma = MagicMock() + if sso_settings is None: + mock_db_record = None + else: + mock_db_record = MagicMock() + mock_db_record.sso_settings = sso_settings + mock_prisma.db.litellm_ssoconfig.find_unique = AsyncMock(return_value=mock_db_record) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma) + + # The resolver decrypts stored values via decrypt_value_helper; make it an + # identity so the plaintext fixtures round-trip. + monkeypatch.setattr( + "litellm.proxy.config_resolvers.sso.decrypt_value_helper", + lambda value, key, exception_type="error", return_original_value=False: value, + ) + + def test_get_sso_settings_falls_back_to_process_env( + self, mock_proxy_config, mock_auth, monkeypatch + ): + """ + Regression for LIT-4165. + + SSO configured purely as process env vars (helm/terraform, no UI writes) + logs users in successfully, because ui_sso.py resolves every setting from + os.environ. /get/sso_settings read only the sso_config table though, so + the Admin UI showed "not configured" for a working SSO deployment and hid + the Edit/Delete controls behind an empty-state placeholder. + """ + self._mock_sso_db_record(monkeypatch, None) + monkeypatch.setenv("GENERIC_CLIENT_ID", "env-client-id") + monkeypatch.setenv("GENERIC_CLIENT_SECRET", "env-client-secret-value") + monkeypatch.setenv("GENERIC_AUTHORIZATION_ENDPOINT", "https://idp.example.com/authorize") + monkeypatch.setenv("GENERIC_TOKEN_ENDPOINT", "https://idp.example.com/token") + monkeypatch.setenv("GENERIC_USERINFO_ENDPOINT", "https://idp.example.com/userinfo") + monkeypatch.setenv("GENERIC_SCOPE", "openid email profile groups") + monkeypatch.setenv("PROXY_BASE_URL", "https://gateway.example.com") + + response = client.get("/get/sso_settings") + + assert response.status_code == 200 + values = response.json()["values"] + + # Every one of these was None before the fix, despite SSO working. + assert values["generic_client_id"] == "env-client-id" + assert values["generic_authorization_endpoint"] == "https://idp.example.com/authorize" + assert values["generic_token_endpoint"] == "https://idp.example.com/token" + assert values["generic_userinfo_endpoint"] == "https://idp.example.com/userinfo" + assert values["generic_scope"] == "openid email profile groups" + assert values["proxy_base_url"] == "https://gateway.example.com" + + # An env-sourced secret is masked exactly like a stored one. + assert values["generic_client_secret"] not in (None, "env-client-secret-value") + assert "*" in values["generic_client_secret"] + + def test_get_sso_settings_does_not_mutate_os_environ( + self, mock_proxy_config, mock_auth, monkeypatch + ): + """A GET must not write os.environ. The legacy read path decrypted DB + values straight into the environment, so opening the settings page + repopulated env and masked any consumer that stopped reading it.""" + self._mock_sso_db_record(monkeypatch, {"generic_client_id": "db-only-id"}) + monkeypatch.delenv("GENERIC_CLIENT_ID", raising=False) + + response = client.get("/get/sso_settings") + + assert response.status_code == 200 + assert response.json()["values"]["generic_client_id"] == "db-only-id" + # The DB value must NOT have leaked into the process environment. + assert "GENERIC_CLIENT_ID" not in os.environ + + def test_get_sso_settings_prefers_stored_over_process_env( + self, mock_proxy_config, mock_auth, monkeypatch + ): + """A stored value wins; only fields absent from the row fall back to env.""" + self._mock_sso_db_record(monkeypatch, {"generic_client_id": "stored-client-id"}) + monkeypatch.setenv("GENERIC_CLIENT_ID", "env-client-id") + monkeypatch.setenv("GENERIC_TOKEN_ENDPOINT", "https://idp.example.com/token") + + response = client.get("/get/sso_settings") + + assert response.status_code == 200 + values = response.json()["values"] + assert values["generic_client_id"] == "stored-client-id" + assert values["generic_token_endpoint"] == "https://idp.example.com/token" + + def test_get_sso_settings_blank_stored_value_falls_back_to_process_env( + self, mock_proxy_config, mock_auth, monkeypatch + ): + """ + Blank means absent. update_sso_settings clears the env var for a blank + field, so a blank row entry cannot describe a live setting; os.environ is + the effective config and is what the UI must report. + """ + self._mock_sso_db_record(monkeypatch, {"generic_client_id": " ", "generic_token_endpoint": ""}) + monkeypatch.setenv("GENERIC_CLIENT_ID", "env-client-id") + monkeypatch.setenv("GENERIC_TOKEN_ENDPOINT", "https://idp.example.com/token") + + response = client.get("/get/sso_settings") + + assert response.status_code == 200 + values = response.json()["values"] + assert values["generic_client_id"] == "env-client-id" + assert values["generic_token_endpoint"] == "https://idp.example.com/token" + + def test_get_sso_settings_unset_everywhere_reports_source( + self, mock_proxy_config, mock_auth, monkeypatch + ): + """A field set in neither source is unset (or its effective default), + and provenance reports which.""" + self._mock_sso_db_record(monkeypatch, None) + for env_var in ( + "GENERIC_CLIENT_ID", + "GENERIC_CLIENT_SECRET", + "GENERIC_TOKEN_ENDPOINT", + "GENERIC_SCOPE", + "GOOGLE_CLIENT_ID", + "MICROSOFT_CLIENT_ID", + "PROXY_BASE_URL", + ): + monkeypatch.delenv(env_var, raising=False) + + response = client.get("/get/sso_settings") + + assert response.status_code == 200 + body = response.json() + values = body["values"] + provenance = body["provenance"] + assert values["generic_client_id"] is None + assert provenance["generic_client_id"] == "unset" + assert values["generic_client_secret"] is None + assert values["google_client_id"] is None + # generic_scope carries the same effective default the login path applies, + # so the settings page shows the scope logins would actually request. + assert values["generic_scope"] == "openid email profile" + assert provenance["generic_scope"] == "default" + def test_update_sso_settings(self, mock_proxy_config, mock_auth, monkeypatch): """Test updating the SSO settings to the dedicated database table""" import json @@ -477,6 +617,75 @@ class TestProxySettingEndpoints: create_sso_settings = json.loads(create_data["sso_settings"]) assert create_sso_settings["google_client_id"] == "new_google_client_id" + def test_update_sso_settings_maps_saml_fields_to_env_vars( + self, mock_proxy_config, mock_auth, monkeypatch + ): + """SAML settings entered in the admin UI must be applied as the SAML_* env + vars the SAML handler reads, and the allow-unsolicited toggle must map to + the 'true'/'false' string the handler expects.""" + import json + import os + from unittest.mock import AsyncMock, MagicMock + + monkeypatch.setenv("LITELLM_SALT_KEY", "test_salt_key") + monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", True) + + mock_prisma = MagicMock() + mock_prisma.db.litellm_ssoconfig.find_unique = AsyncMock(return_value=None) + mock_prisma.db.litellm_ssoconfig.upsert = AsyncMock() + mock_prisma.db.litellm_config = MagicMock() + mock_prisma.db.litellm_config.find_unique = AsyncMock(return_value=None) + mock_prisma.db.litellm_config.update = AsyncMock() + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma) + + from litellm.proxy.proxy_server import proxy_config + + monkeypatch.setattr( + proxy_config, + "_encrypt_env_variables", + lambda environment_variables: environment_variables, + ) + + for var in ( + "SAML_IDP_METADATA_URL", + "SAML_IDP_METADATA_XML", + "SAML_SP_ENTITY_ID", + "SAML_ALLOW_UNSOLICITED", + ): + monkeypatch.delenv(var, raising=False) + + new_sso_settings = { + "saml_idp_metadata_url": "https://idp.example.com/metadata", + "saml_sp_entity_id": "https://proxy.example.com/sso/saml/metadata", + "saml_allow_unsolicited": "true", + "proxy_base_url": "https://proxy.example.com", + "user_email": "admin@example.com", + } + + try: + response = client.patch("/update/sso_settings", json=new_sso_settings) + + assert response.status_code == 200 + + assert os.environ.get("SAML_IDP_METADATA_URL") == "https://idp.example.com/metadata" + assert os.environ.get("SAML_SP_ENTITY_ID") == "https://proxy.example.com/sso/saml/metadata" + assert os.environ.get("SAML_ALLOW_UNSOLICITED") == "true" + assert "SAML_IDP_METADATA_XML" not in os.environ + + stored = json.loads( + mock_prisma.db.litellm_ssoconfig.upsert.call_args.kwargs["data"]["create"]["sso_settings"] + ) + assert stored["saml_idp_metadata_url"] == "https://idp.example.com/metadata" + assert stored["saml_allow_unsolicited"] == "true" + finally: + for var in ( + "SAML_IDP_METADATA_URL", + "SAML_IDP_METADATA_XML", + "SAML_SP_ENTITY_ID", + "SAML_ALLOW_UNSOLICITED", + ): + os.environ.pop(var, None) + def test_update_sso_settings_audits_when_env_cleanup_fails( self, mock_proxy_config, mock_auth, monkeypatch ): @@ -944,6 +1153,88 @@ class TestProxySettingEndpoints: assert data["values"]["logo_url"] == "https://example.com/logo.png" assert data["values"]["favicon_url"] == "https://example.com/favicon.ico" + def test_get_ui_theme_settings_falls_back_to_process_env( + self, mock_proxy_config, monkeypatch + ): + """Branding supplied only as process env vars must surface in the read. + + A deployment that sets UI_LOGO_PATH / LITELLM_FAVICON_URL via IaC and + never touches the UI has no stored ui_theme_config, yet the branding is + live, so the settings page must reflect it rather than reading blank. + """ + monkeypatch.delenv("UI_LOGO_PATH", raising=False) + monkeypatch.delenv("LITELLM_FAVICON_URL", raising=False) + monkeypatch.setenv("UI_LOGO_PATH", "https://cdn.example.com/logo.png") + monkeypatch.setenv("LITELLM_FAVICON_URL", "https://cdn.example.com/favicon.ico") + + response = client.get("/get/ui_theme_settings") + + assert response.status_code == 200 + values = response.json()["values"] + assert values["logo_url"] == "https://cdn.example.com/logo.png" + assert values["favicon_url"] == "https://cdn.example.com/favicon.ico" + + def test_get_ui_theme_settings_stored_value_wins_over_env( + self, mock_auth, monkeypatch + ): + """A stored ui_theme_config field outranks the env var for that field. + + The env fallback only fills fields the stored config leaves blank, so the + UI-driven flow is unchanged while an unstored field still resolves. + """ + from litellm.proxy.proxy_server import proxy_config + + stored_config = { + "litellm_settings": { + "ui_theme_config": {"logo_url": "https://db.example.com/logo.png"} + } + } + + async def mock_get_config(): + return stored_config + + monkeypatch.setattr(proxy_config, "get_config", mock_get_config) + monkeypatch.setenv("UI_LOGO_PATH", "https://env.example.com/logo.png") + monkeypatch.setenv("LITELLM_FAVICON_URL", "https://env.example.com/favicon.ico") + + response = client.get("/get/ui_theme_settings") + + assert response.status_code == 200 + values = response.json()["values"] + assert values["logo_url"] == "https://db.example.com/logo.png" + assert values["favicon_url"] == "https://env.example.com/favicon.ico" + + def test_get_ui_theme_settings_reports_unset_when_absent_everywhere( + self, mock_proxy_config, monkeypatch + ): + """A field set in neither the stored config nor the env stays null.""" + monkeypatch.delenv("UI_LOGO_PATH", raising=False) + monkeypatch.delenv("LITELLM_FAVICON_URL", raising=False) + + response = client.get("/get/ui_theme_settings") + + assert response.status_code == 200 + values = response.json()["values"] + assert values["logo_url"] is None + assert values["favicon_url"] is None + + def test_get_ui_theme_settings_does_not_disclose_local_path_env_value( + self, mock_proxy_config, monkeypatch + ): + """This endpoint is public, so an env-configured local filesystem branding + path must never be surfaced to anonymous callers; only public http(s) URLs. + """ + monkeypatch.setenv("UI_LOGO_PATH", "/mnt/secret/internal/logo.png") + monkeypatch.setenv("LITELLM_FAVICON_URL", "file:///etc/favicon.ico") + + response = client.get("/get/ui_theme_settings") + + assert response.status_code == 200 + values = response.json()["values"] + # the local path / file scheme is withheld rather than disclosed + assert values["logo_url"] is None + assert values["favicon_url"] is None + def test_get_ui_settings(self, mock_auth, monkeypatch): """Test retrieving UI settings with allowlist sanitization""" from unittest.mock import AsyncMock, MagicMock @@ -1381,19 +1672,20 @@ class TestProxySettingEndpoints: monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma) - # Mock the decryption method to return decrypted values - def mock_decrypt_and_set(environment_variables): - return { - "google_client_id": "decrypted_google_id", - "google_client_secret": "decrypted_google_secret", - "microsoft_client_id": "decrypted_microsoft_id", - "proxy_base_url": "https://decrypted.example.com", - } + # The resolver decrypts each stored value via decrypt_value_helper; map + # the ciphertext fixtures to their plaintext. + decrypted_by_ciphertext = { + "encrypted_google_id": "decrypted_google_id", + "encrypted_google_secret": "decrypted_google_secret", + "encrypted_microsoft_id": "decrypted_microsoft_id", + "encrypted_proxy_url": "https://decrypted.example.com", + } - from litellm.proxy.proxy_server import proxy_config + def mock_decrypt(value, key, exception_type="error", return_original_value=False): + return decrypted_by_ciphertext.get(value, value) monkeypatch.setattr( - proxy_config, "_decrypt_and_set_db_env_variables", mock_decrypt_and_set + "litellm.proxy.config_resolvers.sso.decrypt_value_helper", mock_decrypt ) response = client.get("/get/sso_settings") diff --git a/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py b/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py index 64c14abfd83..5c711fc6c34 100644 --- a/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py +++ b/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py @@ -626,6 +626,7 @@ def _moderation_guardrail() -> MagicMock: cb.should_run_guardrail = MagicMock(return_value=True) cb.async_moderation_hook = AsyncMock(return_value=None) cb.async_post_call_success_hook = AsyncMock(return_value=None) + cb.run_in_parallel = False return cb diff --git a/tests/test_litellm/proxy/utils/proxy_logging/test_post_call_success_hook.py b/tests/test_litellm/proxy/utils/proxy_logging/test_post_call_success_hook.py index 6a339b37a80..715d66db181 100644 --- a/tests/test_litellm/proxy/utils/proxy_logging/test_post_call_success_hook.py +++ b/tests/test_litellm/proxy/utils/proxy_logging/test_post_call_success_hook.py @@ -28,6 +28,7 @@ def _make_guardrail(name="g", should_run=True, override=None): cb.event_hook = GuardrailEventHooks.post_call cb.should_run_guardrail = MagicMock(return_value=should_run) cb.async_post_call_success_hook = AsyncMock(return_value=override) + cb.run_in_parallel = False return cb diff --git a/tests/test_litellm/repositories/test_repositories.py b/tests/test_litellm/repositories/test_repositories.py index af2eea823f4..6308faf8fc7 100644 --- a/tests/test_litellm/repositories/test_repositories.py +++ b/tests/test_litellm/repositories/test_repositories.py @@ -5,7 +5,7 @@ Tests for gateway repository layer. import json from datetime import datetime from typing import Any, Dict, List, Optional -from unittest.mock import MagicMock, patch +from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -499,6 +499,42 @@ class TestTeamRepository: assert team.team_id == "team-123" assert team.team_alias == "Engineering" + @pytest.mark.asyncio + @pytest.mark.parametrize( + "raw_value, expected_ids", + [ + ( + [ + {"user_id": "a", "role": "user"}, + {"user_id": "b", "role": "admin"}, + ], + ["a", "b"], + ), + (json.dumps([{"user_id": "a", "role": "user"}]), ["a"]), + ({}, []), + (None, []), + ], + ) + async def test_get_members_with_roles_locked(self, repo, raw_value, expected_ids): + tx = MagicMock() + tx.query_raw = AsyncMock(return_value=[{"members_with_roles": raw_value}]) + + members = await repo.get_members_with_roles_locked(tx, "team-1") + + assert [m.user_id for m in members] == expected_ids + sql = tx.query_raw.call_args.args[0] + assert "FOR UPDATE" in sql + assert tx.query_raw.call_args.args[1] == "team-1" + + @pytest.mark.asyncio + async def test_get_members_with_roles_locked_missing_row(self, repo): + tx = MagicMock() + tx.query_raw = AsyncMock(return_value=[]) + + members = await repo.get_members_with_roles_locked(tx, "missing") + + assert members == [] + @pytest.mark.asyncio async def test_create_team_all_fields(self, repo): team = await repo.create_team( diff --git a/tests/test_litellm/responses/test_responses_prompt_management.py b/tests/test_litellm/responses/test_responses_prompt_management.py index 84e98390268..7044d8384f8 100644 --- a/tests/test_litellm/responses/test_responses_prompt_management.py +++ b/tests/test_litellm/responses/test_responses_prompt_management.py @@ -14,13 +14,19 @@ Covers: """ import asyncio -from typing import List +from typing import List, cast from unittest.mock import AsyncMock, MagicMock, patch import pytest +from litellm.integrations.anthropic_cache_control_hook import ( + AnthropicCacheControlHook, +) from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj -from litellm.types.llms.openai import AllMessageValues +from litellm.types.llms.openai import ( + AllMessageValues, + ResponseInputParam, +) # --------------------------------------------------------------------------- # Helpers @@ -71,6 +77,56 @@ def _patch_responses_dispatch(): ] +def _make_cache_control_case() -> tuple[ + ResponseInputParam, + list[AllMessageValues], + dict[str, object], +]: + system_message = cast( + AllMessageValues, + {"role": "system", "content": "Analyze the request"}, + ) + assistant_message = cast( + AllMessageValues, + { + "type": "message", + "id": "msg_1", + "role": "assistant", + "status": "completed", + "content": [ + { + "type": "output_text", + "text": "The code has a bug", + "annotations": [], + } + ], + }, + ) + user_message = cast( + AllMessageValues, + {"role": "user", "content": "Check for security issues"}, + ) + reasoning_item = { + "type": "reasoning", + "id": "rs_1", + "summary": [], + "encrypted_content": "encrypted", + } + original_input = cast( + ResponseInputParam, + [system_message, reasoning_item, assistant_message, user_message], + ) + _, merged_messages, _ = AnthropicCacheControlHook().get_chat_completion_prompt( + model="azure/gpt-5-codex", + messages=[system_message, assistant_message, user_message], + non_default_params={"cache_control_injection_points": [{"location": "message", "role": "system"}]}, + prompt_id=None, + prompt_variables=None, + dynamic_callback_params={}, + ) + return original_input, merged_messages, reasoning_item + + # --------------------------------------------------------------------------- # Tests # --------------------------------------------------------------------------- @@ -256,6 +312,66 @@ class TestResponsesAPIPromptManagement: assert all(isinstance(m, dict) and "role" in m for m in passed_messages) assert len(passed_messages) == 1 + def test_cache_control_hook_preserves_reasoning_items(self): + original_input, merged_messages, reasoning_item = _make_cache_control_case() + logging_obj = _make_logging_obj( + merged_model="azure/gpt-5-codex", + merged_messages=merged_messages, + ) + + patches = _patch_responses_dispatch() + with patches[0], patches[1], patches[2], patches[3] as mock_handler: + import litellm + + litellm.responses( + input=original_input, + model="azure/gpt-5-codex", + litellm_logging_obj=logging_obj, + cache_control_injection_points=[{"location": "message", "role": "system"}], + ) + + sent_input = mock_handler.call_args.kwargs["input"] + assert [item.get("type") for item in sent_input] == [ + None, + "reasoning", + "message", + None, + ] + assert sent_input[0]["cache_control"] == {"type": "ephemeral"} + assert sent_input[1] == reasoning_item + assert sent_input[2]["id"] == "msg_1" + + def test_all_non_message_input_items_remain_unchanged(self): + reasoning_item = { + "type": "reasoning", + "id": "rs_1", + "summary": [], + "encrypted_content": "encrypted", + } + original_input = cast(ResponseInputParam, [reasoning_item]) + logging_obj = _make_logging_obj( + merged_model="openai/gpt-4o", + merged_messages=[ + cast( + AllMessageValues, + {"role": "system", "content": "Analyze the request"}, + ) + ], + ) + + patches = _patch_responses_dispatch() + with patches[0], patches[1], patches[2], patches[3] as mock_handler: + import litellm + + litellm.responses( + input=original_input, + model="gpt-4o", + prompt_id="all-non-message", + litellm_logging_obj=logging_obj, + ) + + assert mock_handler.call_args.kwargs["input"] == original_input + def test_model_override_re_resolves_provider(self): """[G] When the prompt template overrides the model to a different provider, custom_llm_provider is re-resolved so downstream routing uses the correct provider. @@ -393,3 +509,33 @@ class TestAsyncResponsesAPIPromptManagement: passed_messages = call_kwargs["messages"] assert all(isinstance(m, dict) and "role" in m for m in passed_messages) assert len(passed_messages) == 1 + + @pytest.mark.asyncio + async def test_async_cache_control_hook_preserves_reasoning_items(self): + original_input, merged_messages, reasoning_item = _make_cache_control_case() + logging_obj = _make_logging_obj( + merged_model="azure/gpt-5-codex", + merged_messages=merged_messages, + ) + + patches = _patch_responses_dispatch() + with patches[0], patches[1], patches[2], patches[3] as mock_handler: + import litellm + + await litellm.aresponses( + input=original_input, + model="azure/gpt-5-codex", + litellm_logging_obj=logging_obj, + cache_control_injection_points=[{"location": "message", "role": "system"}], + ) + + sent_input = mock_handler.call_args.kwargs["input"] + assert [item.get("type") for item in sent_input] == [ + None, + "reasoning", + "message", + None, + ] + assert sent_input[0]["cache_control"] == {"type": "ephemeral"} + assert sent_input[1] == reasoning_item + assert sent_input[2]["id"] == "msg_1" diff --git a/tests/test_litellm/responses/test_responses_utils.py b/tests/test_litellm/responses/test_responses_utils.py index 3a75a33fdc7..0141cf5d96a 100644 --- a/tests/test_litellm/responses/test_responses_utils.py +++ b/tests/test_litellm/responses/test_responses_utils.py @@ -369,6 +369,32 @@ class TestResponseAPILoggingUtils: assert result.completion_tokens_details.image_tokens == 272 assert result.completion_tokens_details.text_tokens == 100 + def test_transform_response_api_usage_maps_cache_write_tokens(self): + """Responses API (/v1/responses) cache-write tokens must survive the usage transform. + + gpt-5.6 returns usage.input_tokens_details.cache_write_tokens (an extra field + not typed on InputTokensDetails). Before the fix the transform rebuilt the token + details and dropped it, leaving the cache-creation metric empty (LIT-4633). + """ + usage = { + "input_tokens": 10062, + "output_tokens": 16, + "total_tokens": 10078, + "input_tokens_details": { + "cached_tokens": 0, + "cache_write_tokens": 10059, + }, + } + + result = ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage( + usage + ) + + assert result.prompt_tokens_details is not None + assert result.prompt_tokens_details.cache_write_tokens == 10059 + assert result.prompt_tokens_details.cache_creation_tokens == 10059 + assert result.prompt_tokens_details.cached_tokens == 0 + def test_transform_response_api_usage_mixed_details(self): """Test transformation handles mixed token details (cached + image + audio).""" # Setup - hypothetical usage with mixed token types diff --git a/tests/test_litellm/test_claude_opus_5_config.py b/tests/test_litellm/test_claude_opus_5_config.py new file mode 100644 index 00000000000..84021a83a5a --- /dev/null +++ b/tests/test_litellm/test_claude_opus_5_config.py @@ -0,0 +1,277 @@ +""" +Validate Claude Opus 5 model configuration entries. + +Opus 5 carries Opus 4.8's pricing ($5 / $25 per MTok) and the gen-5 adaptive +thinking profile, but differs from 4.8 in two ways that are behavior-bearing in +LiteLLM: the cacheable-prefix minimum drops to 512 tokens, and Bedrock's Opus 5 +validator accepts the full effort ladder, so the entries must not carry the +``bedrock_output_config_effort_ceiling`` that silently clamps ``max`` to +``xhigh`` on 4.8. The cost-map entries are also what populate +``litellm.anthropic_models`` at import, which is what lets a bare +``claude-opus-5`` name resolve to the ``anthropic`` provider (and match an +``anthropic/*`` wildcard deployment). +""" + +import json +import os + +import pytest + +import litellm +from litellm.constants import BEDROCK_CONVERSE_MODELS +from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap + +REPO_ROOT = os.path.join(os.path.dirname(__file__), "../..") + +ALL_OPUS_5_VARIANTS = ( + "claude-opus-5", + "anthropic.claude-opus-5", + "global.anthropic.claude-opus-5", + "us.anthropic.claude-opus-5", + "eu.anthropic.claude-opus-5", + "au.anthropic.claude-opus-5", + "jp.anthropic.claude-opus-5", + "vertex_ai/claude-opus-5", + "vertex_ai/claude-opus-5@default", + "azure_ai/claude-opus-5", +) + +BEDROCK_OPUS_5_VARIANTS = ( + "anthropic.claude-opus-5", + "global.anthropic.claude-opus-5", + "us.anthropic.claude-opus-5", + "eu.anthropic.claude-opus-5", + "au.anthropic.claude-opus-5", + "jp.anthropic.claude-opus-5", +) + + +def _load_root_cost_map() -> dict: + json_path = os.path.join(REPO_ROOT, "model_prices_and_context_window.json") + with open(json_path) as f: + return json.load(f) + + +@pytest.fixture +def local_model_cost_map(monkeypatch): + """Force the bundled backup cost map so assertions don't depend on the + network-fetched ``main`` copy (which lags this branch until merge).""" + original_model_cost = litellm.model_cost + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + litellm.model_cost = litellm.get_model_cost_map(url="") + litellm.get_model_info.cache_clear() + try: + yield + finally: + litellm.model_cost = original_model_cost + litellm.get_model_info.cache_clear() + + +def test_opus_5_pricing_and_capabilities(): + model_data = _load_root_cost_map() + + expected_providers = { + "claude-opus-5": "anthropic", + "anthropic.claude-opus-5": "bedrock_converse", + "vertex_ai/claude-opus-5": "vertex_ai-anthropic_models", + "azure_ai/claude-opus-5": "azure_ai", + } + + for model_name, provider in expected_providers.items(): + assert model_name in model_data, f"Missing model entry: {model_name}" + info = model_data[model_name] + + assert info["litellm_provider"] == provider + assert info["mode"] == "chat" + assert info["max_input_tokens"] == 1000000 + assert info["max_output_tokens"] == 128000 + assert info["max_tokens"] == 128000 + + # Opus 5 ships at Opus 4.8's rates: $5 / $25 per MTok, with the standard + # 1.25x cache-write, 2x 1-hour cache-write, and 0.1x cache-read multipliers. + assert info["input_cost_per_token"] == 5e-06 + assert info["output_cost_per_token"] == 2.5e-05 + assert info["cache_creation_input_token_cost"] == 6.25e-06 + assert info["cache_creation_input_token_cost_above_1hr"] == 1e-05 + assert info["cache_read_input_token_cost"] == 5e-07 + + # Flat rate across the full 1M window, no long-context premium. + assert "input_cost_per_token_above_200k_tokens" not in info + assert "output_cost_per_token_above_200k_tokens" not in info + + # gen-5 adaptive-thinking profile: effort-driven, no sampling params, no + # assistant prefill. + assert info["supports_adaptive_thinking"] is True + assert info["supports_reasoning"] is True + assert info["supports_sampling_params"] is False + assert info["supports_assistant_prefill"] is False + assert info["supports_xhigh_reasoning_effort"] is True + assert info["supports_max_reasoning_effort"] is True + + assert info["supports_function_calling"] is True + assert info["supports_prompt_caching"] is True + assert info["supports_tool_choice"] is True + assert info["supports_vision"] is True + + +def test_opus_5_bedrock_regional_pricing(): + """Global/base endpoints use base pricing; the us./eu./au./jp. regional + cross-region inference profiles carry a 10% premium.""" + model_data = _load_root_cost_map() + + base_pricing = { + "input_cost_per_token": 5e-06, + "output_cost_per_token": 2.5e-05, + "cache_creation_input_token_cost": 6.25e-06, + "cache_creation_input_token_cost_above_1hr": 1e-05, + "cache_read_input_token_cost": 5e-07, + } + regional_pricing = { + "input_cost_per_token": 5.5e-06, + "output_cost_per_token": 2.75e-05, + "cache_creation_input_token_cost": 6.875e-06, + "cache_creation_input_token_cost_above_1hr": 1.1e-05, + "cache_read_input_token_cost": 5.5e-07, + } + + expected = { + "anthropic.claude-opus-5": base_pricing, + "global.anthropic.claude-opus-5": base_pricing, + "us.anthropic.claude-opus-5": regional_pricing, + "eu.anthropic.claude-opus-5": regional_pricing, + "au.anthropic.claude-opus-5": regional_pricing, + "jp.anthropic.claude-opus-5": regional_pricing, + } + + for model_name, pricing in expected.items(): + assert model_name in model_data, f"Missing model entry: {model_name}" + info = model_data[model_name] + assert info["litellm_provider"] == "bedrock_converse" + for key, value in pricing.items(): + assert info[key] == value, f"{model_name}.{key} = {info[key]}, want {value}" + + +@pytest.mark.parametrize("model_name", BEDROCK_OPUS_5_VARIANTS) +def test_opus_5_bedrock_entries_declare_no_effort_ceiling(model_name): + """Bedrock accepts every effort level for Opus 5, so no clamp belongs here. + + Opus 4.7/4.8 carry ``bedrock_output_config_effort_ceiling: "xhigh"``, which + is what ``normalize_bedrock_opus_output_config_effort`` reads to rewrite a + caller's effort down. Verified against Bedrock on 2026-07-24 that + ``output_config.effort="max"`` returns 200 for the Opus 5 profiles, so the + ceiling is deliberately absent; adding one back would silently downgrade + requests. + + This asserts the cost-map entry rather than calling the normalizer because + ``_BEDROCK_OUTPUT_CONFIG_EFFORT_ORDER`` currently ranks ``max`` (3) below + ``xhigh`` (4), so an ``xhigh`` ceiling never clamps ``max`` and a behavioral + assertion would pass either way. Keeping the entry clean means Opus 5 stays + correct once that ordering is fixed.""" + info = _load_root_cost_map()[model_name] + assert "bedrock_output_config_effort_ceiling" not in info + + +@pytest.mark.parametrize("model_name", BEDROCK_OPUS_5_VARIANTS) +def test_opus_5_bedrock_rejects_strict_tools(model_name, local_model_cost_map): + """Bedrock Converse routes Opus through a validator that rejects + ``toolSpec.strict`` (``tools.0.custom.strict: Extra inputs are not + permitted``), same as Opus 4.7/4.8; verified against Bedrock on 2026-07-24. + Without the flag LiteLLM forwards ``strict`` and every tool call 400s.""" + from litellm.llms.bedrock.common_utils import bedrock_converse_supports_strict_tools + + assert bedrock_converse_supports_strict_tools(model_name) is False + + +def test_opus_5_prompt_cache_minimum_is_512(local_model_cost_map): + """Opus 5 halves the cacheable-prefix minimum (Opus 4.8 is 1024). + + The router's prompt-caching deployment check reads this value, so a stale + 1024 would route prompts of 512-1023 tokens away from a warm Opus 5 + deployment even though they cache fine.""" + from litellm.utils import get_prompt_cache_min_tokens + + assert get_prompt_cache_min_tokens(model="claude-opus-5") == 512 + assert get_prompt_cache_min_tokens(model="us.anthropic.claude-opus-5") == 512 + + +def test_opus_5_supports_fast_mode(local_model_cost_map): + """Fast mode is Opus 5 on the first-party API at $10 / $50 per MTok, i.e. 2x + base. ``supports_speed`` gates whether ``speed="fast"`` is forwarded at all, + and ``provider_specific_entry.fast`` is what prices the response.""" + from litellm.llms.anthropic.chat.transformation import AnthropicConfig + from litellm.llms.anthropic.cost_calculation import ( + cost_per_token as anthropic_cost_per_token, + ) + from litellm.types.utils import Usage + + assert ( + AnthropicConfig._model_supports_speed_param("claude-opus-5", "anthropic") is True + ) + + usage = Usage(prompt_tokens=1000, completion_tokens=500, total_tokens=1500) + usage.speed = "fast" + prompt_cost, completion_cost = anthropic_cost_per_token( + model="claude-opus-5", usage=usage + ) + assert prompt_cost == pytest.approx(1000 * 5e-06 * 2.0) + assert completion_cost == pytest.approx(500 * 2.5e-05 * 2.0) + + +def test_opus_5_present_in_bundled_backup(): + """The bundled backup is the runtime fallback (and what tests load with + ``LITELLM_LOCAL_MODEL_COST_MAP=True``); it must carry the same entries as the + root cost map, otherwise the model resolves on one path but not the other.""" + backup = GetModelCostMap.load_local_model_cost_map() + for model_name in ALL_OPUS_5_VARIANTS: + assert model_name in backup, f"Missing from backup cost map: {model_name}" + + +def test_opus_5_registered_for_bedrock_converse(): + assert "anthropic.claude-opus-5" in BEDROCK_CONVERSE_MODELS + + +def test_opus_5_provider_resolves_via_model_info(local_model_cost_map): + """Regression: ``claude-opus-5`` must resolve to provider ``anthropic``. + + Without the cost-map entry the model is unknown to LiteLLM, so it cannot be + tied to the ``anthropic`` provider and an ``anthropic/*`` wildcard deployment + would not match it.""" + info = litellm.get_model_info(model="claude-opus-5") + assert info["litellm_provider"] == "anthropic" + assert info["max_input_tokens"] == 1000000 + assert info["max_output_tokens"] == 128000 + + +@pytest.mark.parametrize( + "cost_map", + [_load_root_cost_map(), GetModelCostMap.load_local_model_cost_map()], + ids=["root", "bundled_backup"], +) +def test_opus_5_all_variants_carry_adaptive_thinking_flag(cost_map): + """Every Opus 5 entry must advertise ``supports_adaptive_thinking``. + + Adaptive-thinking detection is cost-map driven, so a single variant missing + the flag silently sends the legacy ``thinking.type='enabled'`` shape, which + Opus 5 rejects with a 400.""" + variants = [k for k in cost_map if "claude-opus-5" in k] + assert variants, "no claude-opus-5 entries found in cost map" + missing = [ + k for k in variants if cost_map[k].get("supports_adaptive_thinking") is not True + ] + assert not missing, f"missing supports_adaptive_thinking: {missing}" + + +@pytest.mark.parametrize( + "cost_map", + [_load_root_cost_map(), GetModelCostMap.load_local_model_cost_map()], + ids=["root", "bundled_backup"], +) +def test_opus_5_all_variants_carry_512_token_cache_minimum(cost_map): + variants = [k for k in cost_map if "claude-opus-5" in k] + assert variants, "no claude-opus-5 entries found in cost map" + wrong = { + k: cost_map[k].get("prompt_cache_min_tokens") + for k in variants + if cost_map[k].get("prompt_cache_min_tokens") != 512 + } + assert not wrong, f"prompt_cache_min_tokens must be 512: {wrong}" diff --git a/tests/test_litellm/test_cost_calculator.py b/tests/test_litellm/test_cost_calculator.py index 9636db4f4cd..276ee96ed65 100644 --- a/tests/test_litellm/test_cost_calculator.py +++ b/tests/test_litellm/test_cost_calculator.py @@ -12,6 +12,7 @@ from pydantic import BaseModel import litellm from litellm.cost_calculator import ( + BaseTokenUsageProcessor, RealtimeAPITokenUsageProcessor, completion_cost, cost_per_token, @@ -3479,3 +3480,32 @@ def test_batch_cost_calculator_cache_creation_falls_back_to_input_rate(): ) assert prompt_cost == pytest.approx((1000 * 3e-6 + 8000 * 3e-7 + 2000 * 3e-6) / 2) + + +def test_combine_usage_objects_sums_mirrored_cache_write_fields_once(): + """ + cache_write_tokens and cache_creation_tokens mirror each other on + PromptTokensDetailsWrapper, so field-iterating aggregation must sum the pair + once: a single 50-token usage stays 50 and two combine to 100, not double. + """ + single = Usage( + prompt_tokens=100, + completion_tokens=10, + total_tokens=110, + prompt_tokens_details=PromptTokensDetailsWrapper(cache_write_tokens=50), + ) + combined = BaseTokenUsageProcessor.combine_usage_objects([single]) + assert combined.prompt_tokens_details is not None + assert combined.prompt_tokens_details.cache_write_tokens == 50 + assert combined.prompt_tokens_details.cache_creation_tokens == 50 + + anthropic_style = Usage( + prompt_tokens=100, + completion_tokens=10, + total_tokens=110, + cache_creation_input_tokens=50, + ) + combined_pair = BaseTokenUsageProcessor.combine_usage_objects([anthropic_style, anthropic_style]) + assert combined_pair.prompt_tokens_details is not None + assert combined_pair.prompt_tokens_details.cache_write_tokens == 100 + assert combined_pair.prompt_tokens_details.cache_creation_tokens == 100 diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index a1a9448cc58..edc0cfed63e 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -17,7 +17,9 @@ from litellm.types.utils import ( Delta, LlmProviders, ModelResponseStream, + PromptTokensDetailsWrapper, StreamingChoices, + Usage, ) from litellm.utils import ( ProviderConfigManager, @@ -34,6 +36,57 @@ from litellm.utils import ( # Adds the parent directory to the system path +def test_usage_openai_cache_write_tokens_populates_both_names(): + """OpenAI reports cache-write tokens as prompt_tokens_details.cache_write_tokens. + The Usage constructor must expose it under both cache_write_tokens (canonical, + OpenAI naming) and cache_creation_tokens (legacy, Anthropic naming).""" + usage = Usage( + prompt_tokens=1000, + completion_tokens=10, + total_tokens=1010, + prompt_tokens_details={"cached_tokens": 0, "cache_write_tokens": 800}, + ) + assert usage.prompt_tokens_details.cache_write_tokens == 800 + assert usage.prompt_tokens_details.cache_creation_tokens == 800 + + +def test_usage_anthropic_cache_creation_maps_to_cache_write_tokens(): + """Anthropic/Bedrock report the top-level cache_creation_input_tokens field. + It must be normalized onto the OpenAI cache_write_tokens name as well as the + legacy cache_creation_tokens name.""" + usage = Usage( + prompt_tokens=500, + completion_tokens=50, + total_tokens=550, + cache_creation_input_tokens=300, + cache_read_input_tokens=120, + ) + assert usage.prompt_tokens_details.cache_write_tokens == 300 + assert usage.prompt_tokens_details.cache_creation_tokens == 300 + assert usage.prompt_tokens_details.cached_tokens == 120 + + +def test_prompt_tokens_details_no_cache_write_tokens_when_absent(): + """A read-only cache hit (no cache write) must not surface cache-write fields.""" + details = PromptTokensDetailsWrapper(cached_tokens=800) + assert details.cached_tokens == 800 + assert not hasattr(details, "cache_write_tokens") + assert not hasattr(details, "cache_creation_tokens") + + +def test_prompt_tokens_details_cache_write_creation_stay_in_sync_on_assignment(): + """Assigning either name after construction must mirror to the other, so a + caller that sets only one field can't leave the pair silently out of sync.""" + details = PromptTokensDetailsWrapper(cache_write_tokens=100) + assert details.cache_write_tokens == details.cache_creation_tokens == 100 + + details.cache_write_tokens = 250 + assert details.cache_write_tokens == details.cache_creation_tokens == 250 + + details.cache_creation_tokens = 375 + assert details.cache_write_tokens == details.cache_creation_tokens == 375 + + @pytest.fixture def local_model_cost_map(monkeypatch): original_model_cost = litellm.model_cost diff --git a/ui/litellm-dashboard/eslint-suppressions.json b/ui/litellm-dashboard/eslint-suppressions.json index 6072c8c725b..4ed6a0767f4 100644 --- a/ui/litellm-dashboard/eslint-suppressions.json +++ b/ui/litellm-dashboard/eslint-suppressions.json @@ -4,20 +4,41 @@ "count": 1 } }, - "src/app/(dashboard)/admin-panel/_components/AdminPanel.tsx": { + "src/app/(dashboard)/access-groups/_components/AccessGroupsModal/AccessGroupBaseForm.tsx": { + "no-restricted-imports": { + "count": 2 + } + }, + "src/app/(dashboard)/access-groups/_components/AccessGroupsModal/AccessGroupCreateModal.tsx": { "no-restricted-imports": { "count": 1 + } + }, + "src/app/(dashboard)/access-groups/_components/AccessGroupsModal/AccessGroupEditModal.tsx": { + "no-restricted-imports": { + "count": 1 + } + }, + "src/app/(dashboard)/admin-panel/_components/AdminPanel.tsx": { + "no-restricted-imports": { + "count": 2 }, "react-hooks/set-state-in-effect": { "count": 1 } }, "src/app/(dashboard)/agents/_components/add_agent_form.tsx": { + "local/filename-pascal-case": { + "count": 1 + }, + "max-lines": { + "count": 1 + }, "no-nested-ternary": { "count": 3 }, "no-restricted-imports": { - "count": 1 + "count": 2 }, "react-hooks/set-state-in-effect": { "count": 2 @@ -27,6 +48,9 @@ } }, "src/app/(dashboard)/agents/_components/agent_card_discovery.tsx": { + "local/filename-pascal-case": { + "count": 1 + }, "react-hooks/refs": { "count": 3 }, @@ -35,44 +59,71 @@ } }, "src/app/(dashboard)/agents/_components/agent_cost_view.tsx": { - "no-restricted-imports": { + "local/filename-pascal-case": { "count": 1 } }, "src/app/(dashboard)/agents/_components/agent_form_fields.tsx": { - "no-nested-ternary": { + "local/filename-pascal-case": { "count": 1 - } - }, - "src/app/(dashboard)/agents/_components/agent_info.tsx": { + }, "no-nested-ternary": { "count": 1 }, "no-restricted-imports": { + "count": 2 + } + }, + "src/app/(dashboard)/agents/_components/agent_info.tsx": { + "local/filename-pascal-case": { "count": 1 }, + "local/no-complex-jsx-arrow": { + "count": 1 + }, + "no-nested-ternary": { + "count": 1 + }, + "no-restricted-imports": { + "count": 2 + }, "react-hooks/immutability": { "count": 1 } }, "src/app/(dashboard)/agents/_components/agent_virtual_keys.tsx": { + "local/filename-pascal-case": { + "count": 1 + }, "no-nested-ternary": { "count": 1 } }, - "src/app/(dashboard)/agents/_components/dynamic_agent_form_fields.tsx": { - "no-nested-ternary": { - "count": 2 + "src/app/(dashboard)/agents/_components/cost_config_fields.tsx": { + "local/filename-pascal-case": { + "count": 1 + }, + "no-restricted-imports": { + "count": 1 } }, - "src/app/(dashboard)/api-reference/_components/APIReferenceView.tsx": { + "src/app/(dashboard)/agents/_components/dynamic_agent_form_fields.tsx": { + "local/filename-pascal-case": { + "count": 1 + }, + "no-nested-ternary": { + "count": 2 + }, "no-restricted-imports": { "count": 1 } }, "src/app/(dashboard)/budgets/_components/budget_modal.tsx": { - "no-restricted-imports": { + "local/filename-pascal-case": { "count": 1 + }, + "no-restricted-imports": { + "count": 2 } }, "src/app/(dashboard)/budgets/_components/budget_panel.test.tsx": { @@ -81,19 +132,28 @@ } }, "src/app/(dashboard)/budgets/_components/budget_panel.tsx": { - "no-restricted-imports": { + "local/filename-pascal-case": { "count": 1 } }, "src/app/(dashboard)/budgets/_components/edit_budget_modal.tsx": { - "no-restricted-imports": { + "local/filename-pascal-case": { "count": 1 + }, + "no-restricted-imports": { + "count": 2 } }, "src/app/(dashboard)/caching/_components/cache_dashboard.tsx": { + "local/filename-pascal-case": { + "count": 1 + }, "no-restricted-imports": { "count": 1 }, + "prefer-const": { + "count": 3 + }, "react-hooks/purity": { "count": 1 }, @@ -102,42 +162,113 @@ } }, "src/app/(dashboard)/caching/_components/cache_health.tsx": { + "local/filename-pascal-case": { + "count": 1 + } + }, + "src/app/(dashboard)/caching/_components/cache_settings/CacheFormField.tsx": { "no-restricted-imports": { "count": 1 } }, - "src/app/(dashboard)/caching/_components/cache_settings/RedisTypeSelector.tsx": { + "src/app/(dashboard)/caching/_components/cache_settings/cacheSettingsFields.ts": { "no-restricted-imports": { "count": 1 } }, "src/app/(dashboard)/caching/_components/cache_settings/index.tsx": { - "no-restricted-imports": { + "local/filename-pascal-case": { "count": 1 }, + "no-restricted-imports": { + "count": 2 + }, "react-hooks/set-state-in-effect": { "count": 1 } }, - "src/app/(dashboard)/cost-tracking/_components/add_margin_form.tsx": { + "src/app/(dashboard)/caching/_components/coordination_redis_settings/CoordinationRedisFormField.tsx": { "no-restricted-imports": { "count": 1 } }, - "src/app/(dashboard)/cost-tracking/_components/add_provider_form.tsx": { + "src/app/(dashboard)/caching/_components/coordination_redis_settings/coordinationRedisFields.ts": { "no-restricted-imports": { "count": 1 } }, - "src/app/(dashboard)/cost-tracking/_components/cost_tracking_settings.tsx": { - "no-nested-ternary": { - "count": 2 + "src/app/(dashboard)/caching/_components/coordination_redis_settings/index.tsx": { + "local/filename-pascal-case": { + "count": 1 }, "no-restricted-imports": { "count": 1 } }, + "src/app/(dashboard)/caching/_components/response_time_indicator.tsx": { + "local/filename-pascal-case": { + "count": 1 + } + }, + "src/app/(dashboard)/cost-optimization/_components/AutorouterTab.tsx": { + "no-restricted-imports": { + "count": 1 + } + }, + "src/app/(dashboard)/cost-optimization/_components/CostOptimizationView.tsx": { + "no-restricted-imports": { + "count": 1 + } + }, + "src/app/(dashboard)/cost-optimization/_components/PromptCompressionTab.tsx": { + "no-restricted-imports": { + "count": 1 + } + }, + "src/app/(dashboard)/cost-optimization/_components/UsageTab.tsx": { + "no-restricted-imports": { + "count": 1 + } + }, + "src/app/(dashboard)/cost-tracking/_components/add_margin_form.tsx": { + "local/filename-pascal-case": { + "count": 1 + }, + "no-restricted-imports": { + "count": 2 + } + }, + "src/app/(dashboard)/cost-tracking/_components/add_provider_form.tsx": { + "local/filename-pascal-case": { + "count": 1 + }, + "no-restricted-imports": { + "count": 2 + } + }, + "src/app/(dashboard)/cost-tracking/_components/cost_tracking_settings.tsx": { + "local/filename-pascal-case": { + "count": 1 + }, + "no-nested-ternary": { + "count": 2 + }, + "no-restricted-imports": { + "count": 2 + } + }, "src/app/(dashboard)/cost-tracking/_components/how_it_works.tsx": { + "local/filename-pascal-case": { + "count": 1 + }, + "no-restricted-imports": { + "count": 1 + } + }, + "src/app/(dashboard)/cost-tracking/_components/pricing_calculator/index.tsx": { + "local/filename-pascal-case": { + "count": 1 + }, "no-restricted-imports": { "count": 1 } @@ -148,8 +279,11 @@ } }, "src/app/(dashboard)/cost-tracking/_components/pricing_calculator/multi_cost_results.tsx": { - "no-restricted-imports": { + "local/filename-pascal-case": { "count": 1 + }, + "no-restricted-imports": { + "count": 2 } }, "src/app/(dashboard)/cost-tracking/_components/pricing_calculator/multi_export_dropdown.test.tsx": { @@ -158,6 +292,9 @@ } }, "src/app/(dashboard)/cost-tracking/_components/pricing_calculator/multi_export_dropdown.tsx": { + "local/filename-pascal-case": { + "count": 1 + }, "no-restricted-imports": { "count": 1 } @@ -173,11 +310,17 @@ } }, "src/app/(dashboard)/cost-tracking/_components/provider_discount_table.tsx": { + "local/filename-pascal-case": { + "count": 1 + }, "no-restricted-imports": { "count": 1 } }, "src/app/(dashboard)/cost-tracking/_components/provider_margin_table.tsx": { + "local/filename-pascal-case": { + "count": 1 + }, "no-restricted-imports": { "count": 1 } @@ -193,13 +336,24 @@ } }, "src/app/(dashboard)/guardrails-monitor/_components/EvaluationSettingsModal.tsx": { + "no-restricted-imports": { + "count": 1 + }, "react-hooks/set-state-in-effect": { "count": 1 } }, + "src/app/(dashboard)/guardrails-monitor/_components/GuardrailConfig.tsx": { + "no-restricted-imports": { + "count": 1 + } + }, "src/app/(dashboard)/guardrails-monitor/_components/GuardrailDetail.tsx": { "no-nested-ternary": { "count": 3 + }, + "no-restricted-imports": { + "count": 1 } }, "src/app/(dashboard)/guardrails-monitor/_components/GuardrailsMonitorView.tsx": { @@ -210,37 +364,64 @@ "src/app/(dashboard)/guardrails-monitor/_components/GuardrailsOverview.tsx": { "no-nested-ternary": { "count": 8 + }, + "no-restricted-imports": { + "count": 2 } }, "src/app/(dashboard)/guardrails/_components/GuardrailTestPanel.tsx": { "no-restricted-imports": { - "count": 1 + "count": 2 } }, "src/app/(dashboard)/guardrails/_components/GuardrailTestPlayground.tsx": { "no-nested-ternary": { "count": 1 - } - }, - "src/app/(dashboard)/guardrails/_components/GuardrailTestResults.tsx": { + }, "no-restricted-imports": { "count": 1 } }, + "src/app/(dashboard)/guardrails/_components/GuardrailTestResults.tsx": { + "no-restricted-imports": { + "count": 2 + } + }, "src/app/(dashboard)/guardrails/_components/GuardrailsPanel.tsx": { + "no-restricted-imports": { + "count": 1 + }, "react-hooks/set-state-in-effect": { "count": 1 } }, "src/app/(dashboard)/guardrails/_components/TeamGuardrailsTab.tsx": { + "local/no-complex-jsx-arrow": { + "count": 1 + }, + "max-lines": { + "count": 1 + }, "no-nested-ternary": { "count": 2 }, + "no-restricted-imports": { + "count": 1 + }, "react-hooks/set-state-in-effect": { "count": 1 } }, "src/app/(dashboard)/guardrails/_components/add_guardrail_form.tsx": { + "local/filename-pascal-case": { + "count": 1 + }, + "max-lines": { + "count": 1 + }, + "no-restricted-imports": { + "count": 1 + }, "react-hooks/set-state-in-effect": { "count": 1 }, @@ -248,22 +429,44 @@ "count": 2 } }, + "src/app/(dashboard)/guardrails/_components/content_filter/CategoryTable.tsx": { + "no-restricted-imports": { + "count": 1 + } + }, "src/app/(dashboard)/guardrails/_components/content_filter/CompetitorIntentConfiguration.tsx": { "no-nested-ternary": { "count": 1 }, + "no-restricted-imports": { + "count": 1 + }, "react-hooks/set-state-in-effect": { "count": 1 } }, "src/app/(dashboard)/guardrails/_components/content_filter/ContentCategoryConfiguration.tsx": { + "local/no-complex-jsx-arrow": { + "count": 1 + }, "no-nested-ternary": { "count": 3 }, + "no-restricted-imports": { + "count": 1 + }, "react-hooks/set-state-in-effect": { "count": 1 } }, + "src/app/(dashboard)/guardrails/_components/content_filter/ContentFilterConfiguration.tsx": { + "local/no-complex-jsx-arrow": { + "count": 3 + }, + "no-restricted-imports": { + "count": 1 + } + }, "src/app/(dashboard)/guardrails/_components/content_filter/ContentFilterDisplay.tsx": { "no-restricted-imports": { "count": 1 @@ -273,51 +476,157 @@ "max-params": { "count": 2 }, + "no-restricted-imports": { + "count": 1 + }, "react-hooks/set-state-in-effect": { "count": 1 } }, + "src/app/(dashboard)/guardrails/_components/content_filter/CustomPatternModal.tsx": { + "no-restricted-imports": { + "count": 1 + } + }, + "src/app/(dashboard)/guardrails/_components/content_filter/KeywordModal.tsx": { + "no-restricted-imports": { + "count": 1 + } + }, + "src/app/(dashboard)/guardrails/_components/content_filter/KeywordTable.tsx": { + "no-restricted-imports": { + "count": 1 + } + }, + "src/app/(dashboard)/guardrails/_components/content_filter/PatternModal.tsx": { + "local/no-complex-jsx-arrow": { + "count": 1 + }, + "no-restricted-imports": { + "count": 1 + } + }, + "src/app/(dashboard)/guardrails/_components/content_filter/PatternTable.tsx": { + "no-restricted-imports": { + "count": 1 + } + }, "src/app/(dashboard)/guardrails/_components/custom_code/CustomCodeModal.tsx": { "no-nested-ternary": { "count": 6 }, "no-restricted-imports": { - "count": 1 + "count": 2 }, "react-hooks/set-state-in-effect": { "count": 1 } }, - "src/app/(dashboard)/guardrails/_components/guardrail_info.tsx": { - "max-params": { + "src/app/(dashboard)/guardrails/_components/guardrailTableColumns.tsx": { + "local/filename-pascal-case": { + "count": 1 + } + }, + "src/app/(dashboard)/guardrails/_components/guardrail_garden.tsx": { + "local/filename-pascal-case": { "count": 1 }, "no-restricted-imports": { "count": 1 + } + }, + "src/app/(dashboard)/guardrails/_components/guardrail_garden_card.tsx": { + "local/filename-pascal-case": { + "count": 1 + } + }, + "src/app/(dashboard)/guardrails/_components/guardrail_garden_detail.tsx": { + "local/filename-pascal-case": { + "count": 1 + }, + "no-restricted-imports": { + "count": 1 + } + }, + "src/app/(dashboard)/guardrails/_components/guardrail_info.tsx": { + "local/filename-pascal-case": { + "count": 1 + }, + "local/no-complex-jsx-arrow": { + "count": 1 + }, + "max-params": { + "count": 1 + }, + "no-restricted-imports": { + "count": 2 }, "react-hooks/set-state-in-effect": { "count": 3 } }, + "src/app/(dashboard)/guardrails/_components/guardrail_info_helpers.tsx": { + "local/filename-pascal-case": { + "count": 1 + } + }, "src/app/(dashboard)/guardrails/_components/guardrail_optional_params.tsx": { + "local/filename-pascal-case": { + "count": 1 + }, "no-nested-ternary": { "count": 5 }, + "no-restricted-imports": { + "count": 1 + }, "react-hooks/set-state-in-effect": { "count": 1 } }, "src/app/(dashboard)/guardrails/_components/guardrail_provider_fields.tsx": { + "local/filename-pascal-case": { + "count": 1 + }, "no-nested-ternary": { "count": 5 }, + "no-restricted-imports": { + "count": 1 + }, "react-hooks/set-state-in-effect": { "count": 1 } }, - "src/app/(dashboard)/guardrails/_components/tool_permission/ToolPermissionRulesEditor.tsx": { + "src/app/(dashboard)/guardrails/_components/guardrail_table.tsx": { + "local/filename-pascal-case": { + "count": 1 + } + }, + "src/app/(dashboard)/guardrails/_components/llm_judge/LLMJudgeFields.tsx": { + "no-restricted-imports": { + "count": 2 + } + }, + "src/app/(dashboard)/guardrails/_components/pii_components.tsx": { + "local/filename-pascal-case": { + "count": 1 + }, "no-restricted-imports": { "count": 1 + } + }, + "src/app/(dashboard)/guardrails/_components/pii_configuration.tsx": { + "local/filename-pascal-case": { + "count": 1 + }, + "no-restricted-imports": { + "count": 1 + } + }, + "src/app/(dashboard)/guardrails/_components/tool_permission/ToolPermissionRulesEditor.tsx": { + "no-restricted-imports": { + "count": 2 }, "react-hooks/purity": { "count": 1 @@ -483,6 +792,21 @@ "count": 2 } }, + "src/app/(dashboard)/hooks/useTeams.tsx": { + "local/filename-pascal-case": { + "count": 1 + } + }, + "src/app/(dashboard)/mcp-servers/_components/DcrBridgeToggle.tsx": { + "no-restricted-imports": { + "count": 1 + } + }, + "src/app/(dashboard)/mcp-servers/_components/EnvVarsSection.tsx": { + "no-restricted-imports": { + "count": 1 + } + }, "src/app/(dashboard)/mcp-servers/_components/MCPLogoSelector.test.tsx": { "unused-imports/no-unused-imports": { "count": 1 @@ -493,6 +817,16 @@ "count": 2 } }, + "src/app/(dashboard)/mcp-servers/_components/MCPPermissionManagement.test.tsx": { + "no-restricted-imports": { + "count": 1 + } + }, + "src/app/(dashboard)/mcp-servers/_components/MCPPermissionManagement.tsx": { + "no-restricted-imports": { + "count": 1 + } + }, "src/app/(dashboard)/mcp-servers/_components/MCPSubmissionsTab.tsx": { "react-hooks/set-state-in-effect": { "count": 1 @@ -503,18 +837,31 @@ "count": 1 }, "no-restricted-imports": { - "count": 1 + "count": 2 }, "react-hooks/set-state-in-effect": { "count": 1 } }, + "src/app/(dashboard)/mcp-servers/_components/OAuthFormFields.test.tsx": { + "no-restricted-imports": { + "count": 1 + } + }, "src/app/(dashboard)/mcp-servers/_components/OAuthFormFields.tsx": { "no-nested-ternary": { "count": 1 }, "no-restricted-imports": { + "count": 2 + } + }, + "src/app/(dashboard)/mcp-servers/_components/OpenAPIFormSection.tsx": { + "local/no-complex-jsx-arrow": { "count": 1 + }, + "no-restricted-imports": { + "count": 2 } }, "src/app/(dashboard)/mcp-servers/_components/OpenAPIQuickPicker.tsx": { @@ -522,12 +869,37 @@ "count": 1 } }, + "src/app/(dashboard)/mcp-servers/_components/PassthroughAuthorizeSection.test.tsx": { + "no-restricted-imports": { + "count": 1 + } + }, + "src/app/(dashboard)/mcp-servers/_components/PassthroughAuthorizeSection.tsx": { + "no-restricted-imports": { + "count": 1 + } + }, + "src/app/(dashboard)/mcp-servers/_components/StdioConfiguration.tsx": { + "no-restricted-imports": { + "count": 1 + } + }, + "src/app/(dashboard)/mcp-servers/_components/TokenEndpointAuthMethodField.tsx": { + "no-restricted-imports": { + "count": 1 + } + }, + "src/app/(dashboard)/mcp-servers/_components/TokenExchangeFormFields.tsx": { + "no-restricted-imports": { + "count": 1 + } + }, "src/app/(dashboard)/mcp-servers/_components/ToolTestPanel.tsx": { "no-nested-ternary": { "count": 3 }, "no-restricted-imports": { - "count": 1 + "count": 2 }, "react-hooks/set-state-in-effect": { "count": 1 @@ -536,56 +908,85 @@ "src/app/(dashboard)/mcp-servers/_components/UserEnvVarsModal.tsx": { "no-nested-ternary": { "count": 2 + }, + "no-restricted-imports": { + "count": 1 } }, "src/app/(dashboard)/mcp-servers/_components/create_mcp_server.tsx": { + "local/filename-pascal-case": { + "count": 1 + }, + "max-lines": { + "count": 1 + }, "no-nested-ternary": { "count": 1 }, "no-restricted-imports": { - "count": 1 + "count": 2 }, "react-hooks/set-state-in-effect": { "count": 4 } }, - "src/app/(dashboard)/mcp-servers/_components/mcp_connect.tsx": { - "no-restricted-imports": { + "src/app/(dashboard)/mcp-servers/_components/index.tsx": { + "local/filename-pascal-case": { "count": 1 + } + }, + "src/app/(dashboard)/mcp-servers/_components/mcp_connect.tsx": { + "local/filename-pascal-case": { + "count": 1 + }, + "no-restricted-imports": { + "count": 2 }, "react-hooks/static-components": { "count": 4 } }, "src/app/(dashboard)/mcp-servers/_components/mcp_connection_status.tsx": { + "local/filename-pascal-case": { + "count": 1 + }, "no-nested-ternary": { "count": 3 - }, - "no-restricted-imports": { - "count": 1 } }, "src/app/(dashboard)/mcp-servers/_components/mcp_discovery.tsx": { + "local/filename-pascal-case": { + "count": 1 + }, + "local/no-complex-jsx-arrow": { + "count": 1 + }, "react-hooks/set-state-in-effect": { "count": 2 } }, "src/app/(dashboard)/mcp-servers/_components/mcp_server_cost_config.tsx": { - "no-restricted-imports": { + "local/filename-pascal-case": { "count": 1 } }, "src/app/(dashboard)/mcp-servers/_components/mcp_server_cost_display.tsx": { - "no-restricted-imports": { + "local/filename-pascal-case": { "count": 1 } }, "src/app/(dashboard)/mcp-servers/_components/mcp_server_edit.tsx": { + "local/filename-pascal-case": { + "count": 1 + }, + "max-lines": { + "count": 1 + }, "no-nested-ternary": { "count": 1 }, "no-restricted-imports": { - "count": 1 + "count": 2 }, "react-hooks/immutability": { "count": 1 @@ -595,15 +996,18 @@ } }, "src/app/(dashboard)/mcp-servers/_components/mcp_server_view.tsx": { - "no-restricted-imports": { + "local/filename-pascal-case": { "count": 1 } }, "src/app/(dashboard)/mcp-servers/_components/mcp_servers.tsx": { - "no-nested-ternary": { + "local/filename-pascal-case": { "count": 1 }, - "no-restricted-imports": { + "local/no-complex-jsx-arrow": { + "count": 3 + }, + "no-nested-ternary": { "count": 1 }, "react-hooks/set-state-in-effect": { @@ -611,43 +1015,32 @@ } }, "src/app/(dashboard)/mcp-servers/_components/mcp_tool_configuration.tsx": { - "no-restricted-imports": { + "local/filename-pascal-case": { "count": 1 } }, "src/app/(dashboard)/mcp-servers/_components/mcp_tools.tsx": { - "no-nested-ternary": { + "local/filename-pascal-case": { "count": 1 }, - "no-restricted-imports": { + "local/no-complex-jsx-arrow": { + "count": 1 + }, + "no-nested-ternary": { "count": 1 }, "react-hooks/set-state-in-effect": { "count": 2 } }, - "src/app/(dashboard)/models-and-endpoints/ModelsAndEndpointsView.tsx": { - "no-restricted-imports": { - "count": 1 - }, - "react-hooks/preserve-manual-memoization": { - "count": 4 - } - }, - "src/app/(dashboard)/models-and-endpoints/components/AllModelsTab.test.tsx": { - "max-params": { - "count": 1 - }, - "unused-imports/no-unused-imports": { + "src/app/(dashboard)/mcp-servers/_components/utils.tsx": { + "local/filename-pascal-case": { "count": 1 } }, - "src/app/(dashboard)/models-and-endpoints/components/AllModelsTab.tsx": { + "src/app/(dashboard)/memory/_components/MemoryEditModal.tsx": { "no-restricted-imports": { "count": 1 - }, - "react-hooks/set-state-in-effect": { - "count": 3 } }, "src/app/(dashboard)/models-and-endpoints/components/ModelRetrySettingsTab.test.tsx": { @@ -660,7 +1053,7 @@ "count": 1 }, "no-restricted-imports": { - "count": 1 + "count": 2 } }, "src/app/(dashboard)/models-and-endpoints/components/PriceDataManagementTab.tsx": { @@ -668,9 +1061,25 @@ "count": 1 } }, - "src/app/(dashboard)/old-usage/_components/usage.tsx": { + "src/app/(dashboard)/models-and-endpoints/layout.tsx": { "no-restricted-imports": { - "count": 2 + "count": 1 + } + }, + "src/app/(dashboard)/models-and-endpoints/utils/modelDataTransformer.ts": { + "prefer-const": { + "count": 6 + } + }, + "src/app/(dashboard)/old-usage/_components/usage.tsx": { + "local/filename-pascal-case": { + "count": 1 + }, + "max-lines": { + "count": 1 + }, + "prefer-const": { + "count": 6 }, "react-hooks/immutability": { "count": 1 @@ -679,9 +1088,19 @@ "count": 1 } }, - "src/app/(dashboard)/playground/components/chat_ui/AdditionalModelSettings.tsx": { + "src/app/(dashboard)/organizations/_components/OrganizationsPanel.tsx": { "no-restricted-imports": { "count": 1 + } + }, + "src/app/(dashboard)/playground/components/chat_ui/A2AMetrics.tsx": { + "no-restricted-imports": { + "count": 1 + } + }, + "src/app/(dashboard)/playground/components/chat_ui/AdditionalModelSettings.tsx": { + "no-restricted-imports": { + "count": 2 }, "react-hooks/set-state-in-effect": { "count": 2 @@ -691,10 +1110,18 @@ "no-nested-ternary": { "count": 2 }, + "no-restricted-imports": { + "count": 1 + }, "react-hooks/set-state-in-effect": { "count": 5 } }, + "src/app/(dashboard)/playground/components/chat_ui/ChatImageUpload.tsx": { + "no-restricted-imports": { + "count": 1 + } + }, "src/app/(dashboard)/playground/components/chat_ui/ChatImageUtils.test.tsx": { "max-nested-callbacks": { "count": 1 @@ -706,10 +1133,19 @@ } }, "src/app/(dashboard)/playground/components/chat_ui/ChatUI.tsx": { + "local/no-complex-jsx-arrow": { + "count": 2 + }, + "max-lines": { + "count": 1 + }, "no-nested-ternary": { "count": 7 }, "no-restricted-imports": { + "count": 2 + }, + "prefer-const": { "count": 1 }, "react-hooks/set-state-in-effect": { @@ -723,11 +1159,19 @@ "no-nested-ternary": { "count": 1 }, + "no-restricted-imports": { + "count": 1 + }, "no-restricted-syntax": { "count": 2 } }, "src/app/(dashboard)/playground/components/chat_ui/CodeInterpreterTool.tsx": { + "no-restricted-imports": { + "count": 2 + } + }, + "src/app/(dashboard)/playground/components/chat_ui/EndpointSelector.tsx": { "no-restricted-imports": { "count": 1 } @@ -736,6 +1180,9 @@ "no-nested-ternary": { "count": 2 }, + "no-restricted-imports": { + "count": 1 + }, "react-hooks/immutability": { "count": 2 }, @@ -743,25 +1190,70 @@ "count": 1 } }, + "src/app/(dashboard)/playground/components/chat_ui/ResponsesImageUpload.tsx": { + "no-restricted-imports": { + "count": 1 + } + }, + "src/app/(dashboard)/playground/components/chat_ui/SearchResultsDisplay.tsx": { + "no-restricted-imports": { + "count": 1 + } + }, + "src/app/(dashboard)/playground/components/chat_ui/SessionManagement.tsx": { + "no-restricted-imports": { + "count": 1 + } + }, "src/app/(dashboard)/playground/components/compareUI/CompareUI.tsx": { + "max-lines": { + "count": 1 + }, "no-nested-ternary": { "count": 4 }, + "no-restricted-imports": { + "count": 1 + }, "react-hooks/set-state-in-effect": { "count": 1 } }, + "src/app/(dashboard)/playground/components/compareUI/components/ComparisonPanel.tsx": { + "local/no-complex-jsx-arrow": { + "count": 2 + }, + "no-restricted-imports": { + "count": 1 + } + }, "src/app/(dashboard)/playground/components/compareUI/components/MessageDisplay.tsx": { "no-nested-ternary": { "count": 1 } }, + "src/app/(dashboard)/playground/components/compareUI/components/MessageInput.tsx": { + "no-restricted-imports": { + "count": 1 + } + }, "src/app/(dashboard)/playground/components/compareUI/components/ModelSelector.tsx": { + "no-restricted-imports": { + "count": 2 + } + }, + "src/app/(dashboard)/playground/components/compareUI/components/UnifiedSelector.tsx": { "no-restricted-imports": { "count": 1 } }, "src/app/(dashboard)/playground/components/complianceUI/ComplianceUI.tsx": { + "local/no-complex-jsx-arrow": { + "count": 2 + }, + "max-lines": { + "count": 1 + }, "no-nested-ternary": { "count": 8 }, @@ -770,6 +1262,9 @@ } }, "src/app/(dashboard)/playground/llm_calls/a2a_send_message.tsx": { + "local/filename-pascal-case": { + "count": 1 + }, "max-params": { "count": 2 }, @@ -778,21 +1273,33 @@ } }, "src/app/(dashboard)/playground/llm_calls/anthropic_messages.tsx": { + "local/filename-pascal-case": { + "count": 1 + }, "max-params": { "count": 1 } }, "src/app/(dashboard)/playground/llm_calls/audio_speech.tsx": { + "local/filename-pascal-case": { + "count": 1 + }, "max-params": { "count": 1 } }, "src/app/(dashboard)/playground/llm_calls/audio_transcriptions.tsx": { + "local/filename-pascal-case": { + "count": 1 + }, "max-params": { "count": 1 } }, "src/app/(dashboard)/playground/llm_calls/embeddings_api.tsx": { + "local/filename-pascal-case": { + "count": 1 + }, "max-params": { "count": 1 }, @@ -801,21 +1308,33 @@ } }, "src/app/(dashboard)/playground/llm_calls/fetch_agents.tsx": { + "local/filename-pascal-case": { + "count": 1 + }, "no-restricted-syntax": { "count": 1 } }, "src/app/(dashboard)/playground/llm_calls/image_edits.tsx": { + "local/filename-pascal-case": { + "count": 1 + }, "max-params": { "count": 1 } }, "src/app/(dashboard)/playground/llm_calls/image_generation.tsx": { + "local/filename-pascal-case": { + "count": 1 + }, "max-params": { "count": 1 } }, "src/app/(dashboard)/playground/llm_calls/interactions_api.tsx": { + "local/filename-pascal-case": { + "count": 1 + }, "max-params": { "count": 1 }, @@ -829,17 +1348,26 @@ } }, "src/app/(dashboard)/policies/_components/add_attachment_form.tsx": { - "no-restricted-imports": { + "local/filename-pascal-case": { "count": 1 }, + "no-restricted-imports": { + "count": 2 + }, "react-hooks/immutability": { "count": 1 } }, "src/app/(dashboard)/policies/_components/add_policy_form.tsx": { - "no-restricted-imports": { + "local/filename-pascal-case": { "count": 1 }, + "no-restricted-imports": { + "count": 2 + }, + "prefer-const": { + "count": 2 + }, "react-hooks/immutability": { "count": 2 }, @@ -848,17 +1376,26 @@ } }, "src/app/(dashboard)/policies/_components/ai_suggestion_modal.tsx": { + "local/filename-pascal-case": { + "count": 1 + }, + "local/no-complex-jsx-arrow": { + "count": 3 + }, + "max-lines": { + "count": 1 + }, "no-nested-ternary": { "count": 10 }, - "no-restricted-imports": { - "count": 1 - }, "react-hooks/immutability": { "count": 1 } }, "src/app/(dashboard)/policies/_components/guardrail_selection_modal.tsx": { + "local/filename-pascal-case": { + "count": 1 + }, "no-nested-ternary": { "count": 1 }, @@ -872,10 +1409,18 @@ } }, "src/app/(dashboard)/policies/_components/impact_popover.tsx": { + "local/filename-pascal-case": { + "count": 1 + }, "no-nested-ternary": { "count": 1 }, "no-restricted-imports": { + "count": 2 + } + }, + "src/app/(dashboard)/policies/_components/impact_preview_alert.tsx": { + "local/filename-pascal-case": { "count": 1 } }, @@ -885,7 +1430,10 @@ } }, "src/app/(dashboard)/policies/_components/index.tsx": { - "no-restricted-imports": { + "local/filename-pascal-case": { + "count": 1 + }, + "local/no-complex-jsx-arrow": { "count": 1 }, "react-hooks/set-state-in-effect": { @@ -893,10 +1441,13 @@ } }, "src/app/(dashboard)/policies/_components/pipeline_flow_builder.tsx": { - "no-nested-ternary": { + "local/filename-pascal-case": { "count": 1 }, - "no-restricted-imports": { + "max-lines": { + "count": 1 + }, + "no-nested-ternary": { "count": 1 }, "react-hooks/set-state-in-effect": { @@ -904,23 +1455,31 @@ } }, "src/app/(dashboard)/policies/_components/policy_info.tsx": { - "no-restricted-imports": { + "local/filename-pascal-case": { "count": 1 }, "react-hooks/set-state-in-effect": { "count": 1 } }, - "src/app/(dashboard)/policies/_components/policy_test_panel.tsx": { - "no-restricted-imports": { + "src/app/(dashboard)/policies/_components/policy_templates.tsx": { + "local/filename-pascal-case": { "count": 1 + } + }, + "src/app/(dashboard)/policies/_components/policy_test_panel.tsx": { + "local/filename-pascal-case": { + "count": 1 + }, + "no-restricted-imports": { + "count": 2 }, "react-hooks/immutability": { "count": 1 } }, "src/app/(dashboard)/policies/_components/template_parameter_modal.tsx": { - "no-restricted-imports": { + "local/filename-pascal-case": { "count": 1 }, "react-hooks/immutability": { @@ -933,34 +1492,74 @@ "src/app/(dashboard)/projects/_components/ProjectDetailsPage.tsx": { "no-nested-ternary": { "count": 3 + }, + "no-restricted-imports": { + "count": 1 } }, "src/app/(dashboard)/projects/_components/ProjectKeysSection.tsx": { + "no-restricted-imports": { + "count": 1 + }, "react-hooks/set-state-in-effect": { "count": 1 } }, + "src/app/(dashboard)/projects/_components/ProjectModals/CreateProjectModal.tsx": { + "no-restricted-imports": { + "count": 1 + } + }, + "src/app/(dashboard)/projects/_components/ProjectModals/EditProjectModal.tsx": { + "no-restricted-imports": { + "count": 1 + } + }, + "src/app/(dashboard)/projects/_components/ProjectModals/ProjectBaseForm.test.tsx": { + "no-restricted-imports": { + "count": 1 + } + }, "src/app/(dashboard)/projects/_components/ProjectModals/ProjectBaseForm.tsx": { + "local/no-complex-jsx-arrow": { + "count": 1 + }, + "no-restricted-imports": { + "count": 2 + }, "react-hooks/set-state-in-effect": { "count": 2 } }, - "src/app/(dashboard)/prompts/_components/add_prompt_form.tsx": { + "src/app/(dashboard)/projects/_components/ProjectsPage.tsx": { "no-restricted-imports": { "count": 1 } }, - "src/app/(dashboard)/prompts/_components/index.tsx": { - "no-nested-ternary": { + "src/app/(dashboard)/prompts/_components/add_prompt_form.tsx": { + "local/filename-pascal-case": { "count": 1 }, "no-restricted-imports": { + "count": 3 + } + }, + "src/app/(dashboard)/prompts/_components/index.tsx": { + "local/filename-pascal-case": { + "count": 1 + }, + "no-nested-ternary": { "count": 1 }, "react-hooks/set-state-in-effect": { "count": 1 } }, + "src/app/(dashboard)/prompts/_components/prompt_editor_view.tsx": { + "local/filename-pascal-case": { + "count": 1 + } + }, "src/app/(dashboard)/prompts/_components/prompt_editor_view/DeveloperMessageCard.tsx": { "no-restricted-imports": { "count": 1 @@ -968,7 +1567,7 @@ }, "src/app/(dashboard)/prompts/_components/prompt_editor_view/ModelConfigCard.tsx": { "no-restricted-imports": { - "count": 1 + "count": 2 } }, "src/app/(dashboard)/prompts/_components/prompt_editor_view/PromptCodeSnippets.tsx": { @@ -976,7 +1575,7 @@ "count": 1 }, "no-restricted-imports": { - "count": 1 + "count": 2 }, "react-hooks/set-state-in-effect": { "count": 1 @@ -984,17 +1583,17 @@ }, "src/app/(dashboard)/prompts/_components/prompt_editor_view/PromptEditorHeader.tsx": { "no-restricted-imports": { - "count": 1 + "count": 2 } }, "src/app/(dashboard)/prompts/_components/prompt_editor_view/PromptMessagesCard.tsx": { "no-restricted-imports": { - "count": 1 + "count": 2 } }, "src/app/(dashboard)/prompts/_components/prompt_editor_view/PublishModal.tsx": { "no-restricted-imports": { - "count": 1 + "count": 2 } }, "src/app/(dashboard)/prompts/_components/prompt_editor_view/ToolsCard.tsx": { @@ -1008,19 +1607,38 @@ } }, "src/app/(dashboard)/prompts/_components/prompt_editor_view/VersionHistorySidePanel.tsx": { + "local/no-complex-jsx-arrow": { + "count": 1 + }, "no-nested-ternary": { "count": 1 }, + "no-restricted-imports": { + "count": 1 + }, "react-hooks/immutability": { "count": 1 } }, "src/app/(dashboard)/prompts/_components/prompt_editor_view/conversation_panel/MessageInput.tsx": { + "no-restricted-imports": { + "count": 2 + } + }, + "src/app/(dashboard)/prompts/_components/prompt_editor_view/conversation_panel/MessageList.tsx": { + "no-restricted-imports": { + "count": 1 + } + }, + "src/app/(dashboard)/prompts/_components/prompt_editor_view/conversation_panel/VariableInput.tsx": { "no-restricted-imports": { "count": 1 } }, "src/app/(dashboard)/prompts/_components/prompt_editor_view/conversation_panel/index.tsx": { + "local/filename-pascal-case": { + "count": 1 + }, "no-restricted-imports": { "count": 1 } @@ -1030,94 +1648,135 @@ "count": 1 } }, + "src/app/(dashboard)/prompts/_components/prompt_editor_view/index.tsx": { + "local/filename-pascal-case": { + "count": 1 + } + }, "src/app/(dashboard)/prompts/_components/prompt_info.tsx": { + "local/filename-pascal-case": { + "count": 1 + }, + "local/no-complex-jsx-arrow": { + "count": 1 + }, "no-nested-ternary": { "count": 3 }, "no-restricted-imports": { - "count": 1 + "count": 2 }, "react-hooks/set-state-in-effect": { "count": 2 } }, + "src/app/(dashboard)/prompts/_components/prompt_utils.tsx": { + "local/filename-pascal-case": { + "count": 1 + } + }, + "src/app/(dashboard)/prompts/_components/tool_modal.tsx": { + "local/filename-pascal-case": { + "count": 1 + }, + "no-restricted-imports": { + "count": 1 + } + }, + "src/app/(dashboard)/prompts/_components/variable_textarea.tsx": { + "local/filename-pascal-case": { + "count": 1 + }, + "no-restricted-imports": { + "count": 1 + } + }, "src/app/(dashboard)/router-settings/_components/general_settings.tsx": { + "local/filename-pascal-case": { + "count": 1 + }, "no-nested-ternary": { "count": 1 }, "no-restricted-imports": { + "count": 3 + }, + "prefer-const": { "count": 2 } }, "src/app/(dashboard)/search-tools/_components/CreateSearchTools.tsx": { "no-restricted-imports": { - "count": 1 + "count": 2 }, "react-hooks/set-state-in-effect": { "count": 1 } }, - "src/app/(dashboard)/search-tools/_components/SearchToolTester.tsx": { - "no-restricted-imports": { - "count": 1 - } - }, - "src/app/(dashboard)/search-tools/_components/SearchToolView.tsx": { - "no-restricted-imports": { - "count": 1 - } - }, "src/app/(dashboard)/search-tools/_components/SearchTools.tsx": { + "local/no-complex-jsx-arrow": { + "count": 2 + }, "no-restricted-imports": { - "count": 1 + "count": 2 }, "react-hooks/static-components": { "count": 1 } }, - "src/app/(dashboard)/skills/_components/ClaudeCodePluginsPanel.tsx": { - "no-restricted-imports": { + "src/app/(dashboard)/search-tools/_components/index.tsx": { + "local/filename-pascal-case": { "count": 1 - }, + } + }, + "src/app/(dashboard)/search-tools/_components/types.tsx": { + "local/filename-pascal-case": { + "count": 1 + } + }, + "src/app/(dashboard)/skills/_components/ClaudeCodePluginsPanel.tsx": { "react-hooks/set-state-in-effect": { "count": 1 } }, "src/app/(dashboard)/skills/_components/add_plugin_form.tsx": { - "no-restricted-imports": { + "local/filename-pascal-case": { "count": 1 + }, + "no-restricted-imports": { + "count": 2 } }, "src/app/(dashboard)/tag-management/_components/components/CreateTagModal.tsx": { "no-restricted-imports": { - "count": 1 + "count": 2 } }, "src/app/(dashboard)/tag-management/_components/index.tsx": { - "no-restricted-imports": { + "local/filename-pascal-case": { "count": 1 }, "react-hooks/set-state-in-effect": { "count": 1 } }, + "src/app/(dashboard)/tag-management/_components/tagTableColumns.tsx": { + "local/filename-pascal-case": { + "count": 1 + } + }, "src/app/(dashboard)/tag-management/_components/tag_info.tsx": { - "no-restricted-imports": { + "local/filename-pascal-case": { "count": 1 }, + "no-restricted-imports": { + "count": 3 + }, "react-hooks/set-state-in-effect": { "count": 1 } }, - "src/app/(dashboard)/transform-request/TransformRequestPanel.tsx": { - "no-restricted-imports": { - "count": 1 - } - }, "src/app/(dashboard)/ui-theme/UIThemeSettings.tsx": { - "no-restricted-imports": { - "count": 1 - }, "no-restricted-syntax": { "count": 3 }, @@ -1128,14 +1787,25 @@ "src/app/(dashboard)/usage/_components/components/EndpointUsage/components/EndpointUsageTable.tsx": { "no-nested-ternary": { "count": 1 + }, + "no-restricted-imports": { + "count": 2 } }, "src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.tsx": { + "local/no-complex-jsx-arrow": { + "count": 2 + }, "no-restricted-imports": { - "count": 1 + "count": 2 } }, "src/app/(dashboard)/usage/_components/components/EntityUsage/SpendByProvider.tsx": { + "no-restricted-imports": { + "count": 2 + } + }, + "src/app/(dashboard)/usage/_components/components/EntityUsage/TopModelView.tsx": { "no-restricted-imports": { "count": 1 } @@ -1144,16 +1814,25 @@ "no-nested-ternary": { "count": 1 }, + "no-restricted-imports": { + "count": 1 + }, "react-hooks/immutability": { "count": 1 } }, "src/app/(dashboard)/usage/_components/components/UsagePageView.tsx": { + "local/no-complex-jsx-arrow": { + "count": 2 + }, + "max-lines": { + "count": 1 + }, "no-nested-ternary": { "count": 1 }, "no-restricted-imports": { - "count": 1 + "count": 2 }, "react-hooks/purity": { "count": 1 @@ -1162,6 +1841,14 @@ "count": 3 } }, + "src/app/(dashboard)/usage/_components/components/UsageViewSelect/UsageViewSelect.tsx": { + "local/no-complex-jsx-arrow": { + "count": 2 + }, + "no-restricted-imports": { + "count": 1 + } + }, "src/app/(dashboard)/usage/_components/hooks/usePaginatedDailyActivity.ts": { "react-hooks/refs": { "count": 1 @@ -1170,13 +1857,29 @@ "count": 1 } }, - "src/app/(dashboard)/users/_components/DefaultUserSettings.tsx": { + "src/app/(dashboard)/users/_components/BulkEditUsers.tsx": { "no-restricted-imports": { "count": 1 + }, + "prefer-const": { + "count": 1 } }, - "src/app/(dashboard)/users/_components/edit_user.tsx": { + "src/app/(dashboard)/users/_components/DefaultUserSettings.tsx": { "no-restricted-imports": { + "count": 2 + } + }, + "src/app/(dashboard)/users/_components/edit_user.tsx": { + "local/filename-pascal-case": { + "count": 1 + }, + "no-restricted-imports": { + "count": 2 + } + }, + "src/app/(dashboard)/users/_components/index.tsx": { + "local/filename-pascal-case": { "count": 1 } }, @@ -1189,15 +1892,24 @@ } }, "src/app/(dashboard)/users/_components/user_edit_view.tsx": { - "no-restricted-imports": { + "local/filename-pascal-case": { "count": 1 }, + "no-restricted-imports": { + "count": 2 + }, "react-hooks/set-state-in-effect": { "count": 1 } }, "src/app/(dashboard)/users/_components/view_users.tsx": { + "local/filename-pascal-case": { + "count": 1 + }, "no-restricted-imports": { + "count": 3 + }, + "prefer-const": { "count": 1 }, "react-hooks/set-state-in-effect": { @@ -1205,14 +1917,25 @@ } }, "src/app/(dashboard)/users/_components/view_users/user_info_view.tsx": { - "no-restricted-imports": { + "local/filename-pascal-case": { "count": 1 }, + "local/no-complex-jsx-arrow": { + "count": 1 + }, + "no-restricted-imports": { + "count": 2 + }, "react-hooks/set-state-in-effect": { "count": 1 } }, "src/app/(dashboard)/vector-stores/_components/CreateVectorStore.tsx": { + "no-restricted-imports": { + "count": 3 + } + }, + "src/app/(dashboard)/vector-stores/_components/S3VectorsConfig.tsx": { "no-restricted-imports": { "count": 1 } @@ -1222,14 +1945,14 @@ "count": 2 }, "no-restricted-imports": { - "count": 1 + "count": 2 }, "react/no-unescaped-entities": { "count": 1 } }, "src/app/(dashboard)/vector-stores/_components/index.tsx": { - "no-restricted-imports": { + "local/filename-pascal-case": { "count": 1 }, "react-hooks/set-state-in-effect": { @@ -1237,9 +1960,12 @@ } }, "src/app/(dashboard)/vector-stores/_components/vector_store_info.tsx": { - "no-restricted-imports": { + "local/filename-pascal-case": { "count": 1 }, + "no-restricted-imports": { + "count": 2 + }, "react-hooks/set-state-in-effect": { "count": 1 } @@ -1261,6 +1987,12 @@ } }, "src/app/login/LoginPage.tsx": { + "local/no-complex-jsx-arrow": { + "count": 1 + }, + "no-restricted-imports": { + "count": 1 + }, "react-hooks/set-state-in-effect": { "count": 2 } @@ -1275,15 +2007,41 @@ "count": 1 } }, + "src/app/onboarding/OnboardingErrorView.tsx": { + "no-restricted-imports": { + "count": 1 + } + }, + "src/app/onboarding/OnboardingFormBody.tsx": { + "no-restricted-imports": { + "count": 1 + } + }, + "src/app/onboarding/OnboardingLoadingView.tsx": { + "no-restricted-imports": { + "count": 1 + } + }, "src/components/AIHub/ModelHubTable.test.tsx": { "max-params": { "count": 1 } }, "src/components/AIHub/ModelHubTable.tsx": { + "max-lines": { + "count": 1 + }, "no-nested-ternary": { "count": 1 }, + "no-restricted-imports": { + "count": 2 + }, + "prefer-const": { + "count": 4 + } + }, + "src/components/AIHub/SkillHubDashboard.tsx": { "no-restricted-imports": { "count": 1 } @@ -1298,7 +2056,7 @@ }, "src/components/AIHub/forms/MakeAgentPublicForm.tsx": { "no-restricted-imports": { - "count": 1 + "count": 2 }, "react-hooks/set-state-in-effect": { "count": 1 @@ -1314,7 +2072,7 @@ "count": 2 }, "no-restricted-imports": { - "count": 1 + "count": 2 }, "react-hooks/set-state-in-effect": { "count": 1 @@ -1322,28 +2080,78 @@ }, "src/components/AIHub/forms/MakeModelPublicForm.tsx": { "no-restricted-imports": { - "count": 1 + "count": 2 }, "react-hooks/set-state-in-effect": { "count": 1 } }, - "src/components/CreateUserButton.tsx": { + "src/components/BetaBadge.tsx": { "no-restricted-imports": { "count": 1 + } + }, + "src/components/CloudZeroCostTracking/CloudZeroCreateModal.tsx": { + "no-restricted-imports": { + "count": 1 + } + }, + "src/components/CloudZeroCostTracking/CloudZeroUpdateModal.tsx": { + "no-restricted-imports": { + "count": 1 + } + }, + "src/components/CreateUserButton.tsx": { + "no-restricted-imports": { + "count": 2 }, "react-hooks/set-state-in-effect": { "count": 1 } }, + "src/components/DebugWarningBanner.tsx": { + "no-restricted-imports": { + "count": 1 + } + }, + "src/components/DeletedKeysPage/DeletedKeysPage.tsx": { + "no-restricted-imports": { + "count": 1 + } + }, + "src/components/DeletedTeamsPage/DeletedTeamsPage.tsx": { + "no-restricted-imports": { + "count": 1 + } + }, + "src/components/DeprecationBanner.tsx": { + "no-restricted-imports": { + "count": 1 + } + }, + "src/components/EntityUsageExport/EntityUsageExportModal.tsx": { + "no-restricted-imports": { + "count": 1 + } + }, + "src/components/EntityUsageExport/ExportFormatSelector.tsx": { + "no-restricted-imports": { + "count": 1 + } + }, "src/components/EntityUsageExport/ExportSummary.tsx": { "no-restricted-imports": { "count": 1 } }, + "src/components/EntityUsageExport/ExportTypeSelector.tsx": { + "no-restricted-imports": { + "count": 1 + } + }, "src/components/EntityUsageExport/UsageExportHeader.tsx": { "no-restricted-imports": { - "count": 2 + "count": 3 } }, "src/components/EntityUsageExport/types.ts": { @@ -1367,11 +2175,17 @@ "src/components/GuardrailSettingsView.tsx": { "no-nested-ternary": { "count": 1 + }, + "no-restricted-imports": { + "count": 1 } }, "src/components/GuardrailsMonitor/LogViewer.tsx": { "no-nested-ternary": { "count": 1 + }, + "no-restricted-imports": { + "count": 1 } }, "src/components/HelpLink.test.tsx": { @@ -1379,8 +2193,13 @@ "count": 1 } }, - "src/components/ModelSelect/PaginatedModelSelect/PaginatedModelSelect.tsx": { - "no-nested-ternary": { + "src/components/LicenseExpiryBanner.tsx": { + "no-restricted-imports": { + "count": 1 + } + }, + "src/components/ModelSelect/ModelSelect.tsx": { + "no-restricted-imports": { "count": 1 } }, @@ -1389,20 +2208,58 @@ "count": 12 } }, - "src/components/Navbar/UserDropdown/UserDropdown.tsx": { - "react-hooks/set-state-in-effect": { + "src/components/Navbar/BlogDropdown/BlogDropdown.tsx": { + "no-restricted-imports": { + "count": 2 + } + }, + "src/components/Navbar/CommunityEngagementButtons/CommunityEngagementButtons.tsx": { + "no-restricted-imports": { "count": 1 } }, - "src/components/SCIM.tsx": { + "src/components/Navbar/NotificationsBell/NotificationsBell.tsx": { "no-restricted-imports": { "count": 1 + } + }, + "src/components/Navbar/UserDropdown/UserDropdown.tsx": { + "no-restricted-imports": { + "count": 2 }, "react-hooks/set-state-in-effect": { "count": 1 } }, + "src/components/Navbar/ViewSwitcher.tsx": { + "no-restricted-imports": { + "count": 2 + } + }, + "src/components/Navbar/WorkerDropdown/WorkerDropdown.tsx": { + "no-restricted-imports": { + "count": 1 + } + }, + "src/components/SCIM.tsx": { + "no-restricted-imports": { + "count": 2 + }, + "react-hooks/set-state-in-effect": { + "count": 1 + } + }, + "src/components/SSOModals.test.tsx": { + "no-restricted-imports": { + "count": 1 + } + }, "src/components/SSOModals.tsx": { + "no-restricted-imports": { + "count": 2 + } + }, + "src/components/Settings/AdminSettings/HashicorpVault/EditHashicorpVaultModal.tsx": { "no-restricted-imports": { "count": 1 } @@ -1410,19 +2267,50 @@ "src/components/Settings/AdminSettings/HashicorpVault/HashicorpVault.tsx": { "no-nested-ternary": { "count": 1 + }, + "no-restricted-imports": { + "count": 1 + } + }, + "src/components/Settings/AdminSettings/HashicorpVault/HashicorpVaultEmptyPlaceholder.tsx": { + "no-restricted-imports": { + "count": 1 + } + }, + "src/components/Settings/AdminSettings/LoggingSettings/LoggingSettings.tsx": { + "no-restricted-imports": { + "count": 1 } }, "src/components/Settings/AdminSettings/MCPSemanticFilterSettings/MCPSemanticFilterSettings.tsx": { "no-nested-ternary": { "count": 1 }, + "no-restricted-imports": { + "count": 1 + }, "react-hooks/set-state-in-effect": { "count": 1 } }, + "src/components/Settings/AdminSettings/PluginSettings/PluginSettings.tsx": { + "no-restricted-imports": { + "count": 1 + } + }, + "src/components/Settings/AdminSettings/SSOSettings/Modals/AddSSOSettingsModal.tsx": { + "no-restricted-imports": { + "count": 1 + } + }, + "src/components/Settings/AdminSettings/SSOSettings/Modals/BaseSSOSettingsForm.test.tsx": { + "no-restricted-imports": { + "count": 1 + } + }, "src/components/Settings/AdminSettings/SSOSettings/Modals/BaseSSOSettingsForm.tsx": { "no-restricted-imports": { - "count": 1 + "count": 2 } }, "src/components/Settings/AdminSettings/SSOSettings/Modals/EditSSOSettingsModal.test.tsx": { @@ -1430,9 +2318,32 @@ "count": 1 } }, + "src/components/Settings/AdminSettings/SSOSettings/Modals/EditSSOSettingsModal.tsx": { + "no-restricted-imports": { + "count": 1 + } + }, "src/components/Settings/AdminSettings/SSOSettings/RedactableField.tsx": { "no-nested-ternary": { "count": 1 + }, + "no-restricted-imports": { + "count": 1 + } + }, + "src/components/Settings/AdminSettings/SSOSettings/RoleMappings.tsx": { + "no-restricted-imports": { + "count": 1 + } + }, + "src/components/Settings/AdminSettings/SSOSettings/SSOSettings.tsx": { + "no-restricted-imports": { + "count": 1 + } + }, + "src/components/Settings/AdminSettings/SSOSettings/SSOSettingsEmptyPlaceholder.tsx": { + "no-restricted-imports": { + "count": 1 } }, "src/components/Settings/AdminSettings/SSOSettings/SSOSettingsLoadingSkeleton.test.tsx": { @@ -1440,7 +2351,15 @@ "count": 4 } }, + "src/components/Settings/AdminSettings/SSOSettings/SSOSettingsLoadingSkeleton.tsx": { + "no-restricted-imports": { + "count": 1 + } + }, "src/components/Settings/AdminSettings/UISettings/PageVisibilitySettings.tsx": { + "no-restricted-imports": { + "count": 1 + }, "react-hooks/set-state-in-render": { "count": 2 } @@ -1448,19 +2367,40 @@ "src/components/Settings/AdminSettings/UISettings/UISettings.tsx": { "no-nested-ternary": { "count": 1 + }, + "no-restricted-imports": { + "count": 1 } }, "src/components/Settings/RouterSettings/Fallbacks/AddFallbacks.tsx": { "no-restricted-imports": { - "count": 1 + "count": 2 }, "react-hooks/set-state-in-effect": { "count": 1 } }, - "src/components/Settings/RouterSettings/Fallbacks/FallbackSelectionForm.tsx": { + "src/components/Settings/RouterSettings/Fallbacks/AddFallbacksModal.tsx": { "no-restricted-imports": { "count": 1 + } + }, + "src/components/Settings/RouterSettings/Fallbacks/EditFallbacks.tsx": { + "no-restricted-imports": { + "count": 1 + } + }, + "src/components/Settings/RouterSettings/Fallbacks/FallbackGroupConfig.tsx": { + "local/no-complex-jsx-arrow": { + "count": 1 + }, + "no-restricted-imports": { + "count": 1 + } + }, + "src/components/Settings/RouterSettings/Fallbacks/FallbackSelectionForm.tsx": { + "no-restricted-imports": { + "count": 2 }, "react-hooks/set-state-in-effect": { "count": 1 @@ -1468,7 +2408,10 @@ }, "src/components/Settings/RouterSettings/Fallbacks/Fallbacks.tsx": { "no-restricted-imports": { - "count": 1 + "count": 2 + }, + "prefer-const": { + "count": 2 } }, "src/components/TeamSSOSettings.test.tsx": { @@ -1476,49 +2419,133 @@ "count": 1 } }, + "src/components/TeamSSOSettings.tsx": { + "no-restricted-imports": { + "count": 1 + } + }, "src/components/Teams.test.tsx": { "max-nested-callbacks": { "count": 4 + }, + "prefer-const": { + "count": 6 } }, "src/components/Teams.tsx": { + "local/no-complex-jsx-arrow": { + "count": 4 + }, + "max-lines": { + "count": 1 + }, "no-nested-ternary": { "count": 2 }, "no-restricted-imports": { - "count": 1 + "count": 2 + }, + "prefer-const": { + "count": 2 }, "react-hooks/set-state-in-effect": { "count": 3 } }, + "src/components/TeamsPage/teamTableColumns.tsx": { + "local/filename-pascal-case": { + "count": 1 + } + }, "src/components/ToolDetail.tsx": { "unused-imports/no-unused-imports": { - "count": 2 + "count": 1 } }, "src/components/UIAccessControlForm.tsx": { + "no-restricted-imports": { + "count": 2 + } + }, + "src/components/UsagePage/components/EntityUsage/TopKeyView.tsx": { + "no-restricted-imports": { + "count": 1 + } + }, + "src/components/UsagePage/components/KeyModelUsageView.tsx": { + "no-restricted-imports": { + "count": 2 + } + }, + "src/components/UsagePage/utils/value_formatters.tsx": { + "local/filename-pascal-case": { + "count": 1 + } + }, + "src/components/VirtualKeysPage/keyTableColumns.tsx": { + "local/filename-pascal-case": { + "count": 1 + }, "no-restricted-imports": { "count": 1 } }, "src/components/activity_metrics.tsx": { + "local/filename-pascal-case": { + "count": 1 + }, "no-nested-ternary": { "count": 1 }, + "no-restricted-imports": { + "count": 2 + } + }, + "src/components/add_model/AdaptiveRoutingConfig.tsx": { "no-restricted-imports": { "count": 1 } }, + "src/components/add_model/AddModelForm.test.tsx": { + "no-restricted-imports": { + "count": 2 + } + }, "src/components/add_model/AddModelForm.tsx": { + "local/no-complex-jsx-arrow": { + "count": 1 + }, "no-nested-ternary": { "count": 1 }, + "no-restricted-imports": { + "count": 4 + } + }, + "src/components/add_model/ClassificationMethodConfig.tsx": { + "no-restricted-imports": { + "count": 1 + } + }, + "src/components/add_model/ComplexityRouterConfig.tsx": { + "no-restricted-imports": { + "count": 1 + } + }, + "src/components/add_model/EscalationKeywords.tsx": { + "no-restricted-imports": { + "count": 1 + } + }, + "src/components/add_model/KeywordTierRules.tsx": { "no-restricted-imports": { "count": 1 } }, "src/components/add_model/RouterConfigBuilder.tsx": { + "no-restricted-imports": { + "count": 1 + }, "react-hooks/purity": { "count": 1 }, @@ -1526,56 +2553,153 @@ "count": 1 } }, - "src/components/add_model/add_auto_router_tab.tsx": { + "src/components/add_model/SemanticKeywordMatching.tsx": { "no-restricted-imports": { "count": 1 } }, + "src/components/add_model/add_auto_router_tab.test.tsx": { + "no-restricted-imports": { + "count": 1 + } + }, + "src/components/add_model/add_auto_router_tab.tsx": { + "local/filename-pascal-case": { + "count": 1 + }, + "no-restricted-imports": { + "count": 3 + } + }, + "src/components/add_model/add_model_modes.tsx": { + "local/filename-pascal-case": { + "count": 1 + } + }, + "src/components/add_model/add_model_tab.test.tsx": { + "no-restricted-imports": { + "count": 2 + } + }, "src/components/add_model/add_model_tab.tsx": { + "local/filename-pascal-case": { + "count": 1 + }, + "no-restricted-imports": { + "count": 4 + } + }, + "src/components/add_model/advanced_settings.tsx": { + "local/filename-pascal-case": { + "count": 1 + }, + "no-restricted-imports": { + "count": 4 + }, + "prefer-const": { + "count": 2 + } + }, + "src/components/add_model/auto_router_connection_test.tsx": { + "local/filename-pascal-case": { + "count": 1 + }, "no-restricted-imports": { "count": 1 } }, - "src/components/add_model/advanced_settings.tsx": { + "src/components/add_model/cache_control_settings.tsx": { + "local/filename-pascal-case": { + "count": 1 + }, + "no-restricted-imports": { + "count": 1 + }, + "prefer-const": { + "count": 1 + } + }, + "src/components/add_model/conditional_public_model_name.test.tsx": { "no-restricted-imports": { "count": 1 } }, "src/components/add_model/conditional_public_model_name.tsx": { - "no-restricted-imports": { + "local/filename-pascal-case": { "count": 1 }, + "local/no-complex-jsx-arrow": { + "count": 1 + }, + "no-restricted-imports": { + "count": 2 + }, "react-hooks/set-state-in-effect": { "count": 2 } }, + "src/components/add_model/handle_add_auto_router_submit.tsx": { + "local/filename-pascal-case": { + "count": 1 + } + }, + "src/components/add_model/handle_add_model_submit.tsx": { + "local/filename-pascal-case": { + "count": 1 + } + }, + "src/components/add_model/litellm_model_name.test.tsx": { + "no-restricted-imports": { + "count": 1 + } + }, "src/components/add_model/litellm_model_name.tsx": { + "local/filename-pascal-case": { + "count": 1 + }, "no-nested-ternary": { "count": 1 }, + "no-restricted-imports": { + "count": 3 + } + }, + "src/components/add_model/model_connection_test.tsx": { + "local/filename-pascal-case": { + "count": 1 + }, + "no-nested-ternary": { + "count": 2 + }, "no-restricted-imports": { "count": 1 } }, - "src/components/add_model/model_connection_test.tsx": { - "no-nested-ternary": { - "count": 2 + "src/components/add_model/provider_specific_fields.test.tsx": { + "no-restricted-imports": { + "count": 1 } }, "src/components/add_model/provider_specific_fields.tsx": { + "local/filename-pascal-case": { + "count": 1 + }, "no-nested-ternary": { "count": 5 }, "no-restricted-imports": { - "count": 1 + "count": 2 }, "react-hooks/immutability": { "count": 3 } }, "src/components/add_pass_through.tsx": { + "local/filename-pascal-case": { + "count": 1 + }, "no-restricted-imports": { - "count": 2 + "count": 3 } }, "src/components/agent_management/AgentSelector.test.tsx": { @@ -1586,30 +2710,60 @@ "count": 1 } }, + "src/components/agent_management/AgentSelector.tsx": { + "no-restricted-imports": { + "count": 1 + }, + "prefer-const": { + "count": 1 + } + }, + "src/components/alerting/alerting_settings.tsx": { + "local/filename-pascal-case": { + "count": 1 + }, + "prefer-const": { + "count": 1 + } + }, "src/components/alerting/dynamic_form.tsx": { + "local/filename-pascal-case": { + "count": 1 + }, "no-nested-ternary": { "count": 4 }, "no-restricted-imports": { - "count": 1 + "count": 2 } }, "src/components/bulk_create_users_button.tsx": { - "no-restricted-imports": { + "local/filename-pascal-case": { "count": 1 }, + "no-restricted-imports": { + "count": 2 + }, "react-hooks/set-state-in-effect": { "count": 1 } }, + "src/components/callback_info_helpers.tsx": { + "local/filename-pascal-case": { + "count": 1 + } + }, "src/components/chat/KeysPanel.tsx": { "no-nested-ternary": { "count": 1 } }, "src/components/chat/MCPAppsPanel.tsx": { + "local/no-complex-jsx-arrow": { + "count": 1 + }, "no-nested-ternary": { - "count": 7 + "count": 6 } }, "src/components/chat/MCPConnectPicker.tsx": { @@ -1627,17 +2781,45 @@ "count": 2 } }, - "src/components/claude_code_plugins/MakeSkillPublicForm.tsx": { + "src/components/chat_ui/MCPEventsDisplay.tsx": { "no-restricted-imports": { "count": 1 + } + }, + "src/components/chat_ui/ReasoningContent.tsx": { + "no-restricted-imports": { + "count": 1 + } + }, + "src/components/chat_ui/ResponseMetrics.tsx": { + "no-restricted-imports": { + "count": 1 + } + }, + "src/components/chat_ui/mode_endpoint_mapping.tsx": { + "local/filename-pascal-case": { + "count": 1 + } + }, + "src/components/claude_code_plugins/MakeSkillPublicForm.tsx": { + "no-restricted-imports": { + "count": 2 }, "react-hooks/set-state-in-effect": { "count": 1 } }, - "src/components/cloudzero_export_modal.tsx": { - "no-restricted-imports": { + "src/components/claude_code_plugins/skill_detail.tsx": { + "local/filename-pascal-case": { "count": 1 + } + }, + "src/components/cloudzero_export_modal.tsx": { + "local/filename-pascal-case": { + "count": 1 + }, + "no-restricted-imports": { + "count": 2 }, "no-restricted-syntax": { "count": 3 @@ -1648,7 +2830,7 @@ }, "src/components/common_components/AccessGroupSelector.tsx": { "no-restricted-imports": { - "count": 1 + "count": 2 } }, "src/components/common_components/AutoRotationView.tsx": { @@ -1656,11 +2838,24 @@ "count": 1 } }, + "src/components/common_components/DefaultProxyAdminTag.tsx": { + "no-restricted-imports": { + "count": 1 + } + }, "src/components/common_components/DeleteResourceModal.tsx": { + "no-restricted-imports": { + "count": 1 + }, "react-hooks/set-state-in-effect": { "count": 1 } }, + "src/components/common_components/DurationSelect.tsx": { + "no-restricted-imports": { + "count": 1 + } + }, "src/components/common_components/Filters/FilterInput.tsx": { "react-hooks/set-state-in-effect": { "count": 1 @@ -1671,11 +2866,29 @@ "count": 1 } }, - "src/components/common_components/KeyLifecycleSettings.tsx": { + "src/components/common_components/IconActionButton/TableIconActionButtons/TableIconActionButton.tsx": { "no-restricted-imports": { "count": 1 } }, + "src/components/common_components/KeyLifecycleSettings.tsx": { + "local/no-complex-jsx-arrow": { + "count": 1 + }, + "no-restricted-imports": { + "count": 2 + } + }, + "src/components/common_components/LabeledField.tsx": { + "no-restricted-imports": { + "count": 1 + } + }, + "src/components/common_components/MemberTable.tsx": { + "no-restricted-imports": { + "count": 2 + } + }, "src/components/common_components/ModelAliasManager.tsx": { "no-restricted-imports": { "count": 1 @@ -1686,41 +2899,93 @@ }, "src/components/common_components/ModelSelector.tsx": { "no-restricted-imports": { - "count": 1 + "count": 2 }, "react-hooks/set-state-in-effect": { "count": 1 } }, + "src/components/common_components/NewBadge.tsx": { + "no-restricted-imports": { + "count": 1 + } + }, + "src/components/common_components/OrganizationDropdown.tsx": { + "local/no-complex-jsx-arrow": { + "count": 1 + }, + "no-restricted-imports": { + "count": 1 + } + }, "src/components/common_components/PassThroughGuardrailsSection.tsx": { "no-restricted-imports": { - "count": 1 + "count": 2 }, "react-hooks/set-state-in-effect": { "count": 1 } }, - "src/components/common_components/PassThroughSecuritySection.tsx": { + "src/components/common_components/PassThroughRoutesSelector.tsx": { "no-restricted-imports": { "count": 1 } }, + "src/components/common_components/PassThroughSecuritySection.tsx": { + "no-restricted-imports": { + "count": 2 + } + }, "src/components/common_components/PremiumLoggingSettings.tsx": { "no-restricted-imports": { "count": 1 } }, + "src/components/common_components/ProjectDropdown.tsx": { + "local/no-complex-jsx-arrow": { + "count": 1 + }, + "no-restricted-imports": { + "count": 1 + } + }, + "src/components/common_components/RateLimitTypeFormItem.test.tsx": { + "no-restricted-imports": { + "count": 1 + } + }, + "src/components/common_components/RateLimitTypeFormItem.tsx": { + "no-restricted-imports": { + "count": 1 + } + }, "src/components/common_components/RouterSettingsAccordion.tsx": { "no-restricted-imports": { "count": 1 } }, + "src/components/common_components/TableHeaderSortDropdown/TableHeaderSortDropdown.tsx": { + "no-restricted-imports": { + "count": 1 + } + }, + "src/components/common_components/budget_duration_dropdown.tsx": { + "local/filename-pascal-case": { + "count": 1 + }, + "no-restricted-imports": { + "count": 1 + } + }, "src/components/common_components/chartUtils.test.tsx": { "no-restricted-imports": { "count": 1 } }, "src/components/common_components/chartUtils.tsx": { + "local/filename-pascal-case": { + "count": 1 + }, "no-nested-ternary": { "count": 1 }, @@ -1729,16 +2994,25 @@ } }, "src/components/common_components/check_openapi_schema.tsx": { - "no-restricted-imports": { + "local/filename-pascal-case": { "count": 1 + }, + "no-restricted-imports": { + "count": 3 } }, "src/components/common_components/fetch_teams.tsx": { + "local/filename-pascal-case": { + "count": 1 + }, "max-params": { "count": 1 } }, "src/components/common_components/simple_table.tsx": { + "local/filename-pascal-case": { + "count": 1 + }, "no-nested-ternary": { "count": 1 }, @@ -1746,21 +3020,51 @@ "count": 1 } }, + "src/components/common_components/team_dropdown.tsx": { + "local/filename-pascal-case": { + "count": 1 + }, + "no-restricted-imports": { + "count": 1 + } + }, + "src/components/common_components/team_multi_select.tsx": { + "local/filename-pascal-case": { + "count": 1 + }, + "no-restricted-imports": { + "count": 1 + } + }, + "src/components/common_components/user_search_modal.tsx": { + "local/filename-pascal-case": { + "count": 1 + }, + "no-restricted-imports": { + "count": 1 + } + }, "src/components/constants.tsx": { + "local/filename-pascal-case": { + "count": 1 + }, "react-hooks/set-state-in-effect": { "count": 1 } }, "src/components/edit_auto_router/edit_auto_router_modal.tsx": { - "no-restricted-imports": { + "local/filename-pascal-case": { "count": 1 }, + "no-restricted-imports": { + "count": 2 + }, "react-hooks/immutability": { "count": 1 } }, "src/components/email_events/email_event_settings.tsx": { - "no-restricted-imports": { + "local/filename-pascal-case": { "count": 1 }, "react-hooks/immutability": { @@ -1768,34 +3072,139 @@ } }, "src/components/email_settings.tsx": { + "local/filename-pascal-case": { + "count": 1 + }, + "prefer-const": { + "count": 1 + } + }, + "src/components/guardrails/GuardrailSelector.tsx": { "no-restricted-imports": { "count": 1 } }, + "src/components/key_info_utils.tsx": { + "local/filename-pascal-case": { + "count": 1 + } + }, + "src/components/key_team_helpers/BudgetFallbacksEditor.tsx": { + "no-restricted-imports": { + "count": 1 + } + }, + "src/components/key_team_helpers/BudgetWindowsEditor.tsx": { + "no-restricted-imports": { + "count": 1 + } + }, + "src/components/key_team_helpers/TagRateLimitEditor.tsx": { + "no-restricted-imports": { + "count": 1 + } + }, + "src/components/key_team_helpers/fetch_available_models_team_key.tsx": { + "local/filename-pascal-case": { + "count": 1 + }, + "prefer-const": { + "count": 1 + } + }, "src/components/key_team_helpers/key_list.tsx": { + "local/filename-pascal-case": { + "count": 1 + }, "react-hooks/set-state-in-effect": { "count": 1 } }, + "src/components/key_team_helpers/transform_key_info.tsx": { + "local/filename-pascal-case": { + "count": 1 + } + }, "src/components/key_value_input.tsx": { + "local/filename-pascal-case": { + "count": 1 + }, "no-restricted-imports": { + "count": 2 + } + }, + "src/components/leftnav.tsx": { + "local/filename-pascal-case": { "count": 1 } }, "src/components/llm_calls/chat_completion.tsx": { + "local/filename-pascal-case": { + "count": 1 + }, "max-params": { "count": 1 }, "no-nested-ternary": { "count": 1 + }, + "prefer-const": { + "count": 1 + } + }, + "src/components/llm_calls/fetch_models.tsx": { + "local/filename-pascal-case": { + "count": 1 } }, "src/components/llm_calls/responses_api.tsx": { + "local/filename-pascal-case": { + "count": 1 + }, "max-params": { "count": 1 } }, + "src/components/logging_credentials/AccessControlFields.tsx": { + "no-restricted-imports": { + "count": 1 + } + }, + "src/components/logging_credentials/EditLoggingCredentialModal.tsx": { + "no-restricted-imports": { + "count": 1 + } + }, + "src/components/logging_credentials/LoggingExportersSelect.tsx": { + "no-restricted-imports": { + "count": 1 + } + }, + "src/components/logging_settings_view.tsx": { + "local/filename-pascal-case": { + "count": 1 + }, + "no-restricted-imports": { + "count": 1 + } + }, + "src/components/mcp_server_management/MCPServerSelector.tsx": { + "local/no-complex-jsx-arrow": { + "count": 1 + }, + "no-restricted-imports": { + "count": 1 + } + }, "src/components/mcp_server_management/MCPToolPermissions.tsx": { + "local/no-complex-jsx-arrow": { + "count": 1 + }, + "no-restricted-imports": { + "count": 2 + } + }, + "src/components/mcp_tools/ByokCredentialModal.tsx": { "no-restricted-imports": { "count": 1 } @@ -1803,6 +3212,9 @@ "src/components/mcp_tools/MCPToolArgumentsForm.tsx": { "no-nested-ternary": { "count": 5 + }, + "no-restricted-imports": { + "count": 1 } }, "src/components/mcp_tools/McpCrudPermissionPanel.tsx": { @@ -1810,33 +3222,69 @@ "count": 3 }, "no-restricted-imports": { + "count": 2 + } + }, + "src/components/mcp_tools/types.tsx": { + "local/filename-pascal-case": { "count": 1 } }, "src/components/model_add/CredentialModal.tsx": { "no-restricted-imports": { - "count": 1 + "count": 3 } }, - "src/components/model_add/reuse_credentials.tsx": { + "src/components/model_add/CredentialsPanel.test.tsx": { "no-restricted-imports": { "count": 1 } }, - "src/components/model_dashboard/all_models_table.tsx": { - "no-nested-ternary": { + "src/components/model_add/CredentialsPanel.tsx": { + "no-restricted-imports": { + "count": 1 + } + }, + "src/components/model_add/credential_form_helpers.test.ts": { + "no-restricted-imports": { + "count": 1 + } + }, + "src/components/model_add/credential_form_helpers.ts": { + "no-restricted-imports": { + "count": 1 + } + }, + "src/components/model_add/reuse_credentials.tsx": { + "local/filename-pascal-case": { + "count": 1 + }, + "no-restricted-imports": { + "count": 2 + } + }, + "src/components/model_dashboard/HealthCheckComponent.tsx": { + "no-restricted-imports": { + "count": 2 + } + }, + "src/components/model_dashboard/ModelSettingsModal/ModelSettingsModal.tsx": { + "no-restricted-imports": { + "count": 1 + } + }, + "src/components/model_filters.tsx": { + "local/filename-pascal-case": { "count": 1 }, "no-restricted-imports": { "count": 1 } }, - "src/components/model_filters.tsx": { - "no-restricted-imports": { - "count": 1 - } - }, "src/components/model_group_alias_settings.tsx": { + "local/filename-pascal-case": { + "count": 1 + }, "no-restricted-imports": { "count": 1 }, @@ -1845,46 +3293,77 @@ } }, "src/components/model_info_view.tsx": { + "local/filename-pascal-case": { + "count": 1 + }, + "local/no-complex-jsx-arrow": { + "count": 1 + }, + "max-lines": { + "count": 1 + }, "no-nested-ternary": { "count": 14 }, "no-restricted-imports": { - "count": 1 + "count": 2 + }, + "prefer-const": { + "count": 5 }, "react-hooks/set-state-in-effect": { "count": 1 } }, - "src/components/molecules/filter.tsx": { - "no-nested-ternary": { - "count": 2 - } - }, - "src/components/molecules/models/columns.test.tsx": { - "no-restricted-imports": { - "count": 1 - }, - "react/display-name": { + "src/components/molecules/cost_optimization_feedback_banner.tsx": { + "local/filename-pascal-case": { "count": 1 } }, - "src/components/molecules/models/columns.tsx": { - "max-params": { + "src/components/molecules/message_manager.tsx": { + "local/filename-pascal-case": { "count": 1 }, - "no-nested-ternary": { - "count": 2 - }, + "no-restricted-imports": { + "count": 2 + } + }, + "src/components/molecules/notifications_manager.test.tsx": { "no-restricted-imports": { "count": 1 } }, + "src/components/molecules/notifications_manager.tsx": { + "local/filename-pascal-case": { + "count": 1 + }, + "no-restricted-imports": { + "count": 3 + } + }, "src/components/navbar.test.tsx": { + "prefer-const": { + "count": 1 + }, "unused-imports/no-unused-imports": { "count": 1 } }, + "src/components/navbar.tsx": { + "local/filename-pascal-case": { + "count": 1 + }, + "no-restricted-imports": { + "count": 1 + } + }, "src/components/networking.tsx": { + "local/filename-pascal-case": { + "count": 1 + }, + "max-lines": { + "count": 1 + }, "max-params": { "count": 23 }, @@ -1893,14 +3372,28 @@ }, "no-restricted-syntax": { "count": 154 + }, + "prefer-const": { + "count": 33 } }, "src/components/object_permissions_view.tsx": { + "local/filename-pascal-case": { + "count": 1 + }, "no-restricted-imports": { "count": 1 } }, "src/components/onboarding_link.tsx": { + "local/filename-pascal-case": { + "count": 1 + }, + "no-restricted-imports": { + "count": 2 + } + }, + "src/components/organisms/RegenerateKeyModal.tsx": { "no-restricted-imports": { "count": 1 } @@ -1914,18 +3407,30 @@ } }, "src/components/organisms/create_key_button.tsx": { - "no-restricted-imports": { + "local/filename-pascal-case": { "count": 1 }, + "local/no-complex-jsx-arrow": { + "count": 2 + }, + "max-lines": { + "count": 1 + }, + "no-restricted-imports": { + "count": 2 + }, + "prefer-const": { + "count": 4 + }, "react-hooks/set-state-in-effect": { "count": 4 } }, "src/components/organization/organization_view.tsx": { - "no-restricted-imports": { + "local/filename-pascal-case": { "count": 1 }, - "unused-imports/no-unused-imports": { + "no-restricted-imports": { "count": 1 } }, @@ -1935,11 +3440,17 @@ } }, "src/components/pass_through_info.tsx": { - "no-restricted-imports": { + "local/filename-pascal-case": { "count": 1 + }, + "no-restricted-imports": { + "count": 2 } }, "src/components/per_user_usage.tsx": { + "local/filename-pascal-case": { + "count": 1 + }, "no-restricted-imports": { "count": 1 }, @@ -1949,7 +3460,7 @@ }, "src/components/permissions/AgentPermissions.tsx": { "no-restricted-imports": { - "count": 1 + "count": 2 } }, "src/components/permissions/MCPServerPermissions.tsx": { @@ -1957,7 +3468,7 @@ "count": 3 }, "no-restricted-imports": { - "count": 1 + "count": 2 } }, "src/components/permissions/VectorStorePermissions.tsx": { @@ -1968,19 +3479,58 @@ "src/components/policies/PolicySelector.tsx": { "no-nested-ternary": { "count": 1 - } - }, - "src/components/price_data_reload.tsx": { - "react-hooks/immutability": { - "count": 2 - } - }, - "src/components/public_model_hub.tsx": { + }, "no-restricted-imports": { "count": 1 } }, + "src/components/price_data_reload.tsx": { + "local/filename-pascal-case": { + "count": 1 + }, + "no-restricted-imports": { + "count": 1 + }, + "react-hooks/immutability": { + "count": 2 + } + }, + "src/components/provider_info_helpers.tsx": { + "local/filename-pascal-case": { + "count": 1 + }, + "prefer-const": { + "count": 3 + } + }, + "src/components/public_model_hub.tsx": { + "local/filename-pascal-case": { + "count": 1 + }, + "max-lines": { + "count": 1 + }, + "no-restricted-imports": { + "count": 2 + } + }, "src/components/query_param_input.tsx": { + "local/filename-pascal-case": { + "count": 1 + }, + "no-restricted-imports": { + "count": 2 + } + }, + "src/components/route_preview.tsx": { + "local/filename-pascal-case": { + "count": 1 + }, + "no-restricted-imports": { + "count": 1 + } + }, + "src/components/router_settings/LatencyBasedConfiguration.tsx": { "no-restricted-imports": { "count": 1 } @@ -1988,9 +3538,49 @@ "src/components/router_settings/ReliabilityRetriesSection.tsx": { "no-nested-ternary": { "count": 1 + }, + "no-restricted-imports": { + "count": 1 + } + }, + "src/components/router_settings/RoutingStrategySelector.tsx": { + "no-restricted-imports": { + "count": 1 + } + }, + "src/components/router_settings/TagFilteringToggle.tsx": { + "no-restricted-imports": { + "count": 1 + } + }, + "src/components/router_settings/index.tsx": { + "local/filename-pascal-case": { + "count": 1 + }, + "no-restricted-imports": { + "count": 1 + }, + "prefer-const": { + "count": 2 + } + }, + "src/components/routing_groups/RoutingGroupModal.tsx": { + "no-restricted-imports": { + "count": 1 + } + }, + "src/components/routing_groups/RoutingGroupsTable.tsx": { + "no-restricted-imports": { + "count": 2 } }, "src/components/routing_groups/index.tsx": { + "local/filename-pascal-case": { + "count": 1 + }, + "no-restricted-imports": { + "count": 1 + }, "react-hooks/preserve-manual-memoization": { "count": 1 } @@ -1998,17 +3588,45 @@ "src/components/search_tools/SearchToolSelector.tsx": { "no-nested-ternary": { "count": 1 - } - }, - "src/components/settings.tsx": { - "no-nested-ternary": { - "count": 2 }, "no-restricted-imports": { "count": 1 } }, + "src/components/settings.test.tsx": { + "no-restricted-imports": { + "count": 1 + } + }, + "src/components/settings.tsx": { + "local/filename-pascal-case": { + "count": 1 + }, + "local/no-complex-jsx-arrow": { + "count": 4 + }, + "max-lines": { + "count": 1 + }, + "no-nested-ternary": { + "count": 2 + }, + "no-restricted-imports": { + "count": 3 + }, + "prefer-const": { + "count": 7 + } + }, + "src/components/shared/CreatedKeyDisplay.tsx": { + "no-restricted-imports": { + "count": 1 + } + }, "src/components/shared/advanced_date_picker.tsx": { + "local/filename-pascal-case": { + "count": 1 + }, "no-restricted-imports": { "count": 1 }, @@ -2016,67 +3634,208 @@ "count": 3 } }, + "src/components/shared/chart_loader.tsx": { + "local/filename-pascal-case": { + "count": 1 + } + }, + "src/components/shared/charts/area_chart.tsx": { + "local/filename-pascal-case": { + "count": 1 + } + }, + "src/components/shared/charts/bar_chart.tsx": { + "local/filename-pascal-case": { + "count": 1 + } + }, + "src/components/shared/charts/chart_legend.tsx": { + "local/filename-pascal-case": { + "count": 1 + } + }, + "src/components/shared/charts/chart_tooltip.tsx": { + "local/filename-pascal-case": { + "count": 1 + } + }, + "src/components/shared/charts/donut_chart.tsx": { + "local/filename-pascal-case": { + "count": 1 + } + }, + "src/components/shared/charts/line_chart.tsx": { + "local/filename-pascal-case": { + "count": 1 + } + }, + "src/components/shared/errorUtils.tsx": { + "local/filename-pascal-case": { + "count": 1 + } + }, + "src/components/shared/form/FormField.tsx": { + "local/no-complex-jsx-arrow": { + "count": 1 + } + }, + "src/components/shared/form/field.tsx": { + "local/filename-pascal-case": { + "count": 1 + } + }, "src/components/shared/numerical_input.tsx": { + "local/filename-pascal-case": { + "count": 1 + }, "no-restricted-imports": { "count": 1 } }, + "src/components/shared/table_cells/cell_tooltip.tsx": { + "local/filename-pascal-case": { + "count": 1 + } + }, + "src/components/shared/table_cells/date_cell.tsx": { + "local/filename-pascal-case": { + "count": 1 + } + }, + "src/components/shared/table_cells/id_cell.tsx": { + "local/filename-pascal-case": { + "count": 1 + } + }, + "src/components/shared/table_cells/identity_cell.tsx": { + "local/filename-pascal-case": { + "count": 1 + } + }, + "src/components/shared/table_cells/models_cell.tsx": { + "local/filename-pascal-case": { + "count": 1 + } + }, + "src/components/shared/table_cells/money_cell.tsx": { + "local/filename-pascal-case": { + "count": 1 + } + }, + "src/components/shared/table_cells/spend_budget_cell.tsx": { + "local/filename-pascal-case": { + "count": 1 + } + }, + "src/components/shared/table_cells/status_badge.tsx": { + "local/filename-pascal-case": { + "count": 1 + } + }, "src/components/shared/usage_date_picker.tsx": { + "local/filename-pascal-case": { + "count": 1 + }, "no-restricted-imports": { "count": 1 } }, + "src/components/tag_management/TagSelector.tsx": { + "no-restricted-imports": { + "count": 1 + } + }, + "src/components/tag_management/types.tsx": { + "local/filename-pascal-case": { + "count": 1 + } + }, "src/components/team/EditMembership.tsx": { "no-nested-ternary": { "count": 1 }, "no-restricted-imports": { - "count": 1 + "count": 2 } }, "src/components/team/LoggingSettings.tsx": { + "no-restricted-imports": { + "count": 2 + } + }, + "src/components/team/MyUserTab.tsx": { "no-restricted-imports": { "count": 1 } }, "src/components/team/TeamInfo.tsx": { + "max-lines": { + "count": 1 + }, "no-nested-ternary": { "count": 3 }, "no-restricted-imports": { - "count": 1 + "count": 2 }, "react-hooks/set-state-in-effect": { "count": 1 } }, + "src/components/team/TeamMemberTab.tsx": { + "local/no-complex-jsx-arrow": { + "count": 1 + }, + "no-restricted-imports": { + "count": 2 + } + }, "src/components/team/TeamVirtualKeysTable.tsx": { "no-nested-ternary": { "count": 1 }, "no-restricted-imports": { - "count": 1 + "count": 2 } }, "src/components/team/member_permissions.tsx": { - "no-restricted-imports": { + "local/filename-pascal-case": { "count": 1 }, + "no-restricted-imports": { + "count": 2 + }, "react-hooks/set-state-in-effect": { "count": 1 } }, + "src/components/team/permission_definitions.tsx": { + "local/filename-pascal-case": { + "count": 1 + } + }, "src/components/team/useMyTeamMember.ts": { "no-restricted-syntax": { "count": 1 } }, + "src/components/templates/KeyInfoHeader.tsx": { + "no-restricted-imports": { + "count": 2 + } + }, "src/components/templates/key_edit_view.tsx": { + "local/filename-pascal-case": { + "count": 1 + }, + "local/no-complex-jsx-arrow": { + "count": 2 + }, "no-nested-ternary": { "count": 2 }, "no-restricted-imports": { - "count": 1 + "count": 2 } }, "src/components/templates/key_info_view.test.tsx": { @@ -2085,41 +3844,263 @@ } }, "src/components/templates/key_info_view.tsx": { + "local/filename-pascal-case": { + "count": 1 + }, + "max-lines": { + "count": 1 + }, "no-nested-ternary": { "count": 1 }, "no-restricted-imports": { - "count": 1 + "count": 2 }, "react-hooks/set-state-in-effect": { "count": 1 } }, - "src/components/user_agent_activity.tsx": { + "src/components/ui/AntDLoadingSpinner.tsx": { "no-restricted-imports": { - "count": 2 + "count": 1 + } + }, + "src/components/ui/alert-dialog.tsx": { + "local/filename-pascal-case": { + "count": 1 + } + }, + "src/components/ui/avatar.tsx": { + "local/filename-pascal-case": { + "count": 1 + } + }, + "src/components/ui/badge.tsx": { + "local/filename-pascal-case": { + "count": 1 + } + }, + "src/components/ui/breadcrumb.tsx": { + "local/filename-pascal-case": { + "count": 1 + } + }, + "src/components/ui/button.tsx": { + "local/filename-pascal-case": { + "count": 1 + } + }, + "src/components/ui/card.tsx": { + "local/filename-pascal-case": { + "count": 1 + } + }, + "src/components/ui/chart.tsx": { + "local/filename-pascal-case": { + "count": 1 + } + }, + "src/components/ui/checkbox.tsx": { + "local/filename-pascal-case": { + "count": 1 + } + }, + "src/components/ui/collapsible.tsx": { + "local/filename-pascal-case": { + "count": 1 + } + }, + "src/components/ui/combobox.tsx": { + "local/filename-pascal-case": { + "count": 1 + } + }, + "src/components/ui/dialog.tsx": { + "local/filename-pascal-case": { + "count": 1 + } + }, + "src/components/ui/dropdown-menu.tsx": { + "local/filename-pascal-case": { + "count": 1 + } + }, + "src/components/ui/hover-card.tsx": { + "local/filename-pascal-case": { + "count": 1 + } + }, + "src/components/ui/input-group.tsx": { + "local/filename-pascal-case": { + "count": 1 + } + }, + "src/components/ui/input.tsx": { + "local/filename-pascal-case": { + "count": 1 + } + }, + "src/components/ui/label.tsx": { + "local/filename-pascal-case": { + "count": 1 + } + }, + "src/components/ui/meter.tsx": { + "local/filename-pascal-case": { + "count": 1 + } + }, + "src/components/ui/popover.tsx": { + "local/filename-pascal-case": { + "count": 1 + } + }, + "src/components/ui/radio-group.tsx": { + "local/filename-pascal-case": { + "count": 1 + } + }, + "src/components/ui/scroll-area.tsx": { + "local/filename-pascal-case": { + "count": 1 + } + }, + "src/components/ui/select.tsx": { + "local/filename-pascal-case": { + "count": 1 + } + }, + "src/components/ui/separator.tsx": { + "local/filename-pascal-case": { + "count": 1 + } + }, + "src/components/ui/sheet.tsx": { + "local/filename-pascal-case": { + "count": 1 + } + }, + "src/components/ui/sidebar.tsx": { + "local/filename-pascal-case": { + "count": 1 + } + }, + "src/components/ui/skeleton.tsx": { + "local/filename-pascal-case": { + "count": 1 + } + }, + "src/components/ui/switch.tsx": { + "local/filename-pascal-case": { + "count": 1 + } + }, + "src/components/ui/table.tsx": { + "local/filename-pascal-case": { + "count": 1 + } + }, + "src/components/ui/tabs.tsx": { + "local/filename-pascal-case": { + "count": 1 + } + }, + "src/components/ui/textarea.tsx": { + "local/filename-pascal-case": { + "count": 1 + } + }, + "src/components/ui/tooltip.tsx": { + "local/filename-pascal-case": { + "count": 1 + } + }, + "src/components/ui/ui-loading-spinner.tsx": { + "local/filename-pascal-case": { + "count": 1 + } + }, + "src/components/update_model_credentials_modal.tsx": { + "local/filename-pascal-case": { + "count": 1 + }, + "no-restricted-imports": { + "count": 1 + } + }, + "src/components/user_agent_activity.tsx": { + "local/filename-pascal-case": { + "count": 1 + }, + "no-restricted-imports": { + "count": 3 }, "react-hooks/set-state-in-effect": { "count": 1 } }, "src/components/user_dashboard.tsx": { + "local/filename-pascal-case": { + "count": 1 + }, "no-restricted-imports": { "count": 1 }, + "prefer-const": { + "count": 1 + }, "react-hooks/set-state-in-effect": { "count": 2 } }, + "src/components/vector_store_management/VectorStoreSelector.test.tsx": { + "no-restricted-imports": { + "count": 1 + } + }, + "src/components/vector_store_management/VectorStoreSelector.tsx": { + "no-restricted-imports": { + "count": 1 + } + }, + "src/components/vector_store_management/types.tsx": { + "local/filename-pascal-case": { + "count": 1 + } + }, + "src/components/vector_store_providers.tsx": { + "local/filename-pascal-case": { + "count": 1 + } + }, + "src/components/view_logs/AuditLogDrawer/AuditLogDrawer.tsx": { + "no-restricted-imports": { + "count": 1 + } + }, + "src/components/view_logs/CostBreakdownViewer.tsx": { + "no-restricted-imports": { + "count": 1 + } + }, "src/components/view_logs/EvalViewer/EvalViewer.tsx": { + "local/no-complex-jsx-arrow": { + "count": 1 + }, "no-nested-ternary": { "count": 1 + }, + "no-restricted-imports": { + "count": 1 } }, "src/components/view_logs/GuardrailViewer/CompliancePanel.tsx": { "no-nested-ternary": { "count": 2 }, + "no-restricted-imports": { + "count": 1 + }, "react-hooks/set-state-in-effect": { "count": 1 } @@ -2132,26 +4113,90 @@ "src/components/view_logs/GuardrailViewer/GuardrailViewer.tsx": { "no-nested-ternary": { "count": 4 + }, + "no-restricted-imports": { + "count": 1 + } + }, + "src/components/view_logs/LogDetailsDrawer/CollapsibleMessage.tsx": { + "no-restricted-imports": { + "count": 1 + } + }, + "src/components/view_logs/LogDetailsDrawer/DrawerHeader.tsx": { + "no-restricted-imports": { + "count": 1 + } + }, + "src/components/view_logs/LogDetailsDrawer/HistoryTree.tsx": { + "no-restricted-imports": { + "count": 1 + } + }, + "src/components/view_logs/LogDetailsDrawer/JsonViewer.tsx": { + "no-restricted-imports": { + "count": 1 } }, "src/components/view_logs/LogDetailsDrawer/LogDetailContent.tsx": { "no-nested-ternary": { "count": 3 + }, + "no-restricted-imports": { + "count": 1 } }, "src/components/view_logs/LogDetailsDrawer/LogDetailsDrawer.tsx": { "no-nested-ternary": { - "count": 3 + "count": 2 + }, + "no-restricted-imports": { + "count": 1 }, "react-hooks/set-state-in-effect": { "count": 2 } }, + "src/components/view_logs/LogDetailsDrawer/OutputCard.tsx": { + "no-restricted-imports": { + "count": 1 + } + }, "src/components/view_logs/LogDetailsDrawer/RealtimePrettyView.test.tsx": { "unused-imports/no-unused-imports": { "count": 2 } }, + "src/components/view_logs/LogDetailsDrawer/RealtimePrettyView.tsx": { + "no-restricted-imports": { + "count": 1 + } + }, + "src/components/view_logs/LogDetailsDrawer/SectionHeader.tsx": { + "no-restricted-imports": { + "count": 1 + } + }, + "src/components/view_logs/LogDetailsDrawer/SimpleMessageBlock.tsx": { + "no-restricted-imports": { + "count": 1 + } + }, + "src/components/view_logs/LogDetailsDrawer/SimpleToolCallBlock.tsx": { + "no-restricted-imports": { + "count": 1 + } + }, + "src/components/view_logs/LogDetailsDrawer/TokenFlow.tsx": { + "no-restricted-imports": { + "count": 1 + } + }, + "src/components/view_logs/LogDetailsDrawer/TruncatedValue.tsx": { + "no-restricted-imports": { + "count": 1 + } + }, "src/components/view_logs/LogDetailsDrawer/prettyMessagesUtils.ts": { "no-nested-ternary": { "count": 1 @@ -2162,29 +4207,83 @@ "count": 2 } }, - "src/components/view_logs/LogsTableToolbar.tsx": { - "no-nested-ternary": { - "count": 4 + "src/components/view_logs/ToolsSection/FormattedToolView.tsx": { + "no-restricted-imports": { + "count": 1 + } + }, + "src/components/view_logs/ToolsSection/ToolExpandedContent.tsx": { + "no-restricted-imports": { + "count": 1 + } + }, + "src/components/view_logs/ToolsSection/ToolItem.tsx": { + "no-restricted-imports": { + "count": 1 + } + }, + "src/components/view_logs/ToolsSection/ToolsSection.tsx": { + "no-restricted-imports": { + "count": 1 + } + }, + "src/components/view_logs/VectorStoreViewer.tsx": { + "no-restricted-imports": { + "count": 1 + } + }, + "src/components/view_logs/columns.tsx": { + "local/filename-pascal-case": { + "count": 1 } }, "src/components/view_logs/index.tsx": { - "no-restricted-imports": { + "local/filename-pascal-case": { "count": 1 }, - "react-hooks/set-state-in-effect": { + "no-restricted-imports": { + "count": 1 + } + }, + "src/components/view_logs/log_filter_logic.tsx": { + "local/filename-pascal-case": { + "count": 1 + } + }, + "src/components/view_logs/logs_utils.tsx": { + "local/filename-pascal-case": { "count": 1 } }, "src/components/view_logs/table.tsx": { + "local/filename-pascal-case": { + "count": 1 + }, "no-nested-ternary": { "count": 2 } }, + "src/components/view_model/model_name_display.tsx": { + "local/filename-pascal-case": { + "count": 1 + } + }, "src/components/view_user_spend.tsx": { + "local/filename-pascal-case": { + "count": 1 + }, + "prefer-const": { + "count": 3 + }, "react-hooks/set-state-in-effect": { "count": 2 } }, + "src/contexts/AntdGlobalProvider.tsx": { + "no-restricted-imports": { + "count": 1 + } + }, "src/contexts/AuthContext.tsx": { "react-hooks/set-state-in-effect": { "count": 1 @@ -2211,11 +4310,17 @@ } }, "src/hooks/useMcpOAuthFlow.tsx": { + "local/filename-pascal-case": { + "count": 1 + }, "react-hooks/set-state-in-effect": { "count": 1 } }, "src/hooks/useTestMCPConnection.tsx": { + "local/filename-pascal-case": { + "count": 1 + }, "no-nested-ternary": { "count": 1 }, @@ -2224,6 +4329,9 @@ } }, "src/hooks/useToolsOAuthFlow.tsx": { + "local/filename-pascal-case": { + "count": 1 + }, "react-hooks/refs": { "count": 1 }, @@ -2232,6 +4340,9 @@ } }, "src/hooks/useUserMcpOAuthFlow.tsx": { + "local/filename-pascal-case": { + "count": 1 + }, "react-hooks/set-state-in-effect": { "count": 1 } diff --git a/ui/litellm-dashboard/eslint.config.mjs b/ui/litellm-dashboard/eslint.config.mjs index 0cf5b4ff655..b5e876bfbba 100644 --- a/ui/litellm-dashboard/eslint.config.mjs +++ b/ui/litellm-dashboard/eslint.config.mjs @@ -19,12 +19,13 @@ const eslintConfig = [ "unused-imports/no-unused-imports": "error", "local/no-large-inline-object-arg": "warn", "local/no-long-condition-chain": "warn", + "local/no-complex-jsx-arrow": ["error", { maxStatements: 2 }], "@typescript-eslint/no-explicit-any": "warn", "no-console": ["warn", { allow: ["warn", "error"] }], "@typescript-eslint/no-unused-vars": "off", "@typescript-eslint/no-unused-expressions": "off", "@typescript-eslint/ban-ts-comment": "off", - "prefer-const": "off", + "prefer-const": "error", "no-empty": "off", "no-prototype-builtins": "off", "no-useless-catch": "off", @@ -51,13 +52,32 @@ const eslintConfig = [ patterns: [ { group: ["@tremor/react", "@tremor/react/*"], - message: "@tremor/react is being phased out; build new UI with antd instead of adding tremor imports.", + message: + "@tremor/react is being phased out; build new UI with shadcn/ui primitives instead of adding tremor imports.", + }, + { + group: ["antd", "antd/*"], + message: + "antd is being phased out; build new UI with shadcn/ui primitives instead of adding antd imports.", }, ], }, ], }, }, + { + files: ["src/**/*.tsx"], + rules: { + "local/filename-pascal-case": "error", + }, + }, + { + files: ["src/**/*.{ts,tsx}"], + ignores: ["src/**/*.test.{ts,tsx}", "src/**/*.spec.{ts,tsx}", "src/data/**"], + rules: { + "max-lines": ["error", { max: 800, skipBlankLines: true, skipComments: true }], + }, + }, { files: ["src/lib/http/**"], rules: { diff --git a/ui/litellm-dashboard/knip.json b/ui/litellm-dashboard/knip.json index afed6b0f90e..48b39e8122d 100644 --- a/ui/litellm-dashboard/knip.json +++ b/ui/litellm-dashboard/knip.json @@ -1,7 +1,7 @@ { "$schema": "https://unpkg.com/knip@5/schema.json", "entry": ["scripts/**/*.{ts,mjs}", "src/components/ui/**/*.{ts,tsx}"], - "project": ["src/**/*.{ts,tsx}", "tests/**/*.{ts,tsx}", "scripts/**/*.{ts,mjs}", "e2e_tests/**/*.ts"], + "project": ["src/**/*.{ts,tsx}", "tests/**/*.{ts,tsx}", "scripts/**/*.{ts,mjs}"], "ignore": ["src/lib/http/schema.d.ts"], "ignoreDependencies": [ "openapi-typescript", @@ -10,14 +10,6 @@ "tailwindcss", "tw-animate-css" ], - "playwright": { - "config": [ - "e2e_tests/playwright.config.ts", - "e2e_tests/serverRootPath.config.ts", - "e2e_tests/migration.serverRootPath.config.ts" - ], - "entry": ["e2e_tests/**/*.spec.ts", "e2e_tests/**/*.setup.ts", "e2e_tests/globalSetup.ts"] - }, "vitest": { "config": ["vitest.config.ts"] }, diff --git a/ui/litellm-dashboard/package-lock.json b/ui/litellm-dashboard/package-lock.json index 742c1e4a63f..5f6b4b889b1 100644 --- a/ui/litellm-dashboard/package-lock.json +++ b/ui/litellm-dashboard/package-lock.json @@ -27,7 +27,7 @@ "jwt-decode": "4.0.0", "lucide-react": "0.513.0", "moment": "2.30.1", - "next": "16.2.6", + "next": "16.2.11", "openai": "4.104.0", "openapi-fetch": "^0.17.0", "openapi-react-query": "^0.5.4", @@ -47,7 +47,6 @@ }, "devDependencies": { "@eslint/js": "9.39.2", - "@playwright/test": "1.58.1", "@tailwindcss/forms": "0.5.11", "@tailwindcss/postcss": "4.3.2", "@testing-library/dom": "10.4.1", @@ -62,7 +61,7 @@ "@vitest/coverage-v8": "3.2.6", "@vitest/ui": "3.2.6", "eslint": "9.39.2", - "eslint-config-next": "16.2.6", + "eslint-config-next": "16.2.11", "eslint-config-prettier": "10.1.8", "eslint-plugin-unused-imports": "4.3.0", "jsdom": "27.4.0", @@ -2245,15 +2244,15 @@ } }, "node_modules/@next/env": { - "version": "16.2.6", - "resolved": "https://registry.npmjs.org/@next/env/-/env-16.2.6.tgz", - "integrity": "sha512-gd8HoHN4ufj73WmR3JmVolrpJR47ILK6LouP5xElPglaVxir6e1a7VzvTvDWkOoPXT9rkkTzyCxBu4yeZfZwcw==", + "version": "16.2.11", + "resolved": "https://registry.npmjs.org/@next/env/-/env-16.2.11.tgz", + "integrity": "sha512-0do5A3BJ2gxWr0ZCMcD6BhW+e595jyxdTl3rXTS6lOtD8ektMiW6CO+EPwt1Eca1DBnm90r/7GdiKWBKxH++DA==", "license": "MIT" }, "node_modules/@next/eslint-plugin-next": { - "version": "16.2.6", - "resolved": "https://registry.npmjs.org/@next/eslint-plugin-next/-/eslint-plugin-next-16.2.6.tgz", - "integrity": "sha512-Z8l6o4JWKUl755x4R+wogD86KPeU+Ckw4K+SYG4kHeOJtRenDeK+OSbGcqZpDtbwn9DsJVdir2UxmwXuinUbUw==", + "version": "16.2.11", + "resolved": "https://registry.npmjs.org/@next/eslint-plugin-next/-/eslint-plugin-next-16.2.11.tgz", + "integrity": "sha512-vMEf/aXOpzFFdtIvFYOnIDPKb0xBbrXONsz83CcKdRrekfxNdL8PNkq5qHqAHSXVlIifnX68LOMaxr3z5PkeLQ==", "dev": true, "license": "MIT", "dependencies": { @@ -2261,9 +2260,9 @@ } }, "node_modules/@next/swc-darwin-arm64": { - "version": "16.2.6", - "resolved": "https://registry.npmjs.org/@next/swc-darwin-arm64/-/swc-darwin-arm64-16.2.6.tgz", - "integrity": "sha512-ZJGkkcNfYgrrMkqOdZ7zoLa1TOy0qpcMfk/z4Mh/FKUz40gVO+HNQWqmLxf67Z5WB64DRp0dhEbyHfel+6sJUg==", + "version": "16.2.11", + "resolved": "https://registry.npmjs.org/@next/swc-darwin-arm64/-/swc-darwin-arm64-16.2.11.tgz", + "integrity": "sha512-wryL4pjKmDwGv2ox6+GZDFxvmtSRLqApBR8kL1j4+vhB7Z5vJC/zAnXpiR9Xkfzl0AS8WLMnsuGV/UKI67/rrw==", "cpu": [ "arm64" ], @@ -2277,9 +2276,9 @@ } }, "node_modules/@next/swc-darwin-x64": { - "version": "16.2.6", - "resolved": "https://registry.npmjs.org/@next/swc-darwin-x64/-/swc-darwin-x64-16.2.6.tgz", - "integrity": "sha512-v/YLBHIY132Ced3puBJ7YJKw1lqsCrgcNo2aRJlCEyQrrCeRJlvGlnmxhPxNQI3KE3N1DN5r9TPNPvka3nq5RQ==", + "version": "16.2.11", + "resolved": "https://registry.npmjs.org/@next/swc-darwin-x64/-/swc-darwin-x64-16.2.11.tgz", + "integrity": "sha512-aZl2j4f/fLyjQvOhv0Oe9UaMAQHolYpKhctsoYzplSumKJKPUmgjcf6545aBtysLTcu994TREd0+pSgNE4ohmg==", "cpu": [ "x64" ], @@ -2293,9 +2292,9 @@ } }, "node_modules/@next/swc-linux-arm64-gnu": { - "version": "16.2.6", - "resolved": "https://registry.npmjs.org/@next/swc-linux-arm64-gnu/-/swc-linux-arm64-gnu-16.2.6.tgz", - "integrity": "sha512-RPOvqlYBbcQjkz9VQQDZ2T2bARIjXZV1KFlt+V2Mr6SW/e4I9fcKsaA0hdyf2FHoTlsV2xnBd5Y912rP/1Ce6w==", + "version": "16.2.11", + "resolved": "https://registry.npmjs.org/@next/swc-linux-arm64-gnu/-/swc-linux-arm64-gnu-16.2.11.tgz", + "integrity": "sha512-5jEriyEnH/LWFy27L2ZG0XaLlyEJIjhsImEsiS9P563PKEVp2BVups/xfOucIrsvVntp11oNcZwjHvaDPYVB5g==", "cpu": [ "arm64" ], @@ -2309,9 +2308,9 @@ } }, "node_modules/@next/swc-linux-arm64-musl": { - "version": "16.2.6", - "resolved": "https://registry.npmjs.org/@next/swc-linux-arm64-musl/-/swc-linux-arm64-musl-16.2.6.tgz", - "integrity": "sha512-URUTu1+dMkxJsPFgm+OeEvq9wf5sujw0EvgYy80TDGHTSLTnIHeqb0Eu8A3sC95IRgjejQL+kC4mw+4yPxiAXA==", + "version": "16.2.11", + "resolved": "https://registry.npmjs.org/@next/swc-linux-arm64-musl/-/swc-linux-arm64-musl-16.2.11.tgz", + "integrity": "sha512-eIjcpx2fnnFSSkZDbTxy74KnokUXDjfoLClpWelfgHLf621aTqswhwXQ7GkD5K5rplrS6LZ/Bj+mVuvzluBOEg==", "cpu": [ "arm64" ], @@ -2325,9 +2324,9 @@ } }, "node_modules/@next/swc-linux-x64-gnu": { - "version": "16.2.6", - "resolved": "https://registry.npmjs.org/@next/swc-linux-x64-gnu/-/swc-linux-x64-gnu-16.2.6.tgz", - "integrity": "sha512-DOj182mPV8G3UkrayLoREM5YEYI+Dk5wv7Ox9xl1fFibAELEsFD0lDPfHIeILlutMMfdyhlzYPELG3peuKaurw==", + "version": "16.2.11", + "resolved": "https://registry.npmjs.org/@next/swc-linux-x64-gnu/-/swc-linux-x64-gnu-16.2.11.tgz", + "integrity": "sha512-8WgzpaWMs46qJT9kiV47cje86L0x/Mu9t8/Gwj+pnbgW3rETVfCnaScPjlYUwNScpOozdcIMHWmAvuZJUonR2w==", "cpu": [ "x64" ], @@ -2341,9 +2340,9 @@ } }, "node_modules/@next/swc-linux-x64-musl": { - "version": "16.2.6", - "resolved": "https://registry.npmjs.org/@next/swc-linux-x64-musl/-/swc-linux-x64-musl-16.2.6.tgz", - "integrity": "sha512-HKQ5SP/V/ub73UvF7n/zeJlxk2kLmtL7Wzrg4WfmkjmNos5onJ2tKu7yZOPdL18A6Svfn3max29ym+ry7NkK4g==", + "version": "16.2.11", + "resolved": "https://registry.npmjs.org/@next/swc-linux-x64-musl/-/swc-linux-x64-musl-16.2.11.tgz", + "integrity": "sha512-I3UgPds7G4ZYnTb/H+5GBGuUT2DhAk6j0mL6A4s63RjFs74wB2hOWP0vaxsK+3NJraExt3eYEPQ/UtT0x/64Nw==", "cpu": [ "x64" ], @@ -2357,9 +2356,9 @@ } }, "node_modules/@next/swc-win32-arm64-msvc": { - "version": "16.2.6", - "resolved": "https://registry.npmjs.org/@next/swc-win32-arm64-msvc/-/swc-win32-arm64-msvc-16.2.6.tgz", - "integrity": "sha512-LZXpTlPyS5v7HhSmnvsLGP3iIYgYOBnc8r8ArlT55sGHV89bR2HlDdBjWQ+PY6SJMmk8TuVGFuxalnP3k/0Dwg==", + "version": "16.2.11", + "resolved": "https://registry.npmjs.org/@next/swc-win32-arm64-msvc/-/swc-win32-arm64-msvc-16.2.11.tgz", + "integrity": "sha512-n89CjtcThnjrwgJMAiI5xbqwLY51zvwC9tSlArmVndAJLYVl9T9UAdlkXTmZvE++idoXe8KdglQlhNRdUp1c6g==", "cpu": [ "arm64" ], @@ -2373,9 +2372,9 @@ } }, "node_modules/@next/swc-win32-x64-msvc": { - "version": "16.2.6", - "resolved": "https://registry.npmjs.org/@next/swc-win32-x64-msvc/-/swc-win32-x64-msvc-16.2.6.tgz", - "integrity": "sha512-F0+4i0h9J6C4eE3EAPWsoCk7UW/dbzOjyzxY0qnDUOYFu6FFmdZ6l97/XdV3/Nz3VYyO7UWjyEJUXkGqcoXfMA==", + "version": "16.2.11", + "resolved": "https://registry.npmjs.org/@next/swc-win32-x64-msvc/-/swc-win32-x64-msvc-16.2.11.tgz", + "integrity": "sha512-md8CLNggS1Dx9pUgApzps5uAf+N8GN9xywzmNx9vHAWo94HtBwCCqkSnhIrdfQe83Dhz8Lfo/20Nb1Zxal092w==", "cpu": [ "x64" ], @@ -2723,8 +2722,9 @@ "version": "1.58.1", "resolved": "https://registry.npmjs.org/@playwright/test/-/test-1.58.1.tgz", "integrity": "sha512-6LdVIUERWxQMmUSSQi0I53GgCBYgM2RpGngCPY7hSeju+VrKjq3lvs7HpJoPbDiY5QM5EYRtRX5fvrinnMAz3w==", - "devOptional": true, "license": "Apache-2.0", + "optional": true, + "peer": true, "dependencies": { "playwright": "1.58.1" }, @@ -3620,6 +3620,72 @@ "node": ">=14.0.0" } }, + "node_modules/@tailwindcss/oxide-wasm32-wasi/node_modules/@emnapi/core": { + "version": "1.11.1", + "dev": true, + "inBundle": true, + "license": "MIT", + "optional": true, + "dependencies": { + "@emnapi/wasi-threads": "1.2.2", + "tslib": "^2.4.0" + } + }, + "node_modules/@tailwindcss/oxide-wasm32-wasi/node_modules/@emnapi/runtime": { + "version": "1.11.1", + "dev": true, + "inBundle": true, + "license": "MIT", + "optional": true, + "dependencies": { + "tslib": "^2.4.0" + } + }, + "node_modules/@tailwindcss/oxide-wasm32-wasi/node_modules/@emnapi/wasi-threads": { + "version": "1.2.2", + "dev": true, + "inBundle": true, + "license": "MIT", + "optional": true, + "dependencies": { + "tslib": "^2.4.0" + } + }, + "node_modules/@tailwindcss/oxide-wasm32-wasi/node_modules/@napi-rs/wasm-runtime": { + "version": "1.1.4", + "dev": true, + "inBundle": true, + "license": "MIT", + "optional": true, + "dependencies": { + "@tybys/wasm-util": "^0.10.1" + }, + "funding": { + "type": "github", + "url": "https://github.com/sponsors/Brooooooklyn" + }, + "peerDependencies": { + "@emnapi/core": "^1.7.1", + "@emnapi/runtime": "^1.7.1" + } + }, + "node_modules/@tailwindcss/oxide-wasm32-wasi/node_modules/@tybys/wasm-util": { + "version": "0.10.2", + "dev": true, + "inBundle": true, + "license": "MIT", + "optional": true, + "dependencies": { + "tslib": "^2.4.0" + } + }, + "node_modules/@tailwindcss/oxide-wasm32-wasi/node_modules/tslib": { + "version": "2.8.1", + "dev": true, + "inBundle": true, + "license": "0BSD", + "optional": true + }, "node_modules/@tailwindcss/oxide-win32-arm64-msvc": { "version": "4.3.2", "resolved": "https://registry.npmjs.org/@tailwindcss/oxide-win32-arm64-msvc/-/oxide-win32-arm64-msvc-4.3.2.tgz", @@ -6642,13 +6708,13 @@ } }, "node_modules/eslint-config-next": { - "version": "16.2.6", - "resolved": "https://registry.npmjs.org/eslint-config-next/-/eslint-config-next-16.2.6.tgz", - "integrity": "sha512-z2ELYSkyrrJ6cuunTU8vhsT/RpouPkjaSah06nVW6Rg2Hpg0Vs8s497/e5s8G8qtdp4ccsiovz5P1rv+5VSW2Q==", + "version": "16.2.11", + "resolved": "https://registry.npmjs.org/eslint-config-next/-/eslint-config-next-16.2.11.tgz", + "integrity": "sha512-FIpbK/dUyxUExchDB7eBg3k+VU8R2iR/Cx9/kqTBUTFv2bOIR9aRrpno4rvAQ9VhiPQAyFKNA2NlZwouGWtclA==", "dev": true, "license": "MIT", "dependencies": { - "@next/eslint-plugin-next": "16.2.6", + "@next/eslint-plugin-next": "16.2.11", "eslint-import-resolver-node": "^0.3.6", "eslint-import-resolver-typescript": "^3.5.2", "eslint-plugin-import": "^2.32.0", @@ -7361,7 +7427,6 @@ "version": "2.3.2", "resolved": "https://registry.npmjs.org/fsevents/-/fsevents-2.3.2.tgz", "integrity": "sha512-xiqMQR4xAeHTuB9uWm+fFRcIOgKBMiOBP+eXiyT7jsgVCq1bkVygt00oASowB7EdtpOHaaPgKt812P9ab+DDKA==", - "dev": true, "hasInstallScript": true, "license": "MIT", "optional": true, @@ -10310,12 +10375,12 @@ "license": "MIT" }, "node_modules/next": { - "version": "16.2.6", - "resolved": "https://registry.npmjs.org/next/-/next-16.2.6.tgz", - "integrity": "sha512-qOVgKJg1+At15NpeUP+eJgCHvTCgXsogweq87Ri/Ix7PkqQHg4sdaXmSFqKlgaIXE4kW0g25LE68W87UANlHtw==", + "version": "16.2.11", + "resolved": "https://registry.npmjs.org/next/-/next-16.2.11.tgz", + "integrity": "sha512-B339zaqbyK8cmxhoAvLrcwoabwCP1wz21zSzfqxqXAemTu2BXnH7tQnfcglKv1vnMUIDBc+Hth7XODQriTZiRQ==", "license": "MIT", "dependencies": { - "@next/env": "16.2.6", + "@next/env": "16.2.11", "@swc/helpers": "0.5.15", "baseline-browser-mapping": "^2.9.19", "caniuse-lite": "^1.0.30001579", @@ -10329,14 +10394,14 @@ "node": ">=20.9.0" }, "optionalDependencies": { - "@next/swc-darwin-arm64": "16.2.6", - "@next/swc-darwin-x64": "16.2.6", - "@next/swc-linux-arm64-gnu": "16.2.6", - "@next/swc-linux-arm64-musl": "16.2.6", - "@next/swc-linux-x64-gnu": "16.2.6", - "@next/swc-linux-x64-musl": "16.2.6", - "@next/swc-win32-arm64-msvc": "16.2.6", - "@next/swc-win32-x64-msvc": "16.2.6", + "@next/swc-darwin-arm64": "16.2.11", + "@next/swc-darwin-x64": "16.2.11", + "@next/swc-linux-arm64-gnu": "16.2.11", + "@next/swc-linux-arm64-musl": "16.2.11", + "@next/swc-linux-x64-gnu": "16.2.11", + "@next/swc-linux-x64-musl": "16.2.11", + "@next/swc-win32-arm64-msvc": "16.2.11", + "@next/swc-win32-x64-msvc": "16.2.11", "sharp": "^0.34.5" }, "peerDependencies": { @@ -10948,8 +11013,9 @@ "version": "1.58.1", "resolved": "https://registry.npmjs.org/playwright/-/playwright-1.58.1.tgz", "integrity": "sha512-+2uTZHxSCcxjvGc5C891LrS1/NlxglGxzrC4seZiVjcYVQfUa87wBL6rTDqzGjuoWNjnBzRqKmF6zRYGMvQUaQ==", - "devOptional": true, "license": "Apache-2.0", + "optional": true, + "peer": true, "dependencies": { "playwright-core": "1.58.1" }, @@ -10967,8 +11033,9 @@ "version": "1.58.1", "resolved": "https://registry.npmjs.org/playwright-core/-/playwright-core-1.58.1.tgz", "integrity": "sha512-bcWzOaTxcW+VOOGBCQgnaKToLJ65d6AqfLVKEWvexyS3AS6rbXl+xdpYRMGSRBClPvyj44njOWoxjNdL/H9UNg==", - "devOptional": true, "license": "Apache-2.0", + "optional": true, + "peer": true, "bin": { "playwright-core": "cli.js" }, diff --git a/ui/litellm-dashboard/package.json b/ui/litellm-dashboard/package.json index 0f54c536297..63bcdfb6076 100644 --- a/ui/litellm-dashboard/package.json +++ b/ui/litellm-dashboard/package.json @@ -14,10 +14,6 @@ "test:coverage": "vitest run --coverage", "format": "prettier --write .", "format:check": "prettier --check .", - "e2e": "playwright test --config e2e_tests/playwright.config.ts", - "e2e:ui": "playwright test --ui --config e2e_tests/playwright.config.ts", - "e2e:migration": "playwright test e2e_tests/tests/migration/migratedPages.spec.ts --config e2e_tests/playwright.config.ts", - "e2e:migration:root": "playwright test --config e2e_tests/migration.serverRootPath.config.ts", "knip": "knip", "knip:ci": "knip --exclude exports,nsExports,types,nsTypes,enumMembers,classMembers,duplicates", "knip:fix": "knip --fix", @@ -43,7 +39,7 @@ "jwt-decode": "4.0.0", "lucide-react": "0.513.0", "moment": "2.30.1", - "next": "16.2.6", + "next": "16.2.11", "openai": "4.104.0", "openapi-fetch": "^0.17.0", "openapi-react-query": "^0.5.4", @@ -63,7 +59,6 @@ }, "devDependencies": { "@eslint/js": "9.39.2", - "@playwright/test": "1.58.1", "@tailwindcss/forms": "0.5.11", "@tailwindcss/postcss": "4.3.2", "@testing-library/dom": "10.4.1", @@ -78,7 +73,7 @@ "@vitest/coverage-v8": "3.2.6", "@vitest/ui": "3.2.6", "eslint": "9.39.2", - "eslint-config-next": "16.2.6", + "eslint-config-next": "16.2.11", "eslint-config-prettier": "10.1.8", "eslint-plugin-unused-imports": "4.3.0", "jsdom": "27.4.0", diff --git a/ui/litellm-dashboard/scripts/eslint-rules/filename-pascal-case.mjs b/ui/litellm-dashboard/scripts/eslint-rules/filename-pascal-case.mjs new file mode 100644 index 00000000000..7477925238e --- /dev/null +++ b/ui/litellm-dashboard/scripts/eslint-rules/filename-pascal-case.mjs @@ -0,0 +1,59 @@ +import { basename } from "path"; + +const NEXT_RESERVED = new Set([ + "page", + "layout", + "route", + "template", + "default", + "loading", + "error", + "global-error", + "not-found", + "middleware", + "instrumentation", + "sitemap", + "robots", + "manifest", + "icon", + "apple-icon", + "favicon", + "opengraph-image", + "twitter-image", +]); + +const PASCAL_CASE = /^[A-Z][A-Za-z0-9]*$/; + +const rule = { + meta: { + type: "suggestion", + docs: { + description: "Require PascalCase filenames for .tsx modules; exempt Next.js reserved files and test/spec files.", + }, + schema: [], + messages: { + notPascalCase: "Filename '{{name}}' should be PascalCase (e.g. '{{suggestion}}.tsx').", + }, + }, + create(context) { + const filename = context.filename; + const stem = basename(filename).replace(/\.tsx$/, ""); + const [head, ...rest] = stem.split("."); + if (rest.includes("test") || rest.includes("spec")) return {}; + if (NEXT_RESERVED.has(head)) return {}; + if (PASCAL_CASE.test(head)) return {}; + const pascalHead = head + .split(/[-_]/) + .filter(Boolean) + .map((part) => part.charAt(0).toUpperCase() + part.slice(1)) + .join(""); + const suggestion = [pascalHead, ...rest].join("."); + return { + Program(node) { + context.report({ node, messageId: "notPascalCase", data: { name: `${stem}.tsx`, suggestion } }); + }, + }; + }, +}; + +export default rule; diff --git a/ui/litellm-dashboard/scripts/eslint-rules/index.mjs b/ui/litellm-dashboard/scripts/eslint-rules/index.mjs index 150ba1d02e9..9e9f901a6df 100644 --- a/ui/litellm-dashboard/scripts/eslint-rules/index.mjs +++ b/ui/litellm-dashboard/scripts/eslint-rules/index.mjs @@ -1,10 +1,14 @@ import noLargeInlineObjectArg from "./no-large-inline-object-arg.mjs"; import noLongConditionChain from "./no-long-condition-chain.mjs"; +import noComplexJsxArrow from "./no-complex-jsx-arrow.mjs"; +import filenamePascalCase from "./filename-pascal-case.mjs"; const plugin = { rules: { "no-large-inline-object-arg": noLargeInlineObjectArg, "no-long-condition-chain": noLongConditionChain, + "no-complex-jsx-arrow": noComplexJsxArrow, + "filename-pascal-case": filenamePascalCase, }, }; diff --git a/ui/litellm-dashboard/scripts/eslint-rules/no-complex-jsx-arrow.mjs b/ui/litellm-dashboard/scripts/eslint-rules/no-complex-jsx-arrow.mjs new file mode 100644 index 00000000000..b3dabe03a21 --- /dev/null +++ b/ui/litellm-dashboard/scripts/eslint-rules/no-complex-jsx-arrow.mjs @@ -0,0 +1,41 @@ +const DEFAULT_MAX_STATEMENTS = 2; + +const isJsxAttributeValue = (node) => { + const parent = node.parent; + if (parent == null) return false; + return parent.type === "JSXExpressionContainer" && parent.parent?.type === "JSXAttribute"; +}; + +const rule = { + meta: { + type: "suggestion", + docs: { + description: + "Disallow arrow functions with block bodies over a few statements passed inline as JSX attributes; extract them into a named handler.", + }, + schema: [ + { + type: "object", + properties: { maxStatements: { type: "integer", minimum: 1 } }, + additionalProperties: false, + }, + ], + messages: { + tooComplex: "Inline JSX arrow handler has {{count}} statements; extract it into a named function (max {{max}}).", + }, + }, + create(context) { + const maxStatements = context.options[0]?.maxStatements ?? DEFAULT_MAX_STATEMENTS; + return { + ArrowFunctionExpression(node) { + if (node.body.type !== "BlockStatement") return; + if (!isJsxAttributeValue(node)) return; + const count = node.body.body.length; + if (count <= maxStatements) return; + context.report({ node, messageId: "tooComplex", data: { count, max: maxStatements } }); + }, + }; + }, +}; + +export default rule; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/access-groups/_components/AccessGroupsDetailsPage.tsx b/ui/litellm-dashboard/src/app/(dashboard)/access-groups/_components/AccessGroupsDetailsPage.tsx index 72a89093bdb..9476a8d98af 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/access-groups/_components/AccessGroupsDetailsPage.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/access-groups/_components/AccessGroupsDetailsPage.tsx @@ -1,68 +1,63 @@ import { useAccessGroupDetails } from "@/app/(dashboard)/hooks/accessGroups/useAccessGroupDetails"; -import { - Button, - Card, - Col, - Descriptions, - Empty, - Flex, - Layout, - List, - Row, - Spin, - Tabs, - Tag, - theme, - Typography, -} from "antd"; import { ArrowLeftIcon, BotIcon, EditIcon, KeyIcon, LayersIcon, ServerIcon, UsersIcon } from "lucide-react"; import { useState } from "react"; import DefaultProxyAdminTag from "@/components/common_components/DefaultProxyAdminTag"; +import CopyButton from "@/components/shared/CopyButton"; +import { Badge } from "@/components/ui/badge"; +import { Button } from "@/components/ui/button"; +import { Card, CardAction, CardContent, CardHeader, CardTitle } from "@/components/ui/card"; +import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs"; +import { UiLoadingSpinner } from "@/components/ui/ui-loading-spinner"; import { AccessGroupEditModal } from "./AccessGroupsModal/AccessGroupEditModal"; -const { Title, Text } = Typography; -const { Content } = Layout; - interface AccessGroupDetailProps { accessGroupId: string; onBack: () => void; } +const MAX_PREVIEW = 5; + +function ResourceList({ ids, emptyMessage }: { ids: string[]; emptyMessage: string }) { + if (ids.length === 0) { + return

{emptyMessage}

; + } + return ( +
+ {ids.map((id) => ( + + + {id} + + + ))} +
+ ); +} + export function AccessGroupDetail({ accessGroupId, onBack }: AccessGroupDetailProps) { const { data: accessGroup, isLoading } = useAccessGroupDetails(accessGroupId); - const { token } = theme.useToken(); const [isEditModalVisible, setIsEditModalVisible] = useState(false); const [showAllKeys, setShowAllKeys] = useState(false); const [showAllTeams, setShowAllTeams] = useState(false); - const MAX_PREVIEW = 5; - if (isLoading) { return ( - - - - - +
+
+ +
+
); } if (!accessGroup) { return ( - - +

Access group not found

+ ); } @@ -75,224 +70,159 @@ export function AccessGroupDetail({ accessGroupId, onBack }: AccessGroupDetailPr const displayedKeys = showAllKeys ? keyIds : keyIds.slice(0, MAX_PREVIEW); const displayedTeams = showAllTeams ? teamIds : teamIds.slice(0, MAX_PREVIEW); - const handleEdit = () => { - setIsEditModalVisible(true); - }; - - const tabItems = [ - { - key: "models", - label: ( - - - Models - {modelIds?.length} - - ), - children: - modelIds?.length > 0 ? ( - ( - - - {id} - - - )} - /> - ) : ( - - ), - }, - { - key: "mcp", - label: ( - - - MCP Servers - {mcpServerIds?.length} - - ), - children: - mcpServerIds?.length > 0 ? ( - ( - - - {id} - - - )} - /> - ) : ( - - ), - }, - { - key: "agents", - label: ( - - - Agents - {agentIds?.length} - - ), - children: - agentIds?.length > 0 ? ( - ( - - - {id} - - - )} - /> - ) : ( - - ), - }, - ]; - return ( - - {/* Header */} -
-
-
- - {accessGroup.access_group_name} - - - ID: {accessGroup.access_group_id} - +

{accessGroup.access_group_name}

+
+ ID: {accessGroup.access_group_id} + +
-
- {/* Group Details */} - - - - {accessGroup.description || "—"} - + + + Group Details + + +
+
Description
+
{accessGroup.description || "—"}
+
Created
+
{new Date(accessGroup.created_at).toLocaleString()} {accessGroup.created_by && ( - -  {"by"}  + <> + by - + )} - - +
+
Last Updated
+
{new Date(accessGroup.updated_at).toLocaleString()} {accessGroup.updated_by && ( - -  {"by"}  + <> + by - + )} - - - - - - {/* Attached Keys & Teams */} - - - - - Attached Keys - {keyIds?.length} - - } - extra={ - keyIds?.length > MAX_PREVIEW ? ( - - ) : null - } - > - {keyIds?.length > 0 ? ( - - {displayedKeys.map((id) => ( - - - {id.length > 20 ? `${id.slice(0, 10)}...${id.slice(-6)}` : id} - - - ))} - - ) : ( - - )} - - - - - - Attached Teams - {teamIds?.length} - - } - extra={ - teamIds?.length > MAX_PREVIEW ? ( - - ) : null - } - > - {teamIds?.length > 0 ? ( - - {displayedTeams.map((id) => ( - - - {id} - - - ))} - - ) : ( - - )} - - - - - {/* Resources Tabs */} - - +
+
+ + + +
+ + + + + Attached Keys + {keyIds.length} + + {keyIds.length > MAX_PREVIEW && ( + + + + )} + + + {keyIds.length > 0 ? ( +
+ {displayedKeys.map((id) => ( + + {id.length > 20 ? `${id.slice(0, 10)}...${id.slice(-6)}` : id} + + ))} +
+ ) : ( +

No keys attached

+ )} + + + + + + + + Attached Teams + {teamIds.length} + + {teamIds.length > MAX_PREVIEW && ( + + + + )} + + + {teamIds.length > 0 ? ( +
+ {displayedTeams.map((id) => ( + + {id} + + ))} +
+ ) : ( +

No teams attached

+ )} + + +
+ + + + + + + + Models + {modelIds.length} + + + + MCP Servers + {mcpServerIds.length} + + + + Agents + {agentIds.length} + + + + + + + + + + + + + - {/* Edit Modal */} setIsEditModalVisible(false)} /> - + ); } diff --git a/ui/litellm-dashboard/src/app/(dashboard)/access-groups/_components/AccessGroupsPage.tsx b/ui/litellm-dashboard/src/app/(dashboard)/access-groups/_components/AccessGroupsPage.tsx index 0de6596f57c..f37acb3d85a 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/access-groups/_components/AccessGroupsPage.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/access-groups/_components/AccessGroupsPage.tsx @@ -1,10 +1,11 @@ import { AccessGroupResponse, useAccessGroups } from "@/app/(dashboard)/hooks/accessGroups/useAccessGroups"; import { useDeleteAccessGroup } from "@/app/(dashboard)/hooks/accessGroups/useDeleteAccessGroup"; -import { PlusOutlined } from "@ant-design/icons"; -import { Button, Flex, Input, Layout, Space, theme, Typography } from "antd"; -import { SearchIcon } from "lucide-react"; +import { Plus, SearchIcon, X } from "lucide-react"; import { useMemo, useState } from "react"; import DeleteResourceModal from "@/components/common_components/DeleteResourceModal"; +import { PageHeader } from "@/components/shared/PageHeader"; +import { Button } from "@/components/ui/button"; +import { InputGroup, InputGroupAddon, InputGroupButton, InputGroupInput } from "@/components/ui/input-group"; import { AccessGroupDetail } from "./AccessGroupsDetailsPage"; import { AccessGroupCreateModal } from "./AccessGroupsModal/AccessGroupCreateModal"; import { AccessGroupsTable } from "./AccessGroupsTable"; @@ -12,9 +13,6 @@ import { AccessGroup } from "./types"; import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; import { isProxyAdminRole } from "@/utils/roles"; -const { Title, Text } = Typography; -const { Content } = Layout; - function mapResponseToAccessGroup(r: AccessGroupResponse): AccessGroup { return { id: r.access_group_id, @@ -33,7 +31,6 @@ function mapResponseToAccessGroup(r: AccessGroupResponse): AccessGroup { } export function AccessGroupsPage() { - const { token } = theme.useToken(); const { userRole } = useAuthorized(); // Admin Viewer follows the read-parity rule: see access groups, no writes. const canModify = isProxyAdminRole(userRole ?? ""); @@ -62,31 +59,41 @@ export function AccessGroupsPage() { } return ( - - - - - Access Groups - - Manage resource permissions for your organization - - {canModify && ( - - )} - - - - } - placeholder="Search groups by name, ID, or description..." - style={{ maxWidth: 400 }} - value={searchText} - onChange={(e) => setSearchText(e.target.value)} - allowClear +
+
+ setIsCreateModalVisible(true)}> + + Create Access Group + + ) : undefined + } /> - +
+ +
+ + + + + setSearchText(e.target.value)} + /> + {searchText && ( + + setSearchText("")}> + + + + )} + +
- +
); } diff --git a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/AgentsPanel.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/AgentsPanel.test.tsx index 441d300436a..d873687378b 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/AgentsPanel.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/AgentsPanel.test.tsx @@ -140,8 +140,9 @@ describe("AgentsPanel", () => { await user.click(await screen.findByTestId("agent-actions-agent-9")); await user.click(await screen.findByTestId("agent-action-delete")); - const modal = await screen.findByRole("dialog"); - await user.click(within(modal).getByRole("button", { name: /^delete$/i })); + const confirmPrompt = await screen.findByText(/are you sure you want to delete agent: Doomed Agent\?/i); + const confirmDialog = confirmPrompt.closest('[role="dialog"],[role="alertdialog"]') as HTMLElement; + await user.click(within(confirmDialog).getByRole("button", { name: /^delete$/i })); await waitFor(() => { expect(networking.deleteAgentCall).toHaveBeenCalledWith("test-token", "agent-9"); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/AgentsPanel.tsx b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/AgentsPanel.tsx index a4a71530c84..4459ee0c377 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/AgentsPanel.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/AgentsPanel.tsx @@ -1,6 +1,5 @@ import React, { useState, useEffect } from "react"; -import { Modal, Alert } from "antd"; -import { Plus } from "lucide-react"; +import { Info, Plus } from "lucide-react"; import { getAgentsList, deleteAgentCall } from "@/components/networking"; import AddAgentForm from "./add_agent_form"; import { isAdminRole } from "@/utils/roles"; @@ -9,6 +8,16 @@ import AgentsTable from "./AgentsTable"; import NotificationsManager from "@/components/molecules/notifications_manager"; import { Agent } from "@/components/agents/types"; import { Team } from "@/components/key_team_helpers/key_list"; +import { Alert, AlertDescription, AlertTitle } from "@/components/shared/Alert"; +import { + AlertDialog, + AlertDialogCancel, + AlertDialogContent, + AlertDialogDescription, + AlertDialogFooter, + AlertDialogHeader, + AlertDialogTitle, +} from "@/components/ui/alert-dialog"; import { Button } from "@/components/ui/button"; interface AgentsPanelProps { @@ -130,17 +139,18 @@ const AgentsPanel: React.FC = ({ accessToken, userRole, teams

Agents

-

+

List of A2A-spec agents that are available to be used in your organization. Go to AI Hub, to make agents public.

- + + + Why do agents need keys? + + Keys scope access to an agent and allow it to call MCP tools. Assign a key when creating an agent or from + the Virtual Keys page. + + {isAdmin && (
+ + + )}
); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/AgentsTable.tsx b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/AgentsTable.tsx index 824ae47f3e6..359cb49b910 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/AgentsTable.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/AgentsTable.tsx @@ -1,13 +1,13 @@ "use client"; import { SortingState } from "@tanstack/react-table"; -import { Tooltip, Switch } from "antd"; -import { CheckCircleOutlined } from "@ant-design/icons"; -import { Bot } from "lucide-react"; +import { Bot, CircleCheck } from "lucide-react"; import React, { useMemo, useState } from "react"; import { Agent } from "@/components/agents/types"; import { DataTable } from "@/components/shared/DataTable"; +import { Switch } from "@/components/ui/switch"; +import { Tooltip, TooltipContent, TooltipProvider, TooltipTrigger } from "@/components/ui/tooltip"; import { getAgentsTableColumns } from "./AgentsTableColumns"; @@ -67,18 +67,27 @@ const AgentsTable: React.FC = ({ size="compact" toolbar={() => (
- -
- - Health Check - + + + + Health Check + +
+ } /> -
- + When enabled, only agents with reachable URLs are shown + +
)} /> diff --git a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_card_discovery.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_card_discovery.test.tsx index 4ee6332c54e..7858bdb1cd4 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_card_discovery.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_card_discovery.test.tsx @@ -126,10 +126,7 @@ describe("AgentCardDiscovery", () => { expect(initialSelection.upstream_url).toBe("https://upstream.example.com"); expect(initialSelection.selected_card.skills).toHaveLength(2); - const summarizeLabel = screen.getByText("Summarize").closest("label"); - expect(summarizeLabel).toBeTruthy(); - const summarizeCheckbox = summarizeLabel!.querySelector("input[type='checkbox']") as HTMLInputElement; - await user.click(summarizeCheckbox); + await user.click(screen.getByRole("checkbox", { name: /Summarize/i })); await waitFor(() => { const latest = onApply.mock.calls.at(-1)?.[0]; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_card_discovery.tsx b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_card_discovery.tsx index e979b2dbe3f..017a9928f8b 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_card_discovery.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_card_discovery.tsx @@ -1,18 +1,20 @@ "use client"; import React, { useCallback, useEffect, useMemo, useRef, useState } from "react"; -import { Alert, Button, Checkbox, Collapse, Empty, Input, Space, Spin, Switch, Tag, Tooltip, Typography } from "antd"; -// Empty is used in the skills panel below. -import { - CheckCircleTwoTone, - InfoCircleOutlined, - LinkOutlined, - ReloadOutlined, - SearchOutlined, -} from "@ant-design/icons"; +import { ChevronDown, CircleAlert, CircleCheck, Info, Link as LinkIcon, RotateCw, Search, X } from "lucide-react"; import { useDebouncedCallback } from "@tanstack/react-pacer/debouncer"; import { DiscoveredAgentCard, discoverAgentCardCall } from "@/components/networking"; +import { Alert, AlertAction, AlertDescription, AlertTitle } from "@/components/shared/Alert"; +import { Badge } from "@/components/ui/badge"; +import { Button } from "@/components/ui/button"; +import { Checkbox } from "@/components/ui/checkbox"; +import { Collapsible, CollapsibleContent, CollapsibleTrigger } from "@/components/ui/collapsible"; +import { Input } from "@/components/ui/input"; +import { Switch } from "@/components/ui/switch"; +import { Textarea } from "@/components/ui/textarea"; +import { Tooltip, TooltipContent, TooltipProvider, TooltipTrigger } from "@/components/ui/tooltip"; +import { UiLoadingSpinner } from "@/components/ui/ui-loading-spinner"; import { ALLOWED_CAPABILITY_KEYS, selectionsFromSavedAgentCard, @@ -20,9 +22,6 @@ import { skillId, } from "./agent_discovery_utils"; -const { Text, Paragraph } = Typography; -const { Panel } = Collapse; - const DISCOVERY_DEBOUNCE_WAIT_MS = 400; export interface DiscoveredAgentCardSelection { @@ -243,102 +242,115 @@ const AgentCardDiscovery: React.FC = ({ const skillCount = card?.skills?.length ?? 0; const selectedSkillCount = selectedSkillIds.size; + const renderDiscoverIcon = () => { + if (loading) return ; + if (card) return ; + return ; + }; + const discoverLabel = card ? "Re-discover" : "Discover"; + return ( -
-
- - Discover from agent URL - - - +
+
+ + Discover from agent URL + + + + + + } + /> + + LiteLLM will fetch /.well-known/agent-card.json from this URL and let you pick which skills and + capabilities to expose through the proxy. + + +
{isParentDriven ? ( <> - +

Using the connection details you entered above. We'll fetch: - -

+

+
{discoveryRequest!.display_url || effectiveUrl || ( - Fill in the fields above first + Fill in the fields above first )}
-
) : ( <> - +

Paste the upstream agent's base URL. We'll try /.well-known/agent-card.json,{" "} /.well-known/agent.json, and /agent.json in order. - +

- +
setManualUrl(e.target.value)} - onPressEnter={handleDiscover} - allowClear + onKeyDown={(e) => { + if (e.key === "Enter") handleDiscover(); + }} disabled={loading} /> - - +
)} {error && ( - setError(null)} - /> + + + Discovery failed + {error} + + + + )} {loading && !card && (
- +
)} {card && ( -
-
- - - Upstream card loaded - {card.version && v{card.version}} - {card.provider?.organization && {card.provider.organization}} - +
+
+ + Upstream card loaded + {card.version && v{card.version}} + {card.provider?.organization && {card.provider.organization}}
-
+
- + setEditedName(e.target.value)} placeholder="Agent name" />
- - Description +