diff --git a/.circleci/config.yml b/.circleci/config.yml index ceeb4b5fbf8..90bab08235a 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -38,7 +38,7 @@ commands: parameters: category: type: enum - enum: ["backend", "client", "redis-compat"] + enum: ["backend", "client"] default: "backend" steps: - run: @@ -339,8 +339,15 @@ commands: Add-Content -Path $PROFILE -Value "`$env:Path = `"$cargoBin;`$env:Path`"" } setup_litellm_test_deps: + parameters: + checkout: + type: boolean + default: true steps: - - checkout + - when: + condition: << parameters.checkout >> + steps: + - checkout - setup_google_dns - install_uv - install_rust @@ -376,13 +383,28 @@ jobs: shell: bash.exe command: bash .circleci/scripts/path_filter.sh windows-release - run: - name: Install Python + name: Install Python 3.11.9 from python.org + no_output_timeout: 6m command: | - choco install python --version=3.11.0 -y --no-progress --force - refreshenv - python --version - environment: - CHOCOLATEY_CONFIRM_ALL: "true" + $ProgressPreference = "SilentlyContinue" + $installer = Join-Path $env:TEMP "python-3.11.9-amd64.exe" + Invoke-WebRequest -Uri https://www.python.org/ftp/python/3.11.9/python-3.11.9-amd64.exe -OutFile $installer -TimeoutSec 120 + $expected = "5ee42c4eee1e6b4464bb23722f90b45303f79442df63083f05322f1785f5fdde" + $actual = (Get-FileHash -Path $installer -Algorithm SHA256).Hash.ToLower() + if ($actual -ne $expected) { + throw "Python installer hash mismatch: expected $expected got $actual" + } + $arguments = "/quiet InstallAllUsers=1 PrependPath=1 Include_test=0 Include_doc=0 Include_launcher=0 TargetDir=C:\Python311" + $process = Start-Process -FilePath $installer -ArgumentList $arguments -PassThru + if (-not $process.WaitForExit(300000)) { + $process.Kill() + throw "Python installer did not finish within 5 minutes" + } + if ($process.ExitCode -ne 0) { + throw "Python installer exited with $($process.ExitCode)" + } + Remove-Item $installer + & C:\Python311\python.exe --version - install_windows_toolchain - run: name: Record the Rust build environment for the release cargo cache key @@ -405,7 +427,7 @@ jobs: environment: UV_HTTP_TIMEOUT: "300" command: | - $env:Path = "$HOME\.cargo\bin;$HOME\.local\bin;$env:Path" + $env:Path = "C:\Python311;C:\Python311\Scripts;$HOME\.cargo\bin;$HOME\.local\bin;$env:Path" uv build --wheel --out-dir dist if ($LASTEXITCODE -ne 0) { exit $LASTEXITCODE @@ -512,11 +534,12 @@ jobs: --set image.repository=litellm-ci \ --set image.tag=${IMAGE_TAG} \ --set image.pullPolicy=Never - # Wait for pod to be ready - echo "Waiting 30 seconds for pod to be ready..." - sleep 30 + if ! kubectl rollout status deployment/litellm --timeout=180s; then + kubectl describe pods -l app.kubernetes.io/name=litellm + kubectl logs -l app.kubernetes.io/name=litellm --tail=200 + exit 1 + fi - # Print pod logs before running tests echo "Printing pod logs..." kubectl logs $(kubectl get pods -l app.kubernetes.io/name=litellm -o jsonpath="{.items[0].metadata.name}") @@ -1272,7 +1295,7 @@ jobs: echo "export MCP_TEST_PEER_PYTHON=$PWD/.venv-mcp-peer/bin/python" >> "$BASH_ENV" - start_postgres: image: postgres:16@sha256:e17e86066e5ef83e0952a9347f5c792b7ece00972e2aa787a6986f471b3dd3d5 - server_args: "-c shared_preload_libraries=pg_stat_statements -c pg_stat_statements.track=all -c pg_stat_statements.max=20000" + server_args: "-c max_connections=300 -c shared_preload_libraries=pg_stat_statements -c pg_stat_statements.track=all -c pg_stat_statements.max=20000" - start_redis - run: name: Run owned integration contracts @@ -1325,12 +1348,27 @@ jobs: default: standard shards: type: integer - machine: - image: ubuntu-2204:2024.04.1 - resource_class: medium + docker: + - image: cimg/python:3.12@sha256:9c796c23c84e84a66a964acb508d39dc5433c81a47e07efd56dccbbc2427e07c + auth: + username: ${DOCKERHUB_USERNAME} + password: ${DOCKERHUB_PASSWORD} + resource_class: small working_directory: ~/project + environment: + UV_PYTHON_DOWNLOADS: never steps: - - setup_litellm_test_deps + - checkout + - install_uv + - run: + name: Install the locked shard verifier and coverage dependencies + command: | + uv export --frozen --all-groups --all-extras --no-emit-project --no-emit-workspace --no-header --no-annotate \ + | awk '/^[^ ]/ { keep = $0 ~ /^(annotated-types|coverage|pydantic|pydantic-core|typing-extensions|typing-inspection)==/ } keep' \ + > /tmp/fan-in-requirements.txt + test "$(grep -c '^[^ ]' /tmp/fan-in-requirements.txt)" -eq 6 + uv venv .venv --python 3.12 + uv pip install --no-config --python .venv --require-hashes -r /tmp/fan-in-requirements.txt - attach_workspace: at: /tmp/integration-shards - run: @@ -1446,7 +1484,8 @@ jobs: steps: - checkout - skip_if_unrelated_changes - - setup_litellm_test_deps + - setup_litellm_test_deps: + checkout: false - start_postgres: db_name: litellm_test image: postgres:16@sha256:e17e86066e5ef83e0952a9347f5c792b7ece00972e2aa787a6986f471b3dd3d5 @@ -1491,108 +1530,6 @@ jobs: - store_artifacts: path: test-results - mcp_integration: - machine: - image: ubuntu-2204:2024.04.1 - resource_class: large - working_directory: ~/project - environment: - COVERAGE_CORE: sysmon - LITELLM_LOCAL_MODEL_COST_MAP: "True" - steps: - - checkout - - skip_if_unrelated_changes - - setup_litellm_test_deps - - run: - name: Generate Prisma client - command: uv run --no-sync prisma generate --schema litellm/proxy/schema.prisma - - run: - name: Install MCP SDK1 peer - command: | - uv venv --python 3.12 .venv-mcp-peer - uv pip install --python .venv-mcp-peer 'mcp==1.28.1' 'langchain-mcp-adapters==0.2.1' - echo "export MCP_TEST_PEER_PYTHON=$PWD/.venv-mcp-peer/bin/python" >> "$BASH_ENV" - - run: - name: Run MCP integration tests - command: | - mkdir -p test-results - env -u OPENAI_API_KEY -u ANTHROPIC_API_KEY \ - timeout --signal=TERM 20m uv run --no-sync pytest \ - tests/unit/experimental_mcp_client tests/unit/proxy/_experimental/mcp_server \ - tests/unit/responses/mcp --tb=short -vv --maxfail=10 -n 2 --dist=loadscope --reruns 0 \ - --reruns-delay 1 --timeout=120 --rerun-except "from pytest-timeout" --durations=20 \ - --cov=./litellm --cov=./enterprise/litellm_enterprise --cov-report=xml:coverage.xml \ - --cov-config=pyproject.toml -o junit_family=xunit1 --junitxml=test-results/junit.xml - - install_codecov_cli - - run: - name: Upload MCP integration coverage - when: always - command: | - [ -f coverage.xml ] || { echo "no coverage.xml produced"; exit 1; } - codecov upload-process --disable-search --fail-on-error -f coverage.xml \ - -F mcp-integration -C "$CIRCLE_SHA1" \ - -n "mcp-integration-${CIRCLE_BUILD_NUM}" --git-service github - - store_test_results: - path: test-results - - store_artifacts: - path: test-results - - redis_compat: - parameters: - redis_py: - type: string - machine: - image: ubuntu-2204:2024.04.1 - resource_class: large - working_directory: ~/project - steps: - - checkout - - skip_if_unrelated_changes: - category: redis-compat - - setup_litellm_test_deps - - run: - name: Pin redis-py version - command: | - uv pip install "redis==<< parameters.redis_py >>" - uv run --no-sync python -c "import redis; assert redis.__version__ == '<< parameters.redis_py >>', redis.__version__; print('redis-py', redis.__version__)" - - run: - name: Install redis-server 7.2.16 - command: | - mkdir -p "$HOME/.local/bin" - cid="$(docker create redis:7.2.16@sha256:0637954999d01b7c9ce9167db2da50656e2590d3b884f1c600c5f63bb6e6773c)" - docker cp "${cid}":/usr/local/bin/redis-server "$HOME/.local/bin/redis-server" - docker rm "$cid" - redis-server --version | grep -q 'v=7.2.16' - - run: - name: Run Redis compatibility tests - command: | - mkdir -p test-results - env -u CASSETTE_REDIS_URL -u AZURE_CLIENT_ID -u AZURE_CLIENT_SECRET -u AZURE_TENANT_ID \ - timeout --signal=TERM 15m uv run --no-sync pytest \ - tests/unit/test_redis.py tests/unit/caching/test_redis_connection_pool.py \ - tests/unit/caching/test_redis_cluster_cache.py tests/unit/caching/test_evicted_client_closer.py \ - tests/integration/sdk/test_redis_cluster_iam_auth.py::test_sync_cluster_authenticates_with_azure_credentials \ - tests/integration/sdk/test_redis_cluster_iam_auth.py::test_sync_cluster_authenticates_with_gcp_credentials \ - --tb=short -vv --reruns 2 --reruns-delay 1 --durations=20 --cov=./litellm \ - --cov-report=xml:coverage.xml -o junit_family=xunit1 --junitxml=test-results/junit.xml - - when: - condition: - equal: ["5.3.1", << parameters.redis_py >>] - steps: - - install_codecov_cli - - run: - name: Upload Redis compatibility coverage - when: always - command: | - [ -f coverage.xml ] || { echo "no coverage.xml produced"; exit 1; } - codecov upload-process --disable-search --fail-on-error -f coverage.xml \ - -F redis-compat -C "$CIRCLE_SHA1" \ - -n "redis-compat-${CIRCLE_BUILD_NUM}" --git-service github - - store_test_results: - path: test-results - - store_artifacts: - path: test-results - workflows: migration_startup_scheduled: triggers: @@ -1692,7 +1629,6 @@ workflows: - integration_contracts: name: integration-mcp suite: mcp - resource_class: xlarge parallelism: 2 filters: *integration_default_run - integration_shard_coverage: @@ -1749,19 +1685,6 @@ workflows: coverage_flag: roi-postgres timeout_minutes: 10 filters: *integration_default_run - - mcp_integration: - name: mcp-integration - filters: *integration_default_run - - redis_compat: - name: redis-compat-<< matrix.redis_py >> - matrix: - parameters: - redis_py: - - "5.3.1" - - "6.4.0" - - "7.4.1" - - "8.0.1" - filters: *integration_default_run - windows_release_wheel: filters: *integration_default_run - base_sdk_install: diff --git a/.circleci/scripts/classify_changes.sh b/.circleci/scripts/classify_changes.sh index cca7f19a02a..d2ab3478397 100755 --- a/.circleci/scripts/classify_changes.sh +++ b/.circleci/scripts/classify_changes.sh @@ -14,7 +14,7 @@ outside_cost_map_set=false while IFS= read -r file || [ -n "$file" ]; do [ -n "$file" ] || continue case "$file" in - litellm/_redis.py | litellm/_redis_credential_provider.py | litellm/caching/redis_cache.py | litellm/caching/evicted_client_closer.py | tests/unit/test_redis.py | tests/integration/sdk/test_redis_cluster_iam_auth.py | tests/unit/caching/test_redis_connection_pool.py | tests/unit/caching/test_redis_cluster_cache.py | tests/unit/caching/test_evicted_client_closer.py | .circleci/config.yml | .circleci/scripts/classify_changes.sh | .circleci/scripts/path_filter.sh | pyproject.toml | uv.lock) + litellm/_redis.py | litellm/_redis_credential_provider.py | litellm/caching/redis_cache.py | litellm/caching/evicted_client_closer.py | tests/unit/test_redis.py | tests/integration/sdk/test_redis_cluster_iam_auth.py | tests/unit/caching/test_redis_connection_pool.py | tests/unit/caching/test_redis_cluster_cache.py | tests/unit/caching/test_evicted_client_closer.py | .github/workflows/test-unit.yml | .github/actions/detect-changes/* | .github/scripts/detect_changes.sh | .circleci/scripts/classify_changes.sh | tests/integration/run.py | pyproject.toml | uv.lock) has_redis_compat=true ;; esac case "$file" in diff --git a/.circleci/scripts/path_filter.sh b/.circleci/scripts/path_filter.sh index 111fa2b57ca..c3cf6e566ff 100755 --- a/.circleci/scripts/path_filter.sh +++ b/.circleci/scripts/path_filter.sh @@ -1,7 +1,7 @@ #!/usr/bin/env bash set -uo pipefail -category="${1:?usage: path_filter.sh }" +category="${1:?usage: path_filter.sh }" here="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" run_full() { diff --git a/.circleci/scripts/run_integration.sh b/.circleci/scripts/run_integration.sh index 0210c2dc6f7..be7e70d491c 100644 --- a/.circleci/scripts/run_integration.sh +++ b/.circleci/scripts/run_integration.sh @@ -128,20 +128,14 @@ fi sudo iptables -L integration_only -n -v -x > "$results/egress-guard.txt" awk '$3 == "REJECT" && $1 > 0 { rejected=1 } END { exit !rejected }' "$results/egress-guard.txt" -isolated_workers=false if [ "$mode" = parity ]; then export INTEGRATION_WORKERS=1 else case "$suite" in - management|accounting) export INTEGRATION_WORKERS=2; isolated_workers=true ;; - providers|extensions) export INTEGRATION_WORKERS=4; isolated_workers=true ;; + management|accounting) export INTEGRATION_WORKERS=2 ;; + providers|extensions) export INTEGRATION_WORKERS=4 ;; esac fi -if [ "$isolated_workers" = false ]; then - setsid env -i PATH="$PATH" HOME="$HOME" PYTHONPATH="$PYTHONPATH" INTEGRATION_RUN_ID="$integration_identity" \ - .venv/bin/python -m integration._support.upstream > "$results/upstream.log" 2>&1 & - upstream_pid=$! -fi if [ "$suite" = cost ]; then export INTEGRATION_WORKERS=8 fi @@ -149,7 +143,16 @@ if [ "$suite" = security ]; then export INTEGRATION_WORKERS=4 fi if [ "$suite" = mcp ]; then - export INTEGRATION_WORKERS=4 INTEGRATION_COVERAGE=1 + export INTEGRATION_WORKERS=2 INTEGRATION_COVERAGE=1 +fi +isolated_workers=false +if [ "${INTEGRATION_WORKERS:-1}" -gt 1 ]; then + isolated_workers=true +fi +if [ "$isolated_workers" = false ]; then + setsid env -i PATH="$PATH" HOME="$HOME" PYTHONPATH="$PYTHONPATH" INTEGRATION_RUN_ID="$integration_identity" \ + .venv/bin/python -m integration._support.upstream > "$results/upstream.log" 2>&1 & + upstream_pid=$! fi coverage_data="$PWD/$results/coverage/data" proxy_command=(.venv/bin/python -m integration._support.proxy) @@ -191,9 +194,6 @@ start_proxy() { > "$results/$log_name" 2>&1 & launched_pid=$! } -if { [ "$suite" = management ] || [ "$suite" = mcp ]; } && [ "$isolated_workers" = true ]; then - export INTEGRATION_PEER_URL=http://127.0.0.1:4001 -fi if [ "$isolated_workers" = false ]; then start_proxy 4000 proxy.log proxy_pid="$launched_pid" @@ -239,13 +239,24 @@ if [ "${CIRCLE_NODE_TOTAL:-1}" -gt 1 ]; then shard_arguments=(--shard-files "$results/node-files.txt") fi +shared_services=() +worker_coverage=() +if [ "$isolated_workers" = false ]; then + shared_services=( + "REDIS_HOST=$REDIS_HOST" "REDIS_PORT=$REDIS_PORT" + "INTEGRATION_PROXY_URL=$INTEGRATION_PROXY_URL" "INTEGRATION_PEER_URL=$INTEGRATION_PEER_URL" + "INTEGRATION_UPSTREAM_URL=$INTEGRATION_UPSTREAM_URL" "INTEGRATION_MASTER_KEY=$INTEGRATION_MASTER_KEY" + ) +elif [ "${INTEGRATION_COVERAGE:-0}" = 1 ]; then + worker_coverage=("INTEGRATION_COVERAGE_FILE=$coverage_data") +fi + env -i PATH="$PATH" HOME="$HOME" PYTHONPATH="$PYTHONPATH" \ INTEGRATION_RUN_ID="$integration_identity" \ - DATABASE_URL="$DATABASE_URL" REDIS_HOST="$REDIS_HOST" REDIS_PORT="$REDIS_PORT" \ - INTEGRATION_PROXY_URL="$INTEGRATION_PROXY_URL" INTEGRATION_PEER_URL="$INTEGRATION_PEER_URL" \ - INTEGRATION_UPSTREAM_URL="$INTEGRATION_UPSTREAM_URL" \ + DATABASE_URL="$DATABASE_URL" \ + "${shared_services[@]}" "${worker_coverage[@]}" \ INTEGRATION_WORKERS="${INTEGRATION_WORKERS:-1}" \ - INTEGRATION_MASTER_KEY="$INTEGRATION_MASTER_KEY" LITELLM_MODE=PRODUCTION \ + LITELLM_MODE=PRODUCTION \ LITELLM_LICENSE="${LITELLM_LICENSE:-}" \ INTEGRATION_SEED="$INTEGRATION_SEED" \ INTEGRATION_ORDER_SEED="$INTEGRATION_ORDER_SEED" \ diff --git a/.github/PULL_REQUEST_TEMPLATE/general.md b/.github/PULL_REQUEST_TEMPLATE/general.md index d9e72d58632..644d1c4a450 100644 --- a/.github/PULL_REQUEST_TEMPLATE/general.md +++ b/.github/PULL_REQUEST_TEMPLATE/general.md @@ -68,7 +68,7 @@ After: the same request comes back with real token counts, so the dashboard show - [ ] The handful of test files covering my change pass locally, e.g. `uv run pytest tests/unit/.py -v`. Leave the suites (`make test-unit-*`, `make test-unit`) to CI: it finishes in ~15 minutes where a laptop takes an hour or more - [ ] My PR passes all required CI/CD checks (e.g., lint, schema.d.ts sync check, etc.) - [ ] My PR's scope is as isolated as possible; it only solves 1 specific problem -- [ ] I have received a Greptile **Confidence Score of at least 4/5** before requesting a maintainer review (Greptile reviews automatically once the PR is opened; only comment `@greptileai` to re-request a review after pushing changes) +- [ ] Every automated review on my PR's latest commit has finished, and I have addressed or replied to each of its findings before requesting a maintainer review ## Delays in PR merge? diff --git a/.github/workflows/test-e2e-redis-chaos.yml b/.github/workflows/test-e2e-redis-chaos.yml index c7412a63334..8f351a7a647 100644 --- a/.github/workflows/test-e2e-redis-chaos.yml +++ b/.github/workflows/test-e2e-redis-chaos.yml @@ -91,12 +91,12 @@ jobs: - name: Run the Redis chaos load test env: - E2E_REDIS_CHAOS: "1" + E2E_PERF_TESTS: "1" LITELLM_PROXY_URL: http://localhost:4000 REDIS_HOST: 127.0.0.1 REDIS_PORT: "6379" run: | - uv run --no-sync pytest tests/e2e/load/test_redis_chaos_e2e.py -v --tb=short -rA -s + uv run --no-sync pytest tests/e2e/perf/test_redis_chaos_e2e.py -v --tb=short -rA -s - name: Show proxy log on failure if: failure() diff --git a/.github/workflows/test-linting.yml b/.github/workflows/test-linting.yml index 5384e873f22..15cda063626 100644 --- a/.github/workflows/test-linting.yml +++ b/.github/workflows/test-linting.yml @@ -615,12 +615,6 @@ jobs: - name: check_migrations_no_data_rewrites run: uv run --no-sync python ./tests/code_coverage_tests/check_migrations_no_data_rewrites.py - - name: check_greptile_files_json - run: uv run --no-sync python ./tests/code_coverage_tests/check_greptile_files_json.py - - - name: test_check_greptile_files_json - run: uv run --no-sync pytest -q --noconftest -p no:cacheprovider -c /dev/null tests/code_coverage_tests/test_check_greptile_files_json.py - - name: check_no_publicly_known_master_key run: uv run --no-sync python ./tests/code_coverage_tests/check_no_publicly_known_master_key.py diff --git a/.github/workflows/test-unit.yml b/.github/workflows/test-unit.yml index 23a340cd436..e9236ee7787 100644 --- a/.github/workflows/test-unit.yml +++ b/.github/workflows/test-unit.yml @@ -794,6 +794,10 @@ jobs: tests/unit/proxy/_experimental/mcp_server/test_discoverable_endpoints.py tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py tests/unit/proxy/_experimental/mcp_server/test_db_credentials.py + tests/unit/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py + tests/unit/proxy/_experimental/mcp_server/test_mcp_oauth_passthrough_tools.py + tests/unit/proxy/_experimental/mcp_server/test_mcp_max_concurrent_requests.py + tests/unit/proxy/_experimental/mcp_server/faults/test_list_outcomes.py tests/unit/proxy/_experimental/mcp_server/outbound_credentials python-version: "3.12" workers: 4 @@ -802,6 +806,28 @@ jobs: job-timeout-minutes: 60 dist: loadscope + - shard: mcp-server + test-path: |- + tests/unit/experimental_mcp_client + tests/unit/proxy/_experimental/mcp_server + tests/unit/responses/mcp + --ignore=tests/unit/experimental_mcp_client/test_mcp_client.py + --ignore=tests/unit/proxy/_experimental/mcp_server/test_capabilities.py + --ignore=tests/unit/proxy/_experimental/mcp_server/test_interactions.py + --ignore=tests/unit/proxy/_experimental/mcp_server/test_mcp_elicitation_handler.py + --ignore=tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py + --ignore=tests/unit/proxy/_experimental/mcp_server/test_mcp_server_tool_calls_and_headers.py + --ignore=tests/unit/proxy/_experimental/mcp_server/test_operations.py + --ignore=tests/unit/proxy/_experimental/mcp_server/test_discoverable_endpoints.py + --ignore=tests/unit/proxy/_experimental/mcp_server/test_db_credentials.py + --ignore=tests/unit/proxy/_experimental/mcp_server/outbound_credentials + python-version: "3.12" + workers: 4 + reruns: 0 + timeout-minutes: 20 + job-timeout-minutes: 60 + dist: loadscope + - shard: custom-logging test-path: |- tests/unit/proxy/test_proxy_server_custom_logger.py @@ -1324,6 +1350,115 @@ jobs: - name: Check consumer types with the installed SDK if: steps.changes.outputs.decision != 'skip' run: uv run --no-sync basedpyright --pythonpath .venv/bin/python --project tests/typing/pyrightconfig.json + redis-compat: + name: redis-py ${{ matrix.redis-py }} compatibility + needs: rust-bridge + if: ${{ !cancelled() }} + permissions: + contents: read + id-token: write + pull-requests: read + runs-on: ubuntu-latest + timeout-minutes: 25 + strategy: + fail-fast: false + matrix: + redis-py: ["5.3.1", "6.4.0", "7.4.1", "8.0.1"] + env: + UV_PYTHON: "3.12" + LITELLM_LOCAL_MODEL_COST_MAP: "True" + REDIS_PY_VERSION: ${{ matrix.redis-py }} + steps: + - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 + timeout-minutes: 3 + with: + persist-credentials: false + + - name: Detect relevant changes + id: changes + timeout-minutes: 2 + uses: ./.github/actions/detect-changes + with: + category: redis-compat + + - name: Set up Python + if: steps.changes.outputs.decision != 'skip' + timeout-minutes: 3 + uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0 + with: + python-version: ${{ env.UV_PYTHON }} + + - name: Set up uv + if: steps.changes.outputs.decision != 'skip' + timeout-minutes: 3 + uses: ./.github/actions/setup-uv-with-retries + with: + version: "0.10.9" + + - name: Cache uv dependencies + if: steps.changes.outputs.decision != 'skip' + timeout-minutes: 5 + uses: ./.github/actions/cache-uv-downloads + + - name: Set up the Rust build + if: steps.changes.outputs.decision != 'skip' + timeout-minutes: 5 + uses: ./.github/actions/rust-bridge + with: + artifact: ${{ needs.rust-bridge.outputs.artifact }} + + - name: Install dependencies with redis-py ${{ matrix.redis-py }} + if: steps.changes.outputs.decision != 'skip' + timeout-minutes: 8 + env: + RUST_BRIDGE_ARTIFACT: ${{ needs.rust-bridge.outputs.artifact }} + run: | + if [ -z "$RUST_BRIDGE_ARTIFACT" ]; then + .github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router --extra saml --extra caching --extra extra_proxy --extra proxy-runtime --extra utils + else + .github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router --extra saml --extra caching --extra extra_proxy --extra proxy-runtime --extra utils --no-install-project + uv pip install --no-deps --python .venv/bin/python rust-bridge-dist/*.whl + cp rust-bridge-dist/litellm/rust_bridge/_native.abi3.so litellm/rust_bridge/_native.abi3.so + fi + uv pip install --python .venv/bin/python "redis==${REDIS_PY_VERSION:?}" + uv run --no-sync python -c 'import os, redis; assert redis.__version__ == os.environ["REDIS_PY_VERSION"], redis.__version__' + + - name: Install redis-server 7.2.16 + if: steps.changes.outputs.decision != 'skip' + timeout-minutes: 5 + run: | + mkdir -p "$RUNNER_TEMP/redis-bin" + cid="$(docker create redis:7.2.16@sha256:0637954999d01b7c9ce9167db2da50656e2590d3b884f1c600c5f63bb6e6773c)" + docker cp "${cid}:/usr/local/bin/redis-server" "$RUNNER_TEMP/redis-bin/redis-server" + docker rm "$cid" + "$RUNNER_TEMP/redis-bin/redis-server" --version | grep -q 'v=7.2.16' + echo "$RUNNER_TEMP/redis-bin" >> "$GITHUB_PATH" + + - name: Run Redis compatibility tests + if: steps.changes.outputs.decision != 'skip' + timeout-minutes: 15 + env: + COVERAGE_CORE: sysmon + run: | + uv run --no-sync pytest \ + tests/unit/test_redis.py \ + tests/unit/caching/test_redis_connection_pool.py \ + tests/unit/caching/test_redis_cluster_cache.py \ + tests/unit/caching/test_evicted_client_closer.py \ + tests/integration/sdk/test_redis_cluster_iam_auth.py \ + --tb=short -vv -rs -p no:rerunfailures -p no:pytest-retry --timeout=120 --durations=20 \ + --cov=./litellm --cov-report=xml:coverage.xml + + - name: Upload redis-py ${{ matrix.redis-py }} coverage to Codecov + if: always() && steps.changes.outputs.decision != 'skip' && matrix.redis-py == '5.3.1' + continue-on-error: true + uses: codecov/codecov-action@0fb7174895f61a3b6b78fc075e0cd60383518dac # v5.5.5 + with: + use_oidc: true + files: coverage.xml + root_dir: ${{ github.workspace }} + flags: redis-compat + fail_ci_if_error: false coverage: name: coverage needs: unit @@ -1369,10 +1504,20 @@ jobs: root_dir: ${{ github.workspace }} flags: unit fail_ci_if_error: false + + - name: Upload MCP server coverage under its own flag + continue-on-error: true + uses: codecov/codecov-action@0fb7174895f61a3b6b78fc075e0cd60383518dac # v5.5.5 + with: + use_oidc: true + files: coverage-reports/coverage-mcp-server-${{ github.run_id }}-${{ github.run_attempt }}/coverage.xml + root_dir: ${{ github.workspace }} + flags: mcp-integration + fail_ci_if_error: false unit-passed: name: unit passed permissions: {} - needs: [rust-bridge, assert-shard-coverage, unit, ui-unit, docs, helm, openai-sdk-compat] + needs: [rust-bridge, assert-shard-coverage, unit, ui-unit, docs, helm, openai-sdk-compat, redis-compat] if: always() runs-on: ubuntu-latest timeout-minutes: 2 diff --git a/.github/workflows/weekly_load_anomaly.yml b/.github/workflows/weekly_load_anomaly.yml index 3e1fca89645..0ba4df89cc4 100644 --- a/.github/workflows/weekly_load_anomaly.yml +++ b/.github/workflows/weekly_load_anomaly.yml @@ -63,7 +63,7 @@ jobs: - 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 & + nohup uv run --no-sync litellm --config tests/e2e/perf/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 @@ -78,7 +78,7 @@ jobs: env: E2E_WEEKLY_ANOMALY: "1" run: | - uv run --no-sync pytest tests/e2e/load/test_weekly_session_anomaly_e2e.py -v --tb=short -rA + uv run --no-sync pytest tests/e2e/perf/test_weekly_session_anomaly_e2e.py -v --tb=short -rA - name: Show proxy log on failure if: failure() diff --git a/.greptile/files.json b/.greptile/files.json deleted file mode 100644 index 148881ac289..00000000000 --- a/.greptile/files.json +++ /dev/null @@ -1,655 +0,0 @@ -{ - "files": [ - { - "path": "litellm/_v2/AGENTS.md", - "description": "Conventions for code under litellm/_v2/", - "scope": [ - "litellm/_v2/**" - ] - }, - { - "path": "litellm/_v2/cache/AGENTS.md", - "description": "Conventions for code under litellm/_v2/cache/", - "scope": [ - "litellm/_v2/cache/**" - ] - }, - { - "path": "litellm/proxy/_experimental/mcp_server/AGENTS.md", - "description": "Conventions for code under litellm/proxy/_experimental/mcp_server/", - "scope": [ - "litellm/proxy/_experimental/mcp_server/**" - ] - }, - { - "path": "litellm/rust_bridge/AGENTS.md", - "description": "Conventions for code under litellm/rust_bridge/", - "scope": [ - "litellm/rust_bridge/**" - ] - }, - { - "path": "litellm/rust_bridge/host/AGENTS.md", - "description": "Conventions for code under litellm/rust_bridge/host/", - "scope": [ - "litellm/rust_bridge/host/**" - ] - }, - { - "path": "litellm/tracing/AGENTS.md", - "description": "Conventions for code under litellm/tracing/", - "scope": [ - "litellm/tracing/**" - ] - }, - { - "path": "litellm-rust/AGENTS.md", - "description": "Conventions for code under litellm-rust/", - "scope": [ - "litellm-rust/**" - ] - }, - { - "path": "litellm-rust/crates/cache-response/AGENTS.md", - "description": "Conventions for code under litellm-rust/crates/cache-response/", - "scope": [ - "litellm-rust/crates/cache-response/**" - ] - }, - { - "path": "litellm-rust/crates/callbacks-legacy-python/AGENTS.md", - "description": "Conventions for code under litellm-rust/crates/callbacks-legacy-python/", - "scope": [ - "litellm-rust/crates/callbacks-legacy-python/**" - ] - }, - { - "path": "litellm-rust/crates/coroutine/AGENTS.md", - "description": "Conventions for code under litellm-rust/crates/coroutine/", - "scope": [ - "litellm-rust/crates/coroutine/**" - ] - }, - { - "path": "litellm-rust/crates/gateway/AGENTS.md", - "description": "Conventions for code under litellm-rust/crates/gateway/", - "scope": [ - "litellm-rust/crates/gateway/**" - ] - }, - { - "path": "litellm-rust/crates/gateway-auth/AGENTS.md", - "description": "Conventions for code under litellm-rust/crates/gateway-auth/", - "scope": [ - "litellm-rust/crates/gateway-auth/**" - ] - }, - { - "path": "litellm-rust/crates/gateway-inference/AGENTS.md", - "description": "Conventions for code under litellm-rust/crates/gateway-inference/", - "scope": [ - "litellm-rust/crates/gateway-inference/**" - ] - }, - { - "path": "litellm-rust/crates/gateway-mcp/AGENTS.md", - "description": "Conventions for code under litellm-rust/crates/gateway-mcp/", - "scope": [ - "litellm-rust/crates/gateway-mcp/**" - ] - }, - { - "path": "litellm-rust/crates/host/AGENTS.md", - "description": "Conventions for code under litellm-rust/crates/host/", - "scope": [ - "litellm-rust/crates/host/**" - ] - }, - { - "path": "litellm-rust/crates/host/src/machine/AGENTS.md", - "description": "Conventions for code under litellm-rust/crates/host/src/machine/", - "scope": [ - "litellm-rust/crates/host/src/machine/**" - ] - }, - { - "path": "litellm-rust/crates/host-http/AGENTS.md", - "description": "Conventions for code under litellm-rust/crates/host-http/", - "scope": [ - "litellm-rust/crates/host-http/**" - ] - }, - { - "path": "litellm-rust/crates/host-native/AGENTS.md", - "description": "Conventions for code under litellm-rust/crates/host-native/", - "scope": [ - "litellm-rust/crates/host-native/**" - ] - }, - { - "path": "litellm-rust/crates/host-python/AGENTS.md", - "description": "Conventions for code under litellm-rust/crates/host-python/", - "scope": [ - "litellm-rust/crates/host-python/**" - ] - }, - { - "path": "litellm-rust/crates/http/AGENTS.md", - "description": "Conventions for code under litellm-rust/crates/http/", - "scope": [ - "litellm-rust/crates/http/**" - ] - }, - { - "path": "litellm-rust/crates/inference/AGENTS.md", - "description": "Conventions for code under litellm-rust/crates/inference/", - "scope": [ - "litellm-rust/crates/inference/**" - ] - }, - { - "path": "litellm-rust/crates/inference-chat/AGENTS.md", - "description": "Conventions for code under litellm-rust/crates/inference-chat/", - "scope": [ - "litellm-rust/crates/inference-chat/**" - ] - }, - { - "path": "litellm-rust/crates/inference-messages/AGENTS.md", - "description": "Conventions for code under litellm-rust/crates/inference-messages/", - "scope": [ - "litellm-rust/crates/inference-messages/**" - ] - }, - { - "path": "litellm-rust/crates/inference-ocr/AGENTS.md", - "description": "Conventions for code under litellm-rust/crates/inference-ocr/", - "scope": [ - "litellm-rust/crates/inference-ocr/**" - ] - }, - { - "path": "litellm-rust/crates/inference-responses/AGENTS.md", - "description": "Conventions for code under litellm-rust/crates/inference-responses/", - "scope": [ - "litellm-rust/crates/inference-responses/**" - ] - }, - { - "path": "litellm-rust/crates/inference-transcription/AGENTS.md", - "description": "Conventions for code under litellm-rust/crates/inference-transcription/", - "scope": [ - "litellm-rust/crates/inference-transcription/**" - ] - }, - { - "path": "litellm-rust/crates/llms/AGENTS.md", - "description": "Conventions for code under litellm-rust/crates/llms/", - "scope": [ - "litellm-rust/crates/llms/**" - ] - }, - { - "path": "litellm-rust/crates/llms/src/anthropic/AGENTS.md", - "description": "Conventions for code under litellm-rust/crates/llms/src/anthropic/", - "scope": [ - "litellm-rust/crates/llms/src/anthropic/**" - ] - }, - { - "path": "litellm-rust/crates/llms/src/anthropic/batches/AGENTS.md", - "description": "Conventions for code under litellm-rust/crates/llms/src/anthropic/batches/", - "scope": [ - "litellm-rust/crates/llms/src/anthropic/batches/**" - ] - }, - { - "path": "litellm-rust/crates/llms/src/anthropic/chat/AGENTS.md", - "description": "Conventions for code under litellm-rust/crates/llms/src/anthropic/chat/", - "scope": [ - "litellm-rust/crates/llms/src/anthropic/chat/**" - ] - }, - { - "path": "litellm-rust/crates/llms/src/anthropic/count_tokens/AGENTS.md", - "description": "Conventions for code under litellm-rust/crates/llms/src/anthropic/count_tokens/", - "scope": [ - "litellm-rust/crates/llms/src/anthropic/count_tokens/**" - ] - }, - { - "path": "litellm-rust/crates/llms/src/anthropic/messages/AGENTS.md", - "description": "Conventions for code under litellm-rust/crates/llms/src/anthropic/messages/", - "scope": [ - "litellm-rust/crates/llms/src/anthropic/messages/**" - ] - }, - { - "path": "litellm-rust/crates/llms/src/aws_textract/ocr/AGENTS.md", - "description": "Conventions for code under litellm-rust/crates/llms/src/aws_textract/ocr/", - "scope": [ - "litellm-rust/crates/llms/src/aws_textract/ocr/**" - ] - }, - { - "path": "litellm-rust/crates/llms/src/azure_ai/messages/AGENTS.md", - "description": "Conventions for code under litellm-rust/crates/llms/src/azure_ai/messages/", - "scope": [ - "litellm-rust/crates/llms/src/azure_ai/messages/**" - ] - }, - { - "path": "litellm-rust/crates/llms/src/azure_ai/ocr/AGENTS.md", - "description": "Conventions for code under litellm-rust/crates/llms/src/azure_ai/ocr/", - "scope": [ - "litellm-rust/crates/llms/src/azure_ai/ocr/**" - ] - }, - { - "path": "litellm-rust/crates/llms/src/base_llm/messages/AGENTS.md", - "description": "Conventions for code under litellm-rust/crates/llms/src/base_llm/messages/", - "scope": [ - "litellm-rust/crates/llms/src/base_llm/messages/**" - ] - }, - { - "path": "litellm-rust/crates/llms/src/bedrock/AGENTS.md", - "description": "Conventions for code under litellm-rust/crates/llms/src/bedrock/", - "scope": [ - "litellm-rust/crates/llms/src/bedrock/**" - ] - }, - { - "path": "litellm-rust/crates/llms/src/bedrock/audio_transcription/AGENTS.md", - "description": "Conventions for code under litellm-rust/crates/llms/src/bedrock/audio_transcription/", - "scope": [ - "litellm-rust/crates/llms/src/bedrock/audio_transcription/**" - ] - }, - { - "path": "litellm-rust/crates/llms/src/bedrock/chat/AGENTS.md", - "description": "Conventions for code under litellm-rust/crates/llms/src/bedrock/chat/", - "scope": [ - "litellm-rust/crates/llms/src/bedrock/chat/**" - ] - }, - { - "path": "litellm-rust/crates/llms/src/bedrock/claude_platform/messages/AGENTS.md", - "description": "Conventions for code under litellm-rust/crates/llms/src/bedrock/claude_platform/messages/", - "scope": [ - "litellm-rust/crates/llms/src/bedrock/claude_platform/messages/**" - ] - }, - { - "path": "litellm-rust/crates/llms/src/bedrock/messages/AGENTS.md", - "description": "Conventions for code under litellm-rust/crates/llms/src/bedrock/messages/", - "scope": [ - "litellm-rust/crates/llms/src/bedrock/messages/**" - ] - }, - { - "path": "litellm-rust/crates/llms/src/bedrock_mantle/messages/AGENTS.md", - "description": "Conventions for code under litellm-rust/crates/llms/src/bedrock_mantle/messages/", - "scope": [ - "litellm-rust/crates/llms/src/bedrock_mantle/messages/**" - ] - }, - { - "path": "litellm-rust/crates/llms/src/cohere/ocr/AGENTS.md", - "description": "Conventions for code under litellm-rust/crates/llms/src/cohere/ocr/", - "scope": [ - "litellm-rust/crates/llms/src/cohere/ocr/**" - ] - }, - { - "path": "litellm-rust/crates/llms/src/deepseek/messages/AGENTS.md", - "description": "Conventions for code under litellm-rust/crates/llms/src/deepseek/messages/", - "scope": [ - "litellm-rust/crates/llms/src/deepseek/messages/**" - ] - }, - { - "path": "litellm-rust/crates/llms/src/minimax/messages/AGENTS.md", - "description": "Conventions for code under litellm-rust/crates/llms/src/minimax/messages/", - "scope": [ - "litellm-rust/crates/llms/src/minimax/messages/**" - ] - }, - { - "path": "litellm-rust/crates/llms/src/mistral/ocr/AGENTS.md", - "description": "Conventions for code under litellm-rust/crates/llms/src/mistral/ocr/", - "scope": [ - "litellm-rust/crates/llms/src/mistral/ocr/**" - ] - }, - { - "path": "litellm-rust/crates/llms/src/openai/responses/AGENTS.md", - "description": "Conventions for code under litellm-rust/crates/llms/src/openai/responses/", - "scope": [ - "litellm-rust/crates/llms/src/openai/responses/**" - ] - }, - { - "path": "litellm-rust/crates/llms/src/openai_like/chat/AGENTS.md", - "description": "Conventions for code under litellm-rust/crates/llms/src/openai_like/chat/", - "scope": [ - "litellm-rust/crates/llms/src/openai_like/chat/**" - ] - }, - { - "path": "litellm-rust/crates/llms/src/openrouter/messages/AGENTS.md", - "description": "Conventions for code under litellm-rust/crates/llms/src/openrouter/messages/", - "scope": [ - "litellm-rust/crates/llms/src/openrouter/messages/**" - ] - }, - { - "path": "litellm-rust/crates/llms/src/reducto/ocr/AGENTS.md", - "description": "Conventions for code under litellm-rust/crates/llms/src/reducto/ocr/", - "scope": [ - "litellm-rust/crates/llms/src/reducto/ocr/**" - ] - }, - { - "path": "litellm-rust/crates/llms/src/tencent/messages/AGENTS.md", - "description": "Conventions for code under litellm-rust/crates/llms/src/tencent/messages/", - "scope": [ - "litellm-rust/crates/llms/src/tencent/messages/**" - ] - }, - { - "path": "litellm-rust/crates/llms/src/vertex_ai/messages/AGENTS.md", - "description": "Conventions for code under litellm-rust/crates/llms/src/vertex_ai/messages/", - "scope": [ - "litellm-rust/crates/llms/src/vertex_ai/messages/**" - ] - }, - { - "path": "litellm-rust/crates/llms/src/vertex_ai/ocr/AGENTS.md", - "description": "Conventions for code under litellm-rust/crates/llms/src/vertex_ai/ocr/", - "scope": [ - "litellm-rust/crates/llms/src/vertex_ai/ocr/**" - ] - }, - { - "path": "litellm-rust/crates/llms-types/AGENTS.md", - "description": "Conventions for code under litellm-rust/crates/llms-types/", - "scope": [ - "litellm-rust/crates/llms-types/**" - ] - }, - { - "path": "litellm-rust/crates/llms-types/src/formats/audio_transcription/AGENTS.md", - "description": "Conventions for code under litellm-rust/crates/llms-types/src/formats/audio_transcription/", - "scope": [ - "litellm-rust/crates/llms-types/src/formats/audio_transcription/**" - ] - }, - { - "path": "litellm-rust/crates/llms-types/src/formats/batches/AGENTS.md", - "description": "Conventions for code under litellm-rust/crates/llms-types/src/formats/batches/", - "scope": [ - "litellm-rust/crates/llms-types/src/formats/batches/**" - ] - }, - { - "path": "litellm-rust/crates/llms-types/src/formats/chat_completions/AGENTS.md", - "description": "Conventions for code under litellm-rust/crates/llms-types/src/formats/chat_completions/", - "scope": [ - "litellm-rust/crates/llms-types/src/formats/chat_completions/**" - ] - }, - { - "path": "litellm-rust/crates/llms-types/src/formats/messages/AGENTS.md", - "description": "Conventions for code under litellm-rust/crates/llms-types/src/formats/messages/", - "scope": [ - "litellm-rust/crates/llms-types/src/formats/messages/**" - ] - }, - { - "path": "litellm-rust/crates/llms-types/src/formats/ocr/AGENTS.md", - "description": "Conventions for code under litellm-rust/crates/llms-types/src/formats/ocr/", - "scope": [ - "litellm-rust/crates/llms-types/src/formats/ocr/**" - ] - }, - { - "path": "litellm-rust/crates/llms-types/src/formats/responses/AGENTS.md", - "description": "Conventions for code under litellm-rust/crates/llms-types/src/formats/responses/", - "scope": [ - "litellm-rust/crates/llms-types/src/formats/responses/**" - ] - }, - { - "path": "litellm-rust/crates/llms-types/src/providers/anthropic/AGENTS.md", - "description": "Conventions for code under litellm-rust/crates/llms-types/src/providers/anthropic/", - "scope": [ - "litellm-rust/crates/llms-types/src/providers/anthropic/**" - ] - }, - { - "path": "litellm-rust/crates/llms-types/src/providers/bedrock/AGENTS.md", - "description": "Conventions for code under litellm-rust/crates/llms-types/src/providers/bedrock/", - "scope": [ - "litellm-rust/crates/llms-types/src/providers/bedrock/**" - ] - }, - { - "path": "litellm-rust/crates/llms-types/src/providers/minimax/AGENTS.md", - "description": "Conventions for code under litellm-rust/crates/llms-types/src/providers/minimax/", - "scope": [ - "litellm-rust/crates/llms-types/src/providers/minimax/**" - ] - }, - { - "path": "litellm-rust/crates/llms-types/src/providers/vertex_ai/AGENTS.md", - "description": "Conventions for code under litellm-rust/crates/llms-types/src/providers/vertex_ai/", - "scope": [ - "litellm-rust/crates/llms-types/src/providers/vertex_ai/**" - ] - }, - { - "path": "litellm-rust/crates/model-catalog/AGENTS.md", - "description": "Conventions for code under litellm-rust/crates/model-catalog/", - "scope": [ - "litellm-rust/crates/model-catalog/**" - ] - }, - { - "path": "litellm-rust/crates/python-bridge/AGENTS.md", - "description": "Conventions for code under litellm-rust/crates/python-bridge/", - "scope": [ - "litellm-rust/crates/python-bridge/**" - ] - }, - { - "path": "litellm-rust/crates/python-bridge/src/cache/AGENTS.md", - "description": "Conventions for code under litellm-rust/crates/python-bridge/src/cache/", - "scope": [ - "litellm-rust/crates/python-bridge/src/cache/**" - ] - }, - { - "path": "litellm-rust/crates/python-bridge/src/cache/native/AGENTS.md", - "description": "Conventions for code under litellm-rust/crates/python-bridge/src/cache/native/", - "scope": [ - "litellm-rust/crates/python-bridge/src/cache/native/**" - ] - }, - { - "path": "litellm-rust/crates/python-bridge/src/cache/python/AGENTS.md", - "description": "Conventions for code under litellm-rust/crates/python-bridge/src/cache/python/", - "scope": [ - "litellm-rust/crates/python-bridge/src/cache/python/**" - ] - }, - { - "path": "litellm-rust/crates/python-bridge/src/routes/AGENTS.md", - "description": "Conventions for code under litellm-rust/crates/python-bridge/src/routes/", - "scope": [ - "litellm-rust/crates/python-bridge/src/routes/**" - ] - }, - { - "path": "litellm-rust/crates/python-compat/AGENTS.md", - "description": "Conventions for code under litellm-rust/crates/python-compat/", - "scope": [ - "litellm-rust/crates/python-compat/**" - ] - }, - { - "path": "litellm-rust/crates/secrets-aws/AGENTS.md", - "description": "Conventions for code under litellm-rust/crates/secrets-aws/", - "scope": [ - "litellm-rust/crates/secrets-aws/**" - ] - }, - { - "path": "litellm-rust/crates/secrets-azure/AGENTS.md", - "description": "Conventions for code under litellm-rust/crates/secrets-azure/", - "scope": [ - "litellm-rust/crates/secrets-azure/**" - ] - }, - { - "path": "litellm-rust/crates/secrets-cyberark/AGENTS.md", - "description": "Conventions for code under litellm-rust/crates/secrets-cyberark/", - "scope": [ - "litellm-rust/crates/secrets-cyberark/**" - ] - }, - { - "path": "litellm-rust/crates/secrets-google/AGENTS.md", - "description": "Conventions for code under litellm-rust/crates/secrets-google/", - "scope": [ - "litellm-rust/crates/secrets-google/**" - ] - }, - { - "path": "litellm-rust/crates/secrets-hashicorp/AGENTS.md", - "description": "Conventions for code under litellm-rust/crates/secrets-hashicorp/", - "scope": [ - "litellm-rust/crates/secrets-hashicorp/**" - ] - }, - { - "path": "litellm-rust/crates/storage-clickhouse/AGENTS.md", - "description": "Conventions for code under litellm-rust/crates/storage-clickhouse/", - "scope": [ - "litellm-rust/crates/storage-clickhouse/**" - ] - }, - { - "path": "litellm-rust/docs/adrs/AGENTS.md", - "description": "Conventions for code under litellm-rust/docs/adrs/", - "scope": [ - "litellm-rust/docs/adrs/**" - ] - }, - { - "path": "tests/AGENTS.md", - "description": "Conventions for code under tests/", - "scope": [ - "tests/**" - ] - }, - { - "path": "tests/e2e/AGENTS.md", - "description": "Conventions for code under tests/e2e/", - "scope": [ - "tests/e2e/**" - ] - }, - { - "path": "tests/e2e_harness/AGENTS.md", - "description": "Conventions for code under tests/e2e_harness/", - "scope": [ - "tests/e2e_harness/**" - ] - }, - { - "path": "tests/integration/AGENTS.md", - "description": "Conventions for code under tests/integration/", - "scope": [ - "tests/integration/**" - ] - }, - { - "path": "tests/integration/translation/AGENTS.md", - "description": "Conventions for code under tests/integration/translation/", - "scope": [ - "tests/integration/translation/**" - ] - }, - { - "path": "tests/rust-python-harness/AGENTS.md", - "description": "Conventions for code under tests/rust-python-harness/", - "scope": [ - "tests/rust-python-harness/**" - ] - }, - { - "path": "tests/rust-python-harness/strategies/e2e_parity/AGENTS.md", - "description": "Conventions for code under tests/rust-python-harness/strategies/e2e_parity/", - "scope": [ - "tests/rust-python-harness/strategies/e2e_parity/**" - ] - }, - { - "path": "tests/rust-python-harness/strategies/trace_parity/AGENTS.md", - "description": "Conventions for code under tests/rust-python-harness/strategies/trace_parity/", - "scope": [ - "tests/rust-python-harness/strategies/trace_parity/**" - ] - }, - { - "path": "tests/rust-python-harness/strategies/unit_tests_parity/AGENTS.md", - "description": "Conventions for code under tests/rust-python-harness/strategies/unit_tests_parity/", - "scope": [ - "tests/rust-python-harness/strategies/unit_tests_parity/**" - ] - }, - { - "path": "tests/rust-python-harness/strategies/unit_tests_rust/AGENTS.md", - "description": "Conventions for code under tests/rust-python-harness/strategies/unit_tests_rust/", - "scope": [ - "tests/rust-python-harness/strategies/unit_tests_rust/**" - ] - }, - { - "path": "tests/test_litellm_rust/AGENTS.md", - "description": "Conventions for code under tests/test_litellm_rust/", - "scope": [ - "tests/test_litellm_rust/**" - ] - }, - { - "path": "tests/unit/AGENTS.md", - "description": "Conventions for code under tests/unit/", - "scope": [ - "tests/unit/**" - ] - }, - { - "path": "tests/unit/rust_bridge/AGENTS.md", - "description": "Conventions for code under tests/unit/rust_bridge/", - "scope": [ - "tests/unit/rust_bridge/**" - ] - }, - { - "path": "ui/litellm-dashboard/AGENTS.md", - "description": "Conventions for code under ui/litellm-dashboard/", - "scope": [ - "ui/litellm-dashboard/**" - ] - }, - { - "path": "ui/litellm-dashboard/src/components/chat/AGENTS.md", - "description": "Conventions for code under ui/litellm-dashboard/src/components/chat/", - "scope": [ - "ui/litellm-dashboard/src/components/chat/**" - ] - } - ] -} diff --git a/AGENTS.md b/AGENTS.md index 7f8e9cedc26..2f254a2dadf 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -130,4 +130,4 @@ Before implementing: Ask yourself: "Would a senior engineer say this is overcomplicated?" If yes, simplify -Before requesting maintainer review, verify the current PR tip passes required CI and code coverage, meets Greptile confidence of at least 4/5, and has acceptable Veria and Bugbot reviews. Inspect warnings and findings, fix actionable issues, and rerun the affected checks and reviewers after changes. Record evidence for any false positive or unavailable review. Never treat a pending or missing bot result as a pass. Do not lower coverage thresholds or raise `ANY_CAPS` to satisfy a check +Before requesting maintainer review, verify the current PR tip passes required CI and code coverage and has acceptable reviews from the review bots on the PR, such as Veria and Bugbot. Inspect warnings and findings, fix actionable issues, and rerun the affected checks and reviewers after changes. Record evidence for any false positive or unavailable review. Never treat a pending or missing bot result as a pass. Do not lower coverage thresholds or raise `ANY_CAPS` to satisfy a check diff --git a/README.md b/README.md index 9987f13d127..3bd0fd54a65 100644 --- a/README.md +++ b/README.md @@ -9,7 +9,7 @@ Deploy to Render Deploy on Railway Deploy on AWS - Deploy on GCP + Deploy on GCP

LiteLLM Proxy Server (AI Gateway) | Hosted Proxy | Enterprise Tier | Website

@@ -342,6 +342,7 @@ Set `LITELLM_PROXY_API_BASE` and `LITELLM_PROXY_API_KEY` and every model call th | [Cohere Chat (`cohere_chat`)](https://docs.litellm.ai/docs/providers/cohere) | ✅ | ✅ | ✅ | | | | | | | | | [CometAPI (`cometapi`)](https://docs.litellm.ai/docs/providers/cometapi) | ✅ | ✅ | ✅ | ✅ | | | | | | | | [CompactifAI (`compactifai`)](https://docs.litellm.ai/docs/providers/compactifai) | ✅ | ✅ | ✅ | | | | | | | | +| [CoralBricks (`coralbricks`)](https://docs.litellm.ai/docs/providers/coralbricks) | ✅ | ✅ | ✅ | | | | | | | | | [Custom (`custom`)](https://docs.litellm.ai/docs/providers/custom_llm_server) | ✅ | ✅ | ✅ | | | | | | | | | [Custom OpenAI (`custom_openai`)](https://docs.litellm.ai/docs/providers/openai_compatible) | ✅ | ✅ | ✅ | | | ✅ | ✅ | ✅ | ✅ | | | [Dashscope (`dashscope`)](https://docs.litellm.ai/docs/providers/dashscope) | ✅ | ✅ | ✅ | ✅ | | | | | | ✅ | @@ -526,9 +527,9 @@ Provider API keys live in AWS Secrets Manager; reference ARNs via `gateway_extra #### GCP — Cloud Run + Cloud SQL + Memorystore + HTTPS LB -[![Open in Cloud Shell](https://gstatic.com/cloudssh/images/open-btn.png)](https://ssh.cloud.google.com/cloudshell/editor?cloudshell_git_repo=https%3A%2F%2Fgithub.com%2FBerriAI%2Flitellm&cloudshell_workspace=terraform%2Flitellm%2Fgcp%2Fexamples%2Fdefault&cloudshell_tutorial=TUTORIAL.md&cloudshell_image=gcr.io/ds-artifacts-cloudshell/deploystack_custom_image&shellonly=true) +[![Open in Cloud Shell](https://gstatic.com/cloudssh/images/open-btn.png)](https://ssh.cloud.google.com/cloudshell/editor?cloudshell_git_repo=https%3A%2F%2Fgithub.com%2FBerriAI%2Fterraform-google-litellm&cloudshell_git_branch=v1.105.0&cloudshell_workspace=examples%2Fdefault&cloudshell_tutorial=TUTORIAL.md&cloudshell_image=gcr.io/ds-artifacts-cloudshell/deploystack_custom_image&shellonly=true) -Real 1-click. Opens Cloud Shell, clones this repo, and walks you through `terraform apply` via a built-in [DeployStack tutorial](./terraform/litellm/gcp/examples/default/TUTORIAL.md) — pick the project, the tutorial sets up the Artifact Registry remote repo, writes `terraform.tfvars` from your answers, and runs apply. +Real 1-click. Opens Cloud Shell, clones the [terraform-google-litellm](https://github.com/BerriAI/terraform-google-litellm) module mirror at the latest stable release (a small download instead of the full litellm repo), and walks you through `terraform apply` via a built-in [DeployStack tutorial](./terraform/litellm/gcp/examples/default/TUTORIAL.md). Pick the project, the tutorial sets up the Artifact Registry remote repo, writes `terraform.tfvars` from your answers, and runs apply [Module page →](https://registry.terraform.io/modules/BerriAI/litellm/google/latest) diff --git a/ci_cd/cost_map_guard.py b/ci_cd/cost_map_guard.py index 14eecf01461..e7b873e3f5c 100644 --- a/ci_cd/cost_map_guard.py +++ b/ci_cd/cost_map_guard.py @@ -56,7 +56,7 @@ def _rendered_schema(cost_map: CostMap) -> str: def _evaluation_output_price_failure( key: str, entry: object, evaluation_output_price_sources: Mapping[str, str] ) -> str | None: - if not isinstance(entry, dict) or entry.get("mode") != "evaluation": + if not isinstance(entry, dict) or entry.get("mode") not in ("evaluation", "decisions"): return None output_cost: Final = entry.get("output_cost_per_token") if ( @@ -67,7 +67,7 @@ def _evaluation_output_price_failure( ): return None return ( - f"{key}: mode evaluation rows bill input only; output_cost_per_token must be 0 unless the vendor " + f"{key}: mode evaluation and decisions rows bill input only; output_cost_per_token must be 0 unless the vendor " "prices output tokens for this model (add the source to EVALUATION_OUTPUT_PRICE_SOURCES)" ) diff --git a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py index 4c06f0b9802..67894fbbf3c 100644 --- a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py +++ b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py @@ -1162,7 +1162,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): ) return False - async def check_file_ids_access(self, file_ids: List[str], user_api_key_dict: UserAPIKeyAuth) -> None: + async def check_file_ids_access(self, file_ids: Sequence[str], user_api_key_dict: UserAPIKeyAuth) -> None: """ Check if the user has access to a list of file IDs. Only checks managed (unified) file IDs. @@ -1269,6 +1269,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): elif call_type == CallTypes.acreate_batch.value: input_file_id = cast(Optional[str], data.get("input_file_id")) if input_file_id: + await self.check_file_ids_access((input_file_id,), user_api_key_dict) model_file_id_mapping = await self.get_model_file_id_mapping( [input_file_id], user_api_key_dict.parent_otel_span ) @@ -1315,11 +1316,12 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): elif retrieve_object_id and accessor_key == "batch_id": await self.enforce_batch_object_access(retrieve_object_id, user_api_key_dict) elif call_type == CallTypes.acreate_fine_tuning_job.value: - input_file_id = cast(Optional[str], data.get("training_file")) - if input_file_id: - model_file_id_mapping = await self.get_model_file_id_mapping( - [input_file_id], user_api_key_dict.parent_otel_span - ) + fine_tuning_file_ids: Final = tuple( + file_id + for file_id in (data.get("training_file"), data.get("validation_file")) + if isinstance(file_id, str) and file_id + ) + await self.check_file_ids_access(fine_tuning_file_ids, user_api_key_dict) return data diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index b0c5b02a245..2f77b09221a 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -4024,6 +4024,7 @@ dependencies = [ "litellm-cache-memory", "litellm-cache-response", "litellm-core-utils", + "litellm-framer", "litellm-host", "litellm-host-native", "litellm-http", @@ -4107,9 +4108,11 @@ name = "litellm-inference-testing" version = "0.1.0" dependencies = [ "futures-util", + "litellm-auth", "litellm-http", "litellm-inference", "litellm-secrets", + "tokio", ] [[package]] diff --git a/litellm-rust/crates/callbacks-legacy-python/src/adapter.rs b/litellm-rust/crates/callbacks-legacy-python/src/adapter.rs index dd3da8991cc..9f56f12a8f9 100644 --- a/litellm-rust/crates/callbacks-legacy-python/src/adapter.rs +++ b/litellm-rust/crates/callbacks-legacy-python/src/adapter.rs @@ -442,7 +442,7 @@ impl LegacyLogging { py: Python<'_>, facts: &litellm_host::interceptors::ExecutionFacts, ) -> PyResult> { - use litellm_host::interceptors::ResultSource; + use litellm_host::interceptors::{Cost, ResultSource}; let logger = self.logger()?.object(py); let params = logger @@ -462,6 +462,12 @@ impl LegacyLogging { ), )?; let details = logger.getattr("model_call_details")?; + let reported_cost = match &facts.cost { + Cost::Reported { amount } => amount.as_f64(), + Cost::CacheHit | Cost::Deferred => None, + }; + logger.setattr("_provider_reported_cost", reported_cost)?; + details.call_method1("pop", ("response_cost", py.None()))?; self.cache_key = match &facts.source { ResultSource::Provider => None, ResultSource::Cache { key } => Some(key.clone()), @@ -958,6 +964,81 @@ on_pre_call = lambda additional_args: None check = lambda: None "; + #[rstest] + #[case::reported(Some(0.37), false)] + #[case::free(Some(0.0), false)] + #[case::deferred(None, false)] + #[case::cached(None, true)] + fn result_ready_hands_only_a_reported_cost_to_the_spend_logging_record( + #[case] reported: Option, + #[case] cache_hit: bool, + #[values(0.001, 7.0)] existing_cost: f64, + ) { + Python::initialize(); + Python::attach(|py| { + let locals = namespace(py, PAYLOAD_LOGGER); + run( + py, + &locals, + c" +logger.litellm_params = {} +logger.optional_params = {} +logger.model_call_details = {'response_cost': 7.0} +logger._provider_reported_cost = None +", + ); + let mut logging = LegacyLogging { + logger: Some(PythonLogger::new(local(&locals, "logger").unbind())), + ..legacy_call(py, &locals, false) + }; + let source = if cache_hit { + litellm_host::interceptors::ResultSource::Cache { + key: "cached-test".into(), + } + } else { + litellm_host::interceptors::ResultSource::Provider + }; + let facts = litellm_host::interceptors::ExecutionFacts { + provider: litellm_host::interceptors::ProviderIdentity { + model: "native-test".into(), + provider: "edenai".into(), + }, + cost: litellm_host::interceptors::Cost::settle( + &source, + reported.and_then(serde_json::Number::from_f64), + ), + source, + }; + let previous = litellm_host::interceptors::ExecutionFacts { + provider: facts.provider.clone(), + cost: litellm_host::interceptors::Cost::Reported { + amount: serde_json::Number::from_f64(9.0).unwrap(), + }, + source: litellm_host::interceptors::ResultSource::Provider, + }; + logging.result_ready(py, &previous).unwrap(); + local(&locals, "logger") + .getattr("model_call_details") + .unwrap() + .set_item("response_cost", existing_cost) + .unwrap(); + let step = logging.result_ready(py, &facts).unwrap(); + assert!(matches!(step, HookStep::Ready(()))); + locals.set_item("reported", reported).unwrap(); + locals.set_item("cache_hit", cache_hit).unwrap(); + run( + py, + &locals, + c" +assert logger.model_call_details.get('response_cost') is None, logger.model_call_details +assert logger._provider_reported_cost == reported, logger._provider_reported_cost +assert logger.model_call_details['cache_hit'] is cache_hit, logger.model_call_details +assert logger.update['custom_llm_provider'] == 'edenai', logger.update +", + ); + }); + } + const DOCUMENT: &str = "data:application/pdf;base64,YWJj"; const EDITED: &str = "data:application/pdf;base64,ZWRpdGVk"; diff --git a/litellm-rust/crates/core-utils/src/get_llm_provider_logic.rs b/litellm-rust/crates/core-utils/src/get_llm_provider_logic.rs index 982831bcc5e..acea48799c0 100644 --- a/litellm-rust/crates/core-utils/src/get_llm_provider_logic.rs +++ b/litellm-rust/crates/core-utils/src/get_llm_provider_logic.rs @@ -22,6 +22,8 @@ pub enum LlmProviders { Cohere, #[strum(serialize = "deepseek")] Deepseek, + #[strum(serialize = "edenai")] + Edenai, #[strum(serialize = "minimax")] Minimax, #[strum(serialize = "mistral")] diff --git a/litellm-rust/crates/host/src/interceptors.rs b/litellm-rust/crates/host/src/interceptors.rs index 3633fc7ee39..f0fd7434c6d 100644 --- a/litellm-rust/crates/host/src/interceptors.rs +++ b/litellm-rust/crates/host/src/interceptors.rs @@ -42,10 +42,32 @@ pub enum ResultSource { Cache { key: String }, } +/// What the call cost, settled once by the route. A host records it; it never re-derives it. +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum Cost { + /// The amount the provider itself billed, taken from its response. + Reported { amount: serde_json::Number }, + /// Served from the cache, so nothing was billed. + CacheHit, + /// No provider figure: the host estimates from usage and its price table. + Deferred, +} + +impl Cost { + pub fn settle(source: &ResultSource, reported: Option) -> Self { + match (source, reported) { + (ResultSource::Cache { .. }, _) => Self::CacheHit, + (ResultSource::Provider, Some(amount)) => Self::Reported { amount }, + (ResultSource::Provider, None) => Self::Deferred, + } + } +} + #[derive(Clone, Debug, PartialEq, Eq)] pub struct ExecutionFacts { pub provider: ProviderIdentity, pub source: ResultSource, + pub cost: Cost, } pub trait Interceptors: Send + Sync { diff --git a/litellm-rust/crates/inference-messages/Cargo.toml b/litellm-rust/crates/inference-messages/Cargo.toml index 6801c758e54..e310a0be9b1 100644 --- a/litellm-rust/crates/inference-messages/Cargo.toml +++ b/litellm-rust/crates/inference-messages/Cargo.toml @@ -26,6 +26,7 @@ serde_json = { workspace = true, features = ["preserve_order"] } tracing.workspace = true [dev-dependencies] +litellm-framer = { workspace = true, features = ["sse"] } litellm-cache-memory.workspace = true litellm-host-native.workspace = true litellm-http = { workspace = true, features = ["test-support"] } diff --git a/litellm-rust/crates/inference-messages/src/common_utils.rs b/litellm-rust/crates/inference-messages/src/common_utils.rs index 91b135a6a4a..66e7e1cfe25 100644 --- a/litellm-rust/crates/inference-messages/src/common_utils.rs +++ b/litellm-rust/crates/inference-messages/src/common_utils.rs @@ -8,6 +8,7 @@ use litellm_llms::{ base_llm::messages::transformation::BaseMessagesConfig, bedrock::messages::invoke_transformations::anthropic_claude3_transformation::BEDROCK_ANTHROPIC_MESSAGES_CONFIG, deepseek::messages::transformation::DEEPSEEK_ANTHROPIC_MESSAGES_CONFIG, + edenai::messages::transformation::EDENAI_MESSAGES_CONFIG, minimax::messages::transformation::MINIMAX_MESSAGES_CONFIG, vertex_ai::messages::transformation::VERTEX_ANTHROPIC_MESSAGES_CONFIG, }; @@ -53,6 +54,7 @@ pub(crate) fn messages_provider(provider: LlmProviders, model: &str) -> Option &BEDROCK_ANTHROPIC_MESSAGES_CONFIG, LlmProviders::BedrockMantle => &BEDROCK_MANTLE_MESSAGES_CONFIG, LlmProviders::Deepseek => &DEEPSEEK_ANTHROPIC_MESSAGES_CONFIG, + LlmProviders::Edenai => &EDENAI_MESSAGES_CONFIG, LlmProviders::Minimax => &MINIMAX_MESSAGES_CONFIG, LlmProviders::VertexAi if model.to_ascii_lowercase().contains("claude") => { &VERTEX_ANTHROPIC_MESSAGES_CONFIG @@ -85,6 +87,7 @@ mod tests { #[case::azure_ai("azure_ai")] #[case::bedrock("bedrock")] #[case::deepseek("deepseek")] + #[case::edenai("edenai")] #[case::minimax("minimax")] #[case::vertex_ai("vertex_ai")] fn provider_keeps_its_python_name(#[case] name: &str) { diff --git a/litellm-rust/crates/inference-messages/src/handler.rs b/litellm-rust/crates/inference-messages/src/handler.rs index 4e883320303..178f96c4c33 100644 --- a/litellm-rust/crates/inference-messages/src/handler.rs +++ b/litellm-rust/crates/inference-messages/src/handler.rs @@ -25,7 +25,7 @@ use litellm_inference::{context::CallContext, outbound::outbound_request}; pub(super) struct ProviderCall { pub identity: ProviderIdentity, pub wire: WireRequest, - provider: super::common_utils::MessagesProvider, + pub(super) provider: super::common_utils::MessagesProvider, signer: Option, timeout: Option, stream: bool, diff --git a/litellm-rust/crates/inference-messages/src/lib.rs b/litellm-rust/crates/inference-messages/src/lib.rs index 6cf7e34e52a..c690c7a2ef4 100644 --- a/litellm-rust/crates/inference-messages/src/lib.rs +++ b/litellm-rust/crates/inference-messages/src/lib.rs @@ -7,7 +7,7 @@ mod types; use futures_util::FutureExt; use litellm_auth::AuthServices; -use litellm_host::interceptors::{ExecutionFacts, Interceptors, ResultSource}; +use litellm_host::interceptors::{Cost, ExecutionFacts, Interceptors, ResultSource}; use litellm_inference::{caching::CallCache, context::CallContext}; use litellm_secrets::source::SecretSource; @@ -87,6 +87,7 @@ impl MessagesRoute { &request.wire, ); let identity = request.identity.clone(); + let config = request.provider.config(); let (output, source) = match cache.lookup().await { Some(hit) => hit, None => ( @@ -94,9 +95,14 @@ impl MessagesRoute { ResultSource::Provider, ), }; + let reported = match &output { + MessagesCallResponse::Complete(response) => config.reported_cost(response), + MessagesCallResponse::Stream { .. } => None, + }; context .result_ready(ExecutionFacts { provider: identity, + cost: Cost::settle(&source, reported), source: source.clone(), }) .await?; diff --git a/litellm-rust/crates/inference-messages/tests/caching.rs b/litellm-rust/crates/inference-messages/tests/caching.rs index e9005123119..32c962091ad 100644 --- a/litellm-rust/crates/inference-messages/tests/caching.rs +++ b/litellm-rust/crates/inference-messages/tests/caching.rs @@ -13,8 +13,8 @@ use litellm_cache_response::{ CacheOptions, CacheScope, ResponseCache, ResponseCacheConfig, ResponseCacheService, ScopedCache, }; use litellm_host::interceptors::{ - ExecutionFacts, Interceptors, ProviderIdentity, RawResponse, RequestContext, ResultSource, - WireRequest, + Cost, ExecutionFacts, Interceptors, ProviderIdentity, RawResponse, RequestContext, + ResultSource, WireRequest, }; use litellm_inference::{ RouteError, @@ -241,3 +241,55 @@ async fn messages_cache_identity_follows_resolved_configuration_and_request_call first.verify().await; second.verify().await; } + +#[rstest] +#[case::reported(json!(0.37), Some(0.37))] +#[case::free(json!(0), Some(0.0))] +#[case::malformed(json!("invalid"), None)] +#[case::missing(Value::Null, None)] +#[tokio::test] +async fn provider_reported_cost_is_delivered_once_and_never_charged_for_a_cache_hit( + cache: Arc, + #[case] cost: Value, + #[case] expected: Option, +) { + use litellm_inference_messages::{MessagesCall, MessagesCallResponse}; + let response = json!({"id": "msg_cost", "type": "message", "role": "assistant", "model": "native-test", + "content": [{"type": "text", "text": "answer"}], "stop_reason": "end_turn", "stop_sequence": null, + "usage": {"input_tokens": 3, "output_tokens": 2}, "cost": cost}); + let upstream = support::upstream([support::json_response(response)]).await; + let hooks = ChangingHooks::default(); + let route = support::messages_route(litellm_inference_testing::no_secrets()) + .with_cache(ScopedCache::new(cache, CacheScope::Shared)); + let request = || { + MessagesCall { + body: serde_json::from_value(json!({"model": "native-test", "messages": [{"role": "user", "content": "hi"}], "max_tokens": 16})).unwrap(), + api_key: Some("eden-key".into()), api_base: Some(upstream.uri()), custom_llm_provider: Some("edenai".into()), + litellm_params: Default::default(), extra_headers: None, provider_specific_header: None, + timeout: None, shaping: Default::default(), + } + }; + let first = route.execute(request(), &hooks, None).await.unwrap(); + let second = route.execute(request(), &hooks, None).await.unwrap(); + let (MessagesCallResponse::Complete(first), MessagesCallResponse::Complete(second)) = + (first, second) + else { + panic!("expected Messages responses"); + }; + assert_eq!(first, second); + assert_eq!(first.extra.get("cost"), Some(&cost)); + assert_eq!(support::received(&upstream).await.len(), 1); + let facts = hooks.facts.lock().unwrap(); + let [provider, cached] = facts.as_slice() else { + panic!("expected provider and cache results"); + }; + assert_eq!(provider.source, ResultSource::Provider); + assert_eq!( + provider.cost, + expected.map_or(Cost::Deferred, |amount| Cost::Reported { + amount: serde_json::Number::from_f64(amount).unwrap() + }) + ); + assert!(matches!(cached.source, ResultSource::Cache { .. })); + assert_eq!(cached.cost, Cost::CacheHit); +} diff --git a/litellm-rust/crates/inference-messages/tests/live.rs b/litellm-rust/crates/inference-messages/tests/live.rs new file mode 100644 index 00000000000..24468598141 --- /dev/null +++ b/litellm-rust/crates/inference-messages/tests/live.rs @@ -0,0 +1,8 @@ +#[path = "live/support.rs"] +mod support; + +#[path = "live/edenai.rs"] +mod edenai; + +#[path = "live/vertex_ai.rs"] +mod vertex_ai; diff --git a/litellm-rust/crates/inference-messages/tests/live/edenai.rs b/litellm-rust/crates/inference-messages/tests/live/edenai.rs new file mode 100644 index 00000000000..28c17679865 --- /dev/null +++ b/litellm-rust/crates/inference-messages/tests/live/edenai.rs @@ -0,0 +1,92 @@ +use litellm_host::interceptors::Cost; +use litellm_inference_messages::MessagesRoute; +use litellm_inference_testing::live::within_deadline; +use rstest::rstest; +use serde_json::{Number, Value, json}; + +use super::support::{LiveCall, assert_text, call, complete, request, route, stream_events}; + +const PROVIDER: &str = "edenai"; + +#[rstest] +#[case::basic(json!({"messages": [{"role": "user", "content": "Say hello."}]}))] +#[case::system_text(json!({ + "system": "You are a helpful assistant.", + "messages": [{"role": "user", "content": "Say hello."}] +}))] +#[case::portable_cache_hint(json!({ + "system": [{"type": "text", "text": "You are a helpful assistant.", "cache_control": {"type": "ephemeral", "ttl": "1h"}}], + "messages": [{"role": "user", "content": "Say hello."}] +}))] +#[ignore = "calls a real provider, requires credentials and a live model"] +#[tokio::test] +async fn completion_and_cost_handoff(route: MessagesRoute, #[case] payload: Value) { + within_deadline(async { + let host = LiveCall::new(PROVIDER); + let message = complete(request(&route, call(PROVIDER, payload), &host).await); + assert_text(&message); + host.assert_provider_result(); + let expected_cost = message + .extra + .get("cost") + .and_then(|value| match value { + Value::Number(number) => number.as_f64(), + Value::String(text) => text.parse::().ok(), + _ => None, + }) + .filter(|cost| *cost >= 0.0) + .and_then(Number::from_f64) + .map_or(Cost::Deferred, |amount| Cost::Reported { amount }); + assert_eq!(host.facts().cost, expected_cost); + println!("{}", serde_json::to_string(&message).unwrap()); + }) + .await; +} + +#[rstest] +#[ignore = "calls a real provider, requires credentials and a live model"] +#[tokio::test] +async fn streaming(route: MessagesRoute) { + within_deadline(async { + let host = LiveCall::new(PROVIDER); + let events = stream_events( + request( + &route, + call( + PROVIDER, + json!({ + "stream": true, "messages": [{"role": "user", "content": "Say hello."}] + }), + ), + &host, + ) + .await, + ) + .await; + assert!(events.iter().any(|event| { + event["delta"]["text"] + .as_str() + .is_some_and(|text| !text.is_empty()) + })); + host.assert_provider_result(); + assert_eq!(host.facts().cost, Cost::Deferred); + }) + .await; +} + +#[rstest] +#[case::complete(false)] +#[case::streaming(true)] +#[ignore = "calls a real provider, requires credentials and a live model"] +#[tokio::test] +async fn tool_round_trip(route: MessagesRoute, #[case] stream: bool) { + super::support::tool_round_trip( + route, + PROVIDER, + stream, + json!({ + "max_tokens": 512, "tool_choice": {"type": "auto"} + }), + ) + .await; +} diff --git a/litellm-rust/crates/inference-messages/tests/live/support.rs b/litellm-rust/crates/inference-messages/tests/live/support.rs new file mode 100644 index 00000000000..81e295e34ad --- /dev/null +++ b/litellm-rust/crates/inference-messages/tests/live/support.rs @@ -0,0 +1,265 @@ +use std::{sync::OnceLock, time::Duration}; + +use futures_util::{FutureExt, TryStreamExt}; +use litellm_host::interceptors::{ + ExecutionFacts, Interceptors, RawResponse, RequestContext, ResultSource, WireRequest, +}; +use litellm_inference_messages::{Error, MessagesCall, MessagesCallResponse, MessagesRoute}; +use litellm_inference_testing::live::{LiveResources, model, within_deadline}; +use litellm_llms_types::formats::messages::MessagesResponse; +use serde_json::{Value, json}; + +pub(super) struct LiveCall { + provider: &'static str, + facts: OnceLock, +} + +impl LiveCall { + pub(super) fn new(provider: &'static str) -> Self { + Self { + provider, + facts: OnceLock::new(), + } + } + + pub(super) fn facts(&self) -> &ExecutionFacts { + self.facts + .get() + .expect("the route must deliver execution facts") + } + + pub(super) fn assert_provider_result(&self) { + let facts = self.facts(); + assert_eq!(facts.provider.provider, self.provider); + assert_eq!(facts.source, ResultSource::Provider); + println!( + "{}", + json!({"route": "messages", "provider": self.provider, "source": "provider"}) + ); + } +} + +impl Interceptors for LiveCall { + async fn result_ready(&self, facts: ExecutionFacts) -> Result<(), Error> { + assert!( + self.facts.set(facts).is_ok(), + "result must be accepted exactly once" + ); + Ok(()) + } + + async fn before_provider_request( + &self, + wire: WireRequest, + context: RequestContext, + ) -> Result { + assert_eq!(context.custom_llm_provider, self.provider); + println!( + "{}", + json!({"route": "messages", "provider": self.provider, "model": context.model}) + ); + Ok(wire) + } + + async fn after_provider_response(&self, _: RawResponse) -> Result<(), Error> { + Ok(()) + } +} + +#[rstest::fixture] +pub(super) fn route() -> MessagesRoute { + let resources = LiveResources::default(); + MessagesRoute::new(resources.http, resources.auth, resources.secrets) +} + +pub(super) async fn request( + route: &MessagesRoute, + call: MessagesCall, + host: &LiveCall, +) -> MessagesCallResponse { + route + .execute(call, host, None) + .boxed() + .await + .expect("live request failed") +} + +pub(super) fn call(provider: &str, payload: Value) -> MessagesCall { + let Value::Object(fields) = payload else { + panic!("live payload must be an object") + }; + let body = serde_json::from_value(Value::Object( + [("max_tokens".into(), json!(128))] + .into_iter() + .chain(fields) + .chain([("model".into(), json!(model("messages", provider)))]) + .collect(), + )) + .expect("live payload must be a valid Messages request"); + MessagesCall { + body, + api_key: None, + api_base: None, + custom_llm_provider: Some(provider.into()), + litellm_params: Default::default(), + extra_headers: None, + provider_specific_header: None, + timeout: Some(Duration::from_secs(60)), + shaping: Default::default(), + } +} + +pub(super) fn complete(response: MessagesCallResponse) -> Box { + match response { + MessagesCallResponse::Complete(message) => message, + MessagesCallResponse::Stream { .. } => panic!("a complete request returned a stream"), + } +} + +pub(super) fn assert_text(message: &MessagesResponse) { + assert!(!message.id.is_empty()); + assert_eq!(message.role, "assistant"); + assert!(message.stop_reason.is_some()); + assert!( + message + .content + .iter() + .any(|block| block["text"].as_str().is_some_and(|text| !text.is_empty())) + ); + assert!( + message + .usage + .as_ref() + .and_then(|usage| usage["output_tokens"].as_u64()) + .is_some_and(|count| count > 0) + ); +} + +pub(super) async fn stream_events(response: MessagesCallResponse) -> Vec { + let MessagesCallResponse::Stream { head, chunks } = response else { + panic!("stream request returned a complete response") + }; + let content_types: Vec<_> = head + .headers + .iter() + .filter(|(name, _)| name.eq_ignore_ascii_case("content-type")) + .map(|(_, value)| value.split(';').next().unwrap().trim()) + .collect(); + assert_eq!(content_types, ["text/event-stream"]); + let frames = litellm_framer::frames(chunks, litellm_framer::sse::SseCodec::default()) + .try_collect::>() + .await + .expect("live SSE framing failed"); + let data_frames = match frames.last() { + Some(frame) if frame.data == "[DONE]" => &frames[..frames.len() - 1], + _ => frames.as_slice(), + }; + let events: Vec = data_frames + .iter() + .map(|frame| serde_json::from_str(&frame.data).expect("SSE data must be JSON")) + .collect(); + assert_eq!( + events.first().map(|event| event["type"].as_str()), + Some(Some("message_start")) + ); + assert_eq!( + events.last().map(|event| event["type"].as_str()), + Some(Some("message_stop")) + ); + assert!(!events.iter().any(|event| event["type"] == "error")); + assert!(events.iter().any(|event| { + event["usage"]["output_tokens"] + .as_u64() + .is_some_and(|count| count > 0) + })); + println!("{}", json!({"events": events})); + events +} + +fn prompt() -> Value { + json!({"role": "user", "content": "Use the echo tool with value live-check, then tell me the result."}) +} + +fn tools() -> Value { + json!([{"name": "echo", "description": "Return a value", "input_schema": { + "type": "object", "properties": {"value": {"type": "string", "enum": ["live-check"]}}, "required": ["value"] + }}]) +} + +fn fragments(events: &[Value], index: u64, field: &str) -> String { + events + .iter() + .filter(|event| event["index"].as_u64() == Some(index)) + .filter_map(|event| event["delta"][field].as_str()) + .collect() +} + +fn streamed_block(start: &Value, events: &[Value]) -> Value { + let index = start["index"] + .as_u64() + .expect("content block must have an index"); + let block = &start["content_block"]; + match block["type"].as_str() { + Some("tool_use") => { + let partial = fragments(events, index, "partial_json"); + let input = match partial.is_empty() { + true => block["input"].clone(), + false => { + serde_json::from_str(&partial).expect("streamed tool arguments must be JSON") + } + }; + json!({"type": "tool_use", "id": block["id"], "name": block["name"], "input": input}) + } + Some("text") => { + json!({"type": "text", "text": format!("{}{}", block["text"].as_str().unwrap_or_default(), fragments(events, index, "text"))}) + } + Some("thinking") => { + json!({"type": "thinking", "thinking": format!("{}{}", block["thinking"].as_str().unwrap_or_default(), fragments(events, index, "thinking")), + "signature": format!("{}{}", block["signature"].as_str().unwrap_or_default(), fragments(events, index, "signature"))}) + } + _ => block.clone(), + } +} + +pub(super) async fn tool_round_trip( + route: MessagesRoute, + provider: &'static str, + stream: bool, + tool_params: Value, +) { + within_deadline(async { + let first_host = LiveCall::new(provider); + let payload = json!({"messages": [prompt()], "tools": tools(), "stream": stream}); + let fields = payload.as_object().expect("tool payload must be an object").clone(); + let params = tool_params.as_object().expect("tool parameters must be an object").clone(); + let first = request(&route, call(provider, Value::Object(fields.into_iter().chain(params).collect())), &first_host).await; + let content = match stream { + false => { + let message = complete(first); + assert_eq!(message.stop_reason.as_deref(), Some("tool_use")); + message.content + } + true => { + let events = stream_events(first).await; + assert!(events.iter().any(|event| event["delta"]["stop_reason"] == "tool_use")); + events.iter().filter(|event| event["type"] == "content_block_start") + .map(|event| streamed_block(event, &events)).collect() + } + }; + let tool = content.iter().find(|block| block["type"] == "tool_use").expect("requested tool must be returned"); + assert_eq!(tool["name"], "echo"); + assert_eq!(tool["input"]["value"], "live-check"); + let tool_id = tool["id"].as_str().filter(|id| !id.is_empty()).expect("tool must have an id"); + first_host.assert_provider_result(); + let second_host = LiveCall::new(provider); + let second = complete(request(&route, call(provider, json!({ + "tools": tools(), "messages": [prompt(), + {"role": "assistant", "content": content}, + {"role": "user", "content": [{"type": "tool_result", "tool_use_id": tool_id, "content": "live-check"}]}] + })), &second_host).await); + assert_text(&second); + second_host.assert_provider_result(); + println!("{}", serde_json::to_string(&second).unwrap()); + }) + .await; +} diff --git a/litellm-rust/crates/inference-messages/tests/live/vertex_ai.rs b/litellm-rust/crates/inference-messages/tests/live/vertex_ai.rs new file mode 100644 index 00000000000..1f46f295e6b --- /dev/null +++ b/litellm-rust/crates/inference-messages/tests/live/vertex_ai.rs @@ -0,0 +1,68 @@ +use litellm_inference_messages::MessagesRoute; +use litellm_inference_testing::live::within_deadline; +use rstest::rstest; +use serde_json::{Value, json}; + +use super::support::{LiveCall, assert_text, call, complete, request, route, stream_events}; + +const PROVIDER: &str = "vertex_ai"; + +#[rstest] +#[case::basic(json!({"messages": [{"role": "user", "content": "Say hello."}]}))] +#[case::system_text(json!({ + "system": "You are a helpful assistant.", + "messages": [{"role": "user", "content": "Say hello."}] +}))] +#[ignore = "calls a real provider, requires credentials and a live model"] +#[tokio::test] +async fn completion(route: MessagesRoute, #[case] payload: Value) { + within_deadline(async { + let host = LiveCall::new(PROVIDER); + let message = complete(request(&route, call(PROVIDER, payload), &host).await); + assert_text(&message); + host.assert_provider_result(); + println!("{}", serde_json::to_string(&message).unwrap()); + }) + .await; +} + +#[rstest] +#[ignore = "calls a real provider, requires credentials and a live model"] +#[tokio::test] +async fn streaming(route: MessagesRoute) { + within_deadline(async { + let host = LiveCall::new(PROVIDER); + let response = request( + &route, + call( + PROVIDER, + json!({"stream": true, "messages": [{"role": "user", "content": "Say hello."}]}), + ), + &host, + ) + .await; + let events = stream_events(response).await; + assert!(events.iter().any(|event| { + event["delta"]["text"] + .as_str() + .is_some_and(|text| !text.is_empty()) + })); + host.assert_provider_result(); + }) + .await; +} + +#[rstest] +#[case::complete(false)] +#[case::streaming(true)] +#[ignore = "calls a real provider, requires credentials and a live model"] +#[tokio::test] +async fn tool_round_trip(route: MessagesRoute, #[case] stream: bool) { + super::support::tool_round_trip( + route, + PROVIDER, + stream, + json!({"tool_choice": {"type": "auto"}}), + ) + .await; +} diff --git a/litellm-rust/crates/inference-messages/tests/messages/request.rs b/litellm-rust/crates/inference-messages/tests/messages/request.rs index 7c02d2a2e5a..aa8cb05b170 100644 --- a/litellm-rust/crates/inference-messages/tests/messages/request.rs +++ b/litellm-rust/crates/inference-messages/tests/messages/request.rs @@ -894,3 +894,38 @@ async fn provider_validation_runs_before_caller_parameter_removal( ); assert!(received(&upstream).await.is_empty()); } + +#[rstest] +#[tokio::test] +async fn edenai_preserves_native_thinking_and_strips_billing_and_cache_extensions( + call: MessagesCall, +) { + let upstream = upstream([message_response()]).await; + let thinking = json!({"type": "enabled", "budget_tokens": 512, "native_field": true}); + run_message(MessagesCall { + custom_llm_provider: Some("edenai".into()), + api_base: Some(format!("{}/v3", upstream.uri())), + api_key: Some("eden-key".into()), + extra_headers: headers([("Anthropic-Beta", "unregistered-native-beta")]), + ..with_fields(call, json!({ + "thinking": thinking, "temperature": 0.25, + "system": [{"type": "text", "text": "x-anthropic-billing-header: hidden"}, + {"type": "text", "text": "keep", "cache_control": {"type": "ephemeral", "ttl": "1h", "scope": "global"}}], + })) + }).await; + let sent = only_request(&upstream).await; + assert_eq!(sent.url.path(), "/v3/v1/messages"); + assert_eq!(sent.header("authorization"), Some("Bearer eden-key")); + assert_eq!( + sent.header("anthropic-beta"), + Some("unregistered-native-beta") + ); + assert_eq!( + sent.json(), + json!({ + "model": MODEL, "max_tokens": 16, "messages": [{"role": "user", "content": "hi"}], + "thinking": thinking, "temperature": 0.25, + "system": [{"type": "text", "text": "keep", "cache_control": {"type": "ephemeral"}}], + }) + ); +} diff --git a/litellm-rust/crates/inference-messages/tests/messages/stream.rs b/litellm-rust/crates/inference-messages/tests/messages/stream.rs index 7deaa4fbd3f..cda688107b5 100644 --- a/litellm-rust/crates/inference-messages/tests/messages/stream.rs +++ b/litellm-rust/crates/inference-messages/tests/messages/stream.rs @@ -348,6 +348,7 @@ async fn the_timeout_covers_a_stalled_stream_body(call: MessagesCall) { #[rstest] #[case::anthropic("anthropic")] #[case::azure_ai("azure_ai")] +#[case::edenai("edenai")] #[tokio::test] async fn the_sdk_returns_stream_headers_and_every_sse_byte( call: MessagesCall, @@ -358,6 +359,7 @@ async fn the_sdk_returns_stream_headers_and_every_sse_byte( .execute( MessagesCall { custom_llm_provider: Some(provider.into()), + api_key: Some("stream-key".into()), litellm_params: Default::default(), ..streaming(call, upstream.uri()) }, diff --git a/litellm-rust/crates/inference-testing/Cargo.toml b/litellm-rust/crates/inference-testing/Cargo.toml index b4b9d6df79a..748b1df0d3b 100644 --- a/litellm-rust/crates/inference-testing/Cargo.toml +++ b/litellm-rust/crates/inference-testing/Cargo.toml @@ -7,7 +7,9 @@ repository.workspace = true publish = false [dependencies] +litellm-auth.workspace = true futures-util.workspace = true litellm-http.workspace = true litellm-inference.workspace = true litellm-secrets.workspace = true +tokio.workspace = true diff --git a/litellm-rust/crates/inference-testing/README.md b/litellm-rust/crates/inference-testing/README.md new file mode 100644 index 00000000000..caf66073c35 --- /dev/null +++ b/litellm-rust/crates/inference-testing/README.md @@ -0,0 +1,46 @@ +# Live provider route tests + +Each route crate owns `tests/live.rs`. Use named `rstest` provider cases for the behaviors that route supports. Shared setup lives in `litellm_inference_testing::live`, with production HTTP settings, authentication services, environment secrets and a deadline + +These tests execute the Rust route against a real provider. They do not prove Python parity, proxy authorization, billing persistence or agent compatibility. Those belong in the corresponding boundary and gateway suites + +All live cases use `#[ignore = "calls a real provider, requires credentials and a live model"]`. Normal Cargo runs report them as ignored. An explicitly selected case fails on missing credentials, unavailable models, provider errors or timeouts. Never return early and count missing coverage as a pass + +Set `LITELLM_LIVE___MODEL` to the latest model available to the test account. The model has no compiled default because vendor availability changes independently of LiteLLM. Credentials use the provider's existing configuration rather than a second test-only key convention. Export credentials in the shell before running tests. The harness does not load `.env` or modify process environment while tests run + +## Messages + +Completion checks text and usage. Streaming consumes SSE to completion and checks text, usage and the terminal event. Tool round trips request a tool call and send its returned ID back with a result. Each case checks provider execution facts and rejects cache results + +Messages setup and assertions live in `inference-messages/tests/live/support.rs`. Provider modules under `tests/live/` select their payloads and tool policy. The shared helpers consume streams completely and reconstruct text, thinking signatures and tool arguments for the follow-up request + +The initial provider cases are `vertex_ai` and `edenai`. A registered ignored case is available coverage, not evidence that it passed against that provider + +Run one provider's cases explicitly + +```bash +LITELLM_LIVE_MESSAGES_VERTEX_AI_MODEL= \ + cargo test --manifest-path litellm-rust/Cargo.toml \ + -p litellm-inference-messages --test live vertex_ai \ + -- --ignored --nocapture --test-threads=1 +``` + +For Vertex AI, set `LITELLM_LIVE_MESSAGES_VERTEX_AI_MODEL` to a Claude model published in the project's location. Credentials, project and location resolve through the production Rust GCP authentication, from `GOOGLE_APPLICATION_CREDENTIALS` or `VERTEXAI_CREDENTIALS`, with `VERTEXAI_PROJECT` and `VERTEXAI_LOCATION` + +Vertex AI cases live in `inference-messages/tests/live/vertex_ai.rs`. They cover plain and system messages, text streaming, and complete and streaming tool round trips. Tool cases use automatic selection because current Claude models reject forced `tool_choice`. Each case still requires a tool call, checks its arguments, and sends the tool result back + +For EdenAI, set `LITELLM_LIVE_MESSAGES_EDENAI_MODEL` and `EDENAI_API_KEY`, and select `edenai`. EdenAI requires a production token for live coverage, because sandbox tokens return dummy responses + +EdenAI cases live in `inference-messages/tests/live/edenai.rs`. They cover system messages, portable cache hints, provider-reported cost handoff, text streaming, and complete and streaming tool round trips. The cost assertion compares the public response with the execution facts instead of pinning a vendor price +EdenAI tool cases use automatic selection because current Claude models reject forced `tool_choice`. The test still requires a tool call and verifies its arguments and follow-up response + +Other tool cases require a model that supports forced `tool_choice`, except where a provider's section says otherwise. A provider rejection fails the selected case instead of skipping it + +Select one behavior by its test name, for example `completion::case_1_basic`. To list every available cell without contacting any provider + +```bash +cargo test --manifest-path litellm-rust/Cargo.toml \ + -p litellm-inference-messages --test live -- --list +``` + +For another route, reuse `LiveResources`, `model` and `within_deadline`. Keep payload construction and response assertions in that route's test file. Add only supported provider behaviors, and document any gap instead of substituting mocks or skipping failures. Run cases serially to keep spending and rate limits predictable diff --git a/litellm-rust/crates/inference-testing/src/lib.rs b/litellm-rust/crates/inference-testing/src/lib.rs index c42a032e1d8..2e20d5a9e91 100644 --- a/litellm-rust/crates/inference-testing/src/lib.rs +++ b/litellm-rust/crates/inference-testing/src/lib.rs @@ -9,6 +9,8 @@ use litellm_secrets::{SecretValue, source::SecretSource}; use litellm_inference::resources::CoreResources; +pub mod live; + pub fn http_pool() -> HttpClientPool { HttpClientPool::new(Arc::new(PublicDnsResolver)) } diff --git a/litellm-rust/crates/inference-testing/src/live.rs b/litellm-rust/crates/inference-testing/src/live.rs new file mode 100644 index 00000000000..b63bfd91a38 --- /dev/null +++ b/litellm-rust/crates/inference-testing/src/live.rs @@ -0,0 +1,53 @@ +use std::{future::Future, sync::Arc, time::Duration}; + +use litellm_http::{ClientVariant, HttpSettings, HttpSettingsLayer, Resolution}; +use litellm_inference::resources::CoreResources; +use litellm_secrets::source::{EnvironmentSecrets, SecretSource}; + +pub const DEADLINE: Duration = Duration::from_secs(120); + +pub struct LiveResources { + pub http: litellm_http::Client, + pub auth: Arc, + pub secrets: Arc, +} + +impl Default for LiveResources { + fn default() -> Self { + let resources = CoreResources::new(Arc::new(super::http_pool())); + let settings = + HttpSettings::from_layers([HttpSettingsLayer::from_environment(&|key: &str| { + std::env::var(key).ok() + })]); + let resolution = Resolution::from(&settings); + let http = resources + .pool + .client(&resolution.config, ClientVariant::Provider) + .expect("live provider HTTP settings must be valid"); + let secrets_http = resources + .pool + .client(&resolution.config, ClientVariant::NoRedirect) + .expect("live secret HTTP settings must be valid"); + Self { + http, + auth: resources.auth, + secrets: Arc::new(EnvironmentSecrets::python_compatible(secrets_http)), + } + } +} + +pub fn model(route: &str, provider: &str) -> String { + let name = format!("LITELLM_LIVE_{route}_{provider}_MODEL").to_ascii_uppercase(); + std::env::var(&name) + .ok() + .filter(|value| !value.trim().is_empty()) + .unwrap_or_else(|| { + panic!("set {name} to an available current model before running this live case") + }) +} + +pub async fn within_deadline(operation: impl Future) -> T { + tokio::time::timeout(DEADLINE, operation) + .await + .expect("live case exceeded its 120 second deadline") +} diff --git a/litellm-rust/crates/inference/src/caching.rs b/litellm-rust/crates/inference/src/caching.rs index d2f4198b910..6dfda20ead4 100644 --- a/litellm-rust/crates/inference/src/caching.rs +++ b/litellm-rust/crates/inference/src/caching.rs @@ -13,7 +13,9 @@ use litellm_cache_response::{ }; use litellm_host::{ call::{CallOutput, OutputOf}, - interceptors::{ExecutionFacts, Interceptors, ProviderIdentity, ResultSource, WireRequest}, + interceptors::{ + Cost, ExecutionFacts, Interceptors, ProviderIdentity, ResultSource, WireRequest, + }, lifecycle::{CallEvent, ExecutionEvent}, observation::ObservationSender, protocol::Protocol, @@ -168,10 +170,12 @@ where None => (provider().await?, ResultSource::Provider), }; let from_provider = source == ResultSource::Provider; + let cost = Cost::settle(&source, None); publish( ExecutionFacts { provider: identity, source, + cost, }, interceptors, observers, @@ -212,6 +216,7 @@ where publish( ExecutionFacts { provider: identity, + cost: Cost::settle(&source, None), source: source.clone(), }, interceptors, diff --git a/litellm-rust/crates/llms/AGENTS.md b/litellm-rust/crates/llms/AGENTS.md index 2dbc74e905c..ec0f6ffe95a 100644 --- a/litellm-rust/crates/llms/AGENTS.md +++ b/litellm-rust/crates/llms/AGENTS.md @@ -24,7 +24,7 @@ - Every file is `# rules` then `# references`, both concise bullets. Split a long `# rules` into `##` topic sections - A folder gets an AGENTS.md only when it adds something: a deviation, an explicit reuse of another provider, or upstream docs no other file lists. A folder without one follows its nearest parent -- Every `src///` folder gets one, since it talks to its own upstream endpoint. `.greptile/files.json` lists them all +- Every `src///` folder gets one, since it talks to its own upstream endpoint - A rule lives in the highest file where it holds. Never restate a parent's rule - Each upstream URL appears in exactly one AGENTS.md, the one owning that contract. Others point to that file by path - Format specs: `litellm-llms-types/src/formats//` diff --git a/litellm-rust/crates/llms/src/base_llm/messages/transformation.rs b/litellm-rust/crates/llms/src/base_llm/messages/transformation.rs index 61398fe4ea9..6f0ea5c8f2a 100644 --- a/litellm-rust/crates/llms/src/base_llm/messages/transformation.rs +++ b/litellm-rust/crates/llms/src/base_llm/messages/transformation.rs @@ -41,6 +41,10 @@ pub trait BaseMessagesConfig: Sync { Ok(response) } + fn reported_cost(&self, _response: &MessagesResponse) -> Option { + None + } + fn secret_names(&self) -> &'static [&'static str]; /// Shapes the forwarded headers and names the credential, the way Python's diff --git a/litellm-rust/crates/llms/src/edenai/messages/mod.rs b/litellm-rust/crates/llms/src/edenai/messages/mod.rs new file mode 100644 index 00000000000..f239b6921fa --- /dev/null +++ b/litellm-rust/crates/llms/src/edenai/messages/mod.rs @@ -0,0 +1 @@ +pub mod transformation; diff --git a/litellm-rust/crates/llms/src/edenai/messages/transformation.rs b/litellm-rust/crates/llms/src/edenai/messages/transformation.rs new file mode 100644 index 00000000000..7fd001ddbbc --- /dev/null +++ b/litellm-rust/crates/llms/src/edenai/messages/transformation.rs @@ -0,0 +1,120 @@ +use litellm_auth::{CredentialPlacement, SecretValue}; +use litellm_core_utils::settings::resolve_non_empty; +use litellm_llms_types::formats::messages::{ + MessagesOptionalParams, MessagesRequest, MessagesResponse, +}; +use litellm_llms_types::providers::anthropic::DEFAULT_HEADERS; +use litellm_router_types::LitellmParams; +use serde_json::Value; + +use crate::{ + Error, + anthropic::{ + common_utils::{filter_billing_headers_from_system, has_anthropic_credential}, + messages::transformation::update_headers_with_anthropic_beta, + }, + base_llm::{ + auth::AuthScheme, + messages::{ + context::MessagesTransformContext, + transformation::{BaseMessagesConfig, Headers, ValidatedEnvironment}, + }, + }, + openai_like::messages::transformation::{complete_messages_url, portable_cache_control}, +}; + +const API_KEY_ENV: &str = "EDENAI_API_KEY"; +const API_BASE_ENV: &str = "EDENAI_API_BASE"; +const DEFAULT_BASE: &str = "https://api.edenai.run/v3"; + +pub struct EdenAIAnthropicMessagesConfig; +pub const EDENAI_MESSAGES_CONFIG: EdenAIAnthropicMessagesConfig = EdenAIAnthropicMessagesConfig; + +impl BaseMessagesConfig for EdenAIAnthropicMessagesConfig { + fn validate_environment( + &self, + headers: Headers, + api_key: Option<&str>, + _model: &str, + _params: &LitellmParams, + env: &dyn Fn(&str) -> Option, + ) -> Result { + let key = resolve_non_empty(api_key, env, &[API_KEY_ENV]).ok_or(Error::Auth( + litellm_auth::Error::MissingApiKey { + provider: "Eden AI", + environment_variable: API_KEY_ENV, + }, + ))?; + if has_anthropic_credential(&headers) { + return Ok(ValidatedEnvironment { + headers, + auth: AuthScheme::Forwarded, + }); + } + Ok(ValidatedEnvironment { + headers, + auth: AuthScheme::Credential { + placement: CredentialPlacement::Bearer, + secret: SecretValue::new(key), + }, + }) + } + + fn get_complete_url( + &self, + api_base: Option<&str>, + _model: &str, + _params: &LitellmParams, + _stream: bool, + env: &dyn Fn(&str) -> Option, + ) -> Result { + let base = resolve_non_empty(api_base, env, &[API_BASE_ENV]) + .unwrap_or_else(|| DEFAULT_BASE.into()); + Ok(complete_messages_url(&base)) + } + + fn transform_anthropic_messages_request( + &self, + request: MessagesRequest, + _context: &MessagesTransformContext, + ) -> Result { + Ok(MessagesRequest { + params: MessagesOptionalParams { + system: request + .params + .system + .and_then(filter_billing_headers_from_system), + ..request.params + }, + ..request + }) + } + + fn secret_names(&self) -> &'static [&'static str] { + &[API_KEY_ENV, API_BASE_ENV] + } + + fn default_headers(&self) -> &'static [(&'static str, &'static str)] { + DEFAULT_HEADERS + } + + fn request_headers(&self, headers: Headers, request: &MessagesRequest) -> Headers { + update_headers_with_anthropic_beta(headers, request) + } + + fn wire_body(&self, body: Value) -> Value { + portable_cache_control(body) + } + + fn reported_cost(&self, response: &MessagesResponse) -> Option { + let cost = match response.extra.get("cost")? { + Value::Number(cost) => cost.as_f64()?, + Value::String(cost) => cost.parse::().ok()?, + _ => return None, + }; + if cost < 0.0 { + return None; + } + serde_json::Number::from_f64(cost) + } +} diff --git a/litellm-rust/crates/llms/src/edenai/mod.rs b/litellm-rust/crates/llms/src/edenai/mod.rs new file mode 100644 index 00000000000..ba63992f3cb --- /dev/null +++ b/litellm-rust/crates/llms/src/edenai/mod.rs @@ -0,0 +1 @@ +pub mod messages; diff --git a/litellm-rust/crates/llms/src/lib.rs b/litellm-rust/crates/llms/src/lib.rs index 1cf69ab5745..eb232404d93 100644 --- a/litellm-rust/crates/llms/src/lib.rs +++ b/litellm-rust/crates/llms/src/lib.rs @@ -6,6 +6,7 @@ pub mod bedrock; pub mod bedrock_mantle; pub mod cohere; pub mod deepseek; +pub mod edenai; mod error; pub mod minimax; pub mod mistral; diff --git a/litellm-rust/crates/llms/src/openai_like/messages/mod.rs b/litellm-rust/crates/llms/src/openai_like/messages/mod.rs new file mode 100644 index 00000000000..f239b6921fa --- /dev/null +++ b/litellm-rust/crates/llms/src/openai_like/messages/mod.rs @@ -0,0 +1 @@ +pub mod transformation; diff --git a/litellm-rust/crates/llms/src/openai_like/messages/transformation.rs b/litellm-rust/crates/llms/src/openai_like/messages/transformation.rs new file mode 100644 index 00000000000..f69bc3d59be --- /dev/null +++ b/litellm-rust/crates/llms/src/openai_like/messages/transformation.rs @@ -0,0 +1,100 @@ +use serde_json::{Map, Value}; + +pub fn complete_messages_url(base: &str) -> String { + let base = base.trim_end_matches('/'); + if base.ends_with("/v1/messages") { + return base.into(); + } + format!("{}/v1/messages", base.strip_suffix("/v1").unwrap_or(base)) +} + +pub fn portable_cache_control(body: Value) -> Value { + let Value::Object(fields) = body else { + return body; + }; + Value::Object( + portable_block(fields) + .into_iter() + .map(|(key, value)| { + let value = match key.as_str() { + "system" | "tools" => map_blocks(value, portable_value), + "messages" => map_blocks(value, portable_message), + _ => value, + }; + (key, value) + }) + .collect(), + ) +} + +fn portable_block(fields: Map) -> Map { + fields + .into_iter() + .filter_map(|(key, value)| { + if key != "cache_control" { + return Some((key, value)); + } + let Value::Object(cache) = value else { + return None; + }; + let kind = cache + .get("type") + .and_then(Value::as_str) + .unwrap_or("ephemeral"); + Some((key, serde_json::json!({"type": kind}))) + }) + .collect() +} + +fn portable_value(value: Value) -> Value { + match value { + Value::Object(fields) => Value::Object(portable_block(fields)), + other => other, + } +} + +fn map_blocks(value: Value, map: fn(Value) -> Value) -> Value { + match value { + Value::Array(blocks) => Value::Array(blocks.into_iter().map(map).collect()), + other => other, + } +} + +fn portable_message(value: Value) -> Value { + let Value::Object(fields) = value else { + return value; + }; + Value::Object( + fields + .into_iter() + .map(|(key, value)| { + let value = if key == "content" { + map_blocks(value, portable_content) + } else { + value + }; + (key, value) + }) + .collect(), + ) +} + +fn portable_content(value: Value) -> Value { + let Value::Object(fields) = value else { + return value; + }; + let tool_result = fields.get("type").and_then(Value::as_str) == Some("tool_result"); + Value::Object( + portable_block(fields) + .into_iter() + .map(|(key, value)| { + let value = if tool_result && key == "content" { + map_blocks(value, portable_value) + } else { + value + }; + (key, value) + }) + .collect(), + ) +} diff --git a/litellm-rust/crates/llms/src/openai_like/mod.rs b/litellm-rust/crates/llms/src/openai_like/mod.rs index df0cc73a5b0..3ea6fc63fa4 100644 --- a/litellm-rust/crates/llms/src/openai_like/mod.rs +++ b/litellm-rust/crates/llms/src/openai_like/mod.rs @@ -1,2 +1,3 @@ pub mod chat; pub mod common_utils; +pub mod messages; diff --git a/litellm-rust/crates/llms/tests/edenai_messages_transformation.rs b/litellm-rust/crates/llms/tests/edenai_messages_transformation.rs new file mode 100644 index 00000000000..72b882f94ca --- /dev/null +++ b/litellm-rust/crates/llms/tests/edenai_messages_transformation.rs @@ -0,0 +1,34 @@ +use litellm_llms::{ + base_llm::messages::transformation::BaseMessagesConfig, + edenai::messages::transformation::EDENAI_MESSAGES_CONFIG, +}; +use litellm_llms_types::formats::messages::MessagesResponse; +use rstest::rstest; +use serde_json::{Value, json}; + +#[rstest] +#[case::number(json!(0.025), Some(0.025))] +#[case::free(json!(0), Some(0.0))] +#[case::text(json!("0.75"), Some(0.75))] +#[case::absent(Value::Null, None)] +#[case::malformed(json!("not-a-cost"), None)] +#[case::negative(json!(-1), None)] +#[case::boolean(json!(true), None)] +#[case::infinite(json!("inf"), None)] +#[case::nan(json!("NaN"), None)] +fn reported_cost_is_a_finite_nonnegative_amount( + #[case] cost: Value, + #[case] expected: Option, +) { + let response: MessagesResponse = serde_json::from_value(json!({ + "id": "msg_test", "type": "message", "role": "assistant", "model": "native-test", + "content": [], "stop_reason": "end_turn", "stop_sequence": null, "cost": cost, + })) + .unwrap(); + assert_eq!( + EDENAI_MESSAGES_CONFIG + .reported_cost(&response) + .and_then(|cost| cost.as_f64()), + expected + ); +} diff --git a/litellm-rust/crates/llms/tests/openai_like_messages_transformation.rs b/litellm-rust/crates/llms/tests/openai_like_messages_transformation.rs new file mode 100644 index 00000000000..824c7d8e5f2 --- /dev/null +++ b/litellm-rust/crates/llms/tests/openai_like_messages_transformation.rs @@ -0,0 +1,64 @@ +use litellm_llms::openai_like::messages::transformation::{ + complete_messages_url, portable_cache_control, +}; +use rstest::rstest; +use serde_json::{Value, json}; + +#[rstest] +#[case::root("https://native.test", "https://native.test/v1/messages")] +#[case::trailing_slash("https://native.test/", "https://native.test/v1/messages")] +#[case::versioned("https://native.test/v1/", "https://native.test/v1/messages")] +#[case::full_url( + "https://native.test/v3/v1/messages/", + "https://native.test/v3/v1/messages" +)] +#[case::gateway_prefix("https://native.test/v3", "https://native.test/v3/v1/messages")] +fn messages_url_appends_the_endpoint_once(#[case] base: &str, #[case] expected: &str) { + assert_eq!(complete_messages_url(base), expected); +} + +#[rstest] +#[case::extended(json!({"type": "ephemeral", "ttl": "1h", "scope": "global"}), Some(json!({"type": "ephemeral"})))] +#[case::missing_type(json!({"ttl": "1h"}), Some(json!({"type": "ephemeral"})))] +#[case::invalid_type(json!({"type": 7}), Some(json!({"type": "ephemeral"})))] +#[case::null(Value::Null, None)] +#[case::scalar(json!(false), None)] +fn portable_cache_only_rewrites_protocol_cache_hints( + #[case] cache: Value, + #[case] expected_cache: Option, +) { + let application_data = json!({"cache_control": cache, "nested": {"cache_control": cache}}); + let body = json!({ + "cache_control": cache, + "model": "native-test", + "system": [{"type": "text", "text": "system", "cache_control": cache}], + "tools": [{"name": "tool", "cache_control": cache, "input_schema": application_data}], + "messages": [{"role": "user", "content": [ + {"type": "tool_use", "input": application_data, "cache_control": cache}, + {"type": "tool_result", "content": [{"type": "text", "text": "result", "cache_control": cache}], "cache_control": cache}, + ]}], + }); + let hint = |value: Value| { + let fields = value.as_object().unwrap().clone(); + Value::Object( + fields + .into_iter() + .chain( + expected_cache + .clone() + .map(|cache| ("cache_control".into(), cache)), + ) + .collect(), + ) + }; + let expected = hint(json!({ + "model": "native-test", + "system": [hint(json!({"type": "text", "text": "system"}))], + "tools": [hint(json!({"name": "tool", "input_schema": application_data}))], + "messages": [{"role": "user", "content": [ + hint(json!({"type": "tool_use", "input": application_data})), + hint(json!({"type": "tool_result", "content": [hint(json!({"type": "text", "text": "result"}))]})), + ]}], + })); + assert_eq!(portable_cache_control(body), expected); +} diff --git a/litellm/__init__.py b/litellm/__init__.py index a33a78fb325..a7ea9ce4583 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -477,6 +477,7 @@ max_end_user_budget_id: Optional[str] = None # pass through unchanged. validate_end_user_id_in_db: bool = False block_requests_for_models_without_pricing: bool = False +personal_key_model_access_from_teams: bool = False disable_end_user_cost_tracking: Optional[bool] = None disable_end_user_cost_tracking_prometheus_only: Optional[bool] = None enable_end_user_cost_tracking_prometheus_only: Optional[bool] = None @@ -721,6 +722,7 @@ aws_polly_models: Set = set() transcribe_models: Set = set() gigachat_models: Set = set() llamagate_models: Set = set() +coralbricks_models: Set[str] = set() reducto_models: Set = set() bedrock_mantle_models: Set = set() @@ -1012,6 +1014,8 @@ def _populate_provider_model_sets(model_cost_map: Mapping[str, Mapping[str, obje gigachat_models.add(key) elif value.get("litellm_provider") == "llamagate": llamagate_models.add(key) + elif value.get("litellm_provider") == "coralbricks": + coralbricks_models.add(key) elif value.get("litellm_provider") == "reducto": reducto_models.add(key) elif value.get("litellm_provider") == "bedrock_mantle": @@ -1146,6 +1150,7 @@ model_list = list( | docker_model_runner_models | reducto_models | bedrock_mantle_models + | coralbricks_models | set(clarifai_models) ) @@ -1265,6 +1270,7 @@ def _build_models_by_provider() -> dict: "transcribe": transcribe_models, "gigachat": gigachat_models, "llamagate": llamagate_models, + "coralbricks": coralbricks_models, "reducto": reducto_models, "bedrock_mantle": bedrock_mantle_models, } @@ -1689,6 +1695,9 @@ if TYPE_CHECKING: from .llms.openrouter.decisions.transformation import ( OpenRouterDecisionsConfig as OpenRouterDecisionsConfig, ) + from .llms.fireworks_ai.decisions.transformation import ( + FireworksAIDecisionsConfig as FireworksAIDecisionsConfig, + ) from .llms.cloudflare.decisions.transformation import ( CloudflareDecisionsConfig as CloudflareDecisionsConfig, ) @@ -2277,6 +2286,7 @@ if TYPE_CHECKING: supports_parallel_function_calling: Callable[..., bool] supports_vision: Callable[..., bool] supports_audio_input: Callable[..., bool] + supports_video_input: Callable[..., bool] supports_audio_output: Callable[..., bool] supports_system_messages: Callable[..., bool] supports_reasoning: Callable[..., bool] diff --git a/litellm/_lazy_imports_registry.py b/litellm/_lazy_imports_registry.py index ce69de3b240..cbcb64c9197 100644 --- a/litellm/_lazy_imports_registry.py +++ b/litellm/_lazy_imports_registry.py @@ -36,6 +36,7 @@ UTILS_NAMES: Final = ( "supports_parallel_function_calling", "supports_vision", "supports_audio_input", + "supports_video_input", "supports_audio_output", "supports_system_messages", "supports_reasoning", @@ -160,6 +161,7 @@ LLM_CONFIG_NAMES: Final = ( "PerplexityDecisionsConfig", "TypeSafeDecisionsConfig", "OpenRouterDecisionsConfig", + "FireworksAIDecisionsConfig", "CloudflareDecisionsConfig", "StrandsDeciderDecisionsConfig", "DatabricksDecisionsConfig", @@ -481,6 +483,7 @@ _UTILS_IMPORT_MAP: Final = { ), "supports_vision": (".utils", "supports_vision"), "supports_audio_input": (".utils", "supports_audio_input"), + "supports_video_input": (".utils", "supports_video_input"), "supports_audio_output": (".utils", "supports_audio_output"), "supports_system_messages": (".utils", "supports_system_messages"), "supports_reasoning": (".utils", "supports_reasoning"), @@ -722,6 +725,7 @@ _LLM_CONFIGS_IMPORT_MAP: Final = { "PerplexityDecisionsConfig": (".llms.perplexity.decisions.transformation", "PerplexityDecisionsConfig"), "TypeSafeDecisionsConfig": (".llms.typesafe.decisions.transformation", "TypeSafeDecisionsConfig"), "OpenRouterDecisionsConfig": (".llms.openrouter.decisions.transformation", "OpenRouterDecisionsConfig"), + "FireworksAIDecisionsConfig": (".llms.fireworks_ai.decisions.transformation", "FireworksAIDecisionsConfig"), "CloudflareDecisionsConfig": (".llms.cloudflare.decisions.transformation", "CloudflareDecisionsConfig"), "StrandsDeciderDecisionsConfig": ( ".llms.strands_decider.decisions.transformation", diff --git a/litellm/caching/caching.py b/litellm/caching/caching.py index b928f33e961..05275d58c7d 100644 --- a/litellm/caching/caching.py +++ b/litellm/caching/caching.py @@ -445,7 +445,7 @@ class Cache: 1. If caching groups are set, then return the caching group as the model https://docs.litellm.ai/docs/routing#caching-across-model-groups 2. Else if a model_group is set, then return the model_group as the model. This is used for all requests sent through the litellm.Router() - 3. Else use the `model` passed in kwargs + 3. Else use the `model` passed in kwargs, plus `model_id` when set (e.g. a SageMaker inference component) """ litellm_params: Final[dict] = kwargs.get("litellm_params", {}) or {} metadata_sources: Final[tuple[dict, ...]] = ( @@ -460,7 +460,10 @@ class Cache: caching_group: Final = next( (group for source in metadata_sources if (group := self._get_caching_group(source, model_group))), None ) - return caching_group or model_group or kwargs["model"] + model_id_suffix: Final = ( + f" model_id: {kwargs['model_id']}" if "model_id" in kwargs and kwargs["model_id"] is not None else "" + ) + return caching_group or model_group or f"{kwargs['model']}{model_id_suffix}" def _get_caching_group(self, metadata: dict, model_group: str | None) -> str | None: caching_groups: Final[list | None] = metadata.get("caching_groups", []) diff --git a/litellm/completion_extras/litellm_responses_transformation/transformation.py b/litellm/completion_extras/litellm_responses_transformation/transformation.py index 3b68b58a6d9..2e4bb47654e 100644 --- a/litellm/completion_extras/litellm_responses_transformation/transformation.py +++ b/litellm/completion_extras/litellm_responses_transformation/transformation.py @@ -74,7 +74,7 @@ if TYPE_CHECKING: ChatCompletionToolReferenceObject, OpenAIMessageContentListBlock, ) - from litellm.types.utils import Choices + from litellm.types.utils import Choices, Message _CHAT_COMPLETION_FIELDS: Final = frozenset((*ModelResponse.model_fields, "usage")) @@ -312,6 +312,67 @@ def _as_chat_reasoning_items( return cast(list[ChatCompletionReasoningItem], list(reasoning_items)) +def _shift_annotation(annotation: ChatCompletionAnnotation, offset: int) -> ChatCompletionAnnotation: + shifted: Final = { + key: value + offset if key in ("start_index", "end_index") and isinstance(value, int) else value + for key, value in annotation.items() + } + return cast(ChatCompletionAnnotation, shifted) + + +def _message_annotations(message: "Message") -> Sequence[ChatCompletionAnnotation]: + annotations: Sequence[ChatCompletionAnnotation] | None = getattr(message, "annotations", None) + return annotations or () + + +def _message_reasoning_content(message: "Message") -> str | None: + reasoning_content: str | None = getattr(message, "reasoning_content", None) + return reasoning_content + + +def _message_reasoning_items(message: "Message") -> Sequence[ChatCompletionReasoningItem]: + reasoning_items: Sequence[ChatCompletionReasoningItem] | None = getattr(message, "reasoning_items", None) + return reasoning_items or () + + +def _shifted_annotations(messages: Sequence["Message"], offsets: Sequence[int]) -> Iterator[ChatCompletionAnnotation]: + for message, offset in zip(messages, offsets): + for annotation in _message_annotations(message): + yield _shift_annotation(annotation, offset) + + +def _fold_choices_into_one( + choices: Sequence["Choices"], + tool_calls: Sequence[Mapping[str, object]], + trailing_reasoning_content: str | None, + trailing_reasoning_item: _BuiltReasoningItem | None, +) -> "Choices": + """Clients that only read choices[0] must still see every text and tool call.""" + from litellm.types.utils import Choices, Message + + messages: Final = tuple(choice.message for choice in choices) + texts: Final = tuple(str(message.content or "") for message in messages) + offsets: Final = (0, *accumulate(len(text) for text in texts)) + annotations: Final = tuple(_shifted_annotations(messages, offsets)) + reasoning_contents: Final = (*(_message_reasoning_content(m) for m in messages), trailing_reasoning_content) + reasoning_items: Final = ( + *chain.from_iterable(_message_reasoning_items(m) for m in messages), + *(() if trailing_reasoning_item is None else (trailing_reasoning_item,)), + ) + return Choices( + message=Message( + role=messages[0].role if messages else "assistant", + content="".join(texts) if messages else None, + annotations=list(annotations) or None, + tool_calls=list(tool_calls) or None, + reasoning_content=" ".join(value for value in reasoning_contents if value) or None, + reasoning_items=_as_chat_reasoning_items(reasoning_items), + ), + finish_reason="tool_calls" if tool_calls else "stop", + index=0, + ) + + _ToolChoiceT = TypeVar("_ToolChoiceT") @@ -958,47 +1019,16 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): else: pass # don't fail request if item in list is not supported - if accumulated_tool_calls and choices: - last_choice: Final = choices[-1] - last_reasoning_content: Final = getattr(last_choice.message, "reasoning_content", None) - last_reasoning_items: Final = getattr(last_choice.message, "reasoning_items", None) - merged_reasoning_content: Final = ( - " ".join(value for value in (last_reasoning_content, reasoning_content) if value) or None - ) - merged_reasoning_items: Final = _as_chat_reasoning_items( - ( - *(last_reasoning_items or ()), - *(() if pending_reasoning_item is None else (pending_reasoning_item,)), - ) - ) - merged_message: Final = Message( - role=last_choice.message.role, - content=last_choice.message.content, - annotations=getattr(last_choice.message, "annotations", None), + if not choices and not accumulated_tool_calls: + return choices + return [ + _fold_choices_into_one( + choices=choices, tool_calls=accumulated_tool_calls, - reasoning_content=merged_reasoning_content, - reasoning_items=merged_reasoning_items, + trailing_reasoning_content=reasoning_content, + trailing_reasoning_item=pending_reasoning_item, ) - return [ - *choices[:-1], - Choices(message=merged_message, finish_reason="tool_calls", index=last_choice.index), - ] - - if accumulated_tool_calls: - msg = Message( - content=None, - tool_calls=accumulated_tool_calls, - reasoning_content=reasoning_content, - reasoning_items=cast( - list[ChatCompletionReasoningItem] | None, - ([pending_reasoning_item] if pending_reasoning_item is not None else None), - ), - ) - choices.append(Choices(message=msg, finish_reason="tool_calls", index=index)) - reasoning_content = None - pending_reasoning_item = None - - return choices + ] @staticmethod def _build_empty_incomplete_choice( diff --git a/litellm/constants.py b/litellm/constants.py index 38d16d61d7b..04ed307330c 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1007,6 +1007,7 @@ openai_compatible_endpoints: Final[list] = [ "https://api.meta.ai/v1", "https://api.sailresearch.com/v1", "https://api.cognition.ai/v1", + "https://inference.coralbricks.ai/v1", "https://api.cortecs.ai/v1", "https://api.scx.ai/v1", "https://api.prisminference.com/v1", @@ -1083,6 +1084,7 @@ openai_compatible_providers: Final[list] = [ "darkbloom", "meta", # Meta Model API (Muse Spark) - JSON-configured provider "cognition", + "coralbricks", "cortecs", "scx-ai", "prism", @@ -1961,6 +1963,7 @@ LITELLM_SETTINGS_SAFE_DB_OVERRIDES: Final = [ "openai_system_messages_first", "max_ui_session_budget", "budget_rollover", + "personal_key_model_access_from_teams", "mcp_tool_search", "turn_off_message_logging", "datadog_params", diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 45d98043144..263d0bfcb81 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -18,6 +18,8 @@ from litellm.constants import ( DEFAULT_MAX_LRU_CACHE_SIZE, DEFAULT_REPLICATE_GPU_PRICE_PER_SECOND, ) +from litellm.litellm_core_utils.cost_discount import resolve_cost_discount +from litellm.litellm_core_utils.decisions_pricing import decisions_pricing_fields, is_decisions_call_type from litellm.litellm_core_utils.hidden_params import HIDDEN_PARAMS_ATTR from litellm.litellm_core_utils.llm_cost_calc.tool_call_cost_tracking import ( StandardBuiltInToolCostTracking, @@ -570,7 +572,22 @@ def cost_per_token( ): # Option 3. if user passed model="bedrock/anthropic.claude-3", use model="anthropic.claude-3" model = model_without_prefix - # see this https://learn.microsoft.com/en-us/azure/ai-services/openai/concepts/models + if is_decisions_call_type(call_type): + decisions_model_info: Final = cached_get_model_info_helper(model=model, custom_llm_provider=custom_llm_provider) + decisions_overrides: Final = decisions_pricing_fields( + litellm.model_cost.get(decisions_model_info["key"], {}) # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # model cost map is dynamically loaded + ) + if decisions_overrides: + return generic_cost_per_token( + model=model, + usage=usage_block, + custom_llm_provider=custom_llm_provider, + service_tier=service_tier, + data_residency=data_residency, + model_info=cast( # cast-ok: _decisions fields override same-named ModelInfo rates + ModelInfo, {**decisions_model_info, **decisions_overrides} + ), + ) if call_type == "speech" or call_type == "aspeech": lyria_generation_cost: Final = ( get_vertex_ai_lyria_generation_cost(model=model_without_prefix) @@ -1124,6 +1141,8 @@ def infer_call_type(call_type: CallTypesLiteral | None, completion_response: obj return "text_completion" elif isinstance(completion_response, LiteLLMSendMessageResponse): return "send_message" + elif isinstance(completion_response, (DecisionsResponse, OpenAIDecisionResponse)): + return "adecisions" return call_type @@ -1134,6 +1153,8 @@ _infer_call_type = infer_call_type def _apply_cost_discount( base_cost: float, custom_llm_provider: str | None, + model: str | None = None, + region_name: str | None = None, ) -> tuple[float, float, float]: """ Apply provider-specific cost discount from module-level config. @@ -1141,17 +1162,20 @@ def _apply_cost_discount( Args: base_cost: The base cost before discount custom_llm_provider: The LLM provider name + model: The model name, used to resolve / discount keys + region_name: The region segment the cost key may carry, so patterns can match the region-stripped name Returns: Tuple of (final_cost, discount_percent, discount_amount) """ original_cost: Final = base_cost - discount_percent = 0.0 - discount_amount = 0.0 + resolved_discount: Final = resolve_cost_discount( + litellm.cost_discount_config, custom_llm_provider, model, region_name + ) + discount_percent: Final = 0.0 if resolved_discount is None else resolved_discount + discount_amount: Final = original_cost * discount_percent - if custom_llm_provider and custom_llm_provider in litellm.cost_discount_config: - discount_percent = litellm.cost_discount_config[custom_llm_provider] - discount_amount = original_cost * discount_percent + if resolved_discount is not None: final_cost: Final = original_cost - discount_amount if verbose_logger.isEnabledFor(logging.DEBUG): @@ -1472,6 +1496,7 @@ def completion_cost( ) explicit_pricing: Final = custom_pricing is True or base_model is not None + custom_priced_request_model: Final = model if custom_pricing is True else None selected_model: Final = select_model_name_for_cost_calc( model=model, completion_response=completion_response, @@ -1726,6 +1751,8 @@ def completion_cost( ) = _apply_cost_discount( base_cost=_final_cost, custom_llm_provider=custom_llm_provider, + model=model, + region_name=region_name if isinstance(region_name, str) else None, ) # Apply margin from module-level config if configured @@ -1785,6 +1812,7 @@ def completion_cost( litellm_logging_obj=litellm_logging_obj, total_time=total_time, ) + model_for_discount = custom_priced_request_model or model # Calculate cost based on prompt_tokens, completion_tokens if ( "togethercomputer" in model or "together_ai" in model or custom_llm_provider == "together_ai" @@ -1893,6 +1921,8 @@ def completion_cost( ) = _apply_cost_discount( base_cost=_final_cost, custom_llm_provider=custom_llm_provider, + model=model_for_discount, + region_name=region_name if isinstance(region_name, str) else None, ) else: discount_percent = 0.0 @@ -2062,7 +2092,7 @@ def response_cost_calculator( getattr(response_object, HIDDEN_PARAMS_ATTR), ) provider_response_cost: Final = get_response_cost_from_hidden_params(hidden_params) - if provider_response_cost is not None: + if provider_response_cost is not None and not custom_pricing: return provider_response_cost response_cost = completion_cost( diff --git a/litellm/decisions/call.py b/litellm/decisions/call.py index 1ad2f006736..d2be14e5529 100644 --- a/litellm/decisions/call.py +++ b/litellm/decisions/call.py @@ -19,11 +19,13 @@ from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler from litellm.llms.openai.decisions.transformation import ir_to_openai_response, openai_request_to_ir from litellm.types.decisions import ( DecisionQuestion, + DecisionsAudio, DecisionsIRRequest, DecisionsIRResponse, DecisionsJSON, DecisionsRequestBody, DecisionsResponse, + DecisionsVideo, OpenAIDecisionInput, OpenAIDecisionQuestion, OpenAIDecisionRequestBody, @@ -35,6 +37,8 @@ from litellm.utils import ProviderConfigManager SystemOneQuestions: TypeAlias = Mapping[str, DecisionQuestion | Mapping[str, object]] OpenAIDecisionQuestions: TypeAlias = Sequence[OpenAIDecisionQuestion | Mapping[str, object]] +DecisionsAudioClips: TypeAlias = Sequence[DecisionsAudio | Mapping[str, object]] +DecisionsVideoClips: TypeAlias = Sequence[DecisionsVideo | Mapping[str, object]] DecisionsRequestFormat: TypeAlias = DecisionsRequestBody | OpenAIDecisionRequestBody _SYSTEMONE_REQUEST_ADAPTER: Final[TypeAdapter[DecisionsRequestBody]] = TypeAdapter(DecisionsRequestBody) @@ -49,19 +53,31 @@ _HANDLER: Final = BaseLLMHTTPHandler() class SystemOneRequestFields: state: DecisionsJSON | None questions: SystemOneQuestions | None + audio: DecisionsAudioClips | None = None + videos: DecisionsVideoClips | None = None def validate(self, safety_identifier: str | None) -> DecisionsRequestBody: - return _SYSTEMONE_REQUEST_ADAPTER.validate_python({"state": self.state, "questions": self.questions}) + return _SYSTEMONE_REQUEST_ADAPTER.validate_python( + {"state": self.state, "questions": self.questions, "audio": self.audio, "videos": self.videos} + ) @dataclass(frozen=True, slots=True) class OpenAIDecisionRequestFields: input: OpenAIDecisionInput | None questions: OpenAIDecisionQuestions | None + audio: DecisionsAudioClips | None = None + videos: DecisionsVideoClips | None = None def validate(self, safety_identifier: str | None) -> OpenAIDecisionRequestBody: return _OPENAI_REQUEST_ADAPTER.validate_python( - {"input": self.input, "questions": self.questions, "safety_identifier": safety_identifier} + { + "input": self.input, + "questions": self.questions, + "safety_identifier": safety_identifier, + "audio": self.audio, + "videos": self.videos, + } ) @@ -122,15 +138,49 @@ def _drops_params(kwargs: Mapping[str, object]) -> bool: return litellm.drop_params is True or normalize_drop_params(kwargs.get("drop_params")) is True -def _request_safety_identifier(safety_identifier: object, kwargs: Mapping[str, object]) -> str | None: - if isinstance(safety_identifier, str) or not _drops_params(kwargs): +def _request_safety_identifier( + safety_identifier: object, provider_config: BaseDecisionsConfig, kwargs: Mapping[str, object] +) -> str | None: + if not _drops_params(kwargs): return _SAFETY_IDENTIFIER_ADAPTER.validate_python(safety_identifier) + if isinstance(safety_identifier, str) and provider_config.supports_safety_identifier: + return safety_identifier return None +def _fields_without_unsupported_media( + fields: RequestFields, + provider_config: BaseDecisionsConfig, + model: str, + custom_llm_provider: str, + kwargs: Mapping[str, object], +) -> RequestFields: + if not _drops_params(kwargs): + return fields + return replace( + fields, + audio=fields.audio if provider_config.supports_audio_input(model, custom_llm_provider) else None, + videos=fields.videos if provider_config.supports_video_input(model, custom_llm_provider) else None, + ) + + @dataclass(frozen=True, slots=True) -class _UnsupportedSafetyIdentifier: - pass +class _UnsupportedParams: + names: tuple[str, ...] + + +def _unsupported_param_names( + ir_request: DecisionsIRRequest, provider_config: BaseDecisionsConfig, model: str, custom_llm_provider: str +) -> tuple[str, ...]: + checks: Final = ( + ( + "safety_identifier", + ir_request.safety_identifier is not None and not provider_config.supports_safety_identifier, + ), + ("audio", bool(ir_request.audio) and not provider_config.supports_audio_input(model, custom_llm_provider)), + ("videos", bool(ir_request.videos) and not provider_config.supports_video_input(model, custom_llm_provider)), + ) + return tuple(name for name, unsupported in checks if unsupported) def _provider_ir_request( @@ -138,14 +188,14 @@ def _provider_ir_request( *, safety_identifier: str | None, provider_config: BaseDecisionsConfig, - kwargs: Mapping[str, object], -) -> DecisionsIRRequest | _UnsupportedSafetyIdentifier: + model: str, + custom_llm_provider: str, +) -> DecisionsIRRequest | _UnsupportedParams: ir_request: Final = _ir_request(request, safety_identifier) - if ir_request.safety_identifier is None or provider_config.supports_safety_identifier: - return ir_request - if _drops_params(kwargs): - return replace(ir_request, safety_identifier=None) - return _UnsupportedSafetyIdentifier() + unsupported: Final = _unsupported_param_names(ir_request, provider_config, model, custom_llm_provider) + if unsupported: + return _UnsupportedParams(unsupported) + return ir_request def prepare_call( @@ -178,8 +228,10 @@ def prepare_call( except ValueError as error: raise litellm.BadRequestError(message=str(error), model=model, llm_provider=provider) from error try: - request_safety_identifier: Final = _request_safety_identifier(safety_identifier, kwargs) - request: Final = fields.validate(request_safety_identifier) + request_safety_identifier: Final = _request_safety_identifier(safety_identifier, provider_config, kwargs) + request: Final = _fields_without_unsupported_media( + fields, provider_config, canonical_model, provider, kwargs + ).validate(request_safety_identifier) except ValidationError as error: raise litellm.BadRequestError( message=f"Invalid Decisions request: {error}", @@ -206,12 +258,13 @@ def prepare_call( request, safety_identifier=request_safety_identifier, provider_config=provider_config, - kwargs=kwargs, + model=canonical_model, + custom_llm_provider=provider, ) - if isinstance(ir_request, _UnsupportedSafetyIdentifier): + if isinstance(ir_request, _UnsupportedParams): raise litellm.UnsupportedParamsError( message=( - f"{provider} does not support parameters: ['safety_identifier'], for model={model}. " + f"{provider} does not support parameters: {list(ir_request.names)}, for model={model}. " "To drop these, set `litellm.drop_params=True` or for proxy:\n\n" "`litellm_settings:\n drop_params: true`\n" ), diff --git a/litellm/decisions/main.py b/litellm/decisions/main.py index 8b13d04770a..921243612c7 100644 --- a/litellm/decisions/main.py +++ b/litellm/decisions/main.py @@ -6,7 +6,9 @@ import httpx from pydantic import TypeAdapter from litellm.decisions.call import ( + DecisionsAudioClips, DecisionsCall, + DecisionsVideoClips, OpenAIDecisionQuestions, OpenAIDecisionRequestFields, RequestFields, @@ -29,11 +31,16 @@ _OPENAI_QUESTIONS_ADAPTER: Final[TypeAdapter[OpenAIDecisionQuestions | None]] = def _request_fields( decision_input: OpenAIDecisionInput | None, questions: OpenAIDecisionQuestions | SystemOneQuestions | None, + audio: DecisionsAudioClips | None, + videos: DecisionsVideoClips | None, kwargs: Mapping[str, object], ) -> RequestFields: if "state" not in kwargs: return OpenAIDecisionRequestFields( - input=decision_input, questions=_OPENAI_QUESTIONS_ADAPTER.validate_python(questions) + input=decision_input, + questions=_OPENAI_QUESTIONS_ADAPTER.validate_python(questions), + audio=audio, + videos=videos, ) warnings.warn( "litellm.decisions(state=...) is deprecated, call litellm.systemone for the System One request format", @@ -43,6 +50,8 @@ def _request_fields( return SystemOneRequestFields( state=_STATE_ADAPTER.validate_python(kwargs["state"]), questions=_SYSTEMONE_QUESTIONS_ADAPTER.validate_python(questions), + audio=audio, + videos=videos, ) @@ -52,6 +61,8 @@ def _prepare_decisions_call( decision_input: OpenAIDecisionInput | None, questions: OpenAIDecisionQuestions | SystemOneQuestions | None, safety_identifier: str | None, + audio: DecisionsAudioClips | None, + videos: DecisionsVideoClips | None, api_key: str | None, api_base: str | None, timeout: float | httpx.Timeout | None, @@ -61,7 +72,7 @@ def _prepare_decisions_call( ) -> DecisionsCall: return prepare_call( model=model, - fields=_request_fields(decision_input, questions, kwargs), + fields=_request_fields(decision_input, questions, audio, videos, kwargs), safety_identifier=safety_identifier, api_key=api_key, api_base=api_base, @@ -83,6 +94,8 @@ async def adecisions( timeout: float | httpx.Timeout | None = None, custom_llm_provider: str | None = None, extra_headers: Mapping[str, str] | None = None, + audio: DecisionsAudioClips | None = None, + videos: DecisionsVideoClips | None = None, **kwargs: object, ) -> OpenAIDecisionResponse | DecisionsResponse: call: Final = _prepare_decisions_call( @@ -90,6 +103,8 @@ async def adecisions( decision_input=input, questions=questions, safety_identifier=safety_identifier, + audio=audio, + videos=videos, api_key=api_key, api_base=api_base, timeout=timeout, @@ -111,6 +126,8 @@ def decisions( timeout: float | httpx.Timeout | None = None, custom_llm_provider: str | None = None, extra_headers: Mapping[str, str] | None = None, + audio: DecisionsAudioClips | None = None, + videos: DecisionsVideoClips | None = None, **kwargs: object, ) -> OpenAIDecisionResponse | DecisionsResponse: call: Final = _prepare_decisions_call( @@ -118,6 +135,8 @@ def decisions( decision_input=input, questions=questions, safety_identifier=safety_identifier, + audio=audio, + videos=videos, api_key=api_key, api_base=api_base, timeout=timeout, diff --git a/litellm/experimental_mcp_client/client.py b/litellm/experimental_mcp_client/client.py index 9c5608076a2..d74c0cea195 100644 --- a/litellm/experimental_mcp_client/client.py +++ b/litellm/experimental_mcp_client/client.py @@ -41,6 +41,7 @@ class _SessionCallbackMap(TypedDict): sampling_callback: ReadOnly[SamplingFnT | None] elicitation_callback: ReadOnly[ElicitationFnT | None] logging_callback: ReadOnly[LoggingFnT | None] + response_hooks: ReadOnly[tuple[Callable[[httpx2.Response], Awaitable[None]], ...]] from mcp.types import ( @@ -96,6 +97,7 @@ from litellm.constants import ( ) from litellm.experimental_mcp_client.tools import list_tools_with_pagination from litellm.llms.custom_httpx.http_handler import get_ssl_configuration +from litellm.proxy._experimental.mcp_server.mcp_context import get_active_mcp_request_ctx from litellm.proxy._experimental.mcp_server.mcp_debug import capture_upstream_error_response from litellm.proxy._experimental.mcp_server.result_conversion import ( age_freshness, @@ -122,6 +124,20 @@ if TYPE_CHECKING: from litellm.proxy._experimental.mcp_server.legacy_callbacks import ElicitationCallback +def _mcp_response_hooks() -> tuple[Callable[[httpx2.Response], Awaitable[None]], ...]: + if get_active_mcp_request_ctx() is None: + return (capture_upstream_error_response,) + from litellm.proxy._experimental.mcp_server.legacy_callbacks import ( + get_scope_response, + record_upstream_tool_authorization, + ) + + return ( + partial(record_upstream_tool_authorization, pending=get_scope_response()), + capture_upstream_error_response, + ) + + def to_basic_auth(auth_value: str) -> str: """Convert auth value to Basic Auth format.""" return base64.b64encode(auth_value.encode("utf-8")).decode() @@ -497,6 +513,7 @@ class MCPClient: "sampling_callback": self._sampling_callback, "elicitation_callback": self._elicitation_callback, "logging_callback": self._logging_callback, + "response_hooks": _mcp_response_hooks(), } def session_settings(self) -> tuple[MCPTransport, MCPUpstreamProtocol, float, bool, bool, bool]: @@ -721,7 +738,9 @@ class MCPClient: write_stream, read_timeout_seconds=self.timeout, message_handler=receive_message, - **session_kwargs, + sampling_callback=session_kwargs["sampling_callback"], + elicitation_callback=session_kwargs["elicitation_callback"], + logging_callback=session_kwargs["logging_callback"], ) session: Final = await session_ctx.__aenter__() try: @@ -931,7 +950,14 @@ class MCPClient: verify=ssl_config, follow_redirects=True, event_hooks=MappingProxyType( - {"response": [capture_upstream_error_response], "request": [guard] if guard else []} + { + "response": list( + self._session_callbacks_override["response_hooks"] + if self._session_callbacks_override is not None + else _mcp_response_hooks() + ), + "request": [guard] if guard else [], + } ), ) @@ -1582,6 +1608,7 @@ class PersistentMCPSession: self._initial_callbacks: Final[_SessionCallbackMap] = client.session_callbacks() self._current_callbacks: _SessionCallbackMap = self._initial_callbacks delegated_callbacks: Final[_SessionCallbackMap] = { + "response_hooks": (self._record_response,), "sampling_callback": ( self._delegate("sampling_callback") if self._initial_callbacks["sampling_callback"] is not None @@ -1599,6 +1626,10 @@ class PersistentMCPSession: client.route_session_callbacks(delegated_callbacks) self._task: Final = loop.create_task(self._serve()) + async def _record_response(self, response: httpx2.Response) -> None: + for hook in self._current_callbacks["response_hooks"]: + await hook(response) + @overload def _delegate(self, name: Literal["sampling_callback"]) -> SamplingFnT: ... diff --git a/litellm/integrations/anthropic_cache_control_hook.py b/litellm/integrations/anthropic_cache_control_hook.py index cb0a3f0603b..8f542a87a3a 100644 --- a/litellm/integrations/anthropic_cache_control_hook.py +++ b/litellm/integrations/anthropic_cache_control_hook.py @@ -326,7 +326,7 @@ class AnthropicCacheControlHook(CustomPromptManagement): openai_dialect: Final = ( stamped_dialect if isinstance(stamped_dialect, bool) - else AnthropicCacheControlHook._targets_openai_prompt_cache_breakpoint( + else AnthropicCacheControlHook.targets_openai_prompt_cache_breakpoint( model, non_default_params.get("custom_llm_provider"), non_default_params.get("api_base") or non_default_params.get("base_url"), @@ -384,7 +384,7 @@ class AnthropicCacheControlHook(CustomPromptManagement): return model, processed_messages, non_default_params @staticmethod - def _targets_openai_prompt_cache_breakpoint( + def targets_openai_prompt_cache_breakpoint( model: str | None, custom_llm_provider: str | None, api_base: object = None, @@ -786,7 +786,7 @@ class AnthropicCacheControlHook(CustomPromptManagement): return AnthropicCacheControlHook._stamped( points, "_litellm_openai_dialect", - AnthropicCacheControlHook._targets_openai_prompt_cache_breakpoint( + AnthropicCacheControlHook.targets_openai_prompt_cache_breakpoint( model, custom_llm_provider, api_base, prompt_cache_options ), ) @@ -1103,7 +1103,7 @@ class AnthropicCacheControlHook(CustomPromptManagement): if not injection_points: return messages, system - openai_dialect: Final = AnthropicCacheControlHook._targets_openai_prompt_cache_breakpoint( + openai_dialect: Final = AnthropicCacheControlHook.targets_openai_prompt_cache_breakpoint( model, custom_llm_provider, api_base, kwargs.get("prompt_cache_options") ) breakpoints_before: Final = AnthropicCacheControlHook.count_request_cache_breakpoints(messages, system) diff --git a/litellm/integrations/otel/README.md b/litellm/integrations/otel/README.md index 4b963ca3415..340bd71c983 100644 --- a/litellm/integrations/otel/README.md +++ b/litellm/integrations/otel/README.md @@ -311,10 +311,15 @@ nothing here imports outside it: `LITELLM_OTEL_EXCLUDED_SERVICES` (comma-separated) or `excluded_services` (a YAML list) under `callback_settings.otel`, naming the datastore services to withhold (`redis`, `postgres`, `batch_write_to_db`, `redis_*`, or their - `db.system.name` spellings `redis` / `postgresql`). Unknown names are logged - as an error and ignored. A span is withheld when its `db.system.name` / + `db.system.name` spellings `redis` / `postgresql`). Unknown global names are logged + as an error and ignored. Team and key callback values are validated when saved, + and invalid values are rejected. A span is withheld when its `db.system.name` / `db.system` attribute is in the set, so request root, auth, guardrail and - model spans can never be excluded. + model spans can never be excluded. Team and key destinations can replace the + global set with `excluded_services` in their own `callback_vars`; `"none"` + means no datastore spans are excluded, while an omitted value follows the + global set. These destination overrides require OTel v2 and are accepted only + on `langfuse_otel`, `arize`, `weave_otel` and `newrelic` callbacks. - [`baggage.py`](./model/baggage.py) — the single definition of which request-identity values are promoted into Baggage (so child spans inherit them) and under which attribute keys. diff --git a/litellm/integrations/otel/model/config.py b/litellm/integrations/otel/model/config.py index 55ef28bf470..cca12b885c1 100644 --- a/litellm/integrations/otel/model/config.py +++ b/litellm/integrations/otel/model/config.py @@ -408,6 +408,23 @@ def excluded_db_systems_from(value: object) -> frozenset[str]: return _normalize_excluded_services(excluded_service_names(value)) +def invalid_excluded_service_names(value: str) -> tuple[str, ...]: + names: Final = tuple(service.strip().lower() for service in value.split(",")) + return tuple( + name + for name in names + if name != "none" and (not name or _resolved_db_system_for_excluded_service(name) is None) + ) + + +def parse_excluded_services(value: object) -> frozenset[str] | None: + if not isinstance(value, str): + return None + if value.strip().lower() == "none": + return frozenset() + return excluded_db_systems_from(value) + + def parse_capture_message_content(value: object) -> str | None: return value if isinstance(value, str) and value in CAPTURE_MESSAGE_CONTENT_VALUES else None @@ -448,9 +465,13 @@ def _normalize_excluded_services(services: frozenset[str]) -> frozenset[str]: def _db_system_for_excluded_service(service: str) -> str | None: - resolved: Final = db_system(service) if service != POSTGRESQL else POSTGRESQL + resolved: Final = _resolved_db_system_for_excluded_service(service) if resolved is None: verbose_logger.error( "excluded_services: %r is not a datastore service; ignored. Allowed: postgres, redis", service ) return resolved + + +def _resolved_db_system_for_excluded_service(service: str) -> str | None: + return db_system(service) if service != POSTGRESQL else POSTGRESQL diff --git a/litellm/integrations/otel/model/destination.py b/litellm/integrations/otel/model/destination.py index b0c22dea292..4226b4361dd 100644 --- a/litellm/integrations/otel/model/destination.py +++ b/litellm/integrations/otel/model/destination.py @@ -57,6 +57,10 @@ class OtelDestination(LiteLLMBaseModel): "Omitted follows the global setting." ), ) + excluded_db_systems: frozenset[str] | None = Field( + default=None, + description="An explicit set replaces the global database-span exclusions; ``None`` follows the global set.", + ) def captures_content(self, default: str) -> bool: setting: Final = self.capture_message_content @@ -75,9 +79,10 @@ class OtelDestination(LiteLLMBaseModel): def cache_key(self) -> tuple[str, tuple[tuple[str, str], ...], tuple[tuple[str, str], ...], str | None]: """Identity for processor reuse, so one destination means one exporter. - ``span_scope``, the sampling rates and ``capture_message_content`` are left out on - purpose: they decide which spans reach the processor and what they carry, not how - the processor exports them, so two views of the same account share one exporter. + ``span_scope``, the sampling rates, ``capture_message_content`` and + ``excluded_db_systems`` are left out on purpose: they decide which spans reach + the processor and what they carry, not how the processor exports them, so two + views of the same account share one exporter. """ return ( self.endpoint, diff --git a/litellm/integrations/otel/plumbing/providers.py b/litellm/integrations/otel/plumbing/providers.py index 4e7d4f231ed..898eab5715b 100644 --- a/litellm/integrations/otel/plumbing/providers.py +++ b/litellm/integrations/otel/plumbing/providers.py @@ -705,7 +705,12 @@ class TenantFanOutSpanProcessor(SpanProcessor): if ( self._operator_already_writes(span, destination, suppressed) or not _in_scope(span, destination.span_scope) - or _is_excluded_database_span(attributes, self._excluded_db_systems) + or _is_excluded_database_span( + attributes, + self._excluded_db_systems + if destination.excluded_db_systems is None + else destination.excluded_db_systems, + ) ): continue if not _is_sampled(destination): diff --git a/litellm/integrations/otel/presets/destinations.py b/litellm/integrations/otel/presets/destinations.py index c4273fbb337..a705fc16bbc 100644 --- a/litellm/integrations/otel/presets/destinations.py +++ b/litellm/integrations/otel/presets/destinations.py @@ -188,6 +188,7 @@ def destination_for( params: StandardCallbackDynamicParams, service_name: str | None = None, capture_message_content: str | None = None, + excluded_db_systems: frozenset[str] | None = None, ) -> OtelDestination | None: """The destination ``params`` names for ``callback_name``, or ``None``. @@ -219,6 +220,7 @@ def destination_for( protocol=protocol, span_scope=_span_scope(callback_name, params), capture_message_content=capture_message_content, + excluded_db_systems=excluded_db_systems, success_sampling_rate=success_sampling_rate, error_sampling_rate=error_sampling_rate, ) diff --git a/litellm/litellm_core_utils/cost_discount.py b/litellm/litellm_core_utils/cost_discount.py new file mode 100644 index 00000000000..a64ad0b20f7 --- /dev/null +++ b/litellm/litellm_core_utils/cost_discount.py @@ -0,0 +1,83 @@ +import fnmatch +import re +from collections.abc import Mapping +from dataclasses import dataclass +from typing import Final + +_GLOB_CHARS: Final = frozenset("*?[") +_GLOB_TOKEN: Final = re.compile( + r"\[!\][^\]]*\]|\[\][^\]]*\]|\[!(?!\])[^\]]*\]|\[(?![!\]])[^\]]*\]|.", + re.DOTALL, +) + + +@dataclass(frozen=True, slots=True) +class CostDiscountKey: + provider: str + model_pattern: str | None + + +def parse_cost_discount_key(key: str) -> CostDiscountKey: + provider, separator, pattern = key.partition("/") + return CostDiscountKey(provider=provider, model_pattern=pattern if separator else None) + + +def _literal_length(pattern: str) -> int: + return sum(1 for token in _GLOB_TOKEN.findall(pattern) if len(token) == 1 and token not in "*?") + + +def _model_names(model: str, provider_prefix: str) -> tuple[str, ...]: + prefix_run: Final = re.match(f"(?:{re.escape(provider_prefix)})*", model) + depth: Final = len(prefix_run.group()) // len(provider_prefix) if prefix_run else 0 + if depth == 0: + return (model,) + return tuple(model[len(provider_prefix) * level :] for level in range(1, depth + 1)) + + +def resolve_cost_discount( + cost_discount_config: Mapping[str, float], + custom_llm_provider: str | None, + model: str | None, + region_name: str | None = None, +) -> float | None: + if not custom_llm_provider: + return None + + provider_prefix: Final = f"{custom_llm_provider}/" + region_prefix: Final = f"{provider_prefix}{region_name}/" if region_name else "" + model_names: Final = ( + () + if model is None + else _model_names(model, provider_prefix) + + ( + _model_names(f"{provider_prefix}{model[len(region_prefix) :]}", provider_prefix) + if region_prefix and model.startswith(region_prefix) + else () + ) + ) + patterns: Final = tuple( + parsed.model_pattern + for parsed in (parse_cost_discount_key(key) for key in cost_discount_config) + if parsed.provider == custom_llm_provider and parsed.model_pattern is not None + ) + + exact_patterns: Final = frozenset(pattern for pattern in patterns if _GLOB_CHARS.isdisjoint(pattern)) + exact: Final = next((name for name in model_names if name in exact_patterns), None) + if exact is not None: + return cost_discount_config[f"{custom_llm_provider}/{exact}"] + prefixed_markers: Final = (provider_prefix,) if not region_name else (provider_prefix, f"{region_name}/") + matches: Final = tuple( + pattern + for pattern in patterns + if not _GLOB_CHARS.isdisjoint(pattern) + and any( + fnmatch.fnmatchcase(name, pattern) + for name in model_names + if "/" in pattern or not name.startswith(prefixed_markers) + ) + ) + if matches: + best_match: Final = max(matches, key=_literal_length) + return cost_discount_config[f"{custom_llm_provider}/{best_match}"] + + return cost_discount_config.get(custom_llm_provider) diff --git a/litellm/litellm_core_utils/decisions_pricing.py b/litellm/litellm_core_utils/decisions_pricing.py new file mode 100644 index 00000000000..9199a8097cc --- /dev/null +++ b/litellm/litellm_core_utils/decisions_pricing.py @@ -0,0 +1,33 @@ +"""Decisions-specific pricing read from ``_decisions`` cost map fields. + +A model that serves both chat and Decisions keeps one cost map row. Rates that +differ for ``/v1/decisions`` and ``/v1/systemone`` live on that row under the +same field name with a ``_decisions`` suffix, e.g. ``output_cost_per_token_decisions``, +and override the plain field for those two routes only. A row without any such +field bills Decisions calls from its plain fields. +""" + +from collections.abc import Mapping +from typing import Final + +DECISIONS_PRICING_SUFFIX: Final = "_decisions" +DECISIONS_CALL_TYPES: Final = frozenset({"decisions", "adecisions", "systemone", "asystemone"}) + + +def is_decisions_call_type(call_type: str) -> bool: + return call_type in DECISIONS_CALL_TYPES + + +def decisions_pricing_fields(row: Mapping[str, object]) -> Mapping[str, object]: + """Return ``{plain_field: value}`` for every ``_decisions`` field on ``row``.""" + return { + key.removesuffix(DECISIONS_PRICING_SUFFIX): value + for key, value in row.items() + if key.endswith(DECISIONS_PRICING_SUFFIX) + } + + +def apply_decisions_pricing(row: Mapping[str, object]) -> Mapping[str, object]: + """Return ``row`` priced for a Decisions call: ``_decisions`` values replace their plain fields.""" + plain: Final = {key: value for key, value in row.items() if not key.endswith(DECISIONS_PRICING_SUFFIX)} + return {**plain, **decisions_pricing_fields(row)} diff --git a/litellm/litellm_core_utils/exception_mapping_utils.py b/litellm/litellm_core_utils/exception_mapping_utils.py index b875f7c09a6..014f0cae54f 100644 --- a/litellm/litellm_core_utils/exception_mapping_utils.py +++ b/litellm/litellm_core_utils/exception_mapping_utils.py @@ -880,21 +880,37 @@ def _map_openai_like_exception( _BEDROCK_MANTLE_CONTEXT_WINDOW_PATTERN: Final = re.compile(r"prompt tokens \((\d+)\) exceed model maximum \((\d+)\)") +_BEDROCK_MANTLE_CONTEXT_LENGTH_EXCEEDED_CODE_PATTERN: Final = re.compile( + r"""\\?["']code\\?["']\s*:\s*\\?["']context_length_exceeded\\?["']""" +) +_BEDROCK_MANTLE_CONTEXT_WINDOW_GENERIC_MESSAGE: Final = ( + "prompt is too long: your prompt exceeds the model's context window" +) + + +def _bedrock_mantle_context_window_response(original_exception: _ProviderHTTPException) -> httpx.Response: + original_response: Final = getattr(original_exception, "response", None) + if not isinstance(original_response, httpx.Response) or getattr(original_response, "_request", None) is None: + return httpx.Response(status_code=400, request=httpx.Request(method="POST", url="https://api.openai.com/v1/")) + if original_response.status_code == 400: + return original_response + return httpx.Response(status_code=400, headers=original_response.headers, request=original_response.request) def _get_bedrock_mantle_context_window_message(error_str: str) -> str | None: """ - Mantle reports context overflow as a structured validation error rather than - the plain-text patterns Bedrock itself uses, so it needs its own detection and a - message clients recognize as context overflow (litellm/litellm#36546). + Mantle reports context overflow as a validation_error carrying the token counts, or as + OpenAI's context_length_exceeded code, either in a 400 body or in a streamed error event + with no error type. Clients such as Claude Code only treat "prompt is too long" as + overflow (litellm/litellm#36546). """ - if "invalid_request_error" not in error_str and "validation_error" not in error_str: - return None match = _BEDROCK_MANTLE_CONTEXT_WINDOW_PATTERN.search(error_str) - if match is None: - return None - prompt_tokens, max_tokens = match.groups() - return f"prompt is too long: {prompt_tokens} tokens > {max_tokens} maximum" + if match is not None: + prompt_tokens, max_tokens = match.groups() + return f"prompt is too long: {prompt_tokens} tokens > {max_tokens} maximum" + if _BEDROCK_MANTLE_CONTEXT_LENGTH_EXCEEDED_CODE_PATTERN.search(error_str) is not None: + return _BEDROCK_MANTLE_CONTEXT_WINDOW_GENERIC_MESSAGE + return None def _map_bedrock_exception( @@ -914,7 +930,7 @@ def _map_bedrock_exception( message=mantle_context_window_message, model=model, llm_provider=custom_llm_provider, - response=getattr(original_exception, "response", None), + response=_bedrock_mantle_context_window_response(original_exception), ) if ( "too many tokens" in error_str diff --git a/litellm/litellm_core_utils/initialize_dynamic_callback_params.py b/litellm/litellm_core_utils/initialize_dynamic_callback_params.py index 2efeead6bb2..5a533054f0b 100644 --- a/litellm/litellm_core_utils/initialize_dynamic_callback_params.py +++ b/litellm/litellm_core_utils/initialize_dynamic_callback_params.py @@ -99,6 +99,18 @@ def validate_capture_message_content_value(value: str) -> None: ) +def validate_excluded_services_value(value: str) -> None: + from litellm.integrations.otel.model.config import invalid_excluded_service_names + + names: Final = tuple(name.strip().lower() for name in value.split(",")) + invalid_names: Final = invalid_excluded_service_names(value) + allowed: Final = "redis, postgres, postgresql, batch_write_to_db, redis_-prefixed service names, or none" + if invalid_names or ("none" in names and len(names) != 1): + raise ValueError( + f"Invalid excluded_services {value!r}: use {allowed}; 'none' must be used alone and names cannot be empty" + ) + + # Hardcoded list of supported callback params to avoid runtime inspection issues with TypedDict _supported_callback_params: Final[tuple[str, ...]] = ( "langfuse_public_key", diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index d6c97b120c9..5395383f8bd 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -5,6 +5,7 @@ import copy import datetime import functools import json +import math import os import re import subprocess @@ -732,6 +733,7 @@ class Logging(LiteLLMLoggingBaseClass): # Initialize cost breakdown field self.cost_breakdown: CostBreakdown | None = None self.billed_token_rates: BilledTokenRates | None = None + self._provider_reported_cost: float | None = None # Init Caching related details self.caching_details: CachingDetails | None = None @@ -1948,7 +1950,16 @@ class Logging(LiteLLMLoggingBaseClass): return None try: - response_cost: Final = litellm.response_cost_calculator(**response_cost_calculator_kwargs) + reported_cost: Final = self._provider_reported_cost + response_cost: Final = ( + reported_cost + if isinstance(reported_cost, float) + and math.isfinite(reported_cost) + and reported_cost >= 0 + and not custom_pricing + and base_model is None + else litellm.response_cost_calculator(**response_cost_calculator_kwargs) + ) verbose_logger.debug("response_cost: %s", response_cost) additional_response_cost: Final[object] = self.model_call_details.get("additional_response_cost") diff --git a/litellm/litellm_core_utils/llm_cost_calc/utils.py b/litellm/litellm_core_utils/llm_cost_calc/utils.py index 73428e08e34..2afb6ed4c58 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/utils.py +++ b/litellm/litellm_core_utils/llm_cost_calc/utils.py @@ -1269,6 +1269,13 @@ def get_provider_specific_geo_multiplier(model_info: ModelInfo, usage: Usage) -> return float(provider_specific_entry.get(inference_geo.lower(), 1.0)) +def get_provider_specific_speed_multiplier(model_info: ModelInfo, usage: Usage) -> float: + provider_specific_entry: Final = model_info.get("provider_specific_entry") + if getattr(usage, "speed", None) != "fast" or not provider_specific_entry: + return 1.0 + return float(provider_specific_entry.get("fast", 1.0)) + + def _resolve_reasoning_token_cost( model_info: ModelInfo, service_tier: str | None, @@ -1633,6 +1640,7 @@ def _cost_map_billed_rates( get_regional_uplift_multiplier(model_info, data_residency) * get_vertex_regional_endpoint_uplift(model_info, vertex_location) * get_provider_specific_geo_multiplier(model_info=model_info, usage=usage) + * get_provider_specific_speed_multiplier(model_info=model_info, usage=usage) ) return BilledTokenRates( input_cost_per_token=prompt_base_cost, diff --git a/litellm/litellm_core_utils/prompt_templates/common_utils.py b/litellm/litellm_core_utils/prompt_templates/common_utils.py index 0243141eda2..591cae6dfa4 100644 --- a/litellm/litellm_core_utils/prompt_templates/common_utils.py +++ b/litellm/litellm_core_utils/prompt_templates/common_utils.py @@ -29,6 +29,7 @@ from litellm.types.llms.openai import ( ChatCompletionThinkingBlock, ChatCompletionToolParam, ChatCompletionUserMessage, + PromptCacheBreakpoint, ) from litellm.types.utils import ( Choices, @@ -1793,6 +1794,16 @@ def with_prompt_cache_breakpoint(target: _MarkedT, marker: object) -> _MarkedT: return cast(_MarkedT, marked) # cast-ok: same block shape as the input plus the marker key +def cache_control_as_prompt_cache_breakpoint(block: _MarkedT) -> _MarkedT: + """Anthropic's ``cache_control`` mark rewritten as OpenAI's ``prompt_cache_breakpoint`` on the same block.""" + if block.get("cache_control") is None: + return block + unmarked: Final = {key: value for key, value in block.items() if key != "cache_control"} + return cast( # cast-ok: same block shape as the input with the marker key renamed + _MarkedT, with_prompt_cache_breakpoint(unmarked, PromptCacheBreakpoint(mode="explicit")) + ) + + LITELLM_INTERNAL_MESSAGE_FIELDS: Final = frozenset({"thinking_blocks", "reasoning_content", "provider_specific_fields"}) diff --git a/litellm/litellm_core_utils/prompt_templates/factory.py b/litellm/litellm_core_utils/prompt_templates/factory.py index 5d6c1ac0e4b..7d635712e42 100644 --- a/litellm/litellm_core_utils/prompt_templates/factory.py +++ b/litellm/litellm_core_utils/prompt_templates/factory.py @@ -5093,7 +5093,7 @@ def make_valid_bedrock_tool_name(input_tool_name: str) -> str: def add_cache_point_tool_block(tool: dict, model: str | None = None) -> BedrockToolBlock | None: from litellm.llms.bedrock.common_utils import ( bedrock_model_accepts_cache_points, - is_claude_4_5_on_bedrock, + supports_bedrock_prompt_cache_ttl, ) cache_control: Final = tool.get("cache_control", None) @@ -5103,7 +5103,7 @@ def add_cache_point_tool_block(tool: dict, model: str | None = None) -> BedrockT cache_point_block: Final[CachePointBlock] = {"type": "default"} if isinstance(cache_control, dict) and "ttl" in cache_control: ttl: Final = cache_control["ttl"] - if ttl in ["5m", "1h"] and model is not None and is_claude_4_5_on_bedrock(model): + if ttl in ["5m", "1h"] and model is not None and supports_bedrock_prompt_cache_ttl(model): cache_point_block["ttl"] = ttl return {"cachePoint": cache_point_block} return None diff --git a/litellm/llms/__init__.py b/litellm/llms/__init__.py index 88a44f38c57..7452a4a20ef 100644 --- a/litellm/llms/__init__.py +++ b/litellm/llms/__init__.py @@ -62,6 +62,15 @@ def get_cost_for_web_search_request(custom_llm_provider: str, usage: "Usage", mo ) return cost_per_web_search_request_vertex_ai(usage=usage, model_info=model_info) + elif custom_llm_provider == "azure_ai": + from litellm.litellm_core_utils.llm_cost_calc.utils import get_web_search_requests_from_usage + + azure_ai_model_key: Final[str] = model_info.get("key", "") if model_info else "" + if "claude" not in azure_ai_model_key.lower() and get_web_search_requests_from_usage(usage) is None: + return None + from .anthropic.cost_calculation import get_cost_for_anthropic_web_search + + return get_cost_for_anthropic_web_search(model_info=model_info, usage=usage) elif custom_llm_provider == "perplexity": # Perplexity handles search costs internally in its own cost calculator # Return 0.0 to indicate costs are already accounted for diff --git a/litellm/llms/anthropic/cost_calculation.py b/litellm/llms/anthropic/cost_calculation.py index 88805df2b08..94c8149eff3 100644 --- a/litellm/llms/anthropic/cost_calculation.py +++ b/litellm/llms/anthropic/cost_calculation.py @@ -10,6 +10,7 @@ from pydantic import ValidationError from litellm.litellm_core_utils.llm_cost_calc.utils import ( generic_cost_per_token, get_provider_specific_geo_multiplier, + get_provider_specific_speed_multiplier, get_web_search_requests_from_usage, ) from litellm.types.llms.base import LiteLLMBaseModel @@ -50,14 +51,8 @@ def cost_per_token( if model_info is not None else litellm.get_model_info(model=model, custom_llm_provider="anthropic") ) - provider_specific_entry: Final = effective_info.get("provider_specific_entry") - geo_multiplier: Final = get_provider_specific_geo_multiplier(model_info=effective_info, usage=usage) - speed_multiplier: Final = ( - provider_specific_entry.get("fast", 1.0) - if provider_specific_entry and getattr(usage, "speed", None) == "fast" - else 1.0 - ) + speed_multiplier: Final = get_provider_specific_speed_multiplier(model_info=effective_info, usage=usage) if speed_multiplier != 1.0: prompt_cost *= speed_multiplier diff --git a/litellm/llms/anthropic/pass_through/adapters/streaming_iterator.py b/litellm/llms/anthropic/pass_through/adapters/streaming_iterator.py index 5f6d644853a..34c18b4f0bd 100644 --- a/litellm/llms/anthropic/pass_through/adapters/streaming_iterator.py +++ b/litellm/llms/anthropic/pass_through/adapters/streaming_iterator.py @@ -19,7 +19,7 @@ from typing_extensions import assert_never from litellm._logging import verbose_logger from litellm._uuid import uuid -from litellm.exceptions import MidStreamFallbackError +from litellm.exceptions import ContextWindowExceededError, MidStreamFallbackError from litellm.litellm_core_utils.hidden_params import set_hidden_params from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.types.llms.anthropic import ( @@ -64,7 +64,7 @@ def _optional_attr_sequence(obj: object, name: str) -> Sequence[object]: def _error_status_and_message(exc: Exception) -> tuple[int, str]: - if isinstance(exc, (BaseLLMException, MidStreamFallbackError)): + if isinstance(exc, (BaseLLMException, MidStreamFallbackError, ContextWindowExceededError)): return exc.status_code, exc.message return 500, str(exc) or "Upstream stream ended before completion" diff --git a/litellm/llms/anthropic/pass_through/responses_adapters/handler.py b/litellm/llms/anthropic/pass_through/responses_adapters/handler.py index f20762789dd..5c5fb71c295 100644 --- a/litellm/llms/anthropic/pass_through/responses_adapters/handler.py +++ b/litellm/llms/anthropic/pass_through/responses_adapters/handler.py @@ -4,10 +4,14 @@ Handler for the Anthropic v1/messages -> OpenAI Responses API path. Used when the target model is an OpenAI or Azure model. """ -from collections.abc import AsyncIterator, Coroutine, Mapping +from collections.abc import AsyncIterator, Coroutine, Mapping, Sequence from typing import Final, TypeAlias +from pydantic import TypeAdapter, ValidationError + import litellm +from litellm.integrations.anthropic_cache_control_hook import AnthropicCacheControlHook +from litellm.litellm_core_utils.prompt_templates.common_utils import cache_control_as_prompt_cache_breakpoint from litellm.types.integrations.custom_logger import CONVERTED_STREAM_KEYS from litellm.types.llms.anthropic import ( AllAnthropicMessageValues, @@ -27,6 +31,7 @@ from .streaming_iterator import AnthropicResponsesStreamWrapper from .transformation import LiteLLMAnthropicToResponsesAPIAdapter AnthropicRequestMessages: TypeAlias = list[AllAnthropicMessageValues] | list[dict[str, object]] +_BLOCK_LIST: Final = TypeAdapter(list[dict[str, object]]) _ADAPTER: Final = LiteLLMAnthropicToResponsesAPIAdapter() @@ -45,6 +50,24 @@ def _provider_returns_encrypted_reasoning(model: str, custom_llm_provider: objec return responses_config is not None and "include" in responses_config.get_supported_openai_params(provider_model) +def _reads_prompt_cache_breakpoint(model: str, forwarded_kwargs: Mapping[str, object]) -> bool: + provider: Final = forwarded_kwargs.get("custom_llm_provider") + return AnthropicCacheControlHook.targets_openai_prompt_cache_breakpoint( + model, + provider if isinstance(provider, str) else None, + forwarded_kwargs.get("api_base") or forwarded_kwargs.get("base_url"), + forwarded_kwargs.get("prompt_cache_options"), + ) + + +def _blocks_with_prompt_cache_breakpoints(content: object) -> Sequence[Mapping[str, object]]: + """The content's blocks with each ``cache_control`` mark rewritten, empty when it is not a list of blocks.""" + try: + return [cache_control_as_prompt_cache_breakpoint(block) for block in _BLOCK_LIST.validate_python(content)] + except ValidationError: + return () + + def _build_responses_kwargs( *, max_tokens: int, @@ -68,10 +91,28 @@ def _build_responses_kwargs( """ Build the kwargs dict to pass directly to litellm.responses() / litellm.aresponses(). """ + forwarded_kwargs: Final = _forwarded_kwargs(extra_kwargs) + reads_prompt_cache_breakpoint: Final = _reads_prompt_cache_breakpoint(model, forwarded_kwargs) + request_messages: Final[AnthropicRequestMessages] = ( + [ + { + **message, + "content": list(_blocks_with_prompt_cache_breakpoints(message.get("content"))) + or message.get("content"), + } + for message in messages + ] + if reads_prompt_cache_breakpoint + else messages + ) + request_system: Final = ( + (list(_blocks_with_prompt_cache_breakpoints(system)) or system) if reads_prompt_cache_breakpoint else system + ) + # Build a typed AnthropicMessagesRequest for the adapter request_data: Final[AnthropicMessagesRequest] = { "model": model, - "messages": messages, + "messages": request_messages, "max_tokens": max_tokens, } if context_management: @@ -80,8 +121,8 @@ def _build_responses_kwargs( request_data["output_config"] = output_config if metadata: request_data["metadata"] = metadata - if system: - request_data["system"] = system + if request_system: + request_data["system"] = request_system if temperature is not None: request_data["temperature"] = temperature if thinking: @@ -96,7 +137,6 @@ def _build_responses_kwargs( request_data["output_format"] = output_format anthropic_request: Final = AnthropicMessagesRequest(**request_data) - forwarded_kwargs: Final = _forwarded_kwargs(extra_kwargs) responses_kwargs: Final = _ADAPTER.translate_request( anthropic_request, include_encrypted_reasoning=_provider_returns_encrypted_reasoning( diff --git a/litellm/llms/anthropic/pass_through/responses_adapters/transformation.py b/litellm/llms/anthropic/pass_through/responses_adapters/transformation.py index 642f1857058..e033f1562cd 100644 --- a/litellm/llms/anthropic/pass_through/responses_adapters/transformation.py +++ b/litellm/llms/anthropic/pass_through/responses_adapters/transformation.py @@ -131,12 +131,18 @@ class LiteLLMAnthropicToResponsesAPIAdapter: def _tool_result_output_value( output_text: str, file_parts: tuple[dict[str, str], ...], # mutable-ok: json content parts + tool_result: Mapping[str, object], ) -> str | list[dict[str, str]]: # mutable-ok: API message payload - """Plain string output, or a part list when document file parts are present.""" - if not file_parts: - return output_text + """Plain string output, or a part list when document file parts or a cache breakpoint are present. + + OpenAI reads a breakpoint only on a part, so a marked result ends on a part that carries it. + """ + prompt_cache_breakpoint: Final = tool_result.get("prompt_cache_breakpoint") text_parts: Final = [{"type": "input_text", "text": output_text}] if output_text else [] - return [*text_parts, *file_parts] + parts: Final = [*text_parts, *file_parts] + if not parts or (not file_parts and prompt_cache_breakpoint is None): + return output_text + return [*parts[:-1], with_prompt_cache_breakpoint(parts[-1], prompt_cache_breakpoint)] @staticmethod def _translate_midturn_system_content_to_responses( @@ -327,7 +333,7 @@ class LiteLLMAnthropicToResponsesAPIAdapter: { "type": "function_call_output", "call_id": tool_use_id, - "output": self._tool_result_output_value(output_text, tool_file_parts), + "output": self._tool_result_output_value(output_text, tool_file_parts, block), } ) if tool_image_parts: diff --git a/litellm/llms/base_llm/decisions/transformation.py b/litellm/llms/base_llm/decisions/transformation.py index 724ca23c2ab..5bdb6ee3c18 100644 --- a/litellm/llms/base_llm/decisions/transformation.py +++ b/litellm/llms/base_llm/decisions/transformation.py @@ -8,6 +8,7 @@ import httpx from pydantic import TypeAdapter, ValidationError from typing_extensions import assert_never +import litellm from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.secret_managers.main import get_secret_str from litellm.types.decisions import ( @@ -47,6 +48,7 @@ from litellm.types.decisions import ( UnsupportedDecisionsRequest, systemone_choice_key, ) +from litellm.types.utils import is_decisions_model_mode _PAYLOAD_ADAPTER: Final[TypeAdapter[object]] = TypeAdapter(object) _SYSTEMONE_RESPONSE_ADAPTER: Final[TypeAdapter[DecisionsResponse]] = TypeAdapter(DecisionsResponse) @@ -100,6 +102,8 @@ def _ir_question(name: str, question: DecisionQuestion) -> DecisionsIRQuestion: def systemone_request_to_ir(request: DecisionsRequestBody) -> DecisionsIRRequest: return DecisionsIRRequest( input=DecisionsIRState(state=request.state), + audio=request.audio, + videos=request.videos, questions=tuple(_ir_question(name, question) for name, question in request.questions.items()), ) @@ -166,8 +170,12 @@ def _systemone_question(question: DecisionsIRQuestion) -> Mapping[str, object]: def ir_to_systemone_request( - model: str, request: DecisionsIRRequest + model: str, request: DecisionsIRRequest, *, supports_audio: bool = False, supports_video: bool = False ) -> Mapping[str, object] | UnsupportedDecisionsRequest: + if request.audio and not supports_audio: + return UnsupportedDecisionsRequest(reason="audio input is not supported by this Decisions model") + if request.videos and not supports_video: + return UnsupportedDecisionsRequest(reason="video input is not supported by this Decisions model") state: Final = _systemone_state(request.input) if isinstance(state, UnsupportedDecisionsRequest): return state @@ -175,6 +183,16 @@ def ir_to_systemone_request( return { "model": model, "state": state, + **( + {"audio": [clip if isinstance(clip, str) else clip.model_dump() for clip in request.audio]} + if request.audio + else {} + ), + **( + {"videos": [clip if isinstance(clip, str) else clip.model_dump() for clip in request.videos]} + if request.videos + else {} + ), "questions": {key: _systemone_question(question) for key, question in keyed_questions}, } @@ -296,6 +314,14 @@ def ir_to_systemone_response(response: DecisionsIRResponse, request: DecisionsIR ) +def _is_decisions_model(model: str, custom_llm_provider: str) -> bool: + try: + model_mode: Final = litellm.get_model_info(model=model, custom_llm_provider=custom_llm_provider).get("mode") + except Exception: # noqa: BLE001 # get_model_info raises bare Exception for an unmapped model, which has no mode + return False + return is_decisions_model_mode(model_mode) + + class BaseDecisionsConfig(ABC): path: str = "/v1/systemone" api_key_env: tuple[str, ...] = () @@ -355,13 +381,28 @@ class BaseDecisionsConfig(ABC): ) -> tuple[Mapping[str, str], bytes | None]: return headers, None + def supports_audio_input(self, model: str, custom_llm_provider: str) -> bool: + return _is_decisions_model(model, custom_llm_provider) and litellm.supports_audio_input( + model=model, custom_llm_provider=custom_llm_provider + ) + + def supports_video_input(self, model: str, custom_llm_provider: str) -> bool: + return _is_decisions_model(model, custom_llm_provider) and litellm.supports_video_input( + model=model, custom_llm_provider=custom_llm_provider + ) + def transform_decisions_request( self, model: str, request: DecisionsIRRequest, custom_llm_provider: str, ) -> Mapping[str, object] | UnsupportedDecisionsRequest: - return ir_to_systemone_request(self.request_model(model), request) + return ir_to_systemone_request( + self.request_model(model), + request, + supports_audio=self.supports_audio_input(model, custom_llm_provider), + supports_video=self.supports_video_input(model, custom_llm_provider), + ) def unwrap_response(self, payload: object) -> object: return payload diff --git a/litellm/llms/bedrock/chat/converse_transformation.py b/litellm/llms/bedrock/chat/converse_transformation.py index d5a54bb1e41..1cdcf0d6e4c 100644 --- a/litellm/llms/bedrock/chat/converse_transformation.py +++ b/litellm/llms/bedrock/chat/converse_transformation.py @@ -110,8 +110,8 @@ from ..common_utils import ( get_anthropic_beta_from_headers, get_bedrock_tool_name, is_bedrock_application_inference_profile_arn, - is_claude_4_5_on_bedrock, normalize_bedrock_opus_output_config_effort, + supports_bedrock_prompt_cache_ttl, ) if TYPE_CHECKING: @@ -1385,7 +1385,7 @@ class AmazonConverseConfig(BaseConfig): cache_point: Final = CachePointBlock(type="default") if isinstance(control, dict) and "ttl" in control: ttl: Final = control["ttl"] - if ttl in ["5m", "1h"] and model is not None and is_claude_4_5_on_bedrock(model): + if ttl in ["5m", "1h"] and model is not None and supports_bedrock_prompt_cache_ttl(model): cache_point["ttl"] = ttl return cache_point diff --git a/litellm/llms/bedrock/common_utils.py b/litellm/llms/bedrock/common_utils.py index 151f4858fdb..9722f9dccb2 100644 --- a/litellm/llms/bedrock/common_utils.py +++ b/litellm/llms/bedrock/common_utils.py @@ -9,6 +9,7 @@ import functools import json import os import re +import warnings from collections.abc import Iterator, Mapping, Sequence from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, TypeAlias @@ -1230,15 +1231,11 @@ def bedrock_supports_tool_search(model: str) -> bool: return AnthropicModelInfo.supports_model_capability(model, "supports_tool_search", "bedrock") -def is_claude_4_5_on_bedrock(model: str) -> bool: +def supports_bedrock_prompt_cache_ttl(model: str) -> bool: """ - Check if the model supports Bedrock prompt caching with an extended '1h' TTL - (in addition to the default 5m TTL). - - Backed by the ``cache_creation_input_token_cost_above_1hr`` field in - ``model_prices_and_context_window.json`` instead of a hardcoded list of - model-name patterns, so newly released models pick up support as soon as - their pricing entry ships, with no code change required here. + Whether Bedrock accepts a ``ttl`` ("5m"/"1h") on prompt ``cache_control`` + for this model, decided by ``cache_creation_input_token_cost_above_1hr`` on + the model's or its base model's cost-map entry. """ return any( (litellm.model_cost.get(candidate) or {}).get("cache_creation_input_token_cost_above_1hr") is not None @@ -1246,6 +1243,15 @@ def is_claude_4_5_on_bedrock(model: str) -> bool: ) +def is_claude_4_5_on_bedrock(model: str) -> bool: + warnings.warn( + "is_claude_4_5_on_bedrock is deprecated, use supports_bedrock_prompt_cache_ttl", + UserWarning, + stacklevel=2, + ) + return supports_bedrock_prompt_cache_ttl(model) + + _BEDROCK_MODEL_VERSION_SUFFIX_RE: Final = re.compile(r"-v\d+(?::\d+)?$") _DEPLOYMENT_MODEL_INFO: Final = TypeAdapter(dict[str, object]) diff --git a/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py b/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py index 26a0302ef61..a89a6c97249 100644 --- a/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py +++ b/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py @@ -34,11 +34,11 @@ from litellm.llms.bedrock.common_utils import ( bedrock_supports_tool_search, ensure_bedrock_anthropic_messages_tool_names, get_anthropic_beta_from_headers, - is_claude_4_5_on_bedrock, normalize_bedrock_opus_output_config_effort, normalize_custom_field_on_tools, normalize_tool_input_schema_types_for_bedrock_invoke, strip_unsupported_bedrock_invoke_output_config_keys, + supports_bedrock_prompt_cache_ttl, tools_without_eager_input_streaming, ) from litellm.llms.bedrock.request_metadata import ( @@ -169,9 +169,9 @@ class AmazonAnthropicClaudeMessagesConfig( anthropic_messages_request: The request dictionary to modify in-place model: The model name to check if it supports ttl """ - is_claude_4_5 = False + supports_ttl = False if model: - is_claude_4_5 = self._is_claude_4_5_on_bedrock(model) + supports_ttl = self._supports_prompt_cache_ttl(model) def _sanitize_cache_control(cache_control: dict) -> None: if not isinstance(cache_control, dict): @@ -181,7 +181,7 @@ class AmazonAnthropicClaudeMessagesConfig( # Remove ttl for models that don't support it if "ttl" in cache_control: ttl: Final = cache_control["ttl"] - if is_claude_4_5 and ttl in ["5m", "1h"]: + if supports_ttl and ttl in ["5m", "1h"]: return cache_control.pop("ttl", None) @@ -386,19 +386,9 @@ class AmazonAnthropicClaudeMessagesConfig( ] return any(pattern in model_lower for pattern in opus_4_5_patterns) - def _is_claude_4_5_on_bedrock(self, model: str) -> bool: - """ - Check if the model is Claude 4.5 on Bedrock. - - Claude Sonnet 4.5, Haiku 4.5, and Opus 4.5 support 1-hour prompt caching. - - Args: - model: The model name - - Returns: - True if the model is Claude 4.5 - """ - return is_claude_4_5_on_bedrock(model) + def _supports_prompt_cache_ttl(self, model: str) -> bool: + """Whether Bedrock accepts a ``ttl`` on prompt ``cache_control`` for ``model``.""" + return supports_bedrock_prompt_cache_ttl(model) def _supports_tool_search_on_bedrock(self, model: str) -> bool: """ diff --git a/litellm/llms/databricks/chat/transformation.py b/litellm/llms/databricks/chat/transformation.py index 50d3b82a989..dceec658824 100644 --- a/litellm/llms/databricks/chat/transformation.py +++ b/litellm/llms/databricks/chat/transformation.py @@ -285,6 +285,7 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig): "tool_choice", "reasoning_effort", "thinking", + "service_tier", ] @staticmethod @@ -697,6 +698,9 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig): model_response.id = completion_response["id"] model_response.created = completion_response["created"] setattr(model_response, "usage", Usage(**completion_response["usage"])) + served_tier: Final = completion_response.get("service_tier") + if isinstance(served_tier, str) and served_tier: + setattr(model_response, "service_tier", served_tier) # noqa: B010 # no declared field model_response.choices = self._transform_dbrx_choices( choices=completion_response["choices"], diff --git a/litellm/llms/fireworks_ai/decisions/transformation.py b/litellm/llms/fireworks_ai/decisions/transformation.py new file mode 100644 index 00000000000..87d402902be --- /dev/null +++ b/litellm/llms/fireworks_ai/decisions/transformation.py @@ -0,0 +1,20 @@ +from dataclasses import replace +from typing import Final + +from litellm.llms.base_llm.decisions.transformation import BaseDecisionsConfig, parse_systemone_response +from litellm.types.decisions import DecisionsIRRequest, DecisionsIRResponse + + +class FireworksAIDecisionsConfig(BaseDecisionsConfig): + api_key_env = ("FIREWORKS_API_KEY", "FIREWORKS_AI_API_KEY") + api_base_env = ("FIREWORKS_API_BASE",) + + def get_default_api_base(self) -> str | None: + return "https://api.fireworks.ai/inference/v1" + + def parse_response(self, payload: object, request: DecisionsIRRequest) -> DecisionsIRResponse: + response: Final = parse_systemone_response(self.unwrap_response(payload), request) + cached: Final = response.usage.extra.get("cached_input_tokens") + if isinstance(cached, bool) or not isinstance(cached, int): + return response + return replace(response, usage=replace(response.usage, cached_tokens=cached)) diff --git a/litellm/llms/openai/decisions/transformation.py b/litellm/llms/openai/decisions/transformation.py index 317e5abeeb5..eb26047d323 100644 --- a/litellm/llms/openai/decisions/transformation.py +++ b/litellm/llms/openai/decisions/transformation.py @@ -88,6 +88,8 @@ def _ir_question(question: OpenAIDecisionQuestion) -> DecisionsIRQuestion: def openai_request_to_ir(request: OpenAIDecisionRequestBody) -> DecisionsIRRequest: return DecisionsIRRequest( input=_ir_input(request.input), + audio=request.audio, + videos=request.videos, questions=tuple(_ir_question(question) for question in request.questions), safety_identifier=request.safety_identifier, ) @@ -164,6 +166,8 @@ def _has_one_option(question: DecisionsIRQuestion) -> bool: def ir_to_openai_request(model: str, request: DecisionsIRRequest) -> Mapping[str, object] | UnsupportedDecisionsRequest: + if request.audio or request.videos: + return UnsupportedDecisionsRequest(reason="audio/video input is not supported by this Decisions provider") if any(_has_one_option(question) for question in request.questions): return _SINGLE_OPTION return { @@ -309,6 +313,12 @@ class OpenAIDecisionsConfig(BaseDecisionsConfig): def resolve_api_key(self, api_key: str | None) -> str | None: return api_key or litellm.api_key or litellm.openai_key or get_secret_str("OPENAI_API_KEY") + def supports_audio_input(self, model: str, custom_llm_provider: str) -> bool: + return False + + def supports_video_input(self, model: str, custom_llm_provider: str) -> bool: + return False + def transform_decisions_request( self, model: str, diff --git a/litellm/llms/openai_like/providers.json b/litellm/llms/openai_like/providers.json index f3aa7768928..d25257b0e81 100644 --- a/litellm/llms/openai_like/providers.json +++ b/litellm/llms/openai_like/providers.json @@ -180,6 +180,12 @@ "api_key_env": "COGNITION_API_KEY", "api_base_env": "COGNITION_API_BASE" }, + "coralbricks": { + "base_url": "https://inference.coralbricks.ai/v1", + "api_key_env": "CORALBRICKS_API_KEY", + "api_base_env": "CORALBRICKS_API_BASE", + "supported_endpoints": ["/v1/chat/completions", "/v1/responses", "/v1/messages"] + }, "cortecs": { "base_url": "https://api.cortecs.ai/v1", "api_key_env": "CORTECS_API_KEY", diff --git a/litellm/llms/pass_through/typesafe_logging_handler.py b/litellm/llms/pass_through/typesafe_logging_handler.py index c24915c7126..397ec408ae9 100644 --- a/litellm/llms/pass_through/typesafe_logging_handler.py +++ b/litellm/llms/pass_through/typesafe_logging_handler.py @@ -6,6 +6,7 @@ import httpx from pydantic import TypeAdapter, ValidationError import litellm +from litellm.litellm_core_utils.decisions_pricing import apply_decisions_pricing from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.litellm_logging import ( get_standard_logging_object_payload, # pyright: ignore[reportUnknownVariableType] # legacy helper has an untyped signature @@ -51,7 +52,9 @@ def _pricing_for(model_keys: tuple[str, ...]) -> _RegistryPricing: continue try: return _REGISTRY_PRICING_ADAPTER.validate_python( - litellm.model_cost[model_key] # pyright: ignore[reportUnknownMemberType] # registry is dynamically typed + apply_decisions_pricing( + litellm.model_cost[model_key] # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # registry is dynamically typed + ) ) except ValidationError: continue diff --git a/litellm/llms/sagemaker/common_utils.py b/litellm/llms/sagemaker/common_utils.py index 7772ca6869a..579309035f0 100644 --- a/litellm/llms/sagemaker/common_utils.py +++ b/litellm/llms/sagemaker/common_utils.py @@ -4,7 +4,7 @@ from collections.abc import AsyncIterator, Iterator, Mapping from typing import Final import httpx -from pydantic import ConfigDict, RootModel, TypeAdapter, with_config +from pydantic import BaseModel, ConfigDict, Field, RootModel, TypeAdapter, ValidationError, with_config from typing_extensions import ReadOnly, Required, TypedDict import litellm @@ -13,7 +13,7 @@ from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.types.utils import GenericStreamingChunk as GChunk from litellm.types.utils import StreamingChatCompletionChunk -_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) +STR_KEYED_MAPPING: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) @with_config(ConfigDict(extra="allow", strict=True, hide_input_in_errors=True)) @@ -54,7 +54,9 @@ def _load_sagemaker_response_stream_shape(): from botocore.model import ServiceModel loader: Final = Loader() - service_dict: Final = _JSON_OBJECT.validate_python(loader.load_service_model("sagemaker-runtime", "service-2")) + service_dict: Final = STR_KEYED_MAPPING.validate_python( + loader.load_service_model("sagemaker-runtime", "service-2") + ) return ServiceModel(service_dict).shape_for("InvokeEndpointWithResponseStreamOutput") except Exception as e: verbose_logger.warning( @@ -76,6 +78,38 @@ def get_sagemaker_response_stream_shape(): return _load_sagemaker_response_stream_shape() +def with_inference_component_hint(error_message: str) -> str: + if "Inference Component Name header is required" in error_message: + return f"{error_message}\n pass in via `litellm.embedding(..., model_id={{InferenceComponentName}})`" + if "Inference Component Name header is not allowed" in error_message: + return f"{error_message}\n remove `model_id` from this deployment, the endpoint has no inference components" + return error_message + + +class _ClientErrorBody(BaseModel): + model_config = ConfigDict(extra="ignore", frozen=True) + message: str | None = Field(default=None, alias="Message") + + +class _ClientErrorMetadata(BaseModel): + model_config = ConfigDict(extra="ignore", frozen=True) + http_status_code: int = Field(default=500, alias="HTTPStatusCode") + + +class _ClientErrorResponse(BaseModel): + model_config = ConfigDict(extra="ignore", frozen=True) + error: _ClientErrorBody = Field(default_factory=_ClientErrorBody, alias="Error") + response_metadata: _ClientErrorMetadata = Field(default_factory=_ClientErrorMetadata, alias="ResponseMetadata") + + +def client_error_details(error: Exception) -> tuple[int, str]: + try: + response: Final = _ClientErrorResponse.model_validate(getattr(error, "response", None)) + except ValidationError: + return 500, str(error) + return response.response_metadata.http_status_code, response.error.message or str(error) + + class SagemakerError(BaseLLMException): def __init__( self, diff --git a/litellm/llms/sagemaker/completion/handler.py b/litellm/llms/sagemaker/completion/handler.py index 069518021fa..b285200ef15 100644 --- a/litellm/llms/sagemaker/completion/handler.py +++ b/litellm/llms/sagemaker/completion/handler.py @@ -4,6 +4,7 @@ from copy import deepcopy from typing import Final, cast import httpx +from typing_extensions import ReadOnly, TypedDict import litellm from litellm._logging import verbose_logger @@ -24,12 +25,26 @@ from litellm.utils import ( get_secret, ) -from ..common_utils import AWSEventStreamDecoder, SagemakerError +from ..common_utils import ( + STR_KEYED_MAPPING, + AWSEventStreamDecoder, + SagemakerError, + client_error_details, + with_inference_component_hint, +) from ..embedding.transformation import SagemakerEmbeddingConfig from .transformation import SagemakerConfig sagemaker_config: Final = SagemakerConfig() + +class _InvokeEndpointArgs(TypedDict): + EndpointName: ReadOnly[str] + ContentType: ReadOnly[str] + Body: ReadOnly[bytes] + CustomAttributes: ReadOnly[str] + + """ SAGEMAKER AUTH Keys/Vars os.environ['AWS_ACCESS_KEY_ID'] = "" @@ -288,10 +303,7 @@ class SagemakerLLM(BaseAWSLLM): raise e except Exception as e: verbose_logger.error("Sagemaker error %s", str(e)) - status_code: Final = getattr(e, "response", {}).get("ResponseMetadata", {}).get("HTTPStatusCode", 500) - error_message = getattr(e, "response", {}).get("Error", {}).get("Message", str(e)) - if "Inference Component Name header is required" in error_message: - error_message += "\n pass in via `litellm.completion(..., model_id={InferenceComponentName})`" + status_code, error_message = client_error_details(e) raise SagemakerError(status_code=status_code, message=error_message) return sagemaker_config.transform_response( @@ -512,10 +524,7 @@ class SagemakerLLM(BaseAWSLLM): ) raise e except Exception as e: - error_message = f"{e}" - if "Inference Component Name header is required" in error_message: - error_message += "\n pass in via `litellm.completion(..., model_id={InferenceComponentName})`" - raise SagemakerError(status_code=500, message=error_message) + raise SagemakerError(status_code=500, message=str(e)) return sagemaker_config.transform_response( model=model, raw_response=response, @@ -574,16 +583,29 @@ class SagemakerLLM(BaseAWSLLM): #### EMBEDDING LOGIC # Transform request based on model type provider_config: Final = SagemakerEmbeddingConfig.get_model_config(model) - request_data: Final = provider_config.transform_embedding_request(model, input, optional_params, {}) + endpoint_name: Final = SagemakerEmbeddingConfig.get_endpoint_name(model) + embedding_params: Final = STR_KEYED_MAPPING.validate_python(optional_params) + inference_component_name: Final = embedding_params.get("model_id") + if inference_component_name is not None and not isinstance(inference_component_name, str): + raise SagemakerError(status_code=400, message="model_id must be the inference component name as a string") + body_params: Final = {k: v for k, v in embedding_params.items() if k != "model_id"} + request_data: Final = provider_config.transform_embedding_request(endpoint_name, input, body_params, {}) data: Final = json.dumps(request_data).encode("utf-8") + invoke_args: Final = _InvokeEndpointArgs( + EndpointName=endpoint_name, + ContentType="application/json", + Body=data, + CustomAttributes="accept_eula=true", + ) ## LOGGING request_str: Final = f""" response = client.invoke_endpoint( - EndpointName={model}, + EndpointName={endpoint_name}, ContentType="application/json", Body=f"{data!r}", # Use !r for safe representation CustomAttributes="accept_eula=true", + InferenceComponentName={inference_component_name}, )""" logging_obj.pre_call( input=input, @@ -592,16 +614,14 @@ class SagemakerLLM(BaseAWSLLM): ) ## EMBEDDING CALL try: - response = client.invoke_endpoint( - EndpointName=model, - ContentType="application/json", - Body=data, - CustomAttributes="accept_eula=true", + response = ( + client.invoke_endpoint(**invoke_args, InferenceComponentName=inference_component_name) + if inference_component_name + else client.invoke_endpoint(**invoke_args) ) except Exception as e: - status_code: Final = getattr(e, "response", {}).get("ResponseMetadata", {}).get("HTTPStatusCode", 500) - error_message: Final = getattr(e, "response", {}).get("Error", {}).get("Message", str(e)) - raise SagemakerError(status_code=status_code, message=error_message) + status_code, error_message = client_error_details(e) + raise SagemakerError(status_code=status_code, message=with_inference_component_hint(error_message)) response = json.loads(response["Body"].read().decode("utf8")) ## LOGGING @@ -628,7 +648,7 @@ class SagemakerLLM(BaseAWSLLM): # Use the request_data that was already transformed above return provider_config.transform_embedding_response( - model=model, + model=endpoint_name, raw_response=mock_response, model_response=model_response, logging_obj=logging_obj, diff --git a/litellm/llms/sagemaker/embedding/openai_transformation.py b/litellm/llms/sagemaker/embedding/openai_transformation.py new file mode 100644 index 00000000000..a3f0ebfbf59 --- /dev/null +++ b/litellm/llms/sagemaker/embedding/openai_transformation.py @@ -0,0 +1,146 @@ +""" +Translate from OpenAI's `/v1/embeddings` to SageMaker's `/invoke` + +For endpoints whose container already serves the OpenAI embeddings API, opted into per +deployment with the `openai/` route (`sagemaker/openai/`): the request +goes out as `{"input": [...]}` plus the caller's extra body fields, and the response is +read as the OpenAI embeddings object (`data[].embedding`, `usage`). +""" + +from collections.abc import Mapping, Sequence +from typing import TYPE_CHECKING, Final + +if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + from litellm.types.llms.openai import AllEmbeddingInputValues + +from httpx._models import Headers, Response +from pydantic import BaseModel, ConfigDict, ValidationError + +from litellm.llms.base_llm.chat.transformation import BaseLLMException +from litellm.llms.base_llm.embedding.transformation import BaseEmbeddingConfig +from litellm.types.llms.openai import AllMessageValues +from litellm.types.utils import EmbeddingResponse, Usage + +from ..common_utils import STR_KEYED_MAPPING, SagemakerError + +SUPPORTED_OPENAI_PARAMS: Final = ("dimensions", "encoding_format", "user") + + +class OpenAIEmbeddingUsage(BaseModel): + model_config = ConfigDict(frozen=True) + prompt_tokens: int = 0 + total_tokens: int | None = None + + +class OpenAIEmbeddingRow(BaseModel): + model_config = ConfigDict(frozen=True) + embedding: tuple[float, ...] | str + index: int | None = None + + +class OpenAIEmbeddingsPayload(BaseModel): + model_config = ConfigDict(frozen=True) + data: tuple[OpenAIEmbeddingRow, ...] + model: str | None = None + usage: OpenAIEmbeddingUsage | None = None + + +class SagemakerOpenAIEmbeddingConfig(BaseEmbeddingConfig): + def get_supported_openai_params(self, model: str) -> list[str]: # mutable-ok: BaseConfig returns a list + return list(SUPPORTED_OPENAI_PARAMS) + + def map_openai_params( + self, + non_default_params: Mapping[str, object], + optional_params: Mapping[str, object], + model: str, + drop_params: bool, + ) -> dict[str, object]: # mutable-ok: callers keep adding to the mapped params + mapped: Final = {k: v for k, v in non_default_params.items() if k in SUPPORTED_OPENAI_PARAMS} + return {**optional_params, **mapped} + + def get_error_class( + self, + error_message: str, + status_code: int, + headers: dict[str, str] | Headers, # mutable-ok: SagemakerError takes the same header types + ) -> BaseLLMException: + return SagemakerError(message=error_message, status_code=status_code, headers=headers) + + def transform_embedding_request( + self, + model: str, + input: "AllEmbeddingInputValues", + optional_params: Mapping[str, object], + headers: Mapping[str, str], + ) -> dict[str, object]: # mutable-ok: BaseEmbeddingConfig returns the request body as a dict + extra_body: Final = optional_params.get("extra_body") + extra_fields: Final[Mapping[str, object]] = ( + STR_KEYED_MAPPING.validate_python(extra_body) if isinstance(extra_body, dict) else {} + ) + top_level_params: Final = {k: v for k, v in optional_params.items() if k != "extra_body"} + return {"input": input, **top_level_params, **extra_fields} + + def transform_embedding_response( + self, + model: str, + raw_response: Response, + model_response: "EmbeddingResponse", + logging_obj: "LiteLLMLoggingObj", + api_key: str | None = None, + request_data: Mapping[str, object] | None = None, + optional_params: Mapping[str, object] | None = None, + litellm_params: Mapping[str, object] | None = None, + ) -> "EmbeddingResponse": + try: + response_data: Final = raw_response.json() + except ValueError as e: + raise SagemakerError( + message=f"Failed to parse response: {e}", + status_code=raw_response.status_code, + ) + + try: + payload: Final = OpenAIEmbeddingsPayload.model_validate(response_data) + except ValidationError as e: + raise SagemakerError( + status_code=500, + message=( + "Unexpected response format. Expected an OpenAI embeddings object with a 'data' list of " + f"{{'embedding': [...]}} rows: {e}" + ), + ) + + prompt_tokens: Final = payload.usage.prompt_tokens if payload.usage else 0 + total_tokens: Final = ( + payload.usage.total_tokens if payload.usage and payload.usage.total_tokens is not None else prompt_tokens + ) + + return model_response.model_copy( + update={ + "object": "list", + "data": [ + { + "object": "embedding", + "index": row.index if row.index is not None else idx, + "embedding": row.embedding if isinstance(row.embedding, str) else list(row.embedding), + } + for idx, row in enumerate(payload.data) + ], + "model": payload.model or model, + "usage": Usage(prompt_tokens=prompt_tokens, completion_tokens=0, total_tokens=total_tokens), + } + ) + + def validate_environment( + self, + headers: Mapping[str, str], + model: str, + messages: Sequence[AllMessageValues], + optional_params: Mapping[str, object], + litellm_params: Mapping[str, object], + api_key: str | None = None, + api_base: str | None = None, + ) -> dict[str, str]: # mutable-ok: BaseConfig returns the request headers as a dict + return {"Content-Type": "application/json"} diff --git a/litellm/llms/sagemaker/embedding/transformation.py b/litellm/llms/sagemaker/embedding/transformation.py index 97940929b09..1395f9b98ec 100644 --- a/litellm/llms/sagemaker/embedding/transformation.py +++ b/litellm/llms/sagemaker/embedding/transformation.py @@ -19,6 +19,20 @@ from litellm.types.utils import EmbeddingResponse, Usage from ..common_utils import SagemakerError from .cohere_transformation import SagemakerCohereEmbeddingConfig +from .openai_transformation import SagemakerOpenAIEmbeddingConfig + +OPENAI_COMPATIBLE_ROUTE: Final = "openai/" + + +def _model_specific_config(model: str) -> BaseEmbeddingConfig | None: + if model.startswith(OPENAI_COMPATIBLE_ROUTE): + return SagemakerOpenAIEmbeddingConfig() + model_lower: Final = model.lower() + if "voyage" in model_lower: + return VoyageEmbeddingConfig() + if "cohere" in model_lower: + return SagemakerCohereEmbeddingConfig() + return None class SagemakerEmbeddingConfig(BaseEmbeddingConfig): @@ -35,25 +49,22 @@ class SagemakerEmbeddingConfig(BaseEmbeddingConfig): Factory method to get the appropriate embedding config based on model type Args: - model: The model name + model: The model name, optionally prefixed with the `openai/` route for + endpoints whose container serves the OpenAI embeddings API Returns: Appropriate embedding config instance """ - model_lower: Final = model.lower() - if "voyage" in model_lower: - return VoyageEmbeddingConfig() - if "cohere" in model_lower: - return SagemakerCohereEmbeddingConfig() - return cls() + specific_config: Final = _model_specific_config(model) + return cls() if specific_config is None else specific_config + + @staticmethod + def get_endpoint_name(model: str) -> str: + return model.removeprefix(OPENAI_COMPATIBLE_ROUTE) def get_supported_openai_params(self, model: str) -> list[str]: - model_lower: Final = model.lower() - if "voyage" in model_lower: - return VoyageEmbeddingConfig().get_supported_openai_params(model) - if "cohere" in model_lower: - return SagemakerCohereEmbeddingConfig().get_supported_openai_params(model) - return [] + specific_config: Final = _model_specific_config(model) + return [] if specific_config is None else specific_config.get_supported_openai_params(model) def map_openai_params( self, 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 93e0a183a2b..b259d194a32 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 @@ -1299,11 +1299,6 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): elif "AUDIO" not in optional_params["responseModalities"]: optional_params["responseModalities"].append("AUDIO") - # Set default temperature to 1.0 for Gemini 3 models if not specified - if VertexGeminiConfig._is_gemini_3_or_newer(model): - if "temperature" not in optional_params: - optional_params["temperature"] = 1.0 - self._drop_search_tools_mixed_with_functions(optional_params) return optional_params @@ -1963,20 +1958,23 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): tool_use_prompt_tokens: Final = usage_metadata.get("toolUsePromptTokenCount") or None - prompt_tokens_details = PromptTokensDetailsWrapper( - cached_tokens=cached_tokens, - audio_tokens=prompt_audio_tokens, - text_tokens=prompt_text_tokens, - image_tokens=prompt_image_tokens, - video_tokens=prompt_video_tokens, - tool_use_tokens=tool_use_prompt_tokens, - ) - billable_tool_use_prompt_tokens: Final = ( 0 if VertexGeminiConfig._response_has_search_grounding(completion_response) else (tool_use_prompt_tokens or 0) ) + billable_text_tokens: Final = ( + prompt_text_tokens if prompt_text_tokens is None else prompt_text_tokens + billable_tool_use_prompt_tokens + ) + + prompt_tokens_details = PromptTokensDetailsWrapper( + cached_tokens=cached_tokens, + audio_tokens=prompt_audio_tokens, + text_tokens=billable_text_tokens, + image_tokens=prompt_image_tokens, + video_tokens=prompt_video_tokens, + tool_use_tokens=tool_use_prompt_tokens, + ) completion_tokens = response_tokens or completion_response["usageMetadata"].get("candidatesTokenCount", 0) if not VertexGeminiConfig.is_candidate_token_count_inclusive(usage_metadata) and reasoning_tokens: diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 777b68ac335..3d3846c28c0 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -16022,6 +16022,26 @@ "prompt_cache_min_tokens": 1024, "source": "https://platform.claude.com/docs/en/about-claude/pricing" }, + "cloudflare/clef-omni": { + "input_cost_per_token": 1.5e-07, + "litellm_provider": "cloudflare", + "max_input_tokens": 64000, + "mode": "evaluation", + "output_cost_per_token": 0.0, + "source": "https://developers.cloudflare.com/workers-ai/models/clef-omni/", + "supports_audio_input": true, + "supports_video_input": true + }, + "cloudflare/@cf/cloudflare/clef-omni": { + "input_cost_per_token": 1.5e-07, + "litellm_provider": "cloudflare", + "max_input_tokens": 64000, + "mode": "evaluation", + "output_cost_per_token": 0.0, + "source": "https://developers.cloudflare.com/workers-ai/models/clef-omni/", + "supports_audio_input": true, + "supports_video_input": true + }, "cloudflare/clef": { "input_cost_per_token": 2.4e-07, "litellm_provider": "cloudflare", @@ -16614,6 +16634,144 @@ "supports_function_calling": true, "supports_tool_choice": true }, + "coralbricks/glm-5.3-fast": { + "cache_creation_input_token_cost": 1.68e-06, + "cache_read_input_token_cost": 0.0, + "input_cost_per_token": 1.12e-06, + "litellm_provider": "coralbricks", + "max_input_tokens": 1048576, + "max_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 4.4e-06, + "source": "https://www.coralbricks.ai/pricing", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses", + "/v1/messages" + ], + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": false + }, + "coralbricks/glm-5.3-fp4": { + "cache_creation_input_token_cost": 1.68e-06, + "cache_read_input_token_cost": 0.0, + "input_cost_per_token": 1.12e-06, + "litellm_provider": "coralbricks", + "max_input_tokens": 1048576, + "max_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 4.4e-06, + "source": "https://www.coralbricks.ai/pricing", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses", + "/v1/messages" + ], + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": false + }, + "coralbricks/glm-5.3": { + "cache_creation_input_token_cost": 1.68e-06, + "cache_read_input_token_cost": 0.0, + "input_cost_per_token": 1.12e-06, + "litellm_provider": "coralbricks", + "max_input_tokens": 1048576, + "max_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 4.4e-06, + "source": "https://www.coralbricks.ai/pricing", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses", + "/v1/messages" + ], + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": false + }, + "coralbricks/deepseek-v4.1-flash-fast": { + "cache_creation_input_token_cost": 9e-08, + "cache_read_input_token_cost": 0.0, + "input_cost_per_token": 1e-08, + "litellm_provider": "coralbricks", + "max_input_tokens": 1048576, + "max_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 1.2e-06, + "source": "https://www.coralbricks.ai/pricing", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses", + "/v1/messages" + ], + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "coralbricks/deepseek-v4.1-flash-fast-fp4": { + "cache_creation_input_token_cost": 9e-08, + "cache_read_input_token_cost": 0.0, + "input_cost_per_token": 1e-08, + "litellm_provider": "coralbricks", + "max_input_tokens": 1048576, + "max_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 1.2e-06, + "source": "https://www.coralbricks.ai/pricing", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses", + "/v1/messages" + ], + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "coralbricks/deepseek-v4.1-flash": { + "cache_creation_input_token_cost": 9e-08, + "cache_read_input_token_cost": 0.0, + "input_cost_per_token": 1e-08, + "litellm_provider": "coralbricks", + "max_input_tokens": 1048576, + "max_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 1.2e-06, + "source": "https://www.coralbricks.ai/pricing", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses", + "/v1/messages" + ], + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true + }, "computer-use-preview": { "deprecation_date": "2026-07-23", "input_cost_per_token": 3e-06, @@ -16644,6 +16802,36 @@ "supports_vision": true, "source": "https://platform.openai.com/docs/models/computer-use-preview" }, + "computer-use-preview-2025-03-11": { + "deprecation_date": "2026-07-23", + "input_cost_per_token": 3e-06, + "litellm_provider": "openai", + "max_input_tokens": 8192, + "max_output_tokens": 1024, + "max_tokens": 1024, + "mode": "chat", + "output_cost_per_token": 1.2e-05, + "supported_endpoints": [ + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": false, + "supports_reasoning": true, + "supports_response_schema": false, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true, + "source": "https://developers.openai.com/api/docs/models/computer-use-preview" + }, "deepseek-chat": { "cache_read_input_token_cost": 2.8e-08, "deprecation_date": "2026-07-24", @@ -33449,6 +33637,44 @@ "supports_tool_choice": true, "supports_vision": false }, + "gpt-audio-mini-2025-10-06": { + "deprecation_date": "2026-07-23", + "input_cost_per_audio_token": 1e-05, + "input_cost_per_token": 6e-07, + "litellm_provider": "openai", + "max_input_tokens": 128000, + "max_output_tokens": 16384, + "max_tokens": 16384, + "mode": "chat", + "output_cost_per_audio_token": 2e-05, + "output_cost_per_token": 2.4e-06, + "source": "https://developers.openai.com/api/docs/models/gpt-audio-mini", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses", + "/v1/realtime", + "/v1/batch" + ], + "supported_modalities": [ + "text", + "audio" + ], + "supported_output_modalities": [ + "text", + "audio" + ], + "supports_audio_input": true, + "supports_audio_output": true, + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_parallel_function_calling": true, + "supports_prompt_caching": false, + "supports_reasoning": false, + "supports_response_schema": false, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": false + }, "gpt-audio-mini-2025-12-15": { "input_cost_per_audio_token": 1e-05, "input_cost_per_token": 6e-07, @@ -33609,6 +33835,34 @@ "supports_vision": true, "supports_web_search": true }, + "gpt-4o-mini-search-preview-2025-03-11": { + "deprecation_date": "2026-07-23", + "cache_read_input_token_cost": 7.5e-08, + "input_cost_per_token": 1.5e-07, + "input_cost_per_token_batches": 7.5e-08, + "litellm_provider": "openai", + "max_input_tokens": 128000, + "max_output_tokens": 16384, + "max_tokens": 16384, + "mode": "chat", + "output_cost_per_token": 6e-07, + "output_cost_per_token_batches": 3e-07, + "search_context_cost_per_query": { + "search_context_size_high": 0.025, + "search_context_size_low": 0.025, + "search_context_size_medium": 0.025 + }, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "source": "https://developers.openai.com/api/docs/models/gpt-4o-mini-search-preview" + }, "gpt-4o-mini-transcribe": { "input_cost_per_audio_token": 1.25e-06, "input_cost_per_token": 1.25e-06, @@ -33689,6 +33943,34 @@ "supports_vision": true, "supports_web_search": true }, + "gpt-4o-search-preview-2025-03-11": { + "deprecation_date": "2026-07-23", + "cache_read_input_token_cost": 1.25e-06, + "input_cost_per_token": 2.5e-06, + "input_cost_per_token_batches": 1.25e-06, + "litellm_provider": "openai", + "max_input_tokens": 128000, + "max_output_tokens": 16384, + "max_tokens": 16384, + "mode": "chat", + "output_cost_per_token": 1e-05, + "output_cost_per_token_batches": 5e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.025, + "search_context_size_low": 0.025, + "search_context_size_medium": 0.025 + }, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "source": "https://developers.openai.com/api/docs/models/gpt-4o-search-preview" + }, "gpt-4o-transcribe": { "input_cost_per_audio_token": 2.5e-06, "input_cost_per_token": 2.5e-06, @@ -34942,13 +35224,17 @@ }, "gpt-6-luna": { "cache_creation_input_token_cost": 1.25e-07, + "cache_creation_input_token_cost_decisions": 1e-07, "cache_creation_input_token_cost_above_272k_tokens": 2.5e-07, + "cache_creation_input_token_cost_above_272k_tokens_decisions": 2e-07, "cache_creation_input_token_cost_above_272k_tokens_flex": 1.25e-07, "cache_creation_input_token_cost_above_272k_tokens_priority": 5e-07, "cache_creation_input_token_cost_flex": 6.25e-08, "cache_creation_input_token_cost_priority": 2.5e-07, "cache_read_input_token_cost": 1e-08, + "cache_read_input_token_cost_decisions": 1e-07, "cache_read_input_token_cost_above_272k_tokens": 2e-08, + "cache_read_input_token_cost_above_272k_tokens_decisions": 2e-07, "cache_read_input_token_cost_above_272k_tokens_flex": 1e-08, "cache_read_input_token_cost_above_272k_tokens_priority": 4e-08, "cache_read_input_token_cost_batches": 5e-09, @@ -34958,7 +35244,9 @@ "cache_read_input_token_cost_flex": 5e-09, "cache_read_input_token_cost_priority": 2e-08, "input_cost_per_token": 1e-07, + "input_cost_per_token_decisions": 1e-07, "input_cost_per_token_above_272k_tokens": 2e-07, + "input_cost_per_token_above_272k_tokens_decisions": 2e-07, "input_cost_per_token_above_272k_tokens_flex": 1e-07, "input_cost_per_token_above_272k_tokens_priority": 4e-07, "input_cost_per_token_batches": 5e-08, @@ -34971,7 +35259,9 @@ "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 5e-07, + "output_cost_per_token_decisions": 0.0, "output_cost_per_token_above_272k_tokens": 7.5e-07, + "output_cost_per_token_above_272k_tokens_decisions": 0.0, "output_cost_per_token_above_272k_tokens_flex": 3.75e-07, "output_cost_per_token_above_272k_tokens_priority": 1.5e-06, "output_cost_per_token_batches": 2.5e-07, @@ -38843,6 +39133,49 @@ ], "supports_audio_input": true }, + "meta/muse-wonder-1.0": { + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 4e-06, + "litellm_provider": "meta", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 2e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.0025, + "search_context_size_low": 0.0025, + "search_context_size_medium": 0.0025 + }, + "source": "https://dev.meta.ai/docs/pricing-rate-limits", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses", + "/v1/messages" + ], + "supported_modalities": [ + "text", + "image", + "video" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_minimal_reasoning_effort": true, + "supports_none_reasoning_effort": false, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "supports_xhigh_reasoning_effort": true + }, "meta_llama/Llama-3.3-70B-Instruct": { "litellm_provider": "meta_llama", "max_input_tokens": 128000, @@ -41481,6 +41814,38 @@ "supports_pdf_input": true, "deprecation_date": "2026-07-23" }, + "o3-deep-research-2025-06-26": { + "deprecation_date": "2026-07-23", + "cache_read_input_token_cost": 2.5e-06, + "input_cost_per_token": 1e-05, + "litellm_provider": "openai", + "max_input_tokens": 200000, + "max_output_tokens": 100000, + "max_tokens": 100000, + "mode": "responses", + "output_cost_per_token": 4e-05, + "source": "https://developers.openai.com/api/docs/models/o3-deep-research", + "supported_endpoints": [ + "/v1/batch", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": false, + "supports_native_streaming": true, + "supports_parallel_function_calling": false, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": false, + "supports_tool_choice": false, + "supports_vision": true, + "supports_pdf_input": true + }, "o3-2025-04-16": { "cache_read_input_token_cost": 5e-07, "cache_read_input_token_cost_flex": 2.5e-07, @@ -41711,6 +42076,38 @@ "supports_pdf_input": true, "deprecation_date": "2026-07-23" }, + "o4-mini-deep-research-2025-06-26": { + "deprecation_date": "2026-07-23", + "cache_read_input_token_cost": 5e-07, + "input_cost_per_token": 2e-06, + "litellm_provider": "openai", + "max_input_tokens": 200000, + "max_output_tokens": 100000, + "max_tokens": 100000, + "mode": "responses", + "output_cost_per_token": 8e-06, + "source": "https://developers.openai.com/api/docs/models/o4-mini-deep-research", + "supported_endpoints": [ + "/v1/batch", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": false, + "supports_native_streaming": true, + "supports_parallel_function_calling": false, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": false, + "supports_tool_choice": false, + "supports_vision": true, + "supports_pdf_input": true + }, "o4-mini-2025-04-16": { "cache_read_input_token_cost": 2.75e-07, "cache_read_input_token_cost_flex": 1.38e-07, @@ -43006,13 +43403,21 @@ "output_cost_per_token": 2.8e-07 }, "perplexity/pplx-decider-v1-27b": { - "input_cost_per_token": 4e-08, + "input_cost_per_token": 2e-08, "litellm_provider": "perplexity", "max_input_tokens": 262144, "mode": "evaluation", "output_cost_per_token": 0.0, "source": "https://docs.perplexity.ai/api-reference/decisions-post" }, + "perplexity/pplx-decider-v1.1-27b": { + "input_cost_per_token": 2e-08, + "litellm_provider": "perplexity", + "max_input_tokens": 262144, + "mode": "evaluation", + "output_cost_per_token": 0.0, + "source": "https://docs.perplexity.ai/docs/getting-started/pricing" + }, "perplexity/sonar": { "input_cost_per_token": 1e-06, "litellm_provider": "perplexity", @@ -51404,6 +51809,70 @@ "litellm_provider": "vertex_ai", "mode": "vector_store" }, + "vertex_ai/semantic-ranker-default": { + "input_cost_per_query": 0.001, + "input_cost_per_token": 0.0, + "litellm_provider": "vertex_ai", + "mode": "rerank", + "output_cost_per_token": 0.0, + "source": "https://cloud.google.com/generative-ai-app-builder/pricing" + }, + "vertex_ai/semantic-ranker-default-002": { + "input_cost_per_query": 0.001, + "input_cost_per_token": 0.0, + "litellm_provider": "vertex_ai", + "mode": "rerank", + "output_cost_per_token": 0.0, + "source": "https://cloud.google.com/generative-ai-app-builder/pricing" + }, + "vertex_ai/semantic-ranker-default-003": { + "input_cost_per_query": 0.001, + "input_cost_per_token": 0.0, + "litellm_provider": "vertex_ai", + "mode": "rerank", + "output_cost_per_token": 0.0, + "source": "https://cloud.google.com/generative-ai-app-builder/pricing" + }, + "vertex_ai/semantic-ranker-default-004": { + "input_cost_per_query": 0.001, + "input_cost_per_token": 0.0, + "litellm_provider": "vertex_ai", + "mode": "rerank", + "output_cost_per_token": 0.0, + "source": "https://cloud.google.com/generative-ai-app-builder/pricing" + }, + "vertex_ai/semantic-ranker-default@latest": { + "input_cost_per_query": 0.001, + "input_cost_per_token": 0.0, + "litellm_provider": "vertex_ai", + "mode": "rerank", + "output_cost_per_token": 0.0, + "source": "https://cloud.google.com/generative-ai-app-builder/pricing" + }, + "vertex_ai/semantic-ranker-fast": { + "input_cost_per_query": 0.001, + "input_cost_per_token": 0.0, + "litellm_provider": "vertex_ai", + "mode": "rerank", + "output_cost_per_token": 0.0, + "source": "https://cloud.google.com/generative-ai-app-builder/pricing" + }, + "vertex_ai/semantic-ranker-fast-004": { + "input_cost_per_query": 0.001, + "input_cost_per_token": 0.0, + "litellm_provider": "vertex_ai", + "mode": "rerank", + "output_cost_per_token": 0.0, + "source": "https://cloud.google.com/generative-ai-app-builder/pricing" + }, + "vertex_ai/semantic-ranker-fast@latest": { + "input_cost_per_query": 0.001, + "input_cost_per_token": 0.0, + "litellm_provider": "vertex_ai", + "mode": "rerank", + "output_cost_per_token": 0.0, + "source": "https://cloud.google.com/generative-ai-app-builder/pricing" + }, "openai/container": { "code_interpreter_cost_per_session": 0.03, "litellm_provider": "openai", @@ -65714,7 +66183,6 @@ "source": "https://api.together.xyz/v1/models" }, "together_ai/together/Tev1-4B-experimental": { - "cache_read_input_token_cost": 4.2e-08, "input_cost_per_token": 4.2e-08, "litellm_provider": "together_ai", "max_input_tokens": 32768, @@ -69135,6 +69603,15 @@ "supports_response_schema": true, "supports_tool_choice": true }, + "fireworks_ai/accounts/fireworks/routers/rune": { + "cache_read_input_token_cost": 8.4e-09, + "input_cost_per_token": 4.2e-08, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 262144, + "mode": "evaluation", + "output_cost_per_token": 0.0, + "source": "https://docs.fireworks.ai" + }, "fireworks_ai/accounts/fireworks/models/ember-1": { "cache_read_input_token_cost": 3e-07, "cache_read_input_token_cost_priority": 3.75e-07, @@ -71966,5 +72443,97 @@ "supports_system_messages": false, "supports_vision": false, "supports_function_calling": false + }, + "vertex_ai/gemini-robotics-er-2-preview": { + "cache_read_input_token_cost": 1e-07, + "cache_read_input_token_cost_batches": 5e-08, + "input_cost_per_audio_token": 1e-06, + "input_cost_per_token": 1e-06, + "input_cost_per_token_batches": 5e-07, + "litellm_provider": "vertex_ai", + "max_input_tokens": 131072, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "chat", + "output_cost_per_reasoning_token": 5e-06, + "output_cost_per_token": 5e-06, + "output_cost_per_token_batches": 2.5e-06, + "search_context_cost_per_query": {"search_context_size_low":0.014,"search_context_size_medium":0.014,"search_context_size_high":0.014}, + "source": "https://ai.google.dev/gemini-api/docs/pricing", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions", + "/v1/batch" + ], + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_input": true, + "supports_audio_output": false, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_url_context": true, + "supports_video_input": true, + "supports_vision": true, + "supports_web_search": true, + "web_search_billing_unit": "per_query" + }, + "gemini-robotics-er-2-preview": { + "cache_read_input_token_cost": 1e-07, + "cache_read_input_token_cost_batches": 5e-08, + "input_cost_per_audio_token": 1e-06, + "input_cost_per_token": 1e-06, + "input_cost_per_token_batches": 5e-07, + "litellm_provider": "vertex_ai-language-models", + "max_input_tokens": 131072, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "chat", + "output_cost_per_reasoning_token": 5e-06, + "output_cost_per_token": 5e-06, + "output_cost_per_token_batches": 2.5e-06, + "search_context_cost_per_query": {"search_context_size_low":0.014,"search_context_size_medium":0.014,"search_context_size_high":0.014}, + "source": "https://ai.google.dev/gemini-api/docs/pricing", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions", + "/v1/batch" + ], + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_input": true, + "supports_audio_output": false, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_url_context": true, + "supports_video_input": true, + "supports_vision": true, + "supports_web_search": true, + "web_search_billing_unit": "per_query" } } diff --git a/litellm/provider_endpoints_support_backup.json b/litellm/provider_endpoints_support_backup.json index a0e75b1f02d..8dbdac37717 100644 --- a/litellm/provider_endpoints_support_backup.json +++ b/litellm/provider_endpoints_support_backup.json @@ -618,6 +618,24 @@ "interactions": true } }, + "coralbricks": { + "display_name": "CoralBricks (`coralbricks`)", + "url": "https://docs.litellm.ai/docs/providers/coralbricks", + "endpoints": { + "chat_completions": true, + "messages": true, + "responses": true, + "embeddings": false, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false, + "a2a": false, + "interactions": false + } + }, "cortecs": { "display_name": "Cortecs (`cortecs`)", "url": "https://docs.litellm.ai/docs/providers/cortecs", diff --git a/litellm/proxy/_experimental/mcp_server/db.py b/litellm/proxy/_experimental/mcp_server/db.py index 823f0ed675c..8ab0d0118f4 100644 --- a/litellm/proxy/_experimental/mcp_server/db.py +++ b/litellm/proxy/_experimental/mcp_server/db.py @@ -1790,7 +1790,7 @@ async def store_user_oauth_credential( payload["refresh_token"] = refresh_token if expires_at: payload["expires_at"] = expires_at - if scopes: + if scopes is not None: payload["scopes"] = scopes # Guard against silently overwriting a BYOK credential with an OAuth token. @@ -2123,9 +2123,7 @@ async def refresh_user_oauth_token( new_refresh_token: Final[str | None] = body.get("refresh_token") or refresh_token raw_scope: Final = body.get("scope") - scopes: list[str] | None = (raw_scope.split() if isinstance(raw_scope, str) and raw_scope else None) or cred.get( - "scopes" - ) + scopes: Final = raw_scope.split() if isinstance(raw_scope, str) else cred.get("scopes") await store_user_oauth_credential( prisma_client=prisma_client, diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index 0253719f410..b9a1a73549a 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -12,6 +12,8 @@ from urllib.parse import parse_qsl, urlencode, urlparse, urlunparse import httpx from fastapi import APIRouter, Depends, Form, HTTPException, Request from fastapi.responses import HTMLResponse, JSONResponse, RedirectResponse, Response +from mcp.client.auth.utils import get_client_metadata_scopes, union_scopes +from mcp.shared.auth import OAuthMetadata from pydantic import ConfigDict, Field, SecretStr, TypeAdapter, ValidationError from litellm._logging import verbose_logger @@ -110,6 +112,7 @@ from litellm.types.mcp import MCPAuth, MCPCredentials from litellm.types.mcp_server.mcp_server_manager import MCPServer, MCPTokenEndpointAuthMethod if TYPE_CHECKING: + from litellm.proxy._experimental.mcp_server.db import OAuthCredentialPayload from litellm.proxy._types import LiteLLM_MCPServerTable # TTL cache for upstream OAuth metadata fetched from pass-through MCP servers. @@ -219,6 +222,8 @@ def encode_state_with_base_url( expected_issuer: str | None = None, authorization_response_iss_parameter_supported: bool = False, oauth_nonce: str | None = None, + managed_scope: str | None = None, + scope_user_id: str | None = None, ) -> str: """ Encode the base_url, original state, and PKCE parameters using encryption. @@ -251,6 +256,7 @@ def encode_state_with_base_url( """ state_data: Final = { "oauth_nonce": oauth_nonce, + **({"managed_scope": managed_scope, "scope_user_id": scope_user_id} if managed_scope is not None else {}), "base_url": base_url, "original_state": original_state, "code_challenge": code_challenge, @@ -290,6 +296,18 @@ def decode_state_hash(encrypted_state: str) -> dict: return state_data +_MANAGED_AUTH_CODE_PREFIX: Final = "llm_scope_" + + +class _ManagedAuthorizationCode(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + upstream_code: str = Field(min_length=1) + mcp_server_id: str = Field(min_length=1) + issuer: str | None = None + requested_scope: str + scope_user_id: str | None = None + + _BRIDGE_AUTH_CODE_PREFIX: Final = "llm_bcode_" @@ -298,6 +316,8 @@ class _BridgeAuthorizationCode(LiteLLMBaseModel): model_config = ConfigDict(frozen=True) oauth_nonce: str | None = None + requested_scope: str | None = None + expected_issuer: str | None = None upstream_code: str = Field(min_length=1) litellm_user_id: str = Field(min_length=1) mcp_server_id: str = Field(min_length=1) @@ -314,6 +334,8 @@ def seal_bridge_authorization_code( litellm_user_id: str, mcp_server_id: str, oauth_nonce: str | None = None, + requested_scope: str | None = None, + expected_issuer: str | None = None, ) -> str: """Seal the upstream authorization code and the SSO-captured litellm user into a gateway authorization code. The DCR client only echoes this opaque value back at the token endpoint; the @@ -328,6 +350,8 @@ def seal_bridge_authorization_code( "litellm_user_id": litellm_user_id, "mcp_server_id": mcp_server_id, "oauth_nonce": oauth_nonce, + "requested_scope": requested_scope, + "expected_issuer": expected_issuer, }, sort_keys=True, ) @@ -643,6 +667,7 @@ async def _store_per_user_token_server_side( token_response: dict[str, Any], identity_binding_proof: str | None = None, cimd_client_id: str | None = None, + requested_scope: str | None = None, ) -> None: """Persist the OAuth token server-side and warm the Redis cache. @@ -667,8 +692,8 @@ async def _store_per_user_token_server_side( expires_in = None refresh_token: Final[str | None] = token_response.get("refresh_token") or None - raw_scope: Final = token_response.get("scope") - scopes: Final[list | None] = raw_scope.split() if isinstance(raw_scope, str) and raw_scope else None + raw_scope: Final[object] = token_response.get("scope", requested_scope) + scopes: Final[list[str] | None] = raw_scope.split() if isinstance(raw_scope, str) else None try: prisma_client: Final = get_prisma_client_or_throw("Database not connected. Cannot store per-user OAuth token.") @@ -954,6 +979,46 @@ async def _resolve_oauth_authorization_user( return user_id +async def _managed_authorization_scope( + request: Request, server: MCPServer, scope: str | None, authorized_user_id: str | None +) -> tuple[str | None, str | None]: + from litellm.proxy._experimental.mcp_server.db import get_user_oauth_credential + from litellm.proxy.proxy_server import prisma_client + + gateway_credential: Final = authorized_user_id is None and await oauth_authorization_uses_gateway_credential( + request + ) + session_user: Final = ( + _session_cookie_user_id(request) if authorized_user_id is None and not gateway_credential else None + ) + scope_user_id: Final = authorized_user_id or ( + await authorize_oauth_credential_request(request, server.server_id) + if gateway_credential + else session_user + if session_user and await _user_can_reach_mcp_server(session_user, server.server_id) + else None + ) + prior_grant: Final = ( + await get_user_oauth_credential(prisma_client, scope_user_id, server.server_id) + if scope_user_id and prisma_client is not None + else None + ) + merged_scope: Final = union_scopes(" ".join(prior_grant.get("scopes") or ()) if prior_grant else None, scope) + selected_scope: Final = ( + get_client_metadata_scopes( + merged_scope, + None, + OAuthMetadata.model_validate(server.authorization_server_metadata) + if server.authorization_server_metadata is not None + else None, + ["authorization_code", "refresh_token"], + ) + if merged_scope is not None + else None + ) + return selected_scope, scope_user_id + + async def authorize_with_server( request: Request, mcp_server: MCPServer, @@ -1029,6 +1094,17 @@ async def authorize_with_server( return subject litellm_user_id = subject + managed_per_user: Final = resolved_server.is_gateway_managed_oauth2 and resolved_server.needs_user_oauth_token + initial_scope: Final = union_scopes( + "openid" if enforce_binding else None, + scope or (" ".join(resolved_server.scopes) if resolved_server.scopes else None), + ) + selected_scope, scope_user_id = ( + await _managed_authorization_scope(request, resolved_server, initial_scope, litellm_user_id) + if managed_per_user + else (initial_scope, None) + ) + oauth_nonce: Final = secrets.token_urlsafe(32) if enforce_binding else None encoded_state: Final = encode_state_with_base_url( base_url=base_url, @@ -1038,7 +1114,11 @@ 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=resolved_server.server_id if (litellm_user_id or ephemeral_dcr_client) else None, + mcp_server_id=resolved_server.server_id + if (litellm_user_id or ephemeral_dcr_client or managed_per_user) + else None, + managed_scope=selected_scope if managed_per_user else None, + scope_user_id=scope_user_id, 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 @@ -1057,13 +1137,8 @@ async def authorize_with_server( } if oauth_nonce: params["nonce"] = oauth_nonce - if scope: - params["scope"] = scope - elif resolved_server.scopes: - params["scope"] = " ".join(resolved_server.scopes) - - if enforce_binding and "openid" not in params.get("scope", "").split(): - params["scope"] = f"openid {params.get('scope', '')}".strip() + if selected_scope: + params["scope"] = selected_scope if code_challenge: params["code_challenge"] = code_challenge @@ -1090,18 +1165,11 @@ def _token_credential_source(mcp_server: MCPServer) -> CredentialSource: return "gateway_stored" if mcp_server.client_id or get_cimd_client_id(mcp_server) else "caller_supplied" -async def _saved_cimd_refresh_client_id( +async def _saved_refresh_credential( server: MCPServer, user_id: str | None, refresh_token: str | None -) -> str | None: - """Reuse only the client identity bound to this caller's presented refresh grant.""" - if ( - not user_id - or not refresh_token - or not server.needs_user_oauth_token - or server.auth_type != MCPAuth.oauth2 - or server.client_id - or server.client_secret - ): +) -> "OAuthCredentialPayload | None": + """Reuse only metadata bound to this caller's presented refresh grant.""" + if not user_id or not refresh_token or not server.needs_user_oauth_token or server.auth_type != MCPAuth.oauth2: return None from litellm.proxy import proxy_server from litellm.proxy._experimental.mcp_server.db import get_user_oauth_credential @@ -1117,7 +1185,33 @@ async def _saved_cimd_refresh_client_id( stored_refresh: Final = credential.get("refresh_token") if not stored_refresh or not secrets.compare_digest(stored_refresh.encode(), refresh_token.encode()): return None - return credential.get("cimd_client_id") + return credential + + +def _open_managed_authorization_code( + code: str | None, server: MCPServer, grant_type: str, user_id: str | None +) -> _ManagedAuthorizationCode | None: + if not isinstance(code, str) or not code.startswith(_MANAGED_AUTH_CODE_PREFIX): + return None + decoded_code: Final = decrypt_value_helper( + code[len(_MANAGED_AUTH_CODE_PREFIX) :], "managed_authorization_code", return_original_value=False + ) + if not isinstance(decoded_code, str): + raise HTTPException(status_code=400, detail="Invalid authorization code") + try: + managed_code: Final = _ManagedAuthorizationCode.model_validate_json(decoded_code) + except ValidationError as exc: + raise HTTPException(status_code=400, detail="Invalid authorization code") from exc + if ( + grant_type != "authorization_code" + or not server.is_gateway_managed_oauth2 + or not server.needs_user_oauth_token + or managed_code.mcp_server_id != server.server_id + or managed_code.issuer != server.issuer + or (managed_code.scope_user_id is not None and managed_code.scope_user_id != user_id) + ): + raise HTTPException(status_code=400, detail="Authorization code does not match this OAuth flow") + return managed_code async def exchange_token_with_server( @@ -1159,10 +1253,18 @@ async def exchange_token_with_server( if resolved_server.needs_user_oauth_token or resolved_server.oauth_identity_binding is not None else None ) - cimd_client_id: Final = ( - await _saved_cimd_refresh_client_id(resolved_server, request_user_id, refresh_token) + managed_code: Final = _open_managed_authorization_code(code, resolved_server, grant_type, request_user_id) + authorization_code: Final = managed_code.upstream_code if managed_code else code + + refresh_credential: Final = ( + await _saved_refresh_credential(resolved_server, request_user_id, refresh_token) if grant_type == "refresh_token" else None + ) + cimd_client_id: Final = ( + refresh_credential.get("cimd_client_id") + if refresh_credential and not resolved_server.client_id and not resolved_server.client_secret + else None ) or get_cimd_client_id(resolved_server) # The id, secret, and token-endpoint auth method must come from the same source. When the @@ -1234,7 +1336,7 @@ async def exchange_token_with_server( ) else: refresh_ownership = None # rebind-ok: grant-specific branches assign one ownership value - if not code: + if not authorization_code: raise HTTPException( status_code=400, detail="code is required for authorization_code grant", @@ -1244,14 +1346,18 @@ async def exchange_token_with_server( # below uses the upstream code, and the mint binds the envelope to the recovered user. Bind the # sealed server to this request so a code minted for one bridge server cannot be spent at another. # A raw upstream code (scripted path) opens to None and the code is used as-is. - bridge_identity = open_bridge_authorization_code(code) - if bridge_identity is not None: - if bridge_identity.mcp_server_id != resolved_server.server_id: - raise HTTPException( - status_code=400, - detail="Authorization code was issued for a different MCP server", - ) - code = bridge_identity.upstream_code + bridge_identity = open_bridge_authorization_code(authorization_code) + if bridge_identity is not None and ( + bridge_identity.mcp_server_id != resolved_server.server_id + or ( + bridge_identity.requested_scope is not None + and bridge_identity.expected_issuer != resolved_server.issuer + ) + ): + raise HTTPException( + status_code=400, + detail="Authorization code was issued for a different MCP server", + ) binding: Final = resolved_server.oauth_identity_binding if binding is not None and binding.mode == "enforce": if bridge_identity is None or not bridge_identity.oauth_nonce: @@ -1273,7 +1379,7 @@ async def exchange_token_with_server( resolved_redirect_uri: Final = redirect_uri if bridge_token_relay else f"{proxy_base_url}/callback" token_data = { "grant_type": "authorization_code", - "code": code, + "code": bridge_identity.upstream_code if bridge_identity else authorization_code, "redirect_uri": resolved_redirect_uri, **token_request.body, } @@ -1353,6 +1459,16 @@ async def exchange_token_with_server( else None ) + requested_scope: Final = ( + managed_code.requested_scope + if managed_code is not None + else bridge_identity.requested_scope + if bridge_identity is not None + else scope or " ".join(refresh_credential.get("scopes") or ()) + if refresh_credential + else None + ) + # Store server-side when the server is configured for per-user OAuth and # the calling client has provided a valid LiteLLM identity. # Errors are non-fatal: the token is still returned to the client. @@ -1379,6 +1495,7 @@ async def exchange_token_with_server( user_id=user_id, token_response=token_response, identity_binding_proof=binding_proof, + requested_scope=requested_scope, **({"cimd_client_id": cimd_client_id} if cimd_client_id is not None else {}), ) else: @@ -1431,8 +1548,10 @@ async def exchange_token_with_server( result["expires_in"] = token_response["expires_in"] if token_response.get("refresh_token"): result["refresh_token"] = token_response["refresh_token"] - if token_response.get("scope"): - result["scope"] = token_response["scope"] + token_fields: Final = TypeAdapter[dict[str, object]](dict[str, object]).validate_python(token_response) + effective_scope: Final = token_fields.get("scope", requested_scope) + if effective_scope is not None: + result["scope"] = effective_scope # RFC 6749 §5.1: token responses must not be cached. return JSONResponse(result, headers=TOKEN_NO_CACHE_HEADERS) @@ -2116,6 +2235,7 @@ async def authorize_mcp_session( code_challenge_method: str | None = None, response_type: str | None = None, resource: str | None = None, + scope: str | None = None, ) -> Response: return aggregate_authorize( request=request, @@ -2127,6 +2247,7 @@ async def authorize_mcp_session( response_type=response_type, session_user_id=_session_cookie_user_id(request), resource=resource, + scope=scope, ) @@ -2169,6 +2290,7 @@ async def authorize( response_type=response_type, session_user_id=_session_cookie_user_id(request), resource=resource, + scope=scope, ) from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager @@ -2281,7 +2403,9 @@ async def token_endpoint( ) -async def _vendor_credential_state(user_id: str, server_id: str) -> VendorCredentialState: +async def _vendor_credential_state( + user_id: str, server_id: str, requested_scopes: tuple[str, ...] = () +) -> VendorCredentialState: """Whether the gateway itself can see a live vendor credential for this user and server. The one reading of "authorized" the connect page displays and the finish step enforces, so @@ -2299,7 +2423,10 @@ async def _vendor_credential_state(user_id: str, server_id: str) -> VendorCreden credential: Final = await get_user_oauth_credential(prisma_client, user_id, server_id) except Exception: # noqa: BLE001 # a credential-read fault must fail the scoped grant closed return "unavailable" - return "absent" if oauth_grant_state(credential) == "absent" else "present" + if oauth_grant_state(credential) == "absent": + return "absent" + granted_scopes: Final = credential.get("scopes") if credential else None + return "present" if set(requested_scopes).issubset(granted_scopes or ()) else "absent" @router.get("/authorize/flow") @@ -2542,6 +2669,11 @@ async def callback( mcp_server_id: Final = state_data.get("mcp_server_id") dcr_client_id: Final = state_data.get("dcr_client_id") dcr_client_secret: Final = state_data.get("dcr_client_secret") + managed_scope: Final = issuer_state.get("managed_scope") + expected_issuer: Final = TypeAdapter[str | None](str | None).validate_python( + issuer_state.get("expected_issuer") + ) + scope_owner: Final = TypeAdapter[str | None](str | None).validate_python(issuer_state.get("scope_user_id")) 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( @@ -2549,6 +2681,8 @@ async def callback( litellm_user_id=litellm_user_id, mcp_server_id=mcp_server_id, oauth_nonce=state_data.get("oauth_nonce"), + requested_scope=managed_scope if isinstance(managed_scope, str) else None, + expected_issuer=expected_issuer, ) 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( @@ -2561,7 +2695,23 @@ async def callback( ), ) - params = {"code": forwarded_code, "state": original_state} + managed_code: Final = ( + _ManagedAuthorizationCode( + upstream_code=forwarded_code, + mcp_server_id=mcp_server_id, + issuer=expected_issuer, + requested_scope=managed_scope, + scope_user_id=scope_owner, + ) + if isinstance(managed_scope, str) and isinstance(mcp_server_id, str) and not litellm_user_id + else None + ) + returned_code: Final = ( + _MANAGED_AUTH_CODE_PREFIX + encrypt_value_helper(managed_code.model_dump_json()) + if managed_code is not None + else forwarded_code + ) + params = {"code": returned_code, "state": original_state} complete_returned_url = _append_query_params(redirect_uri, params) response = RedirectResponse(url=complete_returned_url, status_code=302) _clear_oauth_state_cookie(response, request, state) @@ -2807,11 +2957,16 @@ async def _build_oauth_protected_resource_response( if obo_response is not None: return obo_response + resource_scopes: Final = [ + value + for value in (mcp_server.scopes or () if mcp_server else ()) + if value != "offline_access" or not mcp_server or not mcp_server.is_gateway_managed_oauth2 + ] if mcp_server is not None and mcp_server.advertises_gateway_authorization_server: return { "authorization_servers": [f"{request_base_url}/mcp"], "resource": resource_url, - "scopes_supported": (mcp_server.scopes if mcp_server.scopes else []), + "scopes_supported": resource_scopes, } if mcp_server is None or mcp_server.auth_type != MCPAuth.oauth2_token_exchange: @@ -2822,7 +2977,7 @@ async def _build_oauth_protected_resource_response( (f"{request_base_url}/{mcp_server_name}" if mcp_server_name else f"{request_base_url}") ], "resource": resource_url, - "scopes_supported": (mcp_server.scopes if mcp_server and mcp_server.scopes else []), + "scopes_supported": resource_scopes, } diff --git a/litellm/proxy/_experimental/mcp_server/exceptions.py b/litellm/proxy/_experimental/mcp_server/exceptions.py index c818f6b05bd..9554be3ee0d 100644 --- a/litellm/proxy/_experimental/mcp_server/exceptions.py +++ b/litellm/proxy/_experimental/mcp_server/exceptions.py @@ -4,6 +4,8 @@ from typing import Final from fastapi import HTTPException +from litellm.proxy._experimental.mcp_server.oauth_utils import encode_upstream_scope + class MCPServerURLCredentialsError(HTTPException): """A fixed, sanitized URL-credential migration error safe for operator previews.""" @@ -37,10 +39,14 @@ class MCPUpstreamAuthError(Exception): status_code: int, www_authenticate: str | None, server_name: str, + required_scope: str | None = None, + server_id: str | None = None, ) -> None: self.status_code = status_code self.www_authenticate = www_authenticate self.server_name = server_name + self.required_scope = required_scope + self.server_id = server_id super().__init__(f"Upstream MCP server {server_name!r} returned {status_code}") def to_http_exception( @@ -73,19 +79,39 @@ class MCPUpstreamAuthError(Exception): the client originally targeted, matching the path-aware behaviour of ``get_passthrough_resource_metadata_url`` in ``oauth_utils.py``. """ + aggregate_step_up: Final = ( + self.required_scope is not None + and request_path is not None + and (request_path.rstrip("/") == "/mcp" or request_path.startswith("/toolset/")) + ) + requested_scope: Final = ( + " ".join( + encode_upstream_scope(self.server_id or self.server_name, item) for item in self.required_scope.split() + ) + if aggregate_step_up and self.required_scope + else self.required_scope + ) challenge: str | None = self.www_authenticate - if challenge is None and self.status_code == 401 and base_url: + if challenge is None and (self.status_code == 401 or self.required_scope is not None) and base_url: prefix: Final = base_url.rstrip("/") - if request_path and request_path.startswith(f"/{self.server_name}/mcp"): + if aggregate_step_up: + resource_metadata_url = f"{prefix}/.well-known/oauth-protected-resource" + elif request_path and request_path.startswith(f"/{self.server_name}/mcp"): resource_metadata_url = f"{prefix}/.well-known/oauth-protected-resource/{self.server_name}/mcp" else: resource_metadata_url = f"{prefix}/.well-known/oauth-protected-resource/mcp/{self.server_name}" challenge = f'Bearer resource_metadata="{resource_metadata_url}"' + scoped_challenge: Final = ( + (f'{challenge}, error="insufficient_scope"' if challenge else 'Bearer error="insufficient_scope"') + + (f', scope="{requested_scope}"' if requested_scope else "") + if self.required_scope is not None + else challenge + ) detail: Final = "Forbidden" if self.status_code == 403 else "Unauthorized" return HTTPException( status_code=self.status_code, detail=detail, - headers={"www-authenticate": challenge} if challenge else None, + headers={"www-authenticate": scoped_challenge} if scoped_challenge else None, ) diff --git a/litellm/proxy/_experimental/mcp_server/faults/list_outcomes.py b/litellm/proxy/_experimental/mcp_server/faults/list_outcomes.py index aaeb50981b6..b94ed1ee251 100644 --- a/litellm/proxy/_experimental/mcp_server/faults/list_outcomes.py +++ b/litellm/proxy/_experimental/mcp_server/faults/list_outcomes.py @@ -10,13 +10,15 @@ becomes an outcome, never a second failure. from __future__ import annotations +import re from collections.abc import Iterator from typing import Final, Literal, NamedTuple, NoReturn, TypeAlias +from urllib.request import parse_http_list, parse_keqv_list import httpx import httpx2 from mcp.types import Tool as MCPTool -from pydantic import ConfigDict +from pydantic import ConfigDict, TypeAdapter from typing_extensions import assert_never from litellm.proxy._experimental.mcp_server.exceptions import ( @@ -98,6 +100,29 @@ def upstream_auth_challenge(exc: BaseException) -> tuple[int, str | None] | None return None +def upstream_insufficient_scope(exc: BaseException) -> str | None: + auth: Final = upstream_auth_challenge(exc) + if auth is None or auth[0] != 403 or not auth[1]: + return None + scheme, _, parameters = auth[1].partition(" ") + if scheme.lower() != "bearer": + return None + fields: Final = parse_http_list(parameters) + if any("=" not in field or not field.split("=", 1)[1] for field in fields): + return None + keys: Final = tuple(field.split("=", 1)[0].strip().lower() for field in fields) + if len(set(keys)) != len(keys) or any(re.fullmatch(r"[!#$%&'*+.^_`|~0-9a-z-]+", key) is None for key in keys): + return None + parsed: Final = TypeAdapter(dict[str, str]).validate_python(parse_keqv_list(fields)) + values: Final = {key: value for key, value in zip(keys, parsed.values())} + if values.get("error") != "insufficient_scope": + return None + scope: Final = values.get("scope", "") + if scope and re.fullmatch(r"[!#-\[\]-~]+(?: [!#-\[\]-~]+)*", scope) is None: + return None + return scope + + def raise_classified_list_failure( exc: BaseException, server_name: str, diff --git a/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py b/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py index 5b45c0dd859..a895dfed162 100644 --- a/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py +++ b/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py @@ -59,8 +59,11 @@ from litellm.caching.caching import DualCache from litellm.proxy._experimental.mcp_server.catalog import catalog_operation, global_manager from litellm.proxy._experimental.mcp_server.oauth_utils import ( TOKEN_NO_CACHE_HEADERS, + UPSTREAM_SCOPE_PREFIX, + UpstreamScope, canonical_resource_uri, canonicalize_url_identity, + decode_upstream_scope, get_request_base_url, is_loopback_redirect_host, validate_redirect_uri_shape, @@ -250,7 +253,9 @@ class ConsentTeam(LiteLLMBaseModel): class LookupVendorCredential(Protocol): """Injected read of a user's vendor credential for one server.""" - def __call__(self, user_id: str, server_id: str, /) -> Awaitable[VendorCredentialState]: ... + def __call__( + self, user_id: str, server_id: str, requested_scopes: tuple[str, ...] = (), / + ) -> Awaitable[VendorCredentialState]: ... class LookupServerReachability(Protocol): @@ -273,7 +278,9 @@ async def _refuse_subject_token(subject_token: str, request: Request) -> Subject ) -async def _unavailable_vendor_credential(user_id: str, server_id: str) -> VendorCredentialState: +async def _unavailable_vendor_credential( + user_id: str, server_id: str, requested_scopes: tuple[str, ...] = () +) -> VendorCredentialState: return "unavailable" @@ -310,6 +317,8 @@ class _ConnectFlow(LiteLLMBaseModel): exp: int resource_server_id: str | None = None audience: SessionAudience | None = None + requested_scopes: tuple[str, ...] = () + upstream_scopes: tuple[UpstreamScope, ...] = () class _GatewayAuthCode(LiteLLMBaseModel): @@ -555,6 +564,7 @@ def aggregate_authorize( response_type: str | None, session_user_id: str | None, resource: str | None = None, + scope: str | None = 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 @@ -580,6 +590,13 @@ def aggregate_authorize( if session_user_id is None: return _login_redirect(base_url, request) scoped_server: Final = resolve_scoped_resource_server(request, resource) + scope_tokens: Final = tuple(token for token in (scope or "").split() if token.startswith(UPSTREAM_SCOPE_PREFIX)) + decoded_scopes: Final = tuple(decode_upstream_scope(token) for token in scope_tokens) + if any(item is None for item in decoded_scopes): + return _oauth_error(400, "invalid_scope", "invalid upstream permission scope") + upstream_scopes: Final = tuple(item for item in decoded_scopes if item is not None) + if scoped_server is not None and any(item.server_id != scoped_server.server_id for item in upstream_scopes): + return _oauth_error(400, "invalid_scope", "upstream permission does not match the requested resource") handle: Final = secrets.token_urlsafe(24) flow: Final = _new_connect_flow( session_user_id=session_user_id, @@ -589,6 +606,8 @@ def aggregate_authorize( code_challenge=code_challenge or "", resource_server_id=scoped_server.server_id if scoped_server is not None else None, audience=None, + requested_scopes=tuple((scope or "").split()) if scoped_server is not None and not upstream_scopes else (), + upstream_scopes=upstream_scopes, ) connect_url: Final = _append_query_params(f"{base_url}/ui/connect", (("connect_flow", handle),)) response: Final = RedirectResponse(connect_url, status_code=303) @@ -748,6 +767,8 @@ def _new_connect_flow( code_challenge: str, resource_server_id: str | None, audience: SessionAudience | None, + requested_scopes: tuple[str, ...] = (), + upstream_scopes: tuple[UpstreamScope, ...] = (), ) -> _ConnectFlow: now: Final = datetime.now(timezone.utc) return _ConnectFlow( @@ -760,6 +781,8 @@ def _new_connect_flow( exp=int(now.timestamp()) + CONNECT_FLOW_TTL_SECONDS, resource_server_id=resource_server_id, audience=audience, + requested_scopes=requested_scopes, + upstream_scopes=upstream_scopes, ) @@ -815,20 +838,20 @@ def _open_flow_for( @catalog_operation(global_manager) async def _flow_target( - flow: _ConnectFlow, lookup_server_reachability: LookupServerReachability + user_id: str, server_id: str | None, lookup_server_reachability: LookupServerReachability ) -> tuple[Literal["unscoped", "interactive", "m2m", "stale"], MCPServer | None]: - if flow.resource_server_id is None: + if server_id is None: return "unscoped", None from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( # noqa: PLC0415 # import cycle MCPServerManager, global_mcp_server_manager, ) - server: Final = global_mcp_server_manager.get_mcp_server_by_id(flow.resource_server_id) + server: Final = global_mcp_server_manager.get_mcp_server_by_id(server_id) if ( server is None or not (server.is_gateway_managed_oauth2 or server.advertises_gateway_authorization_server) - or not await lookup_server_reachability(flow.user_id, server.server_id) + or not await lookup_server_reachability(user_id, server.server_id) ): return "stale", None state: Final = ( @@ -847,6 +870,7 @@ class ConnectFlowDescription(TypedDict): server_id: ReadOnly[str | None] server_name: ReadOnly[str | None] connected: ReadOnly[bool | None] + requested_scopes: NotRequired[ReadOnly[tuple[str, ...]]] async def _describe_opened_flow( @@ -854,19 +878,11 @@ async def _describe_opened_flow( lookup_vendor_credential: LookupVendorCredential, lookup_server_reachability: LookupServerReachability, ) -> ConnectFlowDescription | Response: - state, server = await _flow_target(flow, lookup_server_reachability) + if flow.upstream_scopes: + return await _describe_upstream_scope_flow(flow, lookup_vendor_credential, lookup_server_reachability) + state, server = await _flow_target(flow.user_id, flow.resource_server_id, lookup_server_reachability) if state == "interactive" and server is not None: - credential: Final = await lookup_vendor_credential(flow.user_id, server.server_id) - if credential == "unavailable": - return _oauth_error(503, "temporarily_unavailable", _DB_UNAVAILABLE_DESCRIPTION) - interactive_description: Final[ConnectFlowDescription] = { - "state": state, - "client_origin": _origin_only(flow.redirect_uri), - "server_id": server.server_id, - "server_name": server.server_name or server.alias or server.name, - "connected": credential == "present", - } - return interactive_description + return await _describe_interactive_target(flow, server, flow.requested_scopes, lookup_vendor_credential) described: Final[ConnectFlowDescription] = { "state": state, "client_origin": _origin_only(flow.redirect_uri), @@ -877,6 +893,70 @@ async def _describe_opened_flow( return described +async def _describe_interactive_target( + flow: _ConnectFlow, + server: MCPServer, + scopes: tuple[str, ...], + lookup_vendor_credential: LookupVendorCredential, +) -> ConnectFlowDescription | Response: + credential: Final = ( + await lookup_vendor_credential(flow.user_id, server.server_id, scopes) + if scopes + else await lookup_vendor_credential(flow.user_id, server.server_id) + ) + if credential == "unavailable": + return _oauth_error(503, "temporarily_unavailable", _DB_UNAVAILABLE_DESCRIPTION) + described: Final[ConnectFlowDescription] = { + "state": "interactive", + "client_origin": _origin_only(flow.redirect_uri), + "server_id": server.server_id, + "server_name": server.server_name or server.alias or server.name, + "connected": credential == "present", + } + return {**described, "requested_scopes": scopes} if scopes else described + + +async def _describe_upstream_scope_flow( + flow: _ConnectFlow, + lookup_vendor_credential: LookupVendorCredential, + lookup_server_reachability: LookupServerReachability, +) -> ConnectFlowDescription | Response: + server_ids: Final = tuple(dict.fromkeys(item.server_id for item in flow.upstream_scopes)) + targets: Final = tuple( + [await _flow_target(flow.user_id, server_id, lookup_server_reachability) for server_id in server_ids] + ) + if any(state != "interactive" or server is None for state, server in targets): + return _oauth_error(400, "invalid_scope", "an upstream permission target is unavailable") + descriptions: Final = tuple( + [ + await _describe_interactive_target( + flow, + server, + tuple( + dict.fromkeys( + item.permission for item in flow.upstream_scopes if item.server_id == server.server_id + ) + ), + lookup_vendor_credential, + ) + for _, server in targets + if server is not None + ] + ) + for description in descriptions: + if isinstance(description, Response) or not description["connected"]: + return description + if flow.resource_server_id is not None: + return descriptions[0] + return { + "state": "unscoped", + "client_origin": _origin_only(flow.redirect_uri), + "server_id": None, + "server_name": None, + "connected": None, + } + + async def describe_connect_flow( request: Request, flow_handle: str, diff --git a/litellm/proxy/_experimental/mcp_server/legacy_callbacks.py b/litellm/proxy/_experimental/mcp_server/legacy_callbacks.py index 31d1ad8ccc2..440b184e6ee 100644 --- a/litellm/proxy/_experimental/mcp_server/legacy_callbacks.py +++ b/litellm/proxy/_experimental/mcp_server/legacy_callbacks.py @@ -1,6 +1,7 @@ from collections.abc import Mapping from typing import Final, Protocol +import httpx2 from mcp.client.session import ClientRequestContext from mcp.types import ( CreateMessageRequestParams, @@ -14,6 +15,11 @@ from mcp.types import ( from litellm.constants import MCP_CLIENT_TIMEOUT from litellm.proxy._experimental.mcp_server.contracts import OperationContext from litellm.proxy._experimental.mcp_server.mcp_context import get_active_mcp_request_ctx +from litellm.proxy._experimental.mcp_server.scope_challenge import ( + OAuthScopeResponse, + is_tool_call_request, + scope_response_from_request, +) from litellm.proxy._types import UserAPIKeyAuth @@ -100,3 +106,19 @@ def create_elicitation_callback(timeout: float | None = None) -> ElicitationCall ) return callback + + +def get_scope_response() -> OAuthScopeResponse | None: + context: Final = get_active_mcp_request_ctx() + return scope_response_from_request(context.request) if context is not None else None + + +async def record_upstream_tool_authorization(response: httpx2.Response, *, pending: OAuthScopeResponse | None) -> None: + if pending is None or not pending.pending or not response.is_success or response.request.method != "POST": + return + try: + is_tool_call: Final = is_tool_call_request(response.request.content) + except httpx2.RequestNotRead: + return + if is_tool_call: + pending.allow() diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index abfed7dd33c..6beea5bfd71 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -38,6 +38,7 @@ import httpx2 from fastapi import HTTPException from httpx import HTTPStatusError from mcp import ReadResourceResult, Resource +from mcp.shared.auth import OAuthMetadata from mcp.types import CallToolRequestParams as MCPCallToolRequestParams from mcp.types import ( CallToolResult, @@ -57,7 +58,7 @@ from mcp.types import ( ResourceTemplate, ) from mcp.types import Tool as MCPTool -from pydantic import AnyUrl, Field, SecretStr, TypeAdapter +from pydantic import AnyUrl, Field, SecretStr, TypeAdapter, ValidationError from typing_extensions import ReadOnly, assert_never import litellm @@ -103,9 +104,10 @@ from litellm.proxy._experimental.mcp_server.faults.list_outcomes import ( ServerListFault, raise_classified_list_failure, upstream_auth_challenge, + upstream_insufficient_scope, ) from litellm.proxy._experimental.mcp_server.interactions import bind_target -from litellm.proxy._experimental.mcp_server.legacy_callbacks import has_active_mcp_request_context +from litellm.proxy._experimental.mcp_server.legacy_callbacks import get_scope_response, has_active_mcp_request_context from litellm.proxy._experimental.mcp_server.mcp_debug import describe_upstream_http_failure, record_auth_resolution from litellm.proxy._experimental.mcp_server.oauth2_token_cache import ( MCPPerUserTokenCache, @@ -895,6 +897,10 @@ def carry_forward_resolved_oauth_endpoints(new_server: MCPServer, previous_serve new_server.authorization_response_iss_parameter_supported = ( # rebind-ok: publish on the existing rebuild object previous_server.authorization_response_iss_parameter_supported ) + if may_carry and new_server.issuer == previous_server.issuer and new_server.authorization_server_metadata is None: + new_server.authorization_server_metadata = ( # rebind-ok: publish capability on the existing rebuild object + previous_server.authorization_server_metadata + ) if new_server.authorization_url is None and previous_server.authorization_url: new_server.authorization_url = previous_server.authorization_url if may_carry and new_server.token_url is None and previous_server.token_url: @@ -930,7 +936,12 @@ def _restrict_discovery_to_corroborated_authorization_server( return metadata if _endpoints_corroborate_authorization_url(metadata.authorization_url, manual_authorization_url): return metadata - if not metadata.token_url and not metadata.registration_url and not metadata.client_id_metadata_document_supported: + if ( + not metadata.token_url + and not metadata.registration_url + and not metadata.client_id_metadata_document_supported + and metadata.authorization_server_metadata is None + ): return metadata bridge_note: Final = ( " The discovered registration_url is rejected with it, so this dcr_bridge server stays on the" @@ -949,7 +960,12 @@ def _restrict_discovery_to_corroborated_authorization_server( bridge_note, ) return metadata.model_copy( - update={"token_url": None, "registration_url": None, "client_id_metadata_document_supported": False} + update={ + "token_url": None, + "registration_url": None, + "client_id_metadata_document_supported": False, + "authorization_server_metadata": None, + } ) @@ -2086,6 +2102,7 @@ class MCPServerManager: return server.model_copy( update={ "scopes": server.scopes or metadata.scopes, + "authorization_server_metadata": metadata.authorization_server_metadata, "issuer": server.issuer or discovered_issuer, "client_id_metadata_document_supported": metadata.client_id_metadata_document_supported, "authorization_response_iss_parameter_supported": ( @@ -2678,6 +2695,9 @@ class MCPServerManager: scopes=resolved_scopes, configured_scopes=tuple(configured_scopes) if configured_scopes else None, issuer=effective_issuer, + authorization_server_metadata=( + gated_oauth_metadata.authorization_server_metadata if gated_oauth_metadata else None + ), client_id_metadata_document_supported=( gated_oauth_metadata.client_id_metadata_document_supported if gated_oauth_metadata else None ), @@ -3259,6 +3279,9 @@ class MCPServerManager: scopes=resolved_scopes, configured_scopes=configured_scopes, issuer=effective_issuer, + authorization_server_metadata=( + gated_oauth_metadata.authorization_server_metadata if gated_oauth_metadata else None + ), client_id_metadata_document_supported=( gated_oauth_metadata.client_id_metadata_document_supported if gated_oauth_metadata else None ), @@ -5383,7 +5406,16 @@ class MCPServerManager: continue scopes = self._extract_scopes(data.get("scopes_supported")) + try: + authorization_metadata = OAuthMetadata.model_validate(data) + except ValidationError: + authorization_metadata = None metadata = MCPOAuthMetadata( + authorization_server_metadata=( + TypeAdapter(dict[str, object]).validate_python(authorization_metadata.model_dump(mode="json")) + if authorization_metadata is not None + else None + ), scopes=scopes, authorization_url=data.get("authorization_endpoint"), token_url=data.get("token_endpoint"), @@ -6491,11 +6523,12 @@ class MCPServerManager: # oauth2 + delegate_auth_to_upstream (is_oauth_passthrough) is being removed, so it is not # added here even though the list path still relays for it. relays_upstream_auth: Final = mcp_server.is_client_forwarded_token + managed_per_user: Final = mcp_server.is_gateway_managed_oauth2 and mcp_server.needs_user_oauth_token server_label: Final = mcp_server.name or mcp_server.server_name or mcp_server.alias or "" async def _call_tool_via_client(client, params): async with self._outbound_limit(mcp_server, persistent_session): - if not relays_upstream_auth: + if not relays_upstream_auth and not managed_per_user: return await client.call_tool( params, host_progress_callback=host_progress_callback, @@ -6520,8 +6553,17 @@ class MCPServerManager: persistent_session=persistent_session, ) except Exception as e: + required_scope: Final = upstream_insufficient_scope(e) if managed_per_user else None + if required_scope is not None: + raise MCPUpstreamAuthError( + status_code=403, + www_authenticate=None, + server_name=mcp_server.alias or server_label, + required_scope=required_scope, + server_id=mcp_server.server_id, + ) from e auth_info: Final = _extract_upstream_auth_failure(e) - if auth_info is None or auth_info[0] != 401: + if not relays_upstream_auth or auth_info is None or auth_info[0] != 401: # A genuine (non-auth or 403-forbidden) upstream/transport failure. # raise_on_error demoted the client-layer log to debug, so surface it here at # warning level to keep the outage visible; the caller still gets the graceful @@ -6866,6 +6908,15 @@ class MCPServerManager: ) tasks.append(during_hook_task) + scope_response: Final = get_scope_response() + if scope_response is not None and not ( + mcp_server.is_gateway_managed_oauth2 + and mcp_server.needs_user_oauth_token + and not mcp_server.spec_path + and mcp_server.transport in (MCPTransport.http, MCPTransport.sse) + ): + scope_response.allow() + caller_oauth2_headers: Final = oauth2_headers oauth2_headers = await self._resolve_oauth2_headers_for_tool_call(mcp_server, oauth2_headers, user_api_key_auth) diff --git a/litellm/proxy/_experimental/mcp_server/oauth_utils.py b/litellm/proxy/_experimental/mcp_server/oauth_utils.py index 276cf250b30..b8675f749e6 100644 --- a/litellm/proxy/_experimental/mcp_server/oauth_utils.py +++ b/litellm/proxy/_experimental/mcp_server/oauth_utils.py @@ -2,6 +2,8 @@ (BYOK + discoverable / pass-through OAuth proxy).""" import os +import re +from dataclasses import dataclass from ipaddress import ip_address from typing import TYPE_CHECKING, Final, NoReturn from urllib.parse import ParseResult, urlparse, urlsplit, urlunparse, urlunsplit @@ -819,3 +821,32 @@ def build_upstream_oauth2_token_request( if not resource: return client_auth return TokenEndpointClientAuth(headers=client_auth.headers, body={**client_auth.body, "resource": resource}) + + +UPSTREAM_SCOPE_PREFIX: Final = "litellm:mcp_scope:" + + +@dataclass(frozen=True, slots=True) +class UpstreamScope: + server_id: str + permission: str + + +def encode_upstream_scope(server_id: str, permission: str) -> str: + return f"{UPSTREAM_SCOPE_PREFIX}{server_id.encode().hex()}:{permission.encode().hex()}" + + +def decode_upstream_scope(token: str) -> UpstreamScope | None: + encoded: Final = token.removeprefix(UPSTREAM_SCOPE_PREFIX).split(":") + if not token.startswith(UPSTREAM_SCOPE_PREFIX) or len(encoded) != 2: + return None + try: + server_id: Final = bytes.fromhex(encoded[0]).decode() + permission: Final = bytes.fromhex(encoded[1]).decode() + except (ValueError, UnicodeError): + return None + if not server_id or re.fullmatch(r"[!#-\[\]-~]+", permission) is None: + return None + if encode_upstream_scope(server_id, permission) != token: + return None + return UpstreamScope(server_id, permission) diff --git a/litellm/proxy/_experimental/mcp_server/operations.py b/litellm/proxy/_experimental/mcp_server/operations.py index 2f0b8e64c9e..5087bc6df09 100644 --- a/litellm/proxy/_experimental/mcp_server/operations.py +++ b/litellm/proxy/_experimental/mcp_server/operations.py @@ -94,6 +94,7 @@ from litellm.proxy._experimental.mcp_server.interactions import ( seal_continuation, target_digest, ) +from litellm.proxy._experimental.mcp_server.legacy_callbacks import get_scope_response from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( # noqa: F401 # legacy module exports MCPServerManager, _caller_authorization_fans_out, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export @@ -3009,11 +3010,12 @@ async def _execute_mcp_server_tool_call( is_error=True, ) except MCPUpstreamAuthError as e: - # The MCP session manager serializes handler exceptions as JSON-RPC errors, so a - # mid-session tool call cannot emit a raw 401 + WWW-Authenticate the way the REST - # call path and the connect-time preemptive check do. Return an explicit isError - # naming the upstream status (at info level, not a traceback) so the client still - # learns it must re-authenticate upstream and expected pass-through 401s don't spam. + scope_response: Final = get_scope_response() + if scope_response is not None: + scope_response.deny(e) + # The request-owned response adapter relays managed scope challenges before + # Streamable HTTP headers commit. Other auth failures retain the existing + # explicit tool error, including failures after streaming has begun. verbose_logger.info("Upstream auth failure calling MCP tool: HTTP %s", e.status_code) return CallToolResult( content=[ diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/authz_code_refresher.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/authz_code_refresher.py index b926ca8a6a4..a653ef5e5ff 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/authz_code_refresher.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/authz_code_refresher.py @@ -65,7 +65,7 @@ def _parse_expires_in(raw: object) -> int | None: def _parse_scopes(raw: object) -> tuple[str, ...] | None: - return tuple(raw.split()) if isinstance(raw, str) and raw else None + return tuple(raw.split()) if isinstance(raw, str) else None class AuthorizationCodeRefresher: @@ -156,7 +156,8 @@ class AuthorizationCodeRefresher: rotated: Final = body.get("refresh_token") new_refresh: Final = rotated if isinstance(rotated, str) and rotated else token.refresh_token expires_in: Final = _parse_expires_in(body.get("expires_in")) - scopes: Final = _parse_scopes(body.get("scope")) or token.scopes + granted_scopes: Final = _parse_scopes(body.get("scope")) + scopes: Final = granted_scopes if granted_scopes is not None else token.scopes await self._persist( user_id, @@ -164,7 +165,7 @@ class AuthorizationCodeRefresher: access_token, new_refresh, expires_in, - scopes or None, + scopes, **({"identity_binding_proof": binding_proof} if binding_proof is not None else {}), **({"cimd_client_id": cimd_client_id} if cimd_client_id is not None else {}), ) diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/per_user_oauth_store.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/per_user_oauth_store.py index d8094d5facb..2e9f3232c8d 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/per_user_oauth_store.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/per_user_oauth_store.py @@ -88,7 +88,7 @@ async def _persist_credential( access_token=access_token, refresh_token=refresh_token, expires_in=expires_in, - scopes=list(scopes) if scopes else None, + scopes=list(scopes) if scopes is not None else None, skip_byok_guard=True, identity_binding_proof=identity_binding_proof, **({"cimd_client_id": cimd_client_id} if cimd_client_id is not None else {}), diff --git a/litellm/proxy/_experimental/mcp_server/scope_challenge.py b/litellm/proxy/_experimental/mcp_server/scope_challenge.py new file mode 100644 index 00000000000..c326418bee7 --- /dev/null +++ b/litellm/proxy/_experimental/mcp_server/scope_challenge.py @@ -0,0 +1,93 @@ +"""Keep upstream scope challenges ahead of Streamable HTTP response headers.""" + +from __future__ import annotations + +import asyncio +from typing import TYPE_CHECKING, Final + +from pydantic import TypeAdapter, ValidationError +from starlette.requests import Request +from starlette.responses import JSONResponse +from starlette.types import Message, Send + +from litellm.proxy._experimental.mcp_server.exceptions import MCPUpstreamAuthError + +if TYPE_CHECKING: + from mcp.server.context import CallNext, HandlerResult, ServerRequestContext + +SCOPE_RESPONSE_KEY: Final = "litellm.mcp.scope_response" +_REQUEST: Final = TypeAdapter[dict[str, object]](dict[str, object]) + + +def is_tool_call_request(body: bytes) -> bool: + try: + return _REQUEST.validate_json(body).get("method") == "tools/call" + except ValidationError: + return False + + +class OAuthScopeResponse: + def __init__(self, send: Send, base_url: str, request_path: str | None) -> None: + self._send = send + self._base_url = base_url + self._request_path = request_path + self._ready = asyncio.Event() + self._lock = asyncio.Lock() + self._start: Message | None = None + self._challenge: MCPUpstreamAuthError | None = None + self._refused = False + + @property + def pending(self) -> bool: + return not self._ready.is_set() + + def allow(self) -> None: + self._ready.set() + + def deny(self, error: MCPUpstreamAuthError) -> None: + if self.pending and error.status_code == 403 and error.required_scope is not None: + self._challenge = error + self._ready.set() + + async def __call__(self, message: Message) -> None: + if message["type"] == "http.response.start": + self._start = message + if message["status"] != 200: + self.allow() + return + await self._ready.wait() + async with self._lock: + if self._refused: + return + if self._start is not None: + start: Final = self._start + self._start = None + if self._challenge is not None: + error: Final = self._challenge.to_http_exception(self._base_url, self._request_path) + response: Final = JSONResponse( + status_code=error.status_code, content={"detail": error.detail}, headers=error.headers + ) + self._refused = True + await self._send( + {"type": "http.response.start", "status": response.status_code, "headers": response.raw_headers} + ) + await self._send({"type": "http.response.body", "body": bytes(response.body)}) + return + await self._send(start) + await self._send(message) + + +def scope_response_from_request(request: object) -> OAuthScopeResponse | None: + if not isinstance(request, Request): + return None + response: Final = request.scope.get(SCOPE_RESPONSE_KEY) + return response if isinstance(response, OAuthScopeResponse) else None + + +async def finish_scope_response(context: ServerRequestContext[object, object], call_next: CallNext) -> HandlerResult: + response: Final = scope_response_from_request(context.request) + try: + return await call_next(context) + finally: + if response is not None: + response.allow() diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 1916a19a2c4..097ec6e1ee9 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -72,6 +72,12 @@ from litellm.proxy._experimental.mcp_server.oauth_utils import ( redact_mcp_resource_url, well_known_root_suffix, ) +from litellm.proxy._experimental.mcp_server.scope_challenge import ( + SCOPE_RESPONSE_KEY, + OAuthScopeResponse, + finish_scope_response, + is_tool_call_request, +) from litellm.proxy._experimental.mcp_server.ui_session_utils import ( ActingUser, GrantedToolsetIds, @@ -604,7 +610,7 @@ if MCP_AVAILABLE: name=LITELLM_MCP_SERVER_NAME, version=LITELLM_MCP_SERVER_VERSION, ) - server.middleware.append(GatewayVersionPolicy()) + server.middleware.extend((finish_scope_response, GatewayVersionPolicy())) server.create_initialization_options = types.MethodType(_gateway_create_initialization_options, server) sse: Final[SseServerTransport] = SseServerTransport("/sse/messages") @@ -2351,6 +2357,18 @@ if MCP_AVAILABLE: await rejection(scope, receive, send) return + scope_response: Final = ( + OAuthScopeResponse( + send, + get_request_base_url(StarletteRequest(scope)), + get_route_relative_request_path(scope), + ) + if is_tool_call_request(body) or len(body) >= _MCP_ROUTING_PEEK_MAX_BYTES + else None + ) + if scope_response is not None: + scope[SCOPE_RESPONSE_KEY] = scope_response + use_stateful: Final = bool(session_id or is_initialize) target_manager: Final = session_manager_stateful if use_stateful else session_manager_stateless @@ -2468,7 +2486,7 @@ if MCP_AVAILABLE: touch_last_seen=(scope.get("method") or "").upper() != "DELETE", copy_existing_session_auth_context=is_initialize, ) - local_send = send + local_send = scope_response or send if use_stateful and is_initialize: local_send = _wrap_send_with_stateful_session_auth_context( local_send, diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index b04a958cd03..1b8f9484aea 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -10040,6 +10040,44 @@ "schema": { "additionalProperties": true, "properties": { + "audio": { + "default": [], + "description": "Clef-omni only. Up to 4 embedded audio clips, each a base64 data URL or {content_type, base64}", + "items": { + "anyOf": [ + { + "pattern": "^data:audio/[^;,\\s]+;base64,.+$", + "type": "string" + }, + { + "additionalProperties": false, + "properties": { + "base64": { + "description": "Base64-encoded audio bytes", + "minLength": 1, + "title": "Base64", + "type": "string" + }, + "content_type": { + "description": "An audio/* media type, such as audio/wav", + "pattern": "^audio/\\S+$", + "title": "Content Type", + "type": "string" + } + }, + "required": [ + "content_type", + "base64" + ], + "title": "DecisionsBase64Audio", + "type": "object" + } + ] + }, + "maxItems": 4, + "title": "Audio", + "type": "array" + }, "input": { "anyOf": [ { @@ -10345,6 +10383,44 @@ ], "description": "A stable id for your end user. Sent to OpenAI and dropped for other providers", "title": "Safety Identifier" + }, + "videos": { + "default": [], + "description": "Clef-omni only. Up to 2 embedded videos, each a base64 data URL or {content_type, base64}", + "items": { + "anyOf": [ + { + "pattern": "^data:video/[^;,\\s]+;base64,.+$", + "type": "string" + }, + { + "additionalProperties": false, + "properties": { + "base64": { + "description": "Base64-encoded video bytes", + "minLength": 1, + "title": "Base64", + "type": "string" + }, + "content_type": { + "description": "A video/* media type, such as video/mp4", + "pattern": "^video/\\S+$", + "title": "Content Type", + "type": "string" + } + }, + "required": [ + "content_type", + "base64" + ], + "title": "DecisionsBase64Video", + "type": "object" + } + ] + }, + "maxItems": 2, + "title": "Videos", + "type": "array" } }, "required": [ @@ -10419,6 +10495,44 @@ "schema": { "additionalProperties": true, "properties": { + "audio": { + "default": [], + "description": "Clef-omni only. Up to 4 embedded audio clips, each a base64 data URL or {content_type, base64}", + "items": { + "anyOf": [ + { + "pattern": "^data:audio/[^;,\\s]+;base64,.+$", + "type": "string" + }, + { + "additionalProperties": false, + "properties": { + "base64": { + "description": "Base64-encoded audio bytes", + "minLength": 1, + "title": "Base64", + "type": "string" + }, + "content_type": { + "description": "An audio/* media type, such as audio/wav", + "pattern": "^audio/\\S+$", + "title": "Content Type", + "type": "string" + } + }, + "required": [ + "content_type", + "base64" + ], + "title": "DecisionsBase64Audio", + "type": "object" + } + ] + }, + "maxItems": 4, + "title": "Audio", + "type": "array" + }, "model": { "description": "A decision model from the proxy model_list", "title": "Model", @@ -10647,6 +10761,44 @@ ], "description": "The thing being judged, such as a support ticket, a document or a chat transcript", "title": "State" + }, + "videos": { + "default": [], + "description": "Clef-omni only. Up to 2 embedded videos, each a base64 data URL or {content_type, base64}", + "items": { + "anyOf": [ + { + "pattern": "^data:video/[^;,\\s]+;base64,.+$", + "type": "string" + }, + { + "additionalProperties": false, + "properties": { + "base64": { + "description": "Base64-encoded video bytes", + "minLength": 1, + "title": "Base64", + "type": "string" + }, + "content_type": { + "description": "A video/* media type, such as video/mp4", + "pattern": "^video/\\S+$", + "title": "Content Type", + "type": "string" + } + }, + "required": [ + "content_type", + "base64" + ], + "title": "DecisionsBase64Video", + "type": "object" + } + ] + }, + "maxItems": 2, + "title": "Videos", + "type": "array" } }, "required": [ @@ -10736,6 +10888,44 @@ "schema": { "additionalProperties": true, "properties": { + "audio": { + "default": [], + "description": "Clef-omni only. Up to 4 embedded audio clips, each a base64 data URL or {content_type, base64}", + "items": { + "anyOf": [ + { + "pattern": "^data:audio/[^;,\\s]+;base64,.+$", + "type": "string" + }, + { + "additionalProperties": false, + "properties": { + "base64": { + "description": "Base64-encoded audio bytes", + "minLength": 1, + "title": "Base64", + "type": "string" + }, + "content_type": { + "description": "An audio/* media type, such as audio/wav", + "pattern": "^audio/\\S+$", + "title": "Content Type", + "type": "string" + } + }, + "required": [ + "content_type", + "base64" + ], + "title": "DecisionsBase64Audio", + "type": "object" + } + ] + }, + "maxItems": 4, + "title": "Audio", + "type": "array" + }, "input": { "anyOf": [ { @@ -11041,6 +11231,44 @@ ], "description": "A stable id for your end user. Sent to OpenAI and dropped for other providers", "title": "Safety Identifier" + }, + "videos": { + "default": [], + "description": "Clef-omni only. Up to 2 embedded videos, each a base64 data URL or {content_type, base64}", + "items": { + "anyOf": [ + { + "pattern": "^data:video/[^;,\\s]+;base64,.+$", + "type": "string" + }, + { + "additionalProperties": false, + "properties": { + "base64": { + "description": "Base64-encoded video bytes", + "minLength": 1, + "title": "Base64", + "type": "string" + }, + "content_type": { + "description": "A video/* media type, such as video/mp4", + "pattern": "^video/\\S+$", + "title": "Content Type", + "type": "string" + } + }, + "required": [ + "content_type", + "base64" + ], + "title": "DecisionsBase64Video", + "type": "object" + } + ] + }, + "maxItems": 2, + "title": "Videos", + "type": "array" } }, "required": [ @@ -11115,6 +11343,44 @@ "schema": { "additionalProperties": true, "properties": { + "audio": { + "default": [], + "description": "Clef-omni only. Up to 4 embedded audio clips, each a base64 data URL or {content_type, base64}", + "items": { + "anyOf": [ + { + "pattern": "^data:audio/[^;,\\s]+;base64,.+$", + "type": "string" + }, + { + "additionalProperties": false, + "properties": { + "base64": { + "description": "Base64-encoded audio bytes", + "minLength": 1, + "title": "Base64", + "type": "string" + }, + "content_type": { + "description": "An audio/* media type, such as audio/wav", + "pattern": "^audio/\\S+$", + "title": "Content Type", + "type": "string" + } + }, + "required": [ + "content_type", + "base64" + ], + "title": "DecisionsBase64Audio", + "type": "object" + } + ] + }, + "maxItems": 4, + "title": "Audio", + "type": "array" + }, "model": { "description": "A decision model from the proxy model_list", "title": "Model", @@ -11343,6 +11609,44 @@ ], "description": "The thing being judged, such as a support ticket, a document or a chat transcript", "title": "State" + }, + "videos": { + "default": [], + "description": "Clef-omni only. Up to 2 embedded videos, each a base64 data URL or {content_type, base64}", + "items": { + "anyOf": [ + { + "pattern": "^data:video/[^;,\\s]+;base64,.+$", + "type": "string" + }, + { + "additionalProperties": false, + "properties": { + "base64": { + "description": "Base64-encoded video bytes", + "minLength": 1, + "title": "Base64", + "type": "string" + }, + "content_type": { + "description": "A video/* media type, such as video/mp4", + "pattern": "^video/\\S+$", + "title": "Content Type", + "type": "string" + } + }, + "required": [ + "content_type", + "base64" + ], + "title": "DecisionsBase64Video", + "type": "object" + } + ] + }, + "maxItems": 2, + "title": "Videos", + "type": "array" } }, "required": [ @@ -37831,6 +38135,22 @@ ], "title": "Resource" } + }, + { + "in": "query", + "name": "scope", + "required": false, + "schema": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Scope" + } } ], "responses": { diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index b6c81f17761..2c75a608d0d 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -32,6 +32,7 @@ from litellm.constants import DEFAULT_STAGGER_WINDOW_SECONDS, MCP_STDIO_ALLOWED_ from litellm.litellm_core_utils.initialize_dynamic_callback_params import ( validate_arize_otlp_protocol_value, validate_capture_message_content_value, + validate_excluded_services_value, validate_langfuse_environment_value, validate_langfuse_span_scope_value, validate_no_callback_env_reference, @@ -90,6 +91,7 @@ from litellm.types.router_weights import validate_router_settings_dict from litellm.types.secret_managers.main import KeyManagementSystem from litellm.types.utils import ( CAPTURE_MESSAGE_CONTENT_VAR, + EXCLUDED_SERVICES_VAR, AzureSpillover, CallTypes, CostBreakdown, @@ -2046,10 +2048,16 @@ class AddTeamCallback(LiteLLMPydanticObjectBase): @classmethod def validate_callback_vars(cls, values): callback_vars: Final = values.get("callback_vars", {}) - valid_keys: Final = {*StandardCallbackDynamicParams.__annotations__, CAPTURE_MESSAGE_CONTENT_VAR} + valid_keys: Final = { + *StandardCallbackDynamicParams.__annotations__, + CAPTURE_MESSAGE_CONTENT_VAR, + EXCLUDED_SERVICES_VAR, + } for key, value in callback_vars.items(): if key not in valid_keys: raise ValueError(f"Invalid callback variable: {key}. Must be one of {valid_keys}") + if key in (CAPTURE_MESSAGE_CONTENT_VAR, EXCLUDED_SERVICES_VAR) and not isinstance(value, str): + raise ValueError(f"{key} must be a string, got {type(value).__name__}") callback_vars[key] = str(value) validate_no_callback_env_reference(key, callback_vars[key], source="key/team callback metadata") if key == "langfuse_environment": @@ -2060,6 +2068,8 @@ class AddTeamCallback(LiteLLMPydanticObjectBase): validate_arize_otlp_protocol_value(callback_vars[key]) if key == CAPTURE_MESSAGE_CONTENT_VAR: validate_capture_message_content_value(callback_vars[key]) + if key == EXCLUDED_SERVICES_VAR: + validate_excluded_services_value(callback_vars[key]) return values diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 56773789059..9e7fa2c28e8 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -1308,6 +1308,20 @@ async def common_checks( valid_token=valid_token, ) + if _model and valid_token is not None and personal_key_team_model_access_applies(valid_token): + await can_personal_key_call_model_via_teams( + model=( + *((_model,) if isinstance(_model, str) else _model), + *request_fallback_model_names(_typed_request_body(request_body)), + ), + user_object=user_object, + valid_token=valid_token, + llm_router=llm_router, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + # 1.1 - 2.2 - 3.0.2 - 3.0.3: Project checks (blocked, model access, budget) with tracer.trace("litellm.proxy.auth.common_checks.run_project_checks"): await _run_project_checks( @@ -5331,6 +5345,31 @@ async def can_key_call_model( raise +async def _personal_key_owner( + valid_token: UserAPIKeyAuth, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + proxy_logging_obj: ProxyLogging, +) -> LiteLLM_UserTable | None: + if valid_token.user_id is None: + return None + try: + return await get_user_object( + user_id=valid_token.user_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + user_id_upsert=False, + parent_otel_span=valid_token.parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + ) + except Exception as e: # noqa: BLE001 # fail closed: an owner that cannot be loaded grants nothing + verbose_proxy_logger.warning( + "Personal key team model access: owner lookup failed, granting nothing: %s", + e, + ) + return None + + async def can_key_call_resolved_model( model: str, llm_model_list: Sequence[object] | None, @@ -5410,6 +5449,23 @@ async def can_key_call_resolved_model( key_model_aliases=key_model_aliases_for_auth_check(valid_token), ) + if personal_key_team_model_access_applies(valid_token) and valid_token.user_id is not None: + owner: Final = await _personal_key_owner( + valid_token=valid_token, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + await can_personal_key_call_model_via_teams( + model=model, + user_object=owner, + valid_token=valid_token, + llm_router=llm_router, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + if valid_token.project_id is not None: project_object: Final = await get_project_object( project_id=valid_token.project_id, @@ -5674,6 +5730,241 @@ async def can_user_call_model( ) +def personal_key_team_model_access_applies(valid_token: UserAPIKeyAuth | None) -> bool: + return ( + bool(litellm.personal_key_model_access_from_teams) + and valid_token is not None + and valid_token.team_id is None + and valid_token.user_id is not None + and valid_token.api_key != LITELLM_PROXY_MASTER_KEY_ALIAS + and valid_token.user_role != LitellmUserRoles.PROXY_ADMIN + ) + + +async def _load_personal_key_team( + team_id: str, + valid_token: UserAPIKeyAuth, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + proxy_logging_obj: ProxyLogging, +) -> LiteLLM_TeamTableCachedObj | None: + try: + team_object: Final = await get_team_object( + team_id=team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=valid_token.parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + ) + except Exception as e: # noqa: BLE001 # fail closed: a team that cannot be loaded grants nothing + verbose_proxy_logger.warning( + "Personal key team model access: team=%s lookup failed, granting nothing: %s", team_id, e + ) + return None + return None if team_object.blocked is True else team_object + + +async def _load_personal_key_teams( + user_object: LiteLLM_UserTable, + valid_token: UserAPIKeyAuth, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + proxy_logging_obj: ProxyLogging, +) -> tuple[LiteLLM_TeamTableCachedObj | None, ...]: + return tuple( + await asyncio.gather( + *( + _load_personal_key_team( + team_id=team_id, + valid_token=valid_token, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + for team_id in dict.fromkeys(user_object.teams) + ) + ) + ) + + +async def _team_grants_personal_key_model( + model: str, + team_object: LiteLLM_TeamTableCachedObj, + valid_token: UserAPIKeyAuth, + llm_router: Router | None, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + proxy_logging_obj: ProxyLogging, +) -> bool: + key_model_aliases: Final = key_model_aliases_for_auth_check(valid_token) + try: + await can_team_access_model( + model=model, + team_object=team_object, + llm_router=llm_router, + key_model_aliases=key_model_aliases, + prisma_client=prisma_client, + ) + await _check_team_member_model_access( + model=model, + team_object=team_object, + valid_token=valid_token, + llm_router=llm_router, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + key_model_aliases=key_model_aliases, + ) + except ProxyException: + return False + except Exception as e: # noqa: BLE001 # fail closed: a team check that errors grants nothing + verbose_proxy_logger.warning( + "Personal key team model access: team=%s check failed, granting nothing: %s", team_object.team_id, e + ) + return False + return True + + +async def _loaded_teams_allow_model( + model: str, + teams: tuple[LiteLLM_TeamTableCachedObj | None, ...], + valid_token: UserAPIKeyAuth, + llm_router: Router | None, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + proxy_logging_obj: ProxyLogging, +) -> bool: + loaded: Final = tuple(team for team in teams if team is not None) + if not loaded: + return False + grants: Final = await asyncio.gather( + *( + _team_grants_personal_key_model( + model=model, + team_object=team, + valid_token=valid_token, + llm_router=llm_router, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + for team in loaded + ) + ) + return any(grants) + + +async def _personal_key_listing_teams( + user_api_key_dict: UserAPIKeyAuth, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache | None, + proxy_logging_obj: ProxyLogging | None, +) -> tuple[LiteLLM_TeamTableCachedObj | None, ...]: + if ( + user_api_key_dict.user_id is None + or prisma_client is None + or user_api_key_cache is None + or proxy_logging_obj is None + ): + return () + owner: Final = await _personal_key_owner( + valid_token=user_api_key_dict, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + if owner is None: + return () + return await _load_personal_key_teams( + user_object=owner, + valid_token=user_api_key_dict, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + + +async def personal_key_team_visible_models( + models: Sequence[str], + user_api_key_dict: UserAPIKeyAuth, + llm_router: Router | None, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache | None, + proxy_logging_obj: ProxyLogging | None, +) -> tuple[str, ...]: + teams: Final = await _personal_key_listing_teams( + user_api_key_dict=user_api_key_dict, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + return tuple( + [ + model + for model in (models if teams else ()) + if ( + prisma_client is not None + and user_api_key_cache is not None + and proxy_logging_obj is not None + and await _loaded_teams_allow_model( + model=model, + teams=teams, + valid_token=user_api_key_dict, + llm_router=llm_router, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + ) + ] + ) + + +async def can_personal_key_call_model_via_teams( + model: str | Sequence[str], + user_object: LiteLLM_UserTable | None, + valid_token: UserAPIKeyAuth, + llm_router: Router | None, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + proxy_logging_obj: ProxyLogging, +) -> Literal[True]: + requested_models: Final = (model,) if isinstance(model, str) else tuple(model) + teams: Final = ( + await _load_personal_key_teams( + user_object=user_object, + valid_token=valid_token, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + if user_object is not None + else () + ) + for requested_model in requested_models: + if user_object is None or not await _loaded_teams_allow_model( + model=requested_model, + teams=teams, + valid_token=valid_token, + llm_router=llm_router, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ): + raise ModelAccessDeniedProxyException( + message=model_access_denied_client_message(model=requested_model), + internal_message=( + f"Personal key model access is derived from the owner's teams. User={valid_token.user_id}, " + f"teams={user_object.teams if user_object is not None else None}, " + f"model={requested_model}" + ), + type=ProxyErrorTypes.key_model_access_denied, + param="model", + code=status.HTTP_403_FORBIDDEN, + ) + return True + + def _search_tool_names_from_object_permission( object_permission: LiteLLM_ObjectPermissionTable | None, ) -> list[str]: diff --git a/litellm/proxy/auth/route_checks.py b/litellm/proxy/auth/route_checks.py index ebfa3211c64..ae129ca47b5 100644 --- a/litellm/proxy/auth/route_checks.py +++ b/litellm/proxy/auth/route_checks.py @@ -7,6 +7,7 @@ from fastapi import HTTPException, Request, status from litellm._logging import verbose_proxy_logger from litellm.proxy._types import ( + UI_TEAM_ID, CommonProxyErrors, KeyManagementRoutes, LiteLLM_UserTable, @@ -502,6 +503,13 @@ class RouteChecks: """ return route in LiteLLMRoutes.info_routes.value + @staticmethod + def is_dashboard_session_info_route(route: str, team_id: str | None) -> bool: + """ + Check if an Admin UI session key is calling an info route the dashboard polls in the background + """ + return team_id == UI_TEAM_ID and RouteChecks.is_info_route(route) + @staticmethod def _is_azure_openai_route(route: str) -> bool: """ diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 2231706b282..cdb3e685c20 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -3671,12 +3671,23 @@ async def run_centralized_common_checks( # may have a non-admin user_role and would otherwise demote the # caller. The token is the source of truth for these paths — force # the admin user_object whenever the token says PROXY_ADMIN, even - # if a DB row was fetched. + # if a DB row was fetched. Every admin token except the built-in admin + # identity (litellm_proxy_admin_name: the master key and the UI + # username/password login) keeps the row's personal budget (max_budget, + # and user_email for budget alerts) so common_checks enforces it; every + # other field keeps its default. if user_api_key_auth_obj.user_role == LitellmUserRoles.PROXY_ADMIN: + budget_row: Final = ( + user_object + if isinstance(user_object, LiteLLM_UserTable) and user_api_key_auth_obj.user_id != litellm_proxy_admin_name + else None + ) user_object = LiteLLM_UserTable( user_id=user_api_key_auth_obj.user_id or litellm_proxy_admin_name, user_role=LitellmUserRoles.PROXY_ADMIN, spend=user_object.spend if user_object is not None else 0.0, + max_budget=budget_row.max_budget if budget_row is not None else None, + user_email=budget_row.user_email if budget_row is not None else None, object_permission_id=( user_object.object_permission_id if isinstance(user_object, LiteLLM_UserTable) else None ), diff --git a/litellm/proxy/common_utils/callback_config_validation.py b/litellm/proxy/common_utils/callback_config_validation.py index 4956da853c6..302978d17a4 100644 --- a/litellm/proxy/common_utils/callback_config_validation.py +++ b/litellm/proxy/common_utils/callback_config_validation.py @@ -6,11 +6,11 @@ per-integration checks here. """ import math -from collections.abc import Mapping, Sequence +from collections.abc import Callable, Mapping, Sequence from types import MappingProxyType -from typing import Final +from typing import Final, cast # noqa: TID251 # preserve arbitrary keys from validated callback mappings -from litellm.types.utils import CAPTURE_MESSAGE_CONTENT_VAR +from litellm.types.utils import CAPTURE_MESSAGE_CONTENT_VAR, EXCLUDED_SERVICES_VAR _NEWRELIC_CALLBACK: Final = "newrelic" _NEWRELIC_VAR_PREFIX: Final = "newrelic_" @@ -38,9 +38,29 @@ def callback_config_error( ) if langfuse_error is not None: return langfuse_error - capture_error: Final = _capture_message_content_error(callback_name, callback_vars, callback_type) - if capture_error is not None: - return capture_error + from litellm.litellm_core_utils.initialize_dynamic_callback_params import ( + validate_capture_message_content_value, + validate_excluded_services_value, + ) + + destination_setting_error: Final = next( + ( + error + for variable_name, validator in ( + (CAPTURE_MESSAGE_CONTENT_VAR, validate_capture_message_content_value), + (EXCLUDED_SERVICES_VAR, validate_excluded_services_value), + ) + if ( + error := _destination_setting_error( + variable_name, callback_name, callback_vars, callback_type, validator + ) + ) + is not None + ), + None, + ) + if destination_setting_error is not None: + return destination_setting_error if callback_name != _NEWRELIC_CALLBACK: return None return _newrelic_config_error(callback_vars) @@ -85,29 +105,30 @@ def _langfuse_span_scope_error(callback_name: str | None, callback_vars: Mapping return None -def _capture_message_content_error( - callback_name: str | None, callback_vars: Mapping[str, str], callback_type: str | None +def _destination_setting_error( + variable_name: str, + callback_name: str | None, + callback_vars: Mapping[str, str], + callback_type: str | None, + validator: Callable[[str], None], ) -> str | None: - value: Final = callback_vars.get(CAPTURE_MESSAGE_CONTENT_VAR) + value: Final = callback_vars.get(variable_name) if value is None: return None from litellm.integrations.otel.model.config import is_otel_v2_enabled from litellm.integrations.otel.presets.destinations import destination_capable_backends - from litellm.litellm_core_utils.initialize_dynamic_callback_params import ( - validate_capture_message_content_value, - ) supported: Final = sorted(destination_capable_backends()) if callback_name not in supported: - return f"{CAPTURE_MESSAGE_CONTENT_VAR} applies to the OTel v2 callbacks {supported} only, not {callback_name!r}" + return f"{variable_name} applies to the OTel v2 callbacks {supported} only, not {callback_name!r}" if not is_otel_v2_enabled(): - return f"Per-destination {CAPTURE_MESSAGE_CONTENT_VAR} requires the proxy to run with LITELLM_OTEL_V2=true." + return f"Per-destination {variable_name} requires the proxy to run with LITELLM_OTEL_V2=true." try: - validate_capture_message_content_value(value) + validator(value) except ValueError as e: return str(e) if callback_type == "failure": - return f"{CAPTURE_MESSAGE_CONTENT_VAR} needs callback_type 'success' or 'success_and_failure'" + return f"{variable_name} needs callback_type 'success' or 'success_and_failure'" return None @@ -221,25 +242,36 @@ def conflicting_span_scope_error( ) -def conflicting_capture_error( +def conflicting_callback_var_error( callback_name: str, + variable_name: str, callback_vars: Mapping[str, str] | None, stored_entries: Sequence[tuple[str | None, Mapping[str, str]]], ) -> str | None: - incoming: Final = None if callback_vars is None else callback_vars.get(CAPTURE_MESSAGE_CONTENT_VAR) + incoming: Final = None if callback_vars is None else callback_vars.get(variable_name) if incoming is None: return None return next( ( - f"{CAPTURE_MESSAGE_CONTENT_VAR} is already set to {stored!r} by another {callback_name} entry. " + f"{variable_name} is already set to {stored!r} by another {callback_name} entry. " f"Every {callback_name} entry shares one value: remove that entry or send the same value." for name, entry_vars in stored_entries - if name == callback_name and (stored := entry_vars.get(CAPTURE_MESSAGE_CONTENT_VAR)) not in (None, incoming) + if name == callback_name + and (stored := entry_vars.get(variable_name)) is not None + and not _same_callback_var_value(variable_name, stored, incoming) ), None, ) +def _same_callback_var_value(variable_name: str, first: str, second: str) -> bool: + if variable_name != EXCLUDED_SERVICES_VAR: + return first == second + from litellm.integrations.otel.model.config import parse_excluded_services + + return parse_excluded_services(first) == parse_excluded_services(second) + + def logging_metadata_config_error(metadata: Mapping[str, object] | None) -> str | None: """Validate every ``logging`` entry of a team/key metadata payload.""" if not metadata: @@ -247,15 +279,20 @@ def logging_metadata_config_error(metadata: Mapping[str, object] | None) -> str entries: Final = _object_sequence(metadata.get("logging")) if entries is None: return None + entry_error: Final = next( + (error for entry in entries if (error := _logging_entry_error(entry)) is not None), + None, + ) + if entry_error is not None: + return entry_error entry_vars: Final = tuple(_entry_callback_vars(entry) for entry in entries) named_vars: Final = tuple(zip((_entry_callback_name(entry) for entry in entries), entry_vars)) return next( ( error for error in ( - *(_logging_entry_error(entry) for entry in entries), *(conflicting_span_scope_error(entry_vars[i], entry_vars[:i]) for i in range(len(entry_vars))), - *(_conflicting_entry_capture_error(named_vars[i], named_vars[:i]) for i in range(len(named_vars))), + *(_conflicting_entry_setting_error(named_vars[i], named_vars[:i]) for i in range(len(named_vars))), ) if error is not None ), @@ -263,8 +300,8 @@ def logging_metadata_config_error(metadata: Mapping[str, object] | None) -> str ) -def stored_capture_entries(logging_entries: object) -> tuple[tuple[str | None, Mapping[str, str]], ...]: - """The callback name and vars of every stored ``logging`` entry, for ``conflicting_capture_error``.""" +def stored_callback_entries(logging_entries: object) -> tuple[tuple[str | None, Mapping[str, str]], ...]: + """The callback name and vars of every stored ``logging`` entry.""" entries: Final = _object_sequence(logging_entries) or () return tuple((_entry_callback_name(entry), _entry_callback_vars(entry)) for entry in entries) @@ -281,11 +318,21 @@ def _entry_callback_name(entry: object) -> str | None: return None -def _conflicting_entry_capture_error( +def _conflicting_entry_setting_error( entry: tuple[str | None, Mapping[str, str]], earlier: Sequence[tuple[str | None, Mapping[str, str]]] ) -> str | None: callback_name, callback_vars = entry - return None if callback_name is None else conflicting_capture_error(callback_name, callback_vars, earlier) + if callback_name is None: + return None + return next( + ( + error + for variable_name in (CAPTURE_MESSAGE_CONTENT_VAR, EXCLUDED_SERVICES_VAR) + if (error := conflicting_callback_var_error(callback_name, variable_name, callback_vars, earlier)) + is not None + ), + None, + ) def _entry_callback_vars(entry: object) -> Mapping[str, str]: @@ -298,12 +345,33 @@ def _entry_callback_vars(entry: object) -> Mapping[str, str]: def _logging_entry_error(entry: object) -> str | None: if not isinstance(entry, Mapping): return None - callback_name: Final = entry.get("callback_name") - if not isinstance(callback_name, str) or not isinstance(entry.get("callback_vars"), Mapping): + typed_entry: Final[Mapping[object, object]] = ( + cast( # cast-ok: entry is Mapping-checked but may have non-string keys + Mapping[object, object], entry + ) + ) + callback_name: Final = typed_entry.get("callback_name") + callback_vars_raw: Final = typed_entry.get("callback_vars") + if not isinstance(callback_name, str) or not isinstance(callback_vars_raw, Mapping): return None - callback_type_raw: Final = entry.get("callback_type") + callback_vars: Final[Mapping[object, object]] = cast( # cast-ok: Mapping check preserves arbitrary key types + Mapping[object, object], callback_vars_raw + ) + destination_var_error: Final = next( + ( + f"{variable_name} must be a string, got {type(callback_vars[variable_name]).__name__}" + for variable_name in (CAPTURE_MESSAGE_CONTENT_VAR, EXCLUDED_SERVICES_VAR) + if variable_name in callback_vars and not isinstance(callback_vars[variable_name], str) + ), + None, + ) + if destination_var_error is not None: + return destination_var_error + callback_type_raw: Final = typed_entry.get("callback_type") return callback_config_error( - callback_name, _entry_callback_vars(entry), callback_type_raw if isinstance(callback_type_raw, str) else None + callback_name, + _entry_callback_vars(typed_entry), + callback_type_raw if isinstance(callback_type_raw, str) else None, ) diff --git a/litellm/proxy/decisions_endpoints/endpoints.py b/litellm/proxy/decisions_endpoints/endpoints.py index dbb7b69bd8b..6b6b7655719 100644 --- a/litellm/proxy/decisions_endpoints/endpoints.py +++ b/litellm/proxy/decisions_endpoints/endpoints.py @@ -92,8 +92,11 @@ async def _invalid_request( ) +_FIELDS_CHECKED_AFTER_ROUTING: Final = frozenset({"safety_identifier", "audio", "videos"}) + + def _fields_checked_before_routing(data: Mapping[str, object]) -> Mapping[str, object]: - return MappingProxyType({key: value for key, value in data.items() if key != "safety_identifier"}) + return MappingProxyType({key: value for key, value in data.items() if key not in _FIELDS_CHECKED_AFTER_ROUTING}) async def _request_data(request: Request, user_api_key_dict: UserAPIKeyAuth) -> dict[str, object]: diff --git a/litellm/proxy/guardrails/guardrail_hooks/decision_model/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/decision_model/__init__.py index 002b9a2cdc9..a196b5ad303 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/decision_model/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/decision_model/__init__.py @@ -60,6 +60,15 @@ def _streaming_sampling_rate(litellm_params: LitellmParams) -> int: return value +def _streaming_buffer_until_moderated(litellm_params: LitellmParams) -> bool | None: + value: Final[object] = getattr(litellm_params, "streaming_buffer_until_moderated", None) + if value is None: + return None + if not isinstance(value, bool): + raise TypeError(f"decision_model guardrail streaming_buffer_until_moderated must be a boolean, got {value!r}") + return value + + def initialize_guardrail(litellm_params: LitellmParams, guardrail: Guardrail) -> DecisionModelGuardrail: import litellm @@ -85,6 +94,7 @@ def initialize_guardrail(litellm_params: LitellmParams, guardrail: Guardrail) -> max_concurrent_decision_calls=litellm_params.max_concurrent_decision_calls or 8, streaming_end_of_stream_only=_streaming_end_of_stream_only(litellm_params), streaming_sampling_rate=_streaming_sampling_rate(litellm_params), + streaming_buffer_until_moderated=_streaming_buffer_until_moderated(litellm_params), ) litellm.logging_callback_manager.add_litellm_callback( # pyright: ignore[reportUnknownMemberType] # callback manager is untyped _callback diff --git a/litellm/proxy/guardrails/guardrail_hooks/decision_model/decision_model.py b/litellm/proxy/guardrails/guardrail_hooks/decision_model/decision_model.py index c2d022850f9..867b1383682 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/decision_model/decision_model.py +++ b/litellm/proxy/guardrails/guardrail_hooks/decision_model/decision_model.py @@ -113,6 +113,7 @@ class DecisionModelGuardrail(CustomGuardrail): max_concurrent_decision_calls: int = 8, streaming_end_of_stream_only: bool = False, streaming_sampling_rate: int = 5, + streaming_buffer_until_moderated: bool | None = None, ) -> None: super().__init__( # pyright: ignore[reportUnknownMemberType] # base init takes untyped **kwargs guardrail_name=guardrail_name, @@ -131,6 +132,8 @@ class DecisionModelGuardrail(CustomGuardrail): self._max_concurrent_decision_calls = max_concurrent_decision_calls self.streaming_end_of_stream_only = streaming_end_of_stream_only self.streaming_sampling_rate = streaming_sampling_rate + if streaming_buffer_until_moderated is not None: + self.streaming_buffer_until_moderated = streaming_buffer_until_moderated self._semaphores_by_loop: dict[AbstractEventLoop, Semaphore] = {} # mutable-ok: per-loop semaphore registry @classmethod diff --git a/litellm/proxy/guardrails/guardrail_hooks/zscaler_ai_guard/zscaler_ai_guard.py b/litellm/proxy/guardrails/guardrail_hooks/zscaler_ai_guard/zscaler_ai_guard.py index 9380a539ecd..e8f0dc68d83 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/zscaler_ai_guard/zscaler_ai_guard.py +++ b/litellm/proxy/guardrails/guardrail_hooks/zscaler_ai_guard/zscaler_ai_guard.py @@ -4,6 +4,7 @@ # # +-------------------------------------------------------------+ import os +from collections.abc import Mapping from typing import TYPE_CHECKING, Final, Literal, Optional from fastapi import HTTPException @@ -105,6 +106,11 @@ class ZscalerAIGuard(CustomGuardrail): super().update_in_memory_litellm_params(litellm_params) self.timeout = self._resolve_timeout(litellm_params.timeout) + @staticmethod + def _resolve_request_user(request_data: Mapping[str, object]) -> str | None: + request_user: Final = request_data.get("user") + return request_user if isinstance(request_user, str) and request_user else None + @staticmethod def _resolve_metadata_value(request_data: dict | None, key: str) -> str | None: """ @@ -207,6 +213,8 @@ class ZscalerAIGuard(CustomGuardrail): ) verbose_proxy_logger.debug("inside apply_guardrail kwargs: %s", kwargs) + user_for_payload: Final = self._resolve_request_user(request_data) + zscaler_ai_guard_result = None direction: Final = "OUT" if input_type == "response" else "IN" verbose_proxy_logger.debug("direction: %s", direction) @@ -219,6 +227,7 @@ class ZscalerAIGuard(CustomGuardrail): policy_id=policy_id, direction=direction, content=concatenated_text, + user=user_for_payload, **kwargs, ) verbose_proxy_logger.debug("response from zscaler ai guards: %s", zscaler_ai_guard_result) @@ -357,7 +366,14 @@ class ZscalerAIGuard(CustomGuardrail): raise HTTPException(status_code=response.status_code, detail=user_facing_error) async def make_zscaler_ai_guard_api_call( - self, zscaler_ai_guard_url, api_key, policy_id, direction, content, **kwargs + self, + zscaler_ai_guard_url: str, + api_key: str | None, + policy_id, + direction: str, + content: str, + user: str | None = None, + **kwargs, ): """ Makes an API call to the Zscaler AI Guard service and handles retries, errors, and response parsing. @@ -374,6 +390,8 @@ class ZscalerAIGuard(CustomGuardrail): # the policy from headers (e.g., user-api-key-alias) if policy_id is not None and policy_id >= 1: data["policyId"] = policy_id + if user: + data["user"] = user try: response: Final = await self._send_request(zscaler_ai_guard_url, extra_headers, data) return self._handle_response(response, direction) diff --git a/litellm/proxy/hooks/batch_rate_limiter.py b/litellm/proxy/hooks/batch_rate_limiter.py index 17f8bd277dd..5f214816578 100644 --- a/litellm/proxy/hooks/batch_rate_limiter.py +++ b/litellm/proxy/hooks/batch_rate_limiter.py @@ -1128,6 +1128,24 @@ class _PROXY_BatchRateLimiter(CustomLogger): }, ) + @staticmethod + async def _check_managed_file_ownership( + managed_files_obj: object, + file_id: str, + user_api_key_dict: UserAPIKeyAuth, + ) -> None: + """Deny before the input file is read, so a caller cannot make the proxy + download, count, or charge its own limits for another tenant's file.""" + from litellm.proxy.openai_files_endpoints.common_utils import ManagedResourceAccessChecker + + if not isinstance(managed_files_obj, ManagedResourceAccessChecker): + raise HTTPException(status_code=500, detail="Managed resource ownership validation is unavailable.") + if not await managed_files_obj.can_user_call_unified_file_id(file_id, user_api_key_dict): + raise HTTPException( + status_code=403, + detail=f"User {user_api_key_dict.user_id} does not have access to the file {file_id}", + ) + async def _fetch_managed_file_content( self, file_id: str, @@ -1170,6 +1188,8 @@ class _PROXY_BatchRateLimiter(CustomLogger): if llm_router is None: raise ValueError("llm_router not available. Cannot access managed files.") + await self._check_managed_file_ownership(managed_files_obj, file_id, user_api_key_dict) + # Use the managed files hook to get file content # This properly handles user permissions and file ownership file_content: Final = await managed_files_obj.afile_content( diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index 0938d3f35c3..5c31bce6f33 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -186,6 +186,8 @@ class _ProxyDBLogger(CustomLogger): and not ( RouteChecks.is_llm_api_route(route=request_route) or RouteChecks.is_info_route(route=request_route) ) + or request_route is not None + and RouteChecks.is_dashboard_session_info_route(request_route, user_api_key_dict.team_id) ): return diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index fdc2e6d94b8..ba319be01cb 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -228,6 +228,7 @@ from litellm.types.llms.anthropic import ANTHROPIC_API_HEADERS from litellm.types.services import ServiceTypes from litellm.types.utils import ( CAPTURE_MESSAGE_CONTENT_VAR, + EXCLUDED_SERVICES_VAR, CustomPricingLiteLLMParams, LlmProviders, ProviderSpecificHeader, @@ -236,6 +237,8 @@ from litellm.types.utils import ( SupportedCacheControls, ) +_DESTINATION_ONLY_CALLBACK_VARS: Final[frozenset[str]] = frozenset({CAPTURE_MESSAGE_CONTENT_VAR, EXCLUDED_SERVICES_VAR}) + _CALLBACK_CREDENTIAL_KEYS: Final = frozenset(StandardCallbackDynamicParams.__annotations__) | frozenset( {TRUSTED_CALLBACK_VARS_FIELD} ) @@ -976,7 +979,7 @@ def convert_key_logging_metadata_to_callback( team_callback_settings_obj.callbacks.append(data.callback_name) for var, value in data.callback_vars.items(): - if var == CAPTURE_MESSAGE_CONTENT_VAR: + if var in _DESTINATION_ONLY_CALLBACK_VARS: continue # New Relic routing reads these from the trusted-vars overlay with no # callback-name check, so scope them to the newrelic entry: a team that @@ -1195,7 +1198,7 @@ def resolve_tenant_otel_destinations( def _tenant_destination( name: str, callbacks: Sequence[AddTeamCallback], service_name: str | None ) -> "OtelDestination | None": - from litellm.integrations.otel.model.config import parse_capture_message_content + from litellm.integrations.otel.model.config import parse_capture_message_content, parse_excluded_services from litellm.integrations.otel.presets.destinations import destination_for callback_vars: Final = MappingProxyType( @@ -1210,6 +1213,7 @@ def _tenant_destination( _tenant_otel_params(callback_vars), service_name, parse_capture_message_content(callback_vars.get(CAPTURE_MESSAGE_CONTENT_VAR)), + parse_excluded_services(callback_vars.get(EXCLUDED_SERVICES_VAR)), ) diff --git a/litellm/proxy/management_endpoints/cost_tracking_settings.py b/litellm/proxy/management_endpoints/cost_tracking_settings.py index 469d112675e..32a0069e0aa 100644 --- a/litellm/proxy/management_endpoints/cost_tracking_settings.py +++ b/litellm/proxy/management_endpoints/cost_tracking_settings.py @@ -20,6 +20,7 @@ import litellm from litellm._internal_context import current_billing_time, pinned_billing_time from litellm._logging import verbose_proxy_logger from litellm.cost_calculator import completion_cost +from litellm.litellm_core_utils.cost_discount import parse_cost_discount_key from litellm.proxy._types import ( CommonProxyErrors, CostEstimateRequest, @@ -238,11 +239,14 @@ async def update_cost_discount_config( Updates the cost_discount_config in litellm_settings. Discounts should be between 0 and 1 (e.g., 0.05 = 5% discount). + A key may also be /, where the pattern is an + fnmatch glob matched against the model name. Example: ```json { "vertex_ai": 0.05, + "vertex_ai/claude-*": 0.20, "gemini": 0.05, "openai": 0.01 } @@ -267,10 +271,16 @@ async def update_cost_discount_config( ) # Validate that all providers are valid LiteLLM providers - invalid_providers: Final = [] - for provider in cost_discount_config: - if provider not in LlmProvidersSet: - invalid_providers.append(provider) + parsed_keys: Final = tuple((key, parse_cost_discount_key(key)) for key in cost_discount_config) + empty_pattern_key: Final = next((key for key, parsed in parsed_keys if parsed.model_pattern == ""), None) + if empty_pattern_key is not None: + raise HTTPException( + status_code=400, + detail=f"Model pattern for {empty_pattern_key} cannot be empty", + ) + invalid_providers: Final = tuple( + parsed.provider for _, parsed in parsed_keys if parsed.provider not in LlmProvidersSet + ) if invalid_providers: raise HTTPException( diff --git a/litellm/proxy/management_endpoints/fallback_management_endpoints.py b/litellm/proxy/management_endpoints/fallback_management_endpoints.py index 460b282a9f6..6a9a30a7689 100644 --- a/litellm/proxy/management_endpoints/fallback_management_endpoints.py +++ b/litellm/proxy/management_endpoints/fallback_management_endpoints.py @@ -11,10 +11,10 @@ DELETE /fallback/{model} - Delete fallbacks for a specific model # pyright: reportMissingImports=false import json -from collections.abc import Mapping +from collections.abc import Mapping, Sequence from typing import TYPE_CHECKING, Annotated, Final, Literal -from pydantic import Field, TypeAdapter +from pydantic import Field, JsonValue, TypeAdapter from litellm._logging import verbose_proxy_logger from litellm.proxy._types import UserAPIKeyAuth @@ -47,6 +47,7 @@ ROUTER_SETTINGS_PARAM: Final = "router_settings" FallbackRule = dict[str, list[str]] StoredFallback = Annotated[FallbackRule | dict[str, object] | str, Field(union_mode="left_to_right")] _STORED_FALLBACKS: Final = TypeAdapter(list[StoredFallback]) +_JSON_RULES: Final = TypeAdapter(list[JsonValue]) def _rule_covers(entry: StoredFallback, model: str) -> bool: @@ -59,6 +60,12 @@ async def _router_settings_fresh_from_db(proxy_config: "ProxyConfig") -> dict[st return config.get(ROUTER_SETTINGS_PARAM, {}) +def _reject_config_owned_rules(proxy_config: "ProxyConfig", fallback_key: str, rules: Sequence[StoredFallback]) -> None: + proxy_config.reject_config_owned_writes( + section_name=ROUTER_SETTINGS_PARAM, changed_keys={fallback_key: _JSON_RULES.validate_python(rules)} + ) + + async def _persist_router_settings(prisma_client: PrismaClient, router_settings: Mapping[str, object]) -> None: router_settings_json: Final = json.dumps(router_settings) await ConfigRepository(prisma_client).table.upsert( @@ -180,6 +187,8 @@ async def create_fallback( # Add new fallback existing_fallbacks.append({data.model: data.fallback_models}) + _reject_config_owned_rules(proxy_config, fallback_key, existing_fallbacks) + # Update router settings router_settings[fallback_key] = existing_fallbacks @@ -323,21 +332,16 @@ async def delete_fallback( # Get existing fallbacks existing_fallbacks: Final = _STORED_FALLBACKS.validate_python(router_settings.get(fallback_key) or []) - # Find and remove the fallback configuration - fallback_found = False - updated_fallbacks: Final = [] - for rule in existing_fallbacks: - if not _rule_covers(rule, model): - updated_fallbacks.append(rule) - else: - fallback_found = True + updated_fallbacks: Final = [rule for rule in existing_fallbacks if not _rule_covers(rule, model)] - if not fallback_found: + if len(updated_fallbacks) == len(existing_fallbacks): raise HTTPException( status_code=status.HTTP_404_NOT_FOUND, detail={"error": f"No {fallback_type} fallbacks configured for model '{model}'"}, ) + _reject_config_owned_rules(proxy_config, fallback_key, updated_fallbacks) + # Update router settings router_settings[fallback_key] = updated_fallbacks diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 782a7e1db3c..be70136e4fe 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -122,6 +122,7 @@ from litellm.proxy.management_helpers.access_group_key_sync import ( sync_key_update_access_group_membership, ) from litellm.proxy.management_helpers.key_settings_audit import with_settings_updated_at +from litellm.proxy.management_helpers.managed_resource_owner_sync import repoint_key_owned_managed_resources from litellm.proxy.management_helpers.object_permission_utils import ( # noqa: F401 # legacy module exports ObjectPermissionUpsert, _set_object_permission, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export @@ -5957,6 +5958,13 @@ async def _execute_virtual_key_regeneration( await evict_and_broadcast(cache_keys=jwt_mapping_cache_keys, user_api_key_cache=user_api_key_cache) + await repoint_key_owned_managed_resources( + prisma_client=prisma_client, + existing_key_row=key_in_db, + previous_key_token=hashed_api_key, + new_key_token=new_token_hash, + ) + # After credential invalidation, so a failure here can never keep the old key alive. await sync_key_regeneration_access_group_membership( prisma_client=prisma_client, diff --git a/litellm/proxy/management_endpoints/team_callback_endpoints.py b/litellm/proxy/management_endpoints/team_callback_endpoints.py index 3908ab3c927..c2d7033c06c 100644 --- a/litellm/proxy/management_endpoints/team_callback_endpoints.py +++ b/litellm/proxy/management_endpoints/team_callback_endpoints.py @@ -31,10 +31,10 @@ from litellm.proxy._types import ( from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_utils.callback_config_validation import ( callback_config_error, - conflicting_capture_error, + conflicting_callback_var_error, conflicting_span_scope_error, cross_entry_family_error, - stored_capture_entries, + stored_callback_entries, ) from litellm.proxy.common_utils.callback_utils import ( # noqa: F401 # legacy module exports _CALLBACK_VAR_ENCRYPTED_PREFIX, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export @@ -56,6 +56,7 @@ from litellm.proxy.management_endpoints.team_endpoints import ( # noqa: F401 # ) from litellm.proxy.management_helpers.utils import management_endpoint_wrapper from litellm.repositories.team_repository import TeamRepository +from litellm.types.utils import CAPTURE_MESSAGE_CONTENT_VAR, EXCLUDED_SERVICES_VAR router: Final = APIRouter() @@ -293,6 +294,7 @@ async def add_team_callbacks( - capture_message_content: For OTel v2 callbacks, "no_content" or "span_only". An explicit value overrides the global setting for this destination; omitted follows the proxy's own callback for that backend, else the global setting. "span_only" puts prompt and response content on the destination's spans, "no_content" leaves it out + - excluded_services: For OTel v2 callbacks, a comma-separated list of datastore names such as "redis,postgres", or "none" to exclude nothing. An explicit value replaces the global exclusions; omitted follows them - gcs_bucket_name: The name of the GCS bucket - gcs_path_service_account: The path to the GCS service account - langsmith_api_key: The API key for the Langsmith callback @@ -356,19 +358,27 @@ async def add_team_callbacks( # Decrypted, because the checks compare the incoming values against # the stored ones and the credentials are encrypted at rest. decrypted_logging: Final = decrypt_callback_vars(team_metadata).get("logging") - stored_capture: Final = stored_capture_entries(decrypted_logging) + stored_callbacks: Final = stored_callback_entries(decrypted_logging) stored_entries: Final = decrypted_logging if isinstance(decrypted_logging, list) else () stored_entry_vars: Final = [entry.get("callback_vars") or {} for entry in stored_entries] scope_error: Final = conflicting_span_scope_error(data.callback_vars, stored_entry_vars) if scope_error is not None: raise _callback_config_error(scope_error) - capture_error: Final = conflicting_capture_error( - data.callback_name, - data.callback_vars, - stored_capture, + destination_setting_error: Final = next( + ( + error + for variable_name in (CAPTURE_MESSAGE_CONTENT_VAR, EXCLUDED_SERVICES_VAR) + if ( + error := conflicting_callback_var_error( + data.callback_name, variable_name, data.callback_vars, stored_callbacks + ) + ) + is not None + ), + None, ) - if capture_error is not None: - raise _callback_config_error(capture_error) + if destination_setting_error is not None: + raise _callback_config_error(destination_setting_error) # One entry has to own a credential family end to end. The entries are # flattened into one dict before a request reads them, so an entry # naming only a destination would pair with a key written on another diff --git a/litellm/proxy/management_helpers/managed_resource_owner_sync.py b/litellm/proxy/management_helpers/managed_resource_owner_sync.py new file mode 100644 index 00000000000..4171205b52b --- /dev/null +++ b/litellm/proxy/management_helpers/managed_resource_owner_sync.py @@ -0,0 +1,44 @@ +"""Carry key-owned managed resources over to a regenerated key's new hash.""" + +from __future__ import annotations + +from typing import Final, Protocol + +from litellm.proxy._types import LiteLLM_VerificationToken +from litellm.proxy.db.routing_prisma_wrapper import writer_wrapper +from litellm.repositories.table_repositories import ManagedObjectRepository + +_KEY_OWNED_TABLES: Final = ( + "LiteLLM_ManagedFileTable", + "LiteLLM_ManagedObjectTable", + "LiteLLM_ManagedVectorStoreTable", +) + + +class _RawWriter(Protocol): + async def execute_raw(self, query: str, *args: str) -> int: ... + + +def _raw_writer(prisma_client: object) -> _RawWriter: + db: Final = ManagedObjectRepository(prisma_client).prisma_client.db # pyright: ignore[reportAny] # untyped Prisma client + return writer_wrapper(db) # pyright: ignore[reportAny, reportReturnType] # untyped Prisma client behind the pin + + +async def repoint_key_owned_managed_resources( + prisma_client: object, + existing_key_row: LiteLLM_VerificationToken, + previous_key_token: str, + new_key_token: str, +) -> None: + """A key with neither a user_id nor a team_id owns its managed files, batches and + vector stores as ``key:`` (``resolve_resource_owner_id``), so a + regenerated key would otherwise be denied its own resources.""" + if existing_key_row.user_id is not None or existing_key_row.team_id is not None: + return + writer: Final = _raw_writer(prisma_client) + for table in _KEY_OWNED_TABLES: + await writer.execute_raw( + f'UPDATE "{table}" SET "created_by" = $2 WHERE "created_by" = $1', + f"key:{previous_key_token}", + f"key:{new_key_token}", + ) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 1c7cf54ab4d..7fe175131d8 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -18845,6 +18845,14 @@ _GENERAL_SETTINGS_UI_LITELLM_FIELDS: Final[dict[str, GeneralSettingsUILiteLLMFie "forgiving it. Applies to key, user, team, team member, org, tag and end-user budgets." ), }, + "personal_key_model_access_from_teams": { + "type": "Boolean", + "description": ( + "Caps each personal key (a key with no team) to the models its owner's current teams can call, " + "re-evaluated on every request. A user in no team gets no model access. A user in several teams " + "can call any model one of those teams can call. Team keys are unaffected." + ), + }, "max_ui_session_budget": { "type": "Dollar", "default": 1.0, diff --git a/litellm/proxy/public_endpoints/provider_create_fields.json b/litellm/proxy/public_endpoints/provider_create_fields.json index 633e9edf71c..cb7d099cafb 100644 --- a/litellm/proxy/public_endpoints/provider_create_fields.json +++ b/litellm/proxy/public_endpoints/provider_create_fields.json @@ -1033,6 +1033,34 @@ ], "default_model_placeholder": "gpt-3.5-turbo" }, + { + "provider": "CORALBRICKS", + "provider_display_name": "CoralBricks", + "litellm_provider": "coralbricks", + "credential_fields": [ + { + "key": "api_base", + "label": "API Base", + "placeholder": "https://inference.coralbricks.ai/v1", + "tooltip": null, + "required": false, + "field_type": "text", + "options": null, + "default_value": null + }, + { + "key": "api_key", + "label": "API Key", + "placeholder": null, + "tooltip": null, + "required": true, + "field_type": "password", + "options": null, + "default_value": null + } + ], + "default_model_placeholder": "coralbricks/glm-5.3-fast" + }, { "provider": "CORTECS", "provider_display_name": "Cortecs", diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 302f55eadf4..5fa78a44d09 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -3356,6 +3356,7 @@ class ProxyLogging: original_exception=original_exception, error_type=error_type, route=user_api_key_dict.request_route, + team_id=user_api_key_dict.team_id, ): await self._handle_logging_proxy_only_error( request_data=request_data, @@ -3413,6 +3414,7 @@ class ProxyLogging: original_exception: Exception, error_type: ProxyErrorTypes | None = None, route: str | None = None, + team_id: str | None = None, ) -> bool: """ Return True if the error is a Proxy Only LLM API Error @@ -3437,6 +3439,9 @@ class ProxyLogging: return False if not (RouteChecks.is_llm_api_route(route) or RouteChecks.is_info_route(route)): return False + # Dashboard session keys poll info routes in the background; their failures are not user requests. + if RouteChecks.is_dashboard_session_info_route(route, team_id): + return False return isinstance(original_exception, _PROXY_ONLY_LLM_API_ERRORS) or (error_type == ProxyErrorTypes.auth_error) @@ -8967,6 +8972,10 @@ async def get_available_models_for_user( Returns: List of model names available to the user """ + from litellm.proxy.auth.auth_checks import ( + personal_key_team_model_access_applies, + personal_key_team_visible_models, + ) from litellm.proxy.auth.model_checks import ( get_complete_model_list, get_key_models, @@ -9037,7 +9046,7 @@ async def get_available_models_for_user( granted_team_models: Final = (*team_models, *access_group_models) if team_models else team_models # Get complete model list - all_models: Final = get_complete_model_list( + complete_models: Final = get_complete_model_list( key_models=granted_key_models, team_models=granted_team_models, proxy_model_list=proxy_model_list, @@ -9050,6 +9059,20 @@ async def get_available_models_for_user( only_model_access_groups=only_model_access_groups, team_id=effective_team_id, ) + all_models: Final = ( + list( + await personal_key_team_visible_models( + models=complete_models, + user_api_key_dict=user_api_key_dict, + llm_router=llm_router, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + ) + if personal_key_team_model_access_applies(user_api_key_dict) + else complete_models + ) agent_visible: Final = await _agent_access_group_visible_models( user_api_key_dict=user_api_key_dict, diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index b76dd68018f..c2cc876970c 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -500,7 +500,7 @@ class LiteLLMCompletionResponsesConfig: "max_tokens": responses_api_request.get("max_output_tokens"), "stream": stream, "metadata": kwargs.get("metadata"), - "service_tier": kwargs.get("service_tier"), + "service_tier": responses_api_request.get("service_tier"), "safety_identifier": responses_api_request.get("safety_identifier"), "web_search_options": web_search_options, "response_format": response_format, @@ -2402,7 +2402,7 @@ class LiteLLMCompletionResponsesConfig: ), "function": function_dict, "type": "function", - "index": 0, + "index": index, } if provider_specific_fields: diff --git a/litellm/systemone/main.py b/litellm/systemone/main.py index 754b69142f9..4b17f06b562 100644 --- a/litellm/systemone/main.py +++ b/litellm/systemone/main.py @@ -4,7 +4,9 @@ from typing import Final import httpx from litellm.decisions.call import ( + DecisionsAudioClips, DecisionsCall, + DecisionsVideoClips, SystemOneQuestions, SystemOneRequestFields, asend, @@ -22,6 +24,8 @@ def _prepare_systemone_call( state: DecisionsJSON | None, questions: SystemOneQuestions | None, safety_identifier: str | None, + audio: DecisionsAudioClips | None, + videos: DecisionsVideoClips | None, api_key: str | None, api_base: str | None, timeout: float | httpx.Timeout | None, @@ -32,7 +36,7 @@ def _prepare_systemone_call( reject_other_format(model=model, kwargs=kwargs, field_name="input", other_function="decisions") return prepare_call( model=model, - fields=SystemOneRequestFields(state=state, questions=questions), + fields=SystemOneRequestFields(state=state, questions=questions, audio=audio, videos=videos), safety_identifier=safety_identifier, api_key=api_key, api_base=api_base, @@ -59,6 +63,8 @@ async def asystemone( custom_llm_provider: str | None = None, extra_headers: Mapping[str, str] | None = None, safety_identifier: str | None = None, + audio: DecisionsAudioClips | None = None, + videos: DecisionsVideoClips | None = None, **kwargs: object, # kwargs-ok: litellm_params forwarded by @client ) -> DecisionsResponse: call: Final = _prepare_systemone_call( @@ -66,6 +72,8 @@ async def asystemone( state=state, questions=questions, safety_identifier=safety_identifier, + audio=audio, + videos=videos, api_key=api_key, api_base=api_base, timeout=timeout, @@ -87,6 +95,8 @@ def systemone( custom_llm_provider: str | None = None, extra_headers: Mapping[str, str] | None = None, safety_identifier: str | None = None, + audio: DecisionsAudioClips | None = None, + videos: DecisionsVideoClips | None = None, **kwargs: object, # kwargs-ok: litellm_params forwarded by @client ) -> DecisionsResponse: call: Final = _prepare_systemone_call( @@ -94,6 +104,8 @@ def systemone( state=state, questions=questions, safety_identifier=safety_identifier, + audio=audio, + videos=videos, api_key=api_key, api_base=api_base, timeout=timeout, diff --git a/litellm/types/decisions.py b/litellm/types/decisions.py index 3bc90a5109a..bef61bee40d 100644 --- a/litellm/types/decisions.py +++ b/litellm/types/decisions.py @@ -3,7 +3,7 @@ from dataclasses import dataclass, field from types import MappingProxyType from typing import Annotated, Final, Literal, TypeAlias -from pydantic import ConfigDict, Field, PrivateAttr, model_validator, with_config +from pydantic import ConfigDict, Field, PrivateAttr, field_validator, model_validator, with_config from typing_extensions import ReadOnly, Required, TypedDict from litellm.types.llms.base import LiteLLMPydanticObjectBase @@ -67,7 +67,49 @@ DecisionQuestionMap: TypeAlias = Annotated[ ] -class DecisionsRequestBody(LiteLLMPydanticObjectBase): +class DecisionsBase64Audio(LiteLLMPydanticObjectBase): + content_type: Annotated[str, Field(pattern=r"^audio/\S+$", description="An audio/* media type, such as audio/wav")] + base64: Annotated[str, Field(min_length=1, description="Base64-encoded audio bytes")] + + model_config = ConfigDict(extra="forbid", frozen=True) + + +class DecisionsBase64Video(LiteLLMPydanticObjectBase): + content_type: Annotated[str, Field(pattern=r"^video/\S+$", description="A video/* media type, such as video/mp4")] + base64: Annotated[str, Field(min_length=1, description="Base64-encoded video bytes")] + + model_config = ConfigDict(extra="forbid", frozen=True) + + +DecisionsAudio: TypeAlias = Annotated[str, Field(pattern=r"^data:audio/[^;,\s]+;base64,.+$")] | DecisionsBase64Audio +DecisionsVideo: TypeAlias = Annotated[str, Field(pattern=r"^data:video/[^;,\s]+;base64,.+$")] | DecisionsBase64Video + + +class DecisionsMediaBody(LiteLLMPydanticObjectBase): + audio: Annotated[ + tuple[DecisionsAudio, ...], + Field( + max_length=4, + description="Clef-omni only. Up to 4 embedded audio clips, each a base64 data URL or {content_type, base64}", + ), + ] = () + videos: Annotated[ + tuple[DecisionsVideo, ...], + Field( + max_length=2, + description="Clef-omni only. Up to 2 embedded videos, each a base64 data URL or {content_type, base64}", + ), + ] = () + + model_config = ConfigDict(extra="allow", frozen=True) + + @field_validator("audio", "videos", mode="before") + @classmethod + def _null_means_no_clips(cls, value: object) -> object: + return () if value is None else value + + +class DecisionsRequestBody(DecisionsMediaBody): state: Annotated[ DecisionsJSON, Field(description="The thing being judged, such as a support ticket, a document or a chat transcript"), @@ -89,6 +131,8 @@ class DecisionsCallParams(TypedDict, total=False): model: Required[ReadOnly[str]] state: Required[ReadOnly[DecisionsJSON]] questions: Required[ReadOnly[DecisionQuestionMap]] + audio: ReadOnly[Sequence[DecisionsAudio]] + videos: ReadOnly[Sequence[DecisionsVideo]] api_key: ReadOnly[str | None] api_base: ReadOnly[str | None] timeout: ReadOnly[float | None] @@ -255,7 +299,7 @@ OpenAIDecisionQuestion: TypeAlias = Annotated[ ] -class OpenAIDecisionRequestBody(LiteLLMPydanticObjectBase): +class OpenAIDecisionRequestBody(DecisionsMediaBody): input: Annotated[ OpenAIDecisionInput, Field( @@ -426,6 +470,8 @@ class DecisionsIRRequest: input: DecisionsIRState | DecisionsIRMessages questions: tuple[DecisionsIRQuestion, ...] safety_identifier: str | None = None + audio: tuple[DecisionsAudio, ...] = () + videos: tuple[DecisionsVideo, ...] = () @dataclass(frozen=True, slots=True) diff --git a/litellm/types/mcp_server/mcp_server_manager.py b/litellm/types/mcp_server/mcp_server_manager.py index a7027102250..9b37a63adb1 100644 --- a/litellm/types/mcp_server/mcp_server_manager.py +++ b/litellm/types/mcp_server/mcp_server_manager.py @@ -1,4 +1,5 @@ import json +from collections.abc import Mapping from datetime import datetime from typing import Annotated, Any, Final, Literal @@ -29,6 +30,7 @@ MCPInfo = Annotated[dict[str, Any], AfterValidator(_validate_mcp_protocol_metada class MCPOAuthMetadata(LiteLLMBaseModel): + authorization_server_metadata: Mapping[str, object] | None = Field(default=None, exclude=True) scopes: list[str] | None = None """Resource-driven scopes for the authorization request: the RFC 9728 protected-resource ``scopes_supported``, or the ``scope`` from the WWW-Authenticate 401 challenge when the resource @@ -117,6 +119,7 @@ class MCPServer(LiteLLMBaseModel): # scope=="user" values must be supplied per-user. env_vars: list[dict[str, Any]] | None = None # OAuth-specific fields + authorization_server_metadata: Mapping[str, object] | None = Field(default=None, exclude=True) client_id: str | None = None client_secret: str | None = None issuer: str | None = None diff --git a/litellm/types/utils.py b/litellm/types/utils.py index f7385a6d92d..06af506076a 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -195,6 +195,7 @@ class ProviderSpecificModelInfo(TypedDict, total=False): supports_thinking_cache_preservation: ReadOnly[bool | None] supports_computer_use: bool | None supports_audio_input: bool | None + supports_video_input: ReadOnly[bool | None] supports_embedding_image_input: bool | None supports_audio_output: bool | None supports_pdf_input: bool | None @@ -3736,6 +3737,7 @@ ARIZE_OTLP_PROTOCOLS: Final[frozenset[str]] = frozenset(get_args(ArizeOtlpProtoc CAPTURE_MESSAGE_CONTENT_VAR: Final = "capture_message_content" +EXCLUDED_SERVICES_VAR: Final = "excluded_services" class CaptureMessageContent(str): @@ -4308,6 +4310,7 @@ class LlmProviders(str, Enum): LIBERTAI = "libertai" PINSTRIPES = "pinstripes" COGNITION = "cognition" + CORALBRICKS = "coralbricks" CORTECS = "cortecs" SCX_AI = "scx-ai" PRISM = "prism" diff --git a/litellm/utils.py b/litellm/utils.py index 44dc75f45d7..3f10d2a3ebd 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -2975,6 +2975,10 @@ def supports_audio_input(model: str, custom_llm_provider: str | None = None) -> return supports_factory(model=model, custom_llm_provider=custom_llm_provider, key="supports_audio_input") +def supports_video_input(model: str, custom_llm_provider: str | None = None) -> bool: + return supports_factory(model=model, custom_llm_provider=custom_llm_provider, key="supports_video_input") + + def supports_pdf_input(model: str, custom_llm_provider: str | None = None) -> bool: """Check if a given model supports pdf input in a chat completion call""" return supports_factory(model=model, custom_llm_provider=custom_llm_provider, key="supports_pdf_input") @@ -5908,9 +5912,18 @@ def get_potential_model_names(model: str, custom_llm_provider: str | None) -> Po provider_prefixed_model_name = combined_model_name if custom_llm_provider in ("bedrock", "bedrock_converse"): - from litellm.llms.bedrock.common_utils import strip_bedrock_routing_prefix + from litellm.llms.bedrock.common_utils import ( + extract_model_name_from_bedrock_arn, + split_bedrock_region_path, + strip_bedrock_routing_prefix, + ) - split_model = strip_bedrock_routing_prefix(split_model) + region_free_bedrock_model: Final = split_bedrock_region_path(split_model)[1] + split_model = ( + extract_model_name_from_bedrock_arn(region_free_bedrock_model) + if region_free_bedrock_model.startswith("arn:") + else strip_bedrock_routing_prefix(split_model) + ) region_free_split_model: Final = ( _strip_mantle_region_prefix(split_model) if custom_llm_provider == "bedrock_mantle" else split_model @@ -6385,6 +6398,7 @@ def get_model_info_helper( supports_prompt_cache_breakpoint=_model_info.get("supports_prompt_cache_breakpoint", None), supports_thinking_cache_preservation=_model_info.get("supports_thinking_cache_preservation", None), supports_audio_input=_model_info.get("supports_audio_input", None), + supports_video_input=_model_info.get("supports_video_input", None), supports_audio_output=_model_info.get("supports_audio_output", None), supports_pdf_input=_model_info.get("supports_pdf_input", None), supports_embedding_image_input=_model_info.get("supports_embedding_image_input", None), @@ -6535,6 +6549,7 @@ def get_model_info( supports_prompt_caching: Optional[bool] supports_prompt_cache_breakpoint: Optional[bool] supports_audio_input: Optional[bool] + supports_video_input: Optional[bool] supports_audio_output: Optional[bool] supports_pdf_input: Optional[bool] supports_web_search: Optional[bool] @@ -9018,6 +9033,8 @@ class ProviderConfigManager: return litellm.TypeSafeDecisionsConfig() if provider == LlmProviders.OPENROUTER: return litellm.OpenRouterDecisionsConfig() + if provider == LlmProviders.FIREWORKS_AI: + return litellm.FireworksAIDecisionsConfig() if provider == LlmProviders.CLOUDFLARE: return litellm.CloudflareDecisionsConfig() if provider == LlmProviders.STRANDS_DECIDER: diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 777b68ac335..3d3846c28c0 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -16022,6 +16022,26 @@ "prompt_cache_min_tokens": 1024, "source": "https://platform.claude.com/docs/en/about-claude/pricing" }, + "cloudflare/clef-omni": { + "input_cost_per_token": 1.5e-07, + "litellm_provider": "cloudflare", + "max_input_tokens": 64000, + "mode": "evaluation", + "output_cost_per_token": 0.0, + "source": "https://developers.cloudflare.com/workers-ai/models/clef-omni/", + "supports_audio_input": true, + "supports_video_input": true + }, + "cloudflare/@cf/cloudflare/clef-omni": { + "input_cost_per_token": 1.5e-07, + "litellm_provider": "cloudflare", + "max_input_tokens": 64000, + "mode": "evaluation", + "output_cost_per_token": 0.0, + "source": "https://developers.cloudflare.com/workers-ai/models/clef-omni/", + "supports_audio_input": true, + "supports_video_input": true + }, "cloudflare/clef": { "input_cost_per_token": 2.4e-07, "litellm_provider": "cloudflare", @@ -16614,6 +16634,144 @@ "supports_function_calling": true, "supports_tool_choice": true }, + "coralbricks/glm-5.3-fast": { + "cache_creation_input_token_cost": 1.68e-06, + "cache_read_input_token_cost": 0.0, + "input_cost_per_token": 1.12e-06, + "litellm_provider": "coralbricks", + "max_input_tokens": 1048576, + "max_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 4.4e-06, + "source": "https://www.coralbricks.ai/pricing", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses", + "/v1/messages" + ], + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": false + }, + "coralbricks/glm-5.3-fp4": { + "cache_creation_input_token_cost": 1.68e-06, + "cache_read_input_token_cost": 0.0, + "input_cost_per_token": 1.12e-06, + "litellm_provider": "coralbricks", + "max_input_tokens": 1048576, + "max_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 4.4e-06, + "source": "https://www.coralbricks.ai/pricing", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses", + "/v1/messages" + ], + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": false + }, + "coralbricks/glm-5.3": { + "cache_creation_input_token_cost": 1.68e-06, + "cache_read_input_token_cost": 0.0, + "input_cost_per_token": 1.12e-06, + "litellm_provider": "coralbricks", + "max_input_tokens": 1048576, + "max_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 4.4e-06, + "source": "https://www.coralbricks.ai/pricing", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses", + "/v1/messages" + ], + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": false + }, + "coralbricks/deepseek-v4.1-flash-fast": { + "cache_creation_input_token_cost": 9e-08, + "cache_read_input_token_cost": 0.0, + "input_cost_per_token": 1e-08, + "litellm_provider": "coralbricks", + "max_input_tokens": 1048576, + "max_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 1.2e-06, + "source": "https://www.coralbricks.ai/pricing", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses", + "/v1/messages" + ], + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "coralbricks/deepseek-v4.1-flash-fast-fp4": { + "cache_creation_input_token_cost": 9e-08, + "cache_read_input_token_cost": 0.0, + "input_cost_per_token": 1e-08, + "litellm_provider": "coralbricks", + "max_input_tokens": 1048576, + "max_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 1.2e-06, + "source": "https://www.coralbricks.ai/pricing", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses", + "/v1/messages" + ], + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "coralbricks/deepseek-v4.1-flash": { + "cache_creation_input_token_cost": 9e-08, + "cache_read_input_token_cost": 0.0, + "input_cost_per_token": 1e-08, + "litellm_provider": "coralbricks", + "max_input_tokens": 1048576, + "max_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 1.2e-06, + "source": "https://www.coralbricks.ai/pricing", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses", + "/v1/messages" + ], + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true + }, "computer-use-preview": { "deprecation_date": "2026-07-23", "input_cost_per_token": 3e-06, @@ -16644,6 +16802,36 @@ "supports_vision": true, "source": "https://platform.openai.com/docs/models/computer-use-preview" }, + "computer-use-preview-2025-03-11": { + "deprecation_date": "2026-07-23", + "input_cost_per_token": 3e-06, + "litellm_provider": "openai", + "max_input_tokens": 8192, + "max_output_tokens": 1024, + "max_tokens": 1024, + "mode": "chat", + "output_cost_per_token": 1.2e-05, + "supported_endpoints": [ + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": false, + "supports_reasoning": true, + "supports_response_schema": false, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true, + "source": "https://developers.openai.com/api/docs/models/computer-use-preview" + }, "deepseek-chat": { "cache_read_input_token_cost": 2.8e-08, "deprecation_date": "2026-07-24", @@ -33449,6 +33637,44 @@ "supports_tool_choice": true, "supports_vision": false }, + "gpt-audio-mini-2025-10-06": { + "deprecation_date": "2026-07-23", + "input_cost_per_audio_token": 1e-05, + "input_cost_per_token": 6e-07, + "litellm_provider": "openai", + "max_input_tokens": 128000, + "max_output_tokens": 16384, + "max_tokens": 16384, + "mode": "chat", + "output_cost_per_audio_token": 2e-05, + "output_cost_per_token": 2.4e-06, + "source": "https://developers.openai.com/api/docs/models/gpt-audio-mini", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses", + "/v1/realtime", + "/v1/batch" + ], + "supported_modalities": [ + "text", + "audio" + ], + "supported_output_modalities": [ + "text", + "audio" + ], + "supports_audio_input": true, + "supports_audio_output": true, + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_parallel_function_calling": true, + "supports_prompt_caching": false, + "supports_reasoning": false, + "supports_response_schema": false, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": false + }, "gpt-audio-mini-2025-12-15": { "input_cost_per_audio_token": 1e-05, "input_cost_per_token": 6e-07, @@ -33609,6 +33835,34 @@ "supports_vision": true, "supports_web_search": true }, + "gpt-4o-mini-search-preview-2025-03-11": { + "deprecation_date": "2026-07-23", + "cache_read_input_token_cost": 7.5e-08, + "input_cost_per_token": 1.5e-07, + "input_cost_per_token_batches": 7.5e-08, + "litellm_provider": "openai", + "max_input_tokens": 128000, + "max_output_tokens": 16384, + "max_tokens": 16384, + "mode": "chat", + "output_cost_per_token": 6e-07, + "output_cost_per_token_batches": 3e-07, + "search_context_cost_per_query": { + "search_context_size_high": 0.025, + "search_context_size_low": 0.025, + "search_context_size_medium": 0.025 + }, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "source": "https://developers.openai.com/api/docs/models/gpt-4o-mini-search-preview" + }, "gpt-4o-mini-transcribe": { "input_cost_per_audio_token": 1.25e-06, "input_cost_per_token": 1.25e-06, @@ -33689,6 +33943,34 @@ "supports_vision": true, "supports_web_search": true }, + "gpt-4o-search-preview-2025-03-11": { + "deprecation_date": "2026-07-23", + "cache_read_input_token_cost": 1.25e-06, + "input_cost_per_token": 2.5e-06, + "input_cost_per_token_batches": 1.25e-06, + "litellm_provider": "openai", + "max_input_tokens": 128000, + "max_output_tokens": 16384, + "max_tokens": 16384, + "mode": "chat", + "output_cost_per_token": 1e-05, + "output_cost_per_token_batches": 5e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.025, + "search_context_size_low": 0.025, + "search_context_size_medium": 0.025 + }, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "source": "https://developers.openai.com/api/docs/models/gpt-4o-search-preview" + }, "gpt-4o-transcribe": { "input_cost_per_audio_token": 2.5e-06, "input_cost_per_token": 2.5e-06, @@ -34942,13 +35224,17 @@ }, "gpt-6-luna": { "cache_creation_input_token_cost": 1.25e-07, + "cache_creation_input_token_cost_decisions": 1e-07, "cache_creation_input_token_cost_above_272k_tokens": 2.5e-07, + "cache_creation_input_token_cost_above_272k_tokens_decisions": 2e-07, "cache_creation_input_token_cost_above_272k_tokens_flex": 1.25e-07, "cache_creation_input_token_cost_above_272k_tokens_priority": 5e-07, "cache_creation_input_token_cost_flex": 6.25e-08, "cache_creation_input_token_cost_priority": 2.5e-07, "cache_read_input_token_cost": 1e-08, + "cache_read_input_token_cost_decisions": 1e-07, "cache_read_input_token_cost_above_272k_tokens": 2e-08, + "cache_read_input_token_cost_above_272k_tokens_decisions": 2e-07, "cache_read_input_token_cost_above_272k_tokens_flex": 1e-08, "cache_read_input_token_cost_above_272k_tokens_priority": 4e-08, "cache_read_input_token_cost_batches": 5e-09, @@ -34958,7 +35244,9 @@ "cache_read_input_token_cost_flex": 5e-09, "cache_read_input_token_cost_priority": 2e-08, "input_cost_per_token": 1e-07, + "input_cost_per_token_decisions": 1e-07, "input_cost_per_token_above_272k_tokens": 2e-07, + "input_cost_per_token_above_272k_tokens_decisions": 2e-07, "input_cost_per_token_above_272k_tokens_flex": 1e-07, "input_cost_per_token_above_272k_tokens_priority": 4e-07, "input_cost_per_token_batches": 5e-08, @@ -34971,7 +35259,9 @@ "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 5e-07, + "output_cost_per_token_decisions": 0.0, "output_cost_per_token_above_272k_tokens": 7.5e-07, + "output_cost_per_token_above_272k_tokens_decisions": 0.0, "output_cost_per_token_above_272k_tokens_flex": 3.75e-07, "output_cost_per_token_above_272k_tokens_priority": 1.5e-06, "output_cost_per_token_batches": 2.5e-07, @@ -38843,6 +39133,49 @@ ], "supports_audio_input": true }, + "meta/muse-wonder-1.0": { + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 4e-06, + "litellm_provider": "meta", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 2e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.0025, + "search_context_size_low": 0.0025, + "search_context_size_medium": 0.0025 + }, + "source": "https://dev.meta.ai/docs/pricing-rate-limits", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses", + "/v1/messages" + ], + "supported_modalities": [ + "text", + "image", + "video" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_minimal_reasoning_effort": true, + "supports_none_reasoning_effort": false, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "supports_xhigh_reasoning_effort": true + }, "meta_llama/Llama-3.3-70B-Instruct": { "litellm_provider": "meta_llama", "max_input_tokens": 128000, @@ -41481,6 +41814,38 @@ "supports_pdf_input": true, "deprecation_date": "2026-07-23" }, + "o3-deep-research-2025-06-26": { + "deprecation_date": "2026-07-23", + "cache_read_input_token_cost": 2.5e-06, + "input_cost_per_token": 1e-05, + "litellm_provider": "openai", + "max_input_tokens": 200000, + "max_output_tokens": 100000, + "max_tokens": 100000, + "mode": "responses", + "output_cost_per_token": 4e-05, + "source": "https://developers.openai.com/api/docs/models/o3-deep-research", + "supported_endpoints": [ + "/v1/batch", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": false, + "supports_native_streaming": true, + "supports_parallel_function_calling": false, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": false, + "supports_tool_choice": false, + "supports_vision": true, + "supports_pdf_input": true + }, "o3-2025-04-16": { "cache_read_input_token_cost": 5e-07, "cache_read_input_token_cost_flex": 2.5e-07, @@ -41711,6 +42076,38 @@ "supports_pdf_input": true, "deprecation_date": "2026-07-23" }, + "o4-mini-deep-research-2025-06-26": { + "deprecation_date": "2026-07-23", + "cache_read_input_token_cost": 5e-07, + "input_cost_per_token": 2e-06, + "litellm_provider": "openai", + "max_input_tokens": 200000, + "max_output_tokens": 100000, + "max_tokens": 100000, + "mode": "responses", + "output_cost_per_token": 8e-06, + "source": "https://developers.openai.com/api/docs/models/o4-mini-deep-research", + "supported_endpoints": [ + "/v1/batch", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": false, + "supports_native_streaming": true, + "supports_parallel_function_calling": false, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": false, + "supports_tool_choice": false, + "supports_vision": true, + "supports_pdf_input": true + }, "o4-mini-2025-04-16": { "cache_read_input_token_cost": 2.75e-07, "cache_read_input_token_cost_flex": 1.38e-07, @@ -43006,13 +43403,21 @@ "output_cost_per_token": 2.8e-07 }, "perplexity/pplx-decider-v1-27b": { - "input_cost_per_token": 4e-08, + "input_cost_per_token": 2e-08, "litellm_provider": "perplexity", "max_input_tokens": 262144, "mode": "evaluation", "output_cost_per_token": 0.0, "source": "https://docs.perplexity.ai/api-reference/decisions-post" }, + "perplexity/pplx-decider-v1.1-27b": { + "input_cost_per_token": 2e-08, + "litellm_provider": "perplexity", + "max_input_tokens": 262144, + "mode": "evaluation", + "output_cost_per_token": 0.0, + "source": "https://docs.perplexity.ai/docs/getting-started/pricing" + }, "perplexity/sonar": { "input_cost_per_token": 1e-06, "litellm_provider": "perplexity", @@ -51404,6 +51809,70 @@ "litellm_provider": "vertex_ai", "mode": "vector_store" }, + "vertex_ai/semantic-ranker-default": { + "input_cost_per_query": 0.001, + "input_cost_per_token": 0.0, + "litellm_provider": "vertex_ai", + "mode": "rerank", + "output_cost_per_token": 0.0, + "source": "https://cloud.google.com/generative-ai-app-builder/pricing" + }, + "vertex_ai/semantic-ranker-default-002": { + "input_cost_per_query": 0.001, + "input_cost_per_token": 0.0, + "litellm_provider": "vertex_ai", + "mode": "rerank", + "output_cost_per_token": 0.0, + "source": "https://cloud.google.com/generative-ai-app-builder/pricing" + }, + "vertex_ai/semantic-ranker-default-003": { + "input_cost_per_query": 0.001, + "input_cost_per_token": 0.0, + "litellm_provider": "vertex_ai", + "mode": "rerank", + "output_cost_per_token": 0.0, + "source": "https://cloud.google.com/generative-ai-app-builder/pricing" + }, + "vertex_ai/semantic-ranker-default-004": { + "input_cost_per_query": 0.001, + "input_cost_per_token": 0.0, + "litellm_provider": "vertex_ai", + "mode": "rerank", + "output_cost_per_token": 0.0, + "source": "https://cloud.google.com/generative-ai-app-builder/pricing" + }, + "vertex_ai/semantic-ranker-default@latest": { + "input_cost_per_query": 0.001, + "input_cost_per_token": 0.0, + "litellm_provider": "vertex_ai", + "mode": "rerank", + "output_cost_per_token": 0.0, + "source": "https://cloud.google.com/generative-ai-app-builder/pricing" + }, + "vertex_ai/semantic-ranker-fast": { + "input_cost_per_query": 0.001, + "input_cost_per_token": 0.0, + "litellm_provider": "vertex_ai", + "mode": "rerank", + "output_cost_per_token": 0.0, + "source": "https://cloud.google.com/generative-ai-app-builder/pricing" + }, + "vertex_ai/semantic-ranker-fast-004": { + "input_cost_per_query": 0.001, + "input_cost_per_token": 0.0, + "litellm_provider": "vertex_ai", + "mode": "rerank", + "output_cost_per_token": 0.0, + "source": "https://cloud.google.com/generative-ai-app-builder/pricing" + }, + "vertex_ai/semantic-ranker-fast@latest": { + "input_cost_per_query": 0.001, + "input_cost_per_token": 0.0, + "litellm_provider": "vertex_ai", + "mode": "rerank", + "output_cost_per_token": 0.0, + "source": "https://cloud.google.com/generative-ai-app-builder/pricing" + }, "openai/container": { "code_interpreter_cost_per_session": 0.03, "litellm_provider": "openai", @@ -65714,7 +66183,6 @@ "source": "https://api.together.xyz/v1/models" }, "together_ai/together/Tev1-4B-experimental": { - "cache_read_input_token_cost": 4.2e-08, "input_cost_per_token": 4.2e-08, "litellm_provider": "together_ai", "max_input_tokens": 32768, @@ -69135,6 +69603,15 @@ "supports_response_schema": true, "supports_tool_choice": true }, + "fireworks_ai/accounts/fireworks/routers/rune": { + "cache_read_input_token_cost": 8.4e-09, + "input_cost_per_token": 4.2e-08, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 262144, + "mode": "evaluation", + "output_cost_per_token": 0.0, + "source": "https://docs.fireworks.ai" + }, "fireworks_ai/accounts/fireworks/models/ember-1": { "cache_read_input_token_cost": 3e-07, "cache_read_input_token_cost_priority": 3.75e-07, @@ -71966,5 +72443,97 @@ "supports_system_messages": false, "supports_vision": false, "supports_function_calling": false + }, + "vertex_ai/gemini-robotics-er-2-preview": { + "cache_read_input_token_cost": 1e-07, + "cache_read_input_token_cost_batches": 5e-08, + "input_cost_per_audio_token": 1e-06, + "input_cost_per_token": 1e-06, + "input_cost_per_token_batches": 5e-07, + "litellm_provider": "vertex_ai", + "max_input_tokens": 131072, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "chat", + "output_cost_per_reasoning_token": 5e-06, + "output_cost_per_token": 5e-06, + "output_cost_per_token_batches": 2.5e-06, + "search_context_cost_per_query": {"search_context_size_low":0.014,"search_context_size_medium":0.014,"search_context_size_high":0.014}, + "source": "https://ai.google.dev/gemini-api/docs/pricing", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions", + "/v1/batch" + ], + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_input": true, + "supports_audio_output": false, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_url_context": true, + "supports_video_input": true, + "supports_vision": true, + "supports_web_search": true, + "web_search_billing_unit": "per_query" + }, + "gemini-robotics-er-2-preview": { + "cache_read_input_token_cost": 1e-07, + "cache_read_input_token_cost_batches": 5e-08, + "input_cost_per_audio_token": 1e-06, + "input_cost_per_token": 1e-06, + "input_cost_per_token_batches": 5e-07, + "litellm_provider": "vertex_ai-language-models", + "max_input_tokens": 131072, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "chat", + "output_cost_per_reasoning_token": 5e-06, + "output_cost_per_token": 5e-06, + "output_cost_per_token_batches": 2.5e-06, + "search_context_cost_per_query": {"search_context_size_low":0.014,"search_context_size_medium":0.014,"search_context_size_high":0.014}, + "source": "https://ai.google.dev/gemini-api/docs/pricing", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions", + "/v1/batch" + ], + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_input": true, + "supports_audio_output": false, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_url_context": true, + "supports_video_input": true, + "supports_vision": true, + "supports_web_search": true, + "web_search_billing_unit": "per_query" } } diff --git a/model_prices_and_context_window.schema.json b/model_prices_and_context_window.schema.json index 6f7435fb026..9f5db5abcda 100644 --- a/model_prices_and_context_window.schema.json +++ b/model_prices_and_context_window.schema.json @@ -123,6 +123,11 @@ "minimum": 0, "description": "Rate applied once the prompt exceeds the token threshold in the field name." }, + "cache_creation_input_token_cost_above_272k_tokens_decisions": { + "type": "number", + "minimum": 0, + "description": "Rate applied once the prompt exceeds the token threshold in the field name." + }, "cache_creation_input_token_cost_above_272k_tokens_flex": { "type": "number", "minimum": 0, @@ -142,6 +147,10 @@ "type": "number", "minimum": 0 }, + "cache_creation_input_token_cost_decisions": { + "type": "number", + "minimum": 0 + }, "cache_creation_input_token_cost_flex": { "type": "number", "minimum": 0, @@ -204,6 +213,11 @@ "minimum": 0, "description": "Rate applied once the prompt exceeds the token threshold in the field name." }, + "cache_read_input_token_cost_above_272k_tokens_decisions": { + "type": "number", + "minimum": 0, + "description": "Rate applied once the prompt exceeds the token threshold in the field name." + }, "cache_read_input_token_cost_above_272k_tokens_flex": { "type": "number", "minimum": 0, @@ -232,6 +246,10 @@ "type": "number", "minimum": 0 }, + "cache_read_input_token_cost_decisions": { + "type": "number", + "minimum": 0 + }, "cache_read_input_token_cost_flex": { "type": "number", "minimum": 0, @@ -407,6 +425,11 @@ "minimum": 0, "description": "Rate applied once the prompt exceeds the token threshold in the field name." }, + "input_cost_per_token_above_272k_tokens_decisions": { + "type": "number", + "minimum": 0, + "description": "Rate applied once the prompt exceeds the token threshold in the field name." + }, "input_cost_per_token_above_272k_tokens_flex": { "type": "number", "minimum": 0, @@ -440,6 +463,10 @@ "type": "number", "minimum": 0 }, + "input_cost_per_token_decisions": { + "type": "number", + "minimum": 0 + }, "input_cost_per_token_flex": { "type": "number", "minimum": 0, @@ -786,6 +813,11 @@ "minimum": 0, "description": "Rate applied once the prompt exceeds the token threshold in the field name." }, + "output_cost_per_token_above_272k_tokens_decisions": { + "type": "number", + "minimum": 0, + "description": "Rate applied once the prompt exceeds the token threshold in the field name." + }, "output_cost_per_token_above_272k_tokens_flex": { "type": "number", "minimum": 0, @@ -815,6 +847,10 @@ "minimum": 0, "description": "USD per generated token via the provider's batch API." }, + "output_cost_per_token_decisions": { + "type": "number", + "minimum": 0 + }, "output_cost_per_token_flex": { "type": "number", "minimum": 0, diff --git a/provider_endpoints_support.json b/provider_endpoints_support.json index ea617cf8462..8263c05a4d9 100644 --- a/provider_endpoints_support.json +++ b/provider_endpoints_support.json @@ -672,6 +672,24 @@ "interactions": true } }, + "coralbricks": { + "display_name": "CoralBricks (`coralbricks`)", + "url": "https://docs.litellm.ai/docs/providers/coralbricks", + "endpoints": { + "chat_completions": true, + "messages": true, + "responses": true, + "embeddings": false, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false, + "a2a": false, + "interactions": false + } + }, "cortecs": { "display_name": "Cortecs (`cortecs`)", "url": "https://docs.litellm.ai/docs/providers/cortecs", @@ -1078,7 +1096,8 @@ "batches": false, "rerank": true, "a2a": true, - "interactions": true + "interactions": true, + "systemone": true } }, "firecrawl": { diff --git a/pyrightconfig.json b/pyrightconfig.json index 0b2d7897c85..2264c29be37 100644 --- a/pyrightconfig.json +++ b/pyrightconfig.json @@ -5,7 +5,7 @@ "executionEnvironments": [ { "root": "tests/e2e_harness", - "extraPaths": ["tests/e2e", "tests/e2e/batches", "tests/e2e/guardrails", "tests/e2e/load", "tests/e2e/logging"] + "extraPaths": ["tests/e2e", "tests/e2e/batches", "tests/e2e/guardrails", "tests/e2e/perf", "tests/e2e/logging"] } ], "pythonVersion": "3.12", diff --git a/scripts/check_test_quality.py b/scripts/check_test_quality.py index 6f40dd6ea84..09b87039048 100644 --- a/scripts/check_test_quality.py +++ b/scripts/check_test_quality.py @@ -63,6 +63,31 @@ TQ009 A child interpreter spawned as `subprocess.run([sys.executable, ...])` w same working-directory hazard seen from the child's side. Use tests.test_litellm_rust.support.child_interpreter.run_child_interpreter, which also asserts the child resolved the same `litellm.__file__` as the parent. +TQ010 An integration test that signals a process itself, outside the harness in + tests/integration/_support: `os.kill`, `os.killpg`, a SIGKILL / SIGSTOP / SIGCONT + constant handed to any call, or a zero-argument `.kill()` / `.terminate()` / + `.suspend()` / `.resume()` / a `.send_signal(...)` method call. A hand-rolled + worker kill returns as soon as a replacement process exists rather than once it + serves, and a hand-rolled suspend outlives the owned proxy's worker healthcheck, + after which uvicorn kills the suspended worker under the test. Use the owned + proxy's `kill_worker` / `replacing_worker` / `suspended` and the process helpers in + tests/integration/_support/process.py. + Only files under tests/integration, outside tests/integration/_support, are + examined. +TQ011 An integration-test assertion that compares an elapsed time against a numeric + literal (`assert elapsed < 2`, `assert time.monotonic() - started <= 5`, + `assert all(duration < 2 for ...)`). A fixed wall-clock bound measures the + machine as much as the code, so it fails whenever sibling workers load the box. + Assert relative to a measured control or to the injected delay instead, or, + when the time is the behaviour under test (a cap, timeout or cooldown the + product enforces), against that configured value plus explicit slack with + `assert_enforced_within` (tests/integration/_support/timing.py). An operand is an elapsed time when it + reads a clock (`time.monotonic()`, `time.perf_counter()`, `time.time()` and + their `_ns` forms), names something called elapsed / duration / latency, or + names a local the enclosing function bound to a clock reading. The other operand + is a literal when it is built only from numbers, or names a module-level number. + Only `assert` statements are examined, so a polling deadline in a `while` loop is + left alone. Same path scope as TQ010. Every rule is suppressible with `# test-quality-ok: ` on the reported line, following the repo's `*-ok: ` convention. A suppression without a @@ -147,9 +172,24 @@ SUBPROCESS_SPAWNS: Final = frozenset(("run", "Popen", "check_output", "check_cal INTERPRETER_ISOLATION_FLAGS: Final = frozenset(("-I", "-P")) RULE_CODES: Final = frozenset(( - "TQ000", "TQ001", "TQ002", "TQ003", "TQ004", "TQ005", "TQ006", "TQ007", "TQ009", + "TQ000", "TQ001", "TQ002", "TQ003", "TQ004", "TQ005", "TQ006", "TQ007", "TQ009", "TQ010", "TQ011", )) +INTEGRATION_ROOT: Final = ("tests", "integration") +INTEGRATION_HARNESS: Final = "_support" +SIGNAL_SENDERS: Final = frozenset(("os.kill", "os.killpg")) +PROCESS_CONTROL_SIGNALS: Final = frozenset(("SIGKILL", "SIGSTOP", "SIGCONT")) +ZERO_ARGUMENT_PROCESS_CONTROLS: Final = frozenset(("kill", "terminate", "suspend", "resume")) +SIGNAL_METHOD: Final = "send_signal" +CLOCK_READERS: Final = frozenset(( + "time.monotonic", "time.perf_counter", "time.time", + "time.monotonic_ns", "time.perf_counter_ns", "time.time_ns", + "monotonic", "perf_counter", "monotonic_ns", "perf_counter_ns", +)) +ELAPSED_NAME_RE: Final = re.compile(r"elapsed|duration|latenc", re.IGNORECASE) +DURATION_READER: Final = "total_seconds" +DURATION_FOLDS: Final = frozenset(("max", "min", "sum", "abs", "round", "float")) + CREDENTIAL_NAME_RE: Final = re.compile( r"(?:API_KEY|_KEY|TOKEN|SECRET|PASSWORD|CREDENTIAL|DATABASE_URL|ACCESS_KEY_ID)$" ) @@ -687,6 +727,156 @@ def iter_child_interpreter_violations(path: Path, tree: ast.Module) -> Iterator[ ) +def is_integration_test_path(path: Path) -> bool: + parts: Final = path.parts + return any( + parts[index:index + 2] == INTEGRATION_ROOT and parts[index + 2] != INTEGRATION_HARNESS + for index in range(len(parts) - 2) + ) + + +def _names_control_signal(node: ast.expr) -> bool: + if isinstance(node, ast.Attribute): + return node.attr in PROCESS_CONTROL_SIGNALS + return isinstance(node, ast.Name) and node.id in PROCESS_CONTROL_SIGNALS + + +def _signals_a_process(call: ast.Call) -> bool: + if _dotted_name(call.func) in SIGNAL_SENDERS: + return True + if any(_names_control_signal(argument) for argument in (*call.args, *(k.value for k in call.keywords))): + return True + func: Final = call.func + if not isinstance(func, ast.Attribute): + return False + if func.attr == SIGNAL_METHOD: + return True + return func.attr in ZERO_ARGUMENT_PROCESS_CONTROLS and not call.args and not call.keywords + + +def iter_process_signal_violations(path: Path, tree: ast.Module) -> Iterator[Violation]: + if not is_integration_test_path(path): + return + for node in ast.walk(tree): + if isinstance(node, ast.Call) and _signals_a_process(node): + yield Violation( + path, + node.lineno, + "TQ010", + "an integration test signals a process itself; a hand-rolled kill returns before the " + "replacement worker serves and a hand-rolled suspend outlives the worker healthcheck, so use " + "the owned proxy's kill_worker / replacing_worker / suspended, or a tests/integration/_support helper " + f"(suppress: `# {SUPPRESSION_TOKEN}: `)", + ) + + +def _reads_clock(node: ast.expr) -> bool: + if isinstance(node, ast.Call): + return _dotted_name(node.func) in CLOCK_READERS + if isinstance(node, ast.BinOp): + return _reads_clock(node.left) or _reads_clock(node.right) + return False + + +def _bound_values(scope: ast.AST) -> Iterator[tuple[ast.expr, ast.expr]]: + for node in ast.walk(scope): + if isinstance(node, ast.Assign): + yield from ((target, node.value) for target in node.targets) + elif isinstance(node, (ast.AnnAssign, ast.NamedExpr)) and node.value is not None: + yield node.target, node.value + + +def _clock_bound_names(scope: ast.AST) -> frozenset[str]: + return frozenset( + target.id for target, value in _bound_values(scope) if isinstance(target, ast.Name) and _reads_clock(value) + ) + + +def _numeric_constant_names(tree: ast.Module) -> frozenset[str]: + return frozenset( + target.id + for statement in tree.body + for target, value in _bound_values(statement) + if isinstance(target, ast.Name) and _is_numeric_literal(value, frozenset()) + ) + + +def _is_numeric_literal(node: ast.expr, constants: frozenset[str]) -> bool: + if isinstance(node, ast.Constant): + return isinstance(node.value, (int, float)) and not isinstance(node.value, bool) + if isinstance(node, ast.UnaryOp): + return _is_numeric_literal(node.operand, constants) + if isinstance(node, ast.BinOp): + return _is_numeric_literal(node.left, constants) and _is_numeric_literal(node.right, constants) + return isinstance(node, ast.Name) and node.id in constants + + +def _is_elapsed(node: ast.expr, clock_names: frozenset[str]) -> bool: + if isinstance(node, ast.Name): + return node.id in clock_names or ELAPSED_NAME_RE.search(node.id) is not None + if isinstance(node, ast.Attribute): + return ELAPSED_NAME_RE.search(node.attr) is not None + if isinstance(node, ast.BinOp): + return _is_elapsed(node.left, clock_names) or _is_elapsed(node.right, clock_names) + if isinstance(node, ast.Subscript): + return _is_elapsed(node.value, clock_names) + if isinstance(node, (ast.GeneratorExp, ast.ListComp)): + return _is_elapsed(node.elt, clock_names) + if not isinstance(node, ast.Call): + return False + if _dotted_name(node.func) in CLOCK_READERS: + return True + if isinstance(node.func, ast.Attribute) and node.func.attr == DURATION_READER: + return _is_elapsed(node.func.value, clock_names) + return _dotted_name(node.func) in DURATION_FOLDS and any(_is_elapsed(arg, clock_names) for arg in node.args) + + +def _upper_bounded_operands(compare: ast.Compare) -> Iterator[tuple[ast.expr, ast.expr]]: + operands: Final = (compare.left, *compare.comparators) + for left, operator, right in zip(operands, compare.ops, operands[1:]): + if isinstance(operator, (ast.Lt, ast.LtE)): + yield left, right + elif isinstance(operator, (ast.Gt, ast.GtE)): + yield right, left + + +def _asserts_in_scope(node: ast.AST, scope: ast.AST) -> Iterator[tuple[ast.Assert, ast.AST]]: + for child in ast.iter_child_nodes(node): + if isinstance(child, ast.Assert): + yield child, scope + inner: Final = child if isinstance(child, (ast.FunctionDef, ast.AsyncFunctionDef)) else scope + yield from _asserts_in_scope(child, inner) + + +def iter_wall_clock_bound_violations(path: Path, tree: ast.Module) -> Iterator[Violation]: + if not is_integration_test_path(path): + return + constants: Final = _numeric_constant_names(tree) + for statement, scope in _asserts_in_scope(tree, tree): + clock_names: Final = _clock_bound_names(scope) + capped: Final = tuple( + compare + for compare in ast.walk(statement.test) + if isinstance(compare, ast.Compare) + and any( + _is_elapsed(bounded, clock_names) and _is_numeric_literal(bound, constants) + for bounded, bound in _upper_bounded_operands(compare) + ) + ) + yield from ( + Violation( + path, + compare.lineno, + "TQ011", + "an elapsed time is capped by a fixed number, which fails whenever sibling workers load the " + "machine; bound it by a measured control or by the injected delay, or a product-enforced limit by " + "its configured value plus slack with assert_enforced_within " + f"(tests/integration/_support/timing.py) (suppress: `# {SUPPRESSION_TOKEN}: `)", + ) + for compare in capped + ) + + def iter_conftest_inventory_violations(path: Path, tree: ast.Module) -> Iterator[Violation]: if path.name != CONFTEST_NAME: return @@ -724,6 +914,8 @@ def check_file(path: Path) -> tuple[Violation, ...]: *iter_credential_skip_violations(path, tree), *iter_conftest_inventory_violations(path, tree), *iter_child_interpreter_violations(path, tree), + *iter_process_signal_violations(path, tree), + *iter_wall_clock_bound_violations(path, tree), ) if violation.line not in skip ) diff --git a/terraform/litellm/gcp/README.md b/terraform/litellm/gcp/README.md index 4b2f576adc1..c90af26ebf2 100644 --- a/terraform/litellm/gcp/README.md +++ b/terraform/litellm/gcp/README.md @@ -1,6 +1,6 @@ # LiteLLM on GCP (Cloud Run) -[![Open in Cloud Shell](https://gstatic.com/cloudssh/images/open-btn.svg)](https://ssh.cloud.google.com/cloudshell/editor?cloudshell_git_repo=https%3A%2F%2Fgithub.com%2FBerriAI%2Flitellm&cloudshell_workspace=terraform%2Flitellm%2Fgcp%2Fexamples%2Fdefault&cloudshell_tutorial=TUTORIAL.md&cloudshell_image=gcr.io/ds-artifacts-cloudshell/deploystack_custom_image&shellonly=true) +[![Open in Cloud Shell](https://gstatic.com/cloudssh/images/open-btn.svg)](https://ssh.cloud.google.com/cloudshell/editor?cloudshell_git_repo=https%3A%2F%2Fgithub.com%2FBerriAI%2Fterraform-google-litellm&cloudshell_git_branch=v1.105.0&cloudshell_workspace=examples%2Fdefault&cloudshell_tutorial=TUTORIAL.md&cloudshell_image=gcr.io/ds-artifacts-cloudshell/deploystack_custom_image&shellonly=true) The button above opens the [DeployStack](https://github.com/GoogleCloudPlatform/deploystack) installer in Cloud Shell, walks you through `TUTORIAL.md`, and runs `terraform apply` once you've answered the prompts. The rest of this README is the manual / advanced path. diff --git a/terraform/litellm/gcp/examples/default/TUTORIAL.md b/terraform/litellm/gcp/examples/default/TUTORIAL.md index 5c7144619d6..ec82feb739a 100644 --- a/terraform/litellm/gcp/examples/default/TUTORIAL.md +++ b/terraform/litellm/gcp/examples/default/TUTORIAL.md @@ -58,7 +58,7 @@ Skip this step entirely for a trial deploy. ## Run the installer -DeployStack will prompt for project, region, tenant, env, image tag, `image_registry`, and TLS posture, then run `terraform apply`. Open `deploystack.json` if you want to see the prompt definitions first. +DeployStack will prompt for project, region, tenant, env, image tag, `image_registry`, and TLS posture, then run `terraform apply`. Open `deploystack.json` if you want to see the prompt definitions first. ```bash deploystack install @@ -131,4 +131,4 @@ deploystack uninstall -Full configuration reference is in `README.md`, and every input variable on the underlying module lives in `variables.tf`. +Full configuration reference is in `README.md`, and every input variable on the underlying module lives in `variables.tf`. diff --git a/tests/code_coverage_tests/check_greptile_files_json.py b/tests/code_coverage_tests/check_greptile_files_json.py deleted file mode 100644 index a3e9658fea0..00000000000 --- a/tests/code_coverage_tests/check_greptile_files_json.py +++ /dev/null @@ -1,71 +0,0 @@ -from __future__ import annotations - -import dataclasses -import difflib -import json -import os -import subprocess -import sys -from pathlib import Path, PurePosixPath -from typing import Final - -REPO_ROOT: Final = Path(__file__).resolve().parents[2] -FILES_JSON: Final = PurePosixPath(".greptile/files.json") -WRITE_FLAG: Final = "--write" - - -@dataclasses.dataclass(frozen=True, slots=True) -class _ContextFile: - path: str - description: str - scope: tuple[str, ...] - - -def _nested_agents_md_paths(repo_root: Path) -> tuple[PurePosixPath, ...]: - result: Final = subprocess.run( - ["git", "-C", os.fspath(repo_root), "ls-files", "-z", "--", "*/AGENTS.md"], - check=True, - stdout=subprocess.PIPE, - ) - return tuple(sorted(PurePosixPath(os.fsdecode(path)) for path in result.stdout.split(b"\0") if path)) - - -def _context_file(agents_md: PurePosixPath) -> _ContextFile: - return _ContextFile( - path=str(agents_md), - description=f"Conventions for code under {agents_md.parent}/", - scope=(f"{agents_md.parent}/**",), - ) - - -def render(repo_root: Path) -> str: - files: Final = tuple(dataclasses.asdict(_context_file(path)) for path in _nested_agents_md_paths(repo_root)) - return json.dumps({"files": files}, indent=2) + "\n" - - -def main(argv: tuple[str, ...]) -> int: - expected: Final = render(REPO_ROOT) - target: Final = REPO_ROOT / FILES_JSON - if WRITE_FLAG in argv: - target.parent.mkdir(exist_ok=True) - target.write_text(expected) - print(f"Wrote {FILES_JSON}") - return 0 - actual: Final = target.read_text() if target.is_file() else "" - if actual == expected: - print(f"{FILES_JSON} lists every nested AGENTS.md") - return 0 - sys.stdout.writelines( - difflib.unified_diff( - actual.splitlines(keepends=True), expected.splitlines(keepends=True), str(FILES_JSON), "expected" - ) - ) - print( - f"\n{FILES_JSON} is out of sync with the nested AGENTS.md files. Regenerate it with:\n" - f" python {Path(__file__).resolve().relative_to(REPO_ROOT)} {WRITE_FLAG}" - ) - return 1 - - -if __name__ == "__main__": - sys.exit(main(tuple(sys.argv[1:]))) diff --git a/tests/code_coverage_tests/test_check_greptile_files_json.py b/tests/code_coverage_tests/test_check_greptile_files_json.py deleted file mode 100644 index 8fdc070d7ff..00000000000 --- a/tests/code_coverage_tests/test_check_greptile_files_json.py +++ /dev/null @@ -1,75 +0,0 @@ -import json -import shutil -import subprocess -import sys -from pathlib import Path -from typing import Final - -GUARD: Final = Path(__file__).resolve().parent / "check_greptile_files_json.py" -FILES_JSON: Final = Path(".greptile/files.json") - - -def _git(repo: Path, *args: str) -> None: - subprocess.run(["git", "-C", str(repo), *args], check=True, capture_output=True) - - -def _write_agents_md(path: Path) -> None: - path.parent.mkdir(parents=True, exist_ok=True) - path.write_text("# conventions\n") - - -def _repo_with(tmp_path: Path, files: tuple[str, ...]) -> Path: - guard_copy: Final = tmp_path / "tests" / "code_coverage_tests" / GUARD.name - guard_copy.parent.mkdir(parents=True) - shutil.copy(GUARD, guard_copy) - for relative_path in files: - _write_agents_md(tmp_path / relative_path) - _git(tmp_path, "init", "-q") - _git(tmp_path, "add", "-A") - return guard_copy - - -def _run(guard: Path, *args: str) -> subprocess.CompletedProcess[str]: - return subprocess.run([sys.executable, "-I", str(guard), *args], capture_output=True, text=True) - - -def test_write_scopes_each_nested_agents_md_to_its_own_directory(tmp_path: Path) -> None: - guard: Final = _repo_with(tmp_path, ("AGENTS.md", "tests/AGENTS.md", "litellm/_v2/cache/AGENTS.md")) - - assert _run(guard, "--write").returncode == 0 - - written: Final = json.loads((tmp_path / FILES_JSON).read_text()) - assert [(entry["path"], entry["scope"]) for entry in written["files"]] == [ - ("litellm/_v2/cache/AGENTS.md", ["litellm/_v2/cache/**"]), - ("tests/AGENTS.md", ["tests/**"]), - ] - assert _run(guard).returncode == 0 - - -def test_new_agents_md_missing_from_files_json_fails(tmp_path: Path) -> None: - guard: Final = _repo_with(tmp_path, ("tests/AGENTS.md",)) - assert _run(guard, "--write").returncode == 0 - _write_agents_md(tmp_path / "ui" / "AGENTS.md") - _git(tmp_path, "add", "-A") - - result: Final = _run(guard) - - assert result.returncode == 1 - assert '+ "path": "ui/AGENTS.md",' in result.stdout - - -def test_deleted_agents_md_left_in_files_json_fails(tmp_path: Path) -> None: - guard: Final = _repo_with(tmp_path, ("tests/AGENTS.md", "ui/AGENTS.md")) - assert _run(guard, "--write").returncode == 0 - _git(tmp_path, "rm", "-q", "-f", "ui/AGENTS.md") - - result: Final = _run(guard) - - assert result.returncode == 1 - assert '- "path": "ui/AGENTS.md",' in result.stdout - - -def test_missing_files_json_fails(tmp_path: Path) -> None: - guard: Final = _repo_with(tmp_path, ("tests/AGENTS.md",)) - - assert _run(guard).returncode == 1 diff --git a/tests/e2e/AGENTS.md b/tests/e2e/AGENTS.md index 3fd2fb9b9a6..8b9132cf05a 100644 --- a/tests/e2e/AGENTS.md +++ b/tests/e2e/AGENTS.md @@ -44,7 +44,7 @@ Each subdirectory under `tests/e2e/` is one suite, scoped to an endpoint family - `logging/` - logging-integration delivery (datadog and friends) - `security/` - secret handling and log-leak protection - `router/` - routing and reliability behavior (fallbacks, cooldowns) plus the memory tests (`test_reliability_memory_e2e.py`: every worker's RSS as read at collection time, before any test traffic, must sit under a fixed idle budget, the release-gate check for a DB-backed boot that idles near the pod limit the way v1.100.x did, and a few hundred failing requests with retries and fallbacks must not grow proxy RSS past a fixed budget nor store a request snapshot past a fixed size, the release-gate check for the v1.100.0 retry-breadcrumb leak) -- `load/` - performance-category tests, kept OUT of the main suite: throughput/load SLO tests are a different testing category from functional e2e (variance-driven, historically flaky) and live outside this suite until re-implemented as their own pipeline (LIT-5163). Do not add a live load test that runs in the default collection. What lives here: 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, marked `weekly` and deselected unless `E2E_WEEKLY_ANOMALY` is set, driven by `.github/workflows/weekly_load_anomaly.yml`), the Redis chaos test (`test_redis_chaos_e2e.py`, locust load against mock deployments split round robin over `/chat/completions` and `/v1/messages`, one endpoint per simulated user, with `CLIENT PAUSE ALL` on the proxy's Redis mid-run to simulate it being down outright, asserting zero failed requests on every endpoint, budgeting RSS and CPU-per-request as ratios against the same run's healthy phase, and holding p50/p90/p99 latency and log-bytes-per-request to flat ceilings (a ratio cannot bound those two: an open breaker skips Redis instead of waiting on it, so the chaos phase can measure cheaper than baseline while still being far slower than a user should see), needs a proxy booted from `gateway/redis_chaos_ci_config.yml` on the same host with `E2E_PROXY_PID` and `E2E_PROXY_LOG` set, marked `redis_chaos`, deselected unless `E2E_REDIS_CHAOS` is set and excluded from the per-PR selector like the rest of `load/`, driven by `.github/workflows/test-e2e-redis-chaos.yml` and by the Buildkite `e2e-redis-chaos` step in project-releaser, which runs the proxy, Postgres and Valkey co-located with pytest in one pod and sets the opt-in). Its aggregation logic (locust, process usage, session anomaly) is covered by `tests/e2e_harness/load/` +- `perf/` - performance-category tests, kept OUT of the main suite: throughput/load SLO tests are a different testing category from functional e2e (variance-driven, historically flaky) and live outside this suite until re-implemented as their own pipeline (LIT-5163). Do not add a live load test that runs in the default collection. What lives here: 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, marked `weekly` and deselected unless `E2E_WEEKLY_ANOMALY` is set, driven by `.github/workflows/weekly_load_anomaly.yml`), the Redis chaos test (`test_redis_chaos_e2e.py`, locust load against mock deployments split round robin over `/chat/completions` and `/v1/messages`, one endpoint per simulated user, with `CLIENT PAUSE ALL` on the proxy's Redis mid-run to simulate it being down outright, asserting zero failed requests on every endpoint, budgeting RSS and CPU-per-request as ratios against the same run's healthy phase, and holding p50/p90/p99 latency and log-bytes-per-request to flat ceilings (a ratio cannot bound those two: an open breaker skips Redis instead of waiting on it, so the chaos phase can measure cheaper than baseline while still being far slower than a user should see), needs a proxy booted from `gateway/redis_chaos_ci_config.yml` on the same host with `E2E_PROXY_PID` and `E2E_PROXY_LOG` set, marked `redis_chaos`, deselected unless `E2E_PERF_TESTS` is set and excluded from the per-PR selector like the rest of `perf/`, driven by `.github/workflows/test-e2e-redis-chaos.yml` and by the Buildkite `e2e-redis-chaos` step in project-releaser, which runs the proxy, Postgres and Valkey co-located with pytest in one pod and sets the opt-in). New perf tests go here too, marked `perf`, behind the same `E2E_PERF_TESTS` opt-in - `other/` - the holding-pen suite for the `other.*` registry cluster with no home of its own yet: the master-key auth gate, JWT auth (access tokens issued by a real Keycloak realm, `idp.py` plus `idp_realm.json`, whose JWKS the proxy's `JWT_PUBLIC_KEY_URL` points at, see CONTRIBUTING.md for the start command and config block), 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 - `secret_manager/` - the gateway's `key_management_system` against a real secret manager: deployment keys resolved from it (`os.environ/` where the name exists only in the manager) and virtual keys written to and deleted from it. The tests are backend-agnostic and each backend is its own lane, because the setting is global to the proxy: `E2E_SECRET_MANAGER=` opts in and picks the backend from `secret_backends.BACKENDS`, the proxy is booted from `gateway/secret_manager__ci_config.yml` against the live manager, and the tests reach that manager through the backend's `SecretStore` (`secret_store_.py`). A test needing something not every backend does carries `requires_capability(...)` and is deselected on lanes that lack it. `secret_manager/backend.sh up ` runs a backend in Docker and writes the proxy's and the tests' env. Marked `secret_manager`, deselected unless `E2E_SECRET_MANAGER` is set, and kept out of the per-PR selector. Backends today: `hashicorp_vault` and `cyberark` (CyberArk Conjur, which cannot delete, so the delete test is Vault-only) - `gateway/` - proxy configuration only (`litellm-config.yml`). No tests @@ -150,7 +150,7 @@ def test_bare_key_blocks_over_its_own_budget(...) -> None: ... `route` is the endpoint the test is checking: `TEAM_MANAGEMENT` for a `/team/update` test, `SPEND_REPORTING` for a `/spend/logs` test, `MESSAGES` for a test of spend on `/v1/messages`. A test whose chat call only triggers the behavior under test, like the budget block above, leaves it unset, since its steps already name the call -Every pytest test under `tests/e2e/` declares a `Subject` with at least its `domain`, except the `load/` suite, which is kept out of the default collection. The harness's own tests in `tests/e2e_harness/` carry none, since they drive nothing. The fields themselves stay optional, since a test that makes no LLM call has no provider, model or mode to name. Every field is a closed enum, so a typo is a basedpyright error at the call site rather than a property that silently never appears. `providers`, `models` and `capabilities` are tuples even with one member, because one test node routinely drives several: the claude_code matrix runs haiku, sonnet and opus in a single body, and a spend test calls two providers on one key. Declare every provider and every model the test drives, fallbacks included. The three are independent sets with no positional pairing between them (one provider x three models is the common case), and each is deduped and sorted at declaration so the committed run artifacts diff cleanly. `models=("gpt-5.5")` is a str and not a tuple, so anything but a tuple raises a `TypeError` where the decorator runs and shows up as a collection error naming the file. `Subject` is serialized with `dataclasses.asdict`, so a new scalar field needs no serializer edit. Empty fields emit no `` at all. A declared model names the constant the test drives (`CHEAP_ANTHROPIC_MODEL`, the file's own `BACKEND`), never a copy of its value, so the property cannot claim one model while an env override runs another. `e2e_metadata` and its call sites never import litellm, only the stdlib, pytest and pydantic: `Provider` mirrors litellm's `LlmProviders` values instead of importing them, because tests/e2e is shipped to the runner image on its own and a `from litellm...` at module scope would make the litellm package a hard dependency of COLLECTING the suite. `TestProviderMirrorsLitellm` in `tests/code_coverage_tests/test_e2e_metadata.py` fails on drift wherever litellm is importable and skips where it is not, so adding a provider is one line in `e2e_metadata` +Every pytest test under `tests/e2e/` declares a `Subject` with at least its `domain`, except the `perf/` suite, which is kept out of the default collection. The harness's own tests in `tests/e2e_harness/` carry none, since they drive nothing. The fields themselves stay optional, since a test that makes no LLM call has no provider, model or mode to name. Every field is a closed enum, so a typo is a basedpyright error at the call site rather than a property that silently never appears. `providers`, `models` and `capabilities` are tuples even with one member, because one test node routinely drives several: the claude_code matrix runs haiku, sonnet and opus in a single body, and a spend test calls two providers on one key. Declare every provider and every model the test drives, fallbacks included. The three are independent sets with no positional pairing between them (one provider x three models is the common case), and each is deduped and sorted at declaration so the committed run artifacts diff cleanly. `models=("gpt-5.5")` is a str and not a tuple, so anything but a tuple raises a `TypeError` where the decorator runs and shows up as a collection error naming the file. `Subject` is serialized with `dataclasses.asdict`, so a new scalar field needs no serializer edit. Empty fields emit no `` at all. A declared model names the constant the test drives (`CHEAP_ANTHROPIC_MODEL`, the file's own `BACKEND`), never a copy of its value, so the property cannot claim one model while an env override runs another. `e2e_metadata` and its call sites never import litellm, only the stdlib, pytest and pydantic: `Provider` mirrors litellm's `LlmProviders` values instead of importing them, because tests/e2e is shipped to the runner image on its own and a `from litellm...` at module scope would make the litellm package a hard dependency of COLLECTING the suite. `TestProviderMirrorsLitellm` in `tests/code_coverage_tests/test_e2e_metadata.py` fails on drift wherever litellm is importable and skips where it is not, so adding a provider is one line in `e2e_metadata` Declared fields ride out as JUnit `` entries behind the fixed prefix, the same way steps do: each scalar under its field name, and each plural value as a repeated property under its SINGULAR name (`provider`, `model`, `capability`). The results JSON downstream regroups them under the plural key, so `providers`, `models` and `capabilities` are arrays there, `[]` when empty diff --git a/tests/e2e/conftest.py b/tests/e2e/conftest.py index 4477b9584aa..3bb4ab0d3be 100644 --- a/tests/e2e/conftest.py +++ b/tests/e2e/conftest.py @@ -33,10 +33,10 @@ from e2e_config import ( OTEL_TLS_OPT_IN_ENV, OTEL_V2_OPT_IN_ENV, OWNED_GATEWAY_OPT_IN_ENV, + PERF_TESTS_OPT_IN_ENV, PROMPT_CACHING_OPT_IN_ENV, PROVIDER_EDGE_HOST_OPT_IN_ENV, PROXY_BASE_URL, - REDIS_CHAOS_OPT_IN_ENV, SAGEMAKER_NOVA2_LITE_OPT_IN_ENV, SAGEMAKER_NOVA_OPT_IN_ENV, SECRET_MANAGER_OPT_IN_ENV, @@ -69,7 +69,8 @@ OPT_IN_MARKERS: Final = MappingProxyType( "weekly": WEEKLY_ANOMALY_OPT_IN_ENV, "managed_files": MANAGED_FILES_OPT_IN_ENV, "prompt_caching_stack": PROMPT_CACHING_OPT_IN_ENV, - "redis_chaos": REDIS_CHAOS_OPT_IN_ENV, + "redis_chaos": PERF_TESTS_OPT_IN_ENV, + "perf": PERF_TESTS_OPT_IN_ENV, "cli_determinism": CLI_DETERMINISM_OPT_IN_ENV, "mcp_oauth_live": MCP_OAUTH_LIVE_OPT_IN_ENV, "provider_edge_host": PROVIDER_EDGE_HOST_OPT_IN_ENV, @@ -139,7 +140,8 @@ def pytest_configure(config: pytest.Config) -> None: ) config.addinivalue_line( "markers", - "load: heavy throughput/load test; collected last so it never perturbs latency-sensitive suites", + "perf: performance test (throughput, latency, memory) against mocked deployments; collected last so it " + "never perturbs latency-sensitive suites, and deselected unless E2E_PERF_TESTS is set", ) config.addinivalue_line( "markers", @@ -161,7 +163,7 @@ def pytest_configure(config: pytest.Config) -> None: config.addinivalue_line( "markers", "redis_chaos: load test that pauses the proxy's Redis outright mid-run; needs a proxy booted from " - "gateway/redis_chaos_ci_config.yml on the same host, and is deselected unless E2E_REDIS_CHAOS is set", + "gateway/redis_chaos_ci_config.yml on the same host, and is deselected unless E2E_PERF_TESTS is set", ) config.addinivalue_line( "markers", @@ -251,7 +253,7 @@ def pytest_collection_modifyitems(config: pytest.Config, items: list[pytest.Item standard report and these properties for package rollups and coverage drill-down. See junit_properties.py. - Also sort `load`-marked items last so a whole-tree run drives heavy throughput + Also sort `perf`-marked items last so a whole-tree run drives heavy throughput traffic only after the latency-sensitive suites have finished.""" deselected = [item for item in items if _needs_unset_opt_in(item)] if deselected: @@ -264,7 +266,7 @@ def pytest_collection_modifyitems(config: pytest.Config, items: list[pytest.Item items[:] = [item for item in items if item.get_closest_marker("migration_startup") is None] if deselected: deselected[0].config.hook.pytest_deselected(items=deselected) - items.sort(key=lambda item: item.get_closest_marker("load") is not None) + items.sort(key=lambda item: item.get_closest_marker("perf") is not None) @pytest.hookimpl(tryfirst=True) diff --git a/tests/e2e/coverage_registry/other.yaml b/tests/e2e/coverage_registry/other.yaml index 2c8bf4309ae..25fd76c9786 100644 --- a/tests/e2e/coverage_registry/other.yaml +++ b/tests/e2e/coverage_registry/other.yaml @@ -18,6 +18,7 @@ - {id: other.auth.jwt.team_header_alias_binds_team, module: other, tier: P0, area: auth, assertions: [team_header_alias_binds_team], source: "handle_jwt.py JWTAuthManager.resolve_team_from_header / LIT-7181", fail_before_fix: proven, rationale: "x-litellm-team-id carrying the team alias binds and attributes the same team as the team id, so a managed client can pin a stable alias instead of a uuid"} - {id: other.auth.jwt.team_header_non_member_alias_denied, module: other, tier: P0, area: auth, assertions: [team_header_non_member_alias_denied], source: "handle_jwt.py JWTAuthManager.resolve_team_from_header / LIT-7181", rationale: "x-litellm-team-id naming the alias of a team the JWT does not grant is denied 403 with the same body as an unknown value, so the response does not reveal whether that team exists"} - {id: other.auth.jwt.team_model_alias_listed_and_routes, module: other, tier: P1, area: auth, assertions: [team_model_alias_listed_and_routes], source: "proxy_server.py model_list / common_utils/model_listing_utils.py alias_listing_entries / LIT-8515", fail_before_fix: proven, rationale: "A team model_aliases name the JWT caller can complete on is also listed by GET /v1/models for that caller, in the OpenAI and the Anthropic (Claude Code) shapes, next to its target, so a managed client can discover the alias it is meant to send"} +- {id: other.auth.jwt.proxy_admin_personal_budget_enforced, module: other, tier: P0, area: auth, assertions: [proxy_admin_personal_budget_enforced], source: "user_api_key_auth.py _run_centralized_common_checks / auth_checks.py _user_max_budget_check / GH#41226 / LIT-8590", fail_before_fix: proven, rationale: "A JWT whose sub resolves to a proxy_admin user row with a spent max_budget is refused on /chat/completions, the same as a non-admin user; forcing the admin role must not drop the row's budget fields"} - {id: other.auth.model_access_group.wildcard_bare_name_allowed, module: other, tier: P0, area: auth, assertions: [wildcard_bare_name_allowed], source: "auth_checks.py:3232 / LIT-5813", fail_before_fix: proven, rationale: "A grant of a group holding a wildcard deployment covers the bare model names callers actually send, not only the provider-prefixed spelling"} - {id: other.auth.model_access_group.member_allowed, module: other, tier: P0, area: auth, assertions: [member_allowed], source: "auth_checks.py:3232", rationale: "A key whose allow-list is a model access group can call the deployments in that group"} - {id: other.auth.model_access_group.non_member_denied, module: other, tier: P0, area: auth, assertions: [non_member_denied], source: "auth_checks.py:3232", rationale: "That same grant reaches nothing outside the group, including provider models the group's wildcard does not cover"} diff --git a/tests/e2e/e2e_config.py b/tests/e2e/e2e_config.py index fd5c0bae27b..8fc7e50d95c 100644 --- a/tests/e2e/e2e_config.py +++ b/tests/e2e/e2e_config.py @@ -205,7 +205,7 @@ LOAD_MIN_CONCURRENCY_EFFICIENCY = float(os.environ.get("E2E_LOAD_MIN_CONCURRENCY WEEKLY_ANOMALY_OPT_IN_ENV = "E2E_WEEKLY_ANOMALY" MANAGED_FILES_OPT_IN_ENV = "E2E_MANAGED_FILES_STACK" PROMPT_CACHING_OPT_IN_ENV = "E2E_PROMPT_CACHING_STACK" -REDIS_CHAOS_OPT_IN_ENV = "E2E_REDIS_CHAOS" +PERF_TESTS_OPT_IN_ENV: Final = "E2E_PERF_TESTS" CLI_DETERMINISM_OPT_IN_ENV = "E2E_CLI_DETERMINISM" MCP_OAUTH_LIVE_OPT_IN_ENV: Final = "E2E_MCP_OAUTH_LIVE" PROVIDER_EDGE_HOST_OPT_IN_ENV: Final = "E2E_PROVIDER_EDGE_HOST_REACHABLE" diff --git a/tests/e2e/models.py b/tests/e2e/models.py index 42a59af45d2..5b1c5778c10 100644 --- a/tests/e2e/models.py +++ b/tests/e2e/models.py @@ -1665,6 +1665,7 @@ class UserNewBody(BaseModel): user_role: UserRole user_id: str | None = None auto_create_key: bool | None = None + max_budget: float | None = None class UserNewResponse(BaseModel): diff --git a/tests/e2e/other/test_jwt_auth_e2e.py b/tests/e2e/other/test_jwt_auth_e2e.py index 6da3924c284..1fb2e2b1e9f 100644 --- a/tests/e2e/other/test_jwt_auth_e2e.py +++ b/tests/e2e/other/test_jwt_auth_e2e.py @@ -12,7 +12,7 @@ from e2e_config import CHEAP_OPENAI_MODEL, unique_marker from e2e_http import UnauthorizedError, UnknownApiError, unwrap from idp import SHORT_LIVED_CLIENT_ID, WRONG_AUDIENCE_CLIENT_ID, Identity from lifecycle import ResourceManager -from models import ChatBody, ChatMessage, TeamNewBody +from models import ChatBody, ChatMessage, TeamNewBody, UserNewBody from other_client import OtherClient from pydantic import BaseModel from e2e_metadata import Domain, Mode, Provider, Subject, meta @@ -150,6 +150,38 @@ class TestJwtAuth: f"spend row must carry the user from the JWT sub claim {identity.user_id!r}, got {row.user!r}" ) + @pytest.mark.covers("other.auth.jwt.proxy_admin_personal_budget_enforced") + @meta( + Subject( + domain=Domain.PROXY_AUTH, + models=(CHEAP_OPENAI_MODEL,), + mode=Mode.NONSTREAM, + ) + ) + def test_proxy_admin_token_is_blocked_once_its_personal_budget_is_spent( + self, client: OtherClient, resources: ResourceManager + ) -> None: + identity: Final = client.idp.provision_user( + marker=unique_marker(), groups=(), group_ids=(), defer=resources.defer + ) + resources.defer(lambda: client.proxy.delete_user(identity.user_id)) + unwrap( + client.user_new( + UserNewBody( + user_id=identity.user_id, + user_email=f"{identity.username}@example.com", + user_role="proxy_admin", + auto_create_key=False, + max_budget=0.0, + ) + ) + ) + + blocked: Final = client.proxy.chat(client.idp.access_token(identity), _ping()) + assert isinstance(blocked, UnknownApiError), f"proxy_admin JWT must be stopped by its max_budget, got {blocked}" + assert (blocked.status_code, "ExceededBudget" in blocked.body) == (422, True), blocked + assert f"User={identity.user_id}" in blocked.body, f"budget denial must name the JWT's user: {blocked.body}" + @pytest.mark.covers("other.auth.jwt.invalid_signature_denied") @meta( Subject( diff --git a/tests/e2e/load/conftest.py b/tests/e2e/perf/conftest.py similarity index 100% rename from tests/e2e/load/conftest.py rename to tests/e2e/perf/conftest.py diff --git a/tests/e2e/load/load_client.py b/tests/e2e/perf/load_client.py similarity index 100% rename from tests/e2e/load/load_client.py rename to tests/e2e/perf/load_client.py diff --git a/tests/e2e/load/locust_load.py b/tests/e2e/perf/locust_load.py similarity index 100% rename from tests/e2e/load/locust_load.py rename to tests/e2e/perf/locust_load.py diff --git a/tests/e2e/load/locustfile.py b/tests/e2e/perf/locustfile.py similarity index 100% rename from tests/e2e/load/locustfile.py rename to tests/e2e/perf/locustfile.py diff --git a/tests/e2e/load/phase_budget.py b/tests/e2e/perf/phase_budget.py similarity index 100% rename from tests/e2e/load/phase_budget.py rename to tests/e2e/perf/phase_budget.py diff --git a/tests/e2e/load/proxy_usage.py b/tests/e2e/perf/proxy_usage.py similarity index 100% rename from tests/e2e/load/proxy_usage.py rename to tests/e2e/perf/proxy_usage.py diff --git a/tests/e2e/load/session_anomaly.py b/tests/e2e/perf/session_anomaly.py similarity index 100% rename from tests/e2e/load/session_anomaly.py rename to tests/e2e/perf/session_anomaly.py diff --git a/tests/e2e/load/test_redis_chaos_e2e.py b/tests/e2e/perf/test_redis_chaos_e2e.py similarity index 99% rename from tests/e2e/load/test_redis_chaos_e2e.py rename to tests/e2e/perf/test_redis_chaos_e2e.py index 9a9e8082e5b..feb87f97e0a 100644 --- a/tests/e2e/load/test_redis_chaos_e2e.py +++ b/tests/e2e/perf/test_redis_chaos_e2e.py @@ -28,7 +28,7 @@ Needs the proxy on the same host, since RSS and CPU come from psutil on its proc a multi-worker proxy serves /metrics from the prometheus multiprocess collector, which drops the process collector's memory and CPU series. Log bytes are read from the file the proxy's stdout/stderr was redirected to, so the same host requirement covers that too. Deselected -unless E2E_REDIS_CHAOS is set. +unless E2E_PERF_TESTS is set. """ from __future__ import annotations diff --git a/tests/e2e/load/test_weekly_session_anomaly_e2e.py b/tests/e2e/perf/test_weekly_session_anomaly_e2e.py similarity index 96% rename from tests/e2e/load/test_weekly_session_anomaly_e2e.py rename to tests/e2e/perf/test_weekly_session_anomaly_e2e.py index d4ef883702e..64a000810fd 100644 --- a/tests/e2e/load/test_weekly_session_anomaly_e2e.py +++ b/tests/e2e/perf/test_weekly_session_anomaly_e2e.py @@ -3,7 +3,6 @@ from __future__ import annotations from dataclasses import dataclass import pytest - from e2e_config import ( ANOMALY_MAX_ERROR_RATIO, ANOMALY_MAX_KEY_SPEND_USD, @@ -21,7 +20,7 @@ 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] +pytestmark = [pytest.mark.e2e, pytest.mark.weekly] @dataclass(frozen=True, slots=True) @@ -67,9 +66,7 @@ class TestWeeklySessionAnomaly: 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) - ) + 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( diff --git a/tests/e2e/load/weekly_anomaly_config.yml b/tests/e2e/perf/weekly_anomaly_config.yml similarity index 100% rename from tests/e2e/load/weekly_anomaly_config.yml rename to tests/e2e/perf/weekly_anomaly_config.yml diff --git a/tests/e2e/pytest.ini b/tests/e2e/pytest.ini index ff81a593fe8..da3dd72d795 100644 --- a/tests/e2e/pytest.ini +++ b/tests/e2e/pytest.ini @@ -7,12 +7,12 @@ markers = e2e: live test that requires a running proxy and real provider keys meta: typed e2e_metadata.Subject describing what this test drives (domain/route/providers/models/capabilities/mode); attach it with @meta(Subject(...)), never as a bare pytest.mark replayable: edge-wired test whose provider traffic replays from a fixture bundle, so it makes zero provider calls in replay mode; the record/replay CI lane selects it with -m replayable - load: heavy throughput/load test; collected last so it never perturbs latency-sensitive suites + perf: performance test (throughput, latency, memory); collected last so it never perturbs latency-sensitive suites, and deselected unless E2E_PERF_TESTS is set weekly: real-provider anomaly load test that spends real money; deselected unless E2E_WEEKLY_ANOMALY is set managed_files: needs a proxy running with require_managed_files enabled; deselected unless E2E_MANAGED_FILES_STACK is set prompt_caching_stack: needs a proxy running with router_settings.optional_pre_call_checks including prompt_caching; deselected unless E2E_PROMPT_CACHING_STACK is set cli_determinism: drives the real claude CLI for several seconds; deselected unless E2E_CLI_DETERMINISM is set - redis_chaos: load test that pauses the proxy's Redis outright mid-run; needs a proxy booted from gateway/redis_chaos_ci_config.yml on the same host, and is deselected unless E2E_REDIS_CHAOS is set + redis_chaos: load test that pauses the proxy's Redis outright mid-run; needs a proxy booted from gateway/redis_chaos_ci_config.yml on the same host, and is deselected unless E2E_PERF_TESTS is set quiet_stack: measures the proxy itself, so it runs while no other test on this host is hitting the stack; every other test waits for it to finish mcp_oauth_live: real Linear OAuth consent via a captured browser session; deselected unless E2E_MCP_OAUTH_LIVE is set provider_edge_host: routes provider traffic through the pytest host's edge in every fixture mode, so the gateway must reach the pytest host; deselected unless E2E_PROVIDER_EDGE_HOST_REACHABLE is set diff --git a/tests/e2e/ui/tests/integrationCritical/addModelCoralBricks.spec.ts b/tests/e2e/ui/tests/integrationCritical/addModelCoralBricks.spec.ts new file mode 100644 index 00000000000..1fea7eacbe3 --- /dev/null +++ b/tests/e2e/ui/tests/integrationCritical/addModelCoralBricks.spec.ts @@ -0,0 +1,172 @@ +import { test, expect, APIRequestContext, Page as PlaywrightPage } from "@playwright/test"; +import { randomUUID } from "node:crypto"; +import { Page } from "../../fixtures/pages"; +import { navigateToPage } from "../../helpers/navigation"; +import { captureRequestBody } from "../../helpers/roundTrip"; + +const master = process.env.LITELLM_MASTER_KEY ?? "sk-integration-master"; +const headers = { Authorization: `Bearer ${master}` }; +const upstream = (process.env.INTEGRATION_UPSTREAM_URL ?? "http://127.0.0.1:8190").replace(/\/+$/, ""); +const CORALBRICKS_MODEL = "deepseek-v4.1-flash-fast"; +const UI_KEY = "cb_ui_integration_key"; +const UPSTREAM_PROMPT_TOKENS = 20; +const UPSTREAM_COMPLETION_TOKENS = 20; + +type DeploymentRow = { + model_name: string; + litellm_params?: { model?: string; custom_llm_provider?: string }; + model_info?: { id?: string }; +}; + +type CatalogRow = { input_cost_per_token?: number; output_cost_per_token?: number }; + +type SpendRow = { + request_id: string; + custom_llm_provider: string; + model: string; + model_group: string; + prompt_tokens: number; + completion_tokens: number; + spend: number | string; +}; + +const optionNamed = (label: string): RegExp => + new RegExp(`(^|\\s)${label.replace(/[.*+?^${}()|[\]\\]/g, "\\$&")}$`); + +async function loginAsAdmin(page: PlaywrightPage): Promise { + await page.goto("/ui/login"); + await page.getByPlaceholder("Enter your username").fill("admin"); + await page.getByPlaceholder("Enter your password").fill(master); + await page.getByRole("button", { name: "Login", exact: true }).click(); + await expect(page).toHaveURL( + (url) => url.pathname.startsWith("/ui") && !url.pathname.includes("login"), + ); +} + +async function selectProvider(page: PlaywrightPage, providerName: string): Promise { + const providerDropdown = page.getByRole("combobox", { name: "Provider", exact: true }); + await providerDropdown.click(); + await providerDropdown.fill(providerName); + await page.getByRole("option", { name: optionNamed(providerName) }).click(); + await expect(providerDropdown).toHaveValue(providerName); +} + +async function deploymentNamed( + request: APIRequestContext, + modelName: string, +): Promise { + const response = await request.get("/v1/model/info", { headers }); + expect(response.ok(), await response.text()).toBe(true); + const rows = (await response.json()).data as DeploymentRow[]; + return rows.find((row) => row.model_name === modelName); +} + +async function deleteDeployment(request: APIRequestContext, modelName: string): Promise { + const stored = await deploymentNamed(request, modelName); + const id = stored?.model_info?.id; + if (!id) return; + const deleted = await request.post("/model/delete", { headers, data: { id } }); + expect(deleted.ok(), await deleted.text()).toBe(true); +} + +async function chatOnce(request: APIRequestContext, modelName: string): Promise<{ status: number; body: string }> { + const response = await request.post("/v1/chat/completions", { + headers, + data: { model: modelName, messages: [{ role: "user", content: `hello from ${modelName}` }] }, + }); + return { status: response.status(), body: await response.text() }; +} + +async function spendRow(request: APIRequestContext, requestId: string): Promise { + const response = await request.get(`/spend/logs?request_id=${encodeURIComponent(requestId)}`, { headers }); + expect(response.ok(), await response.text()).toBe(true); + const rows = (await response.json()) as SpendRow[]; + return rows[0]; +} + +test("the Add Model form offers CoralBricks and the deployment it creates serves and bills a chat", async ({ + page, + request, +}) => { + const alias = `integration-ui-coralbricks-${randomUUID().replace(/-/g, "").slice(0, 12)}`; + try { + await loginAsAdmin(page); + await navigateToPage(page, Page.Models); + const catalogLoaded = page.waitForResponse( + (response) => response.url().includes("/public/litellm_model_cost_map?catalog_only=true") && response.ok(), + ); + await page.getByRole("tab", { name: "Add Model" }).click(); + const catalog = (await (await catalogLoaded).json()) as Record; + const catalogRow = catalog[`coralbricks/${CORALBRICKS_MODEL}`]; + + await selectProvider(page, "CoralBricks"); + + await page.getByRole("combobox", { name: "Select models" }).click(); + await page.getByRole("option", { name: "Custom Model Name (Enter below)" }).click(); + await page.keyboard.press("Escape"); + await page.getByPlaceholder("Enter custom model name").fill(CORALBRICKS_MODEL); + + const publicName = page.getByRole("row").filter({ hasText: CORALBRICKS_MODEL }).getByRole("textbox"); + await expect(publicName).toHaveValue(CORALBRICKS_MODEL); + await publicName.fill(alias); + + await page.locator("#api_base").fill(`${upstream}/v1`); + await page.locator("#api_key").fill(UI_KEY); + + const created = await captureRequestBody(page, { method: "POST", urlIncludes: "/model/new" }, async () => { + await page.getByTestId("add-model-btn").click(); + }); + expect(created.model_name, "the public name typed into the mapping row").toBe(alias); + expect(created.litellm_params?.model, "the custom model name typed into the form").toBe(CORALBRICKS_MODEL); + expect(created.litellm_params?.custom_llm_provider, "the provider the picker maps CoralBricks to").toBe( + "coralbricks", + ); + expect(created.litellm_params?.api_base).toBe(`${upstream}/v1`); + expect(created.litellm_params?.api_key).toBe(UI_KEY); + await expect(page.getByText("created successfully")).toBeVisible({ timeout: 15_000 }); + + let served = { status: 0, body: "" }; + await expect + .poll( + async () => { + served = await chatOnce(request, alias); + return served.status; + }, + { message: `the deployment added through the UI never served a chat: ${served.body}`, timeout: 60_000 }, + ) + .toBe(200); + const requestId = JSON.parse(served.body).id as string; + expect(requestId, served.body).toBeTruthy(); + + await expect + .poll(async () => (await deploymentNamed(request, alias))?.litellm_params?.custom_llm_provider, { + message: `no worker listed ${alias}`, + timeout: 60_000, + }) + .toBe("coralbricks"); + const inputRate = catalogRow?.input_cost_per_token; + const outputRate = catalogRow?.output_cost_per_token; + expect(inputRate, JSON.stringify(catalogRow)).toBeGreaterThan(0); + expect(outputRate, JSON.stringify(catalogRow)).toBeGreaterThan(0); + const expectedSpend = UPSTREAM_PROMPT_TOKENS * inputRate! + UPSTREAM_COMPLETION_TOKENS * outputRate!; + + let row: SpendRow | undefined; + await expect + .poll( + async () => { + row = await spendRow(request, requestId); + return row?.request_id; + }, + { message: `no spend row for ${requestId}`, timeout: 60_000 }, + ) + .toBe(requestId); + expect(row?.custom_llm_provider, JSON.stringify(row)).toBe("coralbricks"); + expect(row?.model, JSON.stringify(row)).toBe(created.litellm_params.model); + expect(row?.model_group, JSON.stringify(row)).toBe(alias); + expect(row?.prompt_tokens, JSON.stringify(row)).toBe(UPSTREAM_PROMPT_TOKENS); + expect(row?.completion_tokens, JSON.stringify(row)).toBe(UPSTREAM_COMPLETION_TOKENS); + expect(Number(row?.spend), JSON.stringify(row)).toBeCloseTo(expectedSpend, 10); + } finally { + await deleteDeployment(request, alias); + } +}); diff --git a/tests/e2e/ui/tests/integrationCritical/expected.json b/tests/e2e/ui/tests/integrationCritical/expected.json index 117df8d4892..e2b3f32645c 100644 --- a/tests/e2e/ui/tests/integrationCritical/expected.json +++ b/tests/e2e/ui/tests/integrationCritical/expected.json @@ -14,5 +14,6 @@ "tests/e2e/ui/tests/integrationCritical/anthropicFederationCredential.spec.ts::the Add Credential dialog stores each federation identity source with only the values it needs and refuses the invalid ones", "tests/e2e/ui/tests/integrationCritical/anthropicFederationCredential.spec.ts::the Edit Credential dialog sends only the changed values and deletes the keys the new choice leaves behind", "tests/e2e/ui/tests/integrationCritical/anthropicFederationCredential.spec.ts::a proxy admin saves a federation credential from the Add Model tab and the model is created against it", - "tests/e2e/ui/tests/integrationCritical/anthropicFederationCredential.spec.ts::a team admin who is not a proxy admin gets no federation shortcut on the Add Model tab" + "tests/e2e/ui/tests/integrationCritical/anthropicFederationCredential.spec.ts::a team admin who is not a proxy admin gets no federation shortcut on the Add Model tab", + "tests/e2e/ui/tests/integrationCritical/addModelCoralBricks.spec.ts::the Add Model form offers CoralBricks and the deployment it creates serves and bills a chat" ] diff --git a/tests/e2e_harness/AGENTS.md b/tests/e2e_harness/AGENTS.md index ef4a11802c9..65a3b32df96 100644 --- a/tests/e2e_harness/AGENTS.md +++ b/tests/e2e_harness/AGENTS.md @@ -1,6 +1,6 @@ # e2e harness tests -Tests of the harness under `tests/e2e/` (the transport, the clients, the fixture bundle and replay edge, the stack lock, the IdP launcher, the coverage collector, the JUnit properties, the load aggregation helpers and the Claude Code driver), not of the product. They live outside `tests/e2e/` because the Buildkite e2e run copies that folder into the runner image and runs every test in it, so a harness test in there counts as a product test in the nightly numbers. Nothing here needs a proxy, provider keys or the network +Tests of the harness under `tests/e2e/` (the transport, the clients, the fixture bundle and replay edge, the stack lock, the IdP launcher, the coverage collector, the JUnit properties and the Claude Code driver), not of the product. They live outside `tests/e2e/` because the Buildkite e2e run copies that folder into the runner image and runs every test in it, so a harness test in there counts as a product test in the nightly numbers. Nothing here needs a proxy, provider keys or the network The layout mirrors `tests/e2e/`: `test_e2e_http.py` covers `tests/e2e/e2e_http.py`, `logging/test_datadog_reader.py` covers `tests/e2e/logging/datadog_reader.py`, and `claude_code/` covers the driver, builder, probe and version resolver. Put a new harness test under the folder that mirrors the suite folder whose module it covers diff --git a/tests/e2e_harness/load/test_locust_load.py b/tests/e2e_harness/load/test_locust_load.py deleted file mode 100644 index af3e1483099..00000000000 --- a/tests/e2e_harness/load/test_locust_load.py +++ /dev/null @@ -1,232 +0,0 @@ -from __future__ import annotations - -from pathlib import Path -from typing import Final - -from locust_load import ( - LoadError, - LoadResult, - LocustStatEntry, - aggregate_stats, - percentile_seconds, - read_errors, - read_generator_warnings, -) - -_FAILURES_HEADER = "Method,Name,Error,Occurrences,First Seen,Last Seen\n" - - -def _entry( - *, - num_requests: int, - name: str = "/chat/completions", - num_failures: int = 0, - start_time: float = 1000.0, - last_request_timestamp: float = 1010.0, - response_times: dict[int, int] | None = None, -) -> LocustStatEntry: - return LocustStatEntry( - name=name, - num_requests=num_requests, - num_failures=num_failures, - start_time=start_time, - last_request_timestamp=last_request_timestamp, - response_times=response_times if response_times is not None else {50: num_requests}, - ) - - -def _result( - *, - errors: tuple[LoadError, ...] = (), - generator_warnings: tuple[str, ...] = (), -) -> LoadResult: - return LoadResult( - requests=10, - failures=10, - requests_per_second=1.0, - p50_seconds=0.05, - p90_seconds=0.08, - p99_seconds=0.1, - endpoints=(), - errors=errors, - generator_warnings=generator_warnings, - ) - - -class TestPercentiles: - def test_median_is_the_middle_sample_not_the_mean_a_slow_tail_would_drag(self) -> None: - # Nine fast requests and one very slow one: the mean is 1.99s, the median is 20ms. - entry = _entry(num_requests=10, response_times={20: 9, 20000: 1}) - - assert percentile_seconds([entry], 0.5) == 0.02 - - def test_the_tail_percentiles_reach_the_slow_samples_the_median_hides(self) -> None: - # 100 samples: 89 fast, 10 slow, 1 very slow. p50 sits in the fast bucket, p90 in the - # slow one, and p99 lands on the single very slow sample. - entry = _entry(num_requests=100, response_times={20: 89, 500: 10, 20000: 1}) - - assert percentile_seconds([entry], 0.5) == 0.02 - assert percentile_seconds([entry], 0.9) == 0.5 - assert percentile_seconds([entry], 0.99) == 0.5 - assert percentile_seconds([entry], 1.0) == 20.0 - - def test_percentiles_merge_the_histograms_of_every_stats_entry(self) -> None: - # Per entry the median would be 10ms and 90ms; merged, the middle of the five samples is 90ms. - entries = [ - _entry(num_requests=2, response_times={10: 2}), - _entry(num_requests=3, response_times={90: 3}), - ] - - assert percentile_seconds(entries, 0.5) == 0.09 - - def test_an_even_split_takes_the_lower_middle_sample_as_locust_itself_does(self) -> None: - entry = _entry(num_requests=4, response_times={10: 2, 90: 2}) - - assert percentile_seconds([entry], 0.5) == 0.01 - - def test_no_samples_reports_zero_rather_than_dividing_by_an_empty_histogram(self) -> None: - assert percentile_seconds([], 0.5) == 0.0 - - -class TestAggregate: - def test_throughput_spans_the_whole_window_and_latency_comes_from_the_histogram(self) -> None: - entry = _entry( - num_requests=180, - start_time=1000.0, - last_request_timestamp=1060.0, - response_times={57: 180}, - ) - - result = aggregate_stats([entry], (), ()) - - assert result.requests_per_second == 3.0 - assert result.p50_seconds == 0.057 - assert result.p99_seconds == 0.057 - assert result.failure_ratio == 0.0 - - def test_tail_percentiles_come_from_the_slow_end_of_the_histogram(self) -> None: - entry = _entry(num_requests=100, response_times={20: 89, 500: 10, 3000: 1}) - - result = aggregate_stats([entry], (), ()) - - assert result.p50_seconds == 0.02 - assert result.p90_seconds == 0.5 - assert result.p99_seconds == 0.5 - assert result.latency_summary() == "p50 0.020s, p90 0.500s, p99 0.500s" - - def test_throughput_spans_from_the_earliest_start_when_locust_reports_several_entries(self) -> None: - entries = [ - _entry(num_requests=60, start_time=1000.0, last_request_timestamp=1030.0), - _entry(num_requests=60, start_time=1020.0, last_request_timestamp=1060.0), - ] - - result = aggregate_stats(entries, (), ()) - - assert result.requests_per_second == 2.0 - - def test_a_run_that_drove_no_traffic_reports_a_total_failure_ratio(self) -> None: - result = aggregate_stats([], (), ()) - - assert result.requests == 0 - assert result.requests_per_second == 0.0 - assert result.failure_ratio == 1.0 - assert result.endpoints == () - - -class TestPerEndpoint: - def test_each_route_keeps_its_own_requests_failures_and_median(self) -> None: - entries: Final = ( - _entry(name="/chat/completions", num_requests=100, response_times={20: 100}), - _entry(name="/v1/messages", num_requests=40, num_failures=3, response_times={900: 40}), - ) - - result: Final = aggregate_stats(entries, (), ()) - - assert tuple((one.name, one.requests, one.failures, one.p50_seconds) for one in result.endpoints) == ( - ("/chat/completions", 100, 0, 0.02), - ("/v1/messages", 40, 3, 0.9), - ) - - def test_several_stats_entries_for_one_route_fold_into_a_single_row(self) -> None: - entries: Final = ( - _entry(name="/v1/messages", num_requests=10, response_times={30: 10}), - _entry(name="/v1/messages", num_requests=30, num_failures=1, response_times={30: 30}), - ) - - result: Final = aggregate_stats(entries, (), ()) - - assert tuple((one.name, one.requests, one.failures) for one in result.endpoints) == (("/v1/messages", 40, 1),) - - def test_a_route_that_never_ran_is_absent_so_a_one_sided_run_cannot_pass_unnoticed(self) -> None: - result: Final = aggregate_stats((_entry(name="/chat/completions", num_requests=10),), (), ()) - - assert tuple(one.name for one in result.endpoints) == ("/chat/completions",) - - def test_the_summary_names_every_route_with_its_counts(self) -> None: - entries: Final = ( - _entry(name="/chat/completions", num_requests=2, response_times={20: 2}), - _entry(name="/v1/messages", num_requests=1, num_failures=1, response_times={500: 1}), - ) - - result: Final = aggregate_stats(entries, (), ()) - - assert result.endpoint_summary() == ( - "/chat/completions 2 requests, 0 failures, p50 0.020s, /v1/messages 1 requests, 1 failures, p50 0.500s" - ) - - -class TestErrorBreakdown: - def test_locust_failure_rows_become_the_error_breakdown(self, tmp_path: Path) -> None: - failures_csv = tmp_path / "locust_failures.csv" - failures_csv.write_text( - _FAILURES_HEADER - + 'POST,/chat/completions,"LocustBadStatusCode(code=401)",381,2026-07-30 12:42:01,2026-07-30 12:45:00\n' - ) - - assert read_errors(failures_csv) == ( - LoadError(name="/chat/completions", error="LocustBadStatusCode(code=401)", occurrences=381), - ) - - def test_a_run_with_no_failures_writes_no_csv_and_reports_no_errors(self, tmp_path: Path) -> None: - assert read_errors(tmp_path / "locust_failures.csv") == () - - def test_diagnosis_leads_with_the_most_common_error(self) -> None: - result = _result( - errors=( - LoadError(name="/chat/completions", error="ConnectionRefused", occurrences=12), - LoadError(name="/chat/completions", error="LocustBadStatusCode(code=503)", occurrences=43675), - ) - ) - - assert result.diagnosis().startswith("43675x /chat/completions: LocustBadStatusCode(code=503)") - - def test_diagnosis_caps_the_list_and_says_how_many_it_left_out(self) -> None: - result = _result( - errors=tuple( - LoadError(name="/chat/completions", error=f"error-{index}", occurrences=index) for index in range(1, 9) - ) - ) - - assert result.diagnosis().count("x /chat/completions") == 5 - assert "and 3 more distinct errors" in result.diagnosis() - - def test_diagnosis_says_so_when_locust_recorded_nothing(self) -> None: - assert _result().diagnosis() == "locust recorded no error breakdown" - - -class TestGeneratorSaturation: - def test_repeated_cpu_warnings_collapse_to_one_and_reach_the_diagnosis(self) -> None: - stderr = ( - "[2026-07-31 12:47:01] WARNING/locust.runners: CPU usage above 90%!\n" - "[2026-07-31 12:47:02] INFO/locust.main: Run time limit reached\n" - "[2026-07-31 12:47:03] WARNING/locust.runners: CPU usage above 90%!\n" - ) - - warnings = read_generator_warnings(stderr) - - assert len(warnings) == 1 - assert "CPU usage above 90%!" in warnings[0] - assert "CPU usage above 90%!" in _result(generator_warnings=warnings).diagnosis() - - def test_ordinary_locust_chatter_is_not_reported_as_a_warning(self) -> None: - assert read_generator_warnings("[2026-07-31] INFO/locust.main: Shutting down (exit code 0)\n") == () diff --git a/tests/e2e_harness/load/test_phase_budget.py b/tests/e2e_harness/load/test_phase_budget.py deleted file mode 100644 index ea9e56afb0d..00000000000 --- a/tests/e2e_harness/load/test_phase_budget.py +++ /dev/null @@ -1,105 +0,0 @@ -from __future__ import annotations - -from typing import Final - -from phase_budget import AbsoluteBudget, RatioBudget, violations - - -def _budget(*, baseline: float, degraded: float, ceiling: float = 2.0) -> RatioBudget: - return RatioBudget( - name="p99 RSS", baseline=baseline, degraded=degraded, ratio_ceiling=ceiling, unit=" MB", decimals=0 - ) - - -class TestRatioBudget: - def test_growth_within_the_ceiling_is_not_a_violation(self) -> None: - assert _budget(baseline=100, degraded=199).violation() is None - - def test_growth_exactly_at_the_ceiling_is_allowed(self) -> None: - assert _budget(baseline=100, degraded=200).violation() is None - - def test_growth_past_the_ceiling_reports_both_values_and_the_ratio(self) -> None: - violation: Final = _budget(baseline=100, degraded=250).violation() - - assert violation is not None - assert "100 MB" in violation - assert "250 MB" in violation - assert "2.5x" in violation - assert "2.0x allowed" in violation - - def test_shrinking_is_never_a_violation(self) -> None: - assert _budget(baseline=100, degraded=10).violation() is None - - def test_a_missing_baseline_is_a_violation_rather_than_a_silent_pass(self) -> None: - # The trap this guards: 0 as a baseline would make every ratio a division by zero, and - # treating it as "no growth" would pass a run that measured nothing at all. - violation: Final = _budget(baseline=0, degraded=4000).violation() - - assert violation is not None - assert "nothing to compare" in violation - - def test_the_unit_and_decimals_carry_into_the_message(self) -> None: - violation: Final = RatioBudget( - name="p99 latency", baseline=0.16, degraded=9.5, ratio_ceiling=8.0, unit="s", decimals=3 - ).violation() - - assert violation is not None - assert "0.160s" in violation - assert "9.500s" in violation - - -class TestAbsoluteBudget: - def test_a_value_under_the_ceiling_is_not_a_violation(self) -> None: - assert AbsoluteBudget(name="p99 latency", measured=1.2, ceiling=5.0, unit="s", decimals=3).violation() is None - - def test_a_value_exactly_at_the_ceiling_is_allowed(self) -> None: - assert AbsoluteBudget(name="p99 latency", measured=5.0, ceiling=5.0, unit="s", decimals=3).violation() is None - - def test_a_value_past_the_ceiling_reports_the_measurement_and_the_ceiling(self) -> None: - violation: Final = AbsoluteBudget( - name="p99 latency", measured=9.5, ceiling=5.0, unit="s", decimals=3 - ).violation() - - assert violation is not None - assert "9.500s" in violation - assert "5.000s allowed" in violation - - def test_a_flat_ceiling_fails_a_degraded_phase_that_is_cheaper_than_its_baseline(self) -> None: - # The whole reason this shape exists: once the breaker opens, requests skip Redis instead - # of waiting on its socket timeout, so the chaos phase can measure faster than the healthy - # one. A ratio against that baseline passes; the user still waited 9.5s. - assert _budget(baseline=20.0, degraded=9.5, ceiling=2.0).violation() is None - assert AbsoluteBudget(name="p99 latency", measured=9.5, ceiling=5.0, unit="s").violation() is not None - - def test_a_zero_measurement_is_not_a_violation(self) -> None: - assert AbsoluteBudget(name="log bytes per request", measured=0, ceiling=12_000, unit=" B").violation() is None - - -class TestViolations: - def test_every_blown_budget_is_reported_not_just_the_first(self) -> None: - blown: Final = violations( - ( - _budget(baseline=100, degraded=500), - _budget(baseline=100, degraded=120), - RatioBudget(name="CPU per request", baseline=10, degraded=90, ratio_ceiling=6.0, unit=" ms"), - ) - ) - - assert len(blown) == 2 - assert blown[0].startswith("p99 RSS") - assert blown[1].startswith("CPU per request") - - def test_both_budget_shapes_report_together(self) -> None: - blown: Final = violations( - ( - _budget(baseline=100, degraded=500), - AbsoluteBudget(name="p99 latency", measured=9.5, ceiling=5.0, unit="s", decimals=3), - ) - ) - - assert len(blown) == 2 - assert blown[0].startswith("p99 RSS") - assert blown[1].startswith("p99 latency") - - def test_a_run_inside_every_budget_reports_nothing(self) -> None: - assert violations((_budget(baseline=100, degraded=150),)) == () diff --git a/tests/e2e_harness/load/test_proxy_usage.py b/tests/e2e_harness/load/test_proxy_usage.py deleted file mode 100644 index 915c564de50..00000000000 --- a/tests/e2e_harness/load/test_proxy_usage.py +++ /dev/null @@ -1,71 +0,0 @@ -from __future__ import annotations - -from typing import Final - -from proxy_usage import UsageSample, UsageWindow - -_MB: Final = 2**20 - - -def _window(*points: tuple[float, int, float]) -> UsageWindow: - return UsageWindow( - samples=tuple( - UsageSample(elapsed_seconds=elapsed, rss_bytes=rss, cpu_seconds=cpu) for elapsed, rss, cpu in points - ) - ) - - -class TestRssPercentiles: - def test_the_tail_percentiles_reach_the_peak_the_median_hides(self) -> None: - # 100 one-second samples: 89 flat, 10 elevated, 1 spike. The median stays flat, p90 sees the - # elevated plateau, and only the max reaches the spike. - window: Final = _window( - *((float(i), 100 * _MB, float(i)) for i in range(89)), - *((float(89 + i), 300 * _MB, float(89 + i)) for i in range(10)), - (99.0, 900 * _MB, 99.0), - ) - - assert window.rss_percentile(0.5) == 100 * _MB - assert window.rss_percentile(0.9) == 300 * _MB - assert window.rss_percentile(0.99) == 300 * _MB - assert window.rss_percentile(1.0) == 900 * _MB - - def test_an_empty_window_reports_zero_rather_than_indexing_nothing(self) -> None: - assert _window().rss_percentile(0.5) == 0 - - -class TestCpuUtilization: - def test_utilization_is_the_counter_delta_over_the_interval_not_the_counter_itself(self) -> None: - # The counter climbs 0.5 CPU seconds per second, then 4.0 per second: half a core, then four. - window: Final = _window((0.0, _MB, 0.0), (1.0, _MB, 0.5), (2.0, _MB, 1.0), (3.0, _MB, 5.0)) - - p50, p90, p99 = window.cpu_utilization_percentiles() - - assert (p50, p90, p99) == (0.5, 4.0, 4.0) - assert window.cpu_seconds_consumed() == 5.0 - - def test_a_single_sample_has_no_interval_and_reports_zero(self) -> None: - window: Final = _window((0.0, _MB, 3.0)) - - assert window.cpu_utilization_percentiles() == (0.0, 0.0, 0.0) - assert window.cpu_seconds_consumed() == 0.0 - - def test_cost_per_request_separates_runs_that_cores_busy_reports_identically(self) -> None: - # Both windows pin 4 cores for 10 seconds, so utilization cannot tell them apart. The - # second one served a tenth of the traffic for the same CPU, which is the regression shape. - window: Final = _window(*((float(i), _MB, 4.0 * i) for i in range(11))) - - assert window.cpu_utilization_percentiles()[0] == 4.0 - assert window.cpu_seconds_per_request(4000) == 0.01 - assert window.cpu_seconds_per_request(400) == 0.1 - - def test_no_requests_reports_zero_cost_rather_than_dividing_by_zero(self) -> None: - assert _window((0.0, _MB, 0.0), (1.0, _MB, 1.0)).cpu_seconds_per_request(0) == 0.0 - - def test_summary_reports_every_percentile_in_human_units(self) -> None: - window: Final = _window((0.0, 200 * _MB, 0.0), (1.0, 200 * _MB, 1.5), (2.0, 200 * _MB, 3.0)) - - assert window.summary() == ( - "RSS p50 200 MB, p90 200 MB, p99 200 MB; " - "CPU cores busy p50 1.50, p90 1.50, p99 1.50; 3.0 CPU seconds consumed" - ) diff --git a/tests/e2e_harness/load/test_session_anomaly.py b/tests/e2e_harness/load/test_session_anomaly.py deleted file mode 100644 index 7062587352b..00000000000 --- a/tests/e2e_harness/load/test_session_anomaly.py +++ /dev/null @@ -1,141 +0,0 @@ -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_harness/pytest.ini b/tests/e2e_harness/pytest.ini index dd629f0bc26..54d2630acc6 100644 --- a/tests/e2e_harness/pytest.ini +++ b/tests/e2e_harness/pytest.ini @@ -3,7 +3,7 @@ # (`from e2e_http import ...`, `from batch_cleanup import ...`) exactly as the suites # do, so tests/e2e and each suite folder that owns a module under test go on the path. addopts = --strict-markers --strict-config -p no:cacheprovider -pythonpath = ../e2e ../e2e/batches ../e2e/guardrails ../e2e/load ../e2e/logging +pythonpath = ../e2e ../e2e/batches ../e2e/guardrails ../e2e/logging markers = covers(cell_id, *, exercised_on=()): coverage-registry cell(s) a test covers; exercised here only to prove the collector and the JUnit properties read it cli_determinism: drives the real claude CLI for several seconds diff --git a/tests/integration/README.md b/tests/integration/README.md index a90faa0823d..572b80ad5c6 100644 --- a/tests/integration/README.md +++ b/tests/integration/README.md @@ -4,11 +4,13 @@ These tests exercise a running gateway, PostgreSQL and Redis with an owned local The ROI database contracts in `database/test_roi_observed.py` run as plain pytest outside `run_integration.sh` in CircleCI's `roi-database` Postgres job. The job uploads coverage with the `roi-postgres` flag. These contracts own temporary databases and script only the external provider transport +The Redis cluster IAM contracts in `sdk/test_redis_cluster_iam_auth.py` start their own `redis-server` and mock the cloud credential clients, so they run in the GitHub Actions `redis-compat` matrix against every supported redis-py major rather than in `run_integration.sh` + The `cost` group is driven by the `CostTrackingTestCase` literals under `cost_calculation/cases/.py` (gathered by `cost_calculation/catalog.py`), with the scripted cost map in `cost_calculation/cost_map.py`. Each case has a name, contract ID, cost-map model, optional deployment overrides, request body, tagged response and exact or recount expectations. Request bodies use `$MODEL` for the registered proxy model, while responses use `$REQUEST_ID` for the per-run scenario ID. To add a case, add a cost-map entry when the model is new, add the request body and exact provider response data, and add hand-computed expected values. The upstream serves each stored response for any path under `/`, while the test-owned cost map is served over loopback through `LITELLM_MODEL_COST_MAP_URL` -Use `tests/integration/run.py management`, `accounting`, `database`, `providers`, `extensions`, `mcp`, `sdk` or `cost` to run a selected group. The group to directory mapping is the `GROUPS` literal at the top of `run.py`; a new directory needs a `GROUPS` entry and an `OWNED_DIRECTORIES` entry in `_support/manifest.py`. Set `INTEGRATION_WORKERS` above 1 to run a group under pytest-xdist. Management, accounting, providers and extensions give each worker a private database, Redis, upstream, fake provider and proxy, with a peer and read-only database role where required. The pytest controller owns these services and closes them after worker failures. Their cases use pytest-xdist load groups inside the files selected by CircleCI timing splits. Module and class fixtures are grouped automatically, so authors do not need scheduling markers. Function fixtures remain independently schedulable. Tests that own an independent mutable sink for every case should use function fixtures, so a shared module fixture does not serialize them or share pending observations. Nested harness probes discard the parent run's shard, worker and result settings, then supply their own selections. MCP keeps its existing resource ownership and explicit load groups. Automatic group suffixes are removed from JUnit and execution records to preserve case identities. Owned proxy helpers wait for a healthy response from every configured, live worker before yielding. Serial runs keep the supplied deployment, including routing parity runs. Set `INTEGRATION_PROXY_URL`, `INTEGRATION_UPSTREAM_URL`, `INTEGRATION_MASTER_KEY` and `DATABASE_URL` to an isolated test deployment. When that deployment runs more than one proxy worker, set `INTEGRATION_PROXY_WORKERS` to the count so a test that writes a model and then calls it waits out the config reload interval, the only cross-worker convergence bound the wire exposes. The runner selects the new domain directories explicitly; the legacy OCI and sandbox selections remain separate +Use `tests/integration/run.py management`, `accounting`, `database`, `providers`, `extensions`, `mcp`, `sdk` or `cost` to run a selected group. The group to directory mapping is the `GROUPS` literal at the top of `run.py`; a new directory needs a `GROUPS` entry and an `OWNED_DIRECTORIES` entry in `_support/manifest.py`. Set `INTEGRATION_WORKERS` above 1 to run a group under pytest-xdist. Every group run by more than one worker gives each worker a private database, Redis, scripted upstream process, fake provider and proxy, with a peer and read-only database role where required. There is no opt-in list and no shared proxy, upstream, database or Redis in that mode: `run_integration.sh` does not start them and does not hand their addresses to the run, so a test can only reach its own worker's services. The pytest controller owns these services and closes them after worker failures. Their cases use pytest-xdist load groups inside the files selected by CircleCI timing splits. Module and class fixtures are grouped automatically, so authors do not need scheduling markers. Function fixtures remain independently schedulable. Tests that own an independent mutable sink for every case should use function fixtures, so a shared module fixture does not serialize them or share pending observations. Nested harness probes discard the parent run's shard, worker and result settings, then supply their own selections. MCP keeps its explicit load groups. Automatic group suffixes are removed from JUnit and execution records to preserve case identities. Owned proxy helpers wait for a healthy response from every configured, live worker before yielding. Serial runs keep the supplied deployment, including routing parity runs. Set `INTEGRATION_PROXY_URL`, `INTEGRATION_UPSTREAM_URL`, `INTEGRATION_MASTER_KEY` and `DATABASE_URL` to an isolated test deployment. When that deployment runs more than one proxy worker, set `INTEGRATION_PROXY_WORKERS` to the count so a test that writes a model and then calls it waits out the config reload interval, the only cross-worker convergence bound the wire exposes. The runner selects the new domain directories explicitly; the legacy OCI and sandbox selections remain separate -Management also requires `INTEGRATION_PEER_URL`, `REDIS_HOST` and `REDIS_PORT`. CircleCI starts two directly addressed proxy processes sharing only that job's stores. The test-only CLI wrapper supplies enterprise route entitlement, following the existing behavior suite's convention. It does not qualify license validation; run it with one worker and no reload +Management and MCP also require `INTEGRATION_PEER_URL`, `REDIS_HOST` and `REDIS_PORT`. Under several workers each worker gets a second directly addressed proxy sharing only that worker's stores, and a single-worker run uses the two proxies `run_integration.sh` starts. The test-only CLI wrapper supplies enterprise route entitlement, following the existing behavior suite's convention. It does not qualify license validation; run it with one worker and no reload The generated lifecycle models use 20 examples, eight steps, generation and shrinking, with isolated resources per example. HTTP operation caps include generation and shrinking and exempt cleanup. Local qualification defaults to seed 4106601 and canonical order; CircleCI derives exploration and ordering seeds from the checked-out revision and workflow ID. The ordering seed shuffles files, scoped fixture configurations and cases within each configuration. Ordering keeps shared module and class configurations together where their dependencies permit it. Dependent fixtures with their own module parameters can still rebuild an outer configuration, with every case retained @@ -24,7 +26,7 @@ Reuse the existing canned provider handlers through `_support/upstream.py`. It r Owned proxies disable `.env` loading by default. A contract testing dotenv behavior can explicitly set `PYTHON_DOTENV_DISABLED=0` in its owned environment. An explicit scratch writer database drops an inherited reader with a different database name; pass both database URLs when the contract needs a paired reader. Database relays close listeners, active connections and their event loops on teardown. Upstream startup failures reap the child process before returning the failure. Custom subprocesses, direct SQL and background tasks still need an explicit owner in the test -Owned proxy workers configure a 32-thread executor before application startup, preserving additional worker startup hooks. Held synchronous provider calls must not consume a CPU-dependent executor limit before the rest of a fault burst can start. The proxy entrypoint imports LiteLLM inside `main`, so spawned workers begin supervision before importing the application. Worker healthchecks have a five-second deadline; application readiness and graceful teardown retain their separate longer deadlines. A fault that deliberately pauses Python before it can answer worker pings must supply a longer healthcheck through `extra_arguments`. Slow application startup is covered separately by the readiness deadline. Cases with two complete owned proxy lifetimes use a 180-second pytest budget for startup, assertions and teardown; their HTTP, provider and persistence deadlines stay unchanged. Readiness requires responsive worker identities that are still alive and a positively identified owned listener. Unreadable socket records cannot establish ownership. Crash tests own their client tasks with `TaskGroup`, release held peers in `finally` and verify that the intended worker dies. A test requiring traffic on both workers prepares that traffic explicitly instead of assuming OS load balancing +Owned proxy workers configure a 32-thread executor before application startup, preserving additional worker startup hooks. Held synchronous provider calls must not consume a CPU-dependent executor limit before the rest of a fault burst can start. The proxy entrypoint imports LiteLLM inside `main`, so spawned workers begin supervision before importing the application. Worker healthchecks have a five-second deadline by default, so uvicorn kills a worker that stays paused longer than that. Tests never signal proxy processes themselves. `OwnedProxy.kill_worker` SIGKILLs a worker and returns once its replacement answers readiness with its own pid, `replacing_worker` does the same for a kill that has to land mid-burst, and `suspended` pauses a worker or the supervisor, which only works on a proxy launched with `healthcheck=WorkerHealthcheck.LENIENT`. Helper processes that are not the proxy stop through `stop_child`, `stop_processes` and `stop_process_group` and pause through `paused`. `scripts/check_test_quality.py` rule TQ010 rejects a raw `os.kill`, SIGKILL / SIGSTOP / SIGCONT, `.kill()`, `.terminate()`, `.suspend()` or `.send_signal()` anywhere under `tests/integration` outside `_support`, and TQ011 rejects an assertion that caps an elapsed time with a fixed number. Bound timings with `_support/timing.py` instead, relative to a measured control or to the injected delay. `run.py` measures how long this box takes to import the proxy application before it starts pytest, records it in `boot-calibration.json` beside the results and exports it as `INTEGRATION_IMPORT_SECONDS`. Readiness deadlines, graceful stops and the default per-test timeout are multiples of that measurement (`_support/budget.py`), so a slower box gets longer budgets and a two-worker proxy gets twice the readiness deadline of a one-worker one. Readiness requires responsive worker identities that are still alive and a positively identified owned listener. Unreadable socket records cannot establish ownership. Crash tests own their client tasks with `TaskGroup`, release held peers in `finally` and verify that the intended worker dies. A test requiring traffic on both workers prepares that traffic explicitly instead of assuming OS load balancing, for example by sending each half while the other worker is suspended The CircleCI workflow starts its own database and Redis, restricts test-phase egress to its owned services and writes JUnit plus an executed-node manifest. Missing setup, failed cleanup or a selected test with neither a passed call nor a skip fail qualification. Skipped nodes are listed under `skipped` in `execution.json`, so the skip reasons double as the open bug list. `GITHUB_FILES` in `run.py` lists the files excluded from the `run_integration.sh` selection @@ -52,9 +54,9 @@ The sdk shard exercises the SDK's own HTTP clients against local protocol peers The extensions shard uses the built-in generic callback and guardrail transports. It checks callback correlation and credential exclusion, guardrail rewriting and denial, retained OpenAI consumers and A2A wire versions. CircleCI runs it on eight xlarge nodes using its file timing splitter. Each node runs its assigned files with four pytest workers, each owning its database, Redis, upstream and proxy. The dependent coverage job verifies all eight execution records against the candidate case inventory. Tests must not assume a particular set of sibling files -The security shard uses four pytest workers on an xlarge machine, each with a private database and Redis. Route sweeps use a new HTTP client for every call and share only the TLS verification context, so authentication and response cookies cannot cross callers. The canary rig supplies owned OpenAI and Anthropic endpoints and synthetic credentials and rejects attempted external egress +The security shard uses four pytest workers on an xlarge machine, each with its own isolated services like every other multi-worker group. Route sweeps use a new HTTP client for every call and share only the TLS verification context, so authentication and response cookies cannot cross callers. The canary rig supplies owned OpenAI and Anthropic endpoints and synthetic credentials and rejects attempted external egress -The two mcp nodes own separate services and use the same case census as providers and extensions. Their coverage data is combined into one MCP report after both nodes finish. Each node runs the MCP gateway against SDK peers owned by each test (`_support/mcp.py`): streamable HTTP, SSE and stdio peers, an OpenAPI-spec app, and an OAuth 2.1 authorization-server double. Every peer records the requests it receives so a test can assert what reached the peer, not only what the proxy answered. The shard runs with `INTEGRATION_WORKERS` set and with `INTEGRATION_COVERAGE=1`, which starts the proxy under `coverage run --parallel-mode` limited to the MCP modules and stores `coverage.txt` plus an HTML report with the job artifacts. A test that fails because the product is wrong is skipped with `pytest.skip("BUG: ")` so the skip list in `execution.json` is the open MCP bug list +The two mcp nodes own separate services and use the same case census as providers and extensions. Their coverage data is combined into one MCP report after both nodes finish. Each node runs the MCP gateway against SDK peers owned by each test (`_support/mcp.py`): streamable HTTP, SSE and stdio peers, an OpenAPI-spec app, and an OAuth 2.1 authorization-server double. Every peer records the requests it receives so a test can assert what reached the peer, not only what the proxy answered. The shard runs with `INTEGRATION_WORKERS` set and with `INTEGRATION_COVERAGE=1`, which starts every worker's proxy and peer under `coverage run --parallel-mode` limited to the MCP modules and stores `coverage.txt` plus an HTML report with the job artifacts. `mcp_coverage.toml` measures lines through `sys.monitoring`, because the CPython 3.12 trace function that branch measurement needs made a covered proxy take about four times as long to become ready. A test that fails because the product is wrong is skipped with `pytest.skip("BUG: ")` so the skip list in `execution.json` is the open MCP bug list Browser contracts live in `tests/e2e/ui/tests/integrationCritical` and run only through `tests/e2e/ui/integration.config.ts`. The expected browser results are listed in `expected.json` in that directory and checked by `.circleci/scripts/verify_integration_browser.py`. The CircleCI browser shard builds the checked-out dashboard, starts the owned proxy with that build, and verifies one exact browser result without retries or skips. The default Playwright selection excludes this directory. The focused project flow asserts the submitted create and clear values, fresh SQL state and actual blocked/restored serving while preserving model restrictions diff --git a/tests/integration/_support/budget.py b/tests/integration/_support/budget.py new file mode 100644 index 00000000000..c9aa4b0802e --- /dev/null +++ b/tests/integration/_support/budget.py @@ -0,0 +1,64 @@ +import json +import os +import subprocess +import sys +import time +from functools import cache +from pathlib import Path +from typing import Final + +IMPORT_SECONDS_VARIABLE: Final = "INTEGRATION_IMPORT_SECONDS" +READINESS_OVERRIDE_VARIABLE: Final = "INTEGRATION_PROXY_READY_SECONDS" +READINESS_IMPORTS: Final = 12 +CALL_IMPORTS: Final = 6 +PROXY_IMPORT: Final = "import litellm.proxy.proxy_server" + + +def measure_import_seconds(root: Path, environment: dict[str, str]) -> float: + started: Final = time.monotonic() + subprocess.run( + (sys.executable, "-P", "-c", PROXY_IMPORT), cwd=root, env=environment, capture_output=True, check=True + ) + return time.monotonic() - started + + +def boot_units(workers: int) -> int: + return 1 if workers <= 1 else 2 + + +def readiness_for(imported: float, workers: int) -> float: + return READINESS_IMPORTS * imported * boot_units(workers) + + +def case_timeout_for(imported: float) -> float: + return readiness_for(imported, 2) + readiness_for(imported, 1) + CALL_IMPORTS * imported + + +def calibration(imported: float) -> str: + return json.dumps( + { + "import_seconds": imported, + "readiness_seconds_1_worker": readiness_for(imported, 1), + "readiness_seconds_2_workers": readiness_for(imported, 2), + "case_timeout_seconds": case_timeout_for(imported), + }, + indent=2, + ) + + +@cache +def import_seconds() -> float: + exported: Final = os.environ.get(IMPORT_SECONDS_VARIABLE) + if exported: + return float(exported) + root: Final = Path(os.environ.get("INTEGRATION_PROXY_ROOT") or Path(__file__).resolve().parents[3]) + return measure_import_seconds(root, dict(os.environ)) + + +def readiness_seconds(workers: int = 1) -> float: + override: Final = os.environ.get(READINESS_OVERRIDE_VARIABLE) + return float(override) if override else readiness_for(import_seconds(), workers) + + +def case_timeout_seconds() -> float: + return case_timeout_for(import_seconds()) diff --git a/tests/integration/_support/client.py b/tests/integration/_support/client.py index 1c58284e42f..b22c1b1da3a 100644 --- a/tests/integration/_support/client.py +++ b/tests/integration/_support/client.py @@ -240,11 +240,20 @@ class Scenario: info: Final = object_value(self.gateway.get("/key/info", {"key": hashed})["info"]) assert info["status"] == "deleted", f"Deleted key still served as live: {info['status']}" - def delete_model(self, identity: str) -> None: - self.gateway.post("/model/delete", {"id": identity}) + def served_model_ids(self) -> frozenset[JsonValue]: entries: Final = self.gateway.get("/model/info")["data"] assert isinstance(entries, list) - assert all(object_value(object_value(entry)["model_info"])["id"] != identity for entry in entries) + return frozenset(object_value(object_value(entry)["model_info"])["id"] for entry in entries) + + def delete_model(self, identity: str) -> None: + from litellm.constants import PROXY_CONFIG_RELOAD_INTERVAL_SECONDS + + self.gateway.post("/model/delete", {"id": identity}) + eventually( + self.served_model_ids, + lambda served: identity not in served, + seconds=2 * PROXY_CONFIG_RELOAD_INTERVAL_SECONDS, + ) assert read_rows('SELECT model_id FROM "LiteLLM_ProxyModelTable" WHERE model_id = %s', (identity,)) == [] def model(self, *, model_info: Mapping[str, JsonValue] | None = None, **parameters: JsonValue) -> str: diff --git a/tests/integration/_support/database.py b/tests/integration/_support/database.py index 3ad4f205e12..ef7d2ba68e1 100644 --- a/tests/integration/_support/database.py +++ b/tests/integration/_support/database.py @@ -27,8 +27,8 @@ def write_rows(query: LiteralString, parameters: tuple[str, ...], *, database_ur @contextmanager -def scratch_database() -> Generator[str]: - name: Final = f"integration_{uuid.uuid4().hex}" +def scratch_database(prefix: str = "integration") -> Generator[str]: + name: Final = f"{prefix}_{uuid.uuid4().hex}" with psycopg.connect(os.environ["DATABASE_URL"], autocommit=True) as admin: admin.execute(sql.SQL("CREATE DATABASE {}").format(sql.Identifier(name))) try: diff --git a/tests/integration/_support/process.py b/tests/integration/_support/process.py index 4f8ca036489..5479629420a 100644 --- a/tests/integration/_support/process.py +++ b/tests/integration/_support/process.py @@ -1,5 +1,6 @@ import errno import json +import threading import os import re import signal @@ -8,9 +9,10 @@ import subprocess import sys import time import uuid -from collections.abc import Generator, Iterator, Mapping -from contextlib import contextmanager +from collections.abc import Generator, Iterable, Iterator, Mapping +from contextlib import contextmanager, suppress from dataclasses import dataclass +from enum import Enum from datetime import UTC, datetime, timedelta from ipaddress import IPv4Address from itertools import accumulate @@ -28,17 +30,19 @@ from cryptography.hazmat.primitives.asymmetric import rsa from cryptography.x509.oid import NameOID from pydantic import TypeAdapter +from tests.integration._support.budget import import_seconds, readiness_for, readiness_seconds from tests.integration._support.client import GATEWAY_LIMITS, Gateway DB_PUSH: Final = ("--use_prisma_db_push",) MIGRATE_DEPLOY: Final = () LEGACY_MIGRATE_DEPLOY: Final = ("--use_legacy_migration_resolver",) CONFIGURATION: Final = TypeAdapter(dict[str, object]) +MCP_COVERAGE_CONFIG: Final = Path(__file__).resolve().parents[1] / "mcp_coverage.toml" -def proxy_database_environment() -> Mapping[str, str]: - writer: Final = os.environ.get("INTEGRATION_PROXY_DATABASE_URL", "") - reader: Final = os.environ.get("INTEGRATION_PROXY_READ_REPLICA_URL", "") +def proxy_database_environment(inherited: Mapping[str, str] = os.environ) -> Mapping[str, str]: + writer: Final = inherited.get("INTEGRATION_PROXY_DATABASE_URL", "") + reader: Final = inherited.get("INTEGRATION_PROXY_READ_REPLICA_URL", "") return MappingProxyType( { **({"DATABASE_URL": writer} if writer else {}), @@ -66,7 +70,7 @@ def signal_group(group: int, action: int) -> None: def graceful_stop_seconds() -> float: - return max(30.0, float(os.environ.get("INTEGRATION_PROXY_READY_SECONDS", "70"))) + return readiness_for(import_seconds(), 1) def stop_root_process(process: subprocess.Popen[bytes]) -> bool: @@ -80,11 +84,247 @@ def stop_root_process(process: subprocess.Popen[bytes]) -> bool: return True +class WorkerHealthcheck(Enum): + STRICT = "strict" + LENIENT = "lenient" + + +STRICT_HEALTHCHECK_SECONDS: Final = 5 +WORKER_ENTRYPOINT: Final = "spawn_main" + + +def healthcheck_seconds(healthcheck: WorkerHealthcheck) -> int: + match healthcheck: + case WorkerHealthcheck.STRICT: + return STRICT_HEALTHCHECK_SECONDS + case WorkerHealthcheck.LENIENT: + return int(graceful_stop_seconds()) + + +def _is_serving_worker(process: psutil.Process) -> bool: + try: + return _ready_worker_pid(process) is not None and WORKER_ENTRYPOINT in " ".join(process.cmdline()) + except psutil.Error: + return False + + +def _is_gone(process: psutil.Process) -> bool: + try: + return process.status() == psutil.STATUS_ZOMBIE + except psutil.NoSuchProcess: + return True + + +def _wait_until_gone(process: psutil.Process, seconds: float) -> None: + deadline: Final = time.monotonic() + seconds + while not _is_gone(process): + assert time.monotonic() < deadline, f"Process {process.pid} is still running after {seconds:.0f}s" + time.sleep(0.05) + + +def _is_stopped(process: psutil.Process) -> bool: + try: + return process.status() == psutil.STATUS_STOPPED + except psutil.NoSuchProcess: + return False + + +def _suspend(process: psutil.Process) -> None: + process.suspend() + deadline: Final = time.monotonic() + STRICT_HEALTHCHECK_SECONDS + while not _is_stopped(process) and time.monotonic() < deadline: + time.sleep(0.01) + if _is_stopped(process): + return + with suppress(psutil.NoSuchProcess): + process.resume() + raise AssertionError(f"Process {process.pid} did not stop within {STRICT_HEALTHCHECK_SECONDS}s of SIGSTOP") + + +def _is_reaped(process: psutil.Process) -> bool: + try: + process.status() + except psutil.NoSuchProcess: + return True + return False + + +def _wait_until_reaped(process: psutil.Process, supervisor: psutil.Process, seconds: float) -> None: + deadline: Final = time.monotonic() + seconds + while not _is_reaped(process) and not _is_stopped(supervisor): + assert time.monotonic() < deadline, f"The supervisor did not reap worker {process.pid} within {seconds:.0f}s" + time.sleep(0.05) + + +@dataclass(frozen=True, slots=True) +class WorkerKill: + victim: psutil.Process + supervisor: psutil.Process + reap_seconds: float + fired: threading.Event + + def freeze(self) -> psutil.Process: + _suspend(self.victim) + return self.victim + + def __call__(self) -> None: + assert not self.fired.is_set(), f"Worker {self.victim.pid} was already killed" + with suppress(psutil.NoSuchProcess): + self.victim.kill() + _wait_until_gone(self.victim, STRICT_HEALTHCHECK_SECONDS) + _wait_until_reaped(self.victim, self.supervisor, self.reap_seconds) + self.fired.set() + + +def replacement_pid(before: frozenset[int], answered: int | None, serving: frozenset[int]) -> int | None: + return answered if answered is not None and answered not in before and answered in serving else None + + @dataclass(frozen=True, slots=True) class OwnedProxy: gateway: Gateway process: subprocess.Popen[bytes] log: Path + healthcheck: WorkerHealthcheck = WorkerHealthcheck.STRICT + + def worker_pids(self) -> frozenset[int]: + try: + children: Final = psutil.Process(self.process.pid).children() + except psutil.NoSuchProcess: + return frozenset() + return frozenset(child.pid for child in children if _is_serving_worker(child)) + + def _answering_pid(self) -> int | None: + try: + with httpx.Client( + base_url=str(self.gateway.client.base_url), + timeout=2, + trust_env=False, + limits=httpx.Limits(max_keepalive_connections=0), + ) as client: + response: Final = client.get("/health/integration-readiness") + except httpx.TransportError: + return None + identity: Final = response.headers.get("x-integration-worker-pid", "") + return int(identity) if response.status_code == 200 and identity.isdecimal() else None + + def _await_replacement(self, before: frozenset[int]) -> int: + budget: Final = readiness_seconds(1) + deadline: Final = time.monotonic() + budget + while (replacement := replacement_pid(before, self._answering_pid(), self.worker_pids())) is None: + assert self.process.poll() is None, "Owned proxy exited while a worker was being replaced" + assert time.monotonic() < deadline, f"No replacement worker served within {budget:.0f}s" + time.sleep(0.05) + return replacement + + def _victim(self, pid: int) -> tuple[psutil.Process, frozenset[int]]: + before: Final = self.worker_pids() + assert pid in before, f"{pid} is not a serving worker of this proxy: {sorted(before)}" + return psutil.Process(pid), before + + def _worker_kill(self, victim: psutil.Process) -> WorkerKill: + reap_seconds: Final = healthcheck_seconds(self.healthcheck) + STRICT_HEALTHCHECK_SECONDS + return WorkerKill(victim, psutil.Process(self.process.pid), reap_seconds, threading.Event()) + + @contextmanager + def replacing_worker(self, pid: int) -> Iterator[WorkerKill]: + victim, before = self._victim(pid) + kill: Final = self._worker_kill(victim) + yield kill + assert kill.fired.is_set(), f"replacing_worker({pid}) exited without killing the worker" + self._await_replacement(before) + + def kill_worker(self, pid: int) -> int: + victim, before = self._victim(pid) + self._worker_kill(victim)() + return self._await_replacement(before) + + @contextmanager + def suspended(self, pid: int) -> Iterator[psutil.Process]: + assert self.healthcheck is WorkerHealthcheck.LENIENT, ( + "uvicorn kills a suspended worker after the strict worker healthcheck; " + "launch this proxy with healthcheck=WorkerHealthcheck.LENIENT" + ) + assert pid == self.process.pid or pid in self.worker_pids(), f"{pid} is not part of this proxy" + target: Final = psutil.Process(pid) + _suspend(target) + try: + yield target + finally: + with suppress(psutil.NoSuchProcess): + target.resume() + assert not _is_gone(target), ( + f"Process {pid} died while suspended under a {healthcheck_seconds(self.healthcheck)}s worker healthcheck" + ) + + def stall_worker(self, pid: int) -> int: + assert self.healthcheck is WorkerHealthcheck.STRICT, "Only the strict worker healthcheck replaces a stall" + victim, before = self._victim(pid) + victim.suspend() + try: + replacement: Final = self._await_replacement(before) + _wait_until_gone(victim, readiness_seconds(1)) + return replacement + finally: + with suppress(psutil.NoSuchProcess): + victim.kill() + + def begin_shutdown(self) -> None: + self.process.terminate() + + def kill_supervisor(self) -> int: + self.process.kill() + return self.process.wait(timeout=STRICT_HEALTHCHECK_SECONDS) + + +def stop_child(process: subprocess.Popen[bytes] | subprocess.Popen[str], grace_seconds: float = 10) -> int: + if process.poll() is None: + process.terminate() + try: + return process.wait(timeout=grace_seconds) + except subprocess.TimeoutExpired: + process.kill() + return process.wait(timeout=grace_seconds) + + +def stop_processes(processes: Iterable[psutil.Process], grace_seconds: float = 5) -> None: + members: Final = tuple(processes) + for member in members: + with suppress(psutil.NoSuchProcess): + member.terminate() + _stopped, alive = psutil.wait_procs(members, timeout=grace_seconds) + for member in alive: + with suppress(psutil.NoSuchProcess): + member.kill() + psutil.wait_procs(alive, timeout=grace_seconds) + + +def stop_process_group(process: subprocess.Popen[bytes] | subprocess.Popen[str], grace_seconds: float) -> None: + if process.poll() is None: + signal_group(process.pid, signal.SIGTERM) + try: + process.wait(timeout=grace_seconds) + except subprocess.TimeoutExpired: + signal_group(process.pid, signal.SIGKILL) + process.wait(timeout=grace_seconds) + + +def kill_stray_process(pid: int) -> None: + with suppress(psutil.NoSuchProcess): + stray: Final = psutil.Process(pid) + stray.kill() + _wait_until_gone(stray, STRICT_HEALTHCHECK_SECONDS) + + +@contextmanager +def paused(process: subprocess.Popen[bytes] | subprocess.Popen[str]) -> Iterator[None]: + target: Final = psutil.Process(process.pid) + _suspend(target) + try: + yield + finally: + with suppress(psutil.NoSuchProcess): + target.resume() @contextmanager @@ -97,6 +337,8 @@ def owned_proxy( remove_environment: tuple[str, ...] = (), workers: int = 1, database_setup: tuple[str, ...] = DB_PUSH, + coverage_file: Path | None = None, + inherited: Mapping[str, str] = os.environ, ) -> Iterator[Gateway]: with owned_proxy_process( gateway, @@ -106,6 +348,8 @@ def owned_proxy( remove_environment=remove_environment, workers=workers, database_setup=database_setup, + coverage_file=coverage_file, + inherited=inherited, ) as owned: yield owned.gateway @@ -151,8 +395,13 @@ class _Launch: started: float +def _wall_clock(monotonic: float) -> float: + return time.time() - (time.monotonic() - monotonic) + + def _stop_launch(launch: _Launch) -> None: started: Final = time.monotonic() + stop_started_at: Final = time.time() exit_before: Final = launch.process.poll() exception_before: Final = sys.exc_info()[0] try: @@ -169,6 +418,9 @@ def _stop_launch(launch: _Launch) -> None: "exception_at_record": exception_at_record.__name__ if exception_at_record else None, "cleanup_seconds": time.monotonic() - started, "process_lifetime_seconds": time.monotonic() - launch.started, + "spawned_at": _wall_clock(launch.started), + "stop_started_at": stop_started_at, + "stopped_at": time.time(), }, indent=2, ) @@ -209,28 +461,80 @@ def _owns_listener(launch: _Launch) -> bool: return False -def _readiness_pids(launch: _Launch) -> Iterator[int]: +@dataclass(frozen=True, slots=True) +class BootPhases: + prisma_started: float | None = None + prisma_finished: float | None = None + listening: float | None = None + first_ready: float | None = None + ready_pids: frozenset[int] = frozenset() + + +@dataclass(frozen=True, slots=True) +class _Probe: + elapsed: float + listening: bool + prisma: bool + ready_pid: int | None + + +def _runs_prisma(process: psutil.Process) -> bool: + try: + return any(Path(part).name == "prisma" for part in process.cmdline()) + except psutil.Error: + return False + + +def _prisma_running(launch: _Launch) -> bool: + try: + children: Final = psutil.Process(launch.process.pid).children() + except psutil.Error: + return False + return any(_runs_prisma(child) for child in children) + + +def _answer(client: httpx.Client, launch: _Launch) -> tuple[bool, int | None]: + try: + response: Final = client.get("/health/integration-readiness", timeout=2) + except httpx.ConnectError: + return False, None + except httpx.TransportError: + return True, None + identity: Final = response.headers.get("x-integration-worker-pid", "") + ready: Final = response.status_code == 200 and identity.isdecimal() and _owns_listener(launch) + return True, int(identity) if ready else None + + +def _probes(launch: _Launch, workers: int) -> Iterator[_Probe]: with httpx.Client( base_url=f"http://127.0.0.1:{launch.port}", timeout=15, trust_env=False, limits=httpx.Limits(max_keepalive_connections=0), ) as client: - deadline: Final = time.monotonic() + float(os.environ.get("INTEGRATION_PROXY_READY_SECONDS", "70")) + deadline: Final = time.monotonic() + readiness_seconds(workers) while launch.process.poll() is None: - try: - response: Final = client.get("/health/integration-readiness", timeout=2) - identity: Final = response.headers.get("x-integration-worker-pid", "") - if response.status_code == 200 and identity.isdecimal() and _owns_listener(launch): - yield int(identity) - except httpx.TransportError: - pass + elapsed: Final = time.monotonic() - launch.started + listening, ready_pid = _answer(client, launch) + yield _Probe(elapsed, listening, not listening and _prisma_running(launch), ready_pid) assert time.monotonic() < deadline, "Owned proxy readiness deadline exceeded" time.sleep(0.1) -def _ready_workers(pids: frozenset[int], pid: int) -> frozenset[int]: - return pids | frozenset((pid,)) +def _first(recorded: float | None, observed: bool, elapsed: float) -> float | None: + return recorded if recorded is not None or not observed else elapsed + + +def _observe(phases: BootPhases, probe: _Probe) -> BootPhases: + prisma_started: Final = _first(phases.prisma_started, probe.prisma, probe.elapsed) + prisma_finished: Final = _first(phases.prisma_finished, prisma_started is not None, probe.elapsed) + return BootPhases( + prisma_started=prisma_started, + prisma_finished=None if probe.prisma else prisma_finished, + listening=_first(phases.listening, probe.listening, probe.elapsed), + first_ready=_first(phases.first_ready, probe.ready_pid is not None, probe.elapsed), + ready_pids=phases.ready_pids | frozenset(() if probe.ready_pid is None else (probe.ready_pid,)), + ) def _ready_worker_pid(member: psutil.Process) -> int | None: @@ -245,13 +549,24 @@ def _ready_worker_pid(member: psutil.Process) -> int | None: return None -def _wait_until_ready(launch: _Launch, workers: int = 1) -> None: - for ready in accumulate(_readiness_pids(launch), _ready_workers, initial=frozenset[int]()): - members: Final = frozenset( - pid for member in group_members(launch.process.pid) if (pid := _ready_worker_pid(member)) is not None - ) - if len(ready & members) >= max(1, workers): - return +def _serves_all(launch: _Launch, phases: BootPhases, workers: int) -> bool: + if len(phases.ready_pids) < max(1, workers): + return False + members: Final = frozenset( + pid for member in group_members(launch.process.pid) if (pid := _ready_worker_pid(member)) is not None + ) + return len(phases.ready_pids & members) >= max(1, workers) + + +def _wait_until_ready(launch: _Launch, workers: int = 1) -> BootPhases | None: + return next( + ( + phases + for phases in accumulate(_probes(launch, workers), _observe, initial=BootPhases()) + if _serves_all(launch, phases, workers) + ), + None, + ) def _launch_until_bound( @@ -259,7 +574,7 @@ def _launch_until_bound( ) -> _Launch: launch: Final = _launch(command, root, environment, output) try: - _wait_until_ready(launch, workers) + phases: Final = _wait_until_ready(launch, workers) exit_code: Final = launch.process.poll() assert exit_code is None or (attempts > 1 and _lost_port_race(exit_code, launch.log)), ( "Owned proxy exited before readiness" @@ -268,12 +583,21 @@ def _launch_until_bound( _stop_launch(launch) raise if exit_code is None: + assert phases is not None launch.log.with_suffix(".startup.json").write_text( json.dumps( { "root_pid": launch.process.pid, "readiness_seconds": time.monotonic() - launch.started, "logical_cpus": os.cpu_count(), + "workers": workers, + "spawned_at": _wall_clock(launch.started), + "phases": { + "prisma_started": phases.prisma_started, + "prisma_finished": phases.prisma_finished, + "listening": phases.listening, + "first_ready": phases.first_ready, + }, } ) + "\n" @@ -287,14 +611,14 @@ def _proxy_root() -> Path: return Path(os.environ.get("INTEGRATION_PROXY_ROOT") or Path(__file__).resolve().parents[3]) -def worker_configuration(config: Path, directory: Path) -> Path: - if os.environ.get("INTEGRATION_WORKER_ISOLATION") != "1": +def worker_configuration(config: Path, directory: Path, inherited: Mapping[str, str] = os.environ) -> Path: + if inherited.get("INTEGRATION_WORKER_ISOLATION") != "1": return config def endpoint(match: re.Match[str]) -> str: if match.group(1) == "8190": - return os.environ["INTEGRATION_UPSTREAM_URL"] - return f"http://127.0.0.1:{os.environ['INTEGRATION_PROVIDER_PORT']}" + return inherited["INTEGRATION_UPSTREAM_URL"] + return f"http://127.0.0.1:{inherited['INTEGRATION_PROVIDER_PORT']}" contents: Final = re.sub(r"http://127[.]0[.]0[.]1:(8190|8191)(?![0-9])", endpoint, config.read_text()) directory.mkdir(parents=True, exist_ok=True) @@ -317,9 +641,12 @@ def worker_configuration(config: Path, directory: Path) -> Path: def _proxy_environment( - gateway: Gateway, overrides: Mapping[str, str], remove_environment: tuple[str, ...] + gateway: Gateway, + overrides: Mapping[str, str], + remove_environment: tuple[str, ...], + source: Mapping[str, str] = os.environ, ) -> Mapping[str, str]: - inherited: Final = {**os.environ, **proxy_database_environment()} + inherited: Final = {**source, **proxy_database_environment(source)} writer: Final = overrides.get("DATABASE_URL") reader: Final = inherited.get("DATABASE_URL_READ_REPLICA") unpaired_reader: Final = ( @@ -337,7 +664,7 @@ def _proxy_environment( { **{name: value for name, value in inherited.items() if name not in removed}, "LITELLM_MASTER_KEY": gateway.key, - "LITELLM_SALT_KEY": os.environ.get("LITELLM_SALT_KEY", "sk-integration-salt"), + "LITELLM_SALT_KEY": source.get("LITELLM_SALT_KEY", "sk-integration-salt"), "STORE_MODEL_IN_DB": "True", "PYTHON_DOTENV_DISABLED": "1", **overrides, @@ -355,7 +682,12 @@ def _proxy_environment( def setup_only_proxy_run( - gateway: Gateway, overrides: Mapping[str, str], *, config: Path, workers: int + gateway: Gateway, + overrides: Mapping[str, str], + *, + config: Path, + workers: int, + database_setup: tuple[str, ...] = DB_PUSH, ) -> subprocess.CompletedProcess[str]: """The proxy CLI's `--skip_server_startup` pass (the image's setup step), run to completion with the environment an owned proxy gets.""" @@ -368,7 +700,7 @@ def setup_only_proxy_run( str(config), "--num_workers", str(workers), - *DB_PUSH, + *database_setup, "--skip_server_startup", ), cwd=_proxy_root(), @@ -391,23 +723,33 @@ def owned_proxy_process( workers: int = 1, database_setup: tuple[str, ...] = DB_PUSH, extra_arguments: tuple[str, ...] = (), + coverage_file: Path | None = None, + healthcheck: WorkerHealthcheck = WorkerHealthcheck.STRICT, + inherited: Mapping[str, str] = os.environ, ) -> Iterator[OwnedProxy]: + assert "--timeout_worker_healthcheck" not in extra_arguments, "Choose the worker healthcheck with `healthcheck`" root: Final = _proxy_root() - environment: Final = _proxy_environment(gateway, overrides, remove_environment) - output: Final = Path(os.environ.get("INTEGRATION_RESULTS_DIR", str(directory))) + environment: Final = _proxy_environment( + gateway, + {**overrides, **({"COVERAGE_FILE": str(coverage_file)} if coverage_file else {})}, + remove_environment, + inherited, + ) + output: Final = Path(inherited.get("INTEGRATION_RESULTS_DIR", str(directory))) output.mkdir(parents=True, exist_ok=True) command: Final = ( sys.executable, "-m", + *(("coverage", "run", f"--rcfile={MCP_COVERAGE_CONFIG}", "-m") if coverage_file else ()), "integration._support.proxy", "--config", - str(worker_configuration(config or root / "tests/integration/proxy_config.yaml", directory)), + str(worker_configuration(config or root / "tests/integration/proxy_config.yaml", directory, inherited)), "--host", "127.0.0.1", "--num_workers", str(workers), "--timeout_worker_healthcheck", - "5", + str(healthcheck_seconds(healthcheck)), *database_setup, *extra_arguments, ) @@ -417,7 +759,7 @@ def owned_proxy_process( with httpx.Client( base_url=f"http://127.0.0.1:{launch.port}", timeout=15, trust_env=False, limits=GATEWAY_LIMITS ) as client: - yield OwnedProxy(Gateway(client, gateway.key, gateway.upstream_url), process, launch.log) + yield OwnedProxy(Gateway(client, gateway.key, gateway.upstream_url), process, launch.log, healthcheck) finally: _stop_launch(launch) @@ -489,7 +831,7 @@ def refused_boot_log( launch: Final = _launch(command, root, environment, output) try: with httpx.Client(base_url=f"http://127.0.0.1:{launch.port}", timeout=15, trust_env=False) as client: - deadline: Final = time.monotonic() + 70 + deadline: Final = time.monotonic() + readiness_seconds(1) while launch.process.poll() is None: assert not _is_ready(client), ( f"Proxy became ready instead of refusing to boot:\n{launch.log.read_text()}" @@ -545,13 +887,21 @@ class UpstreamSlot: It runs from the harness's own checkout, never ``INTEGRATION_PROXY_ROOT``: the double belongs to the tests, the proxy root only names the litellm under test.""" - __slots__ = ("certificate", "directory", "port", "process", "root") + __slots__ = ("certificate", "directory", "inherited", "port", "process", "root") - def __init__(self, directory: Path, port: int, root: Path, certificate: UpstreamCertificate | None = None) -> None: + def __init__( + self, + directory: Path, + port: int, + root: Path, + certificate: UpstreamCertificate | None = None, + inherited: Mapping[str, str] = os.environ, + ) -> None: self.directory = directory self.port = port self.root = root self.certificate = certificate + self.inherited = inherited self.process: subprocess.Popen[bytes] | None = None @property @@ -566,7 +916,7 @@ class UpstreamSlot: def start(self) -> None: assert self.process is None, "Owned upstream is already running" - output: Final = Path(os.environ.get("INTEGRATION_RESULTS_DIR") or self.directory) + output: Final = Path(self.inherited.get("INTEGRATION_RESULTS_DIR") or self.directory) log_path: Final = output / f"owned-upstream-{self.port}-{uuid.uuid4().hex}.log" with log_path.open("w") as log: process: Final = subprocess.Popen( @@ -579,7 +929,7 @@ class UpstreamSlot: *self._tls_arguments(), ], cwd=self.root, - env=dict(os.environ), + env=dict(self.inherited), stdout=log, stderr=subprocess.STDOUT, start_new_session=True, @@ -615,9 +965,11 @@ def _harness_root() -> Path: @contextmanager -def owned_upstream(directory: Path, *, tls: bool = False) -> Generator[UpstreamSlot]: +def owned_upstream( + directory: Path, *, tls: bool = False, port: int | None = None, inherited: Mapping[str, str] = os.environ +) -> Generator[UpstreamSlot]: certificate: Final = self_signed_certificate(directory) if tls else None - slot: Final = UpstreamSlot(directory, _free_port(), _harness_root(), certificate) + slot: Final = UpstreamSlot(directory, port or _free_port(), _harness_root(), certificate, inherited) try: slot.start() yield slot diff --git a/tests/integration/_support/redis_process.py b/tests/integration/_support/redis_process.py index 7966d105302..f5e7f509eb5 100644 --- a/tests/integration/_support/redis_process.py +++ b/tests/integration/_support/redis_process.py @@ -5,7 +5,7 @@ import socket import subprocess import time import uuid -from collections.abc import Iterator +from collections.abc import Callable, Iterator, Mapping from contextlib import contextmanager from dataclasses import dataclass from pathlib import Path @@ -14,6 +14,19 @@ from typing import Final, TextIO from redis import Redis from redis.exceptions import ConnectionError as RedisConnectionError +PORT_ATTEMPTS: Final = 5 + + +@dataclass(frozen=True, slots=True) +class PortUnavailable: + reason: str + + +def free_loopback_port() -> int: + with socket.socket() as reservation: + reservation.bind(("127.0.0.1", 0)) + return reservation.getsockname()[1] + @dataclass class OwnedRedis: @@ -32,21 +45,54 @@ class OwnedRedis: except RedisConnectionError: return False + def owned_pid(self) -> int | None: + assert self.process is not None + if self.command[0] != "docker": + return self.process.pid + result: Final = subprocess.run( + ["docker", "exec", "redis-cache", "cat", self.pid_file], + capture_output=True, + text=True, + timeout=2, + ) + if result.returncode != 0: + return None + return int(result.stdout) + def start(self) -> None: + unavailable: Final = self.try_start() + assert unavailable is None, unavailable.reason + + def try_start(self) -> PortUnavailable | None: assert self.process is None - assert not self.serves(), f"Another Redis already serves {self.host}:{self.port} before the owned one starts" - self.process = subprocess.Popen(self.command, stdout=self.log, stderr=subprocess.STDOUT, start_new_session=True) + if self.serves(): + return PortUnavailable(f"Another Redis already serves {self.host}:{self.port} before the owned one starts") + process: Final = subprocess.Popen( + self.command, stdout=self.log, stderr=subprocess.STDOUT, start_new_session=True + ) + self.process = process deadline: Final = time.monotonic() + 8 with Redis(host=self.host, port=self.port, socket_connect_timeout=0.2, socket_timeout=0.2) as client: while True: - assert self.process.poll() is None, "Owned Redis exited before readiness" + code: Final = process.poll() + if code is not None: + self.process = None + return PortUnavailable( + f"Owned Redis exited with code {code} before readiness on {self.host}:{self.port}" + ) try: if client.ping(): actual: Final = int(client.info("server")["process_id"]) - expected: Final = self.process.pid if self.command[0] != "docker" else int(subprocess.check_output(["docker", "exec", "redis-cache", "cat", self.pid_file], timeout=2)) - assert actual == expected, "Redis readiness reached a different process" - self.server_pid = actual - return + expected: Final = self.owned_pid() + if actual == expected: + self.server_pid = actual + return None + try: + process.wait(timeout=max(deadline - time.monotonic(), 0)) + except subprocess.TimeoutExpired as error: + raise AssertionError("Redis readiness reached a different process") from error + self.process = None + return PortUnavailable(f"A different Redis answered on {self.host}:{self.port} during startup") except RedisConnectionError: pass assert time.monotonic() < deadline, "Owned Redis readiness deadline exceeded" @@ -59,7 +105,9 @@ class OwnedRedis: try: if self.process.poll() is None: with Redis(host=self.host, port=self.port, socket_connect_timeout=1, socket_timeout=1) as client: - assert int(client.info("server")["process_id"]) == self.server_pid, "Redis ownership changed before shutdown" + assert int(client.info("server")["process_id"]) == self.server_pid, ( + "Redis ownership changed before shutdown" + ) client.shutdown(nosave=True) except Exception as error: failure = error @@ -85,32 +133,96 @@ class OwnedRedis: self.process.send_signal(action) return pid: Final = int(subprocess.check_output(["docker", "exec", "redis-cache", "cat", self.pid_file], timeout=2)) - command: Final = subprocess.check_output(["docker", "exec", "redis-cache", "cat", f"/proc/{pid}/cmdline"], timeout=2) + command: Final = subprocess.check_output( + ["docker", "exec", "redis-cache", "cat", f"/proc/{pid}/cmdline"], timeout=2 + ) assert self.pid_file.encode() in command, "Redis process ownership changed" subprocess.run(["docker", "exec", "redis-cache", "kill", f"-{int(action)}", str(pid)], check=True, timeout=2) + @contextmanager + def paused(self) -> Iterator[None]: + self.signal(signal.SIGSTOP) + try: + yield + finally: + self.signal(signal.SIGCONT) + @contextmanager -def owned_redis(directory: Path) -> Iterator[OwnedRedis]: - binary: Final = shutil.which("redis-server") - with socket.socket() as reservation: - reservation.bind(("127.0.0.1", 0)) - port: Final = reservation.getsockname()[1] - if binary: - host = "127.0.0.1" - prefix = (binary,) - else: - host = subprocess.check_output(["docker", "inspect", "--format", "{{range .NetworkSettings.Networks}}{{.IPAddress}}{{end}}", "redis-cache"], text=True).strip() - assert host, "CircleCI owned Redis container has no address" - prefix = ("docker", "exec", "redis-cache", "redis-server") - output: Final = Path(os.environ.get("INTEGRATION_RESULTS_DIR", str(directory))) - output.mkdir(parents=True, exist_ok=True) - with (output / f"owned-redis-{uuid.uuid4().hex}.log").open("w") as log: - pid_file: Final = str(directory / "owned-redis.pid") if binary else f"/tmp/integration-redis-{uuid.uuid4().hex}.pid" - server: Final = OwnedRedis(host, port, (*prefix, "--port", str(port), "--set-proc-title", "no", "--pidfile", pid_file, "--bind", "0.0.0.0" if not binary else "127.0.0.1", "--protected-mode", "no", "--save", "", "--appendonly", "no"), log, pid_file) - try: - server.start() +def _started( + build: Callable[[int], OwnedRedis], + choose_port: Callable[[], int], + attempts: int, + log_name: str, +) -> Iterator[OwnedRedis]: + server: Final = build(choose_port()) + try: + unavailable: Final = server.try_start() + if unavailable is None: + yield server + return + assert attempts > 1, ( + f"Owned Redis found no usable port after {PORT_ATTEMPTS} attempts: {unavailable.reason} (log {log_name})" + ) + with _started(build, choose_port, attempts - 1, log_name) as started: + yield started + finally: + if server.process is not None: + server.stop() + + +@contextmanager +def owned_redis( + directory: Path, + choose_port: Callable[[], int] = free_loopback_port, + inherited: Mapping[str, str] = os.environ, +) -> Iterator[OwnedRedis]: + binary: Final = shutil.which("redis-server") + if binary: + host: Final = "127.0.0.1" + prefix: Final = (binary,) + else: + host: Final = subprocess.check_output( + [ + "docker", + "inspect", + "--format", + "{{range .NetworkSettings.Networks}}{{.IPAddress}}{{end}}", + "redis-cache", + ], + text=True, + ).strip() + assert host, "CircleCI owned Redis container has no address" + prefix: Final = ("docker", "exec", "redis-cache", "redis-server") + output: Final = Path(inherited.get("INTEGRATION_RESULTS_DIR", str(directory))) + output.mkdir(parents=True, exist_ok=True) + directory.mkdir(parents=True, exist_ok=True) + with (output / f"owned-redis-{uuid.uuid4().hex}.log").open("w") as log: + + def build(port: int) -> OwnedRedis: + pid_file: Final = ( + str(directory / f"owned-redis-{port}.pid") + if binary + else f"/tmp/integration-redis-{uuid.uuid4().hex}.pid" + ) + command: Final = ( + *prefix, + "--port", + str(port), + "--set-proc-title", + "no", + "--pidfile", + pid_file, + "--bind", + "0.0.0.0" if not binary else "127.0.0.1", + "--protected-mode", + "no", + "--save", + "", + "--appendonly", + "no", + ) + return OwnedRedis(host, port, command, log, pid_file) + + with _started(build, choose_port, PORT_ATTEMPTS, log.name) as server: yield server - finally: - if server.process is not None: - server.stop() diff --git a/tests/integration/_support/routing.py b/tests/integration/_support/routing.py index 14f4741367a..1f8f5626031 100644 --- a/tests/integration/_support/routing.py +++ b/tests/integration/_support/routing.py @@ -21,6 +21,7 @@ WRITER_ROLE: Final = "litellm_writer" READER_ROLE: Final = "litellm_reader" ROLES: Final = (READER_ROLE, WRITER_ROLE) DATABASE_NAME: Final = "circle_test" +WORKER_DATABASE_PREFIX: Final = "integration_worker" OBSERVED_FILE: Final = "routing-observed.json" DIFF_FILE: Final = "routing-diff.txt" EITHER_ROLE_FILE: Final = Path(__file__).resolve().parents[1] / "routing" / "either_role.json" @@ -216,10 +217,10 @@ def snapshot(connection: psycopg.Connection[object]) -> Mapping[tuple[str, str], SELECT r.rolname, s.query, s.calls FROM pg_stat_statements s JOIN pg_roles r ON r.oid = s.userid - WHERE s.dbid = (SELECT oid FROM pg_database WHERE datname = %s) + WHERE s.dbid IN (SELECT oid FROM pg_database WHERE datname = %s OR starts_with(datname, %s)) AND r.rolname = ANY(%s) """, - (DATABASE_NAME, list(ROLES)), + (DATABASE_NAME, f"{WORKER_DATABASE_PREFIX}_", list(ROLES)), ).fetchall() return MappingProxyType( { diff --git a/tests/integration/_support/sagemaker_embedding.py b/tests/integration/_support/sagemaker_embedding.py new file mode 100644 index 00000000000..67dda14f664 --- /dev/null +++ b/tests/integration/_support/sagemaker_embedding.py @@ -0,0 +1,210 @@ +import base64 +import json +import struct +from collections.abc import Mapping +from typing import Final +from urllib.parse import unquote + +from integration._support.wire import Reply, Request +from pydantic import JsonValue, TypeAdapter + +ACCESS_KEY: Final = "AKIAINTEGRATION000004" +SECRET_KEY: Final = "synthetic-secret-key-for-testing" +REGION: Final = "us-east-1" +COMPONENT: Final = "integration-embed-component" +SERVED_MODEL: Final = "integration-org/served-embedding-model" +VOYAGE_SERVED_MODEL: Final = "voyage-served-model" +ENDPOINT_PREFIX: Final = "integration-embed-" +COMPONENT_HEADER: Final = "x-amzn-sagemaker-inference-component" +JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +NOT_JSON_BODY: Final = b"upstream gateway error" +EXPECTED_SHAPE: Final = "Expected an OpenAI embeddings object with a 'data' list of {'embedding': [...]} rows" +NON_STRING_MODEL_ID: Final = "model_id must be the inference component name as a string" +IC_REQUIRED: Final = ( + "Inference Component Name header is required for endpoints to which you plan to deploy inference components. " + "Please include Inference Component Name header or consider using SageMaker models." +) +IC_NOT_ALLOWED: Final = ( + "Inference Component Name header is not allowed for endpoints to which you dont plan to deploy inference " + "components. Please remove the Inference Component Name header and try again." +) +PASS_IN_HINT: Final = "pass in via `litellm.embedding(..., model_id={InferenceComponentName})`" +REMOVE_HINT: Final = "remove `model_id` from this deployment, the endpoint has no inference components" + + +def endpoint(scenario: str) -> str: + return f"{ENDPOINT_PREFIX}{scenario}" + + +def openai_route(scenario: str) -> str: + return f"sagemaker/openai/{endpoint(scenario)}" + + +def plain_route(scenario: str) -> str: + return f"sagemaker/{endpoint(scenario)}" + + +def invocations_target(name: str) -> str: + return f"/endpoints/{name}/invocations" + + +def vector(index: int) -> list[float]: + return [0.25, -0.5, 1.0 + index] + + +def inputs(body: Mapping[str, JsonValue]) -> tuple[JsonValue, ...]: + carried: Final = next((body[key] for key in ("input", "inputs", "texts") if key in body), []) + return tuple(carried) if isinstance(carried, list) else (carried,) + + +def embedding_value(index: int, encoding_format: JsonValue) -> JsonValue: + floats: Final = vector(index) + if encoding_format == "base64": + return base64.b64encode(struct.pack(f"<{len(floats)}f", *floats)).decode() + return floats + + +def openai_rows(count: int, encoding_format: JsonValue, indexed: bool) -> list[JsonValue]: + return [ + { + "object": "embedding", + **({"index": index} if indexed else {}), + "embedding": embedding_value(index, encoding_format), + } + for index in range(count) + ] + + +def _openai_usage(scenario: str) -> dict[str, JsonValue]: + if scenario == "nousage": + return {} + if scenario == "promptonly": + return {"usage": {"prompt_tokens": 5}} + return {"usage": {"prompt_tokens": 7, "total_tokens": 7}} + + +def openai_payload(body: Mapping[str, JsonValue], scenario: str = "openai") -> dict[str, JsonValue]: + return { + "object": "list", + "data": openai_rows(len(inputs(body)), body.get("encoding_format"), indexed=scenario != "noindex"), + **({} if scenario == "nomodel" else {"model": SERVED_MODEL}), + **_openai_usage(scenario), + } + + +def json_reply(payload: JsonValue) -> Reply: + return Reply(body=json.dumps(payload).encode()) + + +def aws_error(status: int, code: str, message: str, **fields: JsonValue) -> Reply: + return Reply( + status=status, + body=json.dumps({"__type": code, "Message": message, **fields}).encode(), + headers={"x-amzn-ErrorType": code}, + ) + + +def endpoint_name_of(request: Request) -> str: + parts: Final = request.target.split("/") + return unquote(parts[2]) if len(parts) > 2 else "" + + +def scenario_of(name: str) -> str: + if not name.startswith(ENDPOINT_PREFIX): + return "" + return name.removeprefix(ENDPOINT_PREFIX).split("-")[0] + + +def respond(request: Request) -> Reply: + name: Final = endpoint_name_of(request) + body: Final = JSON_OBJECT.validate_json(request.body) + match scenario_of(name): + case "openai" | "nousage" | "nomodel" | "noindex" | "promptonly": + return json_reply(openai_payload(body, scenario_of(name))) + case "hf": + return json_reply([vector(index) for index in range(len(inputs(body)))]) + case "icrequired": + if COMPONENT_HEADER not in request.headers: + return aws_error(400, "ValidationError", IC_REQUIRED) + return json_reply(openai_payload(body)) + case "noic": + if COMPONENT_HEADER in request.headers: + return aws_error(400, "ValidationError", IC_NOT_ALLOWED) + return json_reply([vector(index) for index in range(len(inputs(body)))]) + case "hfdict": + return json_reply({"embedding": [vector(index) for index in range(len(inputs(body)))]}) + case "voyage": + return json_reply( + { + "object": "list", + "data": openai_rows(len(inputs(body)), None, indexed=True), + "model": VOYAGE_SERVED_MODEL, + "usage": {"total_tokens": 9}, + } + ) + case "cohere": + return json_reply( + { + "id": "integration-cohere-response", + "texts": list(inputs(body)), + "embeddings": [vector(index) for index in range(len(inputs(body)))], + "meta": {"api_version": {"version": "1"}, "billed_units": {"input_tokens": 6}}, + "response_type": "embeddings_floats", + } + ) + case "datadict": + return json_reply({"object": "list", "data": {"embedding": vector(0)}}) + case "hfshape": + return json_reply([vector(0)]) + case "notjson": + return Reply(body=NOT_JSON_BODY, content_type="text/html") + case "err400": + return aws_error( + 400, "ValidationError", "1 validation error detected: Value at 'body' failed to satisfy constraint" + ) + case "err424": + return aws_error( + 424, + "ModelError", + 'Received client error (400) from model with message "{\\"error\\":\\"unexpected field inputs\\"}".', + OriginalStatusCode=400, + OriginalMessage='{"error":"unexpected field inputs"}', + LogStreamArn="arn:aws:logs:us-east-1:000000000000:log-group:/aws/sagemaker/Endpoints/integration", + ) + case "err429": + return aws_error(429, "ModelNotReadyException", "Model is not ready for inference yet") + case "err503": + return aws_error(503, "ServiceUnavailable", "The endpoint is scaling, try again later") + case "err500": + return aws_error(500, "InternalFailure", "An internal failure occurred") + case "drop": + return Reply(drop_connection=True) + case _: + return aws_error(400, "ValidationError", f"Endpoint {name} of account 000000000000 not found.") + + +def litellm_params(model: str, **parameters: JsonValue) -> dict[str, JsonValue]: + return { + "model": model, + "aws_access_key_id": ACCESS_KEY, + "aws_secret_access_key": SECRET_KEY, + "aws_region_name": REGION, + **parameters, + } + + +def deployment( + name: str, model: str, *, model_info: Mapping[str, JsonValue] | None = None, **parameters: JsonValue +) -> dict[str, JsonValue]: + return { + "model_name": name, + "litellm_params": litellm_params(model, **parameters), + **({} if model_info is None else {"model_info": dict(model_info)}), + } + + +def openai_deployment( + name: str, scenario: str, *, model_info: Mapping[str, JsonValue] | None = None, **parameters: JsonValue +) -> dict[str, JsonValue]: + defaults: Final[dict[str, JsonValue]] = {"model_id": COMPONENT, "extra_body": {"model": SERVED_MODEL}} + return deployment(name, openai_route(scenario), model_info=model_info, **{**defaults, **parameters}) diff --git a/tests/integration/_support/timing.py b/tests/integration/_support/timing.py new file mode 100644 index 00000000000..5b16c9240cf --- /dev/null +++ b/tests/integration/_support/timing.py @@ -0,0 +1,40 @@ +import time +from collections.abc import Callable, Iterable +from dataclasses import dataclass +from typing import Final + + +@dataclass(frozen=True, slots=True) +class Timed[T]: + value: T + seconds: float + + +def timed[T](operation: Callable[[], T]) -> Timed[T]: + started: Final = time.monotonic() + value: Final = operation() + return Timed(value, time.monotonic() - started) + + +def assert_not_delayed_by(observed: Iterable[float], *, control: Iterable[float], injected_delay: float) -> None: + worst: Final = max(observed) + baseline: Final = max(control) + assert worst < baseline + injected_delay, ( + f"slowest call took {worst:.2f}s, at least the {baseline:.2f}s control plus the {injected_delay}s injected delay" + ) + + +def assert_cut_short(elapsed: float, *, injected_delay: float) -> None: + assert elapsed < injected_delay, ( + f"returned after {elapsed:.2f}s, no sooner than the {injected_delay}s injected delay" + ) + + +def assert_enforced_within(elapsed: float, *, configured: float, slack: float, not_before: float = 0.0) -> None: + assert not_before <= elapsed <= configured + slack, ( + f"took {elapsed:.2f}s, expected between {not_before}s and the configured {configured}s plus {slack}s of slack" + ) + + +def assert_within_factor_of_control(observed: float, *, control: float, factor: float) -> None: + assert observed <= factor * control, f"took {observed:.2f}s, over {factor}x the {control:.2f}s control" diff --git a/tests/integration/_support/upstream.py b/tests/integration/_support/upstream.py index 6b36d56a757..89a28cffb1d 100644 --- a/tests/integration/_support/upstream.py +++ b/tests/integration/_support/upstream.py @@ -496,6 +496,7 @@ class Provider: .encode(), media_type=response.content_type, status_code=response.status, + headers=response.headers, ) case BinaryResponse(): return Response( @@ -518,13 +519,13 @@ class Provider: ).encode() await asyncio.sleep(response.frame_delay_ms / 1000) - return StreamingResponse(stream(), media_type=response.content_type) + return StreamingResponse(stream(), media_type=response.content_type, headers=response.headers) stream_body: Final = ( ("\n\n".join(response.frames) + "\n\n") .replace("$REQUEST_ID", scenario_id) .replace("$UNIQUE_ID", unique_id) ) - return Response(content=stream_body.encode(), media_type=response.content_type) + return Response(content=stream_body.encode(), media_type=response.content_type, headers=response.headers) case EventStreamResponse(): events: Final = ( tuple( @@ -753,6 +754,7 @@ def main() -> None: port=cast(int, arguments.port), access_log=False, timeout_keep_alive=125, + timeout_graceful_shutdown=5, ssl_certfile=cast(str | None, arguments.ssl_certfile), ssl_keyfile=cast(str | None, arguments.ssl_keyfile), ) diff --git a/tests/integration/_support/workers.py b/tests/integration/_support/workers.py index 2502e62477e..33e086576c5 100644 --- a/tests/integration/_support/workers.py +++ b/tests/integration/_support/workers.py @@ -4,8 +4,10 @@ import os import socket import time import uuid -from collections.abc import Iterator, Mapping -from contextlib import contextmanager, nullcontext +from collections.abc import Callable, Iterator, Mapping +from concurrent.futures import ThreadPoolExecutor +from contextlib import ExitStack, contextmanager, nullcontext +from dataclasses import dataclass from pathlib import Path from types import MappingProxyType from typing import Final, Protocol @@ -22,6 +24,7 @@ from redis.exceptions import TimeoutError as RedisTimeoutError from tests.integration._support.client import Gateway from tests.integration._support.database import scratch_database from tests.integration._support.redis_process import owned_redis +from tests.integration._support.routing import WORKER_DATABASE_PREFIX READINESS_OBJECT: Final = TypeAdapter(dict[str, object]) SUBSCRIBER_COUNTS: Final = TypeAdapter(tuple[tuple[str, int], ...]) @@ -116,12 +119,46 @@ def _replica_urls(database: str) -> Mapping[str, str]: ) +@dataclass(frozen=True, slots=True) +class GatewayProfile: + peer: bool + replica: bool + cost_map: bool + coverage_file: Path | None + + +def gateway_profile(environment: Mapping[str, str]) -> GatewayProfile: + coverage: Final = environment.get("INTEGRATION_COVERAGE_FILE", "") + return GatewayProfile( + peer=environment.get("INTEGRATION_WORKER_PEER") == "1", + replica=bool(environment.get("INTEGRATION_PROXY_READ_REPLICA_URL")), + cost_map=environment.get("INTEGRATION_WORKER_COST_MAP") == "1", + coverage_file=Path(coverage) if coverage else None, + ) + + +def cost_map_environment(upstream: str) -> Mapping[str, str]: + return MappingProxyType( + { + "LITELLM_MODEL_COST_MAP_URL": f"{upstream}/_cost_map", + "MODEL_COST_MAP_MIN_MODEL_COUNT": "1", + "MODEL_COST_MAP_MAX_SHRINK_RATIO": "0", + "GEMINI_API_BASE": upstream, + "ANTHROPIC_API_BASE": upstream, + "GEMINI_API_KEY": "sk-scripted-provider", + "ANTHROPIC_API_KEY": "sk-scripted-provider", + } + ) + + @contextmanager -def worker_services(directory: Path, *, replica: bool = False) -> Iterator[Mapping[str, str]]: +def worker_services( + directory: Path, *, replica: bool = False, inherited: Mapping[str, str] = os.environ +) -> Iterator[Mapping[str, str]]: directory.mkdir(parents=True, exist_ok=True) with ( - scratch_database() as database, - owned_redis(directory) as cache, + scratch_database(WORKER_DATABASE_PREFIX) as database, + owned_redis(directory, inherited=inherited) as cache, ): yield MappingProxyType( { @@ -136,25 +173,47 @@ def worker_services(directory: Path, *, replica: bool = False) -> Iterator[Mappi @contextmanager -def worker_gateway(directory: Path, *, peer: bool, replica: bool) -> Iterator[Mapping[str, str]]: - from tests.integration._support.asgi import asgi_server - from tests.integration._support.process import owned_proxy - from tests.integration._support.upstream import Provider +def worker_gateway( + directory: Path, profile: GatewayProfile, inherited: Mapping[str, str] +) -> Iterator[Mapping[str, str]]: + from tests.integration._support.process import owned_proxy, owned_upstream - port: Final = urlsplit(os.environ["INTEGRATION_UPSTREAM_URL"]).port + upstream: Final = inherited["INTEGRATION_UPSTREAM_URL"] + port: Final = urlsplit(upstream).port assert port is not None key: Final = f"sk-{uuid.uuid4().hex}" - with worker_services(directory, replica=replica) as services: + removed: Final = ("LITELLM_LOCAL_MODEL_COST_MAP",) if profile.cost_map else () + with worker_services(directory, replica=profile.replica, inherited=inherited) as services: overrides: Final = { **services, "DATABASE_URL": services["INTEGRATION_PROXY_DATABASE_URL"], - **({"DATABASE_URL_READ_REPLICA": services["INTEGRATION_PROXY_READ_REPLICA_URL"]} if replica else {}), + **( + {"DATABASE_URL_READ_REPLICA": services["INTEGRATION_PROXY_READ_REPLICA_URL"]} if profile.replica else {} + ), + "LITELLM_ENABLE_MCP_STDIO": "true", + **(cost_map_environment(upstream) if profile.cost_map else {}), } with ( - asgi_server(Provider().app(), port=port, keepalive_seconds=125) as upstream, + owned_upstream(directory, port=port, inherited=inherited), httpx.Client(base_url="http://127.0.0.1:1", trust_env=False) as client, - owned_proxy(Gateway(client, key, upstream), directory, overrides) as primary, - owned_proxy(primary, directory, overrides) if peer else nullcontext(None) as secondary, + owned_proxy( + Gateway(client, key, upstream), + directory, + overrides, + remove_environment=removed, + coverage_file=profile.coverage_file, + inherited=inherited, + ) as primary, + owned_proxy( + primary, + directory, + overrides, + remove_environment=removed, + coverage_file=profile.coverage_file, + inherited=inherited, + ) + if profile.peer + else nullcontext(None) as secondary, ): with Redis(host=services["REDIS_HOST"], port=int(services["REDIS_PORT"]), socket_timeout=2) as cache: wait_for_services(upstream, (primary, secondary) if secondary is not None else (primary,), cache) @@ -169,3 +228,36 @@ def worker_gateway(directory: Path, *, peer: bool, replica: bool) -> Iterator[Ma "INTEGRATION_UPSTREAM_URL": upstream, } ) + + +@dataclass(frozen=True, slots=True) +class WorkerStack: + resources: ExitStack + environment: Mapping[str, str] + + +def worker_stack(results: Path, profile: GatewayProfile) -> WorkerStack: + endpoints: Final = worker_endpoints() + inherited: Final = MappingProxyType({**os.environ, **endpoints, "INTEGRATION_RESULTS_DIR": str(results)}) + resources: Final = ExitStack() + try: + services: Final = resources.enter_context(worker_gateway(results / "runtime", profile, inherited)) + except BaseException: + resources.close() + raise + return WorkerStack(resources, MappingProxyType({**endpoints, **services, "INTEGRATION_RESULTS_DIR": str(results)})) + + +def start_worker_stacks( + directories: tuple[Path, ...], build: Callable[[Path], WorkerStack], resources: ExitStack, concurrency: int +) -> tuple[Mapping[str, str], ...]: + with ThreadPoolExecutor(max_workers=max(1, concurrency), thread_name_prefix="integration-worker-stack") as pool: + futures: Final = tuple(pool.submit(build, directory) for directory in directories) + for future in futures: + if future.exception() is None: + resources.enter_context(future.result().resources) + failure: Final = next((error for future in futures if (error := future.exception()) is not None), None) + if failure is not None: + resources.close() + raise failure + return tuple(future.result().environment for future in futures) diff --git a/tests/integration/authorization/test_alias_chain_budget_waiver.py b/tests/integration/authorization/test_alias_chain_budget_waiver.py index 175d4765a24..0d64415d445 100644 --- a/tests/integration/authorization/test_alias_chain_budget_waiver.py +++ b/tests/integration/authorization/test_alias_chain_budget_waiver.py @@ -2,7 +2,6 @@ import json import math import os import re -import signal import threading import uuid from collections.abc import Callable, Generator, Iterator, Mapping @@ -964,11 +963,11 @@ def test_killing_one_worker_mid_burst_keeps_the_chain_entry_served(rig: ChainRig held_by: Final = MappingProxyType({pid: _open_upstream_connections(pid, wire.url) for pid in workers}) assert sum(held_by.values()) == GATEWAY_BURST, held_by victim_pid, survivor_pid = sorted(workers, key=held_by.__getitem__) - victim: Final = psutil.Process(victim_pid) - victim.suspend() - victim.send_signal(signal.SIGKILL) - release.set() - responses: Final = {marker: future.result() for marker, future in pending.items()} + with owned.replacing_worker(victim_pid) as kill: + kill.freeze() + kill() + release.set() + responses: Final = {marker: future.result() for marker, future in pending.items()} served: Final = frozenset( marker for marker, response in responses.items() if response is not None and response.status_code == 200 ) diff --git a/tests/integration/authorization/test_cli_sso_login_ui_disabled.py b/tests/integration/authorization/test_cli_sso_login_ui_disabled.py index a5e9b2752ea..315a0b20334 100644 --- a/tests/integration/authorization/test_cli_sso_login_ui_disabled.py +++ b/tests/integration/authorization/test_cli_sso_login_ui_disabled.py @@ -2,10 +2,8 @@ from __future__ import annotations import base64 import json -import os import re import secrets -import signal import threading import uuid from collections.abc import Iterator, Mapping, Sequence @@ -523,7 +521,6 @@ def test_login_survives_a_worker_kill(idp: Idp, tmp_path: Path) -> None: link: Final = browser.get(_cli_link(proxy, session.login_id)) at_idp: Final = browser.get(_assert_idp_redirect(proxy, idp, link, session.login_id)) assert at_idp.status_code == 302, f"{at_idp.status_code} {at_idp.text}" - os.kill(victim, signal.SIGKILL) def callback_attempt() -> httpx.Response | None: try: @@ -531,9 +528,11 @@ def test_login_survives_a_worker_kill(idp: Idp, tmp_path: Path) -> None: except httpx.TransportError: return None - callback: Final = eventually( - callback_attempt, lambda response: response is not None and response.status_code == 200, seconds=60 - ) + with owned.replacing_worker(victim) as kill: + kill() + callback: Final = eventually( + callback_attempt, lambda response: response is not None and response.status_code == 200, seconds=60 + ) assert callback is not None done: Final = _complete_in_browser(browser, callback, session.user_code) assert done.status_code == 200 and CLI_SUCCESS_PAGE_TITLE in done.text, f"{done.status_code} {done.text}" diff --git a/tests/integration/authorization/test_hidden_alias_budget_bypass.py b/tests/integration/authorization/test_hidden_alias_budget_bypass.py index eabdb01602e..67ec30cec2a 100644 --- a/tests/integration/authorization/test_hidden_alias_budget_bypass.py +++ b/tests/integration/authorization/test_hidden_alias_budget_bypass.py @@ -1,7 +1,6 @@ import base64 import json import math -import signal import time import uuid from collections.abc import Callable, Iterator, Mapping @@ -1384,7 +1383,7 @@ def test_killing_one_worker_keeps_the_shadowing_alias_refused(rig: AliasRig, tmp settle_candidate(candidate, target, exhausted, BUDGET_EXCEEDED, seconds=120) settle_candidate(candidate, shadow, exhausted, BUDGET_EXCEEDED, seconds=120) victim: Final = before[0] - victim.send_signal(signal.SIGKILL) + owned.kill_worker(victim.pid) eventually(lambda: _workers(root), _replaced(victim), seconds=120) settle_candidate(candidate, shadow, exhausted, BUDGET_EXCEEDED, seconds=120) assert chat_statuses(candidate, shadow, exhausted, GATEWAY_BURST) == {BUDGET_EXCEEDED} diff --git a/tests/integration/authorization/test_hidden_alias_budget_bypass_chaos.py b/tests/integration/authorization/test_hidden_alias_budget_bypass_chaos.py index a7da04f1364..0a9b1cb3c4e 100644 --- a/tests/integration/authorization/test_hidden_alias_budget_bypass_chaos.py +++ b/tests/integration/authorization/test_hidden_alias_budget_bypass_chaos.py @@ -1,6 +1,4 @@ -import os import re -import signal import threading import uuid from collections.abc import Callable, Iterator, Mapping @@ -186,14 +184,15 @@ def test_killed_worker_leaves_the_sibling_serving_hidden_free_aliases(rig: Alias ) settle_candidate(candidate, rig.hidden_free, key, 200, seconds=120) settle_candidate(candidate, rig.hidden_paid, key, BUDGET_EXCEEDED, seconds=120) - os.kill(workers[0], signal.SIGKILL) - eventually( - lambda: _tolerant_status(candidate, rig.hidden_free, key, "kill-probe-" + uuid.uuid4().hex), - lambda found: found == 200, - seconds=60, - ) - assert chat_statuses(candidate, rig.hidden_free, key, 8) == {200} - assert chat_statuses(candidate, rig.hidden_paid, key, 8) == {BUDGET_EXCEEDED} + with owned.replacing_worker(workers[0]) as kill: + kill() + eventually( + lambda: _tolerant_status(candidate, rig.hidden_free, key, "kill-probe-" + uuid.uuid4().hex), + lambda found: found == 200, + seconds=60, + ) + assert chat_statuses(candidate, rig.hidden_free, key, 8) == {200} + assert chat_statuses(candidate, rig.hidden_paid, key, 8) == {BUDGET_EXCEEDED} eventually( lambda: _worker_startups(owned.log), lambda found: len(found[0]) == 3 and found[1] == 3, diff --git a/tests/integration/authorization/test_key_guardrail_opt_out_chaos.py b/tests/integration/authorization/test_key_guardrail_opt_out_chaos.py index f205b3804f6..78b226f619b 100644 --- a/tests/integration/authorization/test_key_guardrail_opt_out_chaos.py +++ b/tests/integration/authorization/test_key_guardrail_opt_out_chaos.py @@ -1,23 +1,20 @@ import json -import os -import signal import socket import threading import time import uuid from concurrent.futures import ThreadPoolExecutor -from hashlib import sha256 from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer from pathlib import Path from queue import SimpleQueue from typing import Final import httpx -import psutil from integration._support.client import Gateway, eventually, string_value from integration._support.database import read_rows from integration._support.process import owned_proxy, owned_proxy_process +from integration._support.timing import assert_cut_short, timed from integration.authorization._guardrail_opt_out import ( chat, guardrail_config, @@ -250,13 +247,8 @@ def test_flag_denial_survives_worker_kill(gateway: Gateway, tmp_path: Path) -> N caller: Final = non_admin_caller(scenario, member, team, model) alias: Final = "audit-kill-" + uuid.uuid4().hex - workers: Final = eventually( - lambda: psutil.Process(owned.process.pid).children(recursive=True), - lambda children: len(children) >= 2, - seconds=30, - ) - victim: Final = workers[0] - os.kill(victim.pid, signal.SIGKILL) + workers: Final = eventually(owned.worker_pids, lambda pids: len(pids) >= 2, seconds=30) + owned.kill_worker(min(workers)) probe: Final = eventually( lambda: candidate.request( @@ -303,17 +295,18 @@ def test_exempt_chats_do_not_wait_on_slow_guardrail_sink(gateway: Gateway, tmp_p model: Final = scenario.model() exempt: Final = scenario.key(models=[model], disable_global_guardrails=True) - started: Final = time.monotonic() with ThreadPoolExecutor(max_workers=10) as pool: - responses: Final = tuple( + calls: Final = tuple( pool.map( - lambda index: chat(candidate, model, exempt, f"slow-sink-{index}-{uuid.uuid4().hex}"), + lambda index: timed( + lambda: chat(candidate, model, exempt, f"slow-sink-{index}-{uuid.uuid4().hex}") + ), range(10), ) ) - elapsed: Final = time.monotonic() - started + responses: Final = tuple(call.value for call in calls) assert all(response.status_code == 200 for response in responses), [ (response.status_code, response.text) for response in responses ] - assert elapsed < sink_delay, f"exempt chats waited on the guardrail sink: {elapsed}s" + assert_cut_short(max(call.seconds for call in calls), injected_delay=sink_delay) assert sink.drain() == () diff --git a/tests/integration/authorization/test_websocket_rejection_logging.py b/tests/integration/authorization/test_websocket_rejection_logging.py index ea0a841a2b5..c12b387bfd0 100644 --- a/tests/integration/authorization/test_websocket_rejection_logging.py +++ b/tests/integration/authorization/test_websocket_rejection_logging.py @@ -304,10 +304,11 @@ async def test_a_rejection_burst_survives_a_killed_worker_without_an_asgi_crash( before: Final = await _burst(owned, shapes, "burst") assert _crash_lines(before) == (), before victim: Final = eventually(lambda: _workers(owned), lambda workers: len(workers) == 2)[0] - victim.kill() - during: Final = await _refused(owned, "/v1/responses", shapes["unknown_key"], f"model={_tag('during')}") - assert during.status == 403, during.window - assert _crash_lines(during.window) == (), during.window + with owned.replacing_worker(victim.pid) as kill: + kill() + during: Final = await _refused(owned, "/v1/responses", shapes["unknown_key"], f"model={_tag('during')}") + assert during.status == 403, during.window + assert _crash_lines(during.window) == (), during.window eventually( lambda: frozenset(worker.pid for worker in _workers(owned)), lambda pids: len(pids) == 2 and victim.pid not in pids, diff --git a/tests/integration/caching/test_response_cache_chaos.py b/tests/integration/caching/test_response_cache_chaos.py index 2e663c97d1f..17637c088a0 100644 --- a/tests/integration/caching/test_response_cache_chaos.py +++ b/tests/integration/caching/test_response_cache_chaos.py @@ -438,14 +438,16 @@ def test_worker_killed_before_burst_leaves_the_survivor_serving_and_never_pins_a before: Final = _workers(root) assert len(before) == WORKERS, [process.pid for process in before] victim: Final = before[0] - victim.kill() - victim.wait(timeout=10) - with httpx.Client(base_url=target, timeout=15, trust_env=False) as fresh: - readiness: Final = fresh.get("/health/readiness") - assert readiness.status_code == 200, f"/health/readiness with worker {victim.pid} dead: {readiness.text}" - burst: Final = Burst(target, key) - burst.start(pool, _burst_calls(fleet, PER_ENDPOINT // 2, STREAMS_PER_ENDPOINT // 2)) - outcomes: Final = burst.outcomes() + with owned.replacing_worker(victim.pid) as kill: + kill() + with httpx.Client(base_url=target, timeout=15, trust_env=False) as fresh: + readiness: Final = fresh.get("/health/readiness") + assert readiness.status_code == 200, ( + f"/health/readiness with worker {victim.pid} dead: {readiness.text}" + ) + burst: Final = Burst(target, key) + burst.start(pool, _burst_calls(fleet, PER_ENDPOINT // 2, STREAMS_PER_ENDPOINT // 2)) + outcomes: Final = burst.outcomes() respawned: Final = eventually( lambda: tuple(process.pid for process in _workers(root)), lambda pids: len(pids) == WORKERS and victim.pid not in pids, diff --git a/tests/integration/compatibility/test_granian_admission_saturation.py b/tests/integration/compatibility/test_granian_admission_saturation.py index 9f7cec44d89..02c349b2c03 100644 --- a/tests/integration/compatibility/test_granian_admission_saturation.py +++ b/tests/integration/compatibility/test_granian_admission_saturation.py @@ -11,6 +11,7 @@ import httpx import pytest from tests._master_key import MASTER_KEY +from tests.integration._support.process import stop_child MAX_IN_FLIGHT: Final = 8 MAX_QUEUED: Final = 8 @@ -102,16 +103,6 @@ general_settings: assert len(rejected) >= REQUESTS - MAX_IN_FLIGHT - MAX_QUEUED, statuses assert all(result[1] for result in rejected), rejected finally: - proxy_process.terminate() - try: - proxy_process.wait(timeout=10) - except subprocess.TimeoutExpired: - proxy_process.kill() - proxy_process.wait() + stop_child(proxy_process) finally: - fake_process.terminate() - try: - fake_process.wait(timeout=10) - except subprocess.TimeoutExpired: - fake_process.kill() - fake_process.wait() + stop_child(fake_process) diff --git a/tests/integration/compatibility/test_integration_worker_cleanup.py b/tests/integration/compatibility/test_integration_worker_cleanup.py index a3d849ffe98..c5f19dfa50f 100644 --- a/tests/integration/compatibility/test_integration_worker_cleanup.py +++ b/tests/integration/compatibility/test_integration_worker_cleanup.py @@ -13,6 +13,7 @@ from psycopg import sql from pydantic import BaseModel, ConfigDict from tests.integration._support.database import read_rows +from tests.integration._support.process import stop_processes FAULT_CASE: Final = """import json, os, pathlib, psutil from urllib.parse import urlsplit @@ -44,13 +45,7 @@ def _survivor(identity: tuple[int, float]) -> psutil.Process | None: def _clean_fault_resources(resources: CrashResources) -> None: - processes: Final = tuple(process for identity in resources.processes if (process := _survivor(identity))) - for process in processes: - process.terminate() - _stopped, alive = psutil.wait_procs(processes, timeout=5) - for process in alive: - process.kill() - psutil.wait_procs(alive, timeout=3) + stop_processes(process for identity in resources.processes if (process := _survivor(identity))) with psycopg.connect(os.environ["DATABASE_URL"], autocommit=True) as admin: admin.execute(sql.SQL("DROP DATABASE IF EXISTS {} WITH (FORCE)").format(sql.Identifier(resources.database))) @@ -65,7 +60,9 @@ def test_worker_crash_fails_the_run_and_releases_its_database_and_processes(tmp_ "PYTHONPATH": f"{root}:{root / 'tests'}:{root / 'tests' / 'e2e'}", "PYTEST_ADDOPTS": "", "INTEGRATION_WORKER_ISOLATION": "1", - "INTEGRATION_PEER_URL": "", + "INTEGRATION_WORKER_PEER": "", + "INTEGRATION_WORKER_COST_MAP": "", + "INTEGRATION_COVERAGE_FILE": "", "INTEGRATION_PROXY_READ_REPLICA_URL": "", "INTEGRATION_RESULTS_DIR": str(tmp_path / "results"), "WORKER_CRASH_RECORD": str(record), diff --git a/tests/integration/compatibility/test_missing_body_param_status.py b/tests/integration/compatibility/test_missing_body_param_status.py index facf8007910..fbdf7027d4b 100644 --- a/tests/integration/compatibility/test_missing_body_param_status.py +++ b/tests/integration/compatibility/test_missing_body_param_status.py @@ -3,7 +3,6 @@ from __future__ import annotations import asyncio import json import os -import signal import socket import subprocess import sys @@ -20,7 +19,7 @@ import psutil import pytest from integration._support.client import JSON_OBJECT, Gateway, Scenario, eventually, object_value, string_value from integration._support.database import read_rows -from integration._support.process import owned_proxy_process +from integration._support.process import owned_proxy_process, paused, stop_child from integration._support.upstream import ScenarioHandle, delete_scenario, register_scenario from openai import AsyncOpenAI, BadRequestError, OpenAI from pydantic import JsonValue @@ -1455,10 +1454,7 @@ def _upstream(directory: Path) -> Iterator[tuple[subprocess.Popen[bytes], str]]: eventually(lambda: _healthy(url), lambda status: status == 200, seconds=30) yield process, url finally: - if process.poll() is None: - process.send_signal(signal.SIGCONT) - process.terminate() - process.wait(timeout=10) + stop_child(process) def _register_owned(url: str, identity: str, scripted: JsonResponse) -> ScenarioHandle: @@ -1478,10 +1474,6 @@ def _delete_owned(handle: ScenarioHandle) -> None: ).raise_for_status() -def _workers(process: subprocess.Popen[bytes]) -> tuple[psutil.Process, ...]: - return tuple(psutil.Process(process.pid).children(recursive=True)) - - def _process_tree_line(process: psutil.Process) -> str: try: return f"{process.pid} {' '.join(process.cmdline())}" @@ -1489,13 +1481,6 @@ def _process_tree_line(process: psutil.Process) -> str: return f"{process.pid} " -def _worker_alive(worker: psutil.Process) -> bool: - try: - return worker.is_running() and worker.status() != psutil.STATUS_ZOMBIE - except psutil.NoSuchProcess: - return False - - def _chat_body(model: str, marker: str, valid: bool) -> dict[str, JsonValue]: return {"model": model, "user": marker, **({"messages": [{"role": "user", "content": marker}]} if valid else {})} @@ -1556,11 +1541,7 @@ def test_upstream_pause_and_worker_kill_preserve_required_body_status( processes: Final = (process_root, *process_root.children(recursive=True)) process_tree: Final = "\n".join(_process_tree_line(process) for process in processes) record_property("owned_proxy_process_tree", process_tree) - workers: Final = eventually( - lambda: _workers(owned.process), - lambda children: len(children) >= 2, - seconds=30, - ) + workers: Final = eventually(owned.worker_pids, lambda pids: len(pids) >= 2, seconds=30) observations: Final = _Observations(url) missing_routes: Final = ( ("aspeech", "input"), @@ -1585,26 +1566,22 @@ def test_upstream_pause_and_worker_kill_preserve_required_body_status( valid_calls: Final = tuple( ("/v1/chat/completions", _chat_body(chat_model, marker, True), marker) for marker in markers ) - upstream.send_signal(signal.SIGSTOP) - try: - with ThreadPoolExecutor(max_workers=10) as pool: - missing_futures: Final = tuple( - pool.submit(_post, candidate, path, body) for path, body, _expected, _identity in missing_calls - ) - missing: Final = tuple( - (call, future.result(timeout=15)) for call, future in zip(missing_calls, missing_futures) - ) - assert all( - response.status_code == 400 and response.json() == expected - for (_path, _body, expected, _identity), response in missing - ), [response.text for _call, response in missing] - paused_statuses: Final = tuple(response.status_code for _call, response in missing) - record_property( - "chaos_paused_missing_status_counts", - str({status: paused_statuses.count(status) for status in sorted(set(paused_statuses))}), - ) - finally: - upstream.send_signal(signal.SIGCONT) + with paused(upstream), ThreadPoolExecutor(max_workers=10) as pool: + missing_futures: Final = tuple( + pool.submit(_post, candidate, path, body) for path, body, _expected, _identity in missing_calls + ) + missing: Final = tuple( + (call, future.result(timeout=15)) for call, future in zip(missing_calls, missing_futures) + ) + assert all( + response.status_code == 400 and response.json() == expected + for (_path, _body, expected, _identity), response in missing + ), [response.text for _call, response in missing] + paused_statuses: Final = tuple(response.status_code for _call, response in missing) + record_property( + "chaos_paused_missing_status_counts", + str({status: paused_statuses.count(status) for status in sorted(set(paused_statuses))}), + ) with ThreadPoolExecutor(max_workers=20) as pool: valid_futures: Final = tuple( pool.submit(_post, candidate, path, body) for path, body, _marker in valid_calls @@ -1650,14 +1627,14 @@ def test_upstream_pause_and_worker_kill_preserve_required_body_status( record_property("chaos_spend_query", 'SELECT request_id FROM "LiteLLM_SpendLogs" WHERE request_id = %s') record_property("chaos_spend_response_id_count", len(request_ids)) record_property("chaos_spend_row_counts", str(tuple(len(rows) for rows in spend_rows))) - workers[0].kill() - eventually(lambda: _worker_alive(workers[0]), lambda alive: not alive, seconds=10) - post_kill: Final = tuple( - _post(candidate, path, body) for path, body, _expected, _identity in missing_calls[:5] - ) - assert all( - response.status_code == 400 and response.json() == expected - for response, (_path, _body, expected, _identity) in zip(post_kill, missing_calls[:5]) - ), [response.text for response in post_kill] - recovered: Final = _post(candidate, "/v1/chat/completions", _chat_body(chat_model, "recovered", True)) - assert recovered.status_code == 200, recovered.text + with owned.replacing_worker(min(workers)) as kill: + kill() + post_kill: Final = tuple( + _post(candidate, path, body) for path, body, _expected, _identity in missing_calls[:5] + ) + assert all( + response.status_code == 400 and response.json() == expected + for response, (_path, _body, expected, _identity) in zip(post_kill, missing_calls[:5]) + ), [response.text for response in post_kill] + recovered: Final = _post(candidate, "/v1/chat/completions", _chat_body(chat_model, "recovered", True)) + assert recovered.status_code == 200, recovered.text diff --git a/tests/integration/compatibility/test_openai_sdk_client_paths_chaos.py b/tests/integration/compatibility/test_openai_sdk_client_paths_chaos.py index 0bcc9fc4340..e9ee4a323cd 100644 --- a/tests/integration/compatibility/test_openai_sdk_client_paths_chaos.py +++ b/tests/integration/compatibility/test_openai_sdk_client_paths_chaos.py @@ -1,7 +1,6 @@ from __future__ import annotations import json -import signal import subprocess import uuid from collections.abc import Mapping @@ -19,7 +18,7 @@ from anthropic import Anthropic from anthropic import APIError as AnthropicApiError from integration._support.client import JSON_OBJECT, Gateway, Scenario, eventually, object_value from integration._support.database import read_rows -from integration._support.process import graceful_stop_seconds, owned_proxy_process, owned_upstream +from integration._support.process import graceful_stop_seconds, owned_proxy_process, owned_upstream, paused from integration._support.upstream import ScenarioHandle, delete_scenario, register_scenario from openai import APIError as OpenAiApiError from openai import OpenAI @@ -261,25 +260,6 @@ def _landed_once(response_ids: tuple[str, ...]) -> tuple[tuple[Mapping[str, Json return rows -def _is_worker(child: psutil.Process) -> bool: - try: - return "spawn_main" in " ".join(child.cmdline()) - except psutil.Error: - return False - - -def _worker_alive(worker: psutil.Process) -> bool: - try: - return worker.is_running() and worker.status() != psutil.STATUS_ZOMBIE - except psutil.NoSuchProcess: - return False - - -def _alive_workers(process: subprocess.Popen[bytes]) -> tuple[psutil.Process, ...]: - children: Final = tuple(child for child in psutil.Process(process.pid).children() if _is_worker(child)) - return tuple(child for child in children if _worker_alive(child)) - - def _process_tree(process: subprocess.Popen[bytes]) -> str: root: Final = psutil.Process(process.pid) return "\n".join(_process_line(member) for member in (root, *root.children(recursive=True))) @@ -313,12 +293,9 @@ def test_c01_upstream_pause_mid_burst_keeps_liveliness_and_lands_every_id_once( with ThreadPoolExecutor(max_workers=len(calls)) as pool: futures: Final = _submit(pool, clients, calls) eventually(partial(_done_count, futures), lambda done: done >= 1, seconds=60) - upstream.send_signal(signal.SIGSTOP) - try: + with paused(upstream): done_at_pause: Final = _done_count(futures) paused_liveliness: Final = tuple(_liveliness(gateway) for _ in range(3)) - finally: - upstream.send_signal(signal.SIGCONT) done_at_resume: Final = _done_count(futures) outcomes: Final = tuple(future.result(timeout=180) for future in futures) record_property("c01_burst_size", len(calls)) @@ -358,27 +335,25 @@ def test_c02_worker_sigkill_mid_burst_keeps_serving_and_respawns( partial(_every_worker_serves_every_kind, owned.gateway, deployments), lambda served: served, seconds=90 ) clients: Final = _clients(owned.gateway, deployments) - workers: Final = eventually( - partial(_alive_workers, owned.process), lambda found: len(found) == 2, seconds=60 - ) + workers: Final = eventually(owned.worker_pids, lambda found: len(found) == 2, seconds=60) record_property("c02_process_tree_before_kill", _process_tree(owned.process)) - victim: Final = workers[0] + victim: Final = min(workers) calls: Final = _burst_calls(_CHAT_KINDS, 15) observations: Final = _Observations(slot.url) with ThreadPoolExecutor(max_workers=len(calls)) as pool: futures: Final = _submit(pool, clients, calls) eventually(partial(_done_count, futures), lambda done: done >= 1, seconds=60) - victim.kill() - eventually(partial(_worker_alive, victim), lambda alive: not alive, seconds=10) - done_at_kill: Final = _done_count(futures) - survivor: Final = eventually( - lambda: _attempt(clients, _Call("chat", f"after-kill-{uuid.uuid4().hex}")), - lambda outcome: isinstance(outcome, _Completed), - seconds=30, - ) + with owned.replacing_worker(victim) as kill: + kill() + done_at_kill: Final = _done_count(futures) + survivor: Final = eventually( + lambda: _attempt(clients, _Call("chat", f"after-kill-{uuid.uuid4().hex}")), + lambda outcome: isinstance(outcome, _Completed), + seconds=30, + ) outcomes: Final = tuple(future.result(timeout=180) for future in futures) record_property("c02_burst_size", len(calls)) - record_property("c02_killed_worker_pid", victim.pid) + record_property("c02_killed_worker_pid", victim) record_property("c02_done_at_kill", done_at_kill) assert isinstance(survivor, _Completed), survivor completed: Final = tuple(outcome for outcome in outcomes if isinstance(outcome, _Completed)) @@ -388,12 +363,12 @@ def test_c02_worker_sigkill_mid_burst_keeps_serving_and_respawns( record_property("c02_inflight_failure_errors", str(tuple(outcome.error for outcome in failed))) assert completed, outcomes respawned: Final = eventually( - partial(_alive_workers, owned.process), - lambda found: len(found) == 2 and victim.pid not in {worker.pid for worker in found}, + owned.worker_pids, + lambda found: len(found) == 2 and victim not in found, seconds=graceful_stop_seconds() + 60, ) record_property("c02_process_tree_after_respawn", _process_tree(owned.process)) - record_property("c02_respawned_worker_pids", str(tuple(worker.pid for worker in respawned))) + record_property("c02_respawned_worker_pids", str(tuple(sorted(respawned)))) after_respawn: Final = eventually( lambda: _attempt(clients, _Call("chat", f"after-respawn-{uuid.uuid4().hex}")), lambda outcome: isinstance(outcome, _Completed), diff --git a/tests/integration/configuration/test_codex_catalog_boot.py b/tests/integration/configuration/test_codex_catalog_boot.py index 717b9471244..4868e43b855 100644 --- a/tests/integration/configuration/test_codex_catalog_boot.py +++ b/tests/integration/configuration/test_codex_catalog_boot.py @@ -1,6 +1,5 @@ import json import os -import signal import stat import subprocess import sys @@ -11,12 +10,11 @@ from pathlib import Path from typing import Final import httpx -import psutil import pytest from pydantic import JsonValue from tests.integration._support.client import JSON_OBJECT, Gateway, Scenario, eventually, object_value, string_value -from tests.integration._support.process import group_members, owned_proxy, owned_proxy_process +from tests.integration._support.process import owned_proxy, owned_proxy_process CLIENT_VERSION: Final = "0.159.3" FRESH_CONNECTION: Final = {"Connection": "close"} @@ -224,16 +222,6 @@ def test_team_owned_deployments_of_one_name_serve_each_team_its_own_upstream(gat assert outsider[mixed]["service_tiers"] == [_generic_tier("flex")], outsider[mixed] -def _worker_pids(root_pid: int) -> frozenset[int]: - def is_worker(process: psutil.Process) -> bool: - try: - return "spawn_main" in " ".join(process.cmdline()) - except (psutil.NoSuchProcess, psutil.AccessDenied, psutil.ZombieProcess): - return False - - return frozenset(process.pid for process in group_members(root_pid) if is_worker(process)) - - def _catalog_outcome(proxy: Gateway, slugs: Sequence[JsonValue]) -> str: try: response: Final = _catalog_response(proxy) @@ -252,22 +240,24 @@ def test_catalog_survives_a_worker_kill(gateway: Gateway, tmp_path: Path) -> Non config: Final = _write_config(tmp_path, gateway.upstream_url, tiered, plain, alias) with owned_proxy_process(gateway, tmp_path, {"STORE_MODEL_IN_DB": "False"}, config=config, workers=2) as owned: candidate: Final = owned.gateway - workers: Final = eventually(lambda: _worker_pids(owned.process.pid), lambda pids: len(pids) == 2, seconds=30) + workers: Final = eventually(owned.worker_pids, lambda pids: len(pids) == 2, seconds=30) victim: Final = min(workers) slugs: Final = sorted(entry["slug"] for entry in _entries(_catalog_response(candidate))) assert slugs == sorted((tiered, plain, alias)), slugs - def attempt(index: int) -> str: - if index == 2: - os.kill(victim, signal.SIGKILL) - return _catalog_outcome(candidate, slugs) + with owned.replacing_worker(victim) as kill: - with ThreadPoolExecutor(max_workers=BURST) as pool: - outcomes: Final = tuple(pool.map(attempt, range(BURST))) + def attempt(index: int) -> str: + if index == 2: + kill() + return _catalog_outcome(candidate, slugs) + + with ThreadPoolExecutor(max_workers=BURST) as pool: + outcomes: Final = tuple(pool.map(attempt, range(BURST))) assert outcomes.count("ok") >= 1, outcomes assert all(outcome == "ok" or outcome.startswith("transport:") for outcome in outcomes), outcomes respawned: Final = eventually( - lambda: _worker_pids(owned.process.pid), + owned.worker_pids, lambda pids: len(pids) == 2 and victim not in pids, seconds=60, ) diff --git a/tests/integration/configuration/test_lazy_routes_flag.py b/tests/integration/configuration/test_lazy_routes_flag.py index 28930ec29b4..7809b357a7e 100644 --- a/tests/integration/configuration/test_lazy_routes_flag.py +++ b/tests/integration/configuration/test_lazy_routes_flag.py @@ -23,7 +23,6 @@ from typing import Final import anthropic import httpx import openai -import psutil import pytest import yaml from pydantic import JsonValue, TypeAdapter @@ -273,17 +272,6 @@ def _status(candidate: Gateway, path: str) -> int: return candidate.request("GET", path).status_code -def _workers(owned: OwnedProxy) -> tuple[psutil.Process, ...]: - return tuple(child for child in psutil.Process(owned.process.pid).children() if _is_worker(child)) - - -def _is_worker(child: psutil.Process) -> bool: - try: - return "spawn_main" in " ".join(child.cmdline()) and child.status() != psutil.STATUS_ZOMBIE - except psutil.Error: - return False - - @pytest.mark.parametrize("spelling", ("1", "Yes", "ON")) def test_every_truthy_spelling_of_the_flag_registers_the_ticket_features_at_startup( gateway: Gateway, tmp_path: Path, spelling: str @@ -457,18 +445,19 @@ def test_disable_lazy_routes_flag_route_table_survives_a_boot_burst_and_a_killed assert statuses == (200,) * len(probes), statuses assert _routed_features(owned.gateway) == at_boot, "the boot burst changed the route table" - victim: Final = eventually(lambda: _workers(owned), lambda workers: len(workers) == 2)[0] - victim.kill() - with httpx.Client(base_url=owned.gateway.client.base_url, timeout=15, trust_env=False) as fresh: - survivor: Final = Gateway(fresh, owned.gateway.key, owned.gateway.upstream_url) - during: Final = tuple(survivor.request("GET", "/v1/mcp/server").status_code for _ in range(10)) + victim: Final = min(eventually(owned.worker_pids, lambda workers: len(workers) == 2)) + with owned.replacing_worker(victim) as kill: + kill() + with httpx.Client(base_url=owned.gateway.client.base_url, timeout=15, trust_env=False) as fresh: + survivor: Final = Gateway(fresh, owned.gateway.key, owned.gateway.upstream_url) + during: Final = tuple(survivor.request("GET", "/v1/mcp/server").status_code for _ in range(10)) assert during == (200,) * 10, during respawned: Final = eventually( - lambda: frozenset(worker.pid for worker in _workers(owned)), - lambda pids: len(pids) == 2 and victim.pid not in pids, + owned.worker_pids, + lambda pids: len(pids) == 2 and victim not in pids, seconds=30, ) - assert f"Child process [{victim.pid}] died" in owned.log.read_text(), respawned + assert f"Child process [{victim}] died" in owned.log.read_text(), respawned tables: Final = tuple(_routed_features(owned.gateway) for _ in range(4)) assert all(table == at_boot for table in tables), "the respawned worker disagrees on the route table" diff --git a/tests/integration/configuration/test_multi_worker_slow_import_boot.py b/tests/integration/configuration/test_multi_worker_slow_import_boot.py index fb544f37a66..966a4a79b96 100644 --- a/tests/integration/configuration/test_multi_worker_slow_import_boot.py +++ b/tests/integration/configuration/test_multi_worker_slow_import_boot.py @@ -5,7 +5,7 @@ from typing import Final import pytest from integration._support.client import Gateway, eventually -from integration._support.process import graceful_stop_seconds, owned_proxy_process +from integration._support.process import WorkerHealthcheck, graceful_stop_seconds, owned_proxy_process WORKERS: Final = 2 DEPLOYMENT_HEALTHCHECK_SECONDS: Final = 5 @@ -38,7 +38,7 @@ def test_owned_proxy_workers_outlast_the_deployments_healthcheck_default(gateway tmp_path, _slow_worker_environment(tmp_path), workers=WORKERS, - extra_arguments=("--timeout_worker_healthcheck", str(int(graceful_stop_seconds()))), + healthcheck=WorkerHealthcheck.LENIENT, ) as owned: started: Final = eventually( lambda: STARTED_WORKER.findall(owned.log.read_text()), diff --git a/tests/integration/conftest.py b/tests/integration/conftest.py index 047347ae962..cb23713e6b2 100644 --- a/tests/integration/conftest.py +++ b/tests/integration/conftest.py @@ -7,6 +7,7 @@ import re from collections.abc import AsyncIterator, Generator, Iterator, Mapping, Sequence from contextlib import ExitStack from copy import copy +from functools import partial from importlib.metadata import version from pathlib import Path from typing import TYPE_CHECKING, Final, Protocol, runtime_checkable @@ -24,7 +25,7 @@ from tests.integration._support.manifest import OWNED_DIRECTORIES from tests.integration._support.ordering import collected_case, fixture_groups, scheduling_groups from tests.integration._support.provider import SharedProvider, shared_provider from tests.integration._support.routing import RoutingPlugin -from tests.integration._support.workers import worker_endpoints, worker_gateway +from tests.integration._support.workers import gateway_profile, start_worker_stacks, worker_stack from tests.integration.run import GITHUB_FILES if TYPE_CHECKING: @@ -84,24 +85,14 @@ def pytest_xdist_setupnodes(config: pytest.Config, specs: Sequence[XSpec]) -> No resources: Final = ExitStack() config.add_cleanup(resources.close) destination: Final = Path(os.environ["INTEGRATION_RESULTS_DIR"]) - peer_needed: Final = bool(os.environ.get("INTEGRATION_PEER_URL")) - replica_needed: Final = bool(os.environ.get("INTEGRATION_PROXY_READ_REPLICA_URL")) - for index, spec in enumerate(specs): - results: Final = destination / f"gw{index}" - endpoints: Final = worker_endpoints() - with pytest.MonkeyPatch.context() as environment: - for name, value in endpoints.items(): - environment.setenv(name, value) - environment.setenv("INTEGRATION_RESULTS_DIR", str(results)) - services: Final = resources.enter_context( - worker_gateway(results / "runtime", peer=peer_needed, replica=replica_needed) - ) - spec.env = { - **STRING_ENVIRONMENT.validate_python(spec.env), - **endpoints, - **services, - "INTEGRATION_RESULTS_DIR": str(results), - } + stacks: Final = start_worker_stacks( + tuple(destination / f"gw{index}" for index in range(len(specs))), + partial(worker_stack, profile=gateway_profile(os.environ)), + resources, + os.cpu_count() or 1, + ) + for spec, stack in zip(specs, stacks, strict=True): + spec.env = {**STRING_ENVIRONMENT.validate_python(spec.env), **stack} class IntegrationReportPlugin: diff --git a/tests/integration/cost_calculation/assertions.py b/tests/integration/cost_calculation/assertions.py index 58a0fe99aab..bf5985a409d 100644 --- a/tests/integration/cost_calculation/assertions.py +++ b/tests/integration/cost_calculation/assertions.py @@ -62,20 +62,25 @@ def assert_breakdown( assert (header is None and omitted_component_allowed) or ( header is not None and approx_equal(float(header), expected_component) ), f"{case_name}: {header_name} {header} != expected {expected_component}" - if response is not None and expected.cost_header and response_content_type == "application/json" and any( - component is not None - for component in ( - expected.cache_read_cost, - expected.cache_creation_cost, - expected.reasoning_cost, - expected.tool_usage_cost, + if ( + response is not None + and expected.cost_header + and response_content_type == "application/json" + and any( + component is not None + for component in ( + expected.cache_read_cost, + expected.cache_creation_cost, + expected.reasoning_cost, + expected.tool_usage_cost, + ) ) ): input_header: str | None = response.headers.get("x-litellm-response-cost-input") output_header: str | None = response.headers.get("x-litellm-response-cost-output") - expected_input_header: float = expected.input_cost - ( - expected.cache_read_cost or 0.0 - ) - (expected.cache_creation_cost or 0.0) + expected_input_header: float = ( + expected.input_cost - (expected.cache_read_cost or 0.0) - (expected.cache_creation_cost or 0.0) + ) assert input_header is not None and approx_equal(float(input_header), expected_input_header), ( f"{case_name}: x-litellm-response-cost-input {input_header} != expected {expected_input_header}" ) diff --git a/tests/integration/cost_calculation/bridge_cost_parity/__init__.py b/tests/integration/cost_calculation/bridge_cost_parity/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/integration/cost_calculation/bridge_cost_parity/case.py b/tests/integration/cost_calculation/bridge_cost_parity/case.py new file mode 100644 index 00000000000..a3970a156a7 --- /dev/null +++ b/tests/integration/cost_calculation/bridge_cost_parity/case.py @@ -0,0 +1,63 @@ +from dataclasses import dataclass +from typing import Final + +from integration.cost_calculation.cost_tracking_case import CostTrackingTestCase, ExactExpected + +ROW_FIELDS: Final = frozenset({"spend", "prompt_tokens", "completion_tokens"}) +_EXPECTED_FIELDS: Final = frozenset(ExactExpected.model_fields) + + +@dataclass(frozen=True, slots=True, kw_only=True) +class ProviderCostQuirk: + """A provider-level reason the bridged leg legitimately bills a different row: the `ExactExpected` fields that + differ, why, and the provider document that says so.""" + + fields: frozenset[str] + reason: str + source: str + + def __post_init__(self) -> None: + unknown: Final = self.fields - _EXPECTED_FIELDS + if not self.fields or unknown: + raise ValueError(f"quirk fields must be non-empty ExactExpected fields, got {sorted(unknown)}") + if not self.source.startswith("https://"): + raise ValueError("a quirk cites the provider document that explains it") + + +@dataclass(frozen=True, slots=True, kw_only=True) +class BridgeCostParityTestCase: + """One stored provider response billed through the provider's own API (`native`) and through a client API that + LiteLLM bridges into it (`bridged`). Both legs bill the one row unless `quirk` names the fields that differ.""" + + native: CostTrackingTestCase + bridged: CostTrackingTestCase + quirk: ProviderCostQuirk | None = None + + def __post_init__(self) -> None: + if self.native.endpoint == self.bridged.endpoint: + raise ValueError(f"{self.id}: native and bridged cases must use different client endpoints") + if self.native.response != self.bridged.response: + raise ValueError(f"{self.id}: native and bridged cases must share one stored provider response") + if (self.native.model, self.native.deployment) != (self.bridged.model, self.bridged.deployment): + raise ValueError(f"{self.id}: native and bridged cases must share model and deployment") + if not isinstance(self.native.expected, ExactExpected) or not isinstance(self.bridged.expected, ExactExpected): + raise ValueError(f"{self.id}: both legs need an ExactExpected row") + native: Final = self.native.expected.model_dump() + bridged: Final = self.bridged.expected.model_dump() + differing: Final = frozenset(field for field in _EXPECTED_FIELDS if native[field] != bridged[field]) + if differing != self.quirk_fields: + raise ValueError( + f"{self.id}: expected rows differ in {sorted(differing)} but the quirk covers {sorted(self.quirk_fields)}" + ) + + @property + def quirk_fields(self) -> frozenset[str]: + return self.quirk.fields if self.quirk is not None else frozenset() + + @property + def compared_row_fields(self) -> frozenset[str]: + return ROW_FIELDS - self.quirk_fields + + @property + def id(self) -> str: + return self.bridged.name diff --git a/tests/integration/cost_calculation/bridge_cost_parity/runner.py b/tests/integration/cost_calculation/bridge_cost_parity/runner.py new file mode 100644 index 00000000000..5300913534b --- /dev/null +++ b/tests/integration/cost_calculation/bridge_cost_parity/runner.py @@ -0,0 +1,12 @@ +from typing import Final + +from integration._support.client import Gateway +from integration.cost_calculation.bridge_cost_parity.case import BridgeCostParityTestCase +from integration.cost_calculation.stream_parity.runner import bill_legs + + +def assert_bridge_cost_parity(case: BridgeCostParityTestCase, gateway: Gateway) -> None: + billed: Final = bill_legs(case.id, (case.native, case.bridged), gateway) + native: Final = billed[case.native.name].model_dump(include=set(case.compared_row_fields)) + bridged: Final = billed[case.bridged.name].model_dump(include=set(case.compared_row_fields)) + assert bridged == native, (case.id, billed[case.native.name], billed[case.bridged.name]) diff --git a/tests/integration/cost_calculation/bridge_cost_parity/test_basic.py b/tests/integration/cost_calculation/bridge_cost_parity/test_basic.py new file mode 100644 index 00000000000..9f0ca44636d --- /dev/null +++ b/tests/integration/cost_calculation/bridge_cost_parity/test_basic.py @@ -0,0 +1,21 @@ +"""A request through a bridged client API bills exactly what the provider's native API bills for the same response. + +A bridge translates one API into another before the provider call and translates the provider's usage back after it +(LIT-7730, LIT-9098). Each case is one stored provider response sent twice through the same key, once through the +provider's own API and once through a client API LiteLLM bridges into it. Both LiteLLM_SpendLogs rows must bill the +case's row, and they must agree with each other except where a documented provider quirk says otherwise. +""" + +import pytest +from integration._support.client import Gateway +from integration.cost_calculation.bridge_cost_parity.case import BridgeCostParityTestCase +from integration.cost_calculation.bridge_cost_parity.runner import assert_bridge_cost_parity +from integration.cost_calculation.catalog import BRIDGE_COST_PARITY_CASES + + +@pytest.mark.parametrize("case", BRIDGE_COST_PARITY_CASES, ids=[case.id for case in BRIDGE_COST_PARITY_CASES]) +@pytest.mark.timeout(180) +def test_bridged_request_bills_the_same_row_as_the_native_request( + case: BridgeCostParityTestCase, gateway: Gateway +) -> None: + assert_bridge_cost_parity(case, gateway) diff --git a/tests/integration/cost_calculation/cases/anthropic.py b/tests/integration/cost_calculation/cases/anthropic.py index 762d9469366..1ea9cb44868 100644 --- a/tests/integration/cost_calculation/cases/anthropic.py +++ b/tests/integration/cost_calculation/cases/anthropic.py @@ -7,6 +7,8 @@ from typing import Final from integration.cost_calculation.cost_tracking_case import ( CostTrackingTestCase, ExactExpected, + FailureDetails, + FailureExpected, JsonResponse, RecountExpected, RecountRates, @@ -1459,6 +1461,313 @@ CLAUDE_OPUS_5_ANTHROPIC_FAST_MODE: Final = CostTrackingTestCase( ), ) +CLAUDE_OPUS_5_ANTHROPIC_FAST_MODE_CACHE: Final = CostTrackingTestCase( + name="claude-opus-5-anthropic_fast_mode_cache", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="claude-opus-5", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + "cache_control": {"type": "ephemeral"}, + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "121e2a9a4452 summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": False, + "allowed_openai_params": [], + }, + response=JsonResponse( + content_type="application/json", + body={ + "id": "msg_$REQUEST_ID", + "type": "message", + "role": "assistant", + "model": "claude-opus-5", + "content": [{"type": "text", "text": "scripted answer 121e2a9a4452"}], + "stop_reason": "end_turn", + "usage": { + "input_tokens": 1840, + "output_tokens": 412, + "cache_read_input_tokens": 400, + "cache_creation_input_tokens": 300, + "cache_creation": {"ephemeral_5m_input_tokens": 100, "ephemeral_1h_input_tokens": 200}, + "speed": "fast", + }, + }, + ), + expected=ExactExpected( + spend=0.13395, + input_cost=0.07215, + output_cost=0.0618, + prompt_tokens=2540, + completion_tokens=412, + cache_read_cost=0.0012, + cache_creation_cost=0.01575, + ), +) + +CLAUDE_OPUS_5_ANTHROPIC_FAST_MODE_CACHE_STREAM: Final = CostTrackingTestCase( + name="claude-opus-5-anthropic_fast_mode_cache-stream", + covers="quota_management.spend_tracking.scripted_wire.logs_cost", + model="claude-opus-5", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + "cache_control": {"type": "ephemeral"}, + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "6e8ee84cf8b8 summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": True, + "stream_options": {"include_usage": True}, + "allowed_openai_params": [], + }, + response=SseResponse( + content_type="text/event-stream", + frames=sse_frames( + { + "type": "message_start", + "message": { + "id": "msg_$REQUEST_ID", + "type": "message", + "role": "assistant", + "model": "claude-opus-5", + "content": [], + "stop_reason": None, + "usage": { + "input_tokens": 1840, + "cache_read_input_tokens": 400, + "cache_creation_input_tokens": 300, + "cache_creation": {"ephemeral_5m_input_tokens": 100, "ephemeral_1h_input_tokens": 200}, + "speed": "fast", + }, + }, + }, + {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}, + { + "type": "content_block_delta", + "index": 0, + "delta": {"type": "text_delta", "text": "scripted answer 6e8ee84cf8b8"}, + }, + {"type": "content_block_stop", "index": 0}, + {"type": "message_delta", "delta": {"stop_reason": "end_turn"}, "usage": {"output_tokens": 412}}, + {"type": "message_stop"}, + ), + ), + expected=ExactExpected( + spend=0.13395, + input_cost=0.07215, + output_cost=0.0618, + prompt_tokens=2540, + completion_tokens=412, + cache_read_cost=0.0012, + cache_creation_cost=0.01575, + ), +) + +CLAUDE_OPUS_5_MESSAGES_FAST_MODE_CACHE: Final = CostTrackingTestCase( + name="claude-opus-5-messages_fast_mode_cache", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="claude-opus-5", + endpoint="/v1/messages", + request={ + "model": "$MODEL", + "max_tokens": 412, + "messages": [ + { + "role": "user", + "content": [ + { + "type": "text", + "text": "6d2006a65289 summarize the attached material in one line", + "cache_control": {"type": "ephemeral"}, + } + ], + } + ], + }, + response=JsonResponse( + content_type="application/json", + body={ + "id": "msg_$REQUEST_ID", + "type": "message", + "role": "assistant", + "model": "claude-opus-5", + "content": [{"type": "text", "text": "scripted response"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": { + "input_tokens": 1840, + "output_tokens": 412, + "cache_read_input_tokens": 400, + "cache_creation_input_tokens": 300, + "cache_creation": {"ephemeral_5m_input_tokens": 100, "ephemeral_1h_input_tokens": 200}, + "speed": "fast", + }, + }, + ), + expected=ExactExpected( + spend=0.13395, + input_cost=0.07215, + output_cost=0.0618, + prompt_tokens=2540, + completion_tokens=412, + cache_read_cost=0.0012, + cache_creation_cost=0.01575, + ), +) + +CLAUDE_OPUS_5_ANTHROPIC_STANDARD_SPEED_CACHE: Final = CostTrackingTestCase( + name="claude-opus-5-anthropic_standard_speed_cache", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="claude-opus-5", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + "cache_control": {"type": "ephemeral"}, + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "2b9904617ede summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": False, + "allowed_openai_params": [], + }, + response=JsonResponse( + content_type="application/json", + body={ + "id": "msg_$REQUEST_ID", + "type": "message", + "role": "assistant", + "model": "claude-opus-5", + "content": [{"type": "text", "text": "scripted answer 2b9904617ede"}], + "stop_reason": "end_turn", + "usage": { + "input_tokens": 1840, + "output_tokens": 412, + "cache_read_input_tokens": 400, + "cache_creation_input_tokens": 300, + "cache_creation": {"ephemeral_5m_input_tokens": 100, "ephemeral_1h_input_tokens": 200}, + "speed": "standard", + }, + }, + ), + expected=ExactExpected( + spend=0.022325, + input_cost=0.012025, + output_cost=0.0103, + prompt_tokens=2540, + completion_tokens=412, + cache_read_cost=0.0002, + cache_creation_cost=0.002625, + ), +) + +CLAUDE_OPUS_5_ANTHROPIC_FAST_MODE_US_INFERENCE_CACHE: Final = CostTrackingTestCase( + name="claude-opus-5-anthropic_fast_mode_us_inference_cache", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="claude-opus-5", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + "cache_control": {"type": "ephemeral"}, + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "04b13c88ad07 summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": False, + "allowed_openai_params": [], + }, + response=JsonResponse( + content_type="application/json", + body={ + "id": "msg_$REQUEST_ID", + "type": "message", + "role": "assistant", + "model": "claude-opus-5", + "content": [{"type": "text", "text": "scripted answer 04b13c88ad07"}], + "stop_reason": "end_turn", + "usage": { + "input_tokens": 1840, + "output_tokens": 412, + "cache_read_input_tokens": 400, + "cache_creation_input_tokens": 300, + "cache_creation": {"ephemeral_5m_input_tokens": 100, "ephemeral_1h_input_tokens": 200}, + "speed": "fast", + "inference_geo": "us", + }, + }, + ), + expected=ExactExpected( + spend=0.147345, + input_cost=0.079365, + output_cost=0.06798, + prompt_tokens=2540, + completion_tokens=412, + cache_read_cost=0.00132, + cache_creation_cost=0.017325, + ), +) + CLAUDE_OPUS_5_ANTHROPIC_US_INFERENCE: Final = CostTrackingTestCase( name="claude-opus-5-anthropic_us_inference", covers="quota_management.spend_tracking.cost_matrix.logs_cost", @@ -3760,6 +4069,2593 @@ CLAUDE_OPUS_5_5_INPUT_TEXT: Final = CostTrackingTestCase( ) +CLAUDE_SONNET_5_5_INPUT_TEXT: Final = CostTrackingTestCase( + name="claude-sonnet-5-5-input_text", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="claude-sonnet-5-5", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "d5078ef4d4ba summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": False, + "allowed_openai_params": [], + }, + response=JsonResponse( + content_type="application/json", + body={ + "model": "claude-sonnet-5-5", + "id": "msg_$REQUEST_ID", + "type": "message", + "role": "assistant", + "content": [{"type": "text", "text": "scripted answer d5078ef4d4ba"}], + "container": None, + "stop_reason": "end_turn", + "stop_sequence": None, + "stop_details": None, + "usage": { + "input_tokens": 1840, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 0, + "cache_creation": {"ephemeral_5m_input_tokens": 0, "ephemeral_1h_input_tokens": 0}, + "output_tokens": 412, + "output_tokens_details": {"thinking_tokens": 0}, + "service_tier": "standard", + "inference_geo": "global", + }, + "diagnostics": None, + }, + ), + expected=ExactExpected( + spend=0.01287, input_cost=0.006072, output_cost=0.006798, prompt_tokens=1840, completion_tokens=412 + ), +) + +CLAUDE_SONNET_5_5_STREAM: Final = CostTrackingTestCase( + name="claude-sonnet-5-5-stream", + covers="quota_management.spend_tracking.scripted_wire.logs_cost", + model="claude-sonnet-5-5", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "a648e59362e1 summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": True, + "stream_options": {"include_usage": True}, + "allowed_openai_params": [], + }, + response=SseResponse( + content_type="text/event-stream", + frames=sse_frames( + { + "type": "message_start", + "message": { + "model": "claude-sonnet-5-5", + "id": "msg_$REQUEST_ID", + "type": "message", + "role": "assistant", + "content": [], + "container": None, + "stop_reason": None, + "stop_sequence": None, + "stop_details": None, + "usage": { + "input_tokens": 1840, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 0, + "cache_creation": {"ephemeral_5m_input_tokens": 0, "ephemeral_1h_input_tokens": 0}, + "output_tokens": 1, + "service_tier": "standard", + "inference_geo": "global", + }, + "diagnostics": None, + }, + }, + {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}, + {"type": "ping"}, + { + "type": "content_block_delta", + "index": 0, + "delta": {"type": "text_delta", "text": "scripted answer a648e59362e1"}, + }, + {"type": "content_block_stop", "index": 0}, + { + "type": "message_delta", + "delta": {"stop_reason": "end_turn", "stop_sequence": None, "stop_details": None, "container": None}, + "usage": { + "input_tokens": 1840, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 0, + "output_tokens": 412, + "output_tokens_details": {"thinking_tokens": 0}, + }, + }, + {"type": "message_stop"}, + ), + ), + expected=ExactExpected( + spend=0.01287, input_cost=0.006072, output_cost=0.006798, prompt_tokens=1840, completion_tokens=412 + ), +) + +CLAUDE_SONNET_5_5_CACHE_READ: Final = CostTrackingTestCase( + name="claude-sonnet-5-5-cache_read", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="claude-sonnet-5-5", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + "cache_control": {"type": "ephemeral"}, + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "d07bcc773205 summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": False, + "allowed_openai_params": [], + }, + response=JsonResponse( + content_type="application/json", + body={ + "model": "claude-sonnet-5-5", + "id": "msg_$REQUEST_ID", + "type": "message", + "role": "assistant", + "content": [{"type": "text", "text": "scripted answer d07bcc773205"}], + "container": None, + "stop_reason": "end_turn", + "stop_sequence": None, + "stop_details": None, + "usage": { + "input_tokens": 640, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 12288, + "cache_creation": {"ephemeral_5m_input_tokens": 0, "ephemeral_1h_input_tokens": 0}, + "output_tokens": 380, + "output_tokens_details": {"thinking_tokens": 0}, + "service_tier": "standard", + "inference_geo": "global", + }, + "diagnostics": None, + }, + ), + expected=ExactExpected( + spend=0.01243704, input_cost=0.00616704, output_cost=0.00627, prompt_tokens=12928, completion_tokens=380 + ), +) + +CLAUDE_SONNET_5_5_CACHE_WRITE_5M: Final = CostTrackingTestCase( + name="claude-sonnet-5-5-cache_write_5m", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="claude-sonnet-5-5", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + "cache_control": {"type": "ephemeral"}, + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "f163bea2f9c0 summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": False, + "allowed_openai_params": [], + }, + response=JsonResponse( + content_type="application/json", + body={ + "model": "claude-sonnet-5-5", + "id": "msg_$REQUEST_ID", + "type": "message", + "role": "assistant", + "content": [{"type": "text", "text": "scripted answer f163bea2f9c0"}], + "container": None, + "stop_reason": "end_turn", + "stop_sequence": None, + "stop_details": None, + "usage": { + "input_tokens": 512, + "cache_creation_input_tokens": 9216, + "cache_read_input_tokens": 0, + "cache_creation": {"ephemeral_5m_input_tokens": 9216, "ephemeral_1h_input_tokens": 0}, + "output_tokens": 350, + "output_tokens_details": {"thinking_tokens": 0}, + "service_tier": "standard", + "inference_geo": "global", + }, + "diagnostics": None, + }, + ), + expected=ExactExpected( + spend=0.0454806, input_cost=0.0397056, output_cost=0.005775, prompt_tokens=9728, completion_tokens=350 + ), +) + +CLAUDE_SONNET_5_5_CACHE_WRITE_1H: Final = CostTrackingTestCase( + name="claude-sonnet-5-5-cache_write_1h", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="claude-sonnet-5-5", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + "cache_control": {"type": "ephemeral", "ttl": "1h"}, + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "c42d91d3c8cb summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": False, + "allowed_openai_params": [], + }, + response=JsonResponse( + content_type="application/json", + body={ + "model": "claude-sonnet-5-5", + "id": "msg_$REQUEST_ID", + "type": "message", + "role": "assistant", + "content": [{"type": "text", "text": "scripted answer c42d91d3c8cb"}], + "container": None, + "stop_reason": "end_turn", + "stop_sequence": None, + "stop_details": None, + "usage": { + "input_tokens": 512, + "cache_creation_input_tokens": 9216, + "cache_read_input_tokens": 0, + "cache_creation": {"ephemeral_5m_input_tokens": 0, "ephemeral_1h_input_tokens": 9216}, + "output_tokens": 350, + "output_tokens_details": {"thinking_tokens": 0}, + "service_tier": "standard", + "inference_geo": "global", + }, + "diagnostics": None, + }, + ), + expected=ExactExpected( + spend=0.0682902, input_cost=0.0625152, output_cost=0.005775, prompt_tokens=9728, completion_tokens=350 + ), +) + +CLAUDE_SONNET_5_5_MESSAGES_INPUT_TEXT: Final = CostTrackingTestCase( + name="claude-sonnet-5-5-messages_input_text", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="claude-sonnet-5-5", + endpoint="/v1/messages", + request={ + "model": "$MODEL", + "max_tokens": 412, + "messages": [{"role": "user", "content": "summarize the attached material in one line"}], + }, + response=JsonResponse( + content_type="application/json", + body={ + "model": "claude-sonnet-5-5", + "id": "msg_$REQUEST_ID", + "type": "message", + "role": "assistant", + "content": [{"type": "text", "text": "scripted response"}], + "container": None, + "stop_reason": "end_turn", + "stop_sequence": None, + "stop_details": None, + "usage": { + "input_tokens": 1840, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 0, + "cache_creation": {"ephemeral_5m_input_tokens": 0, "ephemeral_1h_input_tokens": 0}, + "output_tokens": 412, + "output_tokens_details": {"thinking_tokens": 0}, + "service_tier": "standard", + "inference_geo": "global", + }, + "diagnostics": None, + }, + ), + expected=ExactExpected( + spend=0.01287, input_cost=0.006072, output_cost=0.006798, prompt_tokens=1840, completion_tokens=412 + ), +) + +CLAUDE_HAIKU_5_5_INPUT_TEXT: Final = CostTrackingTestCase( + name="claude-haiku-5-5-input_text", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="claude-haiku-5-5", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "10673f3a1db3 summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": False, + "allowed_openai_params": [], + }, + response=JsonResponse( + content_type="application/json", + body={ + "model": "claude-haiku-5-5", + "id": "msg_$REQUEST_ID", + "type": "message", + "role": "assistant", + "content": [{"type": "text", "text": "scripted answer 10673f3a1db3"}], + "container": None, + "stop_reason": "end_turn", + "stop_sequence": None, + "stop_details": None, + "usage": { + "input_tokens": 1840, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 0, + "cache_creation": {"ephemeral_5m_input_tokens": 0, "ephemeral_1h_input_tokens": 0}, + "output_tokens": 412, + "output_tokens_details": {"thinking_tokens": 0}, + "service_tier": "standard", + "inference_geo": "global", + }, + "diagnostics": None, + }, + ), + expected=ExactExpected( + spend=0.00429, input_cost=0.002024, output_cost=0.002266, prompt_tokens=1840, completion_tokens=412 + ), +) + +CLAUDE_HAIKU_5_5_STREAM: Final = CostTrackingTestCase( + name="claude-haiku-5-5-stream", + covers="quota_management.spend_tracking.scripted_wire.logs_cost", + model="claude-haiku-5-5", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "2bf9bcba4444 summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": True, + "stream_options": {"include_usage": True}, + "allowed_openai_params": [], + }, + response=SseResponse( + content_type="text/event-stream", + frames=sse_frames( + { + "type": "message_start", + "message": { + "model": "claude-haiku-5-5", + "id": "msg_$REQUEST_ID", + "type": "message", + "role": "assistant", + "content": [], + "container": None, + "stop_reason": None, + "stop_sequence": None, + "stop_details": None, + "usage": { + "input_tokens": 1840, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 0, + "cache_creation": {"ephemeral_5m_input_tokens": 0, "ephemeral_1h_input_tokens": 0}, + "output_tokens": 1, + "service_tier": "standard", + "inference_geo": "global", + }, + "diagnostics": None, + }, + }, + {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}, + {"type": "ping"}, + { + "type": "content_block_delta", + "index": 0, + "delta": {"type": "text_delta", "text": "scripted answer 2bf9bcba4444"}, + }, + {"type": "content_block_stop", "index": 0}, + { + "type": "message_delta", + "delta": {"stop_reason": "end_turn", "stop_sequence": None, "stop_details": None, "container": None}, + "usage": { + "input_tokens": 1840, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 0, + "output_tokens": 412, + "output_tokens_details": {"thinking_tokens": 0}, + }, + }, + {"type": "message_stop"}, + ), + ), + expected=ExactExpected( + spend=0.00429, input_cost=0.002024, output_cost=0.002266, prompt_tokens=1840, completion_tokens=412 + ), +) + +CLAUDE_HAIKU_5_5_CACHE_READ: Final = CostTrackingTestCase( + name="claude-haiku-5-5-cache_read", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="claude-haiku-5-5", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + "cache_control": {"type": "ephemeral"}, + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "65dcddea7db0 summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": False, + "allowed_openai_params": [], + }, + response=JsonResponse( + content_type="application/json", + body={ + "model": "claude-haiku-5-5", + "id": "msg_$REQUEST_ID", + "type": "message", + "role": "assistant", + "content": [{"type": "text", "text": "scripted answer 65dcddea7db0"}], + "container": None, + "stop_reason": "end_turn", + "stop_sequence": None, + "stop_details": None, + "usage": { + "input_tokens": 640, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 12288, + "cache_creation": {"ephemeral_5m_input_tokens": 0, "ephemeral_1h_input_tokens": 0}, + "output_tokens": 380, + "output_tokens_details": {"thinking_tokens": 0}, + "service_tier": "standard", + "inference_geo": "global", + }, + "diagnostics": None, + }, + ), + expected=ExactExpected( + spend=0.00414568, input_cost=0.00205568, output_cost=0.00209, prompt_tokens=12928, completion_tokens=380 + ), +) + +CLAUDE_HAIKU_5_5_CACHE_WRITE_5M: Final = CostTrackingTestCase( + name="claude-haiku-5-5-cache_write_5m", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="claude-haiku-5-5", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + "cache_control": {"type": "ephemeral"}, + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "92ac58b8c6ad summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": False, + "allowed_openai_params": [], + }, + response=JsonResponse( + content_type="application/json", + body={ + "model": "claude-haiku-5-5", + "id": "msg_$REQUEST_ID", + "type": "message", + "role": "assistant", + "content": [{"type": "text", "text": "scripted answer 92ac58b8c6ad"}], + "container": None, + "stop_reason": "end_turn", + "stop_sequence": None, + "stop_details": None, + "usage": { + "input_tokens": 512, + "cache_creation_input_tokens": 9216, + "cache_read_input_tokens": 0, + "cache_creation": {"ephemeral_5m_input_tokens": 9216, "ephemeral_1h_input_tokens": 0}, + "output_tokens": 350, + "output_tokens_details": {"thinking_tokens": 0}, + "service_tier": "standard", + "inference_geo": "global", + }, + "diagnostics": None, + }, + ), + expected=ExactExpected( + spend=0.0151602, input_cost=0.0132352, output_cost=0.001925, prompt_tokens=9728, completion_tokens=350 + ), +) + +CLAUDE_HAIKU_5_5_CACHE_WRITE_1H: Final = CostTrackingTestCase( + name="claude-haiku-5-5-cache_write_1h", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="claude-haiku-5-5", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + "cache_control": {"type": "ephemeral", "ttl": "1h"}, + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "3ae40172e262 summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": False, + "allowed_openai_params": [], + }, + response=JsonResponse( + content_type="application/json", + body={ + "model": "claude-haiku-5-5", + "id": "msg_$REQUEST_ID", + "type": "message", + "role": "assistant", + "content": [{"type": "text", "text": "scripted answer 3ae40172e262"}], + "container": None, + "stop_reason": "end_turn", + "stop_sequence": None, + "stop_details": None, + "usage": { + "input_tokens": 512, + "cache_creation_input_tokens": 9216, + "cache_read_input_tokens": 0, + "cache_creation": {"ephemeral_5m_input_tokens": 0, "ephemeral_1h_input_tokens": 9216}, + "output_tokens": 350, + "output_tokens_details": {"thinking_tokens": 0}, + "service_tier": "standard", + "inference_geo": "global", + }, + "diagnostics": None, + }, + ), + expected=ExactExpected( + spend=0.0227634, input_cost=0.0208384, output_cost=0.001925, prompt_tokens=9728, completion_tokens=350 + ), +) + +CLAUDE_HAIKU_5_5_MESSAGES_INPUT_TEXT: Final = CostTrackingTestCase( + name="claude-haiku-5-5-messages_input_text", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="claude-haiku-5-5", + endpoint="/v1/messages", + request={ + "model": "$MODEL", + "max_tokens": 412, + "messages": [{"role": "user", "content": "summarize the attached material in one line"}], + }, + response=JsonResponse( + content_type="application/json", + body={ + "model": "claude-haiku-5-5", + "id": "msg_$REQUEST_ID", + "type": "message", + "role": "assistant", + "content": [{"type": "text", "text": "scripted response"}], + "container": None, + "stop_reason": "end_turn", + "stop_sequence": None, + "stop_details": None, + "usage": { + "input_tokens": 1840, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 0, + "cache_creation": {"ephemeral_5m_input_tokens": 0, "ephemeral_1h_input_tokens": 0}, + "output_tokens": 412, + "output_tokens_details": {"thinking_tokens": 0}, + "service_tier": "standard", + "inference_geo": "global", + }, + "diagnostics": None, + }, + ), + expected=ExactExpected( + spend=0.00429, input_cost=0.002024, output_cost=0.002266, prompt_tokens=1840, completion_tokens=412 + ), +) + +CLAUDE_FABLE_5_INPUT_TEXT: Final = CostTrackingTestCase( + name="claude-fable-5-input_text", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="claude-fable-5", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "5a2c19b20bee summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": False, + "allowed_openai_params": [], + }, + response=JsonResponse( + content_type="application/json", + body={ + "model": "claude-fable-5", + "id": "msg_$REQUEST_ID", + "type": "message", + "role": "assistant", + "content": [{"type": "text", "text": "scripted answer 5a2c19b20bee"}], + "container": None, + "stop_reason": "end_turn", + "stop_sequence": None, + "stop_details": None, + "usage": { + "input_tokens": 1840, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 0, + "cache_creation": {"ephemeral_5m_input_tokens": 0, "ephemeral_1h_input_tokens": 0}, + "output_tokens": 412, + "output_tokens_details": {"thinking_tokens": 0}, + "service_tier": "standard", + "inference_geo": "global", + }, + "diagnostics": None, + }, + ), + expected=ExactExpected( + spend=0.0585, input_cost=0.0276, output_cost=0.0309, prompt_tokens=1840, completion_tokens=412 + ), +) + +CLAUDE_FABLE_5_STREAM: Final = CostTrackingTestCase( + name="claude-fable-5-stream", + covers="quota_management.spend_tracking.scripted_wire.logs_cost", + model="claude-fable-5", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "6c3adfe9ff6d summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": True, + "stream_options": {"include_usage": True}, + "allowed_openai_params": [], + }, + response=SseResponse( + content_type="text/event-stream", + frames=sse_frames( + { + "type": "message_start", + "message": { + "model": "claude-fable-5", + "id": "msg_$REQUEST_ID", + "type": "message", + "role": "assistant", + "content": [], + "container": None, + "stop_reason": None, + "stop_sequence": None, + "stop_details": None, + "usage": { + "input_tokens": 1840, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 0, + "cache_creation": {"ephemeral_5m_input_tokens": 0, "ephemeral_1h_input_tokens": 0}, + "output_tokens": 1, + "service_tier": "standard", + "inference_geo": "global", + }, + "diagnostics": None, + }, + }, + {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}, + {"type": "ping"}, + { + "type": "content_block_delta", + "index": 0, + "delta": {"type": "text_delta", "text": "scripted answer 6c3adfe9ff6d"}, + }, + {"type": "content_block_stop", "index": 0}, + { + "type": "message_delta", + "delta": {"stop_reason": "end_turn", "stop_sequence": None, "stop_details": None, "container": None}, + "usage": { + "input_tokens": 1840, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 0, + "output_tokens": 412, + "output_tokens_details": {"thinking_tokens": 0}, + }, + }, + {"type": "message_stop"}, + ), + ), + expected=ExactExpected( + spend=0.0585, input_cost=0.0276, output_cost=0.0309, prompt_tokens=1840, completion_tokens=412 + ), +) + +CLAUDE_FABLE_5_CACHE_READ: Final = CostTrackingTestCase( + name="claude-fable-5-cache_read", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="claude-fable-5", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + "cache_control": {"type": "ephemeral"}, + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "b928021ea8bb summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": False, + "allowed_openai_params": [], + }, + response=JsonResponse( + content_type="application/json", + body={ + "model": "claude-fable-5", + "id": "msg_$REQUEST_ID", + "type": "message", + "role": "assistant", + "content": [{"type": "text", "text": "scripted answer b928021ea8bb"}], + "container": None, + "stop_reason": "end_turn", + "stop_sequence": None, + "stop_details": None, + "usage": { + "input_tokens": 640, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 12288, + "cache_creation": {"ephemeral_5m_input_tokens": 0, "ephemeral_1h_input_tokens": 0}, + "output_tokens": 380, + "output_tokens_details": {"thinking_tokens": 0}, + "service_tier": "standard", + "inference_geo": "global", + }, + "diagnostics": None, + }, + ), + expected=ExactExpected( + spend=0.056532, input_cost=0.028032, output_cost=0.0285, prompt_tokens=12928, completion_tokens=380 + ), +) + +CLAUDE_FABLE_5_CACHE_WRITE_5M: Final = CostTrackingTestCase( + name="claude-fable-5-cache_write_5m", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="claude-fable-5", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + "cache_control": {"type": "ephemeral"}, + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "fd1ed16af5ca summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": False, + "allowed_openai_params": [], + }, + response=JsonResponse( + content_type="application/json", + body={ + "model": "claude-fable-5", + "id": "msg_$REQUEST_ID", + "type": "message", + "role": "assistant", + "content": [{"type": "text", "text": "scripted answer fd1ed16af5ca"}], + "container": None, + "stop_reason": "end_turn", + "stop_sequence": None, + "stop_details": None, + "usage": { + "input_tokens": 512, + "cache_creation_input_tokens": 9216, + "cache_read_input_tokens": 0, + "cache_creation": {"ephemeral_5m_input_tokens": 9216, "ephemeral_1h_input_tokens": 0}, + "output_tokens": 350, + "output_tokens_details": {"thinking_tokens": 0}, + "service_tier": "standard", + "inference_geo": "global", + }, + "diagnostics": None, + }, + ), + expected=ExactExpected( + spend=0.20673, input_cost=0.18048, output_cost=0.02625, prompt_tokens=9728, completion_tokens=350 + ), +) + +CLAUDE_FABLE_5_CACHE_WRITE_1H: Final = CostTrackingTestCase( + name="claude-fable-5-cache_write_1h", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="claude-fable-5", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + "cache_control": {"type": "ephemeral", "ttl": "1h"}, + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "c373b13470db summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": False, + "allowed_openai_params": [], + }, + response=JsonResponse( + content_type="application/json", + body={ + "model": "claude-fable-5", + "id": "msg_$REQUEST_ID", + "type": "message", + "role": "assistant", + "content": [{"type": "text", "text": "scripted answer c373b13470db"}], + "container": None, + "stop_reason": "end_turn", + "stop_sequence": None, + "stop_details": None, + "usage": { + "input_tokens": 512, + "cache_creation_input_tokens": 9216, + "cache_read_input_tokens": 0, + "cache_creation": {"ephemeral_5m_input_tokens": 0, "ephemeral_1h_input_tokens": 9216}, + "output_tokens": 350, + "output_tokens_details": {"thinking_tokens": 0}, + "service_tier": "standard", + "inference_geo": "global", + }, + "diagnostics": None, + }, + ), + expected=ExactExpected( + spend=0.31041, input_cost=0.28416, output_cost=0.02625, prompt_tokens=9728, completion_tokens=350 + ), +) + +CLAUDE_FABLE_5_MESSAGES_INPUT_TEXT: Final = CostTrackingTestCase( + name="claude-fable-5-messages_input_text", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="claude-fable-5", + endpoint="/v1/messages", + request={ + "model": "$MODEL", + "max_tokens": 412, + "messages": [{"role": "user", "content": "summarize the attached material in one line"}], + }, + response=JsonResponse( + content_type="application/json", + body={ + "model": "claude-fable-5", + "id": "msg_$REQUEST_ID", + "type": "message", + "role": "assistant", + "content": [{"type": "text", "text": "scripted response"}], + "container": None, + "stop_reason": "end_turn", + "stop_sequence": None, + "stop_details": None, + "usage": { + "input_tokens": 1840, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 0, + "cache_creation": {"ephemeral_5m_input_tokens": 0, "ephemeral_1h_input_tokens": 0}, + "output_tokens": 412, + "output_tokens_details": {"thinking_tokens": 0}, + "service_tier": "standard", + "inference_geo": "global", + }, + "diagnostics": None, + }, + ), + expected=ExactExpected( + spend=0.0585, input_cost=0.0276, output_cost=0.0309, prompt_tokens=1840, completion_tokens=412 + ), +) + +CLAUDE_FABLE_5_1_INPUT_TEXT: Final = CostTrackingTestCase( + name="claude-fable-5-1-input_text", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="claude-fable-5-1", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "75b71184a26d summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": False, + "allowed_openai_params": [], + }, + response=JsonResponse( + content_type="application/json", + body={ + "model": "claude-fable-5-1", + "id": "msg_$REQUEST_ID", + "type": "message", + "role": "assistant", + "content": [{"type": "text", "text": "scripted answer 75b71184a26d"}], + "container": None, + "stop_reason": "end_turn", + "stop_sequence": None, + "stop_details": None, + "usage": { + "input_tokens": 1840, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 0, + "cache_creation": {"ephemeral_5m_input_tokens": 0, "ephemeral_1h_input_tokens": 0}, + "output_tokens": 412, + "output_tokens_details": {"thinking_tokens": 0}, + "service_tier": "standard", + "inference_geo": "global", + }, + "diagnostics": None, + }, + ), + expected=ExactExpected( + spend=0.0624, input_cost=0.02944, output_cost=0.03296, prompt_tokens=1840, completion_tokens=412 + ), +) + +CLAUDE_FABLE_5_1_STREAM: Final = CostTrackingTestCase( + name="claude-fable-5-1-stream", + covers="quota_management.spend_tracking.scripted_wire.logs_cost", + model="claude-fable-5-1", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "3d584fa07ae1 summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": True, + "stream_options": {"include_usage": True}, + "allowed_openai_params": [], + }, + response=SseResponse( + content_type="text/event-stream", + frames=sse_frames( + { + "type": "message_start", + "message": { + "model": "claude-fable-5-1", + "id": "msg_$REQUEST_ID", + "type": "message", + "role": "assistant", + "content": [], + "container": None, + "stop_reason": None, + "stop_sequence": None, + "stop_details": None, + "usage": { + "input_tokens": 1840, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 0, + "cache_creation": {"ephemeral_5m_input_tokens": 0, "ephemeral_1h_input_tokens": 0}, + "output_tokens": 1, + "service_tier": "standard", + "inference_geo": "global", + }, + "diagnostics": None, + }, + }, + {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}, + {"type": "ping"}, + { + "type": "content_block_delta", + "index": 0, + "delta": {"type": "text_delta", "text": "scripted answer 3d584fa07ae1"}, + }, + {"type": "content_block_stop", "index": 0}, + { + "type": "message_delta", + "delta": {"stop_reason": "end_turn", "stop_sequence": None, "stop_details": None, "container": None}, + "usage": { + "input_tokens": 1840, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 0, + "output_tokens": 412, + "output_tokens_details": {"thinking_tokens": 0}, + }, + }, + {"type": "message_stop"}, + ), + ), + expected=ExactExpected( + spend=0.0624, input_cost=0.02944, output_cost=0.03296, prompt_tokens=1840, completion_tokens=412 + ), +) + +CLAUDE_FABLE_5_1_CACHE_READ: Final = CostTrackingTestCase( + name="claude-fable-5-1-cache_read", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="claude-fable-5-1", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + "cache_control": {"type": "ephemeral"}, + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "e2c093f0a554 summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": False, + "allowed_openai_params": [], + }, + response=JsonResponse( + content_type="application/json", + body={ + "model": "claude-fable-5-1", + "id": "msg_$REQUEST_ID", + "type": "message", + "role": "assistant", + "content": [{"type": "text", "text": "scripted answer e2c093f0a554"}], + "container": None, + "stop_reason": "end_turn", + "stop_sequence": None, + "stop_details": None, + "usage": { + "input_tokens": 640, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 12288, + "cache_creation": {"ephemeral_5m_input_tokens": 0, "ephemeral_1h_input_tokens": 0}, + "output_tokens": 380, + "output_tokens_details": {"thinking_tokens": 0}, + "service_tier": "standard", + "inference_geo": "global", + }, + "diagnostics": None, + }, + ), + expected=ExactExpected( + spend=0.0603008, input_cost=0.0299008, output_cost=0.0304, prompt_tokens=12928, completion_tokens=380 + ), +) + +CLAUDE_FABLE_5_1_CACHE_WRITE_5M: Final = CostTrackingTestCase( + name="claude-fable-5-1-cache_write_5m", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="claude-fable-5-1", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + "cache_control": {"type": "ephemeral"}, + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "c074c3c72a7d summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": False, + "allowed_openai_params": [], + }, + response=JsonResponse( + content_type="application/json", + body={ + "model": "claude-fable-5-1", + "id": "msg_$REQUEST_ID", + "type": "message", + "role": "assistant", + "content": [{"type": "text", "text": "scripted answer c074c3c72a7d"}], + "container": None, + "stop_reason": "end_turn", + "stop_sequence": None, + "stop_details": None, + "usage": { + "input_tokens": 512, + "cache_creation_input_tokens": 9216, + "cache_read_input_tokens": 0, + "cache_creation": {"ephemeral_5m_input_tokens": 9216, "ephemeral_1h_input_tokens": 0}, + "output_tokens": 350, + "output_tokens_details": {"thinking_tokens": 0}, + "service_tier": "standard", + "inference_geo": "global", + }, + "diagnostics": None, + }, + ), + expected=ExactExpected( + spend=0.220512, input_cost=0.192512, output_cost=0.028, prompt_tokens=9728, completion_tokens=350 + ), +) + +CLAUDE_FABLE_5_1_CACHE_WRITE_1H: Final = CostTrackingTestCase( + name="claude-fable-5-1-cache_write_1h", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="claude-fable-5-1", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + "cache_control": {"type": "ephemeral", "ttl": "1h"}, + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "feab5bd43ecd summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": False, + "allowed_openai_params": [], + }, + response=JsonResponse( + content_type="application/json", + body={ + "model": "claude-fable-5-1", + "id": "msg_$REQUEST_ID", + "type": "message", + "role": "assistant", + "content": [{"type": "text", "text": "scripted answer feab5bd43ecd"}], + "container": None, + "stop_reason": "end_turn", + "stop_sequence": None, + "stop_details": None, + "usage": { + "input_tokens": 512, + "cache_creation_input_tokens": 9216, + "cache_read_input_tokens": 0, + "cache_creation": {"ephemeral_5m_input_tokens": 0, "ephemeral_1h_input_tokens": 9216}, + "output_tokens": 350, + "output_tokens_details": {"thinking_tokens": 0}, + "service_tier": "standard", + "inference_geo": "global", + }, + "diagnostics": None, + }, + ), + expected=ExactExpected( + spend=0.331104, input_cost=0.303104, output_cost=0.028, prompt_tokens=9728, completion_tokens=350 + ), +) + +CLAUDE_FABLE_5_1_MESSAGES_INPUT_TEXT: Final = CostTrackingTestCase( + name="claude-fable-5-1-messages_input_text", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="claude-fable-5-1", + endpoint="/v1/messages", + request={ + "model": "$MODEL", + "max_tokens": 412, + "messages": [{"role": "user", "content": "summarize the attached material in one line"}], + }, + response=JsonResponse( + content_type="application/json", + body={ + "model": "claude-fable-5-1", + "id": "msg_$REQUEST_ID", + "type": "message", + "role": "assistant", + "content": [{"type": "text", "text": "scripted response"}], + "container": None, + "stop_reason": "end_turn", + "stop_sequence": None, + "stop_details": None, + "usage": { + "input_tokens": 1840, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 0, + "cache_creation": {"ephemeral_5m_input_tokens": 0, "ephemeral_1h_input_tokens": 0}, + "output_tokens": 412, + "output_tokens_details": {"thinking_tokens": 0}, + "service_tier": "standard", + "inference_geo": "global", + }, + "diagnostics": None, + }, + ), + expected=ExactExpected( + spend=0.0624, input_cost=0.02944, output_cost=0.03296, prompt_tokens=1840, completion_tokens=412 + ), +) + +CLAUDE_OPUS_5_5_STREAM: Final = CostTrackingTestCase( + name="claude-opus-5-5-stream", + covers="quota_management.spend_tracking.scripted_wire.logs_cost", + model="claude-opus-5-5", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "9d9d47a4b07a summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": True, + "stream_options": {"include_usage": True}, + "allowed_openai_params": [], + }, + response=SseResponse( + content_type="text/event-stream", + frames=sse_frames( + { + "type": "message_start", + "message": { + "model": "claude-opus-5-5", + "id": "msg_$REQUEST_ID", + "type": "message", + "role": "assistant", + "content": [], + "container": None, + "stop_reason": None, + "stop_sequence": None, + "stop_details": None, + "usage": { + "input_tokens": 1840, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 0, + "cache_creation": {"ephemeral_5m_input_tokens": 0, "ephemeral_1h_input_tokens": 0}, + "output_tokens": 1, + "service_tier": "standard", + "inference_geo": "global", + }, + "diagnostics": None, + }, + }, + {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}, + {"type": "ping"}, + { + "type": "content_block_delta", + "index": 0, + "delta": {"type": "text_delta", "text": "scripted answer 9d9d47a4b07a"}, + }, + {"type": "content_block_stop", "index": 0}, + { + "type": "message_delta", + "delta": {"stop_reason": "end_turn", "stop_sequence": None, "stop_details": None, "container": None}, + "usage": { + "input_tokens": 1840, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 0, + "output_tokens": 412, + "output_tokens_details": {"thinking_tokens": 0}, + }, + }, + {"type": "message_stop"}, + ), + ), + expected=ExactExpected( + spend=0.0156, input_cost=0.00736, output_cost=0.00824, prompt_tokens=1840, completion_tokens=412 + ), +) + +CLAUDE_OPUS_5_5_CACHE_READ: Final = CostTrackingTestCase( + name="claude-opus-5-5-cache_read", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="claude-opus-5-5", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + "cache_control": {"type": "ephemeral"}, + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "6c9f4bd8c5cb summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": False, + "allowed_openai_params": [], + }, + response=JsonResponse( + content_type="application/json", + body={ + "model": "claude-opus-5-5", + "id": "msg_$REQUEST_ID", + "type": "message", + "role": "assistant", + "content": [{"type": "text", "text": "scripted answer 6c9f4bd8c5cb"}], + "container": None, + "stop_reason": "end_turn", + "stop_sequence": None, + "stop_details": None, + "usage": { + "input_tokens": 640, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 12288, + "cache_creation": {"ephemeral_5m_input_tokens": 0, "ephemeral_1h_input_tokens": 0}, + "output_tokens": 380, + "output_tokens_details": {"thinking_tokens": 0}, + "service_tier": "standard", + "inference_geo": "global", + }, + "diagnostics": None, + }, + ), + expected=ExactExpected( + spend=0.0126176, input_cost=0.0050176, output_cost=0.0076, prompt_tokens=12928, completion_tokens=380 + ), +) + +CLAUDE_OPUS_5_5_CACHE_WRITE_5M: Final = CostTrackingTestCase( + name="claude-opus-5-5-cache_write_5m", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="claude-opus-5-5", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + "cache_control": {"type": "ephemeral"}, + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "e0df66e44e28 summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": False, + "allowed_openai_params": [], + }, + response=JsonResponse( + content_type="application/json", + body={ + "model": "claude-opus-5-5", + "id": "msg_$REQUEST_ID", + "type": "message", + "role": "assistant", + "content": [{"type": "text", "text": "scripted answer e0df66e44e28"}], + "container": None, + "stop_reason": "end_turn", + "stop_sequence": None, + "stop_details": None, + "usage": { + "input_tokens": 512, + "cache_creation_input_tokens": 9216, + "cache_read_input_tokens": 0, + "cache_creation": {"ephemeral_5m_input_tokens": 9216, "ephemeral_1h_input_tokens": 0}, + "output_tokens": 350, + "output_tokens_details": {"thinking_tokens": 0}, + "service_tier": "standard", + "inference_geo": "global", + }, + "diagnostics": None, + }, + ), + expected=ExactExpected( + spend=0.055128, input_cost=0.048128, output_cost=0.007, prompt_tokens=9728, completion_tokens=350 + ), +) + +CLAUDE_OPUS_5_5_CACHE_WRITE_1H: Final = CostTrackingTestCase( + name="claude-opus-5-5-cache_write_1h", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="claude-opus-5-5", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + "cache_control": {"type": "ephemeral", "ttl": "1h"}, + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "458ab0cbfb35 summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": False, + "allowed_openai_params": [], + }, + response=JsonResponse( + content_type="application/json", + body={ + "model": "claude-opus-5-5", + "id": "msg_$REQUEST_ID", + "type": "message", + "role": "assistant", + "content": [{"type": "text", "text": "scripted answer 458ab0cbfb35"}], + "container": None, + "stop_reason": "end_turn", + "stop_sequence": None, + "stop_details": None, + "usage": { + "input_tokens": 512, + "cache_creation_input_tokens": 9216, + "cache_read_input_tokens": 0, + "cache_creation": {"ephemeral_5m_input_tokens": 0, "ephemeral_1h_input_tokens": 9216}, + "output_tokens": 350, + "output_tokens_details": {"thinking_tokens": 0}, + "service_tier": "standard", + "inference_geo": "global", + }, + "diagnostics": None, + }, + ), + expected=ExactExpected( + spend=0.082776, input_cost=0.075776, output_cost=0.007, prompt_tokens=9728, completion_tokens=350 + ), +) + +CLAUDE_OPUS_5_5_MESSAGES_INPUT_TEXT: Final = CostTrackingTestCase( + name="claude-opus-5-5-messages_input_text", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="claude-opus-5-5", + endpoint="/v1/messages", + request={ + "model": "$MODEL", + "max_tokens": 412, + "messages": [{"role": "user", "content": "summarize the attached material in one line"}], + }, + response=JsonResponse( + content_type="application/json", + body={ + "model": "claude-opus-5-5", + "id": "msg_$REQUEST_ID", + "type": "message", + "role": "assistant", + "content": [{"type": "text", "text": "scripted response"}], + "container": None, + "stop_reason": "end_turn", + "stop_sequence": None, + "stop_details": None, + "usage": { + "input_tokens": 1840, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 0, + "cache_creation": {"ephemeral_5m_input_tokens": 0, "ephemeral_1h_input_tokens": 0}, + "output_tokens": 412, + "output_tokens_details": {"thinking_tokens": 0}, + "service_tier": "standard", + "inference_geo": "global", + }, + "diagnostics": None, + }, + ), + expected=ExactExpected( + spend=0.0156, input_cost=0.00736, output_cost=0.00824, prompt_tokens=1840, completion_tokens=412 + ), +) + +CLAUDE_SONNET_5_5_TOOL_CALL: Final = CostTrackingTestCase( + name="claude-sonnet-5-5-tool_call", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="claude-sonnet-5-5", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "545ca20285d9 summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": False, + "tools": [ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get the current weather and a short forecast for a city.", + "parameters": { + "type": "object", + "properties": { + "city": {"type": "string", "description": "City name"}, + "days": {"type": "integer", "description": "Forecast horizon in days"}, + "units": {"type": "string", "enum": ["metric", "imperial"]}, + }, + "required": ["city"], + }, + }, + } + ], + "tool_choice": "auto", + "allowed_openai_params": ["tool_choice"], + }, + response=JsonResponse( + content_type="application/json", + body={ + "model": "claude-sonnet-5-5", + "id": "msg_$REQUEST_ID", + "type": "message", + "role": "assistant", + "content": [ + { + "type": "tool_use", + "id": "toolu_$REQUEST_ID", + "name": "get_weather", + "input": {"city": "Berlin", "days": 7, "units": "metric"}, + "caller": {"type": "direct"}, + } + ], + "container": None, + "stop_reason": "tool_use", + "stop_sequence": None, + "stop_details": None, + "usage": { + "input_tokens": 1840, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 0, + "cache_creation": {"ephemeral_5m_input_tokens": 0, "ephemeral_1h_input_tokens": 0}, + "output_tokens": 412, + "output_tokens_details": {"thinking_tokens": 0}, + "service_tier": "standard", + "inference_geo": "global", + }, + "diagnostics": None, + }, + ), + expected=ExactExpected( + spend=0.01287, input_cost=0.006072, output_cost=0.006798, prompt_tokens=1840, completion_tokens=412 + ), +) + +CLAUDE_SONNET_5_5_THINKING: Final = CostTrackingTestCase( + name="claude-sonnet-5-5-thinking", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="claude-sonnet-5-5", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "87d4acad824b summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": False, + "reasoning_effort": "medium", + "allowed_openai_params": ["reasoning_effort"], + }, + response=JsonResponse( + content_type="application/json", + body={ + "model": "claude-sonnet-5-5", + "id": "msg_$REQUEST_ID", + "type": "message", + "role": "assistant", + "content": [{"type": "text", "text": "scripted answer 87d4acad824b"}], + "container": None, + "stop_reason": "end_turn", + "stop_sequence": None, + "stop_details": None, + "usage": { + "input_tokens": 1240, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 0, + "cache_creation": {"ephemeral_5m_input_tokens": 0, "ephemeral_1h_input_tokens": 0}, + "output_tokens": 412, + "output_tokens_details": {"thinking_tokens": 0}, + "service_tier": "standard", + "inference_geo": "global", + }, + "diagnostics": None, + }, + ), + expected=ExactExpected( + spend=0.01089, input_cost=0.004092, output_cost=0.006798, prompt_tokens=1240, completion_tokens=412 + ), +) + +CLAUDE_HAIKU_5_5_TOOL_CALL: Final = CostTrackingTestCase( + name="claude-haiku-5-5-tool_call", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="claude-haiku-5-5", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "69e58cd95315 summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": False, + "tools": [ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get the current weather and a short forecast for a city.", + "parameters": { + "type": "object", + "properties": { + "city": {"type": "string", "description": "City name"}, + "days": {"type": "integer", "description": "Forecast horizon in days"}, + "units": {"type": "string", "enum": ["metric", "imperial"]}, + }, + "required": ["city"], + }, + }, + } + ], + "tool_choice": "auto", + "allowed_openai_params": ["tool_choice"], + }, + response=JsonResponse( + content_type="application/json", + body={ + "model": "claude-haiku-5-5", + "id": "msg_$REQUEST_ID", + "type": "message", + "role": "assistant", + "content": [ + { + "type": "tool_use", + "id": "toolu_$REQUEST_ID", + "name": "get_weather", + "input": {"city": "Berlin", "days": 7, "units": "metric"}, + "caller": {"type": "direct"}, + } + ], + "container": None, + "stop_reason": "tool_use", + "stop_sequence": None, + "stop_details": None, + "usage": { + "input_tokens": 1840, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 0, + "cache_creation": {"ephemeral_5m_input_tokens": 0, "ephemeral_1h_input_tokens": 0}, + "output_tokens": 412, + "output_tokens_details": {"thinking_tokens": 0}, + "service_tier": "standard", + "inference_geo": "global", + }, + "diagnostics": None, + }, + ), + expected=ExactExpected( + spend=0.00429, input_cost=0.002024, output_cost=0.002266, prompt_tokens=1840, completion_tokens=412 + ), +) + +CLAUDE_HAIKU_5_5_THINKING: Final = CostTrackingTestCase( + name="claude-haiku-5-5-thinking", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="claude-haiku-5-5", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "eadef5247ccc summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": False, + "reasoning_effort": "medium", + "allowed_openai_params": ["reasoning_effort"], + }, + response=JsonResponse( + content_type="application/json", + body={ + "model": "claude-haiku-5-5", + "id": "msg_$REQUEST_ID", + "type": "message", + "role": "assistant", + "content": [ + { + "type": "thinking", + "thinking": "", + "signature": "CAQS1gQKEAgSGAI4AUIIdGhpbmtpbmcSDOsCmv80YWY0IrUDvBoMog2", + }, + {"type": "text", "text": "scripted answer eadef5247ccc"}, + ], + "container": None, + "stop_reason": "end_turn", + "stop_sequence": None, + "stop_details": None, + "usage": { + "input_tokens": 1240, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 0, + "cache_creation": {"ephemeral_5m_input_tokens": 0, "ephemeral_1h_input_tokens": 0}, + "output_tokens": 4040, + "output_tokens_details": {"thinking_tokens": 3480}, + "service_tier": "standard", + "inference_geo": "global", + }, + "diagnostics": None, + }, + ), + expected=ExactExpected( + spend=0.023584, input_cost=0.001364, output_cost=0.02222, prompt_tokens=1240, completion_tokens=4040 + ), +) + +CLAUDE_FABLE_5_TOOL_CALL: Final = CostTrackingTestCase( + name="claude-fable-5-tool_call", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="claude-fable-5", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "80901cc6dd89 summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": False, + "tools": [ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get the current weather and a short forecast for a city.", + "parameters": { + "type": "object", + "properties": { + "city": {"type": "string", "description": "City name"}, + "days": {"type": "integer", "description": "Forecast horizon in days"}, + "units": {"type": "string", "enum": ["metric", "imperial"]}, + }, + "required": ["city"], + }, + }, + } + ], + "tool_choice": "auto", + "allowed_openai_params": ["tool_choice"], + }, + response=JsonResponse( + content_type="application/json", + body={ + "model": "claude-fable-5", + "id": "msg_$REQUEST_ID", + "type": "message", + "role": "assistant", + "content": [ + { + "type": "tool_use", + "id": "toolu_$REQUEST_ID", + "name": "get_weather", + "input": {"city": "Berlin", "days": 7, "units": "metric"}, + "caller": {"type": "direct"}, + } + ], + "container": None, + "stop_reason": "tool_use", + "stop_sequence": None, + "stop_details": None, + "usage": { + "input_tokens": 1840, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 0, + "cache_creation": {"ephemeral_5m_input_tokens": 0, "ephemeral_1h_input_tokens": 0}, + "output_tokens": 412, + "output_tokens_details": {"thinking_tokens": 0}, + "service_tier": "standard", + "inference_geo": "global", + }, + "diagnostics": None, + }, + ), + expected=ExactExpected( + spend=0.0585, input_cost=0.0276, output_cost=0.0309, prompt_tokens=1840, completion_tokens=412 + ), +) + +CLAUDE_FABLE_5_THINKING: Final = CostTrackingTestCase( + name="claude-fable-5-thinking", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="claude-fable-5", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "d4bb9ddb7e64 summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": False, + "reasoning_effort": "medium", + "allowed_openai_params": ["reasoning_effort"], + }, + response=JsonResponse( + content_type="application/json", + body={ + "model": "claude-fable-5", + "id": "msg_$REQUEST_ID", + "type": "message", + "role": "assistant", + "content": [ + { + "type": "thinking", + "thinking": "", + "signature": "CAQS1gQKEAgSGAI4AUIIdGhpbmtpbmcSDOsCmv80YWY0IrUDvBoMog2", + }, + {"type": "text", "text": "scripted answer d4bb9ddb7e64"}, + ], + "container": None, + "stop_reason": "end_turn", + "stop_sequence": None, + "stop_details": None, + "usage": { + "input_tokens": 1240, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 0, + "cache_creation": {"ephemeral_5m_input_tokens": 0, "ephemeral_1h_input_tokens": 0}, + "output_tokens": 4040, + "output_tokens_details": {"thinking_tokens": 3480}, + "service_tier": "standard", + "inference_geo": "global", + }, + "diagnostics": None, + }, + ), + expected=ExactExpected( + spend=0.3216, input_cost=0.0186, output_cost=0.303, prompt_tokens=1240, completion_tokens=4040 + ), +) + +CLAUDE_FABLE_5_1_TOOL_CALL: Final = CostTrackingTestCase( + name="claude-fable-5-1-tool_call", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="claude-fable-5-1", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "e0fdc87d1c48 summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": False, + "tools": [ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get the current weather and a short forecast for a city.", + "parameters": { + "type": "object", + "properties": { + "city": {"type": "string", "description": "City name"}, + "days": {"type": "integer", "description": "Forecast horizon in days"}, + "units": {"type": "string", "enum": ["metric", "imperial"]}, + }, + "required": ["city"], + }, + }, + } + ], + "tool_choice": "auto", + "allowed_openai_params": ["tool_choice"], + }, + response=JsonResponse( + content_type="application/json", + body={ + "model": "claude-fable-5-1", + "id": "msg_$REQUEST_ID", + "type": "message", + "role": "assistant", + "content": [ + { + "type": "tool_use", + "id": "toolu_$REQUEST_ID", + "name": "get_weather", + "input": {"city": "Berlin", "days": 7, "units": "metric"}, + "caller": {"type": "direct"}, + } + ], + "container": None, + "stop_reason": "tool_use", + "stop_sequence": None, + "stop_details": None, + "usage": { + "input_tokens": 1840, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 0, + "cache_creation": {"ephemeral_5m_input_tokens": 0, "ephemeral_1h_input_tokens": 0}, + "output_tokens": 412, + "output_tokens_details": {"thinking_tokens": 0}, + "service_tier": "standard", + "inference_geo": "global", + }, + "diagnostics": None, + }, + ), + expected=ExactExpected( + spend=0.0624, input_cost=0.02944, output_cost=0.03296, prompt_tokens=1840, completion_tokens=412 + ), +) + +CLAUDE_FABLE_5_1_THINKING: Final = CostTrackingTestCase( + name="claude-fable-5-1-thinking", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="claude-fable-5-1", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "bf3807156155 summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": False, + "reasoning_effort": "medium", + "allowed_openai_params": ["reasoning_effort"], + }, + response=JsonResponse( + content_type="application/json", + body={ + "model": "claude-fable-5-1", + "id": "msg_$REQUEST_ID", + "type": "message", + "role": "assistant", + "content": [{"type": "text", "text": "scripted answer bf3807156155"}], + "container": None, + "stop_reason": "end_turn", + "stop_sequence": None, + "stop_details": None, + "usage": { + "input_tokens": 1240, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 0, + "cache_creation": {"ephemeral_5m_input_tokens": 0, "ephemeral_1h_input_tokens": 0}, + "output_tokens": 412, + "output_tokens_details": {"thinking_tokens": 0}, + "service_tier": "standard", + "inference_geo": "global", + }, + "diagnostics": None, + }, + ), + expected=ExactExpected( + spend=0.0528, input_cost=0.01984, output_cost=0.03296, prompt_tokens=1240, completion_tokens=412 + ), +) + +CLAUDE_OPUS_5_5_TOOL_CALL: Final = CostTrackingTestCase( + name="claude-opus-5-5-tool_call", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="claude-opus-5-5", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "31f57d0ee701 summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": False, + "tools": [ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get the current weather and a short forecast for a city.", + "parameters": { + "type": "object", + "properties": { + "city": {"type": "string", "description": "City name"}, + "days": {"type": "integer", "description": "Forecast horizon in days"}, + "units": {"type": "string", "enum": ["metric", "imperial"]}, + }, + "required": ["city"], + }, + }, + } + ], + "tool_choice": "auto", + "allowed_openai_params": ["tool_choice"], + }, + response=JsonResponse( + content_type="application/json", + body={ + "model": "claude-opus-5-5", + "id": "msg_$REQUEST_ID", + "type": "message", + "role": "assistant", + "content": [ + { + "type": "tool_use", + "id": "toolu_$REQUEST_ID", + "name": "get_weather", + "input": {"city": "Berlin", "days": 7, "units": "metric"}, + "caller": {"type": "direct"}, + } + ], + "container": None, + "stop_reason": "tool_use", + "stop_sequence": None, + "stop_details": None, + "usage": { + "input_tokens": 1840, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 0, + "cache_creation": {"ephemeral_5m_input_tokens": 0, "ephemeral_1h_input_tokens": 0}, + "output_tokens": 412, + "output_tokens_details": {"thinking_tokens": 0}, + "service_tier": "standard", + "inference_geo": "global", + }, + "diagnostics": None, + }, + ), + expected=ExactExpected( + spend=0.0156, input_cost=0.00736, output_cost=0.00824, prompt_tokens=1840, completion_tokens=412 + ), +) + +CLAUDE_OPUS_5_5_THINKING: Final = CostTrackingTestCase( + name="claude-opus-5-5-thinking", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="claude-opus-5-5", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "d4af5bf007f7 summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": False, + "reasoning_effort": "medium", + "allowed_openai_params": ["reasoning_effort"], + }, + response=JsonResponse( + content_type="application/json", + body={ + "model": "claude-opus-5-5", + "id": "msg_$REQUEST_ID", + "type": "message", + "role": "assistant", + "content": [ + { + "type": "thinking", + "thinking": "", + "signature": "CAQS1gQKEAgSGAI4AUIIdGhpbmtpbmcSDOsCmv80YWY0IrUDvBoMog2", + }, + {"type": "text", "text": "scripted answer d4af5bf007f7"}, + ], + "container": None, + "stop_reason": "end_turn", + "stop_sequence": None, + "stop_details": None, + "usage": { + "input_tokens": 1240, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 0, + "cache_creation": {"ephemeral_5m_input_tokens": 0, "ephemeral_1h_input_tokens": 0}, + "output_tokens": 4040, + "output_tokens_details": {"thinking_tokens": 3480}, + "service_tier": "standard", + "inference_geo": "global", + }, + "diagnostics": None, + }, + ), + expected=ExactExpected( + spend=0.08576, input_cost=0.00496, output_cost=0.0808, prompt_tokens=1240, completion_tokens=4040 + ), +) + +CLAUDE_OPUS_5_5_UPSTREAM_404_ZERO_SPEND: Final = CostTrackingTestCase( + name="claude-opus-5-5-upstream_404_zero_spend", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="claude-opus-5-5", + request={ + "model": "$MODEL", + "messages": [{"role": "user", "content": "scripted upstream failure 404"}], + "stream": False, + }, + response=JsonResponse( + content_type="application/json", + body={ + "type": "error", + "error": {"type": "not_found_error", "message": "model: scripted-missing-model"}, + "request_id": "req_$REQUEST_ID", + }, + status=404, + ), + expected=FailureExpected(failure=FailureDetails(status=404)), +) + +CLAUDE_OPUS_5_5_INPUT_TEXT_ROLLUPS: Final = CostTrackingTestCase( + name="claude-opus-5-5-input_text_rollups", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="claude-opus-5-5", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "fed41894bd20 summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": False, + "allowed_openai_params": [], + }, + response=JsonResponse( + content_type="application/json", + body={ + "model": "claude-opus-5-5", + "id": "msg_$UNIQUE_ID", + "type": "message", + "role": "assistant", + "content": [{"type": "text", "text": "scripted answer fed41894bd20"}], + "container": None, + "stop_reason": "end_turn", + "stop_sequence": None, + "stop_details": None, + "usage": { + "input_tokens": 1840, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 0, + "cache_creation": {"ephemeral_5m_input_tokens": 0, "ephemeral_1h_input_tokens": 0}, + "output_tokens": 412, + "output_tokens_details": {"thinking_tokens": 0}, + "service_tier": "standard", + "inference_geo": "global", + }, + "diagnostics": None, + }, + ), + expected=ExactExpected( + spend=0.0156, input_cost=0.00736, output_cost=0.00824, prompt_tokens=1840, completion_tokens=412, rollups=True + ), +) + CASES: Final[tuple[CostTrackingTestCase, ...]] = ( CLAUDE_HAIKU_4_5_INPUT_TEXT, CLAUDE_HAIKU_4_5_CACHE_READ, @@ -3786,6 +6682,11 @@ CASES: Final[tuple[CostTrackingTestCase, ...]] = ( CLAUDE_OPUS_5_TIERED_CACHE_WRITE_ABOVE_200K, CLAUDE_OPUS_5_SERVICE_TIER_PRIORITY, CLAUDE_OPUS_5_ANTHROPIC_FAST_MODE, + CLAUDE_OPUS_5_ANTHROPIC_FAST_MODE_CACHE, + CLAUDE_OPUS_5_ANTHROPIC_FAST_MODE_CACHE_STREAM, + CLAUDE_OPUS_5_MESSAGES_FAST_MODE_CACHE, + CLAUDE_OPUS_5_ANTHROPIC_STANDARD_SPEED_CACHE, + CLAUDE_OPUS_5_ANTHROPIC_FAST_MODE_US_INFERENCE_CACHE, CLAUDE_OPUS_5_ANTHROPIC_US_INFERENCE, CLAUDE_OPUS_5_WEB_SEARCH_MEDIUM, CLAUDE_OPUS_5_STREAM, @@ -3829,6 +6730,47 @@ CASES: Final[tuple[CostTrackingTestCase, ...]] = ( CLAUDE_SONNET_5_PASSTHROUGH_MESSAGES_CACHE_READ, CLAUDE_SONNET_5_CACHE_CREATION_1H_ABOVE_200K, CLAUDE_OPUS_5_5_INPUT_TEXT, + CLAUDE_SONNET_5_5_INPUT_TEXT, + CLAUDE_SONNET_5_5_STREAM, + CLAUDE_SONNET_5_5_CACHE_READ, + CLAUDE_SONNET_5_5_CACHE_WRITE_5M, + CLAUDE_SONNET_5_5_CACHE_WRITE_1H, + CLAUDE_SONNET_5_5_MESSAGES_INPUT_TEXT, + CLAUDE_HAIKU_5_5_INPUT_TEXT, + CLAUDE_HAIKU_5_5_STREAM, + CLAUDE_HAIKU_5_5_CACHE_READ, + CLAUDE_HAIKU_5_5_CACHE_WRITE_5M, + CLAUDE_HAIKU_5_5_CACHE_WRITE_1H, + CLAUDE_HAIKU_5_5_MESSAGES_INPUT_TEXT, + CLAUDE_FABLE_5_INPUT_TEXT, + CLAUDE_FABLE_5_STREAM, + CLAUDE_FABLE_5_CACHE_READ, + CLAUDE_FABLE_5_CACHE_WRITE_5M, + CLAUDE_FABLE_5_CACHE_WRITE_1H, + CLAUDE_FABLE_5_MESSAGES_INPUT_TEXT, + CLAUDE_FABLE_5_1_INPUT_TEXT, + CLAUDE_FABLE_5_1_STREAM, + CLAUDE_FABLE_5_1_CACHE_READ, + CLAUDE_FABLE_5_1_CACHE_WRITE_5M, + CLAUDE_FABLE_5_1_CACHE_WRITE_1H, + CLAUDE_FABLE_5_1_MESSAGES_INPUT_TEXT, + CLAUDE_OPUS_5_5_STREAM, + CLAUDE_OPUS_5_5_CACHE_READ, + CLAUDE_OPUS_5_5_CACHE_WRITE_5M, + CLAUDE_OPUS_5_5_CACHE_WRITE_1H, + CLAUDE_OPUS_5_5_MESSAGES_INPUT_TEXT, + CLAUDE_SONNET_5_5_TOOL_CALL, + CLAUDE_SONNET_5_5_THINKING, + CLAUDE_HAIKU_5_5_TOOL_CALL, + CLAUDE_HAIKU_5_5_THINKING, + CLAUDE_FABLE_5_TOOL_CALL, + CLAUDE_FABLE_5_THINKING, + CLAUDE_FABLE_5_1_TOOL_CALL, + CLAUDE_FABLE_5_1_THINKING, + CLAUDE_OPUS_5_5_TOOL_CALL, + CLAUDE_OPUS_5_5_THINKING, + CLAUDE_OPUS_5_5_UPSTREAM_404_ZERO_SPEND, + CLAUDE_OPUS_5_5_INPUT_TEXT_ROLLUPS, ) PARITY: Final[tuple[StreamParityTestCase, ...]] = ( @@ -3842,9 +6784,17 @@ PARITY: Final[tuple[StreamParityTestCase, ...]] = ( plain=CLAUDE_OPUS_5_RESPONSE_MODEL_OVERRIDE, streamed=CLAUDE_OPUS_5_STREAM_RESPONSE_MODEL_OVERRIDE ), StreamParityTestCase(plain=CLAUDE_OPUS_5_TOOL_CALL, streamed=CLAUDE_OPUS_5_STREAM_TOOL_CALL), + StreamParityTestCase( + plain=CLAUDE_OPUS_5_ANTHROPIC_FAST_MODE_CACHE, streamed=CLAUDE_OPUS_5_ANTHROPIC_FAST_MODE_CACHE_STREAM + ), StreamParityTestCase( plain=CLAUDE_SONNET_5_RESPONSE_MODEL_OVERRIDE, streamed=CLAUDE_SONNET_5_STREAM_RESPONSE_MODEL_OVERRIDE ), StreamParityTestCase(plain=CLAUDE_SONNET_5_TOOL_CALL, streamed=CLAUDE_SONNET_5_STREAM_TOOL_CALL), StreamParityTestCase(plain=CLAUDE_SONNET_5_MESSAGES_INPUT_TEXT, streamed=CLAUDE_SONNET_5_MESSAGES_STREAM), + StreamParityTestCase(plain=CLAUDE_SONNET_5_5_INPUT_TEXT, streamed=CLAUDE_SONNET_5_5_STREAM), + StreamParityTestCase(plain=CLAUDE_HAIKU_5_5_INPUT_TEXT, streamed=CLAUDE_HAIKU_5_5_STREAM), + StreamParityTestCase(plain=CLAUDE_FABLE_5_INPUT_TEXT, streamed=CLAUDE_FABLE_5_STREAM), + StreamParityTestCase(plain=CLAUDE_FABLE_5_1_INPUT_TEXT, streamed=CLAUDE_FABLE_5_1_STREAM), + StreamParityTestCase(plain=CLAUDE_OPUS_5_5_INPUT_TEXT, streamed=CLAUDE_OPUS_5_5_STREAM), ) diff --git a/tests/integration/cost_calculation/cases/azure.py b/tests/integration/cost_calculation/cases/azure.py index 1285dcf98da..b6d6d0496da 100644 --- a/tests/integration/cost_calculation/cases/azure.py +++ b/tests/integration/cost_calculation/cases/azure.py @@ -9,6 +9,8 @@ from integration.cost_calculation.cost_tracking_case import ( CostTrackingTestCase, Deployment, ExactExpected, + FailureDetails, + FailureExpected, JsonResponse, RecountExpected, RecountRates, @@ -2988,6 +2990,2588 @@ AZURE_TEXT_EMBEDDINGS_4_LARGE_DEPLOYMENT: Final = CostTrackingTestCase( ) +_GPT_6_CONTENT_FILTER_RESULTS: Final = { + "hate": {"filtered": False, "severity": "safe"}, + "protected_material_code": {"detected": False, "filtered": False}, + "protected_material_text": {"detected": False, "filtered": False}, + "self_harm": {"filtered": False, "severity": "safe"}, + "sexual": {"filtered": False, "severity": "safe"}, + "violence": {"filtered": False, "severity": "safe"}, +} +_GPT_6_PROMPT_FILTER_RESULTS: Final = [ + { + "prompt_index": 0, + "content_filter_results": { + "hate": {"filtered": False, "severity": "safe"}, + "jailbreak": {"detected": False, "filtered": False}, + "self_harm": {"filtered": False, "severity": "safe"}, + "sexual": {"filtered": False, "severity": "safe"}, + "violence": {"filtered": False, "severity": "safe"}, + }, + } +] + +AZURE_GPT_6_SOL_INPUT_TEXT: Final = CostTrackingTestCase( + name="azure-gpt-6-sol-input_text", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="azure/gpt-6-sol", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "f77f71b0428a summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": False, + "allowed_openai_params": [], + }, + response=JsonResponse( + content_type="application/json", + body={ + "choices": [ + { + "content_filter_results": _GPT_6_CONTENT_FILTER_RESULTS, + "finish_reason": "stop", + "index": 0, + "logprobs": None, + "message": { + "annotations": [], + "content": "scripted answer f77f71b0428a", + "refusal": None, + "role": "assistant", + }, + } + ], + "created": 1791852029, + "id": "chatcmpl-$REQUEST_ID", + "model": "gpt-6-sol-2026-09-22", + "object": "chat.completion", + "prompt_filter_results": _GPT_6_PROMPT_FILTER_RESULTS, + "routing": {"serving_pipereplica": "d20260922011332-060efdec-prefill-r0-dp0-prefill"}, + "service_tier": "default", + "system_fingerprint": None, + "usage": { + "completion_tokens": 412, + "completion_tokens_details": { + "accepted_prediction_tokens": 0, + "audio_tokens": 0, + "reasoning_tokens": 0, + "rejected_prediction_tokens": 0, + }, + "latency_checkpoint": { + "engine_tbt_ms": 4, + "engine_ttft_ms": 131, + "engine_ttlt_ms": 157, + "pre_inference_ms": 55, + "service_tbt_ms": 5, + "service_ttft_ms": 1034, + "service_ttlt_ms": 1058, + "user_visible_ttft_ms": 979, + }, + "prompt_tokens": 1840, + "prompt_tokens_details": {"audio_tokens": 0, "cache_write_tokens": 0, "cached_tokens": 0}, + "total_tokens": 2252, + }, + }, + ), + expected=ExactExpected( + spend=0.00897, input_cost=0.004232, output_cost=0.004738, prompt_tokens=1840, completion_tokens=412 + ), +) + +AZURE_GPT_6_SOL_STREAM: Final = CostTrackingTestCase( + name="azure-gpt-6-sol-stream", + covers="quota_management.spend_tracking.scripted_wire.logs_cost", + model="azure/gpt-6-sol", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "bf171da6c983 summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": True, + "stream_options": {"include_usage": True}, + "allowed_openai_params": [], + }, + response=SseResponse( + content_type="text/event-stream", + frames=sse_frames( + { + "choices": [], + "created": 0, + "id": "", + "model": "", + "object": "", + "prompt_filter_results": _GPT_6_PROMPT_FILTER_RESULTS, + }, + { + "choices": [ + { + "content_filter_results": {}, + "delta": {"content": "", "refusal": None, "role": "assistant"}, + "finish_reason": None, + "index": 0, + "logprobs": None, + } + ], + "created": 1791852023, + "id": "chatcmpl-$REQUEST_ID", + "model": "gpt-6-sol-2026-09-22", + "object": "chat.completion.chunk", + "service_tier": "default", + "system_fingerprint": None, + "usage": None, + }, + { + "choices": [ + { + "content_filter_results": {"protected_material_text": {"detected": False, "filtered": False}}, + "delta": {"content": "scripted answer bf171da6c983"}, + "finish_reason": None, + "index": 0, + "logprobs": None, + } + ], + "created": 1791852023, + "id": "chatcmpl-$REQUEST_ID", + "model": "gpt-6-sol-2026-09-22", + "object": "chat.completion.chunk", + "service_tier": "default", + "system_fingerprint": None, + "usage": None, + }, + { + "choices": [ + {"content_filter_results": {}, "delta": {}, "finish_reason": "stop", "index": 0, "logprobs": None} + ], + "created": 1791852023, + "id": "chatcmpl-$REQUEST_ID", + "model": "gpt-6-sol-2026-09-22", + "object": "chat.completion.chunk", + "service_tier": "default", + "system_fingerprint": None, + "usage": None, + }, + { + "choices": [], + "created": 1791852023, + "id": "chatcmpl-$REQUEST_ID", + "model": "gpt-6-sol-2026-09-22", + "object": "chat.completion.chunk", + "service_tier": "default", + "system_fingerprint": None, + "routing": {"serving_pipereplica": "d20260925225746-6787865d-prefill-r11-dp0-prefill"}, + "usage": { + "completion_tokens": 412, + "completion_tokens_details": { + "accepted_prediction_tokens": 0, + "audio_tokens": 0, + "reasoning_tokens": 0, + "rejected_prediction_tokens": 0, + }, + "latency_checkpoint": { + "engine_tbt_ms": 4, + "engine_ttft_ms": 131, + "engine_ttlt_ms": 157, + "pre_inference_ms": 55, + "service_tbt_ms": 5, + "service_ttft_ms": 1034, + "service_ttlt_ms": 1058, + "user_visible_ttft_ms": 979, + }, + "prompt_tokens": 1840, + "prompt_tokens_details": {"audio_tokens": 0, "cache_write_tokens": 0, "cached_tokens": 0}, + "total_tokens": 2252, + }, + }, + ), + ), + expected=ExactExpected( + spend=0.00897, input_cost=0.004232, output_cost=0.004738, prompt_tokens=1840, completion_tokens=412 + ), +) + +AZURE_GPT_6_SOL_SERVICE_TIER_FLEX: Final = CostTrackingTestCase( + name="azure-gpt-6-sol-service_tier_flex", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="azure/gpt-6-sol", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "54651d96d231 summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": False, + "service_tier": "flex", + "allowed_openai_params": [], + }, + response=JsonResponse( + content_type="application/json", + body={ + "choices": [ + { + "content_filter_results": _GPT_6_CONTENT_FILTER_RESULTS, + "finish_reason": "stop", + "index": 0, + "logprobs": None, + "message": { + "annotations": [], + "content": "scripted answer 54651d96d231", + "refusal": None, + "role": "assistant", + }, + } + ], + "created": 1791852029, + "id": "chatcmpl-$REQUEST_ID", + "model": "gpt-6-sol-2026-09-22", + "object": "chat.completion", + "prompt_filter_results": _GPT_6_PROMPT_FILTER_RESULTS, + "routing": {"serving_pipereplica": "d20260922011332-060efdec-prefill-r0-dp0-prefill"}, + "service_tier": "flex", + "system_fingerprint": None, + "usage": { + "completion_tokens": 412, + "completion_tokens_details": { + "accepted_prediction_tokens": 0, + "audio_tokens": 0, + "reasoning_tokens": 0, + "rejected_prediction_tokens": 0, + }, + "latency_checkpoint": { + "engine_tbt_ms": 4, + "engine_ttft_ms": 131, + "engine_ttlt_ms": 157, + "pre_inference_ms": 55, + "service_tbt_ms": 5, + "service_ttft_ms": 1034, + "service_ttlt_ms": 1058, + "user_visible_ttft_ms": 979, + }, + "prompt_tokens": 1840, + "prompt_tokens_details": {"audio_tokens": 0, "cache_write_tokens": 0, "cached_tokens": 0}, + "total_tokens": 2252, + }, + }, + ), + expected=ExactExpected( + spend=0.004485, input_cost=0.002116, output_cost=0.002369, prompt_tokens=1840, completion_tokens=412 + ), +) + +AZURE_GPT_6_SOL_SERVICE_TIER_PRIORITY: Final = CostTrackingTestCase( + name="azure-gpt-6-sol-service_tier_priority", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="azure/gpt-6-sol", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "6add01b22373 summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": False, + "service_tier": "priority", + "allowed_openai_params": [], + }, + response=JsonResponse( + content_type="application/json", + body={ + "choices": [ + { + "content_filter_results": _GPT_6_CONTENT_FILTER_RESULTS, + "finish_reason": "stop", + "index": 0, + "logprobs": None, + "message": { + "annotations": [], + "content": "scripted answer 6add01b22373", + "refusal": None, + "role": "assistant", + }, + } + ], + "created": 1791852029, + "id": "chatcmpl-$REQUEST_ID", + "model": "gpt-6-sol-2026-09-22", + "object": "chat.completion", + "prompt_filter_results": _GPT_6_PROMPT_FILTER_RESULTS, + "routing": {"serving_pipereplica": "d20260922011332-060efdec-prefill-r0-dp0-prefill"}, + "service_tier": "priority", + "system_fingerprint": None, + "usage": { + "completion_tokens": 412, + "completion_tokens_details": { + "accepted_prediction_tokens": 0, + "audio_tokens": 0, + "reasoning_tokens": 0, + "rejected_prediction_tokens": 0, + }, + "latency_checkpoint": { + "engine_tbt_ms": 4, + "engine_ttft_ms": 131, + "engine_ttlt_ms": 157, + "pre_inference_ms": 55, + "service_tbt_ms": 5, + "service_ttft_ms": 1034, + "service_ttlt_ms": 1058, + "user_visible_ttft_ms": 979, + }, + "prompt_tokens": 1840, + "prompt_tokens_details": {"audio_tokens": 0, "cache_write_tokens": 0, "cached_tokens": 0}, + "total_tokens": 2252, + }, + }, + ), + expected=ExactExpected( + spend=0.01794, input_cost=0.008464, output_cost=0.009476, prompt_tokens=1840, completion_tokens=412 + ), +) + +AZURE_GPT_6_LUNA_INPUT_TEXT: Final = CostTrackingTestCase( + name="azure-gpt-6-luna-input_text", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="azure/gpt-6-luna", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "a6a8b0e9d64c summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": False, + "allowed_openai_params": [], + }, + response=JsonResponse( + content_type="application/json", + body={ + "choices": [ + { + "content_filter_results": _GPT_6_CONTENT_FILTER_RESULTS, + "finish_reason": "stop", + "index": 0, + "logprobs": None, + "message": { + "annotations": [], + "content": "scripted answer a6a8b0e9d64c", + "refusal": None, + "role": "assistant", + }, + } + ], + "created": 1791852029, + "id": "chatcmpl-$REQUEST_ID", + "model": "gpt-6-luna-2026-09-22", + "object": "chat.completion", + "prompt_filter_results": _GPT_6_PROMPT_FILTER_RESULTS, + "routing": {"serving_pipereplica": "d20260922011332-060efdec-prefill-r0-dp0-prefill"}, + "service_tier": "default", + "system_fingerprint": None, + "usage": { + "completion_tokens": 412, + "completion_tokens_details": { + "accepted_prediction_tokens": 0, + "audio_tokens": 0, + "reasoning_tokens": 0, + "rejected_prediction_tokens": 0, + }, + "latency_checkpoint": { + "engine_tbt_ms": 4, + "engine_ttft_ms": 131, + "engine_ttlt_ms": 157, + "pre_inference_ms": 55, + "service_tbt_ms": 5, + "service_ttft_ms": 1034, + "service_ttlt_ms": 1058, + "user_visible_ttft_ms": 979, + }, + "prompt_tokens": 1840, + "prompt_tokens_details": {"audio_tokens": 0, "cache_write_tokens": 0, "cached_tokens": 0}, + "total_tokens": 2252, + }, + }, + ), + expected=ExactExpected( + spend=0.000507, input_cost=0.0002392, output_cost=0.0002678, prompt_tokens=1840, completion_tokens=412 + ), +) + +AZURE_GPT_6_LUNA_STREAM: Final = CostTrackingTestCase( + name="azure-gpt-6-luna-stream", + covers="quota_management.spend_tracking.scripted_wire.logs_cost", + model="azure/gpt-6-luna", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "b22d4ad4da33 summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": True, + "stream_options": {"include_usage": True}, + "allowed_openai_params": [], + }, + response=SseResponse( + content_type="text/event-stream", + frames=sse_frames( + { + "choices": [], + "created": 0, + "id": "", + "model": "", + "object": "", + "prompt_filter_results": _GPT_6_PROMPT_FILTER_RESULTS, + }, + { + "choices": [ + { + "content_filter_results": {}, + "delta": {"content": "", "refusal": None, "role": "assistant"}, + "finish_reason": None, + "index": 0, + "logprobs": None, + } + ], + "created": 1791852023, + "id": "chatcmpl-$REQUEST_ID", + "model": "gpt-6-luna-2026-09-22", + "object": "chat.completion.chunk", + "service_tier": "default", + "system_fingerprint": None, + "usage": None, + }, + { + "choices": [ + { + "content_filter_results": {"protected_material_text": {"detected": False, "filtered": False}}, + "delta": {"content": "scripted answer b22d4ad4da33"}, + "finish_reason": None, + "index": 0, + "logprobs": None, + } + ], + "created": 1791852023, + "id": "chatcmpl-$REQUEST_ID", + "model": "gpt-6-luna-2026-09-22", + "object": "chat.completion.chunk", + "service_tier": "default", + "system_fingerprint": None, + "usage": None, + }, + { + "choices": [ + {"content_filter_results": {}, "delta": {}, "finish_reason": "stop", "index": 0, "logprobs": None} + ], + "created": 1791852023, + "id": "chatcmpl-$REQUEST_ID", + "model": "gpt-6-luna-2026-09-22", + "object": "chat.completion.chunk", + "service_tier": "default", + "system_fingerprint": None, + "usage": None, + }, + { + "choices": [], + "created": 1791852023, + "id": "chatcmpl-$REQUEST_ID", + "model": "gpt-6-luna-2026-09-22", + "object": "chat.completion.chunk", + "service_tier": "default", + "system_fingerprint": None, + "routing": {"serving_pipereplica": "d20260925225746-6787865d-prefill-r11-dp0-prefill"}, + "usage": { + "completion_tokens": 412, + "completion_tokens_details": { + "accepted_prediction_tokens": 0, + "audio_tokens": 0, + "reasoning_tokens": 0, + "rejected_prediction_tokens": 0, + }, + "latency_checkpoint": { + "engine_tbt_ms": 4, + "engine_ttft_ms": 131, + "engine_ttlt_ms": 157, + "pre_inference_ms": 55, + "service_tbt_ms": 5, + "service_ttft_ms": 1034, + "service_ttlt_ms": 1058, + "user_visible_ttft_ms": 979, + }, + "prompt_tokens": 1840, + "prompt_tokens_details": {"audio_tokens": 0, "cache_write_tokens": 0, "cached_tokens": 0}, + "total_tokens": 2252, + }, + }, + ), + ), + expected=ExactExpected( + spend=0.000507, input_cost=0.0002392, output_cost=0.0002678, prompt_tokens=1840, completion_tokens=412 + ), +) + +AZURE_GPT_6_LUNA_SERVICE_TIER_FLEX: Final = CostTrackingTestCase( + name="azure-gpt-6-luna-service_tier_flex", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="azure/gpt-6-luna", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "895f84528ea4 summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": False, + "service_tier": "flex", + "allowed_openai_params": [], + }, + response=JsonResponse( + content_type="application/json", + body={ + "choices": [ + { + "content_filter_results": _GPT_6_CONTENT_FILTER_RESULTS, + "finish_reason": "stop", + "index": 0, + "logprobs": None, + "message": { + "annotations": [], + "content": "scripted answer 895f84528ea4", + "refusal": None, + "role": "assistant", + }, + } + ], + "created": 1791852029, + "id": "chatcmpl-$REQUEST_ID", + "model": "gpt-6-luna-2026-09-22", + "object": "chat.completion", + "prompt_filter_results": _GPT_6_PROMPT_FILTER_RESULTS, + "routing": {"serving_pipereplica": "d20260922011332-060efdec-prefill-r0-dp0-prefill"}, + "service_tier": "flex", + "system_fingerprint": None, + "usage": { + "completion_tokens": 412, + "completion_tokens_details": { + "accepted_prediction_tokens": 0, + "audio_tokens": 0, + "reasoning_tokens": 0, + "rejected_prediction_tokens": 0, + }, + "latency_checkpoint": { + "engine_tbt_ms": 4, + "engine_ttft_ms": 131, + "engine_ttlt_ms": 157, + "pre_inference_ms": 55, + "service_tbt_ms": 5, + "service_ttft_ms": 1034, + "service_ttlt_ms": 1058, + "user_visible_ttft_ms": 979, + }, + "prompt_tokens": 1840, + "prompt_tokens_details": {"audio_tokens": 0, "cache_write_tokens": 0, "cached_tokens": 0}, + "total_tokens": 2252, + }, + }, + ), + expected=ExactExpected( + spend=0.0002535, input_cost=0.0001196, output_cost=0.0001339, prompt_tokens=1840, completion_tokens=412 + ), +) + +AZURE_GPT_6_LUNA_SERVICE_TIER_PRIORITY: Final = CostTrackingTestCase( + name="azure-gpt-6-luna-service_tier_priority", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="azure/gpt-6-luna", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "705c026752fc summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": False, + "service_tier": "priority", + "allowed_openai_params": [], + }, + response=JsonResponse( + content_type="application/json", + body={ + "choices": [ + { + "content_filter_results": _GPT_6_CONTENT_FILTER_RESULTS, + "finish_reason": "stop", + "index": 0, + "logprobs": None, + "message": { + "annotations": [], + "content": "scripted answer 705c026752fc", + "refusal": None, + "role": "assistant", + }, + } + ], + "created": 1791852029, + "id": "chatcmpl-$REQUEST_ID", + "model": "gpt-6-luna-2026-09-22", + "object": "chat.completion", + "prompt_filter_results": _GPT_6_PROMPT_FILTER_RESULTS, + "routing": {"serving_pipereplica": "d20260922011332-060efdec-prefill-r0-dp0-prefill"}, + "service_tier": "default", + "system_fingerprint": None, + "usage": { + "completion_tokens": 412, + "completion_tokens_details": { + "accepted_prediction_tokens": 0, + "audio_tokens": 0, + "reasoning_tokens": 0, + "rejected_prediction_tokens": 0, + }, + "latency_checkpoint": { + "engine_tbt_ms": 4, + "engine_ttft_ms": 131, + "engine_ttlt_ms": 157, + "pre_inference_ms": 55, + "service_tbt_ms": 5, + "service_ttft_ms": 1034, + "service_ttlt_ms": 1058, + "user_visible_ttft_ms": 979, + }, + "prompt_tokens": 1840, + "prompt_tokens_details": {"audio_tokens": 0, "cache_write_tokens": 0, "cached_tokens": 0}, + "total_tokens": 2252, + }, + }, + ), + expected=ExactExpected( + spend=0.000507, input_cost=0.0002392, output_cost=0.0002678, prompt_tokens=1840, completion_tokens=412 + ), +) + +AZURE_GPT_6_ASTRA_INPUT_TEXT: Final = CostTrackingTestCase( + name="azure-gpt-6-astra-input_text", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="azure/gpt-6-astra", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "41ccdf63f9d7 summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": False, + "allowed_openai_params": [], + }, + response=JsonResponse( + content_type="application/json", + body={ + "choices": [ + { + "content_filter_results": _GPT_6_CONTENT_FILTER_RESULTS, + "finish_reason": "stop", + "index": 0, + "logprobs": None, + "message": { + "annotations": [], + "content": "scripted answer 41ccdf63f9d7", + "refusal": None, + "role": "assistant", + }, + } + ], + "created": 1791852029, + "id": "chatcmpl-$REQUEST_ID", + "model": "gpt-6-astra-2026-09-03", + "object": "chat.completion", + "prompt_filter_results": _GPT_6_PROMPT_FILTER_RESULTS, + "routing": {"serving_pipereplica": "d20260922011332-060efdec-prefill-r0-dp0-prefill"}, + "service_tier": "default", + "system_fingerprint": None, + "usage": { + "completion_tokens": 412, + "completion_tokens_details": { + "accepted_prediction_tokens": 0, + "audio_tokens": 0, + "reasoning_tokens": 0, + "rejected_prediction_tokens": 0, + }, + "latency_checkpoint": { + "engine_tbt_ms": 4, + "engine_ttft_ms": 131, + "engine_ttlt_ms": 157, + "pre_inference_ms": 55, + "service_tbt_ms": 5, + "service_ttft_ms": 1034, + "service_ttlt_ms": 1058, + "user_visible_ttft_ms": 979, + }, + "prompt_tokens": 1840, + "prompt_tokens_details": {"audio_tokens": 0, "cache_write_tokens": 0, "cached_tokens": 0}, + "total_tokens": 2252, + }, + }, + ), + expected=ExactExpected( + spend=0.04875, input_cost=0.023, output_cost=0.02575, prompt_tokens=1840, completion_tokens=412 + ), +) + +AZURE_GPT_6_ASTRA_STREAM: Final = CostTrackingTestCase( + name="azure-gpt-6-astra-stream", + covers="quota_management.spend_tracking.scripted_wire.logs_cost", + model="azure/gpt-6-astra", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "480ff10f91fb summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": True, + "stream_options": {"include_usage": True}, + "allowed_openai_params": [], + }, + response=SseResponse( + content_type="text/event-stream", + frames=sse_frames( + { + "choices": [], + "created": 0, + "id": "", + "model": "", + "object": "", + "prompt_filter_results": _GPT_6_PROMPT_FILTER_RESULTS, + }, + { + "choices": [ + { + "content_filter_results": {}, + "delta": {"content": "", "refusal": None, "role": "assistant"}, + "finish_reason": None, + "index": 0, + "logprobs": None, + } + ], + "created": 1791852023, + "id": "chatcmpl-$REQUEST_ID", + "model": "gpt-6-astra-2026-09-03", + "object": "chat.completion.chunk", + "service_tier": "default", + "system_fingerprint": None, + "usage": None, + }, + { + "choices": [ + { + "content_filter_results": {"protected_material_text": {"detected": False, "filtered": False}}, + "delta": {"content": "scripted answer 480ff10f91fb"}, + "finish_reason": None, + "index": 0, + "logprobs": None, + } + ], + "created": 1791852023, + "id": "chatcmpl-$REQUEST_ID", + "model": "gpt-6-astra-2026-09-03", + "object": "chat.completion.chunk", + "service_tier": "default", + "system_fingerprint": None, + "usage": None, + }, + { + "choices": [ + {"content_filter_results": {}, "delta": {}, "finish_reason": "stop", "index": 0, "logprobs": None} + ], + "created": 1791852023, + "id": "chatcmpl-$REQUEST_ID", + "model": "gpt-6-astra-2026-09-03", + "object": "chat.completion.chunk", + "service_tier": "default", + "system_fingerprint": None, + "usage": None, + }, + { + "choices": [], + "created": 1791852023, + "id": "chatcmpl-$REQUEST_ID", + "model": "gpt-6-astra-2026-09-03", + "object": "chat.completion.chunk", + "service_tier": "default", + "system_fingerprint": None, + "routing": {"serving_pipereplica": "d20260925225746-6787865d-prefill-r11-dp0-prefill"}, + "usage": { + "completion_tokens": 412, + "completion_tokens_details": { + "accepted_prediction_tokens": 0, + "audio_tokens": 0, + "reasoning_tokens": 0, + "rejected_prediction_tokens": 0, + }, + "latency_checkpoint": { + "engine_tbt_ms": 4, + "engine_ttft_ms": 131, + "engine_ttlt_ms": 157, + "pre_inference_ms": 55, + "service_tbt_ms": 5, + "service_ttft_ms": 1034, + "service_ttlt_ms": 1058, + "user_visible_ttft_ms": 979, + }, + "prompt_tokens": 1840, + "prompt_tokens_details": {"audio_tokens": 0, "cache_write_tokens": 0, "cached_tokens": 0}, + "total_tokens": 2252, + }, + }, + ), + ), + expected=ExactExpected( + spend=0.04875, input_cost=0.023, output_cost=0.02575, prompt_tokens=1840, completion_tokens=412 + ), +) + +AZURE_GPT_6_ASTRA_SERVICE_TIER_FLEX: Final = CostTrackingTestCase( + name="azure-gpt-6-astra-service_tier_flex", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="azure/gpt-6-astra", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "6ff9027e3f7f summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": False, + "service_tier": "flex", + "allowed_openai_params": [], + }, + response=JsonResponse( + content_type="application/json", + body={ + "choices": [ + { + "content_filter_results": _GPT_6_CONTENT_FILTER_RESULTS, + "finish_reason": "stop", + "index": 0, + "logprobs": None, + "message": { + "annotations": [], + "content": "scripted answer 6ff9027e3f7f", + "refusal": None, + "role": "assistant", + }, + } + ], + "created": 1791852029, + "id": "chatcmpl-$REQUEST_ID", + "model": "gpt-6-astra-2026-09-03", + "object": "chat.completion", + "prompt_filter_results": _GPT_6_PROMPT_FILTER_RESULTS, + "routing": {"serving_pipereplica": "d20260922011332-060efdec-prefill-r0-dp0-prefill"}, + "service_tier": "flex", + "system_fingerprint": None, + "usage": { + "completion_tokens": 412, + "completion_tokens_details": { + "accepted_prediction_tokens": 0, + "audio_tokens": 0, + "reasoning_tokens": 0, + "rejected_prediction_tokens": 0, + }, + "latency_checkpoint": { + "engine_tbt_ms": 4, + "engine_ttft_ms": 131, + "engine_ttlt_ms": 157, + "pre_inference_ms": 55, + "service_tbt_ms": 5, + "service_ttft_ms": 1034, + "service_ttlt_ms": 1058, + "user_visible_ttft_ms": 979, + }, + "prompt_tokens": 1840, + "prompt_tokens_details": {"audio_tokens": 0, "cache_write_tokens": 0, "cached_tokens": 0}, + "total_tokens": 2252, + }, + }, + ), + expected=ExactExpected( + spend=0.024375, input_cost=0.0115, output_cost=0.012875, prompt_tokens=1840, completion_tokens=412 + ), +) + +AZURE_GPT_6_ASTRA_SERVICE_TIER_PRIORITY: Final = CostTrackingTestCase( + name="azure-gpt-6-astra-service_tier_priority", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="azure/gpt-6-astra", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "9a0c038a2ad8 summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": False, + "service_tier": "priority", + "allowed_openai_params": [], + }, + response=JsonResponse( + content_type="application/json", + body={ + "choices": [ + { + "content_filter_results": _GPT_6_CONTENT_FILTER_RESULTS, + "finish_reason": "stop", + "index": 0, + "logprobs": None, + "message": { + "annotations": [], + "content": "scripted answer 9a0c038a2ad8", + "refusal": None, + "role": "assistant", + }, + } + ], + "created": 1791852029, + "id": "chatcmpl-$REQUEST_ID", + "model": "gpt-6-astra-2026-09-03", + "object": "chat.completion", + "prompt_filter_results": _GPT_6_PROMPT_FILTER_RESULTS, + "routing": {"serving_pipereplica": "d20260922011332-060efdec-prefill-r0-dp0-prefill"}, + "service_tier": "default", + "system_fingerprint": None, + "usage": { + "completion_tokens": 412, + "completion_tokens_details": { + "accepted_prediction_tokens": 0, + "audio_tokens": 0, + "reasoning_tokens": 0, + "rejected_prediction_tokens": 0, + }, + "latency_checkpoint": { + "engine_tbt_ms": 4, + "engine_ttft_ms": 131, + "engine_ttlt_ms": 157, + "pre_inference_ms": 55, + "service_tbt_ms": 5, + "service_ttft_ms": 1034, + "service_ttlt_ms": 1058, + "user_visible_ttft_ms": 979, + }, + "prompt_tokens": 1840, + "prompt_tokens_details": {"audio_tokens": 0, "cache_write_tokens": 0, "cached_tokens": 0}, + "total_tokens": 2252, + }, + }, + ), + expected=ExactExpected( + spend=0.04875, input_cost=0.023, output_cost=0.02575, prompt_tokens=1840, completion_tokens=412 + ), +) + +AZURE_GPT_6_1_SOL_INPUT_TEXT: Final = CostTrackingTestCase( + name="azure-gpt-6.1-sol-input_text", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="azure/gpt-6.1-sol", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "89ad6904966e summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": False, + "allowed_openai_params": [], + }, + response=JsonResponse( + content_type="application/json", + body={ + "choices": [ + { + "content_filter_results": _GPT_6_CONTENT_FILTER_RESULTS, + "finish_reason": "stop", + "index": 0, + "logprobs": None, + "message": { + "annotations": [], + "content": "scripted answer 89ad6904966e", + "refusal": None, + "role": "assistant", + }, + } + ], + "created": 1791852029, + "id": "chatcmpl-$REQUEST_ID", + "model": "gpt-6.1-sol-2026-09-29", + "object": "chat.completion", + "prompt_filter_results": _GPT_6_PROMPT_FILTER_RESULTS, + "routing": {"serving_pipereplica": "d20260922011332-060efdec-prefill-r0-dp0-prefill"}, + "service_tier": "default", + "system_fingerprint": None, + "usage": { + "completion_tokens": 412, + "completion_tokens_details": { + "accepted_prediction_tokens": 0, + "audio_tokens": 0, + "reasoning_tokens": 0, + "rejected_prediction_tokens": 0, + }, + "latency_checkpoint": { + "engine_tbt_ms": 4, + "engine_ttft_ms": 131, + "engine_ttlt_ms": 157, + "pre_inference_ms": 55, + "service_tbt_ms": 5, + "service_ttft_ms": 1034, + "service_ttlt_ms": 1058, + "user_visible_ttft_ms": 979, + }, + "prompt_tokens": 1840, + "prompt_tokens_details": {"audio_tokens": 0, "cache_write_tokens": 0, "cached_tokens": 0}, + "total_tokens": 2252, + }, + }, + ), + expected=ExactExpected( + spend=0.00936, input_cost=0.004416, output_cost=0.004944, prompt_tokens=1840, completion_tokens=412 + ), +) + +AZURE_GPT_6_1_SOL_STREAM: Final = CostTrackingTestCase( + name="azure-gpt-6.1-sol-stream", + covers="quota_management.spend_tracking.scripted_wire.logs_cost", + model="azure/gpt-6.1-sol", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "493cd994b1de summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": True, + "stream_options": {"include_usage": True}, + "allowed_openai_params": [], + }, + response=SseResponse( + content_type="text/event-stream", + frames=sse_frames( + { + "choices": [], + "created": 0, + "id": "", + "model": "", + "object": "", + "prompt_filter_results": _GPT_6_PROMPT_FILTER_RESULTS, + }, + { + "choices": [ + { + "content_filter_results": {}, + "delta": {"content": "", "refusal": None, "role": "assistant"}, + "finish_reason": None, + "index": 0, + "logprobs": None, + } + ], + "created": 1791852023, + "id": "chatcmpl-$REQUEST_ID", + "model": "gpt-6.1-sol-2026-09-29", + "object": "chat.completion.chunk", + "service_tier": "default", + "system_fingerprint": None, + "usage": None, + }, + { + "choices": [ + { + "content_filter_results": {"protected_material_text": {"detected": False, "filtered": False}}, + "delta": {"content": "scripted answer 493cd994b1de"}, + "finish_reason": None, + "index": 0, + "logprobs": None, + } + ], + "created": 1791852023, + "id": "chatcmpl-$REQUEST_ID", + "model": "gpt-6.1-sol-2026-09-29", + "object": "chat.completion.chunk", + "service_tier": "default", + "system_fingerprint": None, + "usage": None, + }, + { + "choices": [ + {"content_filter_results": {}, "delta": {}, "finish_reason": "stop", "index": 0, "logprobs": None} + ], + "created": 1791852023, + "id": "chatcmpl-$REQUEST_ID", + "model": "gpt-6.1-sol-2026-09-29", + "object": "chat.completion.chunk", + "service_tier": "default", + "system_fingerprint": None, + "usage": None, + }, + { + "choices": [], + "created": 1791852023, + "id": "chatcmpl-$REQUEST_ID", + "model": "gpt-6.1-sol-2026-09-29", + "object": "chat.completion.chunk", + "service_tier": "default", + "system_fingerprint": None, + "routing": {"serving_pipereplica": "d20260925225746-6787865d-prefill-r11-dp0-prefill"}, + "usage": { + "completion_tokens": 412, + "completion_tokens_details": { + "accepted_prediction_tokens": 0, + "audio_tokens": 0, + "reasoning_tokens": 0, + "rejected_prediction_tokens": 0, + }, + "latency_checkpoint": { + "engine_tbt_ms": 4, + "engine_ttft_ms": 131, + "engine_ttlt_ms": 157, + "pre_inference_ms": 55, + "service_tbt_ms": 5, + "service_ttft_ms": 1034, + "service_ttlt_ms": 1058, + "user_visible_ttft_ms": 979, + }, + "prompt_tokens": 1840, + "prompt_tokens_details": {"audio_tokens": 0, "cache_write_tokens": 0, "cached_tokens": 0}, + "total_tokens": 2252, + }, + }, + ), + ), + expected=ExactExpected( + spend=0.00936, input_cost=0.004416, output_cost=0.004944, prompt_tokens=1840, completion_tokens=412 + ), +) + +AZURE_GPT_6_1_SOL_SERVICE_TIER_FLEX: Final = CostTrackingTestCase( + name="azure-gpt-6.1-sol-service_tier_flex", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="azure/gpt-6.1-sol", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "960e0935140f summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": False, + "service_tier": "flex", + "allowed_openai_params": [], + }, + response=JsonResponse( + content_type="application/json", + body={ + "choices": [ + { + "content_filter_results": _GPT_6_CONTENT_FILTER_RESULTS, + "finish_reason": "stop", + "index": 0, + "logprobs": None, + "message": { + "annotations": [], + "content": "scripted answer 960e0935140f", + "refusal": None, + "role": "assistant", + }, + } + ], + "created": 1791852029, + "id": "chatcmpl-$REQUEST_ID", + "model": "gpt-6.1-sol-2026-09-29", + "object": "chat.completion", + "prompt_filter_results": _GPT_6_PROMPT_FILTER_RESULTS, + "routing": {"serving_pipereplica": "d20260922011332-060efdec-prefill-r0-dp0-prefill"}, + "service_tier": "flex", + "system_fingerprint": None, + "usage": { + "completion_tokens": 412, + "completion_tokens_details": { + "accepted_prediction_tokens": 0, + "audio_tokens": 0, + "reasoning_tokens": 0, + "rejected_prediction_tokens": 0, + }, + "latency_checkpoint": { + "engine_tbt_ms": 4, + "engine_ttft_ms": 131, + "engine_ttlt_ms": 157, + "pre_inference_ms": 55, + "service_tbt_ms": 5, + "service_ttft_ms": 1034, + "service_ttlt_ms": 1058, + "user_visible_ttft_ms": 979, + }, + "prompt_tokens": 1840, + "prompt_tokens_details": {"audio_tokens": 0, "cache_write_tokens": 0, "cached_tokens": 0}, + "total_tokens": 2252, + }, + }, + ), + expected=ExactExpected( + spend=0.00468, input_cost=0.002208, output_cost=0.002472, prompt_tokens=1840, completion_tokens=412 + ), +) + +AZURE_GPT_6_1_SOL_SERVICE_TIER_PRIORITY: Final = CostTrackingTestCase( + name="azure-gpt-6.1-sol-service_tier_priority", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="azure/gpt-6.1-sol", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "6b36e2dc1f33 summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": False, + "service_tier": "priority", + "allowed_openai_params": [], + }, + response=JsonResponse( + content_type="application/json", + body={ + "choices": [ + { + "content_filter_results": _GPT_6_CONTENT_FILTER_RESULTS, + "finish_reason": "stop", + "index": 0, + "logprobs": None, + "message": { + "annotations": [], + "content": "scripted answer 6b36e2dc1f33", + "refusal": None, + "role": "assistant", + }, + } + ], + "created": 1791852029, + "id": "chatcmpl-$REQUEST_ID", + "model": "gpt-6.1-sol-2026-09-29", + "object": "chat.completion", + "prompt_filter_results": _GPT_6_PROMPT_FILTER_RESULTS, + "routing": {"serving_pipereplica": "d20260922011332-060efdec-prefill-r0-dp0-prefill"}, + "service_tier": "default", + "system_fingerprint": None, + "usage": { + "completion_tokens": 412, + "completion_tokens_details": { + "accepted_prediction_tokens": 0, + "audio_tokens": 0, + "reasoning_tokens": 0, + "rejected_prediction_tokens": 0, + }, + "latency_checkpoint": { + "engine_tbt_ms": 4, + "engine_ttft_ms": 131, + "engine_ttlt_ms": 157, + "pre_inference_ms": 55, + "service_tbt_ms": 5, + "service_ttft_ms": 1034, + "service_ttlt_ms": 1058, + "user_visible_ttft_ms": 979, + }, + "prompt_tokens": 1840, + "prompt_tokens_details": {"audio_tokens": 0, "cache_write_tokens": 0, "cached_tokens": 0}, + "total_tokens": 2252, + }, + }, + ), + expected=ExactExpected( + spend=0.00936, input_cost=0.004416, output_cost=0.004944, prompt_tokens=1840, completion_tokens=412 + ), +) + +AZURE_GPT_6_SOL_REASONING: Final = CostTrackingTestCase( + name="azure/gpt-6-sol-reasoning", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="azure/gpt-6-sol", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "5c0e79234ed5 summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": False, + "reasoning_effort": "medium", + "allowed_openai_params": ["reasoning_effort"], + }, + response=JsonResponse( + content_type="application/json", + body={ + "choices": [ + { + "content_filter_results": _GPT_6_CONTENT_FILTER_RESULTS, + "finish_reason": "stop", + "index": 0, + "logprobs": None, + "message": { + "annotations": [], + "content": "scripted answer 5c0e79234ed5", + "refusal": None, + "role": "assistant", + }, + } + ], + "created": 1791855585, + "id": "chatcmpl-$REQUEST_ID", + "model": "gpt-6-sol-2026-09-22", + "object": "chat.completion", + "prompt_filter_results": _GPT_6_PROMPT_FILTER_RESULTS, + "routing": {"serving_pipereplica": "d20260925225746-6787865d-prefill-r15-dp0-prefill"}, + "service_tier": "default", + "system_fingerprint": None, + "usage": { + "completion_tokens": 4040, + "completion_tokens_details": { + "accepted_prediction_tokens": 0, + "audio_tokens": 0, + "reasoning_tokens": 3480, + "rejected_prediction_tokens": 0, + }, + "latency_checkpoint": { + "engine_tbt_ms": 6, + "engine_ttft_ms": 152, + "engine_ttlt_ms": 305, + "pre_inference_ms": 102, + "service_tbt_ms": 9, + "service_ttft_ms": 389, + "service_ttlt_ms": 608, + "user_visible_ttft_ms": 288, + }, + "prompt_tokens": 1240, + "prompt_tokens_details": {"audio_tokens": 0, "cache_write_tokens": 0, "cached_tokens": 0}, + "total_tokens": 5280, + }, + }, + ), + expected=ExactExpected( + spend=0.049312, input_cost=0.002852, output_cost=0.04646, prompt_tokens=1240, completion_tokens=4040 + ), +) + +AZURE_GPT_6_SOL_STREAM_NO_USAGE: Final = CostTrackingTestCase( + name="azure/gpt-6-sol-stream_no_usage", + covers="quota_management.spend_tracking.scripted_wire.logs_cost", + model="azure/gpt-6-sol", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "29fbe99c40db summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": True, + "allowed_openai_params": [], + }, + response=SseResponse( + content_type="text/event-stream", + frames=sse_frames( + { + "choices": [], + "created": 0, + "id": "", + "model": "", + "object": "", + "prompt_filter_results": _GPT_6_PROMPT_FILTER_RESULTS, + }, + { + "choices": [ + { + "content_filter_results": {}, + "delta": {"content": "", "refusal": None, "role": "assistant"}, + "finish_reason": None, + "index": 0, + "logprobs": None, + } + ], + "created": 1791855587, + "id": "chatcmpl-$REQUEST_ID", + "model": "gpt-6-sol-2026-09-22", + "object": "chat.completion.chunk", + "service_tier": "default", + "system_fingerprint": None, + "obfuscation": "5A3jk", + }, + { + "choices": [ + { + "content_filter_results": {"protected_material_text": {"detected": False, "filtered": False}}, + "delta": {"content": "scripted answer 29fbe99c40db"}, + "finish_reason": None, + "index": 0, + "logprobs": None, + } + ], + "created": 1791855587, + "id": "chatcmpl-$REQUEST_ID", + "model": "gpt-6-sol-2026-09-22", + "object": "chat.completion.chunk", + "service_tier": "default", + "system_fingerprint": None, + "obfuscation": "Q2x", + }, + { + "choices": [ + {"content_filter_results": {}, "delta": {}, "finish_reason": "stop", "index": 0, "logprobs": None} + ], + "created": 1791855587, + "id": "chatcmpl-$REQUEST_ID", + "model": "gpt-6-sol-2026-09-22", + "object": "chat.completion.chunk", + "service_tier": "default", + "system_fingerprint": None, + "obfuscation": "Uz14j7Ajk4CYaTB", + }, + ), + ), + expected=RecountExpected(recount=RecountRates(input_cost_per_token=2.3e-06, output_cost_per_token=1.15e-05)), +) + +AZURE_GPT_6_LUNA_REASONING: Final = CostTrackingTestCase( + name="azure/gpt-6-luna-reasoning", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="azure/gpt-6-luna", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "21c501e3e521 summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": False, + "reasoning_effort": "medium", + "allowed_openai_params": ["reasoning_effort"], + }, + response=JsonResponse( + content_type="application/json", + body={ + "choices": [ + { + "content_filter_results": _GPT_6_CONTENT_FILTER_RESULTS, + "finish_reason": "stop", + "index": 0, + "logprobs": None, + "message": { + "annotations": [], + "content": "scripted answer 21c501e3e521", + "refusal": None, + "role": "assistant", + }, + } + ], + "created": 1791855585, + "id": "chatcmpl-$REQUEST_ID", + "model": "gpt-6-luna-2026-09-22", + "object": "chat.completion", + "prompt_filter_results": _GPT_6_PROMPT_FILTER_RESULTS, + "routing": {"serving_pipereplica": "d20260925225746-6787865d-prefill-r15-dp0-prefill"}, + "service_tier": "default", + "system_fingerprint": None, + "usage": { + "completion_tokens": 4040, + "completion_tokens_details": { + "accepted_prediction_tokens": 0, + "audio_tokens": 0, + "reasoning_tokens": 3480, + "rejected_prediction_tokens": 0, + }, + "latency_checkpoint": { + "engine_tbt_ms": 6, + "engine_ttft_ms": 152, + "engine_ttlt_ms": 305, + "pre_inference_ms": 102, + "service_tbt_ms": 9, + "service_ttft_ms": 389, + "service_ttlt_ms": 608, + "user_visible_ttft_ms": 288, + }, + "prompt_tokens": 1240, + "prompt_tokens_details": {"audio_tokens": 0, "cache_write_tokens": 0, "cached_tokens": 0}, + "total_tokens": 5280, + }, + }, + ), + expected=ExactExpected( + spend=0.0027872, input_cost=0.0001612, output_cost=0.002626, prompt_tokens=1240, completion_tokens=4040 + ), +) + +AZURE_GPT_6_LUNA_STREAM_NO_USAGE: Final = CostTrackingTestCase( + name="azure/gpt-6-luna-stream_no_usage", + covers="quota_management.spend_tracking.scripted_wire.logs_cost", + model="azure/gpt-6-luna", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "fef24d05a468 summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": True, + "allowed_openai_params": [], + }, + response=SseResponse( + content_type="text/event-stream", + frames=sse_frames( + { + "choices": [], + "created": 0, + "id": "", + "model": "", + "object": "", + "prompt_filter_results": _GPT_6_PROMPT_FILTER_RESULTS, + }, + { + "choices": [ + { + "content_filter_results": {}, + "delta": {"content": "", "refusal": None, "role": "assistant"}, + "finish_reason": None, + "index": 0, + "logprobs": None, + } + ], + "created": 1791855587, + "id": "chatcmpl-$REQUEST_ID", + "model": "gpt-6-luna-2026-09-22", + "object": "chat.completion.chunk", + "service_tier": "default", + "system_fingerprint": None, + "obfuscation": "5A3jk", + }, + { + "choices": [ + { + "content_filter_results": {"protected_material_text": {"detected": False, "filtered": False}}, + "delta": {"content": "scripted answer fef24d05a468"}, + "finish_reason": None, + "index": 0, + "logprobs": None, + } + ], + "created": 1791855587, + "id": "chatcmpl-$REQUEST_ID", + "model": "gpt-6-luna-2026-09-22", + "object": "chat.completion.chunk", + "service_tier": "default", + "system_fingerprint": None, + "obfuscation": "Q2x", + }, + { + "choices": [ + {"content_filter_results": {}, "delta": {}, "finish_reason": "stop", "index": 0, "logprobs": None} + ], + "created": 1791855587, + "id": "chatcmpl-$REQUEST_ID", + "model": "gpt-6-luna-2026-09-22", + "object": "chat.completion.chunk", + "service_tier": "default", + "system_fingerprint": None, + "obfuscation": "Uz14j7Ajk4CYaTB", + }, + ), + ), + expected=RecountExpected(recount=RecountRates(input_cost_per_token=1.3e-07, output_cost_per_token=6.5e-07)), +) + +AZURE_GPT_6_ASTRA_REASONING: Final = CostTrackingTestCase( + name="azure/gpt-6-astra-reasoning", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="azure/gpt-6-astra", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "d7f4aaa662f3 summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": False, + "reasoning_effort": "medium", + "allowed_openai_params": ["reasoning_effort"], + }, + response=JsonResponse( + content_type="application/json", + body={ + "choices": [ + { + "content_filter_results": _GPT_6_CONTENT_FILTER_RESULTS, + "finish_reason": "stop", + "index": 0, + "logprobs": None, + "message": { + "annotations": [], + "content": "scripted answer d7f4aaa662f3", + "refusal": None, + "role": "assistant", + }, + } + ], + "created": 1791855585, + "id": "chatcmpl-$REQUEST_ID", + "model": "gpt-6-astra-2026-09-03", + "object": "chat.completion", + "prompt_filter_results": _GPT_6_PROMPT_FILTER_RESULTS, + "routing": {"serving_pipereplica": "d20260925225746-6787865d-prefill-r15-dp0-prefill"}, + "service_tier": "default", + "system_fingerprint": None, + "usage": { + "completion_tokens": 4040, + "completion_tokens_details": { + "accepted_prediction_tokens": 0, + "audio_tokens": 0, + "reasoning_tokens": 3480, + "rejected_prediction_tokens": 0, + }, + "latency_checkpoint": { + "engine_tbt_ms": 6, + "engine_ttft_ms": 152, + "engine_ttlt_ms": 305, + "pre_inference_ms": 102, + "service_tbt_ms": 9, + "service_ttft_ms": 389, + "service_ttlt_ms": 608, + "user_visible_ttft_ms": 288, + }, + "prompt_tokens": 1240, + "prompt_tokens_details": {"audio_tokens": 0, "cache_write_tokens": 0, "cached_tokens": 0}, + "total_tokens": 5280, + }, + }, + ), + expected=ExactExpected( + spend=0.268, input_cost=0.0155, output_cost=0.2525, prompt_tokens=1240, completion_tokens=4040 + ), +) + +AZURE_GPT_6_ASTRA_STREAM_NO_USAGE: Final = CostTrackingTestCase( + name="azure/gpt-6-astra-stream_no_usage", + covers="quota_management.spend_tracking.scripted_wire.logs_cost", + model="azure/gpt-6-astra", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "65e182029270 summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": True, + "allowed_openai_params": [], + }, + response=SseResponse( + content_type="text/event-stream", + frames=sse_frames( + { + "choices": [], + "created": 0, + "id": "", + "model": "", + "object": "", + "prompt_filter_results": _GPT_6_PROMPT_FILTER_RESULTS, + }, + { + "choices": [ + { + "content_filter_results": {}, + "delta": {"content": "", "refusal": None, "role": "assistant"}, + "finish_reason": None, + "index": 0, + "logprobs": None, + } + ], + "created": 1791855587, + "id": "chatcmpl-$REQUEST_ID", + "model": "gpt-6-astra-2026-09-03", + "object": "chat.completion.chunk", + "service_tier": "default", + "system_fingerprint": None, + "obfuscation": "5A3jk", + }, + { + "choices": [ + { + "content_filter_results": {"protected_material_text": {"detected": False, "filtered": False}}, + "delta": {"content": "scripted answer 65e182029270"}, + "finish_reason": None, + "index": 0, + "logprobs": None, + } + ], + "created": 1791855587, + "id": "chatcmpl-$REQUEST_ID", + "model": "gpt-6-astra-2026-09-03", + "object": "chat.completion.chunk", + "service_tier": "default", + "system_fingerprint": None, + "obfuscation": "Q2x", + }, + { + "choices": [ + {"content_filter_results": {}, "delta": {}, "finish_reason": "stop", "index": 0, "logprobs": None} + ], + "created": 1791855587, + "id": "chatcmpl-$REQUEST_ID", + "model": "gpt-6-astra-2026-09-03", + "object": "chat.completion.chunk", + "service_tier": "default", + "system_fingerprint": None, + "obfuscation": "Uz14j7Ajk4CYaTB", + }, + ), + ), + expected=RecountExpected(recount=RecountRates(input_cost_per_token=1.25e-05, output_cost_per_token=6.25e-05)), +) + +AZURE_GPT_6_1_SOL_REASONING: Final = CostTrackingTestCase( + name="azure/gpt-6.1-sol-reasoning", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="azure/gpt-6.1-sol", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "8ac8b47d17d4 summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": False, + "reasoning_effort": "medium", + "allowed_openai_params": ["reasoning_effort"], + }, + response=JsonResponse( + content_type="application/json", + body={ + "choices": [ + { + "content_filter_results": _GPT_6_CONTENT_FILTER_RESULTS, + "finish_reason": "stop", + "index": 0, + "logprobs": None, + "message": { + "annotations": [], + "content": "scripted answer 8ac8b47d17d4", + "refusal": None, + "role": "assistant", + }, + } + ], + "created": 1791855585, + "id": "chatcmpl-$REQUEST_ID", + "model": "gpt-6.1-sol-2026-09-29", + "object": "chat.completion", + "prompt_filter_results": _GPT_6_PROMPT_FILTER_RESULTS, + "routing": {"serving_pipereplica": "d20260925225746-6787865d-prefill-r15-dp0-prefill"}, + "service_tier": "default", + "system_fingerprint": None, + "usage": { + "completion_tokens": 4040, + "completion_tokens_details": { + "accepted_prediction_tokens": 0, + "audio_tokens": 0, + "reasoning_tokens": 3480, + "rejected_prediction_tokens": 0, + }, + "latency_checkpoint": { + "engine_tbt_ms": 6, + "engine_ttft_ms": 152, + "engine_ttlt_ms": 305, + "pre_inference_ms": 102, + "service_tbt_ms": 9, + "service_ttft_ms": 389, + "service_ttlt_ms": 608, + "user_visible_ttft_ms": 288, + }, + "prompt_tokens": 1240, + "prompt_tokens_details": {"audio_tokens": 0, "cache_write_tokens": 0, "cached_tokens": 0}, + "total_tokens": 5280, + }, + }, + ), + expected=ExactExpected( + spend=0.051456, input_cost=0.002976, output_cost=0.04848, prompt_tokens=1240, completion_tokens=4040 + ), +) + +AZURE_GPT_6_1_SOL_STREAM_NO_USAGE: Final = CostTrackingTestCase( + name="azure/gpt-6.1-sol-stream_no_usage", + covers="quota_management.spend_tracking.scripted_wire.logs_cost", + model="azure/gpt-6.1-sol", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "20e483d62a14 summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": True, + "allowed_openai_params": [], + }, + response=SseResponse( + content_type="text/event-stream", + frames=sse_frames( + { + "choices": [], + "created": 0, + "id": "", + "model": "", + "object": "", + "prompt_filter_results": _GPT_6_PROMPT_FILTER_RESULTS, + }, + { + "choices": [ + { + "content_filter_results": {}, + "delta": {"content": "", "refusal": None, "role": "assistant"}, + "finish_reason": None, + "index": 0, + "logprobs": None, + } + ], + "created": 1791855587, + "id": "chatcmpl-$REQUEST_ID", + "model": "gpt-6.1-sol-2026-09-29", + "object": "chat.completion.chunk", + "service_tier": "default", + "system_fingerprint": None, + "obfuscation": "5A3jk", + }, + { + "choices": [ + { + "content_filter_results": {"protected_material_text": {"detected": False, "filtered": False}}, + "delta": {"content": "scripted answer 20e483d62a14"}, + "finish_reason": None, + "index": 0, + "logprobs": None, + } + ], + "created": 1791855587, + "id": "chatcmpl-$REQUEST_ID", + "model": "gpt-6.1-sol-2026-09-29", + "object": "chat.completion.chunk", + "service_tier": "default", + "system_fingerprint": None, + "obfuscation": "Q2x", + }, + { + "choices": [ + {"content_filter_results": {}, "delta": {}, "finish_reason": "stop", "index": 0, "logprobs": None} + ], + "created": 1791855587, + "id": "chatcmpl-$REQUEST_ID", + "model": "gpt-6.1-sol-2026-09-29", + "object": "chat.completion.chunk", + "service_tier": "default", + "system_fingerprint": None, + "obfuscation": "Uz14j7Ajk4CYaTB", + }, + ), + ), + expected=RecountExpected(recount=RecountRates(input_cost_per_token=2.4e-06, output_cost_per_token=1.2e-05)), +) + +AZURE_GPT_6_SOL_UPSTREAM_404_ZERO_SPEND: Final = CostTrackingTestCase( + name="azure/gpt-6-sol-upstream_404_zero_spend", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="azure/gpt-6-sol", + request={ + "model": "$MODEL", + "messages": [{"role": "user", "content": "scripted upstream failure 404"}], + "stream": False, + }, + response=JsonResponse( + content_type="application/json", + body={ + "error": { + "type": "invalid_request_error", + "code": "DeploymentNotFound", + "message": "The API deployment for this resource does not exist. If you created the deployment within the last 5 minutes, please wait a moment and try again.", + } + }, + status=404, + ), + expected=FailureExpected(failure=FailureDetails(status=404)), +) + +AZURE_GPT_6_SOL_INPUT_TEXT_ROLLUPS: Final = CostTrackingTestCase( + name="azure/gpt-6-sol-input_text_rollups", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="azure/gpt-6-sol", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "c01175bb4811 summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": False, + "allowed_openai_params": [], + }, + response=JsonResponse( + content_type="application/json", + body={ + "choices": [ + { + "content_filter_results": _GPT_6_CONTENT_FILTER_RESULTS, + "finish_reason": "stop", + "index": 0, + "logprobs": None, + "message": { + "annotations": [], + "content": "scripted answer c01175bb4811", + "refusal": None, + "role": "assistant", + }, + } + ], + "created": 1791855585, + "id": "chatcmpl-$UNIQUE_ID", + "model": "gpt-6-sol-2026-09-22", + "object": "chat.completion", + "prompt_filter_results": _GPT_6_PROMPT_FILTER_RESULTS, + "routing": {"serving_pipereplica": "d20260925225746-6787865d-prefill-r15-dp0-prefill"}, + "service_tier": "default", + "system_fingerprint": None, + "usage": { + "completion_tokens": 412, + "completion_tokens_details": { + "accepted_prediction_tokens": 0, + "audio_tokens": 0, + "reasoning_tokens": 0, + "rejected_prediction_tokens": 0, + }, + "latency_checkpoint": { + "engine_tbt_ms": 6, + "engine_ttft_ms": 152, + "engine_ttlt_ms": 305, + "pre_inference_ms": 102, + "service_tbt_ms": 9, + "service_ttft_ms": 389, + "service_ttlt_ms": 608, + "user_visible_ttft_ms": 288, + }, + "prompt_tokens": 1840, + "prompt_tokens_details": {"audio_tokens": 0, "cache_write_tokens": 0, "cached_tokens": 0}, + "total_tokens": 2252, + }, + }, + ), + expected=ExactExpected( + spend=0.00897, + input_cost=0.004232, + output_cost=0.004738, + prompt_tokens=1840, + completion_tokens=412, + rollups=True, + ), +) + +AZURE_GPT_6_SOL_TOOL_CALL: Final = CostTrackingTestCase( + name="azure/gpt-6-sol-tool_call", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="azure/gpt-6-sol", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "59601dc5d1c2 summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": False, + "tools": [ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get the current weather and a short forecast for a city.", + "parameters": { + "type": "object", + "properties": { + "city": {"type": "string", "description": "City name"}, + "days": {"type": "integer", "description": "Forecast horizon in days"}, + "units": {"type": "string", "enum": ["metric", "imperial"]}, + }, + "required": ["city"], + }, + }, + } + ], + "tool_choice": "auto", + "reasoning_effort": "none", + "allowed_openai_params": ["tool_choice", "reasoning_effort"], + }, + response=JsonResponse( + content_type="application/json", + body={ + "choices": [ + { + "content_filter_results": _GPT_6_CONTENT_FILTER_RESULTS, + "finish_reason": "tool_calls", + "index": 0, + "logprobs": None, + "message": { + "annotations": [], + "content": None, + "refusal": None, + "role": "assistant", + "tool_calls": [ + { + "function": { + "arguments": '{"city":"Berlin","days":7,"units":"metric"}', + "name": "get_weather", + }, + "id": "call_$REQUEST_ID", + "type": "function", + } + ], + }, + } + ], + "created": 1791855585, + "id": "chatcmpl-$REQUEST_ID", + "model": "gpt-6-sol-2026-09-22", + "object": "chat.completion", + "prompt_filter_results": _GPT_6_PROMPT_FILTER_RESULTS, + "routing": {"serving_pipereplica": "d20260925225746-6787865d-prefill-r15-dp0-prefill"}, + "service_tier": "default", + "system_fingerprint": None, + "usage": { + "completion_tokens": 412, + "completion_tokens_details": { + "accepted_prediction_tokens": 0, + "audio_tokens": 0, + "reasoning_tokens": 0, + "rejected_prediction_tokens": 0, + }, + "latency_checkpoint": { + "engine_tbt_ms": 6, + "engine_ttft_ms": 152, + "engine_ttlt_ms": 305, + "pre_inference_ms": 102, + "service_tbt_ms": 9, + "service_ttft_ms": 389, + "service_ttlt_ms": 608, + "user_visible_ttft_ms": 288, + }, + "prompt_tokens": 1840, + "prompt_tokens_details": {"audio_tokens": 0, "cache_write_tokens": 0, "cached_tokens": 0}, + "total_tokens": 2252, + }, + }, + ), + expected=ExactExpected( + spend=0.00897, input_cost=0.004232, output_cost=0.004738, prompt_tokens=1840, completion_tokens=412 + ), +) + +AZURE_GPT_6_LUNA_TOOL_CALL: Final = CostTrackingTestCase( + name="azure/gpt-6-luna-tool_call", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="azure/gpt-6-luna", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "2dc3f74a8aa3 summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": False, + "tools": [ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get the current weather and a short forecast for a city.", + "parameters": { + "type": "object", + "properties": { + "city": {"type": "string", "description": "City name"}, + "days": {"type": "integer", "description": "Forecast horizon in days"}, + "units": {"type": "string", "enum": ["metric", "imperial"]}, + }, + "required": ["city"], + }, + }, + } + ], + "tool_choice": "auto", + "reasoning_effort": "none", + "allowed_openai_params": ["tool_choice", "reasoning_effort"], + }, + response=JsonResponse( + content_type="application/json", + body={ + "choices": [ + { + "content_filter_results": _GPT_6_CONTENT_FILTER_RESULTS, + "finish_reason": "tool_calls", + "index": 0, + "logprobs": None, + "message": { + "annotations": [], + "content": None, + "refusal": None, + "role": "assistant", + "tool_calls": [ + { + "function": { + "arguments": '{"city":"Berlin","days":7,"units":"metric"}', + "name": "get_weather", + }, + "id": "call_$REQUEST_ID", + "type": "function", + } + ], + }, + } + ], + "created": 1791855585, + "id": "chatcmpl-$REQUEST_ID", + "model": "gpt-6-luna-2026-09-22", + "object": "chat.completion", + "prompt_filter_results": _GPT_6_PROMPT_FILTER_RESULTS, + "routing": {"serving_pipereplica": "d20260925225746-6787865d-prefill-r15-dp0-prefill"}, + "service_tier": "default", + "system_fingerprint": None, + "usage": { + "completion_tokens": 412, + "completion_tokens_details": { + "accepted_prediction_tokens": 0, + "audio_tokens": 0, + "reasoning_tokens": 0, + "rejected_prediction_tokens": 0, + }, + "latency_checkpoint": { + "engine_tbt_ms": 6, + "engine_ttft_ms": 152, + "engine_ttlt_ms": 305, + "pre_inference_ms": 102, + "service_tbt_ms": 9, + "service_ttft_ms": 389, + "service_ttlt_ms": 608, + "user_visible_ttft_ms": 288, + }, + "prompt_tokens": 1840, + "prompt_tokens_details": {"audio_tokens": 0, "cache_write_tokens": 0, "cached_tokens": 0}, + "total_tokens": 2252, + }, + }, + ), + expected=ExactExpected( + spend=0.000507, input_cost=0.0002392, output_cost=0.0002678, prompt_tokens=1840, completion_tokens=412 + ), +) + CASES: Final[tuple[CostTrackingTestCase, ...]] = ( AZURE_GPT_5_4_MINI_INPUT_TEXT, AZURE_GPT_5_4_MINI_CACHE_READ, @@ -3031,6 +5615,34 @@ CASES: Final[tuple[CostTrackingTestCase, ...]] = ( AZURE_TTS_NEXT_SPEECH_DEPLOYMENT, AZURE_TEXT_EMBEDDINGS_4_LARGE_DEPLOYMENT, AZURE_PINNED_GPT_5_4_MINI_STREAM, + AZURE_GPT_6_SOL_INPUT_TEXT, + AZURE_GPT_6_SOL_STREAM, + AZURE_GPT_6_SOL_SERVICE_TIER_FLEX, + AZURE_GPT_6_SOL_SERVICE_TIER_PRIORITY, + AZURE_GPT_6_LUNA_INPUT_TEXT, + AZURE_GPT_6_LUNA_STREAM, + AZURE_GPT_6_LUNA_SERVICE_TIER_FLEX, + AZURE_GPT_6_LUNA_SERVICE_TIER_PRIORITY, + AZURE_GPT_6_ASTRA_INPUT_TEXT, + AZURE_GPT_6_ASTRA_STREAM, + AZURE_GPT_6_ASTRA_SERVICE_TIER_FLEX, + AZURE_GPT_6_ASTRA_SERVICE_TIER_PRIORITY, + AZURE_GPT_6_1_SOL_INPUT_TEXT, + AZURE_GPT_6_1_SOL_STREAM, + AZURE_GPT_6_1_SOL_SERVICE_TIER_FLEX, + AZURE_GPT_6_1_SOL_SERVICE_TIER_PRIORITY, + AZURE_GPT_6_SOL_REASONING, + AZURE_GPT_6_SOL_STREAM_NO_USAGE, + AZURE_GPT_6_LUNA_REASONING, + AZURE_GPT_6_LUNA_STREAM_NO_USAGE, + AZURE_GPT_6_ASTRA_REASONING, + AZURE_GPT_6_ASTRA_STREAM_NO_USAGE, + AZURE_GPT_6_1_SOL_REASONING, + AZURE_GPT_6_1_SOL_STREAM_NO_USAGE, + AZURE_GPT_6_SOL_UPSTREAM_404_ZERO_SPEND, + AZURE_GPT_6_SOL_INPUT_TEXT_ROLLUPS, + AZURE_GPT_6_SOL_TOOL_CALL, + AZURE_GPT_6_LUNA_TOOL_CALL, ) PARITY: Final[tuple[StreamParityTestCase, ...]] = ( @@ -3039,4 +5651,8 @@ PARITY: Final[tuple[StreamParityTestCase, ...]] = ( plain=AZURE_GPT_5_6_RESPONSE_MODEL_OVERRIDE, streamed=AZURE_GPT_5_6_STREAM_RESPONSE_MODEL_OVERRIDE ), StreamParityTestCase(plain=AZURE_GPT_5_6_TOOL_CALL, streamed=AZURE_GPT_5_6_STREAM_TOOL_CALL), + StreamParityTestCase(plain=AZURE_GPT_6_SOL_INPUT_TEXT, streamed=AZURE_GPT_6_SOL_STREAM), + StreamParityTestCase(plain=AZURE_GPT_6_LUNA_INPUT_TEXT, streamed=AZURE_GPT_6_LUNA_STREAM), + StreamParityTestCase(plain=AZURE_GPT_6_ASTRA_INPUT_TEXT, streamed=AZURE_GPT_6_ASTRA_STREAM), + StreamParityTestCase(plain=AZURE_GPT_6_1_SOL_INPUT_TEXT, streamed=AZURE_GPT_6_1_SOL_STREAM), ) diff --git a/tests/integration/cost_calculation/cases/bedrock_converse.py b/tests/integration/cost_calculation/cases/bedrock_converse.py index 7987ecfee76..7e33b700a48 100644 --- a/tests/integration/cost_calculation/cases/bedrock_converse.py +++ b/tests/integration/cost_calculation/cases/bedrock_converse.py @@ -10,13 +10,14 @@ from integration.cost_calculation.cost_tracking_case import ( EventStreamEvent, EventStreamResponse, ExactExpected, + FailureDetails, + FailureExpected, JsonResponse, RecountExpected, RecountRates, ) from integration.cost_calculation.stream_parity.case import StreamParityTestCase - # anthropic.claude-sonnet-5-v1:0 ANTHROPIC_CLAUDE_SONNET_5_V1_0_INPUT_TEXT: Final = CostTrackingTestCase( name="anthropic.claude-sonnet-5-v1:0-input_text", @@ -1096,6 +1097,119 @@ BEDROCK_CONVERSE_APAC_BARE_FALLBACK: Final = CostTrackingTestCase( ) +BEDROCK_CONVERSE_PROFILE_ARN_REGIONAL_KEY: Final = CostTrackingTestCase( + name="bedrock-converse-profile-arn-regional-key", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="eu.anthropic.claude-sonnet-5-v1:0", + deployment=Deployment( + model="bedrock/converse/arn:aws:bedrock:eu-west-1:123456789012:inference-profile/eu.anthropic.claude-sonnet-5-v1:0" + ), + request={ + "model": "$MODEL", + "messages": [{"role": "user", "content": "bedrock-converse-profile-arn-regional-key"}], + "stream": False, + }, + response=JsonResponse( + content_type="application/json", + body={ + "output": {"message": {"role": "assistant", "content": [{"text": "scripted response"}]}}, + "stopReason": "end_turn", + "usage": {"inputTokens": 1840, "outputTokens": 412, "totalTokens": 2252}, + "metrics": {"latencyMs": 1}, + }, + ), + expected=ExactExpected( + spend=0.01326, input_cost=0.006256, output_cost=0.007004, prompt_tokens=1840, completion_tokens=412 + ), +) + +BEDROCK_CONVERSE_PROFILE_ARN_REGIONAL_KEY_STREAM: Final = CostTrackingTestCase( + name="bedrock-converse-profile-arn-regional-key-stream", + covers="quota_management.spend_tracking.scripted_wire.logs_cost", + model="eu.anthropic.claude-sonnet-5-v1:0", + deployment=Deployment( + model="bedrock/converse/arn:aws:bedrock:eu-west-1:123456789012:inference-profile/eu.anthropic.claude-sonnet-5-v1:0" + ), + request={ + "model": "$MODEL", + "messages": [{"role": "user", "content": "bedrock-converse-profile-arn-regional-key-stream"}], + "stream": True, + "stream_options": {"include_usage": True}, + }, + response=EventStreamResponse( + content_type="application/vnd.amazon.eventstream", + events=( + EventStreamEvent(event_type="messageStart", payload={"role": "assistant"}), + EventStreamEvent( + event_type="contentBlockDelta", payload={"delta": {"text": "scripted response"}, "contentBlockIndex": 0} + ), + EventStreamEvent(event_type="contentBlockStop", payload={"contentBlockIndex": 0}), + EventStreamEvent(event_type="messageStop", payload={"stopReason": "end_turn"}), + EventStreamEvent( + event_type="metadata", + payload={ + "usage": {"inputTokens": 1840, "outputTokens": 412, "totalTokens": 2252}, + "metrics": {"latencyMs": 1}, + }, + ), + ), + ), + expected=ExactExpected( + spend=0.01326, input_cost=0.006256, output_cost=0.007004, prompt_tokens=1840, completion_tokens=412 + ), +) + +BEDROCK_CONVERSE_PROFILE_ARN_REGIONAL_KEY_MESSAGES: Final = CostTrackingTestCase( + name="bedrock-converse-profile-arn-regional-key-messages", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="eu.anthropic.claude-sonnet-5-v1:0", + endpoint="/v1/messages", + deployment=Deployment( + model="bedrock/converse/arn:aws:bedrock:eu-west-1:123456789012:inference-profile/eu.anthropic.claude-sonnet-5-v1:0" + ), + request={ + "model": "$MODEL", + "max_tokens": 412, + "messages": [{"role": "user", "content": "bedrock-converse-profile-arn-regional-key-messages"}], + }, + response=JsonResponse( + content_type="application/json", + body={ + "output": {"message": {"role": "assistant", "content": [{"text": "scripted response"}]}}, + "stopReason": "end_turn", + "usage": {"inputTokens": 1840, "outputTokens": 412, "totalTokens": 2252}, + "metrics": {"latencyMs": 1}, + }, + ), + expected=ExactExpected( + spend=0.01326, input_cost=0.006256, output_cost=0.007004, prompt_tokens=1840, completion_tokens=412 + ), +) + +BEDROCK_CONVERSE_PROFILE_ARN_REGIONAL_KEY_RESPONSES: Final = CostTrackingTestCase( + name="bedrock-converse-profile-arn-regional-key-responses", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="eu.anthropic.claude-sonnet-5-v1:0", + endpoint="/v1/responses", + deployment=Deployment( + model="bedrock/converse/arn:aws:bedrock:eu-west-1:123456789012:inference-profile/eu.anthropic.claude-sonnet-5-v1:0" + ), + request={"model": "$MODEL", "input": "bedrock-converse-profile-arn-regional-key-responses"}, + response=JsonResponse( + content_type="application/json", + body={ + "output": {"message": {"role": "assistant", "content": [{"text": "scripted response"}]}}, + "stopReason": "end_turn", + "usage": {"inputTokens": 1840, "outputTokens": 412, "totalTokens": 2252}, + "metrics": {"latencyMs": 1}, + }, + ), + expected=ExactExpected( + spend=0.01326, input_cost=0.006256, output_cost=0.007004, prompt_tokens=1840, completion_tokens=412 + ), +) + + # meta.llama4-maverick-17b-instruct-v1:0 META_LLAMA4_MAVERICK_17B_INSTRUCT_V1_0_INPUT_TEXT: Final = CostTrackingTestCase( name="meta.llama4-maverick-17b-instruct-v1:0-input_text", @@ -2991,6 +3105,2825 @@ US_ANTHROPIC_CLAUDE_OPUS_5_5_INPUT_TEXT: Final = CostTrackingTestCase( ) +US_ANTHROPIC_CLAUDE_SONNET_5_INPUT_TEXT: Final = CostTrackingTestCase( + name="us.anthropic.claude-sonnet-5-input_text", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="us.anthropic.claude-sonnet-5", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "5517b299b268 summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": False, + "allowed_openai_params": [], + }, + response=JsonResponse( + content_type="application/json", + body={ + "metrics": {"latencyMs": 878}, + "output": {"message": {"content": [{"text": "scripted answer 5517b299b268"}], "role": "assistant"}}, + "stopReason": "end_turn", + "usage": { + "cacheReadInputTokenCount": 0, + "cacheReadInputTokens": 0, + "cacheWriteInputTokenCount": 0, + "cacheWriteInputTokens": 0, + "inputTokens": 1840, + "outputTokens": 412, + "serverToolUsage": {}, + "totalTokens": 2252, + }, + }, + ), + expected=ExactExpected( + spend=0.01404, input_cost=0.006624, output_cost=0.007416, prompt_tokens=1840, completion_tokens=412 + ), +) + +US_ANTHROPIC_CLAUDE_SONNET_5_STREAM: Final = CostTrackingTestCase( + name="us.anthropic.claude-sonnet-5-stream", + covers="quota_management.spend_tracking.scripted_wire.logs_cost", + model="us.anthropic.claude-sonnet-5", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "00121895862b summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": True, + "stream_options": {"include_usage": True}, + "allowed_openai_params": [], + }, + response=EventStreamResponse( + content_type="application/vnd.amazon.eventstream", + events=( + EventStreamEvent(event_type="messageStart", payload={"role": "assistant"}), + EventStreamEvent( + event_type="contentBlockDelta", + payload={"contentBlockIndex": 0, "delta": {"text": "scripted answer 00121895862b"}}, + ), + EventStreamEvent(event_type="contentBlockStop", payload={"contentBlockIndex": 0}), + EventStreamEvent(event_type="messageStop", payload={"stopReason": "end_turn"}), + EventStreamEvent( + event_type="metadata", + payload={ + "metrics": {"latencyMs": 883}, + "usage": {"inputTokens": 1840, "outputTokens": 412, "serverToolUsage": {}, "totalTokens": 2252}, + }, + ), + ), + ), + expected=ExactExpected( + spend=0.01404, input_cost=0.006624, output_cost=0.007416, prompt_tokens=1840, completion_tokens=412 + ), +) + +US_ANTHROPIC_CLAUDE_SONNET_5_CACHE_READ: Final = CostTrackingTestCase( + name="us.anthropic.claude-sonnet-5-cache_read", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="us.anthropic.claude-sonnet-5", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + "cache_control": {"type": "ephemeral"}, + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "c90a8b9a2eb4 summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": False, + "allowed_openai_params": [], + }, + response=JsonResponse( + content_type="application/json", + body={ + "metrics": {"latencyMs": 878}, + "output": {"message": {"content": [{"text": "scripted answer c90a8b9a2eb4"}], "role": "assistant"}}, + "stopReason": "end_turn", + "usage": { + "inputTokens": 640, + "outputTokens": 380, + "totalTokens": 13308, + "cacheReadInputTokens": 12288, + "cacheWriteInputTokens": 0, + }, + }, + ), + expected=ExactExpected( + spend=0.01356768, input_cost=0.00672768, output_cost=0.00684, prompt_tokens=12928, completion_tokens=380 + ), +) + +US_ANTHROPIC_CLAUDE_SONNET_5_CACHE_WRITE_5M: Final = CostTrackingTestCase( + name="us.anthropic.claude-sonnet-5-cache_write_5m", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="us.anthropic.claude-sonnet-5", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + "cache_control": {"type": "ephemeral"}, + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "ffd613a5af4a summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": False, + "allowed_openai_params": [], + }, + response=JsonResponse( + content_type="application/json", + body={ + "metrics": {"latencyMs": 878}, + "output": {"message": {"content": [{"text": "scripted answer ffd613a5af4a"}], "role": "assistant"}}, + "stopReason": "end_turn", + "usage": { + "inputTokens": 512, + "outputTokens": 350, + "totalTokens": 10078, + "cacheReadInputTokens": 0, + "cacheWriteInputTokens": 9216, + "cacheDetails": [{"ttl": "5m", "inputTokens": 9216}], + }, + }, + ), + expected=ExactExpected( + spend=0.0496152, input_cost=0.0433152, output_cost=0.0063, prompt_tokens=9728, completion_tokens=350 + ), +) + +US_ANTHROPIC_CLAUDE_SONNET_5_CACHE_WRITE_1H: Final = CostTrackingTestCase( + name="us.anthropic.claude-sonnet-5-cache_write_1h", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="us.anthropic.claude-sonnet-5", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + "cache_control": {"type": "ephemeral", "ttl": "1h"}, + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "402b2e3c8f52 summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": False, + "allowed_openai_params": [], + }, + response=JsonResponse( + content_type="application/json", + body={ + "metrics": {"latencyMs": 878}, + "output": {"message": {"content": [{"text": "scripted answer 402b2e3c8f52"}], "role": "assistant"}}, + "stopReason": "end_turn", + "usage": { + "inputTokens": 512, + "outputTokens": 350, + "totalTokens": 10078, + "cacheReadInputTokens": 0, + "cacheWriteInputTokens": 9216, + "cacheDetails": [{"ttl": "1h", "inputTokens": 9216}], + }, + }, + ), + expected=ExactExpected( + spend=0.0744984, input_cost=0.0681984, output_cost=0.0063, prompt_tokens=9728, completion_tokens=350 + ), +) + +US_ANTHROPIC_CLAUDE_SONNET_5_5_INPUT_TEXT: Final = CostTrackingTestCase( + name="us.anthropic.claude-sonnet-5-5-input_text", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="us.anthropic.claude-sonnet-5-5", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "1a50e720f40a summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": False, + "allowed_openai_params": [], + }, + response=JsonResponse( + content_type="application/json", + body={ + "metrics": {"latencyMs": 878}, + "output": {"message": {"content": [{"text": "scripted answer 1a50e720f40a"}], "role": "assistant"}}, + "stopReason": "end_turn", + "usage": { + "cacheReadInputTokenCount": 0, + "cacheReadInputTokens": 0, + "cacheWriteInputTokenCount": 0, + "cacheWriteInputTokens": 0, + "inputTokens": 1840, + "outputTokens": 412, + "serverToolUsage": {}, + "totalTokens": 2252, + }, + }, + ), + expected=ExactExpected( + spend=0.01443, input_cost=0.006808, output_cost=0.007622, prompt_tokens=1840, completion_tokens=412 + ), +) + +US_ANTHROPIC_CLAUDE_SONNET_5_5_STREAM: Final = CostTrackingTestCase( + name="us.anthropic.claude-sonnet-5-5-stream", + covers="quota_management.spend_tracking.scripted_wire.logs_cost", + model="us.anthropic.claude-sonnet-5-5", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "00d4a73ea61d summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": True, + "stream_options": {"include_usage": True}, + "allowed_openai_params": [], + }, + response=EventStreamResponse( + content_type="application/vnd.amazon.eventstream", + events=( + EventStreamEvent(event_type="messageStart", payload={"role": "assistant"}), + EventStreamEvent( + event_type="contentBlockDelta", + payload={"contentBlockIndex": 0, "delta": {"text": "scripted answer 00d4a73ea61d"}}, + ), + EventStreamEvent(event_type="contentBlockStop", payload={"contentBlockIndex": 0}), + EventStreamEvent(event_type="messageStop", payload={"stopReason": "end_turn"}), + EventStreamEvent( + event_type="metadata", + payload={ + "metrics": {"latencyMs": 883}, + "usage": {"inputTokens": 1840, "outputTokens": 412, "serverToolUsage": {}, "totalTokens": 2252}, + }, + ), + ), + ), + expected=ExactExpected( + spend=0.01443, input_cost=0.006808, output_cost=0.007622, prompt_tokens=1840, completion_tokens=412 + ), +) + +US_ANTHROPIC_CLAUDE_SONNET_5_5_CACHE_READ: Final = CostTrackingTestCase( + name="us.anthropic.claude-sonnet-5-5-cache_read", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="us.anthropic.claude-sonnet-5-5", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + "cache_control": {"type": "ephemeral"}, + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "5eff7f685f14 summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": False, + "allowed_openai_params": [], + }, + response=JsonResponse( + content_type="application/json", + body={ + "metrics": {"latencyMs": 878}, + "output": {"message": {"content": [{"text": "scripted answer 5eff7f685f14"}], "role": "assistant"}}, + "stopReason": "end_turn", + "usage": { + "inputTokens": 640, + "outputTokens": 380, + "totalTokens": 13308, + "cacheReadInputTokens": 12288, + "cacheWriteInputTokens": 0, + }, + }, + ), + expected=ExactExpected( + spend=0.01394456, input_cost=0.00691456, output_cost=0.00703, prompt_tokens=12928, completion_tokens=380 + ), +) + +US_ANTHROPIC_CLAUDE_SONNET_5_5_CACHE_WRITE_5M: Final = CostTrackingTestCase( + name="us.anthropic.claude-sonnet-5-5-cache_write_5m", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="us.anthropic.claude-sonnet-5-5", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + "cache_control": {"type": "ephemeral"}, + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "eda648e63c9f summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": False, + "allowed_openai_params": [], + }, + response=JsonResponse( + content_type="application/json", + body={ + "metrics": {"latencyMs": 878}, + "output": {"message": {"content": [{"text": "scripted answer eda648e63c9f"}], "role": "assistant"}}, + "stopReason": "end_turn", + "usage": { + "inputTokens": 512, + "outputTokens": 350, + "totalTokens": 10078, + "cacheReadInputTokens": 0, + "cacheWriteInputTokens": 9216, + "cacheDetails": [{"ttl": "5m", "inputTokens": 9216}], + }, + }, + ), + expected=ExactExpected( + spend=0.0509934, input_cost=0.0445184, output_cost=0.006475, prompt_tokens=9728, completion_tokens=350 + ), +) + +US_ANTHROPIC_CLAUDE_SONNET_5_5_CACHE_WRITE_1H: Final = CostTrackingTestCase( + name="us.anthropic.claude-sonnet-5-5-cache_write_1h", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="us.anthropic.claude-sonnet-5-5", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + "cache_control": {"type": "ephemeral", "ttl": "1h"}, + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "1799189351b9 summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": False, + "allowed_openai_params": [], + }, + response=JsonResponse( + content_type="application/json", + body={ + "metrics": {"latencyMs": 878}, + "output": {"message": {"content": [{"text": "scripted answer 1799189351b9"}], "role": "assistant"}}, + "stopReason": "end_turn", + "usage": { + "inputTokens": 512, + "outputTokens": 350, + "totalTokens": 10078, + "cacheReadInputTokens": 0, + "cacheWriteInputTokens": 9216, + "cacheDetails": [{"ttl": "1h", "inputTokens": 9216}], + }, + }, + ), + expected=ExactExpected( + spend=0.0765678, input_cost=0.0700928, output_cost=0.006475, prompt_tokens=9728, completion_tokens=350 + ), +) + +US_ANTHROPIC_CLAUDE_OPUS_5_INPUT_TEXT: Final = CostTrackingTestCase( + name="us.anthropic.claude-opus-5-input_text", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="us.anthropic.claude-opus-5", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "e9aaa98730fc summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": False, + "allowed_openai_params": [], + }, + response=JsonResponse( + content_type="application/json", + body={ + "metrics": {"latencyMs": 878}, + "output": {"message": {"content": [{"text": "scripted answer e9aaa98730fc"}], "role": "assistant"}}, + "stopReason": "end_turn", + "usage": { + "cacheReadInputTokenCount": 0, + "cacheReadInputTokens": 0, + "cacheWriteInputTokenCount": 0, + "cacheWriteInputTokens": 0, + "inputTokens": 1840, + "outputTokens": 412, + "serverToolUsage": {}, + "totalTokens": 2252, + }, + }, + ), + expected=ExactExpected( + spend=0.0234, input_cost=0.01104, output_cost=0.01236, prompt_tokens=1840, completion_tokens=412 + ), +) + +US_ANTHROPIC_CLAUDE_OPUS_5_STREAM: Final = CostTrackingTestCase( + name="us.anthropic.claude-opus-5-stream", + covers="quota_management.spend_tracking.scripted_wire.logs_cost", + model="us.anthropic.claude-opus-5", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "236c9ecadc46 summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": True, + "stream_options": {"include_usage": True}, + "allowed_openai_params": [], + }, + response=EventStreamResponse( + content_type="application/vnd.amazon.eventstream", + events=( + EventStreamEvent(event_type="messageStart", payload={"role": "assistant"}), + EventStreamEvent( + event_type="contentBlockDelta", + payload={"contentBlockIndex": 0, "delta": {"text": "scripted answer 236c9ecadc46"}}, + ), + EventStreamEvent(event_type="contentBlockStop", payload={"contentBlockIndex": 0}), + EventStreamEvent(event_type="messageStop", payload={"stopReason": "end_turn"}), + EventStreamEvent( + event_type="metadata", + payload={ + "metrics": {"latencyMs": 883}, + "usage": {"inputTokens": 1840, "outputTokens": 412, "serverToolUsage": {}, "totalTokens": 2252}, + }, + ), + ), + ), + expected=ExactExpected( + spend=0.0234, input_cost=0.01104, output_cost=0.01236, prompt_tokens=1840, completion_tokens=412 + ), +) + +US_ANTHROPIC_CLAUDE_OPUS_5_CACHE_READ: Final = CostTrackingTestCase( + name="us.anthropic.claude-opus-5-cache_read", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="us.anthropic.claude-opus-5", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + "cache_control": {"type": "ephemeral"}, + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "58afbb027002 summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": False, + "allowed_openai_params": [], + }, + response=JsonResponse( + content_type="application/json", + body={ + "metrics": {"latencyMs": 878}, + "output": {"message": {"content": [{"text": "scripted answer 58afbb027002"}], "role": "assistant"}}, + "stopReason": "end_turn", + "usage": { + "inputTokens": 640, + "outputTokens": 380, + "totalTokens": 13308, + "cacheReadInputTokens": 12288, + "cacheWriteInputTokens": 0, + }, + }, + ), + expected=ExactExpected( + spend=0.0226128, input_cost=0.0112128, output_cost=0.0114, prompt_tokens=12928, completion_tokens=380 + ), +) + +US_ANTHROPIC_CLAUDE_OPUS_5_CACHE_WRITE_5M: Final = CostTrackingTestCase( + name="us.anthropic.claude-opus-5-cache_write_5m", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="us.anthropic.claude-opus-5", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + "cache_control": {"type": "ephemeral"}, + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "16c6329484ca summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": False, + "allowed_openai_params": [], + }, + response=JsonResponse( + content_type="application/json", + body={ + "metrics": {"latencyMs": 878}, + "output": {"message": {"content": [{"text": "scripted answer 16c6329484ca"}], "role": "assistant"}}, + "stopReason": "end_turn", + "usage": { + "inputTokens": 512, + "outputTokens": 350, + "totalTokens": 10078, + "cacheReadInputTokens": 0, + "cacheWriteInputTokens": 9216, + "cacheDetails": [{"ttl": "5m", "inputTokens": 9216}], + }, + }, + ), + expected=ExactExpected( + spend=0.082692, input_cost=0.072192, output_cost=0.0105, prompt_tokens=9728, completion_tokens=350 + ), +) + +US_ANTHROPIC_CLAUDE_OPUS_5_CACHE_WRITE_1H: Final = CostTrackingTestCase( + name="us.anthropic.claude-opus-5-cache_write_1h", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="us.anthropic.claude-opus-5", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + "cache_control": {"type": "ephemeral", "ttl": "1h"}, + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "78000b5822e0 summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": False, + "allowed_openai_params": [], + }, + response=JsonResponse( + content_type="application/json", + body={ + "metrics": {"latencyMs": 878}, + "output": {"message": {"content": [{"text": "scripted answer 78000b5822e0"}], "role": "assistant"}}, + "stopReason": "end_turn", + "usage": { + "inputTokens": 512, + "outputTokens": 350, + "totalTokens": 10078, + "cacheReadInputTokens": 0, + "cacheWriteInputTokens": 9216, + "cacheDetails": [{"ttl": "1h", "inputTokens": 9216}], + }, + }, + ), + expected=ExactExpected( + spend=0.124164, input_cost=0.113664, output_cost=0.0105, prompt_tokens=9728, completion_tokens=350 + ), +) + +US_ANTHROPIC_CLAUDE_OPUS_5_5_STREAM: Final = CostTrackingTestCase( + name="us.anthropic.claude-opus-5-5-stream", + covers="quota_management.spend_tracking.scripted_wire.logs_cost", + model="us.anthropic.claude-opus-5-5", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "cd0638ad2925 summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": True, + "stream_options": {"include_usage": True}, + "allowed_openai_params": [], + }, + response=EventStreamResponse( + content_type="application/vnd.amazon.eventstream", + events=( + EventStreamEvent(event_type="messageStart", payload={"role": "assistant"}), + EventStreamEvent( + event_type="contentBlockDelta", + payload={"contentBlockIndex": 0, "delta": {"text": "scripted answer cd0638ad2925"}}, + ), + EventStreamEvent(event_type="contentBlockStop", payload={"contentBlockIndex": 0}), + EventStreamEvent(event_type="messageStop", payload={"stopReason": "end_turn"}), + EventStreamEvent( + event_type="metadata", + payload={ + "metrics": {"latencyMs": 883}, + "usage": {"inputTokens": 1840, "outputTokens": 412, "serverToolUsage": {}, "totalTokens": 2252}, + }, + ), + ), + ), + expected=ExactExpected( + spend=0.01716, input_cost=0.008096, output_cost=0.009064, prompt_tokens=1840, completion_tokens=412 + ), +) + +US_ANTHROPIC_CLAUDE_OPUS_5_5_CACHE_READ: Final = CostTrackingTestCase( + name="us.anthropic.claude-opus-5-5-cache_read", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="us.anthropic.claude-opus-5-5", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + "cache_control": {"type": "ephemeral"}, + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "bbd4da2e326b summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": False, + "allowed_openai_params": [], + }, + response=JsonResponse( + content_type="application/json", + body={ + "metrics": {"latencyMs": 878}, + "output": {"message": {"content": [{"text": "scripted answer bbd4da2e326b"}], "role": "assistant"}}, + "stopReason": "end_turn", + "usage": { + "inputTokens": 640, + "outputTokens": 380, + "totalTokens": 13308, + "cacheReadInputTokens": 12288, + "cacheWriteInputTokens": 0, + }, + }, + ), + expected=ExactExpected( + spend=0.01387936, input_cost=0.00551936, output_cost=0.00836, prompt_tokens=12928, completion_tokens=380 + ), +) + +US_ANTHROPIC_CLAUDE_OPUS_5_5_CACHE_WRITE_5M: Final = CostTrackingTestCase( + name="us.anthropic.claude-opus-5-5-cache_write_5m", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="us.anthropic.claude-opus-5-5", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + "cache_control": {"type": "ephemeral"}, + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "0bd639d53999 summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": False, + "allowed_openai_params": [], + }, + response=JsonResponse( + content_type="application/json", + body={ + "metrics": {"latencyMs": 878}, + "output": {"message": {"content": [{"text": "scripted answer 0bd639d53999"}], "role": "assistant"}}, + "stopReason": "end_turn", + "usage": { + "inputTokens": 512, + "outputTokens": 350, + "totalTokens": 10078, + "cacheReadInputTokens": 0, + "cacheWriteInputTokens": 9216, + "cacheDetails": [{"ttl": "5m", "inputTokens": 9216}], + }, + }, + ), + expected=ExactExpected( + spend=0.0606408, input_cost=0.0529408, output_cost=0.0077, prompt_tokens=9728, completion_tokens=350 + ), +) + +US_ANTHROPIC_CLAUDE_OPUS_5_5_CACHE_WRITE_1H: Final = CostTrackingTestCase( + name="us.anthropic.claude-opus-5-5-cache_write_1h", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="us.anthropic.claude-opus-5-5", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + "cache_control": {"type": "ephemeral", "ttl": "1h"}, + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "74362dbc328e summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": False, + "allowed_openai_params": [], + }, + response=JsonResponse( + content_type="application/json", + body={ + "metrics": {"latencyMs": 878}, + "output": {"message": {"content": [{"text": "scripted answer 74362dbc328e"}], "role": "assistant"}}, + "stopReason": "end_turn", + "usage": { + "inputTokens": 512, + "outputTokens": 350, + "totalTokens": 10078, + "cacheReadInputTokens": 0, + "cacheWriteInputTokens": 9216, + "cacheDetails": [{"ttl": "1h", "inputTokens": 9216}], + }, + }, + ), + expected=ExactExpected( + spend=0.0910536, input_cost=0.0833536, output_cost=0.0077, prompt_tokens=9728, completion_tokens=350 + ), +) + +US_ANTHROPIC_CLAUDE_HAIKU_5_5_INPUT_TEXT: Final = CostTrackingTestCase( + name="us.anthropic.claude-haiku-5-5-input_text", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="us.anthropic.claude-haiku-5-5", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "8ff4d0ccb5a5 summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": False, + "allowed_openai_params": [], + }, + response=JsonResponse( + content_type="application/json", + body={ + "metrics": {"latencyMs": 878}, + "output": {"message": {"content": [{"text": "scripted answer 8ff4d0ccb5a5"}], "role": "assistant"}}, + "stopReason": "end_turn", + "usage": { + "cacheReadInputTokenCount": 0, + "cacheReadInputTokens": 0, + "cacheWriteInputTokenCount": 0, + "cacheWriteInputTokens": 0, + "inputTokens": 1840, + "outputTokens": 412, + "serverToolUsage": {}, + "totalTokens": 2252, + }, + }, + ), + expected=ExactExpected( + spend=0.00468, input_cost=0.002208, output_cost=0.002472, prompt_tokens=1840, completion_tokens=412 + ), +) + +US_ANTHROPIC_CLAUDE_HAIKU_5_5_STREAM: Final = CostTrackingTestCase( + name="us.anthropic.claude-haiku-5-5-stream", + covers="quota_management.spend_tracking.scripted_wire.logs_cost", + model="us.anthropic.claude-haiku-5-5", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "4f49296715c0 summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": True, + "stream_options": {"include_usage": True}, + "allowed_openai_params": [], + }, + response=EventStreamResponse( + content_type="application/vnd.amazon.eventstream", + events=( + EventStreamEvent(event_type="messageStart", payload={"role": "assistant"}), + EventStreamEvent( + event_type="contentBlockDelta", + payload={"contentBlockIndex": 0, "delta": {"text": "scripted answer 4f49296715c0"}}, + ), + EventStreamEvent(event_type="contentBlockStop", payload={"contentBlockIndex": 0}), + EventStreamEvent(event_type="messageStop", payload={"stopReason": "end_turn"}), + EventStreamEvent( + event_type="metadata", + payload={ + "metrics": {"latencyMs": 883}, + "usage": {"inputTokens": 1840, "outputTokens": 412, "serverToolUsage": {}, "totalTokens": 2252}, + }, + ), + ), + ), + expected=ExactExpected( + spend=0.00468, input_cost=0.002208, output_cost=0.002472, prompt_tokens=1840, completion_tokens=412 + ), +) + +US_ANTHROPIC_CLAUDE_HAIKU_5_5_CACHE_READ: Final = CostTrackingTestCase( + name="us.anthropic.claude-haiku-5-5-cache_read", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="us.anthropic.claude-haiku-5-5", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + "cache_control": {"type": "ephemeral"}, + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "cf757e3e24df summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": False, + "allowed_openai_params": [], + }, + response=JsonResponse( + content_type="application/json", + body={ + "metrics": {"latencyMs": 878}, + "output": {"message": {"content": [{"text": "scripted answer cf757e3e24df"}], "role": "assistant"}}, + "stopReason": "end_turn", + "usage": { + "inputTokens": 640, + "outputTokens": 380, + "totalTokens": 13308, + "cacheReadInputTokens": 12288, + "cacheWriteInputTokens": 0, + }, + }, + ), + expected=ExactExpected( + spend=0.00452256, input_cost=0.00224256, output_cost=0.00228, prompt_tokens=12928, completion_tokens=380 + ), +) + +US_ANTHROPIC_CLAUDE_HAIKU_5_5_CACHE_WRITE_5M: Final = CostTrackingTestCase( + name="us.anthropic.claude-haiku-5-5-cache_write_5m", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="us.anthropic.claude-haiku-5-5", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + "cache_control": {"type": "ephemeral"}, + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "3f2676ede7bd summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": False, + "allowed_openai_params": [], + }, + response=JsonResponse( + content_type="application/json", + body={ + "metrics": {"latencyMs": 878}, + "output": {"message": {"content": [{"text": "scripted answer 3f2676ede7bd"}], "role": "assistant"}}, + "stopReason": "end_turn", + "usage": { + "inputTokens": 512, + "outputTokens": 350, + "totalTokens": 10078, + "cacheReadInputTokens": 0, + "cacheWriteInputTokens": 9216, + "cacheDetails": [{"ttl": "5m", "inputTokens": 9216}], + }, + }, + ), + expected=ExactExpected( + spend=0.0165384, input_cost=0.0144384, output_cost=0.0021, prompt_tokens=9728, completion_tokens=350 + ), +) + +US_ANTHROPIC_CLAUDE_HAIKU_5_5_CACHE_WRITE_1H: Final = CostTrackingTestCase( + name="us.anthropic.claude-haiku-5-5-cache_write_1h", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="us.anthropic.claude-haiku-5-5", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + "cache_control": {"type": "ephemeral", "ttl": "1h"}, + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "f88c14519a7a summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": False, + "allowed_openai_params": [], + }, + response=JsonResponse( + content_type="application/json", + body={ + "metrics": {"latencyMs": 878}, + "output": {"message": {"content": [{"text": "scripted answer f88c14519a7a"}], "role": "assistant"}}, + "stopReason": "end_turn", + "usage": { + "inputTokens": 512, + "outputTokens": 350, + "totalTokens": 10078, + "cacheReadInputTokens": 0, + "cacheWriteInputTokens": 9216, + "cacheDetails": [{"ttl": "1h", "inputTokens": 9216}], + }, + }, + ), + expected=ExactExpected( + spend=0.0248328, input_cost=0.0227328, output_cost=0.0021, prompt_tokens=9728, completion_tokens=350 + ), +) + +US_ANTHROPIC_CLAUDE_FABLE_5_INPUT_TEXT: Final = CostTrackingTestCase( + name="us.anthropic.claude-fable-5-input_text", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="us.anthropic.claude-fable-5", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "8818a06860ae summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": False, + "allowed_openai_params": [], + }, + response=JsonResponse( + content_type="application/json", + body={ + "metrics": {"latencyMs": 878}, + "output": {"message": {"content": [{"text": "scripted answer 8818a06860ae"}], "role": "assistant"}}, + "stopReason": "end_turn", + "usage": { + "cacheReadInputTokenCount": 0, + "cacheReadInputTokens": 0, + "cacheWriteInputTokenCount": 0, + "cacheWriteInputTokens": 0, + "inputTokens": 1840, + "outputTokens": 412, + "serverToolUsage": {}, + "totalTokens": 2252, + }, + }, + ), + expected=ExactExpected( + spend=0.0663, input_cost=0.03128, output_cost=0.03502, prompt_tokens=1840, completion_tokens=412 + ), +) + +US_ANTHROPIC_CLAUDE_FABLE_5_STREAM: Final = CostTrackingTestCase( + name="us.anthropic.claude-fable-5-stream", + covers="quota_management.spend_tracking.scripted_wire.logs_cost", + model="us.anthropic.claude-fable-5", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "5f82eb56d1a9 summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": True, + "stream_options": {"include_usage": True}, + "allowed_openai_params": [], + }, + response=EventStreamResponse( + content_type="application/vnd.amazon.eventstream", + events=( + EventStreamEvent(event_type="messageStart", payload={"role": "assistant"}), + EventStreamEvent( + event_type="contentBlockDelta", + payload={"contentBlockIndex": 0, "delta": {"text": "scripted answer 5f82eb56d1a9"}}, + ), + EventStreamEvent(event_type="contentBlockStop", payload={"contentBlockIndex": 0}), + EventStreamEvent(event_type="messageStop", payload={"stopReason": "end_turn"}), + EventStreamEvent( + event_type="metadata", + payload={ + "metrics": {"latencyMs": 883}, + "usage": {"inputTokens": 1840, "outputTokens": 412, "serverToolUsage": {}, "totalTokens": 2252}, + }, + ), + ), + ), + expected=ExactExpected( + spend=0.0663, input_cost=0.03128, output_cost=0.03502, prompt_tokens=1840, completion_tokens=412 + ), +) + +US_ANTHROPIC_CLAUDE_FABLE_5_CACHE_READ: Final = CostTrackingTestCase( + name="us.anthropic.claude-fable-5-cache_read", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="us.anthropic.claude-fable-5", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + "cache_control": {"type": "ephemeral"}, + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "aa578bd8dff7 summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": False, + "allowed_openai_params": [], + }, + response=JsonResponse( + content_type="application/json", + body={ + "metrics": {"latencyMs": 878}, + "output": {"message": {"content": [{"text": "scripted answer aa578bd8dff7"}], "role": "assistant"}}, + "stopReason": "end_turn", + "usage": { + "inputTokens": 640, + "outputTokens": 380, + "totalTokens": 13308, + "cacheReadInputTokens": 12288, + "cacheWriteInputTokens": 0, + }, + }, + ), + expected=ExactExpected( + spend=0.0640696, input_cost=0.0317696, output_cost=0.0323, prompt_tokens=12928, completion_tokens=380 + ), +) + +US_ANTHROPIC_CLAUDE_FABLE_5_CACHE_WRITE_5M: Final = CostTrackingTestCase( + name="us.anthropic.claude-fable-5-cache_write_5m", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="us.anthropic.claude-fable-5", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + "cache_control": {"type": "ephemeral"}, + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "44d5ae527d01 summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": False, + "allowed_openai_params": [], + }, + response=JsonResponse( + content_type="application/json", + body={ + "metrics": {"latencyMs": 878}, + "output": {"message": {"content": [{"text": "scripted answer 44d5ae527d01"}], "role": "assistant"}}, + "stopReason": "end_turn", + "usage": { + "inputTokens": 512, + "outputTokens": 350, + "totalTokens": 10078, + "cacheReadInputTokens": 0, + "cacheWriteInputTokens": 9216, + "cacheDetails": [{"ttl": "5m", "inputTokens": 9216}], + }, + }, + ), + expected=ExactExpected( + spend=0.234294, input_cost=0.204544, output_cost=0.02975, prompt_tokens=9728, completion_tokens=350 + ), +) + +US_ANTHROPIC_CLAUDE_FABLE_5_CACHE_WRITE_1H: Final = CostTrackingTestCase( + name="us.anthropic.claude-fable-5-cache_write_1h", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="us.anthropic.claude-fable-5", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + "cache_control": {"type": "ephemeral", "ttl": "1h"}, + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "ab69d8f80811 summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": False, + "allowed_openai_params": [], + }, + response=JsonResponse( + content_type="application/json", + body={ + "metrics": {"latencyMs": 878}, + "output": {"message": {"content": [{"text": "scripted answer ab69d8f80811"}], "role": "assistant"}}, + "stopReason": "end_turn", + "usage": { + "inputTokens": 512, + "outputTokens": 350, + "totalTokens": 10078, + "cacheReadInputTokens": 0, + "cacheWriteInputTokens": 9216, + "cacheDetails": [{"ttl": "1h", "inputTokens": 9216}], + }, + }, + ), + expected=ExactExpected( + spend=0.351798, input_cost=0.322048, output_cost=0.02975, prompt_tokens=9728, completion_tokens=350 + ), +) + +US_ANTHROPIC_CLAUDE_FABLE_5_1_INPUT_TEXT: Final = CostTrackingTestCase( + name="us.anthropic.claude-fable-5-1-input_text", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="us.anthropic.claude-fable-5-1", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "a5b0ca93d2e1 summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": False, + "allowed_openai_params": [], + }, + response=JsonResponse( + content_type="application/json", + body={ + "metrics": {"latencyMs": 878}, + "output": {"message": {"content": [{"text": "scripted answer a5b0ca93d2e1"}], "role": "assistant"}}, + "stopReason": "end_turn", + "usage": { + "cacheReadInputTokenCount": 0, + "cacheReadInputTokens": 0, + "cacheWriteInputTokenCount": 0, + "cacheWriteInputTokens": 0, + "inputTokens": 1840, + "outputTokens": 412, + "serverToolUsage": {}, + "totalTokens": 2252, + }, + }, + ), + expected=ExactExpected( + spend=0.0702, input_cost=0.03312, output_cost=0.03708, prompt_tokens=1840, completion_tokens=412 + ), +) + +US_ANTHROPIC_CLAUDE_FABLE_5_1_STREAM: Final = CostTrackingTestCase( + name="us.anthropic.claude-fable-5-1-stream", + covers="quota_management.spend_tracking.scripted_wire.logs_cost", + model="us.anthropic.claude-fable-5-1", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "d62b383256ff summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": True, + "stream_options": {"include_usage": True}, + "allowed_openai_params": [], + }, + response=EventStreamResponse( + content_type="application/vnd.amazon.eventstream", + events=( + EventStreamEvent(event_type="messageStart", payload={"role": "assistant"}), + EventStreamEvent( + event_type="contentBlockDelta", + payload={"contentBlockIndex": 0, "delta": {"text": "scripted answer d62b383256ff"}}, + ), + EventStreamEvent(event_type="contentBlockStop", payload={"contentBlockIndex": 0}), + EventStreamEvent(event_type="messageStop", payload={"stopReason": "end_turn"}), + EventStreamEvent( + event_type="metadata", + payload={ + "metrics": {"latencyMs": 883}, + "usage": {"inputTokens": 1840, "outputTokens": 412, "serverToolUsage": {}, "totalTokens": 2252}, + }, + ), + ), + ), + expected=ExactExpected( + spend=0.0702, input_cost=0.03312, output_cost=0.03708, prompt_tokens=1840, completion_tokens=412 + ), +) + +US_ANTHROPIC_CLAUDE_FABLE_5_1_CACHE_READ: Final = CostTrackingTestCase( + name="us.anthropic.claude-fable-5-1-cache_read", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="us.anthropic.claude-fable-5-1", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + "cache_control": {"type": "ephemeral"}, + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "1a5a397f1134 summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": False, + "allowed_openai_params": [], + }, + response=JsonResponse( + content_type="application/json", + body={ + "metrics": {"latencyMs": 878}, + "output": {"message": {"content": [{"text": "scripted answer 1a5a397f1134"}], "role": "assistant"}}, + "stopReason": "end_turn", + "usage": { + "inputTokens": 640, + "outputTokens": 380, + "totalTokens": 13308, + "cacheReadInputTokens": 12288, + "cacheWriteInputTokens": 0, + }, + }, + ), + expected=ExactExpected( + spend=0.0678384, input_cost=0.0336384, output_cost=0.0342, prompt_tokens=12928, completion_tokens=380 + ), +) + +US_ANTHROPIC_CLAUDE_FABLE_5_1_CACHE_WRITE_5M: Final = CostTrackingTestCase( + name="us.anthropic.claude-fable-5-1-cache_write_5m", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="us.anthropic.claude-fable-5-1", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + "cache_control": {"type": "ephemeral"}, + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "3481bca67e1a summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": False, + "allowed_openai_params": [], + }, + response=JsonResponse( + content_type="application/json", + body={ + "metrics": {"latencyMs": 878}, + "output": {"message": {"content": [{"text": "scripted answer 3481bca67e1a"}], "role": "assistant"}}, + "stopReason": "end_turn", + "usage": { + "inputTokens": 512, + "outputTokens": 350, + "totalTokens": 10078, + "cacheReadInputTokens": 0, + "cacheWriteInputTokens": 9216, + "cacheDetails": [{"ttl": "5m", "inputTokens": 9216}], + }, + }, + ), + expected=ExactExpected( + spend=0.248076, input_cost=0.216576, output_cost=0.0315, prompt_tokens=9728, completion_tokens=350 + ), +) + +US_ANTHROPIC_CLAUDE_FABLE_5_1_CACHE_WRITE_1H: Final = CostTrackingTestCase( + name="us.anthropic.claude-fable-5-1-cache_write_1h", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="us.anthropic.claude-fable-5-1", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + "cache_control": {"type": "ephemeral", "ttl": "1h"}, + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "3f037d5317b0 summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": False, + "allowed_openai_params": [], + }, + response=JsonResponse( + content_type="application/json", + body={ + "metrics": {"latencyMs": 878}, + "output": {"message": {"content": [{"text": "scripted answer 3f037d5317b0"}], "role": "assistant"}}, + "stopReason": "end_turn", + "usage": { + "inputTokens": 512, + "outputTokens": 350, + "totalTokens": 10078, + "cacheReadInputTokens": 0, + "cacheWriteInputTokens": 9216, + "cacheDetails": [{"ttl": "1h", "inputTokens": 9216}], + }, + }, + ), + expected=ExactExpected( + spend=0.372492, input_cost=0.340992, output_cost=0.0315, prompt_tokens=9728, completion_tokens=350 + ), +) + +US_ANTHROPIC_CLAUDE_SONNET_5_TOOL_CALL: Final = CostTrackingTestCase( + name="us.anthropic.claude-sonnet-5-tool_call", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="us.anthropic.claude-sonnet-5", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "c87a264a36d3 summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": False, + "tools": [ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get the current weather and a short forecast for a city.", + "parameters": { + "type": "object", + "properties": { + "city": {"type": "string", "description": "City name"}, + "days": {"type": "integer", "description": "Forecast horizon in days"}, + "units": {"type": "string", "enum": ["metric", "imperial"]}, + }, + "required": ["city"], + }, + }, + } + ], + "tool_choice": "auto", + "allowed_openai_params": ["tool_choice"], + }, + response=JsonResponse( + content_type="application/json", + body={ + "output": { + "message": { + "role": "assistant", + "content": [ + { + "reasoningContent": { + "reasoningText": { + "text": "", + "signature": "CAQS1gQKEAgSGAI4AUIIdGhpbmtpbmcSDOsCmv80YWY0IrUDvBoMog2", + } + } + }, + { + "toolUse": { + "toolUseId": "tooluse_$REQUEST_ID", + "name": "get_weather", + "input": {"city": "Berlin", "days": 7, "units": "metric"}, + "type": "tool_use", + } + }, + ], + } + }, + "stopReason": "tool_use", + "usage": { + "inputTokens": 1840, + "outputTokens": 412, + "totalTokens": 2252, + "cacheReadInputTokens": 0, + "cacheWriteInputTokens": 0, + }, + "metrics": {"latencyMs": 2071}, + }, + ), + expected=ExactExpected( + spend=0.01404, input_cost=0.006624, output_cost=0.007416, prompt_tokens=1840, completion_tokens=412 + ), +) + +US_ANTHROPIC_CLAUDE_SONNET_5_THINKING: Final = CostTrackingTestCase( + name="us.anthropic.claude-sonnet-5-thinking", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="us.anthropic.claude-sonnet-5", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "2736f5db99fa summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": False, + "reasoning_effort": "medium", + "allowed_openai_params": ["reasoning_effort"], + }, + response=JsonResponse( + content_type="application/json", + body={ + "output": { + "message": { + "role": "assistant", + "content": [ + { + "reasoningContent": { + "reasoningText": { + "text": "", + "signature": "CAQS1gQKEAgSGAI4AUIIdGhpbmtpbmcSDOsCmv80YWY0IrUDvBoMog2", + } + } + }, + {"text": "scripted answer 2736f5db99fa"}, + ], + } + }, + "stopReason": "end_turn", + "usage": { + "inputTokens": 1240, + "outputTokens": 4040, + "totalTokens": 5280, + "cacheReadInputTokens": 0, + "cacheWriteInputTokens": 0, + }, + "metrics": {"latencyMs": 2071}, + }, + ), + expected=ExactExpected( + spend=0.077184, input_cost=0.004464, output_cost=0.07272, prompt_tokens=1240, completion_tokens=4040 + ), +) + +US_ANTHROPIC_CLAUDE_SONNET_5_5_TOOL_CALL: Final = CostTrackingTestCase( + name="us.anthropic.claude-sonnet-5-5-tool_call", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="us.anthropic.claude-sonnet-5-5", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "a19e69e4777a summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": False, + "tools": [ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get the current weather and a short forecast for a city.", + "parameters": { + "type": "object", + "properties": { + "city": {"type": "string", "description": "City name"}, + "days": {"type": "integer", "description": "Forecast horizon in days"}, + "units": {"type": "string", "enum": ["metric", "imperial"]}, + }, + "required": ["city"], + }, + }, + } + ], + "tool_choice": "auto", + "allowed_openai_params": ["tool_choice"], + }, + response=JsonResponse( + content_type="application/json", + body={ + "output": { + "message": { + "role": "assistant", + "content": [ + { + "toolUse": { + "toolUseId": "tooluse_$REQUEST_ID", + "name": "get_weather", + "input": {"city": "Berlin", "days": 7, "units": "metric"}, + "type": "tool_use", + } + } + ], + } + }, + "stopReason": "tool_use", + "usage": { + "inputTokens": 1840, + "outputTokens": 412, + "totalTokens": 2252, + "cacheReadInputTokens": 0, + "cacheWriteInputTokens": 0, + }, + "metrics": {"latencyMs": 2071}, + }, + ), + expected=ExactExpected( + spend=0.01443, input_cost=0.006808, output_cost=0.007622, prompt_tokens=1840, completion_tokens=412 + ), +) + +US_ANTHROPIC_CLAUDE_SONNET_5_5_THINKING: Final = CostTrackingTestCase( + name="us.anthropic.claude-sonnet-5-5-thinking", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="us.anthropic.claude-sonnet-5-5", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "fe9c9714d04d summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": False, + "reasoning_effort": "medium", + "allowed_openai_params": ["reasoning_effort"], + }, + response=JsonResponse( + content_type="application/json", + body={ + "output": {"message": {"role": "assistant", "content": [{"text": "scripted answer fe9c9714d04d"}]}}, + "stopReason": "end_turn", + "usage": { + "inputTokens": 1240, + "outputTokens": 412, + "totalTokens": 1652, + "cacheReadInputTokens": 0, + "cacheWriteInputTokens": 0, + }, + "metrics": {"latencyMs": 2071}, + }, + ), + expected=ExactExpected( + spend=0.01221, input_cost=0.004588, output_cost=0.007622, prompt_tokens=1240, completion_tokens=412 + ), +) + +US_ANTHROPIC_CLAUDE_OPUS_5_TOOL_CALL: Final = CostTrackingTestCase( + name="us.anthropic.claude-opus-5-tool_call", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="us.anthropic.claude-opus-5", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "d2235a5023bf summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": False, + "tools": [ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get the current weather and a short forecast for a city.", + "parameters": { + "type": "object", + "properties": { + "city": {"type": "string", "description": "City name"}, + "days": {"type": "integer", "description": "Forecast horizon in days"}, + "units": {"type": "string", "enum": ["metric", "imperial"]}, + }, + "required": ["city"], + }, + }, + } + ], + "tool_choice": "auto", + "allowed_openai_params": ["tool_choice"], + }, + response=JsonResponse( + content_type="application/json", + body={ + "output": { + "message": { + "role": "assistant", + "content": [ + { + "toolUse": { + "toolUseId": "tooluse_$REQUEST_ID", + "name": "get_weather", + "input": {"city": "Berlin", "days": 7, "units": "metric"}, + "type": "tool_use", + } + } + ], + } + }, + "stopReason": "tool_use", + "usage": { + "inputTokens": 1840, + "outputTokens": 412, + "totalTokens": 2252, + "cacheReadInputTokens": 0, + "cacheWriteInputTokens": 0, + }, + "metrics": {"latencyMs": 2071}, + }, + ), + expected=ExactExpected( + spend=0.0234, input_cost=0.01104, output_cost=0.01236, prompt_tokens=1840, completion_tokens=412 + ), +) + +US_ANTHROPIC_CLAUDE_OPUS_5_THINKING: Final = CostTrackingTestCase( + name="us.anthropic.claude-opus-5-thinking", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="us.anthropic.claude-opus-5", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "50c5c1415303 summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": False, + "reasoning_effort": "medium", + "allowed_openai_params": ["reasoning_effort"], + }, + response=JsonResponse( + content_type="application/json", + body={ + "output": { + "message": { + "role": "assistant", + "content": [ + { + "reasoningContent": { + "reasoningText": { + "text": "", + "signature": "CAQS1gQKEAgSGAI4AUIIdGhpbmtpbmcSDOsCmv80YWY0IrUDvBoMog2", + } + } + }, + {"text": "scripted answer 50c5c1415303"}, + ], + } + }, + "stopReason": "end_turn", + "usage": { + "inputTokens": 1240, + "outputTokens": 4040, + "totalTokens": 5280, + "cacheReadInputTokens": 0, + "cacheWriteInputTokens": 0, + }, + "metrics": {"latencyMs": 2071}, + }, + ), + expected=ExactExpected( + spend=0.12864, input_cost=0.00744, output_cost=0.1212, prompt_tokens=1240, completion_tokens=4040 + ), +) + +US_ANTHROPIC_CLAUDE_OPUS_5_5_TOOL_CALL: Final = CostTrackingTestCase( + name="us.anthropic.claude-opus-5-5-tool_call", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="us.anthropic.claude-opus-5-5", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "6095fb901387 summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": False, + "tools": [ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get the current weather and a short forecast for a city.", + "parameters": { + "type": "object", + "properties": { + "city": {"type": "string", "description": "City name"}, + "days": {"type": "integer", "description": "Forecast horizon in days"}, + "units": {"type": "string", "enum": ["metric", "imperial"]}, + }, + "required": ["city"], + }, + }, + } + ], + "tool_choice": "auto", + "allowed_openai_params": ["tool_choice"], + }, + response=JsonResponse( + content_type="application/json", + body={ + "output": { + "message": { + "role": "assistant", + "content": [ + { + "toolUse": { + "toolUseId": "tooluse_$REQUEST_ID", + "name": "get_weather", + "input": {"city": "Berlin", "days": 7, "units": "metric"}, + "type": "tool_use", + } + } + ], + } + }, + "stopReason": "tool_use", + "usage": { + "inputTokens": 1840, + "outputTokens": 412, + "totalTokens": 2252, + "cacheReadInputTokens": 0, + "cacheWriteInputTokens": 0, + }, + "metrics": {"latencyMs": 2071}, + }, + ), + expected=ExactExpected( + spend=0.01716, input_cost=0.008096, output_cost=0.009064, prompt_tokens=1840, completion_tokens=412 + ), +) + +US_ANTHROPIC_CLAUDE_OPUS_5_5_THINKING: Final = CostTrackingTestCase( + name="us.anthropic.claude-opus-5-5-thinking", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="us.anthropic.claude-opus-5-5", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "2ba98156d243 summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": False, + "reasoning_effort": "medium", + "allowed_openai_params": ["reasoning_effort"], + }, + response=JsonResponse( + content_type="application/json", + body={ + "output": { + "message": { + "role": "assistant", + "content": [ + { + "reasoningContent": { + "reasoningText": { + "text": "", + "signature": "CAQS1gQKEAgSGAI4AUIIdGhpbmtpbmcSDOsCmv80YWY0IrUDvBoMog2", + } + } + }, + {"text": "scripted answer 2ba98156d243"}, + ], + } + }, + "stopReason": "end_turn", + "usage": { + "inputTokens": 1240, + "outputTokens": 4040, + "totalTokens": 5280, + "cacheReadInputTokens": 0, + "cacheWriteInputTokens": 0, + }, + "metrics": {"latencyMs": 2071}, + }, + ), + expected=ExactExpected( + spend=0.094336, input_cost=0.005456, output_cost=0.08888, prompt_tokens=1240, completion_tokens=4040 + ), +) + +US_ANTHROPIC_CLAUDE_HAIKU_5_5_TOOL_CALL: Final = CostTrackingTestCase( + name="us.anthropic.claude-haiku-5-5-tool_call", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="us.anthropic.claude-haiku-5-5", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "7d7caefb1f0d summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": False, + "tools": [ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get the current weather and a short forecast for a city.", + "parameters": { + "type": "object", + "properties": { + "city": {"type": "string", "description": "City name"}, + "days": {"type": "integer", "description": "Forecast horizon in days"}, + "units": {"type": "string", "enum": ["metric", "imperial"]}, + }, + "required": ["city"], + }, + }, + } + ], + "tool_choice": "auto", + "allowed_openai_params": ["tool_choice"], + }, + response=JsonResponse( + content_type="application/json", + body={ + "output": { + "message": { + "role": "assistant", + "content": [ + { + "toolUse": { + "toolUseId": "tooluse_$REQUEST_ID", + "name": "get_weather", + "input": {"city": "Berlin", "days": 7, "units": "metric"}, + "type": "tool_use", + } + } + ], + } + }, + "stopReason": "tool_use", + "usage": { + "inputTokens": 1840, + "outputTokens": 412, + "totalTokens": 2252, + "cacheReadInputTokens": 0, + "cacheWriteInputTokens": 0, + }, + "metrics": {"latencyMs": 2071}, + }, + ), + expected=ExactExpected( + spend=0.00468, input_cost=0.002208, output_cost=0.002472, prompt_tokens=1840, completion_tokens=412 + ), +) + +US_ANTHROPIC_CLAUDE_HAIKU_5_5_THINKING: Final = CostTrackingTestCase( + name="us.anthropic.claude-haiku-5-5-thinking", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="us.anthropic.claude-haiku-5-5", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "90a6bb1996a5 summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": False, + "reasoning_effort": "medium", + "allowed_openai_params": ["reasoning_effort"], + }, + response=JsonResponse( + content_type="application/json", + body={ + "output": {"message": {"role": "assistant", "content": [{"text": "scripted answer 90a6bb1996a5"}]}}, + "stopReason": "end_turn", + "usage": { + "inputTokens": 1240, + "outputTokens": 412, + "totalTokens": 1652, + "cacheReadInputTokens": 0, + "cacheWriteInputTokens": 0, + }, + "metrics": {"latencyMs": 2071}, + }, + ), + expected=ExactExpected( + spend=0.00396, input_cost=0.001488, output_cost=0.002472, prompt_tokens=1240, completion_tokens=412 + ), +) + +US_ANTHROPIC_CLAUDE_FABLE_5_TOOL_CALL: Final = CostTrackingTestCase( + name="us.anthropic.claude-fable-5-tool_call", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="us.anthropic.claude-fable-5", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "a51000416e04 summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": False, + "tools": [ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get the current weather and a short forecast for a city.", + "parameters": { + "type": "object", + "properties": { + "city": {"type": "string", "description": "City name"}, + "days": {"type": "integer", "description": "Forecast horizon in days"}, + "units": {"type": "string", "enum": ["metric", "imperial"]}, + }, + "required": ["city"], + }, + }, + } + ], + "tool_choice": "auto", + "allowed_openai_params": ["tool_choice"], + }, + response=JsonResponse( + content_type="application/json", + body={ + "output": { + "message": { + "role": "assistant", + "content": [ + { + "toolUse": { + "toolUseId": "tooluse_$REQUEST_ID", + "name": "get_weather", + "input": {"city": "Berlin", "days": 7, "units": "metric"}, + "type": "tool_use", + } + } + ], + } + }, + "stopReason": "tool_use", + "usage": { + "inputTokens": 1840, + "outputTokens": 412, + "totalTokens": 2252, + "cacheReadInputTokens": 0, + "cacheWriteInputTokens": 0, + }, + "metrics": {"latencyMs": 2071}, + }, + ), + expected=ExactExpected( + spend=0.0663, input_cost=0.03128, output_cost=0.03502, prompt_tokens=1840, completion_tokens=412 + ), +) + +US_ANTHROPIC_CLAUDE_FABLE_5_THINKING: Final = CostTrackingTestCase( + name="us.anthropic.claude-fable-5-thinking", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="us.anthropic.claude-fable-5", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "24888a505ee9 summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": False, + "reasoning_effort": "medium", + "allowed_openai_params": ["reasoning_effort"], + }, + response=JsonResponse( + content_type="application/json", + body={ + "output": { + "message": { + "role": "assistant", + "content": [ + { + "reasoningContent": { + "reasoningText": { + "text": "", + "signature": "CAQS1gQKEAgSGAI4AUIIdGhpbmtpbmcSDOsCmv80YWY0IrUDvBoMog2", + } + } + }, + {"text": "scripted answer 24888a505ee9"}, + ], + } + }, + "stopReason": "end_turn", + "usage": { + "inputTokens": 1240, + "outputTokens": 4040, + "totalTokens": 5280, + "cacheReadInputTokens": 0, + "cacheWriteInputTokens": 0, + }, + "metrics": {"latencyMs": 2071}, + }, + ), + expected=ExactExpected( + spend=0.36448, input_cost=0.02108, output_cost=0.3434, prompt_tokens=1240, completion_tokens=4040 + ), +) + +US_ANTHROPIC_CLAUDE_FABLE_5_1_TOOL_CALL: Final = CostTrackingTestCase( + name="us.anthropic.claude-fable-5-1-tool_call", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="us.anthropic.claude-fable-5-1", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "7f97ab74e808 summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": False, + "tools": [ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get the current weather and a short forecast for a city.", + "parameters": { + "type": "object", + "properties": { + "city": {"type": "string", "description": "City name"}, + "days": {"type": "integer", "description": "Forecast horizon in days"}, + "units": {"type": "string", "enum": ["metric", "imperial"]}, + }, + "required": ["city"], + }, + }, + } + ], + "tool_choice": "auto", + "allowed_openai_params": ["tool_choice"], + }, + response=JsonResponse( + content_type="application/json", + body={ + "output": { + "message": { + "role": "assistant", + "content": [ + { + "toolUse": { + "toolUseId": "tooluse_$REQUEST_ID", + "name": "get_weather", + "input": {"city": "Berlin", "days": 7, "units": "metric"}, + "type": "tool_use", + } + } + ], + } + }, + "stopReason": "tool_use", + "usage": { + "inputTokens": 1840, + "outputTokens": 412, + "totalTokens": 2252, + "cacheReadInputTokens": 0, + "cacheWriteInputTokens": 0, + }, + "metrics": {"latencyMs": 2071}, + }, + ), + expected=ExactExpected( + spend=0.0702, input_cost=0.03312, output_cost=0.03708, prompt_tokens=1840, completion_tokens=412 + ), +) + +US_ANTHROPIC_CLAUDE_FABLE_5_1_THINKING: Final = CostTrackingTestCase( + name="us.anthropic.claude-fable-5-1-thinking", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="us.anthropic.claude-fable-5-1", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "bca700fd52a2 summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": False, + "reasoning_effort": "medium", + "allowed_openai_params": ["reasoning_effort"], + }, + response=JsonResponse( + content_type="application/json", + body={ + "output": {"message": {"role": "assistant", "content": [{"text": "scripted answer bca700fd52a2"}]}}, + "stopReason": "end_turn", + "usage": { + "inputTokens": 1240, + "outputTokens": 412, + "totalTokens": 1652, + "cacheReadInputTokens": 0, + "cacheWriteInputTokens": 0, + }, + "metrics": {"latencyMs": 2071}, + }, + ), + expected=ExactExpected( + spend=0.0594, input_cost=0.02232, output_cost=0.03708, prompt_tokens=1240, completion_tokens=412 + ), +) + +US_ANTHROPIC_CLAUDE_OPUS_5_5_UPSTREAM_400_ZERO_SPEND: Final = CostTrackingTestCase( + name="us.anthropic.claude-opus-5-5-upstream_400_zero_spend", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="us.anthropic.claude-opus-5-5", + request={ + "model": "$MODEL", + "messages": [{"role": "user", "content": "scripted upstream failure 400"}], + "stream": False, + }, + response=JsonResponse( + content_type="application/json", + body={"message": "The provided model identifier is invalid."}, + status=400, + ), + expected=FailureExpected(failure=FailureDetails(status=400)), +) + +US_ANTHROPIC_CLAUDE_OPUS_5_5_INPUT_TEXT_ROLLUPS: Final = CostTrackingTestCase( + name="us.anthropic.claude-opus-5-5-input_text_rollups", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="us.anthropic.claude-opus-5-5", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "75e5a866d0e3 summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": False, + "allowed_openai_params": [], + }, + response=JsonResponse( + content_type="application/json", + body={ + "output": { + "message": {"role": "assistant", "content": [{"text": "scripted answer 75e5a866d0e3 $UNIQUE_ID"}]} + }, + "stopReason": "end_turn", + "usage": { + "inputTokens": 1840, + "outputTokens": 412, + "totalTokens": 2252, + "cacheReadInputTokens": 0, + "cacheWriteInputTokens": 0, + }, + "metrics": {"latencyMs": 2071}, + }, + ), + expected=ExactExpected( + spend=0.01716, + input_cost=0.008096, + output_cost=0.009064, + prompt_tokens=1840, + completion_tokens=412, + rollups=True, + ), +) + CASES: Final[tuple[CostTrackingTestCase, ...]] = ( ANTHROPIC_CLAUDE_SONNET_5_V1_0_INPUT_TEXT, ANTHROPIC_CLAUDE_SONNET_5_V1_0_CACHE_READ, @@ -3044,12 +5977,69 @@ CASES: Final[tuple[CostTrackingTestCase, ...]] = ( BEDROCK_CONVERSE_PROFILE_BASE_MODEL, BEDROCK_CONVERSE_EU_REGIONAL_KEY, BEDROCK_CONVERSE_APAC_BARE_FALLBACK, + BEDROCK_CONVERSE_PROFILE_ARN_REGIONAL_KEY, + BEDROCK_CONVERSE_PROFILE_ARN_REGIONAL_KEY_STREAM, + BEDROCK_CONVERSE_PROFILE_ARN_REGIONAL_KEY_MESSAGES, + BEDROCK_CONVERSE_PROFILE_ARN_REGIONAL_KEY_RESPONSES, BEDROCK_CONVERSE_NOVA_2_PRO, BEDROCK_CONVERSE_MISTRAL_LARGE_3_STREAM, US_ANTHROPIC_CLAUDE_OPUS_5_5_INPUT_TEXT, + US_ANTHROPIC_CLAUDE_SONNET_5_INPUT_TEXT, + US_ANTHROPIC_CLAUDE_SONNET_5_STREAM, + US_ANTHROPIC_CLAUDE_SONNET_5_CACHE_READ, + US_ANTHROPIC_CLAUDE_SONNET_5_CACHE_WRITE_5M, + US_ANTHROPIC_CLAUDE_SONNET_5_CACHE_WRITE_1H, + US_ANTHROPIC_CLAUDE_SONNET_5_5_INPUT_TEXT, + US_ANTHROPIC_CLAUDE_SONNET_5_5_STREAM, + US_ANTHROPIC_CLAUDE_SONNET_5_5_CACHE_READ, + US_ANTHROPIC_CLAUDE_SONNET_5_5_CACHE_WRITE_5M, + US_ANTHROPIC_CLAUDE_SONNET_5_5_CACHE_WRITE_1H, + US_ANTHROPIC_CLAUDE_OPUS_5_INPUT_TEXT, + US_ANTHROPIC_CLAUDE_OPUS_5_STREAM, + US_ANTHROPIC_CLAUDE_OPUS_5_CACHE_READ, + US_ANTHROPIC_CLAUDE_OPUS_5_CACHE_WRITE_5M, + US_ANTHROPIC_CLAUDE_OPUS_5_CACHE_WRITE_1H, + US_ANTHROPIC_CLAUDE_OPUS_5_5_STREAM, + US_ANTHROPIC_CLAUDE_OPUS_5_5_CACHE_READ, + US_ANTHROPIC_CLAUDE_OPUS_5_5_CACHE_WRITE_5M, + US_ANTHROPIC_CLAUDE_OPUS_5_5_CACHE_WRITE_1H, + US_ANTHROPIC_CLAUDE_HAIKU_5_5_INPUT_TEXT, + US_ANTHROPIC_CLAUDE_HAIKU_5_5_STREAM, + US_ANTHROPIC_CLAUDE_HAIKU_5_5_CACHE_READ, + US_ANTHROPIC_CLAUDE_HAIKU_5_5_CACHE_WRITE_5M, + US_ANTHROPIC_CLAUDE_HAIKU_5_5_CACHE_WRITE_1H, + US_ANTHROPIC_CLAUDE_FABLE_5_INPUT_TEXT, + US_ANTHROPIC_CLAUDE_FABLE_5_STREAM, + US_ANTHROPIC_CLAUDE_FABLE_5_CACHE_READ, + US_ANTHROPIC_CLAUDE_FABLE_5_CACHE_WRITE_5M, + US_ANTHROPIC_CLAUDE_FABLE_5_CACHE_WRITE_1H, + US_ANTHROPIC_CLAUDE_FABLE_5_1_INPUT_TEXT, + US_ANTHROPIC_CLAUDE_FABLE_5_1_STREAM, + US_ANTHROPIC_CLAUDE_FABLE_5_1_CACHE_READ, + US_ANTHROPIC_CLAUDE_FABLE_5_1_CACHE_WRITE_5M, + US_ANTHROPIC_CLAUDE_FABLE_5_1_CACHE_WRITE_1H, + US_ANTHROPIC_CLAUDE_SONNET_5_TOOL_CALL, + US_ANTHROPIC_CLAUDE_SONNET_5_THINKING, + US_ANTHROPIC_CLAUDE_SONNET_5_5_TOOL_CALL, + US_ANTHROPIC_CLAUDE_SONNET_5_5_THINKING, + US_ANTHROPIC_CLAUDE_OPUS_5_TOOL_CALL, + US_ANTHROPIC_CLAUDE_OPUS_5_THINKING, + US_ANTHROPIC_CLAUDE_OPUS_5_5_TOOL_CALL, + US_ANTHROPIC_CLAUDE_OPUS_5_5_THINKING, + US_ANTHROPIC_CLAUDE_HAIKU_5_5_TOOL_CALL, + US_ANTHROPIC_CLAUDE_HAIKU_5_5_THINKING, + US_ANTHROPIC_CLAUDE_FABLE_5_TOOL_CALL, + US_ANTHROPIC_CLAUDE_FABLE_5_THINKING, + US_ANTHROPIC_CLAUDE_FABLE_5_1_TOOL_CALL, + US_ANTHROPIC_CLAUDE_FABLE_5_1_THINKING, + US_ANTHROPIC_CLAUDE_OPUS_5_5_UPSTREAM_400_ZERO_SPEND, + US_ANTHROPIC_CLAUDE_OPUS_5_5_INPUT_TEXT_ROLLUPS, ) PARITY: Final[tuple[StreamParityTestCase, ...]] = ( + StreamParityTestCase( + plain=BEDROCK_CONVERSE_PROFILE_ARN_REGIONAL_KEY, streamed=BEDROCK_CONVERSE_PROFILE_ARN_REGIONAL_KEY_STREAM + ), StreamParityTestCase( plain=ANTHROPIC_CLAUDE_SONNET_5_V1_0_TOOL_CALL, streamed=ANTHROPIC_CLAUDE_SONNET_5_V1_0_STREAM_TOOL_CALL ), @@ -3060,4 +6050,13 @@ PARITY: Final[tuple[StreamParityTestCase, ...]] = ( StreamParityTestCase( plain=US_ANTHROPIC_CLAUDE_OPUS_5_V1_0_TOOL_CALL, streamed=US_ANTHROPIC_CLAUDE_OPUS_5_V1_0_STREAM_TOOL_CALL ), + StreamParityTestCase(plain=US_ANTHROPIC_CLAUDE_SONNET_5_INPUT_TEXT, streamed=US_ANTHROPIC_CLAUDE_SONNET_5_STREAM), + StreamParityTestCase( + plain=US_ANTHROPIC_CLAUDE_SONNET_5_5_INPUT_TEXT, streamed=US_ANTHROPIC_CLAUDE_SONNET_5_5_STREAM + ), + StreamParityTestCase(plain=US_ANTHROPIC_CLAUDE_OPUS_5_INPUT_TEXT, streamed=US_ANTHROPIC_CLAUDE_OPUS_5_STREAM), + StreamParityTestCase(plain=US_ANTHROPIC_CLAUDE_OPUS_5_5_INPUT_TEXT, streamed=US_ANTHROPIC_CLAUDE_OPUS_5_5_STREAM), + StreamParityTestCase(plain=US_ANTHROPIC_CLAUDE_HAIKU_5_5_INPUT_TEXT, streamed=US_ANTHROPIC_CLAUDE_HAIKU_5_5_STREAM), + StreamParityTestCase(plain=US_ANTHROPIC_CLAUDE_FABLE_5_INPUT_TEXT, streamed=US_ANTHROPIC_CLAUDE_FABLE_5_STREAM), + StreamParityTestCase(plain=US_ANTHROPIC_CLAUDE_FABLE_5_1_INPUT_TEXT, streamed=US_ANTHROPIC_CLAUDE_FABLE_5_1_STREAM), ) diff --git a/tests/integration/cost_calculation/cases/decisions.py b/tests/integration/cost_calculation/cases/decisions.py new file mode 100644 index 00000000000..23c71ea7ad4 --- /dev/null +++ b/tests/integration/cost_calculation/cases/decisions.py @@ -0,0 +1,529 @@ +"""Decisions cost cases: every Decisions provider called through /v1/decisions and /v1/systemone. + +Request and response bodies are verbatim captures of what litellm sent to and got back from the +real provider on 2026-10-10 (recorded through a forwarding proxy in front of each API), so each +case replays the provider's own usage shape against the scripted upstream. Spend is the recorded +usage at the suite cost map's rates; OpenRouter has no per-token rate and bills the `usage.cost` +it reports. +""" + +from __future__ import annotations + +from typing import Final + +from integration.cost_calculation.cost_tracking_case import CostTrackingTestCase, ExactExpected, JsonResponse +from integration.cost_calculation.stream_parity.case import StreamParityTestCase + + +# gpt-6-luna +OPENAI_GPT_6_LUNA_DECISIONS: Final = CostTrackingTestCase( + name="gpt-6-luna-decisions", + covers="quota_management.spend_tracking.decisions.logs_cost", + model="gpt-6-luna", + endpoint="/v1/decisions", + request={ + "model": "$MODEL", + "input": "My order arrived broken and I want a refund today.", + "questions": [ + {"type": "predicate", "name": "is_refund", "instructions": "Is the customer asking for a refund?"} + ], + }, + upstream_body={ + "model": "gpt-6-luna", + "input": "My order arrived broken and I want a refund today.", + "questions": [ + {"type": "predicate", "name": "is_refund", "instructions": "Is the customer asking for a refund?"} + ], + }, + response=JsonResponse( + content_type="application/json", + body={ + "model": "gpt-6-luna", + "answers": [{"type": "predicate", "name": "is_refund", "probability": 1.0}], + "usage": { + "input_tokens": 167, + "input_tokens_details": {"cached_tokens": 0, "cache_write_tokens": 0}, + "output_tokens": 0, + "output_tokens_details": {"reasoning_tokens": 0}, + "total_tokens": 167, + }, + }, + ), + expected=ExactExpected( + spend=1.67e-05, + input_cost=1.67e-05, + output_cost=0.0, + prompt_tokens=167, + completion_tokens=0, + ), +) +OPENAI_GPT_6_LUNA_SYSTEMONE: Final = CostTrackingTestCase( + name="gpt-6-luna-systemone", + covers="quota_management.spend_tracking.systemone.logs_cost", + model="gpt-6-luna", + endpoint="/v1/systemone", + request={ + "model": "$MODEL", + "state": "My order arrived broken and I want a refund today.", + "questions": {"is_refund": {"type": "noul", "instructions": "Is the customer asking for a refund?"}}, + }, + upstream_body={ + "model": "gpt-6-luna", + "input": "My order arrived broken and I want a refund today.", + "questions": [ + {"type": "predicate", "name": "is_refund", "instructions": "Is the customer asking for a refund?"} + ], + }, + response=JsonResponse( + content_type="application/json", + body={ + "model": "gpt-6-luna", + "answers": [{"type": "predicate", "name": "is_refund", "probability": 1.0}], + "usage": { + "input_tokens": 167, + "input_tokens_details": {"cached_tokens": 0, "cache_write_tokens": 0}, + "output_tokens": 0, + "output_tokens_details": {"reasoning_tokens": 0}, + "total_tokens": 167, + }, + }, + ), + expected=ExactExpected( + spend=1.67e-05, + input_cost=1.67e-05, + output_cost=0.0, + prompt_tokens=167, + completion_tokens=0, + ), +) + +# typesafe/jev-latest +TYPESAFE_JEV_LATEST_DECISIONS: Final = CostTrackingTestCase( + name="typesafe/jev-latest-decisions", + covers="quota_management.spend_tracking.decisions.logs_cost", + model="typesafe/jev-latest", + endpoint="/v1/decisions", + request={ + "model": "$MODEL", + "input": "My order arrived broken and I want a refund today.", + "questions": [ + {"type": "predicate", "name": "is_refund", "instructions": "Is the customer asking for a refund?"} + ], + }, + upstream_body={ + "model": "jev-latest", + "state": "My order arrived broken and I want a refund today.", + "questions": {"is_refund": {"type": "noul", "instructions": "Is the customer asking for a refund?"}}, + }, + response=JsonResponse( + content_type="application/json", + body={ + "model": "jev-1.13.0", + "answers": {"is_refund": {"type": "noul", "noul": 0.99}}, + "usage": {"input_tokens": 285, "output_tokens": 22}, + }, + ), + expected=ExactExpected( + spend=1.197e-05, + input_cost=1.197e-05, + output_cost=0.0, + prompt_tokens=285, + completion_tokens=22, + ), +) +TYPESAFE_JEV_LATEST_SYSTEMONE: Final = CostTrackingTestCase( + name="typesafe/jev-latest-systemone", + covers="quota_management.spend_tracking.systemone.logs_cost", + model="typesafe/jev-latest", + endpoint="/v1/systemone", + request={ + "model": "$MODEL", + "state": "My order arrived broken and I want a refund today.", + "questions": {"is_refund": {"type": "noul", "instructions": "Is the customer asking for a refund?"}}, + }, + upstream_body={ + "model": "jev-latest", + "state": "My order arrived broken and I want a refund today.", + "questions": {"is_refund": {"type": "noul", "instructions": "Is the customer asking for a refund?"}}, + }, + response=JsonResponse( + content_type="application/json", + body={ + "model": "jev-1.13.0", + "answers": {"is_refund": {"type": "noul", "noul": 0.99}}, + "usage": {"input_tokens": 285, "output_tokens": 22}, + }, + ), + expected=ExactExpected( + spend=1.197e-05, + input_cost=1.197e-05, + output_cost=0.0, + prompt_tokens=285, + completion_tokens=22, + ), +) + +# perplexity/pplx-decider-v1-27b +PERPLEXITY_PPLX_DECIDER_V1_27B_DECISIONS: Final = CostTrackingTestCase( + name="perplexity/pplx-decider-v1-27b-decisions", + covers="quota_management.spend_tracking.decisions.logs_cost", + model="perplexity/pplx-decider-v1-27b", + endpoint="/v1/decisions", + request={ + "model": "$MODEL", + "input": "My order arrived broken and I want a refund today.", + "questions": [ + {"type": "predicate", "name": "is_refund", "instructions": "Is the customer asking for a refund?"} + ], + }, + upstream_body={ + "model": "pplx-decider-v1-27b", + "state": "My order arrived broken and I want a refund today.", + "questions": {"is_refund": {"type": "noul", "instructions": "Is the customer asking for a refund?"}}, + }, + response=JsonResponse( + content_type="application/json", + body={ + "model": "pplx-decider-v1-27b", + "answers": {"is_refund": {"type": "noul", "noul": 0.9997319426472651}}, + "usage": {"input_tokens": 98, "output_tokens": 1}, + }, + ), + expected=ExactExpected( + spend=1.96e-06, + input_cost=1.96e-06, + output_cost=0.0, + prompt_tokens=98, + completion_tokens=1, + ), +) +PERPLEXITY_PPLX_DECIDER_V1_27B_SYSTEMONE: Final = CostTrackingTestCase( + name="perplexity/pplx-decider-v1-27b-systemone", + covers="quota_management.spend_tracking.systemone.logs_cost", + model="perplexity/pplx-decider-v1-27b", + endpoint="/v1/systemone", + request={ + "model": "$MODEL", + "state": "My order arrived broken and I want a refund today.", + "questions": {"is_refund": {"type": "noul", "instructions": "Is the customer asking for a refund?"}}, + }, + upstream_body={ + "model": "pplx-decider-v1-27b", + "state": "My order arrived broken and I want a refund today.", + "questions": {"is_refund": {"type": "noul", "instructions": "Is the customer asking for a refund?"}}, + }, + response=JsonResponse( + content_type="application/json", + body={ + "model": "pplx-decider-v1-27b", + "answers": {"is_refund": {"type": "noul", "noul": 0.9997319426472651}}, + "usage": {"input_tokens": 98, "output_tokens": 1}, + }, + ), + expected=ExactExpected( + spend=1.96e-06, + input_cost=1.96e-06, + output_cost=0.0, + prompt_tokens=98, + completion_tokens=1, + ), +) + +# cloudflare/@cf/cloudflare/clef-flash +CLOUDFLARE_CLEF_FLASH_DECISIONS: Final = CostTrackingTestCase( + name="cloudflare/@cf/cloudflare/clef-flash-decisions", + covers="quota_management.spend_tracking.decisions.logs_cost", + model="cloudflare/@cf/cloudflare/clef-flash", + endpoint="/v1/decisions", + request={ + "model": "$MODEL", + "input": "My order arrived broken and I want a refund today.", + "questions": [ + {"type": "predicate", "name": "is_refund", "instructions": "Is the customer asking for a refund?"} + ], + }, + upstream_body={ + "model": "clef-flash", + "state": "My order arrived broken and I want a refund today.", + "questions": {"is_refund": {"type": "noul", "instructions": "Is the customer asking for a refund?"}}, + }, + response=JsonResponse( + content_type="application/json", + body={ + "result": { + "model": "clef-flash", + "answers": {"is_refund": {"type": "noul", "noul": 0.9693}}, + "usage": {"input_tokens": 157, "output_tokens": 0}, + }, + "success": True, + "errors": [], + "messages": [], + }, + ), + expected=ExactExpected( + spend=1.413e-05, + input_cost=1.413e-05, + output_cost=0.0, + prompt_tokens=157, + completion_tokens=0, + ), +) +CLOUDFLARE_CLEF_FLASH_SYSTEMONE: Final = CostTrackingTestCase( + name="cloudflare/@cf/cloudflare/clef-flash-systemone", + covers="quota_management.spend_tracking.systemone.logs_cost", + model="cloudflare/@cf/cloudflare/clef-flash", + endpoint="/v1/systemone", + request={ + "model": "$MODEL", + "state": "My order arrived broken and I want a refund today.", + "questions": {"is_refund": {"type": "noul", "instructions": "Is the customer asking for a refund?"}}, + }, + upstream_body={ + "model": "clef-flash", + "state": "My order arrived broken and I want a refund today.", + "questions": {"is_refund": {"type": "noul", "instructions": "Is the customer asking for a refund?"}}, + }, + response=JsonResponse( + content_type="application/json", + body={ + "result": { + "model": "clef-flash", + "answers": {"is_refund": {"type": "noul", "noul": 0.9693}}, + "usage": {"input_tokens": 157, "output_tokens": 0}, + }, + "success": True, + "errors": [], + "messages": [], + }, + ), + expected=ExactExpected( + spend=1.413e-05, + input_cost=1.413e-05, + output_cost=0.0, + prompt_tokens=157, + completion_tokens=0, + ), +) + +# databricks/databricks-openjev-qwen35-4b +DATABRICKS_OPENJEV_QWEN35_4B_DECISIONS: Final = CostTrackingTestCase( + name="databricks/databricks-openjev-qwen35-4b-decisions", + covers="quota_management.spend_tracking.decisions.logs_cost", + model="databricks/databricks-openjev-qwen35-4b", + endpoint="/v1/decisions", + request={ + "model": "$MODEL", + "input": "My order arrived broken and I want a refund today.", + "questions": [ + {"type": "predicate", "name": "is_refund", "instructions": "Is the customer asking for a refund?"} + ], + }, + upstream_body={ + "model": "databricks-openjev-qwen35-4b", + "state": "My order arrived broken and I want a refund today.", + "questions": {"is_refund": {"type": "noul", "instructions": "Is the customer asking for a refund?"}}, + }, + response=JsonResponse( + content_type="application/json", + body={ + "model": "databricks-openjev-qwen35-4b", + "answers": {"is_refund": {"type": "noul", "noul": 0.9046505248181417}}, + "usage": {"input_tokens": 45, "output_tokens": 0}, + }, + ), + expected=ExactExpected( + spend=0.0, + input_cost=0.0, + output_cost=0.0, + prompt_tokens=45, + completion_tokens=0, + ), +) +DATABRICKS_OPENJEV_QWEN35_4B_SYSTEMONE: Final = CostTrackingTestCase( + name="databricks/databricks-openjev-qwen35-4b-systemone", + covers="quota_management.spend_tracking.systemone.logs_cost", + model="databricks/databricks-openjev-qwen35-4b", + endpoint="/v1/systemone", + request={ + "model": "$MODEL", + "state": "My order arrived broken and I want a refund today.", + "questions": {"is_refund": {"type": "noul", "instructions": "Is the customer asking for a refund?"}}, + }, + upstream_body={ + "model": "databricks-openjev-qwen35-4b", + "state": "My order arrived broken and I want a refund today.", + "questions": {"is_refund": {"type": "noul", "instructions": "Is the customer asking for a refund?"}}, + }, + response=JsonResponse( + content_type="application/json", + body={ + "model": "databricks-openjev-qwen35-4b", + "answers": {"is_refund": {"type": "noul", "noul": 0.9046505248181417}}, + "usage": {"input_tokens": 45, "output_tokens": 0}, + }, + ), + expected=ExactExpected( + spend=0.0, + input_cost=0.0, + output_cost=0.0, + prompt_tokens=45, + completion_tokens=0, + ), +) + +# openrouter/typesafe/jev-1.13 +OPENROUTER_TYPESAFE_JEV_1_13_DECISIONS: Final = CostTrackingTestCase( + name="openrouter/typesafe/jev-1.13-decisions", + covers="quota_management.spend_tracking.decisions.logs_cost", + model="openrouter/typesafe/jev-1.13", + endpoint="/v1/decisions", + request={ + "model": "$MODEL", + "input": "My order arrived broken and I want a refund today.", + "questions": [ + {"type": "predicate", "name": "is_refund", "instructions": "Is the customer asking for a refund?"} + ], + }, + upstream_body={ + "model": "typesafe/jev-1.13", + "state": "My order arrived broken and I want a refund today.", + "questions": {"is_refund": {"type": "noul", "instructions": "Is the customer asking for a refund?"}}, + }, + response=JsonResponse( + content_type="application/json", + body={ + "model": "typesafe/jev-1.13-20260917", + "answers": {"is_refund": {"type": "noul", "noul": 0.99}}, + "usage": {"input_tokens": 285, "output_tokens": 22, "cost": 1.197e-05}, + "id": "gen-dec-$REQUEST_ID", + "provider": "TypeSafe", + }, + ), + expected=ExactExpected( + spend=1.197e-05, + input_cost=1.197e-05, + output_cost=0.0, + prompt_tokens=285, + completion_tokens=22, + breakdown_persisted=False, + ), +) +OPENROUTER_TYPESAFE_JEV_1_13_SYSTEMONE: Final = CostTrackingTestCase( + name="openrouter/typesafe/jev-1.13-systemone", + covers="quota_management.spend_tracking.systemone.logs_cost", + model="openrouter/typesafe/jev-1.13", + endpoint="/v1/systemone", + request={ + "model": "$MODEL", + "state": "My order arrived broken and I want a refund today.", + "questions": {"is_refund": {"type": "noul", "instructions": "Is the customer asking for a refund?"}}, + }, + upstream_body={ + "model": "typesafe/jev-1.13", + "state": "My order arrived broken and I want a refund today.", + "questions": {"is_refund": {"type": "noul", "instructions": "Is the customer asking for a refund?"}}, + }, + response=JsonResponse( + content_type="application/json", + body={ + "model": "typesafe/jev-1.13-20260917", + "answers": {"is_refund": {"type": "noul", "noul": 0.99}}, + "usage": {"input_tokens": 285, "output_tokens": 22, "cost": 1.197e-05}, + "id": "gen-dec-$REQUEST_ID", + "provider": "TypeSafe", + }, + ), + expected=ExactExpected( + spend=1.197e-05, + input_cost=1.197e-05, + output_cost=0.0, + prompt_tokens=285, + completion_tokens=22, + breakdown_persisted=False, + ), +) + +# fireworks_ai/accounts/fireworks/routers/rune +FIREWORKS_AI_RUNE_DECISIONS: Final = CostTrackingTestCase( + name="fireworks_ai/accounts/fireworks/routers/rune-decisions", + covers="quota_management.spend_tracking.decisions.logs_cost", + model="fireworks_ai/accounts/fireworks/routers/rune", + endpoint="/v1/decisions", + request={ + "model": "$MODEL", + "input": "My order arrived broken and I want a refund today.", + "questions": [ + {"type": "predicate", "name": "is_refund", "instructions": "Is the customer asking for a refund?"} + ], + }, + upstream_body={ + "model": "accounts/fireworks/routers/rune", + "state": "My order arrived broken and I want a refund today.", + "questions": {"is_refund": {"type": "noul", "instructions": "Is the customer asking for a refund?"}}, + }, + response=JsonResponse( + content_type="application/json", + body={ + "model": "accounts/fireworks/routers/rune", + "answers": {"is_refund": {"type": "noul", "noul": 0.99}}, + "usage": {"input_tokens": 300, "cached_input_tokens": 200, "output_tokens": 0}, + }, + ), + expected=ExactExpected( + spend=5.88e-06, + input_cost=5.88e-06, + output_cost=0.0, + prompt_tokens=300, + completion_tokens=0, + ), +) +FIREWORKS_AI_RUNE_SYSTEMONE: Final = CostTrackingTestCase( + name="fireworks_ai/accounts/fireworks/routers/rune-systemone", + covers="quota_management.spend_tracking.systemone.logs_cost", + model="fireworks_ai/accounts/fireworks/routers/rune", + endpoint="/v1/systemone", + request={ + "model": "$MODEL", + "state": "My order arrived broken and I want a refund today.", + "questions": {"is_refund": {"type": "noul", "instructions": "Is the customer asking for a refund?"}}, + }, + upstream_body={ + "model": "accounts/fireworks/routers/rune", + "state": "My order arrived broken and I want a refund today.", + "questions": {"is_refund": {"type": "noul", "instructions": "Is the customer asking for a refund?"}}, + }, + response=JsonResponse( + content_type="application/json", + body={ + "model": "accounts/fireworks/routers/rune", + "answers": {"is_refund": {"type": "noul", "noul": 0.99}}, + "usage": {"input_tokens": 300, "cached_input_tokens": 200, "output_tokens": 0}, + }, + ), + expected=ExactExpected( + spend=5.88e-06, + input_cost=5.88e-06, + output_cost=0.0, + prompt_tokens=300, + completion_tokens=0, + ), +) + +CASES: Final[tuple[CostTrackingTestCase, ...]] = ( + OPENAI_GPT_6_LUNA_DECISIONS, + OPENAI_GPT_6_LUNA_SYSTEMONE, + TYPESAFE_JEV_LATEST_DECISIONS, + TYPESAFE_JEV_LATEST_SYSTEMONE, + PERPLEXITY_PPLX_DECIDER_V1_27B_DECISIONS, + PERPLEXITY_PPLX_DECIDER_V1_27B_SYSTEMONE, + CLOUDFLARE_CLEF_FLASH_DECISIONS, + CLOUDFLARE_CLEF_FLASH_SYSTEMONE, + DATABRICKS_OPENJEV_QWEN35_4B_DECISIONS, + DATABRICKS_OPENJEV_QWEN35_4B_SYSTEMONE, + OPENROUTER_TYPESAFE_JEV_1_13_DECISIONS, + OPENROUTER_TYPESAFE_JEV_1_13_SYSTEMONE, + FIREWORKS_AI_RUNE_DECISIONS, + FIREWORKS_AI_RUNE_SYSTEMONE, +) + +PARITY: Final[tuple[StreamParityTestCase, ...]] = () diff --git a/tests/integration/cost_calculation/cases/misc.py b/tests/integration/cost_calculation/cases/misc.py index d1c17b11adb..24b57c3e93c 100644 --- a/tests/integration/cost_calculation/cases/misc.py +++ b/tests/integration/cost_calculation/cases/misc.py @@ -1,4 +1,4 @@ -"""azure_ai, bedrock, cohere, cohere_chat, dashscope, deepgram, deepseek, groq, mistral, openrouter, perplexity, text-completion-openai, xai cost tracking cases, one CostTrackingTestCase literal per request shape (moved from cost_tracking_cases.json). +"""azure_ai, bedrock, cohere, cohere_chat, dashscope, deepgram, deepseek, groq, mistral, openrouter, text-completion-openai, xai cost tracking cases, one CostTrackingTestCase literal per request shape (moved from cost_tracking_cases.json). PARITY pairs the streamed case with its plain twin where both bill the same row.""" @@ -740,6 +740,535 @@ AZURE_AI_GPT_5_4_MINI_LATEST_STREAM: Final = CostTrackingTestCase( ) +# azure_ai/claude-haiku-4-5 (Azure AI Foundry Anthropic, web search billed per server_tool_use.web_search_requests) +AZURE_AI_CLAUDE_HAIKU_4_5_MESSAGES_WEB_SEARCH: Final = CostTrackingTestCase( + name="azure-ai-claude-haiku-4-5-messages_web_search_two_requests", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="azure_ai/claude-haiku-4-5", + endpoint="/v1/messages", + request={ + "model": "$MODEL", + "max_tokens": 412, + "messages": [{"role": "user", "content": "azure-ai-claude-haiku-4-5-messages_web_search_two_requests"}], + "tools": [{"type": "web_search_20250305", "name": "web_search", "max_uses": 3}], + }, + response=JsonResponse( + content_type="application/json", + body={ + "id": "msg_$UNIQUE_ID", + "type": "message", + "role": "assistant", + "model": "claude-haiku-4-5", + "content": [{"type": "text", "text": "scripted answer a1b2c3d4e5f6"}], + "stop_reason": "end_turn", + "usage": {"input_tokens": 1840, "output_tokens": 412, "server_tool_use": {"web_search_requests": 2}}, + }, + ), + expected=ExactExpected( + spend=0.0239, + input_cost=0.00184, + output_cost=0.00206, + prompt_tokens=1840, + completion_tokens=412, + tool_usage_cost=0.02, + rollups=True, + ), +) + +AZURE_AI_CLAUDE_HAIKU_4_5_MESSAGES_WEB_SEARCH_STREAM: Final = CostTrackingTestCase( + name="azure-ai-claude-haiku-4-5-messages_web_search_two_requests_stream", + covers="quota_management.spend_tracking.scripted_wire.logs_cost", + model="azure_ai/claude-haiku-4-5", + endpoint="/v1/messages", + request={ + "model": "$MODEL", + "max_tokens": 412, + "messages": [{"role": "user", "content": "azure-ai-claude-haiku-4-5-messages_web_search_two_requests_stream"}], + "tools": [{"type": "web_search_20250305", "name": "web_search", "max_uses": 3}], + "stream": True, + }, + response=SseResponse( + content_type="text/event-stream", + frames=( + 'event: message_start\ndata: {"type":"message_start","message":{"id":"msg_$REQUEST_ID","type":"message","role":"assistant","model":"claude-haiku-4-5","content":[],"stop_reason":null,"stop_sequence":null,"usage":{"input_tokens":1840}}}', + 'event: content_block_start\ndata: {"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}', + 'event: content_block_delta\ndata: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"scripted response"}}', + 'event: content_block_stop\ndata: {"type":"content_block_stop","index":0}', + 'event: message_delta\ndata: {"type":"message_delta","delta":{"stop_reason":"end_turn"},"usage":{"output_tokens":412,"server_tool_use":{"web_search_requests":2}}}', + 'event: message_stop\ndata: {"type":"message_stop"}', + ), + ), + expected=ExactExpected( + spend=0.0239, + input_cost=0.00184, + output_cost=0.00206, + prompt_tokens=1840, + completion_tokens=412, + tool_usage_cost=0.02, + ), +) + +AZURE_AI_CLAUDE_HAIKU_4_5_MESSAGES_WEB_SEARCH_ONE_REQUEST: Final = CostTrackingTestCase( + name="azure-ai-claude-haiku-4-5-messages_web_search_one_request", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="azure_ai/claude-haiku-4-5", + endpoint="/v1/messages", + request={ + "model": "$MODEL", + "max_tokens": 412, + "messages": [{"role": "user", "content": "azure-ai-claude-haiku-4-5-messages_web_search_one_request"}], + "tools": [{"type": "web_search_20250305", "name": "web_search", "max_uses": 3}], + }, + response=JsonResponse( + content_type="application/json", + body={ + "id": "msg_$REQUEST_ID", + "type": "message", + "role": "assistant", + "model": "claude-haiku-4-5", + "content": [{"type": "text", "text": "scripted answer 0f1e2d3c4b5a"}], + "stop_reason": "end_turn", + "usage": {"input_tokens": 1840, "output_tokens": 412, "server_tool_use": {"web_search_requests": 1}}, + }, + ), + expected=ExactExpected( + spend=0.0139, + input_cost=0.00184, + output_cost=0.00206, + prompt_tokens=1840, + completion_tokens=412, + tool_usage_cost=0.01, + ), +) + +AZURE_AI_CLAUDE_HAIKU_4_5_MESSAGES_NO_SEARCH: Final = CostTrackingTestCase( + name="azure-ai-claude-haiku-4-5-messages_no_server_tool_use", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="azure_ai/claude-haiku-4-5", + endpoint="/v1/messages", + request={ + "model": "$MODEL", + "max_tokens": 412, + "messages": [{"role": "user", "content": "azure-ai-claude-haiku-4-5-messages_no_server_tool_use"}], + "tools": [{"type": "web_search_20250305", "name": "web_search", "max_uses": 3}], + }, + response=JsonResponse( + content_type="application/json", + body={ + "id": "msg_$REQUEST_ID", + "type": "message", + "role": "assistant", + "model": "claude-haiku-4-5", + "content": [{"type": "text", "text": "scripted answer 9a8b7c6d5e4f"}], + "stop_reason": "end_turn", + "usage": {"input_tokens": 1840, "output_tokens": 412}, + }, + ), + expected=ExactExpected( + spend=0.0039, + input_cost=0.00184, + output_cost=0.00206, + prompt_tokens=1840, + completion_tokens=412, + tool_usage_cost=0.0, + ), +) + +AZURE_AI_CLAUDE_HAIKU_4_5_CHAT_WEB_SEARCH: Final = CostTrackingTestCase( + name="azure-ai-claude-haiku-4-5-chat_web_search_three_requests", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="azure_ai/claude-haiku-4-5", + request={ + "model": "$MODEL", + "messages": [{"role": "user", "content": "azure-ai-claude-haiku-4-5-chat_web_search_three_requests"}], + "stream": False, + "tools": [{"type": "web_search_20250305", "name": "web_search", "max_uses": 5}], + "allowed_openai_params": ["web_search_options"], + }, + response=JsonResponse( + content_type="application/json", + body={ + "id": "msg_$REQUEST_ID", + "type": "message", + "role": "assistant", + "model": "claude-haiku-4-5", + "content": [{"type": "text", "text": "scripted answer 5e4f3a2b1c0d"}], + "stop_reason": "end_turn", + "usage": {"input_tokens": 1840, "output_tokens": 412, "server_tool_use": {"web_search_requests": 3}}, + }, + ), + expected=ExactExpected( + spend=0.0339, + input_cost=0.00184, + output_cost=0.00206, + prompt_tokens=1840, + completion_tokens=412, + tool_usage_cost=0.03, + ), +) + +AZURE_AI_CLAUDE_HAIKU_4_5_CHAT_WEB_SEARCH_STREAM: Final = CostTrackingTestCase( + name="azure-ai-claude-haiku-4-5-chat_web_search_two_requests_stream", + covers="quota_management.spend_tracking.scripted_wire.logs_cost", + model="azure_ai/claude-haiku-4-5", + request={ + "model": "$MODEL", + "messages": [{"role": "user", "content": "azure-ai-claude-haiku-4-5-chat_web_search_two_requests_stream"}], + "stream": True, + "tools": [{"type": "web_search_20250305", "name": "web_search", "max_uses": 5}], + "allowed_openai_params": ["web_search_options"], + }, + response=SseResponse( + content_type="text/event-stream", + frames=sse_frames( + { + "type": "message_start", + "message": { + "id": "msg_$REQUEST_ID", + "type": "message", + "role": "assistant", + "model": "claude-haiku-4-5", + "content": [], + "stop_reason": None, + "usage": {"input_tokens": 1840}, + }, + }, + {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}, + { + "type": "content_block_delta", + "index": 0, + "delta": {"type": "text_delta", "text": "scripted answer 2c3d4e5f6a7b"}, + }, + {"type": "content_block_stop", "index": 0}, + { + "type": "message_delta", + "delta": {"stop_reason": "end_turn"}, + "usage": {"output_tokens": 412, "server_tool_use": {"web_search_requests": 2}}, + }, + {"type": "message_stop"}, + ), + ), + expected=ExactExpected( + spend=0.0239, + input_cost=0.00184, + output_cost=0.00206, + prompt_tokens=1840, + completion_tokens=412, + tool_usage_cost=0.02, + ), +) + + +AZURE_AI_CLAUDE_HAIKU_4_5_RESPONSES_WEB_SEARCH: Final = CostTrackingTestCase( + name="azure-ai-claude-haiku-4-5-responses_web_search_two_requests", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="azure_ai/claude-haiku-4-5", + endpoint="/v1/responses", + request={ + "model": "$MODEL", + "input": "azure-ai-claude-haiku-4-5-responses_web_search_two_requests", + "tools": [{"type": "web_search"}], + }, + response=JsonResponse( + content_type="application/json", + body={ + "id": "msg_$REQUEST_ID", + "type": "message", + "role": "assistant", + "model": "claude-haiku-4-5", + "content": [{"type": "text", "text": "scripted answer 7d6e5f4a3b2c"}], + "stop_reason": "end_turn", + "usage": {"input_tokens": 1840, "output_tokens": 412, "server_tool_use": {"web_search_requests": 2}}, + }, + ), + expected=ExactExpected( + spend=0.0239, + input_cost=0.00184, + output_cost=0.00206, + prompt_tokens=1840, + completion_tokens=412, + tool_usage_cost=0.02, + ), +) + + +# azure_ai/claude-sonnet-4-5 (public row has no search price, the deployment owns input, output and search rates) +AZURE_AI_CLAUDE_SONNET_4_5_DEPLOYMENT_MESSAGES_WEB_SEARCH: Final = CostTrackingTestCase( + name="azure-ai-claude-sonnet-4-5-deployment_rates_messages_web_search_two_requests", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="azure_ai/claude-sonnet-4-5", + endpoint="/v1/messages", + deployment=Deployment(input_cost_per_token=1e-06, output_cost_per_token=5e-06, search_cost_per_query=0.01), + request={ + "model": "$MODEL", + "max_tokens": 412, + "messages": [{"role": "user", "content": "azure-ai-claude-sonnet-4-5-deployment_rates_messages_web_search_two_requests"}], + "tools": [{"type": "web_search_20250305", "name": "web_search", "max_uses": 3}], + }, + response=JsonResponse( + content_type="application/json", + body={ + "id": "msg_$UNIQUE_ID", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-4-5", + "content": [{"type": "text", "text": "scripted answer b1c2d3e4f5a6"}], + "stop_reason": "end_turn", + "usage": {"input_tokens": 1840, "output_tokens": 412, "server_tool_use": {"web_search_requests": 2}}, + }, + ), + expected=ExactExpected( + spend=0.0239, + input_cost=0.00184, + output_cost=0.00206, + prompt_tokens=1840, + completion_tokens=412, + tool_usage_cost=0.02, + rollups=True, + ), +) + +AZURE_AI_CLAUDE_SONNET_4_5_DEPLOYMENT_MESSAGES_WEB_SEARCH_STREAM: Final = CostTrackingTestCase( + name="azure-ai-claude-sonnet-4-5-deployment_rates_messages_web_search_two_requests_stream", + covers="quota_management.spend_tracking.scripted_wire.logs_cost", + model="azure_ai/claude-sonnet-4-5", + endpoint="/v1/messages", + deployment=Deployment(input_cost_per_token=1e-06, output_cost_per_token=5e-06, search_cost_per_query=0.01), + request={ + "model": "$MODEL", + "max_tokens": 412, + "messages": [{"role": "user", "content": "azure-ai-claude-sonnet-4-5-deployment_rates_messages_web_search_two_requests_stream"}], + "stream": True, + "tools": [{"type": "web_search_20250305", "name": "web_search", "max_uses": 3}], + }, + response=SseResponse( + content_type="text/event-stream", + frames=sse_frames( + { + "type": "message_start", + "message": { + "id": "msg_$REQUEST_ID", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-4-5", + "content": [], + "stop_reason": None, + "usage": {"input_tokens": 1840}, + }, + }, + {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}, + { + "type": "content_block_delta", + "index": 0, + "delta": {"type": "text_delta", "text": "scripted answer c2d3e4f5a6b1"}, + }, + {"type": "content_block_stop", "index": 0}, + { + "type": "message_delta", + "delta": {"stop_reason": "end_turn"}, + "usage": {"output_tokens": 412, "server_tool_use": {"web_search_requests": 2}}, + }, + {"type": "message_stop"}, + ), + ), + expected=ExactExpected( + spend=0.0239, + input_cost=0.00184, + output_cost=0.00206, + prompt_tokens=1840, + completion_tokens=412, + tool_usage_cost=0.02, + ), +) + +AZURE_AI_CLAUDE_SONNET_4_5_DEPLOYMENT_MESSAGES_WEB_SEARCH_ZERO_RATE: Final = CostTrackingTestCase( + name="azure-ai-claude-sonnet-4-5-deployment_zero_search_rate_messages_two_requests", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="azure_ai/claude-sonnet-4-5", + endpoint="/v1/messages", + deployment=Deployment(input_cost_per_token=1e-06, output_cost_per_token=5e-06, search_cost_per_query=0.0), + request={ + "model": "$MODEL", + "max_tokens": 412, + "messages": [{"role": "user", "content": "azure-ai-claude-sonnet-4-5-deployment_zero_search_rate_messages_two_requests"}], + "tools": [{"type": "web_search_20250305", "name": "web_search", "max_uses": 3}], + }, + response=JsonResponse( + content_type="application/json", + body={ + "id": "msg_$REQUEST_ID", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-4-5", + "content": [{"type": "text", "text": "scripted answer d3e4f5a6b1c2"}], + "stop_reason": "end_turn", + "usage": {"input_tokens": 1840, "output_tokens": 412, "server_tool_use": {"web_search_requests": 2}}, + }, + ), + expected=ExactExpected( + spend=0.0039, + input_cost=0.00184, + output_cost=0.00206, + prompt_tokens=1840, + completion_tokens=412, + tool_usage_cost=0.0, + ), +) + +AZURE_AI_CLAUDE_SONNET_4_5_DEPLOYMENT_MESSAGES_NO_SEARCH: Final = CostTrackingTestCase( + name="azure-ai-claude-sonnet-4-5-deployment_rates_messages_no_server_tool_use", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="azure_ai/claude-sonnet-4-5", + endpoint="/v1/messages", + deployment=Deployment(input_cost_per_token=1e-06, output_cost_per_token=5e-06, search_cost_per_query=0.01), + request={ + "model": "$MODEL", + "max_tokens": 412, + "messages": [{"role": "user", "content": "azure-ai-claude-sonnet-4-5-deployment_rates_messages_no_server_tool_use"}], + "tools": [{"type": "web_search_20250305", "name": "web_search", "max_uses": 3}], + }, + response=JsonResponse( + content_type="application/json", + body={ + "id": "msg_$REQUEST_ID", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-4-5", + "content": [{"type": "text", "text": "scripted answer e4f5a6b1c2d3"}], + "stop_reason": "end_turn", + "usage": {"input_tokens": 1840, "output_tokens": 412}, + }, + ), + expected=ExactExpected( + spend=0.0039, + input_cost=0.00184, + output_cost=0.00206, + prompt_tokens=1840, + completion_tokens=412, + tool_usage_cost=0.0, + ), +) + +AZURE_AI_CLAUDE_SONNET_4_5_DEPLOYMENT_CHAT_WEB_SEARCH: Final = CostTrackingTestCase( + name="azure-ai-claude-sonnet-4-5-deployment_rates_chat_web_search_three_requests", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="azure_ai/claude-sonnet-4-5", + deployment=Deployment(input_cost_per_token=1e-06, output_cost_per_token=5e-06, search_cost_per_query=0.01), + request={ + "model": "$MODEL", + "messages": [{"role": "user", "content": "azure-ai-claude-sonnet-4-5-deployment_rates_chat_web_search_three_requests"}], + "stream": False, + "tools": [{"type": "web_search_20250305", "name": "web_search", "max_uses": 5}], + "allowed_openai_params": ["web_search_options"], + }, + response=JsonResponse( + content_type="application/json", + body={ + "id": "msg_$REQUEST_ID", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-4-5", + "content": [{"type": "text", "text": "scripted answer f5a6b1c2d3e4"}], + "stop_reason": "end_turn", + "usage": {"input_tokens": 1840, "output_tokens": 412, "server_tool_use": {"web_search_requests": 3}}, + }, + ), + expected=ExactExpected( + spend=0.0339, + input_cost=0.00184, + output_cost=0.00206, + prompt_tokens=1840, + completion_tokens=412, + tool_usage_cost=0.03, + ), +) + +AZURE_AI_CLAUDE_SONNET_4_5_DEPLOYMENT_CHAT_WEB_SEARCH_STREAM: Final = CostTrackingTestCase( + name="azure-ai-claude-sonnet-4-5-deployment_rates_chat_web_search_two_requests_stream", + covers="quota_management.spend_tracking.scripted_wire.logs_cost", + model="azure_ai/claude-sonnet-4-5", + deployment=Deployment(input_cost_per_token=1e-06, output_cost_per_token=5e-06, search_cost_per_query=0.01), + request={ + "model": "$MODEL", + "messages": [{"role": "user", "content": "azure-ai-claude-sonnet-4-5-deployment_rates_chat_web_search_two_requests_stream"}], + "stream": True, + "tools": [{"type": "web_search_20250305", "name": "web_search", "max_uses": 3}], + "allowed_openai_params": ["web_search_options"], + }, + response=SseResponse( + content_type="text/event-stream", + frames=sse_frames( + { + "type": "message_start", + "message": { + "id": "msg_$REQUEST_ID", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-4-5", + "content": [], + "stop_reason": None, + "usage": {"input_tokens": 1840}, + }, + }, + {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}, + { + "type": "content_block_delta", + "index": 0, + "delta": {"type": "text_delta", "text": "scripted answer a6b1c2d3e4f5"}, + }, + {"type": "content_block_stop", "index": 0}, + { + "type": "message_delta", + "delta": {"stop_reason": "end_turn"}, + "usage": {"output_tokens": 412, "server_tool_use": {"web_search_requests": 2}}, + }, + {"type": "message_stop"}, + ), + ), + expected=ExactExpected( + spend=0.0239, + input_cost=0.00184, + output_cost=0.00206, + prompt_tokens=1840, + completion_tokens=412, + tool_usage_cost=0.02, + ), +) + +AZURE_AI_CLAUDE_SONNET_4_5_DEPLOYMENT_RESPONSES_WEB_SEARCH: Final = CostTrackingTestCase( + name="azure-ai-claude-sonnet-4-5-deployment_rates_responses_web_search_two_requests", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="azure_ai/claude-sonnet-4-5", + endpoint="/v1/responses", + deployment=Deployment(input_cost_per_token=1e-06, output_cost_per_token=5e-06, search_cost_per_query=0.01), + request={ + "model": "$MODEL", + "input": "azure-ai-claude-sonnet-4-5-deployment_rates_responses_web_search_two_requests", + "tools": [{"type": "web_search"}], + }, + response=JsonResponse( + content_type="application/json", + body={ + "id": "msg_$REQUEST_ID", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-4-5", + "content": [{"type": "text", "text": "scripted answer b1c2d3e4f5a7"}], + "stop_reason": "end_turn", + "usage": {"input_tokens": 1840, "output_tokens": 412, "server_tool_use": {"web_search_requests": 2}}, + }, + ), + expected=ExactExpected( + spend=0.0239, + input_cost=0.00184, + output_cost=0.00206, + prompt_tokens=1840, + completion_tokens=412, + tool_usage_cost=0.02, + ), +) + + # groq/qwen/qwen3.8-27b GROQ_QWEN_3_8_JSON: Final = CostTrackingTestCase( name="groq-qwen-3.8-json", @@ -840,31 +1369,6 @@ MISTRAL_MEDIUM_2604_JSON: Final = CostTrackingTestCase( ) -# perplexity/pplx-decider-v1-27b -PERPLEXITY_PPLX_DECIDER_V1_27B_DECISIONS: Final = CostTrackingTestCase( - name="perplexity/pplx-decider-v1-27b-decisions", - covers="quota_management.spend_tracking.decisions_costs", - model="perplexity/pplx-decider-v1-27b", - endpoint="/v1/systemone", - request={ - "model": "$MODEL", - "state": {"source": "cost-tracking"}, - "questions": {"is_defect": {"type": "noul", "instructions": "Is this a defect?"}}, - }, - response=JsonResponse( - content_type="application/json", - body={ - "model": "pplx-decider-v1-27b", - "answers": {"is_defect": {"type": "noul", "noul": 0.9}}, - "usage": {"input_tokens": 367, "output_tokens": 3}, - }, - ), - expected=ExactExpected( - spend=1.468e-05, input_cost=1.468e-05, output_cost=0.0, prompt_tokens=367, completion_tokens=3 - ), -) - - # xai/grok-4.7 XAI_GROK_4_7_INPUT_TEXT: Final = CostTrackingTestCase( name="xai-grok-4.7-input_text", @@ -921,11 +1425,24 @@ CASES: Final[tuple[CostTrackingTestCase, ...]] = ( BEDROCK_INVOKE_HAIKU_STREAM, AZURE_AI_GPT_5_4_MINI_LATEST, AZURE_AI_GPT_5_4_MINI_LATEST_STREAM, + AZURE_AI_CLAUDE_HAIKU_4_5_MESSAGES_WEB_SEARCH, + AZURE_AI_CLAUDE_HAIKU_4_5_MESSAGES_WEB_SEARCH_STREAM, + AZURE_AI_CLAUDE_HAIKU_4_5_MESSAGES_WEB_SEARCH_ONE_REQUEST, + AZURE_AI_CLAUDE_HAIKU_4_5_MESSAGES_NO_SEARCH, + AZURE_AI_CLAUDE_HAIKU_4_5_CHAT_WEB_SEARCH, + AZURE_AI_CLAUDE_HAIKU_4_5_CHAT_WEB_SEARCH_STREAM, + AZURE_AI_CLAUDE_HAIKU_4_5_RESPONSES_WEB_SEARCH, + AZURE_AI_CLAUDE_SONNET_4_5_DEPLOYMENT_MESSAGES_WEB_SEARCH, + AZURE_AI_CLAUDE_SONNET_4_5_DEPLOYMENT_MESSAGES_WEB_SEARCH_STREAM, + AZURE_AI_CLAUDE_SONNET_4_5_DEPLOYMENT_MESSAGES_WEB_SEARCH_ZERO_RATE, + AZURE_AI_CLAUDE_SONNET_4_5_DEPLOYMENT_MESSAGES_NO_SEARCH, + AZURE_AI_CLAUDE_SONNET_4_5_DEPLOYMENT_CHAT_WEB_SEARCH, + AZURE_AI_CLAUDE_SONNET_4_5_DEPLOYMENT_CHAT_WEB_SEARCH_STREAM, + AZURE_AI_CLAUDE_SONNET_4_5_DEPLOYMENT_RESPONSES_WEB_SEARCH, GROQ_QWEN_3_8_JSON, GROQ_QWEN_3_8_STREAM_X_GROQ_RECOUNT, COHERE_COMMAND_A_V2_TOKENS, MISTRAL_MEDIUM_2604_JSON, - PERPLEXITY_PPLX_DECIDER_V1_27B_DECISIONS, XAI_GROK_4_7_INPUT_TEXT, ) diff --git a/tests/integration/cost_calculation/cases/openai.py b/tests/integration/cost_calculation/cases/openai.py index c4e88166388..6268ec2b463 100644 --- a/tests/integration/cost_calculation/cases/openai.py +++ b/tests/integration/cost_calculation/cases/openai.py @@ -2,9 +2,10 @@ PARITY pairs the streamed case with its plain twin where both bill the same row.""" -from typing import Final +from typing import Final, cast from integration.cost_calculation.cost_tracking_case import ( + COST_MAP, BinaryResponse, CostTrackingTestCase, Deployment, @@ -18,7 +19,9 @@ from integration.cost_calculation.cost_tracking_case import ( SseResponse, WavUpload, ) +from pydantic import JsonValue from integration.cost_calculation.stream_parity.case import StreamParityTestCase, sse_frames +from integration.cost_calculation.bridge_cost_parity.case import BridgeCostParityTestCase # gpt-5.3-codex @@ -7160,6 +7163,201 @@ OPENAI_DEPLOYMENT_PRICING_OVERRIDE: Final = CostTrackingTestCase( ), ) +_CC_CUSTOM_MODEL_COST_HEADER_FRAMES: Final = sse_frames( + { + "id": "chatcmpl-$REQUEST_ID", + "object": "chat.completion.chunk", + "created": 1700000000, + "model": "cc-custom-model", + "choices": [ + {"index": 0, "delta": {"role": "assistant", "content": "scripted response"}, "finish_reason": None} + ], + }, + { + "id": "chatcmpl-$REQUEST_ID", + "object": "chat.completion.chunk", + "created": 1700000000, + "model": "cc-custom-model", + "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}], + }, + { + "id": "chatcmpl-$REQUEST_ID", + "object": "chat.completion.chunk", + "created": 1700000000, + "model": "cc-custom-model", + "choices": [], + "usage": {"prompt_tokens": 10, "completion_tokens": 253, "total_tokens": 263}, + }, + done=True, +) + +_CC_CUSTOM_MODEL_RESPONSES_BODY: Final[dict[str, JsonValue]] = { + "id": "resp_$REQUEST_ID", + "object": "response", + "status": "completed", + "created_at": 1700000000, + "model": "cc-custom-model", + "output": [ + { + "type": "message", + "id": "msg_$REQUEST_ID", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "scripted response", "annotations": []}], + } + ], + "usage": { + "input_tokens": 10, + "output_tokens": 253, + "total_tokens": 263, + "input_tokens_details": {"cached_tokens": 0}, + "output_tokens_details": {"reasoning_tokens": 0}, + }, +} + +OPENAI_DEPLOYMENT_ZERO_PRICING_IGNORES_UPSTREAM_COST_HEADER: Final = CostTrackingTestCase( + name="openai-deployment-zero-pricing-ignores-upstream-cost-header", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="gpt-5.6", + deployment=Deployment(model="openai/cc-custom-model", input_cost_per_token=0.0, output_cost_per_token=0.0), + request={ + "model": "$MODEL", + "messages": [{"role": "user", "content": "openai-deployment-zero-pricing-ignores-upstream-cost-header"}], + "stream": False, + }, + response=JsonResponse( + content_type="application/json", + body={ + "id": "chatcmpl-$REQUEST_ID", + "object": "chat.completion", + "model": "cc-custom-model", + "choices": [ + {"index": 0, "message": {"role": "assistant", "content": "scripted response"}, "finish_reason": "stop"} + ], + "usage": {"prompt_tokens": 10, "completion_tokens": 253, "total_tokens": 263}, + }, + headers={"x-litellm-response-cost": "0.000413"}, + ), + expected=ExactExpected(spend=0.0, input_cost=0.0, output_cost=0.0, prompt_tokens=10, completion_tokens=253), +) + +OPENAI_DEPLOYMENT_ZERO_PRICING_IGNORES_UPSTREAM_COST_HEADER_STREAM: Final = CostTrackingTestCase( + name="openai-deployment-zero-pricing-ignores-upstream-cost-header-stream", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="gpt-5.6", + deployment=Deployment(model="openai/cc-custom-model", input_cost_per_token=0.0, output_cost_per_token=0.0), + request={ + "model": "$MODEL", + "messages": [{"role": "user", "content": "openai-deployment-zero-pricing-ignores-upstream-cost-header"}], + "stream": True, + "stream_options": {"include_usage": True}, + }, + response=SseResponse( + content_type="text/event-stream", + frames=_CC_CUSTOM_MODEL_COST_HEADER_FRAMES, + headers={"x-litellm-response-cost": "0.000413"}, + ), + expected=ExactExpected(spend=0.0, input_cost=0.0, output_cost=0.0, prompt_tokens=10, completion_tokens=253), +) + +OPENAI_DEPLOYMENT_ZERO_PRICING_IGNORES_UPSTREAM_COST_HEADER_MESSAGES: Final = CostTrackingTestCase( + name="openai-deployment-zero-pricing-ignores-upstream-cost-header-messages", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="gpt-5.6", + endpoint="/v1/messages", + deployment=Deployment(model="openai/cc-custom-model", input_cost_per_token=0.0, output_cost_per_token=0.0), + request={ + "model": "$MODEL", + "max_tokens": 300, + "messages": [{"role": "user", "content": "openai-deployment-zero-pricing-ignores-upstream-cost-header"}], + }, + response=JsonResponse( + content_type="application/json", + body=_CC_CUSTOM_MODEL_RESPONSES_BODY, + headers={"x-litellm-response-cost": "0.000413"}, + ), + expected=ExactExpected(spend=0.0, input_cost=0.0, output_cost=0.0, prompt_tokens=10, completion_tokens=253), +) + +OPENAI_DEPLOYMENT_ZERO_PRICING_IGNORES_UPSTREAM_COST_HEADER_RESPONSES: Final = CostTrackingTestCase( + name="openai-deployment-zero-pricing-ignores-upstream-cost-header-responses", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="gpt-5.6", + endpoint="/v1/responses", + deployment=Deployment(model="openai/cc-custom-model", input_cost_per_token=0.0, output_cost_per_token=0.0), + request={ + "model": "$MODEL", + "input": "openai-deployment-zero-pricing-ignores-upstream-cost-header", + "stream": False, + }, + response=JsonResponse( + content_type="application/json", + body=_CC_CUSTOM_MODEL_RESPONSES_BODY, + headers={"x-litellm-response-cost": "0.000413"}, + ), + expected=ExactExpected(spend=0.0, input_cost=0.0, output_cost=0.0, prompt_tokens=10, completion_tokens=253), +) + +OPENAI_DEPLOYMENT_PRICING_IGNORES_UPSTREAM_COST_HEADER: Final = CostTrackingTestCase( + name="openai-deployment-pricing-ignores-upstream-cost-header", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="gpt-5.6", + deployment=Deployment(model="openai/cc-custom-model", input_cost_per_token=0.001, output_cost_per_token=0.002), + request={ + "model": "$MODEL", + "messages": [{"role": "user", "content": "openai-deployment-pricing-ignores-upstream-cost-header"}], + "stream": False, + }, + response=JsonResponse( + content_type="application/json", + body={ + "id": "chatcmpl-$UNIQUE_ID", + "object": "chat.completion", + "model": "cc-custom-model", + "choices": [ + {"index": 0, "message": {"role": "assistant", "content": "scripted response"}, "finish_reason": "stop"} + ], + "usage": {"prompt_tokens": 10, "completion_tokens": 253, "total_tokens": 263}, + }, + headers={"x-litellm-response-cost": "0.000413"}, + ), + expected=ExactExpected( + spend=0.516, input_cost=0.01, output_cost=0.506, prompt_tokens=10, completion_tokens=253, rollups=True + ), +) + +GPT_5_6_UPSTREAM_COST_HEADER_BILLS_HEADER_VALUE: Final = CostTrackingTestCase( + name="gpt-5.6-upstream-cost-header-bills-header-value", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="gpt-5.6", + request={ + "model": "$MODEL", + "messages": [{"role": "user", "content": "gpt-5.6-upstream-cost-header-bills-header-value"}], + "stream": False, + }, + response=JsonResponse( + content_type="application/json", + body={ + "id": "chatcmpl-$REQUEST_ID", + "object": "chat.completion", + "model": "gpt-5.6", + "choices": [ + {"index": 0, "message": {"role": "assistant", "content": "scripted response"}, "finish_reason": "stop"} + ], + "usage": {"prompt_tokens": 10, "completion_tokens": 253, "total_tokens": 263}, + }, + headers={"x-litellm-response-cost": "0.000413"}, + ), + expected=ExactExpected( + spend=0.000413, + input_cost=0.0, + output_cost=0.000413, + prompt_tokens=10, + completion_tokens=253, + breakdown_persisted=False, + ), +) + GPT_5_6_UPSTREAM_400_ZERO_SPEND: Final = CostTrackingTestCase( name="gpt-5.6-upstream_400_zero_spend", covers="quota_management.spend_tracking.cost_matrix.logs_cost", @@ -7848,6 +8046,2504 @@ GPT_6_1_SOL_INPUT_TEXT: Final = CostTrackingTestCase( ) +_OPENAI_GPT_6_LUNA_DECISIONS_INPUT_TOKENS: Final = 367 +_OPENAI_GPT_6_LUNA_DECISIONS_OUTPUT_TOKENS: Final = 3 +_OPENAI_GPT_6_LUNA_DECISIONS_INPUT_RATE: Final = cast( + float, + COST_MAP["gpt-6-luna"].input_cost_per_token_decisions, +) +_OPENAI_GPT_6_LUNA_DECISIONS_SPEND: Final = ( + _OPENAI_GPT_6_LUNA_DECISIONS_INPUT_TOKENS * _OPENAI_GPT_6_LUNA_DECISIONS_INPUT_RATE +) +_OPENAI_GPT_6_LUNA_DECISIONS_UPSTREAM_BODY: Final[dict[str, JsonValue]] = { + "model": "gpt-6-luna", + "input": "The package arrived broken.", + "questions": [{"type": "predicate", "name": "is_defect", "instructions": "Is this a defect?"}], +} +_OPENAI_GPT_6_LUNA_DECISIONS_RESPONSE: Final = JsonResponse( + content_type="application/json", + body={ + "model": "gpt-6-luna", + "answers": [{"type": "predicate", "name": "is_defect", "probability": 0.9}], + "usage": { + "input_tokens": _OPENAI_GPT_6_LUNA_DECISIONS_INPUT_TOKENS, + "input_tokens_details": {"cached_tokens": 256, "cache_write_tokens": 64}, + "output_tokens": _OPENAI_GPT_6_LUNA_DECISIONS_OUTPUT_TOKENS, + "output_tokens_details": {"reasoning_tokens": 0}, + "total_tokens": _OPENAI_GPT_6_LUNA_DECISIONS_INPUT_TOKENS + _OPENAI_GPT_6_LUNA_DECISIONS_OUTPUT_TOKENS, + }, + }, +) + +OPENAI_GPT_6_LUNA_DECISIONS: Final = CostTrackingTestCase( + name="openai-gpt-6-luna-decisions-evaluation-pricing", + covers="quota_management.spend_tracking.decisions_costs", + model="gpt-6-luna", + endpoint="/v1/decisions", + deployment=Deployment(model="openai/gpt-6-luna"), + request={ + "model": "$MODEL", + "input": "The package arrived broken.", + "questions": [{"type": "predicate", "name": "is_defect", "instructions": "Is this a defect?"}], + }, + upstream_body=_OPENAI_GPT_6_LUNA_DECISIONS_UPSTREAM_BODY, + response=_OPENAI_GPT_6_LUNA_DECISIONS_RESPONSE, + expected=ExactExpected( + spend=_OPENAI_GPT_6_LUNA_DECISIONS_SPEND, + input_cost=_OPENAI_GPT_6_LUNA_DECISIONS_SPEND, + output_cost=0.0, + prompt_tokens=_OPENAI_GPT_6_LUNA_DECISIONS_INPUT_TOKENS, + completion_tokens=_OPENAI_GPT_6_LUNA_DECISIONS_OUTPUT_TOKENS, + ), +) + +OPENAI_GPT_6_LUNA_SYSTEMONE: Final = CostTrackingTestCase( + name="openai-gpt-6-luna-systemone-evaluation-pricing", + covers="quota_management.spend_tracking.decisions_costs", + model="gpt-6-luna", + endpoint="/v1/systemone", + deployment=Deployment(model="openai/gpt-6-luna"), + request={ + "model": "$MODEL", + "state": "The package arrived broken.", + "questions": {"is_defect": {"type": "noul", "instructions": "Is this a defect?"}}, + }, + upstream_body=_OPENAI_GPT_6_LUNA_DECISIONS_UPSTREAM_BODY, + response=_OPENAI_GPT_6_LUNA_DECISIONS_RESPONSE, + expected=ExactExpected( + spend=_OPENAI_GPT_6_LUNA_DECISIONS_SPEND, + input_cost=_OPENAI_GPT_6_LUNA_DECISIONS_SPEND, + output_cost=0.0, + prompt_tokens=_OPENAI_GPT_6_LUNA_DECISIONS_INPUT_TOKENS, + completion_tokens=_OPENAI_GPT_6_LUNA_DECISIONS_OUTPUT_TOKENS, + ), +) + + +# Above 272k input tokens OpenAI bills the whole request at the long-context column. For decisions +# that column is the `_above_272k_tokens_decisions` fields: 2e-07 for input, cached and cache-write +# tokens alike, output still 0. Mirrors the live 320,162-token call that billed 0.0640324. +_OPENAI_GPT_6_LUNA_DECISIONS_LONG_INPUT_TOKENS: Final = 320162 +_OPENAI_GPT_6_LUNA_DECISIONS_LONG_CACHED_TOKENS: Final = 100000 +_OPENAI_GPT_6_LUNA_DECISIONS_LONG_CACHE_WRITE_TOKENS: Final = 10000 +_OPENAI_GPT_6_LUNA_DECISIONS_LONG_OUTPUT_TOKENS: Final = 3 +_OPENAI_GPT_6_LUNA_DECISIONS_LONG_INPUT_RATE: Final = cast( + float, + COST_MAP["gpt-6-luna"].input_cost_per_token_above_272k_tokens_decisions, +) +_OPENAI_GPT_6_LUNA_DECISIONS_LONG_CACHE_READ_RATE: Final = cast( + float, + COST_MAP["gpt-6-luna"].cache_read_input_token_cost_above_272k_tokens_decisions, +) +_OPENAI_GPT_6_LUNA_DECISIONS_LONG_CACHE_WRITE_RATE: Final = cast( + float, + COST_MAP["gpt-6-luna"].cache_creation_input_token_cost_above_272k_tokens_decisions, +) +_OPENAI_GPT_6_LUNA_DECISIONS_LONG_SPEND: Final = ( + ( + _OPENAI_GPT_6_LUNA_DECISIONS_LONG_INPUT_TOKENS + - _OPENAI_GPT_6_LUNA_DECISIONS_LONG_CACHED_TOKENS + - _OPENAI_GPT_6_LUNA_DECISIONS_LONG_CACHE_WRITE_TOKENS + ) + * _OPENAI_GPT_6_LUNA_DECISIONS_LONG_INPUT_RATE + + _OPENAI_GPT_6_LUNA_DECISIONS_LONG_CACHED_TOKENS * _OPENAI_GPT_6_LUNA_DECISIONS_LONG_CACHE_READ_RATE + + _OPENAI_GPT_6_LUNA_DECISIONS_LONG_CACHE_WRITE_TOKENS * _OPENAI_GPT_6_LUNA_DECISIONS_LONG_CACHE_WRITE_RATE +) +_OPENAI_GPT_6_LUNA_DECISIONS_LONG_RESPONSE: Final = JsonResponse( + content_type="application/json", + body={ + "model": "gpt-6-luna", + "answers": [{"type": "predicate", "name": "is_defect", "probability": 0.84}], + "usage": { + "input_tokens": _OPENAI_GPT_6_LUNA_DECISIONS_LONG_INPUT_TOKENS, + "input_tokens_details": { + "cached_tokens": _OPENAI_GPT_6_LUNA_DECISIONS_LONG_CACHED_TOKENS, + "cache_write_tokens": _OPENAI_GPT_6_LUNA_DECISIONS_LONG_CACHE_WRITE_TOKENS, + }, + "output_tokens": _OPENAI_GPT_6_LUNA_DECISIONS_LONG_OUTPUT_TOKENS, + "output_tokens_details": {"reasoning_tokens": 0}, + "total_tokens": _OPENAI_GPT_6_LUNA_DECISIONS_LONG_INPUT_TOKENS + + _OPENAI_GPT_6_LUNA_DECISIONS_LONG_OUTPUT_TOKENS, + }, + }, +) + +OPENAI_GPT_6_LUNA_DECISIONS_LONG_CONTEXT: Final = CostTrackingTestCase( + name="openai-gpt-6-luna-decisions-above-272k-pricing", + covers="quota_management.spend_tracking.decisions_costs", + model="gpt-6-luna", + endpoint="/v1/decisions", + deployment=Deployment(model="openai/gpt-6-luna"), + request={ + "model": "$MODEL", + "input": "The package arrived broken.", + "questions": [{"type": "predicate", "name": "is_defect", "instructions": "Is this a defect?"}], + }, + upstream_body=_OPENAI_GPT_6_LUNA_DECISIONS_UPSTREAM_BODY, + response=_OPENAI_GPT_6_LUNA_DECISIONS_LONG_RESPONSE, + expected=ExactExpected( + spend=_OPENAI_GPT_6_LUNA_DECISIONS_LONG_SPEND, + input_cost=_OPENAI_GPT_6_LUNA_DECISIONS_LONG_SPEND, + output_cost=0.0, + prompt_tokens=_OPENAI_GPT_6_LUNA_DECISIONS_LONG_INPUT_TOKENS, + completion_tokens=_OPENAI_GPT_6_LUNA_DECISIONS_LONG_OUTPUT_TOKENS, + ), +) + +OPENAI_GPT_6_LUNA_SYSTEMONE_LONG_CONTEXT: Final = CostTrackingTestCase( + name="openai-gpt-6-luna-systemone-above-272k-pricing", + covers="quota_management.spend_tracking.decisions_costs", + model="gpt-6-luna", + endpoint="/v1/systemone", + deployment=Deployment(model="openai/gpt-6-luna"), + request={ + "model": "$MODEL", + "state": "The package arrived broken.", + "questions": {"is_defect": {"type": "noul", "instructions": "Is this a defect?"}}, + }, + upstream_body=_OPENAI_GPT_6_LUNA_DECISIONS_UPSTREAM_BODY, + response=_OPENAI_GPT_6_LUNA_DECISIONS_LONG_RESPONSE, + expected=ExactExpected( + spend=_OPENAI_GPT_6_LUNA_DECISIONS_LONG_SPEND, + input_cost=_OPENAI_GPT_6_LUNA_DECISIONS_LONG_SPEND, + output_cost=0.0, + prompt_tokens=_OPENAI_GPT_6_LUNA_DECISIONS_LONG_INPUT_TOKENS, + completion_tokens=_OPENAI_GPT_6_LUNA_DECISIONS_LONG_OUTPUT_TOKENS, + ), +) + + +GPT_6_SOL_INPUT_TEXT: Final = CostTrackingTestCase( + name="gpt-6-sol-input_text", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="gpt-6-sol", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "d36c648eb5c2 summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": False, + "allowed_openai_params": [], + }, + response=JsonResponse( + content_type="application/json", + body={ + "id": "chatcmpl-$REQUEST_ID", + "object": "chat.completion", + "created": 1791851507, + "model": "gpt-6-sol", + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": "scripted answer d36c648eb5c2", + "refusal": None, + "annotations": [], + }, + "finish_reason": "stop", + } + ], + "usage": { + "prompt_tokens": 1840, + "completion_tokens": 412, + "total_tokens": 2252, + "prompt_tokens_details": {"cached_tokens": 0, "cache_write_tokens": 0, "audio_tokens": 0}, + "completion_tokens_details": { + "reasoning_tokens": 0, + "audio_tokens": 0, + "accepted_prediction_tokens": 0, + "rejected_prediction_tokens": 0, + }, + }, + "service_tier": "default", + "system_fingerprint": None, + }, + ), + expected=ExactExpected( + spend=0.00858, input_cost=0.004048, output_cost=0.004532, prompt_tokens=1840, completion_tokens=412 + ), +) + +GPT_6_SOL_STREAM: Final = CostTrackingTestCase( + name="gpt-6-sol-stream", + covers="quota_management.spend_tracking.scripted_wire.logs_cost", + model="gpt-6-sol", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "b807c011e149 summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": True, + "stream_options": {"include_usage": True}, + "allowed_openai_params": [], + }, + response=SseResponse( + content_type="text/event-stream", + frames=sse_frames( + { + "id": "chatcmpl-$REQUEST_ID", + "object": "chat.completion.chunk", + "created": 1791851658, + "model": "gpt-6-sol", + "service_tier": "default", + "system_fingerprint": None, + "choices": [ + {"index": 0, "delta": {"role": "assistant", "content": "", "refusal": None}, "finish_reason": None} + ], + "usage": None, + }, + { + "id": "chatcmpl-$REQUEST_ID", + "object": "chat.completion.chunk", + "created": 1791851658, + "model": "gpt-6-sol", + "service_tier": "default", + "system_fingerprint": None, + "choices": [{"index": 0, "delta": {"content": "scripted answer b807c011e149"}, "finish_reason": None}], + "usage": None, + }, + { + "id": "chatcmpl-$REQUEST_ID", + "object": "chat.completion.chunk", + "created": 1791851658, + "model": "gpt-6-sol", + "service_tier": "default", + "system_fingerprint": None, + "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}], + "usage": None, + }, + { + "id": "chatcmpl-$REQUEST_ID", + "object": "chat.completion.chunk", + "created": 1791851658, + "model": "gpt-6-sol", + "service_tier": "default", + "system_fingerprint": None, + "choices": [], + "usage": { + "prompt_tokens": 1840, + "completion_tokens": 412, + "total_tokens": 2252, + "prompt_tokens_details": {"cached_tokens": 0, "cache_write_tokens": 0, "audio_tokens": 0}, + "completion_tokens_details": { + "reasoning_tokens": 0, + "audio_tokens": 0, + "accepted_prediction_tokens": 0, + "rejected_prediction_tokens": 0, + }, + }, + }, + ), + ), + expected=ExactExpected( + spend=0.00858, input_cost=0.004048, output_cost=0.004532, prompt_tokens=1840, completion_tokens=412 + ), +) + +GPT_6_SOL_SERVICE_TIER_FLEX: Final = CostTrackingTestCase( + name="gpt-6-sol-service_tier_flex", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="gpt-6-sol", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "42b013573e14 summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": False, + "service_tier": "flex", + "allowed_openai_params": [], + }, + response=JsonResponse( + content_type="application/json", + body={ + "id": "chatcmpl-$REQUEST_ID", + "object": "chat.completion", + "created": 1791851507, + "model": "gpt-6-sol", + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": "scripted answer 42b013573e14", + "refusal": None, + "annotations": [], + }, + "finish_reason": "stop", + } + ], + "usage": { + "prompt_tokens": 1840, + "completion_tokens": 412, + "total_tokens": 2252, + "prompt_tokens_details": {"cached_tokens": 0, "cache_write_tokens": 0, "audio_tokens": 0}, + "completion_tokens_details": { + "reasoning_tokens": 0, + "audio_tokens": 0, + "accepted_prediction_tokens": 0, + "rejected_prediction_tokens": 0, + }, + }, + "service_tier": "flex", + "system_fingerprint": None, + }, + ), + expected=ExactExpected( + spend=0.00429, input_cost=0.002024, output_cost=0.002266, prompt_tokens=1840, completion_tokens=412 + ), +) + +GPT_6_SOL_SERVICE_TIER_PRIORITY: Final = CostTrackingTestCase( + name="gpt-6-sol-service_tier_priority", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="gpt-6-sol", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "f4ef2ca58612 summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": False, + "service_tier": "priority", + "allowed_openai_params": [], + }, + response=JsonResponse( + content_type="application/json", + body={ + "id": "chatcmpl-$REQUEST_ID", + "object": "chat.completion", + "created": 1791851507, + "model": "gpt-6-sol", + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": "scripted answer f4ef2ca58612", + "refusal": None, + "annotations": [], + }, + "finish_reason": "stop", + } + ], + "usage": { + "prompt_tokens": 1840, + "completion_tokens": 412, + "total_tokens": 2252, + "prompt_tokens_details": {"cached_tokens": 0, "cache_write_tokens": 0, "audio_tokens": 0}, + "completion_tokens_details": { + "reasoning_tokens": 0, + "audio_tokens": 0, + "accepted_prediction_tokens": 0, + "rejected_prediction_tokens": 0, + }, + }, + "service_tier": "fast", + "system_fingerprint": None, + }, + ), + expected=ExactExpected( + spend=0.01716, input_cost=0.008096, output_cost=0.009064, prompt_tokens=1840, completion_tokens=412 + ), +) + +GPT_6_SOL_RESPONSES_NATIVE_JSON: Final = CostTrackingTestCase( + name="gpt-6-sol-responses_native_json", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="gpt-6-sol", + endpoint="/v1/responses", + request={"model": "$MODEL", "input": "responses native fixture", "stream": False}, + response=JsonResponse( + content_type="application/json", + body={ + "id": "resp_$REQUEST_ID", + "object": "response", + "created_at": 1791851574, + "status": "completed", + "background": False, + "billing": {"payer": "developer"}, + "completed_at": 1791851574, + "error": None, + "incomplete_details": None, + "instructions": None, + "max_output_tokens": 412, + "model": "gpt-6-sol", + "output": [ + { + "id": "msg_$REQUEST_ID", + "type": "message", + "status": "completed", + "content": [ + {"type": "output_text", "annotations": [], "logprobs": [], "text": "scripted response"} + ], + "phase": "final_answer", + "role": "assistant", + } + ], + "parallel_tool_calls": True, + "previous_response_id": None, + "prompt_cache_key": None, + "prompt_cache_retention": "24h", + "reasoning": {"context": "all_turns", "effort": "medium", "mode": "standard", "summary": None}, + "service_tier": "default", + "store": True, + "temperature": 1.0, + "text": {"format": {"type": "text"}, "verbosity": "medium"}, + "tool_choice": "auto", + "tool_usage": { + "image_gen": { + "input_tokens": 0, + "input_tokens_details": {"image_tokens": 0, "text_tokens": 0}, + "output_tokens": 0, + "output_tokens_details": {"image_tokens": 0, "text_tokens": 0}, + "total_tokens": 0, + }, + "web_search": {"num_requests": 0}, + }, + "tools": [], + "top_logprobs": 0, + "top_p": 0.98, + "truncation": "disabled", + "usage": { + "input_tokens": 1840, + "input_tokens_details": {"cache_write_tokens": 0, "cached_tokens": 0}, + "output_tokens": 412, + "output_tokens_details": {"reasoning_tokens": 0}, + "total_tokens": 2252, + }, + "user": None, + "metadata": {}, + }, + ), + expected=ExactExpected( + spend=0.00858, input_cost=0.004048, output_cost=0.004532, prompt_tokens=1840, completion_tokens=412 + ), +) + +GPT_6_SOL_MESSAGES_BRIDGED_JSON: Final = GPT_6_SOL_RESPONSES_NATIVE_JSON.model_copy( + update={ + "name": "gpt-6-sol-messages_bridged_json", + "endpoint": "/v1/messages", + "request": { + "model": "$MODEL", + "max_tokens": 412, + "messages": [{"role": "user", "content": "responses native fixture"}], + }, + } +) + +GPT_6_LUNA_INPUT_TEXT: Final = CostTrackingTestCase( + name="gpt-6-luna-input_text", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="gpt-6-luna", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "694ee3f43b69 summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": False, + "allowed_openai_params": [], + }, + response=JsonResponse( + content_type="application/json", + body={ + "id": "chatcmpl-$REQUEST_ID", + "object": "chat.completion", + "created": 1791851507, + "model": "gpt-6-luna", + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": "scripted answer 694ee3f43b69", + "refusal": None, + "annotations": [], + }, + "finish_reason": "stop", + } + ], + "usage": { + "prompt_tokens": 1840, + "completion_tokens": 412, + "total_tokens": 2252, + "prompt_tokens_details": {"cached_tokens": 0, "cache_write_tokens": 0, "audio_tokens": 0}, + "completion_tokens_details": { + "reasoning_tokens": 0, + "audio_tokens": 0, + "accepted_prediction_tokens": 0, + "rejected_prediction_tokens": 0, + }, + }, + "service_tier": "default", + "system_fingerprint": None, + }, + ), + expected=ExactExpected( + spend=0.00039, input_cost=0.000184, output_cost=0.000206, prompt_tokens=1840, completion_tokens=412 + ), +) + +GPT_6_LUNA_STREAM: Final = CostTrackingTestCase( + name="gpt-6-luna-stream", + covers="quota_management.spend_tracking.scripted_wire.logs_cost", + model="gpt-6-luna", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "f3f11b8d2818 summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": True, + "stream_options": {"include_usage": True}, + "allowed_openai_params": [], + }, + response=SseResponse( + content_type="text/event-stream", + frames=sse_frames( + { + "id": "chatcmpl-$REQUEST_ID", + "object": "chat.completion.chunk", + "created": 1791851658, + "model": "gpt-6-luna", + "service_tier": "default", + "system_fingerprint": None, + "choices": [ + {"index": 0, "delta": {"role": "assistant", "content": "", "refusal": None}, "finish_reason": None} + ], + "usage": None, + }, + { + "id": "chatcmpl-$REQUEST_ID", + "object": "chat.completion.chunk", + "created": 1791851658, + "model": "gpt-6-luna", + "service_tier": "default", + "system_fingerprint": None, + "choices": [{"index": 0, "delta": {"content": "scripted answer f3f11b8d2818"}, "finish_reason": None}], + "usage": None, + }, + { + "id": "chatcmpl-$REQUEST_ID", + "object": "chat.completion.chunk", + "created": 1791851658, + "model": "gpt-6-luna", + "service_tier": "default", + "system_fingerprint": None, + "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}], + "usage": None, + }, + { + "id": "chatcmpl-$REQUEST_ID", + "object": "chat.completion.chunk", + "created": 1791851658, + "model": "gpt-6-luna", + "service_tier": "default", + "system_fingerprint": None, + "choices": [], + "usage": { + "prompt_tokens": 1840, + "completion_tokens": 412, + "total_tokens": 2252, + "prompt_tokens_details": {"cached_tokens": 0, "cache_write_tokens": 0, "audio_tokens": 0}, + "completion_tokens_details": { + "reasoning_tokens": 0, + "audio_tokens": 0, + "accepted_prediction_tokens": 0, + "rejected_prediction_tokens": 0, + }, + }, + }, + ), + ), + expected=ExactExpected( + spend=0.00039, input_cost=0.000184, output_cost=0.000206, prompt_tokens=1840, completion_tokens=412 + ), +) + +GPT_6_LUNA_SERVICE_TIER_FLEX: Final = CostTrackingTestCase( + name="gpt-6-luna-service_tier_flex", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="gpt-6-luna", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "dd297fa791f8 summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": False, + "service_tier": "flex", + "allowed_openai_params": [], + }, + response=JsonResponse( + content_type="application/json", + body={ + "id": "chatcmpl-$REQUEST_ID", + "object": "chat.completion", + "created": 1791851507, + "model": "gpt-6-luna", + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": "scripted answer dd297fa791f8", + "refusal": None, + "annotations": [], + }, + "finish_reason": "stop", + } + ], + "usage": { + "prompt_tokens": 1840, + "completion_tokens": 412, + "total_tokens": 2252, + "prompt_tokens_details": {"cached_tokens": 0, "cache_write_tokens": 0, "audio_tokens": 0}, + "completion_tokens_details": { + "reasoning_tokens": 0, + "audio_tokens": 0, + "accepted_prediction_tokens": 0, + "rejected_prediction_tokens": 0, + }, + }, + "service_tier": "flex", + "system_fingerprint": None, + }, + ), + expected=ExactExpected( + spend=0.000195, input_cost=0.000092, output_cost=0.000103, prompt_tokens=1840, completion_tokens=412 + ), +) + +GPT_6_LUNA_SERVICE_TIER_PRIORITY: Final = CostTrackingTestCase( + name="gpt-6-luna-service_tier_priority", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="gpt-6-luna", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "c6908010aebc summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": False, + "service_tier": "priority", + "allowed_openai_params": [], + }, + response=JsonResponse( + content_type="application/json", + body={ + "id": "chatcmpl-$REQUEST_ID", + "object": "chat.completion", + "created": 1791851507, + "model": "gpt-6-luna", + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": "scripted answer c6908010aebc", + "refusal": None, + "annotations": [], + }, + "finish_reason": "stop", + } + ], + "usage": { + "prompt_tokens": 1840, + "completion_tokens": 412, + "total_tokens": 2252, + "prompt_tokens_details": {"cached_tokens": 0, "cache_write_tokens": 0, "audio_tokens": 0}, + "completion_tokens_details": { + "reasoning_tokens": 0, + "audio_tokens": 0, + "accepted_prediction_tokens": 0, + "rejected_prediction_tokens": 0, + }, + }, + "service_tier": "fast", + "system_fingerprint": None, + }, + ), + expected=ExactExpected( + spend=0.00078, input_cost=0.000368, output_cost=0.000412, prompt_tokens=1840, completion_tokens=412 + ), +) + +GPT_6_LUNA_RESPONSES_NATIVE_JSON: Final = CostTrackingTestCase( + name="gpt-6-luna-responses_native_json", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="gpt-6-luna", + endpoint="/v1/responses", + request={"model": "$MODEL", "input": "responses native fixture", "stream": False}, + response=JsonResponse( + content_type="application/json", + body={ + "id": "resp_$REQUEST_ID", + "object": "response", + "created_at": 1791851574, + "status": "completed", + "background": False, + "billing": {"payer": "developer"}, + "completed_at": 1791851574, + "error": None, + "incomplete_details": None, + "instructions": None, + "max_output_tokens": 412, + "model": "gpt-6-luna", + "output": [ + { + "id": "msg_$REQUEST_ID", + "type": "message", + "status": "completed", + "content": [ + {"type": "output_text", "annotations": [], "logprobs": [], "text": "scripted response"} + ], + "phase": "final_answer", + "role": "assistant", + } + ], + "parallel_tool_calls": True, + "previous_response_id": None, + "prompt_cache_key": None, + "prompt_cache_retention": "24h", + "reasoning": {"context": "all_turns", "effort": "medium", "mode": "standard", "summary": None}, + "service_tier": "default", + "store": True, + "temperature": 1.0, + "text": {"format": {"type": "text"}, "verbosity": "medium"}, + "tool_choice": "auto", + "tool_usage": { + "image_gen": { + "input_tokens": 0, + "input_tokens_details": {"image_tokens": 0, "text_tokens": 0}, + "output_tokens": 0, + "output_tokens_details": {"image_tokens": 0, "text_tokens": 0}, + "total_tokens": 0, + }, + "web_search": {"num_requests": 0}, + }, + "tools": [], + "top_logprobs": 0, + "top_p": 0.98, + "truncation": "disabled", + "usage": { + "input_tokens": 1840, + "input_tokens_details": {"cache_write_tokens": 0, "cached_tokens": 0}, + "output_tokens": 412, + "output_tokens_details": {"reasoning_tokens": 0}, + "total_tokens": 2252, + }, + "user": None, + "metadata": {}, + }, + ), + expected=ExactExpected( + spend=0.00039, input_cost=0.000184, output_cost=0.000206, prompt_tokens=1840, completion_tokens=412 + ), +) + +GPT_6_ASTRA_INPUT_TEXT: Final = CostTrackingTestCase( + name="gpt-6-astra-input_text", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="gpt-6-astra", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "682dbc7dce00 summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": False, + "allowed_openai_params": [], + }, + response=JsonResponse( + content_type="application/json", + body={ + "id": "chatcmpl-$REQUEST_ID", + "object": "chat.completion", + "created": 1791851507, + "model": "gpt-6-astra", + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": "scripted answer 682dbc7dce00", + "refusal": None, + "annotations": [], + }, + "finish_reason": "stop", + } + ], + "usage": { + "prompt_tokens": 1840, + "completion_tokens": 412, + "total_tokens": 2252, + "prompt_tokens_details": {"cached_tokens": 0, "cache_write_tokens": 0, "audio_tokens": 0}, + "completion_tokens_details": { + "reasoning_tokens": 0, + "audio_tokens": 0, + "accepted_prediction_tokens": 0, + "rejected_prediction_tokens": 0, + }, + }, + "service_tier": "default", + "system_fingerprint": None, + }, + ), + expected=ExactExpected( + spend=0.0468, input_cost=0.02208, output_cost=0.02472, prompt_tokens=1840, completion_tokens=412 + ), +) + +GPT_6_ASTRA_STREAM: Final = CostTrackingTestCase( + name="gpt-6-astra-stream", + covers="quota_management.spend_tracking.scripted_wire.logs_cost", + model="gpt-6-astra", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "867507e67c53 summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": True, + "stream_options": {"include_usage": True}, + "allowed_openai_params": [], + }, + response=SseResponse( + content_type="text/event-stream", + frames=sse_frames( + { + "id": "chatcmpl-$REQUEST_ID", + "object": "chat.completion.chunk", + "created": 1791851658, + "model": "gpt-6-astra", + "service_tier": "default", + "system_fingerprint": None, + "choices": [ + {"index": 0, "delta": {"role": "assistant", "content": "", "refusal": None}, "finish_reason": None} + ], + "usage": None, + }, + { + "id": "chatcmpl-$REQUEST_ID", + "object": "chat.completion.chunk", + "created": 1791851658, + "model": "gpt-6-astra", + "service_tier": "default", + "system_fingerprint": None, + "choices": [{"index": 0, "delta": {"content": "scripted answer 867507e67c53"}, "finish_reason": None}], + "usage": None, + }, + { + "id": "chatcmpl-$REQUEST_ID", + "object": "chat.completion.chunk", + "created": 1791851658, + "model": "gpt-6-astra", + "service_tier": "default", + "system_fingerprint": None, + "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}], + "usage": None, + }, + { + "id": "chatcmpl-$REQUEST_ID", + "object": "chat.completion.chunk", + "created": 1791851658, + "model": "gpt-6-astra", + "service_tier": "default", + "system_fingerprint": None, + "choices": [], + "usage": { + "prompt_tokens": 1840, + "completion_tokens": 412, + "total_tokens": 2252, + "prompt_tokens_details": {"cached_tokens": 0, "cache_write_tokens": 0, "audio_tokens": 0}, + "completion_tokens_details": { + "reasoning_tokens": 0, + "audio_tokens": 0, + "accepted_prediction_tokens": 0, + "rejected_prediction_tokens": 0, + }, + }, + }, + ), + ), + expected=ExactExpected( + spend=0.0468, input_cost=0.02208, output_cost=0.02472, prompt_tokens=1840, completion_tokens=412 + ), +) + +GPT_6_ASTRA_SERVICE_TIER_FLEX: Final = CostTrackingTestCase( + name="gpt-6-astra-service_tier_flex", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="gpt-6-astra", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "72512aa3c14a summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": False, + "service_tier": "flex", + "allowed_openai_params": [], + }, + response=JsonResponse( + content_type="application/json", + body={ + "id": "chatcmpl-$REQUEST_ID", + "object": "chat.completion", + "created": 1791851507, + "model": "gpt-6-astra", + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": "scripted answer 72512aa3c14a", + "refusal": None, + "annotations": [], + }, + "finish_reason": "stop", + } + ], + "usage": { + "prompt_tokens": 1840, + "completion_tokens": 412, + "total_tokens": 2252, + "prompt_tokens_details": {"cached_tokens": 0, "cache_write_tokens": 0, "audio_tokens": 0}, + "completion_tokens_details": { + "reasoning_tokens": 0, + "audio_tokens": 0, + "accepted_prediction_tokens": 0, + "rejected_prediction_tokens": 0, + }, + }, + "service_tier": "flex", + "system_fingerprint": None, + }, + ), + expected=ExactExpected( + spend=0.0234, input_cost=0.01104, output_cost=0.01236, prompt_tokens=1840, completion_tokens=412 + ), +) + +GPT_6_ASTRA_SERVICE_TIER_PRIORITY: Final = CostTrackingTestCase( + name="gpt-6-astra-service_tier_priority", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="gpt-6-astra", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "53b84772ae29 summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": False, + "service_tier": "priority", + "allowed_openai_params": [], + }, + response=JsonResponse( + content_type="application/json", + body={ + "id": "chatcmpl-$REQUEST_ID", + "object": "chat.completion", + "created": 1791851507, + "model": "gpt-6-astra", + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": "scripted answer 53b84772ae29", + "refusal": None, + "annotations": [], + }, + "finish_reason": "stop", + } + ], + "usage": { + "prompt_tokens": 1840, + "completion_tokens": 412, + "total_tokens": 2252, + "prompt_tokens_details": {"cached_tokens": 0, "cache_write_tokens": 0, "audio_tokens": 0}, + "completion_tokens_details": { + "reasoning_tokens": 0, + "audio_tokens": 0, + "accepted_prediction_tokens": 0, + "rejected_prediction_tokens": 0, + }, + }, + "service_tier": "fast", + "system_fingerprint": None, + }, + ), + expected=ExactExpected( + spend=0.0936, input_cost=0.04416, output_cost=0.04944, prompt_tokens=1840, completion_tokens=412 + ), +) + +GPT_6_ASTRA_RESPONSES_NATIVE_JSON: Final = CostTrackingTestCase( + name="gpt-6-astra-responses_native_json", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="gpt-6-astra", + endpoint="/v1/responses", + request={"model": "$MODEL", "input": "responses native fixture", "stream": False}, + response=JsonResponse( + content_type="application/json", + body={ + "id": "resp_$REQUEST_ID", + "object": "response", + "created_at": 1791851574, + "status": "completed", + "background": False, + "billing": {"payer": "developer"}, + "completed_at": 1791851574, + "error": None, + "incomplete_details": None, + "instructions": None, + "max_output_tokens": 412, + "model": "gpt-6-astra", + "output": [ + { + "id": "msg_$REQUEST_ID", + "type": "message", + "status": "completed", + "content": [ + {"type": "output_text", "annotations": [], "logprobs": [], "text": "scripted response"} + ], + "phase": "final_answer", + "role": "assistant", + } + ], + "parallel_tool_calls": True, + "previous_response_id": None, + "prompt_cache_key": None, + "prompt_cache_retention": "24h", + "reasoning": {"context": "all_turns", "effort": "medium", "mode": "standard", "summary": None}, + "service_tier": "default", + "store": True, + "temperature": 1.0, + "text": {"format": {"type": "text"}, "verbosity": "medium"}, + "tool_choice": "auto", + "tool_usage": { + "image_gen": { + "input_tokens": 0, + "input_tokens_details": {"image_tokens": 0, "text_tokens": 0}, + "output_tokens": 0, + "output_tokens_details": {"image_tokens": 0, "text_tokens": 0}, + "total_tokens": 0, + }, + "web_search": {"num_requests": 0}, + }, + "tools": [], + "top_logprobs": 0, + "top_p": 0.98, + "truncation": "disabled", + "usage": { + "input_tokens": 1840, + "input_tokens_details": {"cache_write_tokens": 0, "cached_tokens": 0}, + "output_tokens": 412, + "output_tokens_details": {"reasoning_tokens": 0}, + "total_tokens": 2252, + }, + "user": None, + "metadata": {}, + }, + ), + expected=ExactExpected( + spend=0.0468, input_cost=0.02208, output_cost=0.02472, prompt_tokens=1840, completion_tokens=412 + ), +) + +GPT_6_1_SOL_STREAM: Final = CostTrackingTestCase( + name="gpt-6.1-sol-stream", + covers="quota_management.spend_tracking.scripted_wire.logs_cost", + model="gpt-6.1-sol", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "fe2e0078a155 summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": True, + "stream_options": {"include_usage": True}, + "allowed_openai_params": [], + }, + response=SseResponse( + content_type="text/event-stream", + frames=sse_frames( + { + "id": "chatcmpl-$REQUEST_ID", + "object": "chat.completion.chunk", + "created": 1791851658, + "model": "gpt-6.1-sol", + "service_tier": "default", + "system_fingerprint": None, + "choices": [ + {"index": 0, "delta": {"role": "assistant", "content": "", "refusal": None}, "finish_reason": None} + ], + "usage": None, + }, + { + "id": "chatcmpl-$REQUEST_ID", + "object": "chat.completion.chunk", + "created": 1791851658, + "model": "gpt-6.1-sol", + "service_tier": "default", + "system_fingerprint": None, + "choices": [{"index": 0, "delta": {"content": "scripted answer fe2e0078a155"}, "finish_reason": None}], + "usage": None, + }, + { + "id": "chatcmpl-$REQUEST_ID", + "object": "chat.completion.chunk", + "created": 1791851658, + "model": "gpt-6.1-sol", + "service_tier": "default", + "system_fingerprint": None, + "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}], + "usage": None, + }, + { + "id": "chatcmpl-$REQUEST_ID", + "object": "chat.completion.chunk", + "created": 1791851658, + "model": "gpt-6.1-sol", + "service_tier": "default", + "system_fingerprint": None, + "choices": [], + "usage": { + "prompt_tokens": 1840, + "completion_tokens": 412, + "total_tokens": 2252, + "prompt_tokens_details": {"cached_tokens": 0, "cache_write_tokens": 0, "audio_tokens": 0}, + "completion_tokens_details": { + "reasoning_tokens": 0, + "audio_tokens": 0, + "accepted_prediction_tokens": 0, + "rejected_prediction_tokens": 0, + }, + }, + }, + ), + ), + expected=ExactExpected( + spend=0.0078, input_cost=0.00368, output_cost=0.00412, prompt_tokens=1840, completion_tokens=412 + ), +) + +GPT_6_1_SOL_SERVICE_TIER_FLEX: Final = CostTrackingTestCase( + name="gpt-6.1-sol-service_tier_flex", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="gpt-6.1-sol", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "28201777a291 summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": False, + "service_tier": "flex", + "allowed_openai_params": [], + }, + response=JsonResponse( + content_type="application/json", + body={ + "id": "chatcmpl-$REQUEST_ID", + "object": "chat.completion", + "created": 1791851507, + "model": "gpt-6.1-sol", + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": "scripted answer 28201777a291", + "refusal": None, + "annotations": [], + }, + "finish_reason": "stop", + } + ], + "usage": { + "prompt_tokens": 1840, + "completion_tokens": 412, + "total_tokens": 2252, + "prompt_tokens_details": {"cached_tokens": 0, "cache_write_tokens": 0, "audio_tokens": 0}, + "completion_tokens_details": { + "reasoning_tokens": 0, + "audio_tokens": 0, + "accepted_prediction_tokens": 0, + "rejected_prediction_tokens": 0, + }, + }, + "service_tier": "flex", + "system_fingerprint": None, + }, + ), + expected=ExactExpected( + spend=0.0039, input_cost=0.00184, output_cost=0.00206, prompt_tokens=1840, completion_tokens=412 + ), +) + +GPT_6_1_SOL_SERVICE_TIER_PRIORITY: Final = CostTrackingTestCase( + name="gpt-6.1-sol-service_tier_priority", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="gpt-6.1-sol", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "27c98f26478e summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": False, + "service_tier": "priority", + "allowed_openai_params": [], + }, + response=JsonResponse( + content_type="application/json", + body={ + "id": "chatcmpl-$REQUEST_ID", + "object": "chat.completion", + "created": 1791851507, + "model": "gpt-6.1-sol", + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": "scripted answer 27c98f26478e", + "refusal": None, + "annotations": [], + }, + "finish_reason": "stop", + } + ], + "usage": { + "prompt_tokens": 1840, + "completion_tokens": 412, + "total_tokens": 2252, + "prompt_tokens_details": {"cached_tokens": 0, "cache_write_tokens": 0, "audio_tokens": 0}, + "completion_tokens_details": { + "reasoning_tokens": 0, + "audio_tokens": 0, + "accepted_prediction_tokens": 0, + "rejected_prediction_tokens": 0, + }, + }, + "service_tier": "fast", + "system_fingerprint": None, + }, + ), + expected=ExactExpected( + spend=0.0156, input_cost=0.00736, output_cost=0.00824, prompt_tokens=1840, completion_tokens=412 + ), +) + +GPT_6_1_SOL_RESPONSES_NATIVE_JSON: Final = CostTrackingTestCase( + name="gpt-6.1-sol-responses_native_json", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="gpt-6.1-sol", + endpoint="/v1/responses", + request={"model": "$MODEL", "input": "responses native fixture", "stream": False}, + response=JsonResponse( + content_type="application/json", + body={ + "id": "resp_$REQUEST_ID", + "object": "response", + "created_at": 1791851574, + "status": "completed", + "background": False, + "billing": {"payer": "developer"}, + "completed_at": 1791851574, + "error": None, + "incomplete_details": None, + "instructions": None, + "max_output_tokens": 412, + "model": "gpt-6.1-sol", + "output": [ + { + "id": "msg_$REQUEST_ID", + "type": "message", + "status": "completed", + "content": [ + {"type": "output_text", "annotations": [], "logprobs": [], "text": "scripted response"} + ], + "phase": "final_answer", + "role": "assistant", + } + ], + "parallel_tool_calls": True, + "previous_response_id": None, + "prompt_cache_key": None, + "prompt_cache_retention": "24h", + "reasoning": {"context": "all_turns", "effort": "medium", "mode": "standard", "summary": None}, + "service_tier": "default", + "store": True, + "temperature": 1.0, + "text": {"format": {"type": "text"}, "verbosity": "medium"}, + "tool_choice": "auto", + "tool_usage": { + "image_gen": { + "input_tokens": 0, + "input_tokens_details": {"image_tokens": 0, "text_tokens": 0}, + "output_tokens": 0, + "output_tokens_details": {"image_tokens": 0, "text_tokens": 0}, + "total_tokens": 0, + }, + "web_search": {"num_requests": 0}, + }, + "tools": [], + "top_logprobs": 0, + "top_p": 0.98, + "truncation": "disabled", + "usage": { + "input_tokens": 1840, + "input_tokens_details": {"cache_write_tokens": 0, "cached_tokens": 0}, + "output_tokens": 412, + "output_tokens_details": {"reasoning_tokens": 0}, + "total_tokens": 2252, + }, + "user": None, + "metadata": {}, + }, + ), + expected=ExactExpected( + spend=0.0078, input_cost=0.00368, output_cost=0.00412, prompt_tokens=1840, completion_tokens=412 + ), +) + +GPT_6_SOL_TOOL_CALL: Final = CostTrackingTestCase( + name="gpt-6-sol-tool_call", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="gpt-6-sol", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "b2ce2f555743 summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": False, + "tools": [ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get the current weather and a short forecast for a city.", + "parameters": { + "type": "object", + "properties": { + "city": {"type": "string", "description": "City name"}, + "days": {"type": "integer", "description": "Forecast horizon in days"}, + "units": {"type": "string", "enum": ["metric", "imperial"]}, + }, + "required": ["city"], + }, + }, + } + ], + "tool_choice": "auto", + "reasoning_effort": "none", + "allowed_openai_params": ["tool_choice", "reasoning_effort"], + }, + response=JsonResponse( + content_type="application/json", + body={ + "id": "chatcmpl-$REQUEST_ID", + "object": "chat.completion", + "created": 1791855452, + "model": "gpt-6-sol", + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_$REQUEST_ID", + "type": "function", + "function": { + "name": "get_weather", + "arguments": '{"city":"Berlin","days":7,"units":"metric"}', + }, + } + ], + "refusal": None, + "annotations": [], + }, + "finish_reason": "tool_calls", + } + ], + "usage": { + "prompt_tokens": 1840, + "completion_tokens": 412, + "total_tokens": 2252, + "prompt_tokens_details": {"cached_tokens": 0, "cache_write_tokens": 0, "audio_tokens": 0}, + "completion_tokens_details": { + "reasoning_tokens": 0, + "audio_tokens": 0, + "accepted_prediction_tokens": 0, + "rejected_prediction_tokens": 0, + }, + }, + "service_tier": "default", + "system_fingerprint": None, + }, + ), + expected=ExactExpected( + spend=0.00858, input_cost=0.004048, output_cost=0.004532, prompt_tokens=1840, completion_tokens=412 + ), +) + +GPT_6_SOL_REASONING: Final = CostTrackingTestCase( + name="gpt-6-sol-reasoning", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="gpt-6-sol", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "3646649ee64b summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": False, + "reasoning_effort": "medium", + "allowed_openai_params": ["reasoning_effort"], + }, + response=JsonResponse( + content_type="application/json", + body={ + "id": "chatcmpl-$REQUEST_ID", + "object": "chat.completion", + "created": 1791855452, + "model": "gpt-6-sol", + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": "scripted answer 3646649ee64b", + "refusal": None, + "annotations": [], + }, + "finish_reason": "stop", + } + ], + "usage": { + "prompt_tokens": 1240, + "completion_tokens": 4040, + "total_tokens": 5280, + "prompt_tokens_details": {"cached_tokens": 0, "cache_write_tokens": 0, "audio_tokens": 0}, + "completion_tokens_details": { + "reasoning_tokens": 3480, + "audio_tokens": 0, + "accepted_prediction_tokens": 0, + "rejected_prediction_tokens": 0, + }, + }, + "service_tier": "default", + "system_fingerprint": None, + }, + ), + expected=ExactExpected( + spend=0.047168, input_cost=0.002728, output_cost=0.04444, prompt_tokens=1240, completion_tokens=4040 + ), +) + +GPT_6_SOL_STREAM_NO_USAGE: Final = CostTrackingTestCase( + name="gpt-6-sol-stream_no_usage", + covers="quota_management.spend_tracking.scripted_wire.logs_cost", + model="gpt-6-sol", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "e026eaee0c4e summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": True, + "allowed_openai_params": [], + }, + response=SseResponse( + content_type="text/event-stream", + frames=sse_frames( + { + "id": "chatcmpl-$REQUEST_ID", + "object": "chat.completion.chunk", + "created": 1791855376, + "model": "gpt-6-sol", + "service_tier": "default", + "system_fingerprint": None, + "choices": [ + {"index": 0, "delta": {"role": "assistant", "content": "", "refusal": None}, "finish_reason": None} + ], + "obfuscation": "5A3jk", + }, + { + "id": "chatcmpl-$REQUEST_ID", + "object": "chat.completion.chunk", + "created": 1791855376, + "model": "gpt-6-sol", + "service_tier": "default", + "system_fingerprint": None, + "choices": [{"index": 0, "delta": {"content": "scripted answer e026eaee0c4e"}, "finish_reason": None}], + "obfuscation": "Q2x", + }, + { + "id": "chatcmpl-$REQUEST_ID", + "object": "chat.completion.chunk", + "created": 1791855376, + "model": "gpt-6-sol", + "service_tier": "default", + "system_fingerprint": None, + "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}], + "obfuscation": "v37dpurTBq", + }, + ), + ), + expected=RecountExpected(recount=RecountRates(input_cost_per_token=2.2e-06, output_cost_per_token=1.1e-05)), +) + +GPT_6_LUNA_TOOL_CALL: Final = CostTrackingTestCase( + name="gpt-6-luna-tool_call", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="gpt-6-luna", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "29a8fbc85239 summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": False, + "tools": [ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get the current weather and a short forecast for a city.", + "parameters": { + "type": "object", + "properties": { + "city": {"type": "string", "description": "City name"}, + "days": {"type": "integer", "description": "Forecast horizon in days"}, + "units": {"type": "string", "enum": ["metric", "imperial"]}, + }, + "required": ["city"], + }, + }, + } + ], + "tool_choice": "auto", + "reasoning_effort": "none", + "allowed_openai_params": ["tool_choice", "reasoning_effort"], + }, + response=JsonResponse( + content_type="application/json", + body={ + "id": "chatcmpl-$REQUEST_ID", + "object": "chat.completion", + "created": 1791855452, + "model": "gpt-6-luna", + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_$REQUEST_ID", + "type": "function", + "function": { + "name": "get_weather", + "arguments": '{"city":"Berlin","days":7,"units":"metric"}', + }, + } + ], + "refusal": None, + "annotations": [], + }, + "finish_reason": "tool_calls", + } + ], + "usage": { + "prompt_tokens": 1840, + "completion_tokens": 412, + "total_tokens": 2252, + "prompt_tokens_details": {"cached_tokens": 0, "cache_write_tokens": 0, "audio_tokens": 0}, + "completion_tokens_details": { + "reasoning_tokens": 0, + "audio_tokens": 0, + "accepted_prediction_tokens": 0, + "rejected_prediction_tokens": 0, + }, + }, + "service_tier": "default", + "system_fingerprint": None, + }, + ), + expected=ExactExpected( + spend=0.00039, input_cost=0.000184, output_cost=0.000206, prompt_tokens=1840, completion_tokens=412 + ), +) + +GPT_6_LUNA_REASONING: Final = CostTrackingTestCase( + name="gpt-6-luna-reasoning", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="gpt-6-luna", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "ae40dff71266 summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": False, + "reasoning_effort": "medium", + "allowed_openai_params": ["reasoning_effort"], + }, + response=JsonResponse( + content_type="application/json", + body={ + "id": "chatcmpl-$REQUEST_ID", + "object": "chat.completion", + "created": 1791855452, + "model": "gpt-6-luna", + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": "scripted answer ae40dff71266", + "refusal": None, + "annotations": [], + }, + "finish_reason": "stop", + } + ], + "usage": { + "prompt_tokens": 1240, + "completion_tokens": 4040, + "total_tokens": 5280, + "prompt_tokens_details": {"cached_tokens": 0, "cache_write_tokens": 0, "audio_tokens": 0}, + "completion_tokens_details": { + "reasoning_tokens": 3480, + "audio_tokens": 0, + "accepted_prediction_tokens": 0, + "rejected_prediction_tokens": 0, + }, + }, + "service_tier": "default", + "system_fingerprint": None, + }, + ), + expected=ExactExpected( + spend=0.002144, input_cost=0.000124, output_cost=0.00202, prompt_tokens=1240, completion_tokens=4040 + ), +) + +GPT_6_LUNA_STREAM_NO_USAGE: Final = CostTrackingTestCase( + name="gpt-6-luna-stream_no_usage", + covers="quota_management.spend_tracking.scripted_wire.logs_cost", + model="gpt-6-luna", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "e8f84b81110b summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": True, + "allowed_openai_params": [], + }, + response=SseResponse( + content_type="text/event-stream", + frames=sse_frames( + { + "id": "chatcmpl-$REQUEST_ID", + "object": "chat.completion.chunk", + "created": 1791855376, + "model": "gpt-6-luna", + "service_tier": "default", + "system_fingerprint": None, + "choices": [ + {"index": 0, "delta": {"role": "assistant", "content": "", "refusal": None}, "finish_reason": None} + ], + "obfuscation": "5A3jk", + }, + { + "id": "chatcmpl-$REQUEST_ID", + "object": "chat.completion.chunk", + "created": 1791855376, + "model": "gpt-6-luna", + "service_tier": "default", + "system_fingerprint": None, + "choices": [{"index": 0, "delta": {"content": "scripted answer e8f84b81110b"}, "finish_reason": None}], + "obfuscation": "Q2x", + }, + { + "id": "chatcmpl-$REQUEST_ID", + "object": "chat.completion.chunk", + "created": 1791855376, + "model": "gpt-6-luna", + "service_tier": "default", + "system_fingerprint": None, + "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}], + "obfuscation": "v37dpurTBq", + }, + ), + ), + expected=RecountExpected(recount=RecountRates(input_cost_per_token=1e-07, output_cost_per_token=5e-07)), +) + +GPT_6_ASTRA_REASONING: Final = CostTrackingTestCase( + name="gpt-6-astra-reasoning", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="gpt-6-astra", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "5ff48be12807 summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": False, + "reasoning_effort": "medium", + "allowed_openai_params": ["reasoning_effort"], + }, + response=JsonResponse( + content_type="application/json", + body={ + "id": "chatcmpl-$REQUEST_ID", + "object": "chat.completion", + "created": 1791855452, + "model": "gpt-6-astra", + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": "scripted answer 5ff48be12807", + "refusal": None, + "annotations": [], + }, + "finish_reason": "stop", + } + ], + "usage": { + "prompt_tokens": 1240, + "completion_tokens": 4040, + "total_tokens": 5280, + "prompt_tokens_details": {"cached_tokens": 0, "cache_write_tokens": 0, "audio_tokens": 0}, + "completion_tokens_details": { + "reasoning_tokens": 3480, + "audio_tokens": 0, + "accepted_prediction_tokens": 0, + "rejected_prediction_tokens": 0, + }, + }, + "service_tier": "default", + "system_fingerprint": None, + }, + ), + expected=ExactExpected( + spend=0.25728, input_cost=0.01488, output_cost=0.2424, prompt_tokens=1240, completion_tokens=4040 + ), +) + +GPT_6_ASTRA_STREAM_NO_USAGE: Final = CostTrackingTestCase( + name="gpt-6-astra-stream_no_usage", + covers="quota_management.spend_tracking.scripted_wire.logs_cost", + model="gpt-6-astra", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "8508ebd1992c summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": True, + "allowed_openai_params": [], + }, + response=SseResponse( + content_type="text/event-stream", + frames=sse_frames( + { + "id": "chatcmpl-$REQUEST_ID", + "object": "chat.completion.chunk", + "created": 1791855376, + "model": "gpt-6-astra", + "service_tier": "default", + "system_fingerprint": None, + "choices": [ + {"index": 0, "delta": {"role": "assistant", "content": "", "refusal": None}, "finish_reason": None} + ], + "obfuscation": "5A3jk", + }, + { + "id": "chatcmpl-$REQUEST_ID", + "object": "chat.completion.chunk", + "created": 1791855376, + "model": "gpt-6-astra", + "service_tier": "default", + "system_fingerprint": None, + "choices": [{"index": 0, "delta": {"content": "scripted answer 8508ebd1992c"}, "finish_reason": None}], + "obfuscation": "Q2x", + }, + { + "id": "chatcmpl-$REQUEST_ID", + "object": "chat.completion.chunk", + "created": 1791855376, + "model": "gpt-6-astra", + "service_tier": "default", + "system_fingerprint": None, + "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}], + "obfuscation": "v37dpurTBq", + }, + ), + ), + expected=RecountExpected(recount=RecountRates(input_cost_per_token=1.2e-05, output_cost_per_token=6e-05)), +) + +GPT_6_1_SOL_REASONING: Final = CostTrackingTestCase( + name="gpt-6.1-sol-reasoning", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="gpt-6.1-sol", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "6450817d8662 summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": False, + "reasoning_effort": "medium", + "allowed_openai_params": ["reasoning_effort"], + }, + response=JsonResponse( + content_type="application/json", + body={ + "id": "chatcmpl-$REQUEST_ID", + "object": "chat.completion", + "created": 1791855452, + "model": "gpt-6.1-sol", + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": "scripted answer 6450817d8662", + "refusal": None, + "annotations": [], + }, + "finish_reason": "stop", + } + ], + "usage": { + "prompt_tokens": 1240, + "completion_tokens": 4040, + "total_tokens": 5280, + "prompt_tokens_details": {"cached_tokens": 0, "cache_write_tokens": 0, "audio_tokens": 0}, + "completion_tokens_details": { + "reasoning_tokens": 3480, + "audio_tokens": 0, + "accepted_prediction_tokens": 0, + "rejected_prediction_tokens": 0, + }, + }, + "service_tier": "default", + "system_fingerprint": None, + }, + ), + expected=ExactExpected( + spend=0.04288, input_cost=0.00248, output_cost=0.0404, prompt_tokens=1240, completion_tokens=4040 + ), +) + +GPT_6_1_SOL_STREAM_NO_USAGE: Final = CostTrackingTestCase( + name="gpt-6.1-sol-stream_no_usage", + covers="quota_management.spend_tracking.scripted_wire.logs_cost", + model="gpt-6.1-sol", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "8b76805f8920 summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": True, + "allowed_openai_params": [], + }, + response=SseResponse( + content_type="text/event-stream", + frames=sse_frames( + { + "id": "chatcmpl-$REQUEST_ID", + "object": "chat.completion.chunk", + "created": 1791855376, + "model": "gpt-6.1-sol", + "service_tier": "default", + "system_fingerprint": None, + "choices": [ + {"index": 0, "delta": {"role": "assistant", "content": "", "refusal": None}, "finish_reason": None} + ], + "obfuscation": "5A3jk", + }, + { + "id": "chatcmpl-$REQUEST_ID", + "object": "chat.completion.chunk", + "created": 1791855376, + "model": "gpt-6.1-sol", + "service_tier": "default", + "system_fingerprint": None, + "choices": [{"index": 0, "delta": {"content": "scripted answer 8b76805f8920"}, "finish_reason": None}], + "obfuscation": "Q2x", + }, + { + "id": "chatcmpl-$REQUEST_ID", + "object": "chat.completion.chunk", + "created": 1791855376, + "model": "gpt-6.1-sol", + "service_tier": "default", + "system_fingerprint": None, + "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}], + "obfuscation": "v37dpurTBq", + }, + ), + ), + expected=RecountExpected(recount=RecountRates(input_cost_per_token=2e-06, output_cost_per_token=1e-05)), +) + +GPT_6_SOL_UPSTREAM_404_ZERO_SPEND: Final = CostTrackingTestCase( + name="gpt-6-sol-upstream_404_zero_spend", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="gpt-6-sol", + request={ + "model": "$MODEL", + "messages": [{"role": "user", "content": "scripted upstream failure 404"}], + "stream": False, + }, + response=JsonResponse( + content_type="application/json", + body={ + "error": { + "message": "The model `scripted-missing-model` does not exist or you do not have access to it.", + "type": "invalid_request_error", + "param": None, + "code": "model_not_found", + } + }, + status=404, + ), + expected=FailureExpected(failure=FailureDetails(status=404)), +) + +GPT_6_SOL_INPUT_TEXT_ROLLUPS: Final = CostTrackingTestCase( + name="gpt-6-sol-input_text_rollups", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="gpt-6-sol", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "0a5d8710e0d8 summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": False, + "allowed_openai_params": [], + }, + response=JsonResponse( + content_type="application/json", + body={ + "id": "chatcmpl-$UNIQUE_ID", + "object": "chat.completion", + "created": 1791855452, + "model": "gpt-6-sol", + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": "scripted answer 0a5d8710e0d8", + "refusal": None, + "annotations": [], + }, + "finish_reason": "stop", + } + ], + "usage": { + "prompt_tokens": 1840, + "completion_tokens": 412, + "total_tokens": 2252, + "prompt_tokens_details": {"cached_tokens": 0, "cache_write_tokens": 0, "audio_tokens": 0}, + "completion_tokens_details": { + "reasoning_tokens": 0, + "audio_tokens": 0, + "accepted_prediction_tokens": 0, + "rejected_prediction_tokens": 0, + }, + }, + "service_tier": "default", + "system_fingerprint": None, + }, + ), + expected=ExactExpected( + spend=0.00858, + input_cost=0.004048, + output_cost=0.004532, + prompt_tokens=1840, + completion_tokens=412, + rollups=True, + ), +) + +GPT_6_LUNA_MESSAGES_BRIDGED_JSON: Final = GPT_6_LUNA_RESPONSES_NATIVE_JSON.model_copy( + update={ + "name": "gpt-6-luna-messages_bridged_json", + "endpoint": "/v1/messages", + "request": { + "model": "$MODEL", + "max_tokens": 412, + "messages": [{"role": "user", "content": "responses native fixture"}], + }, + } +) + +GPT_6_ASTRA_MESSAGES_BRIDGED_JSON: Final = GPT_6_ASTRA_RESPONSES_NATIVE_JSON.model_copy( + update={ + "name": "gpt-6-astra-messages_bridged_json", + "endpoint": "/v1/messages", + "request": { + "model": "$MODEL", + "max_tokens": 412, + "messages": [{"role": "user", "content": "responses native fixture"}], + }, + } +) + +GPT_6_1_SOL_MESSAGES_BRIDGED_JSON: Final = GPT_6_1_SOL_RESPONSES_NATIVE_JSON.model_copy( + update={ + "name": "gpt-6.1-sol-messages_bridged_json", + "endpoint": "/v1/messages", + "request": { + "model": "$MODEL", + "max_tokens": 412, + "messages": [{"role": "user", "content": "responses native fixture"}], + }, + } +) + CASES: Final[tuple[CostTrackingTestCase, ...]] = ( GPT_5_3_CODEX_INPUT_TEXT, GPT_5_3_CODEX_CACHE_READ, @@ -7965,6 +10661,12 @@ CASES: Final[tuple[CostTrackingTestCase, ...]] = ( GPT_5_6_RESPONSES_SERVICE_TIER_FLEX, GPT_5_6_RESPONSES_SERVICE_TIER_PRIORITY, OPENAI_DEPLOYMENT_PRICING_OVERRIDE, + OPENAI_DEPLOYMENT_ZERO_PRICING_IGNORES_UPSTREAM_COST_HEADER, + OPENAI_DEPLOYMENT_ZERO_PRICING_IGNORES_UPSTREAM_COST_HEADER_STREAM, + OPENAI_DEPLOYMENT_ZERO_PRICING_IGNORES_UPSTREAM_COST_HEADER_MESSAGES, + OPENAI_DEPLOYMENT_ZERO_PRICING_IGNORES_UPSTREAM_COST_HEADER_RESPONSES, + OPENAI_DEPLOYMENT_PRICING_IGNORES_UPSTREAM_COST_HEADER, + GPT_5_6_UPSTREAM_COST_HEADER_BILLS_HEADER_VALUE, GPT_5_6_UPSTREAM_400_ZERO_SPEND, GPT_5_6_UPSTREAM_401_ZERO_SPEND, GPT_5_6_UPSTREAM_500_STREAM_REQUEST_ZERO_SPEND, @@ -7980,6 +10682,42 @@ CASES: Final[tuple[CostTrackingTestCase, ...]] = ( GPT_5_6_CHAT_REQUEST_TO_EMBEDDING_ENTRY, GPT_5_6_CLIENT_DISCONNECT_MID_STREAM, GPT_6_1_SOL_INPUT_TEXT, + OPENAI_GPT_6_LUNA_DECISIONS, + OPENAI_GPT_6_LUNA_SYSTEMONE, + OPENAI_GPT_6_LUNA_DECISIONS_LONG_CONTEXT, + OPENAI_GPT_6_LUNA_SYSTEMONE_LONG_CONTEXT, + GPT_6_SOL_INPUT_TEXT, + GPT_6_SOL_STREAM, + GPT_6_SOL_SERVICE_TIER_FLEX, + GPT_6_SOL_SERVICE_TIER_PRIORITY, + GPT_6_SOL_RESPONSES_NATIVE_JSON, + GPT_6_LUNA_INPUT_TEXT, + GPT_6_LUNA_STREAM, + GPT_6_LUNA_SERVICE_TIER_FLEX, + GPT_6_LUNA_SERVICE_TIER_PRIORITY, + GPT_6_LUNA_RESPONSES_NATIVE_JSON, + GPT_6_ASTRA_INPUT_TEXT, + GPT_6_ASTRA_STREAM, + GPT_6_ASTRA_SERVICE_TIER_FLEX, + GPT_6_ASTRA_SERVICE_TIER_PRIORITY, + GPT_6_ASTRA_RESPONSES_NATIVE_JSON, + GPT_6_1_SOL_STREAM, + GPT_6_1_SOL_SERVICE_TIER_FLEX, + GPT_6_1_SOL_SERVICE_TIER_PRIORITY, + GPT_6_1_SOL_RESPONSES_NATIVE_JSON, + GPT_6_SOL_MESSAGES_BRIDGED_JSON, + GPT_6_SOL_TOOL_CALL, + GPT_6_SOL_REASONING, + GPT_6_SOL_STREAM_NO_USAGE, + GPT_6_LUNA_TOOL_CALL, + GPT_6_LUNA_REASONING, + GPT_6_LUNA_STREAM_NO_USAGE, + GPT_6_ASTRA_REASONING, + GPT_6_ASTRA_STREAM_NO_USAGE, + GPT_6_1_SOL_REASONING, + GPT_6_1_SOL_STREAM_NO_USAGE, + GPT_6_SOL_UPSTREAM_404_ZERO_SPEND, + GPT_6_SOL_INPUT_TEXT_ROLLUPS, ) PARITY: Final[tuple[StreamParityTestCase, ...]] = ( @@ -8004,4 +10742,15 @@ PARITY: Final[tuple[StreamParityTestCase, ...]] = ( StreamParityTestCase(plain=GPT_5_5_PRO_TOOL_CALL, streamed=GPT_5_5_PRO_STREAM_TOOL_CALL), StreamParityTestCase(plain=GPT_5_6_RESPONSE_MODEL_OVERRIDE, streamed=GPT_5_6_STREAM_RESPONSE_MODEL_OVERRIDE), StreamParityTestCase(plain=GPT_5_6_TOOL_CALL, streamed=GPT_5_6_STREAM_TOOL_CALL), + StreamParityTestCase(plain=GPT_6_SOL_INPUT_TEXT, streamed=GPT_6_SOL_STREAM), + StreamParityTestCase(plain=GPT_6_LUNA_INPUT_TEXT, streamed=GPT_6_LUNA_STREAM), + StreamParityTestCase(plain=GPT_6_ASTRA_INPUT_TEXT, streamed=GPT_6_ASTRA_STREAM), + StreamParityTestCase(plain=GPT_6_1_SOL_INPUT_TEXT, streamed=GPT_6_1_SOL_STREAM), +) + +BRIDGE_COST_PARITY: Final[tuple[BridgeCostParityTestCase, ...]] = ( + BridgeCostParityTestCase(native=GPT_6_SOL_RESPONSES_NATIVE_JSON, bridged=GPT_6_SOL_MESSAGES_BRIDGED_JSON), + BridgeCostParityTestCase(native=GPT_6_LUNA_RESPONSES_NATIVE_JSON, bridged=GPT_6_LUNA_MESSAGES_BRIDGED_JSON), + BridgeCostParityTestCase(native=GPT_6_ASTRA_RESPONSES_NATIVE_JSON, bridged=GPT_6_ASTRA_MESSAGES_BRIDGED_JSON), + BridgeCostParityTestCase(native=GPT_6_1_SOL_RESPONSES_NATIVE_JSON, bridged=GPT_6_1_SOL_MESSAGES_BRIDGED_JSON), ) diff --git a/tests/integration/cost_calculation/catalog.py b/tests/integration/cost_calculation/catalog.py index fe7367aed8a..e83ba186cb2 100644 --- a/tests/integration/cost_calculation/catalog.py +++ b/tests/integration/cost_calculation/catalog.py @@ -10,6 +10,7 @@ from integration.cost_calculation.cases import ( azure, batch_realtime, bedrock_converse, + decisions, fireworks_ai, gemini, misc, @@ -29,6 +30,7 @@ from integration.cost_calculation.cost_tracking_case import ( RecountExpected, SseResponse, ) +from integration.cost_calculation.bridge_cost_parity.case import BridgeCostParityTestCase from integration.cost_calculation.stream_parity.case import StreamParityTestCase _PROVIDER_MODULES: Final = ( @@ -36,6 +38,7 @@ _PROVIDER_MODULES: Final = ( anthropic, gemini, bedrock_converse, + decisions, vertex_ai, azure, fireworks_ai, @@ -48,7 +51,10 @@ PARITY_CASES: Final[tuple[StreamParityTestCase, ...]] = tuple( ) BATCH_CASES: Final[tuple[BatchCostCase, ...]] = batch_realtime.BATCH_CASES REALTIME_CASES: Final[tuple[RealtimeCostCase, ...]] = batch_realtime.REALTIME_CASES -_PARITY_LEGS: Final = frozenset(leg.name for pair in PARITY_CASES for leg in (pair.plain, pair.streamed)) +BRIDGE_COST_PARITY_CASES: Final[tuple[BridgeCostParityTestCase, ...]] = openai.BRIDGE_COST_PARITY +_PARITY_LEGS: Final = frozenset(leg.name for pair in PARITY_CASES for leg in (pair.plain, pair.streamed)) | frozenset( + leg.name for pair in BRIDGE_COST_PARITY_CASES for leg in (pair.native, pair.bridged) +) STANDALONE_CASES: Final[tuple[CostTrackingTestCase, ...]] = tuple( case for case in CASES if case.name not in _PARITY_LEGS ) diff --git a/tests/integration/cost_calculation/conftest.py b/tests/integration/cost_calculation/conftest.py index 7a75320a70f..574d8aaf9af 100644 --- a/tests/integration/cost_calculation/conftest.py +++ b/tests/integration/cost_calculation/conftest.py @@ -86,10 +86,7 @@ def approx_equal(actual: float, expected: float) -> bool: def assert_total_is_sum_of_components(row: CostRow, breakdown: CostBreakdown, context: str) -> None: - total: Final = sum( - cost or 0.0 - for cost in (breakdown.input_cost, breakdown.output_cost, breakdown.tool_usage_cost) - ) + total: Final = sum(cost or 0.0 for cost in (breakdown.input_cost, breakdown.output_cost, breakdown.tool_usage_cost)) assert breakdown.total_cost is not None and approx_equal(breakdown.total_cost, total), ( f"{context}: total_cost {breakdown.total_cost} != input_cost {breakdown.input_cost} " f"+ output_cost {breakdown.output_cost} + tool_usage_cost {breakdown.tool_usage_cost} " @@ -112,7 +109,7 @@ def poll_cost_row(key: str) -> CostRow: def read() -> CostRow | None: rows: Final = read_rows( - 'SELECT spend, status, metadata, prompt_tokens, completion_tokens, model_id, call_type ' + "SELECT spend, status, metadata, prompt_tokens, completion_tokens, model_id, call_type " 'FROM "LiteLLM_SpendLogs" WHERE api_key=%s', (digest,), ) @@ -126,7 +123,7 @@ def poll_cost_row(key: str) -> CostRow: def read_rows_now(key: str) -> tuple[CostRow, ...]: digest: Final = sha256(key.encode()).hexdigest() rows: Final = read_rows( - 'SELECT spend, status, metadata, prompt_tokens, completion_tokens, model_id, call_type ' + "SELECT spend, status, metadata, prompt_tokens, completion_tokens, model_id, call_type " 'FROM "LiteLLM_SpendLogs" WHERE api_key=%s ORDER BY "startTime"', (digest,), ) @@ -178,12 +175,12 @@ def poll_rollups( (end_user_id,), ) daily_user_rows: Final = read_rows( - 'SELECT spend, prompt_tokens, completion_tokens, api_requests ' + "SELECT spend, prompt_tokens, completion_tokens, api_requests " 'FROM "LiteLLM_DailyUserSpend" WHERE user_id=%s AND api_key=%s AND date=CURRENT_DATE::text', (user_id, digest), ) daily_team_rows: Final = read_rows( - 'SELECT spend, prompt_tokens, completion_tokens, api_requests ' + "SELECT spend, prompt_tokens, completion_tokens, api_requests " 'FROM "LiteLLM_DailyTeamSpend" WHERE team_id=%s AND api_key=%s AND date=CURRENT_DATE::text', (team_id, digest), ) @@ -232,11 +229,7 @@ def poll_failure_row(key: str) -> FailureRow: (digest,), ) return next( - ( - parsed - for row in rows - if (parsed := FailureRow.model_validate(row)).status == "failure" - ), + (parsed for row in rows if (parsed := FailureRow.model_validate(row)).status == "failure"), None, ) @@ -247,11 +240,15 @@ def poll_failure_row(key: str) -> FailureRow: @functools.cache def _vertex_private_key_pem() -> str: - return rsa.generate_private_key(public_exponent=65537, key_size=2048).private_bytes( - serialization.Encoding.PEM, - serialization.PrivateFormat.PKCS8, - serialization.NoEncryption(), - ).decode() + return ( + rsa.generate_private_key(public_exponent=65537, key_size=2048) + .private_bytes( + serialization.Encoding.PEM, + serialization.PrivateFormat.PKCS8, + serialization.NoEncryption(), + ) + .decode() + ) def _vertex_service_account_json(url: str) -> str: @@ -310,6 +307,17 @@ def register_scenario_deployment( if case.deployment is not None else {} ), + **( + { + "search_context_cost_per_query": { + "search_context_size_low": case.deployment.search_cost_per_query, + "search_context_size_medium": case.deployment.search_cost_per_query, + "search_context_size_high": case.deployment.search_cost_per_query, + } + } + if case.deployment is not None and case.deployment.search_cost_per_query is not None + else {} + ), **( {"vertex_credentials": _vertex_service_account_json(control_url)} if case.rates.litellm_provider.startswith("vertex_ai") @@ -318,15 +326,13 @@ def register_scenario_deployment( } created: Final = scenario.gateway.post( "/model/new", - JSON_OBJECT.validate_python({ - "model_name": registered_model_name, - "litellm_params": parameters, - "model_info": ( - {"base_model": case.base_model} - if case.base_model is not None - else {} - ), - }), + JSON_OBJECT.validate_python( + { + "model_name": registered_model_name, + "litellm_params": parameters, + "model_info": ({"base_model": case.base_model} if case.base_model is not None else {}), + } + ), ) identity: Final = string_value(object_value(created["model_info"])["id"]) scenario.cleanups.callback(scenario.delete_model, identity) diff --git a/tests/integration/cost_calculation/cost_map.py b/tests/integration/cost_calculation/cost_map.py index 2b597aaadb8..ac51f96a0ec 100644 --- a/tests/integration/cost_calculation/cost_map.py +++ b/tests/integration/cost_calculation/cost_map.py @@ -32,6 +32,31 @@ COST_MAP_ENTRIES: Final[dict[str, dict[str, JsonValue]]] = { "input_cost_per_token": 3.2e-06, "output_cost_per_token": 1.6e-05, }, + "gpt-6-luna": { + "output_cost_per_token_priority": 1e-06, + "input_cost_per_token_priority": 2e-07, + "output_cost_per_token_flex": 2.5e-07, + "input_cost_per_token_flex": 5e-08, + "cache_creation_input_token_cost": 1.25e-07, + "cache_creation_input_token_cost_decisions": 1e-07, + "cache_creation_input_token_cost_above_272k_tokens": 2.5e-07, + "cache_creation_input_token_cost_above_272k_tokens_decisions": 2e-07, + "cache_read_input_token_cost": 1e-08, + "cache_read_input_token_cost_decisions": 1e-07, + "cache_read_input_token_cost_above_272k_tokens": 2e-08, + "cache_read_input_token_cost_above_272k_tokens_decisions": 2e-07, + "input_cost_per_token": 1e-07, + "input_cost_per_token_decisions": 1e-07, + "input_cost_per_token_above_272k_tokens": 2e-07, + "input_cost_per_token_above_272k_tokens_decisions": 2e-07, + "litellm_provider": "openai", + "max_input_tokens": 922000, + "mode": "chat", + "output_cost_per_token": 5e-07, + "output_cost_per_token_decisions": 0.0, + "output_cost_per_token_above_272k_tokens": 7.5e-07, + "output_cost_per_token_above_272k_tokens_decisions": 0.0, + }, "perplexity/sonar-next": { "litellm_provider": "perplexity", "mode": "chat", @@ -43,8 +68,39 @@ COST_MAP_ENTRIES: Final[dict[str, dict[str, JsonValue]]] = { "search_context_size_high": 0.012, }, }, + "typesafe/jev-latest": { + "litellm_provider": "typesafe", + "mode": "evaluation", + "input_cost_per_token": 4.2e-08, + "output_cost_per_token": 0.0, + }, + "fireworks_ai/accounts/fireworks/routers/rune": { + "litellm_provider": "fireworks_ai", + "mode": "evaluation", + "input_cost_per_token": 4.2e-08, + "cache_read_input_token_cost": 8.4e-09, + "output_cost_per_token": 0.0, + }, + "cloudflare/@cf/cloudflare/clef-flash": { + "litellm_provider": "cloudflare", + "mode": "evaluation", + "input_cost_per_token": 9e-08, + "output_cost_per_token": 0.0, + }, + "databricks/databricks-openjev-qwen35-4b": { + "litellm_provider": "databricks", + "mode": "evaluation", + "input_cost_per_token": 0.0, + "output_cost_per_token": 0.0, + }, + "openrouter/typesafe/jev-1.13": { + "litellm_provider": "openrouter", + "mode": "evaluation", + "input_cost_per_token": 0.0, + "output_cost_per_token": 0.0, + }, "perplexity/pplx-decider-v1-27b": { - "input_cost_per_token": 4e-08, + "input_cost_per_token": 2e-08, "litellm_provider": "perplexity", "max_input_tokens": 262144, "mode": "evaluation", @@ -552,6 +608,24 @@ COST_MAP_ENTRIES: Final[dict[str, dict[str, JsonValue]]] = { "input_cost_per_token": 7.5e-07, "output_cost_per_token": 4.5e-06, }, + "azure_ai/claude-haiku-4-5": { + "litellm_provider": "azure_ai", + "mode": "chat", + "input_cost_per_token": 1e-06, + "output_cost_per_token": 5e-06, + "search_context_cost_per_query": { + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01, + "search_context_size_high": 0.01, + }, + "web_search_billing_unit": "per_query", + }, + "azure_ai/claude-sonnet-4-5": { + "litellm_provider": "azure_ai", + "mode": "chat", + "input_cost_per_token": 3e-06, + "output_cost_per_token": 1.5e-05, + }, "groq/qwen/qwen3.8-27b": { "litellm_provider": "groq", "mode": "chat", @@ -723,6 +797,7 @@ COST_MAP_ENTRIES: Final[dict[str, dict[str, JsonValue]]] = { "supports_function_calling": True, }, "us.anthropic.claude-opus-5-5": { + "cache_creation_input_token_cost_above_1hr": 8.8e-06, "litellm_provider": "bedrock_converse", "mode": "chat", "input_cost_per_token": 4.4e-06, @@ -735,6 +810,7 @@ COST_MAP_ENTRIES: Final[dict[str, dict[str, JsonValue]]] = { "supports_function_calling": True, }, "claude-opus-5-5": { + "cache_creation_input_token_cost_above_1hr": 8e-06, "litellm_provider": "anthropic", "mode": "chat", "input_cost_per_token": 4e-06, @@ -747,6 +823,10 @@ COST_MAP_ENTRIES: Final[dict[str, dict[str, JsonValue]]] = { "supports_function_calling": True, }, "gpt-6.1-sol": { + "output_cost_per_token_priority": 2e-05, + "input_cost_per_token_priority": 4e-06, + "output_cost_per_token_flex": 5e-06, + "input_cost_per_token_flex": 1e-06, "litellm_provider": "openai", "mode": "chat", "input_cost_per_token": 2e-06, @@ -757,4 +837,228 @@ COST_MAP_ENTRIES: Final[dict[str, dict[str, JsonValue]]] = { "max_output_tokens": 128000, "supports_function_calling": True, }, + "claude-sonnet-5-5": { + "cache_creation_input_token_cost": 4.125e-06, + "cache_creation_input_token_cost_above_1hr": 6.6e-06, + "cache_read_input_token_cost": 3.3e-07, + "input_cost_per_token": 3.3e-06, + "litellm_provider": "anthropic", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 1.65e-05, + "supports_function_calling": True, + }, + "claude-haiku-5-5": { + "cache_creation_input_token_cost": 1.375e-06, + "cache_creation_input_token_cost_above_1hr": 2.2e-06, + "cache_read_input_token_cost": 1.1e-07, + "input_cost_per_token": 1.1e-06, + "litellm_provider": "anthropic", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 5.5e-06, + "supports_function_calling": True, + }, + "claude-fable-5": { + "cache_creation_input_token_cost": 1.875e-05, + "cache_creation_input_token_cost_above_1hr": 3e-05, + "cache_read_input_token_cost": 1.5e-06, + "input_cost_per_token": 1.5e-05, + "litellm_provider": "anthropic", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 7.5e-05, + "supports_function_calling": True, + }, + "claude-fable-5-1": { + "cache_creation_input_token_cost": 2e-05, + "cache_creation_input_token_cost_above_1hr": 3.2e-05, + "cache_read_input_token_cost": 1.6e-06, + "input_cost_per_token": 1.6e-05, + "litellm_provider": "anthropic", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 8e-05, + "supports_function_calling": True, + }, + "us.anthropic.claude-sonnet-5": { + "cache_creation_input_token_cost": 4.5e-06, + "cache_creation_input_token_cost_above_1hr": 7.2e-06, + "cache_read_input_token_cost": 3.6e-07, + "input_cost_per_token": 3.6e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 1.8e-05, + "supports_function_calling": True, + }, + "us.anthropic.claude-sonnet-5-5": { + "cache_creation_input_token_cost": 4.625e-06, + "cache_creation_input_token_cost_above_1hr": 7.4e-06, + "cache_read_input_token_cost": 3.7e-07, + "input_cost_per_token": 3.7e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 1.85e-05, + "supports_function_calling": True, + }, + "us.anthropic.claude-opus-5": { + "cache_creation_input_token_cost": 7.5e-06, + "cache_creation_input_token_cost_above_1hr": 1.2e-05, + "cache_read_input_token_cost": 6e-07, + "input_cost_per_token": 6e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 3e-05, + "supports_function_calling": True, + }, + "us.anthropic.claude-haiku-5-5": { + "cache_creation_input_token_cost": 1.5e-06, + "cache_creation_input_token_cost_above_1hr": 2.4e-06, + "cache_read_input_token_cost": 1.2e-07, + "input_cost_per_token": 1.2e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 6e-06, + "supports_function_calling": True, + }, + "us.anthropic.claude-fable-5": { + "cache_creation_input_token_cost": 2.125e-05, + "cache_creation_input_token_cost_above_1hr": 3.4e-05, + "cache_read_input_token_cost": 1.7e-06, + "input_cost_per_token": 1.7e-05, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 8.5e-05, + "supports_function_calling": True, + }, + "us.anthropic.claude-fable-5-1": { + "cache_creation_input_token_cost": 2.25e-05, + "cache_creation_input_token_cost_above_1hr": 3.6e-05, + "cache_read_input_token_cost": 1.8e-06, + "input_cost_per_token": 1.8e-05, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 9e-05, + "supports_function_calling": True, + }, + "gpt-6-sol": { + "cache_read_input_token_cost": 2.2e-07, + "input_cost_per_token": 2.2e-06, + "input_cost_per_token_flex": 1.1e-06, + "input_cost_per_token_priority": 4.4e-06, + "litellm_provider": "openai", + "max_input_tokens": 922000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 1.1e-05, + "output_cost_per_token_flex": 5.5e-06, + "output_cost_per_token_priority": 2.2e-05, + "supports_function_calling": True, + }, + "gpt-6-astra": { + "cache_read_input_token_cost": 1.2e-06, + "input_cost_per_token": 1.2e-05, + "input_cost_per_token_flex": 6e-06, + "input_cost_per_token_priority": 2.4e-05, + "litellm_provider": "openai", + "max_input_tokens": 922000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 6e-05, + "output_cost_per_token_flex": 3e-05, + "output_cost_per_token_priority": 0.00012, + "supports_function_calling": True, + }, + "azure/gpt-6-sol": { + "cache_read_input_token_cost": 2.3e-07, + "input_cost_per_token": 2.3e-06, + "input_cost_per_token_flex": 1.15e-06, + "input_cost_per_token_priority": 4.6e-06, + "litellm_provider": "azure", + "max_input_tokens": 922000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 1.15e-05, + "output_cost_per_token_flex": 5.75e-06, + "output_cost_per_token_priority": 2.3e-05, + "supports_function_calling": True, + "supports_none_reasoning_effort": True, + }, + "azure/gpt-6-luna": { + "cache_read_input_token_cost": 1.3e-08, + "input_cost_per_token": 1.3e-07, + "input_cost_per_token_flex": 6.5e-08, + "input_cost_per_token_priority": 2.6e-07, + "litellm_provider": "azure", + "max_input_tokens": 922000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 6.5e-07, + "output_cost_per_token_flex": 3.25e-07, + "output_cost_per_token_priority": 1.3e-06, + "supports_function_calling": True, + "supports_none_reasoning_effort": True, + }, + "azure/gpt-6-astra": { + "cache_read_input_token_cost": 1.25e-06, + "input_cost_per_token": 1.25e-05, + "input_cost_per_token_flex": 6.25e-06, + "input_cost_per_token_priority": 2.5e-05, + "litellm_provider": "azure", + "max_input_tokens": 922000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 6.25e-05, + "output_cost_per_token_flex": 3.125e-05, + "output_cost_per_token_priority": 0.000125, + "supports_function_calling": True, + "supports_none_reasoning_effort": True, + }, + "azure/gpt-6.1-sol": { + "cache_read_input_token_cost": 2.4e-07, + "input_cost_per_token": 2.4e-06, + "input_cost_per_token_flex": 1.2e-06, + "input_cost_per_token_priority": 4.8e-06, + "litellm_provider": "azure", + "max_input_tokens": 922000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 1.2e-05, + "output_cost_per_token_flex": 6e-06, + "output_cost_per_token_priority": 2.4e-05, + "supports_function_calling": True, + "supports_none_reasoning_effort": True, + }, } diff --git a/tests/integration/cost_calculation/cost_tracking_case.py b/tests/integration/cost_calculation/cost_tracking_case.py index a146401f3ae..714c2bf3730 100644 --- a/tests/integration/cost_calculation/cost_tracking_case.py +++ b/tests/integration/cost_calculation/cost_tracking_case.py @@ -43,9 +43,22 @@ class CostMapEntry(BaseModel): max_input_tokens: int | None = None max_output_tokens: int | None = None supports_function_calling: bool | None = None + supports_none_reasoning_effort: bool | None = None input_cost_per_token: float | None = None + input_cost_per_token_decisions: float | None = None input_cost_per_query: float | None = None output_cost_per_token: float | None = None + output_cost_per_token_decisions: float | None = None + cache_read_input_token_cost_decisions: float | None = None + cache_creation_input_token_cost_decisions: float | None = None + input_cost_per_token_above_272k_tokens: float | None = None + input_cost_per_token_above_272k_tokens_decisions: float | None = None + output_cost_per_token_above_272k_tokens: float | None = None + output_cost_per_token_above_272k_tokens_decisions: float | None = None + cache_read_input_token_cost_above_272k_tokens: float | None = None + cache_read_input_token_cost_above_272k_tokens_decisions: float | None = None + cache_creation_input_token_cost_above_272k_tokens: float | None = None + cache_creation_input_token_cost_above_272k_tokens_decisions: float | None = None input_cost_per_token_batches: float | None = None output_cost_per_token_batches: float | None = None output_cost_per_image_token_batches: float | None = None @@ -93,6 +106,7 @@ class Deployment(BaseModel): base_model: str | None = None input_cost_per_token: float | None = None output_cost_per_token: float | None = None + search_cost_per_query: float | None = None class WavUpload(BaseModel): @@ -117,6 +131,7 @@ class JsonResponse(BaseModel): content_type: Literal["application/json"] body: dict[str, JsonValue] status: int = 200 + headers: dict[str, str] = Field(default_factory=dict) class SseResponse(BaseModel): @@ -125,6 +140,7 @@ class SseResponse(BaseModel): content_type: Literal["text/event-stream"] frames: tuple[str, ...] frame_delay_ms: int = Field(default=0, ge=0) + headers: dict[str, str] = Field(default_factory=dict) class EventStreamEvent(BaseModel): @@ -250,6 +266,7 @@ class CostTrackingTestCase(BaseModel): "/v1/images/generations", "/v1/images/edits", "/v1/systemone", + "/v1/decisions", ] | Annotated[str, Field(pattern=r"^/(gemini|anthropic|bedrock)/")] ) = "/v1/chat/completions" @@ -260,6 +277,7 @@ class CostTrackingTestCase(BaseModel): expected: Expected fallback_from: StoredResponse | None = None disconnect_after_frames: int | None = Field(default=None, ge=1) + upstream_body: dict[str, JsonValue] | None = None @property def rates(self) -> CostMapEntry: @@ -315,6 +333,8 @@ class CostTrackingTestCase(BaseModel): def reports_provider_cost(self) -> bool: if not isinstance(self.response, JsonResponse): return False + if "x-litellm-response-cost" in self.response.headers: + return True usage: Final = self.response.body.get("usage") return isinstance(usage, dict) and isinstance(usage.get("cost"), (int, float)) @@ -504,6 +524,9 @@ _PROVIDER_PREFIXES: Final[Mapping[str, str]] = MappingProxyType( "dashscope": "", "openrouter": "", "perplexity": "", + "typesafe": "typesafe", + "cloudflare": "cloudflare", + "databricks": "databricks", "deepseek": "", "xai": "", "azure_ai": "azure_ai", @@ -548,6 +571,9 @@ _LITELLM_PARAMS: Final[Mapping[str, Mapping[str, str]]] = MappingProxyType( "openai": MappingProxyType({}), "dashscope": MappingProxyType({}), "openrouter": MappingProxyType({}), + "typesafe": MappingProxyType({}), + "cloudflare": MappingProxyType({}), + "databricks": MappingProxyType({}), "perplexity": MappingProxyType({}), "deepseek": MappingProxyType({}), "xai": MappingProxyType({}), diff --git a/tests/integration/cost_calculation/stream_parity/runner.py b/tests/integration/cost_calculation/stream_parity/runner.py index 17ce7058c14..044e082e6d0 100644 --- a/tests/integration/cost_calculation/stream_parity/runner.py +++ b/tests/integration/cost_calculation/stream_parity/runner.py @@ -1,16 +1,19 @@ from collections import Counter +from collections.abc import Mapping from hashlib import sha256 +from types import MappingProxyType from typing import Final from integration._support.client import Gateway, Scenario from integration.cost_calculation.assertions import assert_exact from integration.cost_calculation.conftest import ( + CostRow, poll_rows, poll_rows_where, read_rows_now, register_scenario_deployment, ) -from integration.cost_calculation.cost_tracking_case import ExactExpected +from integration.cost_calculation.cost_tracking_case import CostTrackingTestCase, ExactExpected from integration.cost_calculation.stream_parity.bases import openai_responses from integration.cost_calculation.stream_parity.case import StreamParityTestCase @@ -28,43 +31,49 @@ def flush_spend_logs(scenario: Scenario, gateway: Gateway, key: str) -> str: return flush.identity -def assert_stream_parity(case: StreamParityTestCase, gateway: Gateway) -> None: - expected: Final = case.plain.expected - assert isinstance(expected, ExactExpected), case.id +def bill_legs(case_id: str, legs: tuple[CostTrackingTestCase, ...], gateway: Gateway) -> Mapping[str, CostRow]: + """Send every leg through one key, wait for exactly one SpendLogs row per leg, and check each row against its leg's + expected row.""" with gateway.scenario() as scenario: key: Final = scenario.key() deployments: Final = { leg.name: register_scenario_deployment(scenario, leg, sha256(leg.name.encode()).hexdigest()[:12], key) - for leg in (case.plain, case.streamed) + for leg in legs } responses: Final = { leg.name: gateway.request( "POST", leg.endpoint, {**leg.request, "model": deployments[leg.name].model_name}, key=key ) - for leg in (case.plain, case.streamed) + for leg in legs } - for leg in (case.plain, case.streamed): + for leg in legs: assert responses[leg.name].status_code == 200, (leg.name, responses[leg.name].text[:400]) - poll_rows(key, 2) + poll_rows(key, len(legs)) flush_identity: Final = flush_spend_logs(scenario, gateway, key) rows: Final = read_rows_now(key) rows_per_deployment: Final = Counter(row.model_id for row in rows) - expected_rows: Final = {deployments[leg.name].identity: 1 for leg in (case.plain, case.streamed)} + expected_rows: Final = {deployments[leg.name].identity: 1 for leg in legs} assert rows_per_deployment == {**expected_rows, flush_identity: 1}, ( - f"{case.id}: expected exactly one SpendLogs row per leg plus the flush row, got {dict(rows_per_deployment)}", + f"{case_id}: expected exactly one SpendLogs row per leg plus the flush row, got {dict(rows_per_deployment)}", rows, ) billed: Final = { - leg.name: next(row for row in rows if row.model_id == deployments[leg.name].identity) - for leg in (case.plain, case.streamed) + leg.name: next(row for row in rows if row.model_id == deployments[leg.name].identity) for leg in legs } - for leg in (case.plain, case.streamed): + for leg in legs: + expected: Final = leg.expected + assert isinstance(expected, ExactExpected), leg.name assert_exact(leg.name, leg.response.content_type, expected, billed[leg.name], responses[leg.name]) - plain: Final = billed[case.plain.name] - streamed: Final = billed[case.streamed.name] - assert (streamed.spend, streamed.prompt_tokens, streamed.completion_tokens) == ( - plain.spend, - plain.prompt_tokens, - plain.completion_tokens, - ), (plain, streamed) + return MappingProxyType(billed) + + +def assert_stream_parity(case: StreamParityTestCase, gateway: Gateway) -> None: + billed: Final = bill_legs(case.id, (case.plain, case.streamed), gateway) + plain: Final = billed[case.plain.name] + streamed: Final = billed[case.streamed.name] + assert (streamed.spend, streamed.prompt_tokens, streamed.completion_tokens) == ( + plain.spend, + plain.prompt_tokens, + plain.completion_tokens, + ), (plain, streamed) diff --git a/tests/integration/cost_calculation/test_cost_tracking.py b/tests/integration/cost_calculation/test_cost_tracking.py index a84c31e810d..9db55589fb3 100644 --- a/tests/integration/cost_calculation/test_cost_tracking.py +++ b/tests/integration/cost_calculation/test_cost_tracking.py @@ -252,32 +252,19 @@ def test_case_bills_expected_cost(gateway: Gateway, case: CostTrackingTestCase) assert row.spend == 0, f"{case.name}: failure spend was {row.spend}" return assert response.is_success, f"{case.name}: proxy returned {response.status_code}: {response.text[:400]}" - if case.endpoint == "/v1/systemone": + if case.upstream_body is not None: observed: Final = JSON_OBJECT.validate_json( httpx.get(f"{gateway.upstream_url}/__observations", timeout=5, trust_env=False).content ) - decision_observations: Final = tuple( - value + upstream_bodies: Final = tuple( + value.get("body") for value in observed["requests"] - if isinstance(value, dict) and value.get("path") == f"/{scenario_id}/v1/decisions" + if isinstance(value, dict) + and value.get("method") == "POST" + and str(value.get("path", "")).startswith(f"/{scenario_id}/") ) - assert decision_observations == ( - { - "path": f"/{scenario_id}/v1/decisions", - "authorization": "Bearer sk-scripted-provider", - "method": "POST", - "api_key": "", - "body": { - "model": "pplx-decider-v1-27b", - "state": {"source": "cost-tracking"}, - "questions": { - "is_defect": { - "type": "noul", - "instructions": "Is this a defect?", - } - }, - }, - }, + assert upstream_bodies == (case.upstream_body,), ( + f"{case.name}: upstream received {upstream_bodies}, expected exactly {case.upstream_body}" ) if case.response.content_type == "text/event-stream": _assert_stream_has_no_error(response.text) diff --git a/tests/integration/database/migrations/test_autorouter_baseline_state.py b/tests/integration/database/migrations/test_autorouter_baseline_state.py index 629089eda02..729c34e1bb6 100644 --- a/tests/integration/database/migrations/test_autorouter_baseline_state.py +++ b/tests/integration/database/migrations/test_autorouter_baseline_state.py @@ -18,7 +18,9 @@ from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.proxy.db.baseline_accounting import BaselineAccountingStore from litellm.proxy.db.routing_prisma_wrapper import RoutingPrismaWrapper from litellm.proxy.utils import PrismaClient, ProxyLogging +from tests.integration._support.timing import assert_enforced_within +_PROJECTION_LOCK_TIMEOUT_SECONDS: Final = 1 _MIGRATION: Final = Path(__file__).parents[4] / ( "litellm-proxy-extras/litellm_proxy_extras/migrations/20260915010000_add_autorouter_baseline_state/migration.sql" ) @@ -95,7 +97,7 @@ async def test_locked_projection_is_bounded_and_cancellation_propagates( connection.execute('LOCK TABLE "LiteLLM_AutoRouterBaselineComparison" IN ACCESS EXCLUSIVE MODE') started: Final = time.monotonic() assert await store.project("scope") == "unavailable" - assert time.monotonic() - started < 2 + assert_enforced_within(time.monotonic() - started, configured=_PROJECTION_LOCK_TIMEOUT_SECONDS, slack=1) pending: Final = asyncio.create_task(store.project("scope")) await asyncio.sleep(0.01) pending.cancel() diff --git a/tests/integration/database/migrations/test_prisma_toolchain.py b/tests/integration/database/migrations/test_prisma_toolchain.py index fecb1fece91..bf111e7d6b9 100644 --- a/tests/integration/database/migrations/test_prisma_toolchain.py +++ b/tests/integration/database/migrations/test_prisma_toolchain.py @@ -17,7 +17,6 @@ the proxy gave up after four identical timeouts. import json import logging import os -import signal import sys import time from pathlib import Path @@ -33,8 +32,13 @@ from litellm_proxy_extras.prisma_toolchain import ( from litellm_proxy_extras.utils import ProxyExtrasDBManager from tests._process_helpers import process_is_gone +from tests.integration._support.process import kill_stray_process +from tests.integration._support.timing import assert_cut_short, assert_enforced_within +HUNG_BOOTSTRAP_SECONDS = 30 +HUNG_DEPLOY_SECONDS = 60 +COMMAND_TIMEOUT_SECONDS = 1 FAKE_PRISMA = """#!{python} import json import os @@ -159,8 +163,8 @@ def test_bootstrap_is_not_bounded_by_the_per_command_timeout( def test_bootstrap_stops_at_its_own_timeout( toolchain_env: tuple[Path, Path], monkeypatch: pytest.MonkeyPatch ) -> None: - monkeypatch.setenv(PRISMA_BOOTSTRAP_TIMEOUT_ENV_VAR, "1") - monkeypatch.setenv("FAKE_PRISMA_SLEEP", "30") + monkeypatch.setenv(PRISMA_BOOTSTRAP_TIMEOUT_ENV_VAR, str(COMMAND_TIMEOUT_SECONDS)) + monkeypatch.setenv("FAKE_PRISMA_SLEEP", str(HUNG_BOOTSTRAP_SECONDS)) started = time.monotonic() result = ensure_prisma_toolchain( @@ -169,7 +173,8 @@ def test_bootstrap_stops_at_its_own_timeout( elapsed = time.monotonic() - started assert result.ready is False - assert elapsed < 15 + assert_cut_short(elapsed, injected_delay=HUNG_BOOTSTRAP_SECONDS) + assert_enforced_within(elapsed, configured=COMMAND_TIMEOUT_SECONDS, slack=14) def test_setup_database_prepares_the_toolchain_before_migrating( @@ -210,8 +215,8 @@ def test_migrate_deploy_stops_at_its_own_timeout( """The deploy budget still bounds a deploy that hangs, so boot cannot wait forever.""" _, log_path = toolchain_env monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@localhost:9/x") - monkeypatch.setenv(PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR, "1") - monkeypatch.setenv("FAKE_PRISMA_FIRST_DEPLOY_SLEEP", "60") + monkeypatch.setenv(PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR, str(COMMAND_TIMEOUT_SECONDS)) + monkeypatch.setenv("FAKE_PRISMA_FIRST_DEPLOY_SLEEP", str(HUNG_DEPLOY_SECONDS)) monkeypatch.setenv("FAKE_PRISMA_LATER_DEPLOY_STDERR", "Error: P3018 permission denied for schema public") started = time.monotonic() @@ -220,7 +225,8 @@ def test_migrate_deploy_stops_at_its_own_timeout( elapsed = time.monotonic() - started assert len(_deploy_calls(log_path)) == 2 - assert elapsed < 30 + assert_cut_short(elapsed, injected_delay=HUNG_DEPLOY_SECONDS) + assert_enforced_within(elapsed, configured=2 * COMMAND_TIMEOUT_SECONDS, slack=28) def test_a_timed_out_migrate_deploy_takes_its_process_tree_with_it( @@ -230,7 +236,7 @@ def test_a_timed_out_migrate_deploy_takes_its_process_tree_with_it( _, log_path = toolchain_env pidfile = tmp_path / "grandchild.pid" monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@localhost:9/x") - monkeypatch.setenv(PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR, "1") + monkeypatch.setenv(PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR, str(COMMAND_TIMEOUT_SECONDS)) monkeypatch.setenv("FAKE_PRISMA_FIRST_DEPLOY_SLEEP", "60") monkeypatch.setenv("FAKE_PRISMA_LATER_DEPLOY_STDERR", "Error: P3018 permission denied for schema public") monkeypatch.setenv("FAKE_PRISMA_GRANDCHILD_PIDFILE", str(pidfile)) @@ -243,10 +249,7 @@ def test_a_timed_out_migrate_deploy_takes_its_process_tree_with_it( assert len(_deploy_calls(log_path)) == 2 assert process_is_gone(grandchild_pid, within_seconds=5) finally: - try: - os.kill(grandchild_pid, signal.SIGKILL) - except ProcessLookupError: - pass + kill_stray_process(grandchild_pid) def test_db_push_timeout_hint_names_the_per_command_budget( diff --git a/tests/integration/database/test_harness_lifecycle.py b/tests/integration/database/test_harness_lifecycle.py index bee5261ea0e..0d604e03f03 100644 --- a/tests/integration/database/test_harness_lifecycle.py +++ b/tests/integration/database/test_harness_lifecycle.py @@ -31,6 +31,7 @@ from typing_extensions import ReadOnly, TypedDict from tests.integration._support import process as process_support from tests.integration._support.async_server import LoopbackServer +from tests.integration._support.budget import case_timeout_seconds from tests.integration._support.client import Gateway, eventually from tests.integration._support.database_relay import ( DatabaseRelay, @@ -54,6 +55,17 @@ class ExitRecord(TypedDict): Relay = DatabaseRelay | DroppedConnectionRelay | HeldStatementRelay RelayFactory = Callable[[str, bytes], AbstractContextManager[tuple[Relay, str]]] STARTED_WORKER: Final = re.compile(r"Started server process \[(\d+)\]") +NESTED_RUN_SECONDS: Final = 60.0 +WORKER_STACK_BOOT_SECONDS: Final = 30.0 +OUTER_TEST_SLACK_SECONDS: Final = 30.0 + + +def _nested_run_seconds(workers: int) -> float: + return NESTED_RUN_SECONDS + (workers * WORKER_STACK_BOOT_SECONDS if workers > 1 else 0.0) + + +def _nested_run(*values: object, workers: int) -> object: + return pytest.param(*values, marks=pytest.mark.timeout(_nested_run_seconds(workers) + OUTER_TEST_SLACK_SECONDS)) def _nested_pytest_environment(environment: Mapping[str, str]) -> dict[str, str]: @@ -62,6 +74,7 @@ def _nested_pytest_environment(environment: Mapping[str, str]) -> dict[str, str] "INTEGRATION_RESULTS_DIR", "INTEGRATION_ROUTING", "INTEGRATION_WORKER_ISOLATION", + "INTEGRATION_COVERAGE_FILE", "INTEGRATION_CASE_SHARDS", "INTEGRATION_CASE_SHARD_INDEX", "CIRCLE_NODE_TOTAL", @@ -90,8 +103,7 @@ def test_owned_readiness_rejects_a_healthy_listener_from_another_process( process_support._wait_until_ready(launch) assert process.poll() is None finally: - process.kill() - process.wait(timeout=3) + process_support.stop_child(process, grace_seconds=3) @pytest.mark.parametrize("factory", (database_relay, dropped_connection_relay, held_statement_relay)) @@ -165,7 +177,7 @@ async def test_owned_runtime_starts_the_entire_held_provider_burst_on_a_small_ru release.set() -@pytest.mark.timeout(90) +@pytest.mark.timeout(case_timeout_seconds()) def test_slow_worker_startup_survives_healthchecks_and_a_stalled_worker_is_replaced( gateway: Gateway, tmp_path: Path ) -> None: @@ -193,23 +205,13 @@ def test_slow_worker_startup_survives_healthchecks_and_a_stalled_worker_is_repla "Proxy yielded before both workers finished startup" ) assert owned.log.read_text().count("Application startup complete.") == 2 - victim: Final = psutil.Process(workers[0]) - victim.suspend() - try: - replacement: Final = eventually( - lambda: tuple(int(pid) for pid in STARTED_WORKER.findall(owned.log.read_text())), - lambda pids: len(pids) == 3 and owned.log.read_text().count("Application startup complete.") == 3, - seconds=30, - ) - assert replacement[-1] not in workers - assert not psutil.pid_exists(victim.pid), "Supervisor did not reap the stalled worker" - response: Final = owned.gateway.request("GET", "/health/readiness") - assert response.status_code == 200, response.text - finally: - try: - victim.kill() - except psutil.NoSuchProcess: - pass + replacement: Final = owned.stall_worker(workers[0]) + assert replacement not in workers + assert tuple(int(pid) for pid in STARTED_WORKER.findall(owned.log.read_text()))[-1] == replacement + assert owned.log.read_text().count("Application startup complete.") == 3 + assert not psutil.pid_exists(workers[0]), "Supervisor did not reap the stalled worker" + response: Final = owned.gateway.request("GET", "/health/readiness") + assert response.status_code == 200, response.text def test_interrupted_cleanup_reaps_the_owned_root_and_its_child() -> None: @@ -249,8 +251,7 @@ def test_interrupted_cleanup_reaps_the_owned_root_and_its_child() -> None: assert process.poll() is not None eventually(lambda: process_support.group_members(process.pid), lambda members: not members, seconds=5) finally: - process_support.signal_group(process.pid, signal.SIGKILL) - process.wait(timeout=3) + process_support.stop_process_group(process, 3) @pytest.mark.parametrize("explicit_dotenv", (False, True)) @@ -451,7 +452,10 @@ def test_native_file_selection_reuses_fixtures_and_qualifies_an_empty_shard( ) -@pytest.mark.parametrize(("workers", "seed"), ((1, 1848224111), (2, 0), (2, 984872314))) +@pytest.mark.parametrize( + ("workers", "seed"), + tuple(_nested_run(workers, seed, workers=workers) for workers, seed in ((1, 1848224111), (2, 0), (2, 984872314))), +) @pytest.mark.parametrize("class_parameters", (False, True)) def test_mixed_fixture_scopes_reuse_outer_and_inner_configurations( tmp_path: Path, workers: int, seed: int, class_parameters: bool @@ -511,7 +515,7 @@ def test_mixed_fixture_scopes_reuse_outer_and_inner_configurations( env=_nested_pytest_environment(os.environ), capture_output=True, text=True, - timeout=60, + timeout=_nested_run_seconds(workers), check=False, ) assert result.returncode == 0, result.stdout + result.stderr @@ -536,16 +540,21 @@ def test_mixed_fixture_scopes_reuse_outer_and_inner_configurations( assert len(executed) == len(frozenset(executed)) == expected_passes -@pytest.mark.parametrize("workers", (1, 2)) +@pytest.mark.parametrize("workers", tuple(_nested_run(workers, workers=workers) for workers in (1, 2))) def test_direct_and_dependent_consumers_reuse_the_same_module_parameter(tmp_path: Path, workers: int) -> None: source: Final = Path(__file__).resolve().parents[3] root: Final = tmp_path / "probe-root" for relative in ( "pyproject.toml", "tests/__init__.py", + "tests/_fake_openai_endpoint_server.py", "tests/integration/__init__.py", "tests/integration/conftest.py", "tests/integration/run.py", + "tests/integration/proxy_config.yaml", + "tests/integration/mcp_coverage.toml", + "tests/integration/cost_calculation/cost_map.py", + "tests/integration/cost_calculation/cost_tracking_case.py", ): destination: Final = root / relative destination.parent.mkdir(parents=True, exist_ok=True) @@ -609,7 +618,7 @@ def test_direct_and_dependent_consumers_reuse_the_same_module_parameter(tmp_path ), capture_output=True, text=True, - timeout=60, + timeout=_nested_run_seconds(workers), check=False, ) assert result.returncode == 0, result.stdout + result.stderr diff --git a/tests/integration/database/test_v1_migration_error_logging.py b/tests/integration/database/test_v1_migration_error_logging.py index c3a3d8c43d6..43879d371b3 100644 --- a/tests/integration/database/test_v1_migration_error_logging.py +++ b/tests/integration/database/test_v1_migration_error_logging.py @@ -3,7 +3,6 @@ from __future__ import annotations import json import os import re -import signal import socket import socketserver import subprocess @@ -24,6 +23,7 @@ from psycopg.types.json import Jsonb from tests.integration._support.client import eventually from tests.integration._support.process import _free_port as free_port +from tests.integration._support.process import stop_process_group REPO_ROOT: Final = Path(__file__).resolve().parents[3] MIGRATIONS_DIR: Final = REPO_ROOT / "litellm-proxy-extras" / "litellm_proxy_extras" / "migrations" @@ -299,18 +299,6 @@ def _gated_postgres_forwarder(port: int, target: tuple[str, int]) -> Iterator[Ca thread.join(timeout=10) -def _stop_process(process: subprocess.Popen[str]) -> None: - if process.poll() is None: - with suppress(ProcessLookupError): - os.killpg(process.pid, signal.SIGTERM) - try: - process.wait(timeout=30) - except subprocess.TimeoutExpired: - with suppress(ProcessLookupError): - os.killpg(process.pid, signal.SIGKILL) - process.wait(timeout=30) - - def _migration_log_has_p1001_or_process_exited(output_path: Path, process: subprocess.Popen[str]) -> tuple[bool, bool]: p1001_logged: Final = any("P1001" in line for line in error_lines(output_path.read_text())) process_exited: Final = process.poll() is not None @@ -422,7 +410,7 @@ def test_unreachable_database_recovers_after_postgres_forwarder_starts( open_forwarder() completed_returncode: Final = process.wait(timeout=240) finally: - _stop_process(process) + stop_process_group(process, 30) output: Final = output_path.read_text() errors: Final = error_lines(output) p1001_errors: Final = tuple(line for line in errors if "P1001" in line) diff --git a/tests/integration/management/test_batch_file_usage_caps_chaos.py b/tests/integration/management/test_batch_file_usage_caps_chaos.py index 0ec3ab597e0..5f31fa94230 100644 --- a/tests/integration/management/test_batch_file_usage_caps_chaos.py +++ b/tests/integration/management/test_batch_file_usage_caps_chaos.py @@ -1,6 +1,5 @@ import asyncio import re -import signal import time from collections.abc import Callable, Coroutine, Iterator, Mapping from contextlib import contextmanager @@ -220,8 +219,7 @@ async def test_downloads_keep_answering_while_redis_hangs(chaos: Chaos) -> None: key: Final = _key(scenario, **{DOWNLOADS: DOWNLOAD_CAP}) file_id: Final = f"file-{marker()}" minute_ends: Final = window_end(MINUTE_SECONDS, MINUTE_ROOM_SECONDS) - chaos.redis.signal(signal.SIGSTOP) - try: + with chaos.redis.paused(): burst: Final = asyncio.create_task( _download_burst(str(chaos.candidate.client.base_url), key, file_id, BURST) ) @@ -229,8 +227,6 @@ async def test_downloads_keep_answering_while_redis_hangs(chaos: Chaos) -> None: timed, partial(chaos.candidate.request, "GET", "/health/liveliness") ) during: Final = await burst - finally: - chaos.redis.signal(signal.SIGCONT) assert time.time() < minute_ends assert liveliness.response.status_code == 200, liveliness.response.text assert liveliness.after - liveliness.before < 5, liveliness @@ -305,21 +301,24 @@ async def test_a_killed_worker_keeps_its_upload_slots_used_and_the_sibling_servi held_by: Final = {pid: _open_upstream_connections(pid, files.url) for pid in workers} assert sum(held_by.values()) == HELD_UPLOADS, held_by victim_pid, survivor_pid = sorted(workers, key=held_by.__getitem__) - victim: Final = psutil.Process(victim_pid) - victim.suspend() - victim.send_signal(signal.SIGKILL) - release.set() - served: Final = await burst - assert _statuses(served) == [200] * held_by[survivor_pid], (held_by, [response.text for response in served]) - assert_upload_limited( - timed(partial(upload, candidate, key, batch_file(marker(), 1))), - day_ends, - HELD_UPLOADS, - THIS_KEY, - IN_KEY, - ) - assert len(uploads_seen(files.drain(), mark)) == HELD_UPLOADS - assert candidate.request("GET", "/health/liveliness").status_code == 200 + with owned.replacing_worker(victim_pid) as kill: + kill.freeze() + await asyncio.to_thread(kill) + release.set() + served: Final = await burst + assert _statuses(served) == [200] * held_by[survivor_pid], ( + held_by, + [response.text for response in served], + ) + assert_upload_limited( + timed(partial(upload, candidate, key, batch_file(marker(), 1))), + day_ends, + HELD_UPLOADS, + THIS_KEY, + IN_KEY, + ) + assert len(uploads_seen(files.drain(), mark)) == HELD_UPLOADS + assert candidate.request("GET", "/health/liveliness").status_code == 200 def test_the_daily_count_survives_a_proxy_restart(chaos: Chaos) -> None: diff --git a/tests/integration/management/test_credential_federation_values.py b/tests/integration/management/test_credential_federation_values.py index e17fb1154d8..99be9f0cb8e 100644 --- a/tests/integration/management/test_credential_federation_values.py +++ b/tests/integration/management/test_credential_federation_values.py @@ -10,7 +10,7 @@ from functools import partial from itertools import product from pathlib import Path from types import MappingProxyType -from typing import Final, TypeVar +from typing import Final, Protocol, TypeVar, cast from urllib.parse import parse_qs import anthropic @@ -50,6 +50,11 @@ from tests.integration._support.wire import Reply, Request, Wire, wire_server T = TypeVar("T") + +class _ClientNetworkStream(Protocol): + def get_extra_info(self, info: str) -> object | None: ... + + SOURCES: Final = ("token_file", "secret_reference", "internal_issuer", "keycloak", "environment") CLIENTS: Final = ("chat", "messages", "messages_stream") @@ -771,6 +776,62 @@ def _workers(owned: OwnedProxy) -> tuple[psutil.Process, ...]: return tuple(child for child in psutil.Process(owned.process.pid).children() if _is_worker(child)) +def _worker_pids(owned: OwnedProxy) -> frozenset[int]: + return frozenset(child.pid for child in _workers(owned)) + + +def _client_address(response: httpx.Response) -> object | None: + stream: Final = cast(_ClientNetworkStream | None, response.extensions.get("network_stream")) + return stream.get_extra_info("client_addr") if stream is not None else None + + +def _credential_read_on_worker(owned: OwnedProxy, name: str, _: int) -> int | None: + try: + with httpx.Client( + base_url=owned.gateway.client.base_url, + timeout=15, + trust_env=False, + limits=httpx.Limits(max_connections=1), + ) as client: + before: Final = client.get("/health/integration-readiness") + credential: Final = client.get( + f"/credentials/by_name/{name}", + headers={"Authorization": f"Bearer {owned.gateway.key}"}, + ) + after: Final = client.get("/health/integration-readiness") + addresses: Final = tuple(_client_address(response) for response in (before, credential, after)) + except (httpx.TransportError, OSError): + return None + before_pid: Final = cast(str | None, before.headers.get("x-integration-worker-pid")) or "" + after_pid: Final = cast(str | None, after.headers.get("x-integration-worker-pid")) or "" + if ( + before.status_code != 200 + or credential.status_code != 200 + or after.status_code != 200 + or not before_pid.isdecimal() + or not after_pid.isdecimal() + or before_pid != after_pid + or ( + any(address is not None for address in addresses) + and not all(address == addresses[0] for address in addresses) + ) + ): + return None + return int(before_pid) + + +def _every_worker_holds(owned: OwnedProxy, name: str) -> bool: + with ThreadPoolExecutor(max_workers=8) as pool: + samples: Final = tuple(pool.map(partial(_credential_read_on_worker, owned, name), range(16))) + holders: Final = frozenset(pid for pid in samples if pid is not None) + workers: Final = _worker_pids(owned) + return len(workers) == 2 and workers <= holders + + +def _loaded_on_every_worker(owned: OwnedProxy, name: str) -> None: + assert eventually(lambda: _every_worker_holds(owned, name), lambda held: held, seconds=70) + + @pytest.mark.timeout(240) def test_worker_kill_mid_credential_burst(gateway: Gateway, tmp_path: Path) -> None: directory: Final = tmp_path.resolve() @@ -801,9 +862,9 @@ def test_worker_kill_mid_credential_burst(gateway: Gateway, tmp_path: Path) -> N return name, response.status_code victim: Final = eventually(lambda: _workers(owned), lambda workers: len(workers) == 2)[0] - with ThreadPoolExecutor(max_workers=8) as pool: + with owned.replacing_worker(victim.pid) as kill, ThreadPoolExecutor(max_workers=8) as pool: futures: Final = tuple(pool.submit(create, name) for name in names) - victim.kill() + kill() outcomes: Final = tuple(future.result() for future in futures) assert all(status == 200 or isinstance(status, str) for _, status in outcomes), outcomes landed: Final = tuple(name for name, status in outcomes if status == 200) @@ -1043,6 +1104,7 @@ def test_connection_probe_names_the_missing_id(federation: FederationRig) -> Non with federation.owned.gateway.scenario() as scenario: misconfigured: Final = _misconfigured(federation, scenario, _MISCONFIGURED["inline_token_without_org"]) _deployments_visible(federation.owned.gateway, (misconfigured.model,)) + _loaded_on_every_worker(federation.owned, misconfigured.credential) response: Final = federation.owned.gateway.request( "POST", "/health/test_connection", diff --git a/tests/integration/management/test_fallback_management_public_names.py b/tests/integration/management/test_fallback_management_public_names.py index b54465c1fd0..75dad3ae428 100644 --- a/tests/integration/management/test_fallback_management_public_names.py +++ b/tests/integration/management/test_fallback_management_public_names.py @@ -11,9 +11,22 @@ import psutil import pytest from pydantic import JsonValue -from tests.integration._support.client import Gateway, Scenario, eventually, object_value, string_value +from tests.integration._support.client import ( + Gateway, + Scenario, + eventually, + gateway_from_environment, + object_value, + string_value, +) from tests.integration._support.database import read_rows -from tests.integration._support.process import OwnedProxy, owned_proxy_process +from tests.integration._support.process import ( + MIGRATE_DEPLOY, + OwnedProxy, + owned_proxy_process, + setup_only_proxy_run, + worker_configuration, +) from tests.integration._support.redis_process import owned_redis pytestmark: Final = pytest.mark.timeout(300) @@ -21,6 +34,20 @@ pytestmark: Final = pytest.mark.timeout(300) PROVIDER_KEY: Final = "integration-provider-key" EVICTION_WARNING: Final = "config cache eviction of router_settings failed" REFUSED: Final = frozenset({401, 403}) +PROXY_CONFIG: Final = Path(__file__).resolve().parents[1] / "proxy_config.yaml" + + +@pytest.fixture(scope="module") +def migrated_database(tmp_path_factory: pytest.TempPathFactory) -> None: + with gateway_from_environment() as gateway: + completed: Final = setup_only_proxy_run( + gateway, + {}, + config=worker_configuration(PROXY_CONFIG, tmp_path_factory.mktemp("migrate")), + workers=1, + database_setup=MIGRATE_DEPLOY, + ) + assert completed.returncode == 0, completed.stdout[-3000:] + completed.stderr[-3000:] @pytest.fixture @@ -334,6 +361,7 @@ def test_context_window_and_content_policy_types_persist(gateway: Gateway) -> No assert policy in covered, covered +@pytest.mark.usefixtures("migrated_database") def test_fallback_write_survives_a_redis_outage(gateway: Gateway, tmp_path: Path) -> None: with owned_redis(tmp_path) as cache: overrides: Final = {"REDIS_HOST": cache.host, "REDIS_PORT": str(cache.port)} @@ -357,6 +385,7 @@ def test_fallback_write_survives_a_redis_outage(gateway: Gateway, tmp_path: Path cache.start() +@pytest.mark.usefixtures("migrated_database") def test_fallback_write_survives_a_worker_kill(gateway: Gateway, tmp_path: Path) -> None: with owned_proxy_process(gateway, tmp_path, {}, workers=2, database_setup=()) as owned: with owned.gateway.scenario() as scenario: @@ -368,12 +397,15 @@ def test_fallback_write_survives_a_worker_kill(gateway: Gateway, tmp_path: Path) _wait_until_served(owned.gateway, model) assert _create_fallback(owned.gateway, scenario, before, [target]).status_code == 200 victim: Final = eventually(lambda: _workers(owned), lambda workers: len(workers) == 2)[0] - victim.kill() fresh: Final = scenario.cleanups.enter_context( httpx.Client(base_url=owned.gateway.client.base_url, timeout=15, trust_env=False) ) survivor: Final = Gateway(fresh, owned.gateway.key, owned.gateway.upstream_url) - degraded: Final = tuple(_create_fallback(survivor, scenario, model, [target]) for model in (after, last)) + with owned.replacing_worker(victim.pid) as kill: + kill() + degraded: Final = tuple( + _create_fallback(survivor, scenario, model, [target]) for model in (after, last) + ) assert all(response.status_code == 200 for response in degraded), [r.text for r in degraded] covered: Final = _covered_models(_stored_fallbacks()) assert {before, after, last} <= covered, covered diff --git a/tests/integration/management/test_managed_file_create_ownership.py b/tests/integration/management/test_managed_file_create_ownership.py new file mode 100644 index 00000000000..a14f1fced19 --- /dev/null +++ b/tests/integration/management/test_managed_file_create_ownership.py @@ -0,0 +1,271 @@ +from __future__ import annotations + +import base64 +import json +import uuid +from dataclasses import dataclass +from typing import Final + +import httpx +import pytest +from integration._support.client import JSON_OBJECT, Gateway, Scenario, string_value +from integration._support.upstream import ScenarioHandle, delete_scenario, register_scenario +from integration.cost_calculation.cost_tracking_case import JsonResponse, RoutedResponse, TextResponse +from pydantic import JsonValue + +_FILE: Final = { + "id": "file-in-$REQUEST_ID", + "object": "file", + "purpose": "batch", + "bytes": 100, + "created_at": 1, + "filename": "input.jsonl", + "status": "processed", +} +_BATCH: Final = { + "id": "batch-$REQUEST_ID", + "object": "batch", + "endpoint": "/v1/chat/completions", + "errors": None, + "input_file_id": "file-in-$REQUEST_ID", + "completion_window": "24h", + "status": "validating", + "output_file_id": None, + "error_file_id": None, + "created_at": 1, + "expires_at": 1, + "request_counts": {"total": 1, "completed": 0, "failed": 0}, + "metadata": None, +} +_JOB: Final = { + "id": "ftjob-$REQUEST_ID", + "object": "fine_tuning.job", + "model": "gpt-4o-mini-2024-07-18", + "created_at": 1, + "fine_tuned_model": None, + "organization_id": "org-integration", + "result_files": [], + "status": "validating_files", + "validation_file": None, + "training_file": "file-in-$REQUEST_ID", + "hyperparameters": {"n_epochs": "auto"}, + "seed": 1, +} + + +@dataclass(frozen=True, slots=True) +class _Tenants: + model: str + handle: ScenarioHandle + owner_key: str + teammate_key: str + outsider_key: str + owner_file_id: str + + +def _routes() -> RoutedResponse: + return RoutedResponse( + content_type="application/x-routed", + routes={ + "POST /files": JsonResponse(content_type="application/json", body=_FILE), + "POST /batches": JsonResponse(content_type="application/json", body=_BATCH), + "POST /fine_tuning/jobs": JsonResponse(content_type="application/json", body=_JOB), + }, + ) + + +def _upload(scenario: Scenario, key: str, model: str, purpose: str) -> str: + line: Final = json.dumps( + { + "custom_id": "req-1", + "method": "POST", + "url": "/v1/chat/completions", + "body": {"model": model, "messages": [{"role": "user", "content": "ping"}]}, + } + ).encode() + uploaded: Final = scenario.gateway.request_multipart( + "/v1/files", + {"purpose": purpose, "target_model_names": model}, + {"file": ("input.jsonl", line + b"\n", "application/jsonl")}, + key=key, + ) + assert uploaded.status_code == 200, uploaded.text + return string_value(JSON_OBJECT.validate_json(uploaded.content)["id"]) + + +def _tenants(scenario: Scenario, purpose: str) -> _Tenants: + handle: Final = register_scenario(f"managed-file-create-ownership-{uuid.uuid4().hex}", _routes()) + scenario.cleanups.callback(delete_scenario, handle) + model: Final = scenario.model(api_base=handle.api_base()) + team_a: Final = scenario.team(models=[model]) + team_b: Final = scenario.team(models=[model]) + owner_key: Final = scenario.key(user_id=scenario.user(user_role="internal_user"), team_id=team_a, models=[model]) + teammate_key: Final = scenario.key(user_id=scenario.user(user_role="internal_user"), team_id=team_a, models=[model]) + outsider_key: Final = scenario.key(user_id=scenario.user(user_role="internal_user"), team_id=team_b, models=[model]) + return _Tenants(model, handle, owner_key, teammate_key, outsider_key, _upload(scenario, owner_key, model, purpose)) + + +def _upstream_creates(gateway: Gateway, handle: ScenarioHandle, path: str) -> int: + response: Final = httpx.get(f"{gateway.upstream_url}/__observations", timeout=15, trust_env=False) + assert response.status_code == 200, response.text + requests: Final = JSON_OBJECT.validate_json(response.content)["requests"] + assert isinstance(requests, list), response.text + return sum( + 1 + for request in requests + if isinstance(request, dict) + and request.get("method") == "POST" + and handle.scenario_id in str(request.get("path")) + and str(request.get("path")).endswith(path) + ) + + +def _create_batch(gateway: Gateway, key: str, model: str, input_file_id: JsonValue) -> httpx.Response: + return gateway.request( + "POST", + "/v1/batches", + { + "input_file_id": input_file_id, + "endpoint": "/v1/chat/completions", + "completion_window": "24h", + "model": model, + }, + key=key, + ) + + +def _create_job(gateway: Gateway, key: str, model: str, files: dict[str, JsonValue]) -> httpx.Response: + return gateway.request("POST", "/v1/fine_tuning/jobs", {"model": model, **files}, key=key) + + +def test_batch_create_from_another_teams_managed_file_is_403_and_never_reaches_the_provider(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + tenants: Final = _tenants(scenario, "batch") + denied: Final = _create_batch(gateway, tenants.outsider_key, tenants.model, tenants.owner_file_id) + assert denied.status_code == 403, denied.text + assert tenants.owner_file_id in denied.text, denied.text + assert _upstream_creates(gateway, tenants.handle, "/batches") == 0 + + +def test_batch_create_from_own_or_teammates_managed_file_still_dispatches(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + tenants: Final = _tenants(scenario, "batch") + owner: Final = _create_batch(gateway, tenants.owner_key, tenants.model, tenants.owner_file_id) + teammate: Final = _create_batch(gateway, tenants.teammate_key, tenants.model, tenants.owner_file_id) + admin: Final = _create_batch(gateway, gateway.key, tenants.model, tenants.owner_file_id) + assert (owner.status_code, teammate.status_code, admin.status_code) == (200, 200, 200), ( + owner.text, + teammate.text, + admin.text, + ) + assert JSON_OBJECT.validate_json(owner.content)["input_file_id"] == tenants.owner_file_id, owner.text + assert _upstream_creates(gateway, tenants.handle, "/batches") == 3 + + +def test_batch_create_from_a_managed_file_id_with_no_owner_row_is_404(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + tenants: Final = _tenants(scenario, "batch") + forged_raw: Final = ( + f"litellm_proxy:application/octet-stream;unified_id,{uuid.uuid4()};target_model_names,{tenants.model};" + "llm_output_file_id,file-forged;llm_output_file_model_id,model-forged" + ) + forged: Final = base64.urlsafe_b64encode(forged_raw.encode()).decode().rstrip("=") + missing: Final = _create_batch(gateway, tenants.outsider_key, tenants.model, forged) + assert missing.status_code == 404, missing.text + assert _upstream_creates(gateway, tenants.handle, "/batches") == 0 + + +@pytest.mark.parametrize("field", ["training_file", "validation_file"]) +def test_fine_tuning_create_from_another_teams_managed_file_is_403(gateway: Gateway, field: str) -> None: + with gateway.scenario() as scenario: + tenants: Final = _tenants(scenario, "fine-tune") + own_training: Final = _upload(scenario, tenants.outsider_key, tenants.model, "fine-tune") + files: Final[dict[str, JsonValue]] = {"training_file": own_training, field: tenants.owner_file_id} + denied: Final = _create_job(gateway, tenants.outsider_key, tenants.model, files) + assert denied.status_code == 403, denied.text + assert _upstream_creates(gateway, tenants.handle, "/fine_tuning/jobs") == 0 + + +def test_fine_tuning_create_from_own_managed_files_still_dispatches(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + tenants: Final = _tenants(scenario, "fine-tune") + validation: Final = _upload(scenario, tenants.teammate_key, tenants.model, "fine-tune") + created: Final = _create_job( + gateway, + tenants.owner_key, + tenants.model, + {"training_file": tenants.owner_file_id, "validation_file": validation}, + ) + assert created.status_code == 200, created.text + assert _upstream_creates(gateway, tenants.handle, "/fine_tuning/jobs") == 1 + + +def _upstream_calls(gateway: Gateway, handle: ScenarioHandle, method: str, suffix: str) -> int: + response: Final = httpx.get(f"{gateway.upstream_url}/__observations", timeout=15, trust_env=False) + assert response.status_code == 200, response.text + requests: Final = JSON_OBJECT.validate_json(response.content)["requests"] + assert isinstance(requests, list), response.text + return sum( + 1 + for request in requests + if isinstance(request, dict) + and request.get("method") == method + and handle.scenario_id in str(request.get("path")) + and str(request.get("path")).endswith(suffix) + ) + + +def test_batch_create_from_another_teams_file_is_403_before_the_rate_limiter_reads_it(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + scenario_id: Final = f"managed-file-create-ownership-{uuid.uuid4().hex}" + handle: Final = register_scenario(scenario_id, _routes()) + scenario.cleanups.callback(delete_scenario, handle) + model: Final = scenario.model(api_base=handle.api_base()) + line: Final = json.dumps( + { + "custom_id": "req-1", + "method": "POST", + "url": "/v1/chat/completions", + "body": {"model": model, "messages": [{"role": "user", "content": "ping"}], "max_tokens": 5}, + } + ) + register_scenario( + scenario_id, + RoutedResponse( + content_type="application/x-routed", + routes={ + **_routes().routes, + "GET /files/file-in-$REQUEST_ID/content": TextResponse( + content_type="application/jsonl", body=line + "\n" + ), + }, + ), + ) + team_a: Final = scenario.team(models=[model]) + team_b: Final = scenario.team(models=[model]) + owner: Final = scenario.key(user_id=scenario.user(user_role="internal_user"), team_id=team_a) + outsider: Final = scenario.key(user_id=scenario.user(user_role="internal_user"), team_id=team_b, tpm_limit=1) + owner_file_id: Final = _upload(scenario, owner, model, "batch") + denied: Final = _create_batch(gateway, outsider, model, owner_file_id) + assert denied.status_code == 403, denied.text + assert owner_file_id in denied.text, denied.text + assert _upstream_calls(gateway, handle, "GET", "/content") == 0 + + +def test_regenerated_key_without_user_or_team_keeps_its_managed_files(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + tenants: Final = _tenants(scenario, "batch") + original: Final = string_value(gateway.post("/key/generate", {})["key"]) + file_id: Final = _upload(scenario, original, tenants.model, "batch") + regenerated: Final = string_value(gateway.post("/key/regenerate", {"key": original})["key"]) + scenario.cleanups.callback(scenario.delete_key, regenerated) + batch: Final = _create_batch(gateway, regenerated, tenants.model, file_id) + retrieved: Final = gateway.request("GET", f"/v1/files/{file_id}", key=regenerated) + stranger: Final = _create_batch( + gateway, string_value(gateway.post("/key/generate", {})["key"]), tenants.model, file_id + ) + assert (batch.status_code, retrieved.status_code, stranger.status_code) == (200, 200, 403), ( + batch.text, + retrieved.text, + stranger.text, + ) diff --git a/tests/integration/management/test_project_lifecycle.py b/tests/integration/management/test_project_lifecycle.py index d2f45d853a4..d503daf6a0f 100644 --- a/tests/integration/management/test_project_lifecycle.py +++ b/tests/integration/management/test_project_lifecycle.py @@ -1,7 +1,6 @@ import json import os import re -import signal import threading from collections.abc import Iterator from concurrent.futures import ThreadPoolExecutor @@ -14,6 +13,7 @@ import httpx import psutil import psycopg import pytest +import yaml from integration._support.client import ( JSON_OBJECT, Gateway, @@ -24,7 +24,7 @@ from integration._support.client import ( string_value, ) from integration._support.database import read_rows, write_rows -from integration._support.process import group_members, owned_proxy, owned_proxy_process +from integration._support.process import CONFIGURATION, OwnedProxy, group_members, owned_proxy, owned_proxy_process from integration._support.wire import Reply, Request, wire_server from pydantic import JsonValue @@ -32,9 +32,26 @@ from litellm.models.user import LiteLLM_UserTable from litellm.proxy.auth.auth_checks import ExperimentalUIJWTToken _OWNED_PROXY_SALT_KEY: Final = "sk-integration-salt" +_OWNERSHIP_POOL_LIMIT: Final = 4 +_CONNECTION_BUDGET: Final = re.compile( + r"Postgres connection budget per pod = \d+ worker\(s\) x \(?writer connection_limit (\d+)" +) _OWNERSHIP_MARKER: Final = re.compile(rb"ownership-[0-9a-f]{32}") +def _ownership_config(directory: Path) -> Path: + source: Final = Path(__file__).resolve().parents[1] / "proxy_config.yaml" + settings: Final = CONFIGURATION.validate_python(yaml.safe_load(source.read_text())) + general: Final = CONFIGURATION.validate_python(settings["general_settings"]) + configured: Final = { + **settings, + "general_settings": {**general, "database_connection_pool_limit": _OWNERSHIP_POOL_LIMIT}, + } + destination: Final = directory / "ownership-proxy-config.yaml" + destination.write_text(yaml.safe_dump(configured, sort_keys=False)) + return destination + + def _project_rows(project_id: str) -> list[dict[str, JsonValue]]: return read_rows( 'SELECT p.project_id, p.project_alias, p.description, p.team_id, p.models, p.blocked, ' @@ -806,7 +823,12 @@ def test_ownership_rejections_during_concurrent_traffic_burst(ownership_gateway: assert _project_state_rows(project_a) == project_before -def test_ownership_checks_wait_for_row_locks_without_failing_readiness(ownership_gateway: Gateway) -> None: +def test_ownership_checks_wait_for_row_locks_without_failing_readiness( + ownership_gateway: Gateway, ownership_proxy: OwnedProxy +) -> None: + budget: Final = _CONNECTION_BUDGET.search(ownership_proxy.log.read_text()) + assert budget is not None, ownership_proxy.log.read_text() + assert int(budget.group(1)) == _OWNERSHIP_POOL_LIMIT with ownership_gateway.scenario() as scenario: model: Final = scenario.model() team: Final = scenario.team(models=[model]) @@ -935,14 +957,12 @@ def test_ownership_enforced_after_worker_kill(ownership_gateway: Gateway, tmp_pa assert len(workers) == 2, tuple(process.pid for process in workers) victim: Final = workers[0] survivor: Final = workers[1] - with ThreadPoolExecutor(max_workers=24) as pool: + with owned.replacing_worker(victim.pid) as kill, ThreadPoolExecutor(max_workers=24) as pool: futures: Final = tuple(pool.submit(serve, index) for index in range(24)) try: assert started.wait(timeout=10) - os.kill(victim.pid, signal.SIGKILL) + kill() release_stream.set() - psutil.wait_procs((victim,), timeout=10) - assert not psutil.pid_exists(victim.pid), victim.pid finally: release_stream.set() burst: Final = tuple(future.result(timeout=90) for future in futures) @@ -986,15 +1006,22 @@ def test_ownership_enforced_after_worker_kill(ownership_gateway: Gateway, tmp_pa @pytest.fixture(scope="module") -def ownership_gateway(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Gateway]: +def ownership_proxy(tmp_path_factory: pytest.TempPathFactory) -> Iterator[OwnedProxy]: + directory: Final = tmp_path_factory.mktemp("project-team-ownership") with gateway_from_environment() as gateway: - with owned_proxy( + with owned_proxy_process( gateway, - tmp_path_factory.mktemp("project-team-ownership"), + directory, {"LITELLM_SALT_KEY": _OWNED_PROXY_SALT_KEY}, + config=_ownership_config(directory), workers=2, - ) as candidate: - yield candidate + ) as owned: + yield owned + + +@pytest.fixture(scope="module") +def ownership_gateway(ownership_proxy: OwnedProxy) -> Gateway: + return ownership_proxy.gateway @pytest.mark.covers("mgmt.project.new.real_route_persists") diff --git a/tests/integration/management/test_scim_group_pathless_patch.py b/tests/integration/management/test_scim_group_pathless_patch.py index 5ed784d2bb8..efd3afba5b3 100644 --- a/tests/integration/management/test_scim_group_pathless_patch.py +++ b/tests/integration/management/test_scim_group_pathless_patch.py @@ -1,5 +1,3 @@ -import os -import signal import threading import uuid from collections.abc import Callable, Mapping, Sequence @@ -11,7 +9,6 @@ from types import MappingProxyType from typing import Final import httpx -import psutil import pytest from integration._support.client import Gateway, Scenario, eventually, object_value, string_value from integration._support.process import owned_proxy_process @@ -522,21 +519,17 @@ def _patch_status(candidate: Gateway, team: str) -> int | None: return None -def _serving_workers(root: int) -> tuple[psutil.Process, ...]: - return tuple(child for child in psutil.Process(root).children() if child.children()) - - @pytest.mark.timeout(240) def test_worker_kill_mid_burst_keeps_serving_and_leaves_no_empty_key(gateway: Gateway, tmp_path: Path) -> None: with owned_proxy_process(gateway, tmp_path, {}, workers=2) as owned, gateway.scenario() as scenario: teams: Final = _burst_teams(scenario) - workers: Final = eventually( - lambda: _serving_workers(owned.process.pid), lambda children: len(children) >= 2, seconds=30 - ) - attempts: Final = _burst(owned.gateway, tuple(teams), lambda: os.kill(workers[0].pid, signal.SIGKILL)) - assert not workers[0].is_running() or workers[0].status() == psutil.STATUS_ZOMBIE, workers[0] - probe: Final = scenario.team() - eventually(lambda: _patch_status(owned.gateway, probe), lambda status_code: status_code == 200, seconds=60) + workers: Final = eventually(owned.worker_pids, lambda pids: len(pids) >= 2, seconds=30) + with owned.replacing_worker(min(workers)) as kill: + attempts: Final = _burst(owned.gateway, tuple(teams), kill) + probe: Final = scenario.team() + eventually( + lambda: _patch_status(owned.gateway, probe), lambda status_code: status_code == 200, seconds=60 + ) _assert_burst_reflected(owned.gateway, teams, attempts) @@ -545,5 +538,5 @@ def test_rolling_restart_mid_burst_drains_without_empty_keys(gateway: Gateway, t with gateway.scenario() as scenario, owned_proxy_process(gateway, tmp_path, {}, workers=2) as replacement: teams: Final = _burst_teams(scenario) with owned_proxy_process(gateway, tmp_path, {}, workers=2) as retiring: - attempts: Final = _burst(retiring.gateway, tuple(teams), retiring.process.terminate) + attempts: Final = _burst(retiring.gateway, tuple(teams), retiring.begin_shutdown) _assert_burst_reflected(replacement.gateway, teams, attempts) diff --git a/tests/integration/management/test_team_delete_chaos.py b/tests/integration/management/test_team_delete_chaos.py index ebe515f59ec..9df4d1bb973 100644 --- a/tests/integration/management/test_team_delete_chaos.py +++ b/tests/integration/management/test_team_delete_chaos.py @@ -382,21 +382,21 @@ def test_worker_killed_mid_burst_leaves_every_team_fully_deleted_after_retry( fleet: Final = _fleet(owned.gateway, scenario) pinned: Final = fleet[0] burst: Final = Burst(owned.gateway) - with _holding_team_lock(pinned.team_id) as holder_pid: - burst.start(pool, fleet) - _await_pinned_delete_blocked(burst, pinned, holder_pid, record_property) - before: Final = _workers(root) - assert len(before) == WORKERS, [process.pid for process in before] - victim: Final = before[0] - victim.kill() # SIGKILL with the pinned delete blocked: the worker cannot finish its in-flight deletes - victim.wait(timeout=10) - with httpx.Client(base_url=str(owned.gateway.client.base_url), timeout=15, trust_env=False) as fresh: - readiness: Final = fresh.get("/health/readiness") - assert readiness.status_code == 200, ( - f"/health/readiness with worker {victim.pid} dead: {readiness.status_code} {readiness.text}" - ) - # The lock is released: the pinned delete finishes on the survivor, or was dropped with the victim. - outcomes: Final = burst.outcomes() + before: Final = _workers(root) + assert len(before) == WORKERS, [process.pid for process in before] + victim: Final = before[0] + with owned.replacing_worker(victim.pid) as kill: + with _holding_team_lock(pinned.team_id) as holder_pid: + burst.start(pool, fleet) + _await_pinned_delete_blocked(burst, pinned, holder_pid, record_property) + kill() # SIGKILL with the pinned delete blocked: the worker cannot finish its in-flight deletes + with httpx.Client(base_url=str(owned.gateway.client.base_url), timeout=15, trust_env=False) as fresh: + readiness: Final = fresh.get("/health/readiness") + assert readiness.status_code == 200, ( + f"/health/readiness with worker {victim.pid} dead: {readiness.status_code} {readiness.text}" + ) + # The lock is released: the pinned delete finishes on the survivor, or was dropped with the victim. + outcomes: Final = burst.outcomes() respawned: Final = eventually( lambda: tuple(process.pid for process in _workers(root)), lambda pids: len(pids) == WORKERS and victim.pid not in pids, @@ -480,7 +480,7 @@ def test_proxy_terminated_mid_burst_then_restarted_leaves_every_team_fully_delet with _holding_team_lock(pinned.team_id) as holder_pid: burst.start(pool, fleet) _await_pinned_delete_blocked(burst, pinned, holder_pid, record_property) - doomed.process.terminate() # SIGTERM with the pinned delete blocked: uvicorn stops accepting and drains + doomed.begin_shutdown() # SIGTERM with the pinned delete blocked: uvicorn stops accepting and drains assert burst.pending(pinned.team_id), f"{pinned.team_id}: delete answered while its team lock was held" # The lock is released: the drain lets the pinned delete finish before the proxy exits. doomed.process.wait(timeout=120) diff --git a/tests/integration/management/test_team_metadata_validation.py b/tests/integration/management/test_team_metadata_validation.py index b1e30a0d86f..f02f1e66a15 100644 --- a/tests/integration/management/test_team_metadata_validation.py +++ b/tests/integration/management/test_team_metadata_validation.py @@ -10,7 +10,7 @@ import pytest from pydantic import JsonValue from tests.integration._support.client import Gateway, eventually, gateway_from_environment, object_value, string_value -from tests.integration._support.process import owned_proxy +from tests.integration._support.process import owned_proxy, stop_child _SERVICE_URL: Final = "http://127.0.0.1:9414" _CONFIG: Final = Path(__file__).with_name("team_metadata_validation_proxy_config.yaml") @@ -87,9 +87,7 @@ def cost_center_service() -> Iterator[None]: assert health is not None and health.status_code == 200 yield finally: - if process.poll() is None: - process.terminate() - process.wait(timeout=5) + stop_child(process, grace_seconds=5) @pytest.fixture(scope="module") diff --git a/tests/integration/management/test_user_updates_wedged_coordination_redis.py b/tests/integration/management/test_user_updates_wedged_coordination_redis.py index 0715059744e..79b1f3166e3 100644 --- a/tests/integration/management/test_user_updates_wedged_coordination_redis.py +++ b/tests/integration/management/test_user_updates_wedged_coordination_redis.py @@ -1,5 +1,4 @@ import os -import signal import time import uuid from collections.abc import Callable, Mapping @@ -8,7 +7,6 @@ from pathlib import Path from typing import Final from urllib.parse import urlsplit, urlunsplit -import psutil import psycopg import pytest from psycopg import sql @@ -16,12 +14,14 @@ from pydantic import JsonValue from redis import Redis from redis.client import PubSub +from tests.integration._support.budget import case_timeout_seconds from tests.integration._support.client import JSON_OBJECT, Gateway, eventually, object_value, string_value -from tests.integration._support.process import owned_proxy +from tests.integration._support.process import owned_proxy_process from tests.integration._support.redis_process import owned_redis _USERS: Final = 60 _BURST: Final = 30 +_WEDGED_PUBLISHES_PER_BURST_USER: Final = 2 _HANDLER_BUDGET_SECONDS: Final = 0.75 _BULK_BUDGET_SECONDS: Final = 2.0 _CHANNEL: Final = "litellm_proxy.auth_cache_invalidation" @@ -52,26 +52,6 @@ def _received(pubsub: PubSub) -> tuple[dict[str, JsonValue], ...]: messages.append(JSON_OBJECT.validate_json(data)) -def _worker_pid(port: int) -> int: - for process in psutil.process_iter(): - parent = process.parent() - if parent is None: - continue - try: - cmdline = parent.cmdline() - own_cmdline = process.cmdline() - except (psutil.NoSuchProcess, psutil.AccessDenied): - continue - if ( - "integration._support.proxy" in cmdline - and "--port" in cmdline - and str(port) in cmdline - and not any("prisma" in part for part in own_cmdline) - ): - return process.pid - raise AssertionError(f"no uvicorn worker found under the owned proxy on port {port}") - - def _burst_call( index: int, users: tuple[str, ...], key: str, team_id: str, customer_id: str ) -> tuple[str, dict[str, JsonValue]]: @@ -88,7 +68,7 @@ def _burst_call( return "/customer/update", {"user_id": customer_id, "max_budget": 7.0 + index} -@pytest.mark.timeout(240) +@pytest.mark.timeout(case_timeout_seconds()) @pytest.mark.covers( "mgmt.user.update.budget_change_returns_promptly_with_wedged_coordination_redis", "mgmt.user.bulk_update.budget_change_returns_promptly_with_wedged_coordination_redis", @@ -112,7 +92,7 @@ def test_user_budget_updates_return_promptly_while_coordination_redis_is_wedged( prior_logs: Final = frozenset(results_dir.glob("owned-proxy-*.log")) with ( owned_redis(tmp_path) as coordination, - owned_proxy( + owned_proxy_process( gateway, tmp_path, { @@ -123,9 +103,10 @@ def test_user_budget_updates_return_promptly_while_coordination_redis_is_wedged( config=Path("tests/integration/coordination_redis_proxy_config.yaml"), remove_environment=("DATABASE_URL_READ_REPLICA",), workers=2, - ) as candidate, + ) as owned, Redis(host=coordination.host, port=coordination.port, socket_timeout=1) as subscriber_client, ): + candidate: Final = owned.gateway pubsub: Final = subscriber_client.pubsub() pubsub.subscribe(_CHANNEL) received: list[dict[str, JsonValue]] = [] @@ -168,8 +149,8 @@ def test_user_budget_updates_return_promptly_while_coordination_redis_is_wedged( assert timings["h2_healthy_control"] < _HANDLER_BUDGET_SECONDS, ( f"healthy control update took {timings['h2_healthy_control']:.3f}s" ) - coordination.signal(signal.SIGSTOP) - try: + with owned.replacing_worker(min(owned.worker_pids())) as kill, coordination.paused(): + kill() timings["s2_wedged_control"] = _timed_post( candidate, "/user/update", {"user_id": users[0], "tpm_limit": 1000} ) @@ -225,22 +206,17 @@ def test_user_budget_updates_return_promptly_while_coordination_redis_is_wedged( f"/user/update with an invalid body took {timings['s8_invalid_body']:.3f}s" ) - def burst_request(path: str, body: Mapping[str, JsonValue]) -> tuple[object, float]: + def burst_request(path: str, body: Mapping[str, JsonValue]) -> tuple[int, float]: started: Final = time.monotonic() - try: - response: Final = candidate.client.request( - "POST", - path, - json=body, - headers={"Authorization": f"Bearer {candidate.key}"}, - timeout=60, - ) - return response.status_code, time.monotonic() - started - except Exception as error: # noqa: BLE001 # the killed worker drops in-flight requests - return error, time.monotonic() - started + response: Final = candidate.client.request( + "POST", + path, + json=body, + headers={"Authorization": f"Bearer {candidate.key}"}, + timeout=60, + ) + return response.status_code, time.monotonic() - started - port: Final = candidate.client.base_url.port - assert port is not None, f"owned proxy client has no port: {candidate.client.base_url}" with ThreadPoolExecutor(_BURST) as pool: futures: Final = [ pool.submit( @@ -249,15 +225,9 @@ def test_user_budget_updates_return_promptly_while_coordination_redis_is_wedged( ) for i in range(_BURST) ] - os.kill(_worker_pid(port), signal.SIGKILL) results: Final = [future.result() for future in futures] - responses: Final = [(status, elapsed) for status, elapsed in results if isinstance(status, int)] - failures: Final = [status for status, _elapsed in responses if status != 200] + failures: Final = [status for status, _elapsed in results if status != 200] assert not failures, f"burst responses that were not 200: {failures}" - transport_errors: Final = [status for status, _elapsed in results if not isinstance(status, int)] - assert len(transport_errors) <= 3, ( - f"{len(transport_errors)} requests raised transport errors: {transport_errors!r}" - ) elapsed_sorted: Final = sorted( elapsed for i, (status, elapsed) in enumerate(results) if i % 5 in (0, 1, 4) and status == 200 ) @@ -276,39 +246,38 @@ def test_user_budget_updates_return_promptly_while_coordination_redis_is_wedged( assert timings["c1_survivor"] < _HANDLER_BUDGET_SECONDS, ( f"control update on the surviving worker took {timings['c1_survivor']:.3f}s" ) - finally: - coordination.signal(signal.SIGCONT) - wedged_keys: Final = {users[i] for i in range(_BURST) if i % 5 == 0 and i != 0} | {f"team_id:{team_id}"} + wedged_keys: Final = {users[i] for i in range(_BURST) if i % 5 == 0 and i != 0} | { + f"team_id:{team_id}" + } - def proxy_log() -> str: - return "".join( - path.read_text() for path in results_dir.glob("owned-proxy-*.log") if path not in prior_logs + def proxy_log() -> str: + return "".join( + path.read_text() for path in results_dir.glob("owned-proxy-*.log") if path not in prior_logs + ) + + team_wedged_key: Final = f"team_id:{team_id}" + eventually( + proxy_log, + lambda text: ( + all( + text.count(f"publish for {wedged_key} failed") >= _WEDGED_PUBLISHES_PER_BURST_USER + for wedged_key in wedged_keys + if wedged_key != team_wedged_key + ) + and ( + f"publish for {team_wedged_key} failed" in text + or f"internal usage cache entry {team_wedged_key}" in text + ) + ), + seconds=45, ) - - team_wedged_key: Final = f"team_id:{team_id}" - eventually( - proxy_log, - lambda text: ( - all( - f"publish for {wedged_key} failed" in text - for wedged_key in wedged_keys - if wedged_key != team_wedged_key - ) - and ( - f"publish for {team_wedged_key} failed" in text - or f"internal usage cache entry {team_wedged_key}" in text - ) - ), - seconds=45, - ) marker: Final = len(received) drained() recovered_keys: Final = {str(message.get("cache_key")) for message in received[marker:]} assert recovered_keys.isdisjoint(wedged_keys), ( f"wedged publishes unexpectedly landed after recovery: {sorted(recovered_keys & wedged_keys)}" ) - coordination.signal(signal.SIGSTOP) - try: + with coordination.paused(): timings["r1b_short_wedge_a"] = _timed_post( candidate, "/user/update", {"user_id": users[4], "max_budget": 15.0} ) @@ -321,8 +290,6 @@ def test_user_budget_updates_return_promptly_while_coordination_redis_is_wedged( assert timings["r1b_short_wedge_b"] < _HANDLER_BUDGET_SECONDS, ( f"/user/update inside a short wedge took {timings['r1b_short_wedge_b']:.3f}s" ) - finally: - coordination.signal(signal.SIGCONT) eventually( drained, lambda messages: {str(message.get("cache_key")) for message in messages} >= {users[4], users[5]}, diff --git a/tests/integration/management/test_vector_store_file_list_managed_ids.py b/tests/integration/management/test_vector_store_file_list_managed_ids.py index c6124519c9e..13a485f59d4 100644 --- a/tests/integration/management/test_vector_store_file_list_managed_ids.py +++ b/tests/integration/management/test_vector_store_file_list_managed_ids.py @@ -2,7 +2,6 @@ import base64 import hashlib import json import re -import signal import threading import uuid from collections.abc import Callable, Generator, Mapping @@ -743,11 +742,11 @@ def test_worker_sigkill_mid_burst_leaves_the_sibling_mapping_ids(gateway: Gatewa held_by: Final = MappingProxyType({pid: _open_connections_to(pid, wire.url) for pid in workers}) assert sum(held_by.values()) == 20, held_by victim_pid, survivor_pid = sorted(workers, key=held_by.__getitem__) - victim: Final = psutil.Process(victim_pid) - victim.suspend() - victim.send_signal(signal.SIGKILL) - release.set() - served: Final = tuple(result for future in burst if (result := future.result()) is not None) + with owned.replacing_worker(victim_pid) as kill: + kill.freeze() + kill() + release.set() + served: Final = tuple(result for future in burst if (result := future.result()) is not None) assert held_by[survivor_pid] >= 1, held_by assert len(served) == held_by[survivor_pid], (held_by, len(served)) for response in served: diff --git a/tests/integration/mcp/test_mcp_credentials.py b/tests/integration/mcp/test_mcp_credentials.py index ee36581acdd..63f1b0024c0 100644 --- a/tests/integration/mcp/test_mcp_credentials.py +++ b/tests/integration/mcp/test_mcp_credentials.py @@ -631,16 +631,23 @@ def test_oauth_passthrough_probe_uses_server_token_without_admission_key( with mcp_peer() as peer, gateway.scenario() as scenario: alias: Final = "probe" + uuid.uuid4().hex[:8] identity: Final = register_mcp( - scenario, peer, alias, auth_type="none", oauth_passthrough=True, extra_headers=["Authorization"] + scenario, + peer, + alias, + auth_type="none", + oauth_passthrough=True, + extra_headers=["Authorization", "x-integration-caller"], ) key: Final = scenario.key(object_permission={"mcp_servers": [identity]}) upstream_token: Final = "upstream-" + uuid.uuid4().hex + caller_trace: Final = "caller-" + uuid.uuid4().hex caller: Final = McpCaller( gateway, key, entry, alias, headers={ + "x-integration-caller": caller_trace, "Authorization": f"Bearer {upstream_token if forwarded_token else key}", f"x-mcp-{alias}-authorization": "Bearer expired-token" if forwarded_token @@ -651,18 +658,28 @@ def test_oauth_passthrough_probe_uses_server_token_without_admission_key( outcome: Final = caller.call(f"{alias}-add", ADD, identity if entry != "server_mcp" else None) assert outcome.ok, outcome.raw observed: Final = peer.drain() + forwarded: Final = f"Bearer {upstream_token}".encode() probes: Final = tuple( request for request in observed if TypeAdapter(dict[str, object]).validate_python(request["body"]).get("id") == "litellm-mcp-auth-probe" ) assert probes, "expected an upstream initialize authentication probe" - assert all(_header(probe, b"authorization") == f"Bearer {upstream_token}".encode() for probe in probes) - assert len(tool_calls(observed)) == 1 - foreign: Final = tuple( - request for request in observed if _header(request, b"authorization") != f"Bearer {upstream_token}".encode() + assert all(_header(probe, b"authorization") == forwarded for probe in probes) + callers_own: Final = tuple( + request for request in observed if _header(request, b"x-integration-caller") == caller_trace.encode() ) + assert len(tool_calls(observed)) == 1 + assert len(tool_calls(callers_own)) == 1, "the tool call lost the caller's forwarded headers" + attributed: Final = probes + callers_own + foreign: Final = tuple(request for request in attributed if _header(request, b"authorization") != forwarded) assert foreign == (), foreign + credentialed_background: Final = tuple( + request + for request in observed + if request not in attributed and _header(request, b"authorization") is not None + ) + assert credentialed_background == (), credentialed_background header_sets: Final = tuple( TypeAdapter(dict[bytes, bytes]).validate_python(request["headers"]) for request in observed ) diff --git a/tests/integration/mcp/test_mcp_oauth_flows.py b/tests/integration/mcp/test_mcp_oauth_flows.py index e95cdb902e8..0bd299f3bbd 100644 --- a/tests/integration/mcp/test_mcp_oauth_flows.py +++ b/tests/integration/mcp/test_mcp_oauth_flows.py @@ -1,11 +1,13 @@ import base64 import hashlib +import json import re import secrets import textwrap import time import uuid from collections.abc import Iterator +from contextlib import contextmanager from dataclasses import dataclass from pathlib import Path from typing import Final @@ -23,15 +25,19 @@ from integration._support.mcp import ( McpCaller, McpPeer, Outcome, + ScriptedTool, _outcome_from_rpc, call_tool, mcp_peer, register_mcp, + scripted_peer, + text_result, tool_calls, ) from integration._support.mcp_grants import create_toolset from integration._support.oauth_server import AuthorizationServer, oauth_server from integration._support.process import owned_proxy +from integration._support.wire import Reply, Request, Wire, wire_server ADD: Final = {"a": 2, "b": 3} CLIENT_REDIRECT: Final = "http://127.0.0.1:9/cb" @@ -240,7 +246,8 @@ class _Pkce: def _authorize_through_gateway( - gateway: Gateway, auth: AuthorizationServer, alias: str, key: str, client_id: str, pkce: _Pkce + gateway: Gateway, auth: AuthorizationServer, alias: str, key: str, client_id: str, pkce: _Pkce, + scope: str = "tools.call", ) -> str: started: Final = gateway.client.get( f"/{alias}/authorize", @@ -251,7 +258,7 @@ def _authorize_through_gateway( "state": "client-state", "code_challenge": pkce.challenge, "code_challenge_method": "S256", - "scope": "tools.call", + "scope": scope, }, headers={"x-litellm-api-key": key}, ) @@ -575,3 +582,140 @@ def test_token_exchange_callers_with_different_subject_tokens_own_separate_listi assert second.list_tools().ok assert _echoed_description(second.call(f"{alias}-add", probe)) == "Add two integers" assert tool_calls(peer.drain()) == (), "a blocked probe reached the peer" + + +@contextmanager +def _scope_enforcing_peer(auth: AuthorizationServer) -> Iterator[tuple[McpPeer, Wire]]: + with scripted_peer( + ScriptedTool("read_op", lambda _: text_result("read completed")), + ScriptedTool("write_op", lambda _: text_result("write completed")), + ScriptedTool("forbid_op", lambda _: text_result("must not execute")), + ) as inner: + def enforce(request: Request) -> Reply: + body: Final = json.loads(request.body) if request.body else {} + if body.get("method") == "tools/call": + name: Final = body["params"]["name"] + if name == "forbid_op": + return Reply(status=403, body=b"denied") + required: Final = "tools.write" if name == "write_op" else "tools.read" + token: Final = request.headers.get("authorization", "").removeprefix("Bearer ") + with auth.lock: + granted: Final = auth.access_tokens.get(token, {}).get("scope", "") + if required not in granted.split(): + return Reply(status=403, headers={"WWW-Authenticate": f'Bearer error="insufficient_scope", scope="{required}"'}) + forwarded: Final = httpx.request(request.method, inner.url, content=request.body, headers=dict(request.headers)) + return Reply(status=forwarded.status_code, body=forwarded.content, content_type=forwarded.headers.get("content-type", "application/json")) + + with wire_server(enforce) as wire: + yield McpPeer(wire.url + "/mcp", inner.calls), wire + + +@pytest.mark.parametrize("entry", ["rest", "server_mcp"]) +def test_managed_scope_step_up_preserves_read_and_write_without_retry(gateway: Gateway, entry: EntryPoint) -> None: + with oauth_server(scopes=("tools.read", "tools.write")) as auth, _scope_enforcing_peer(auth) as (peer, wire), gateway.scenario() as scenario: + alias: Final = "scope" + uuid.uuid4().hex[:8] + identity: Final = _register_oauth( + scenario, peer, auth, alias, auth_type="oauth2", oauth2_flow="authorization_code", + credentials={"client_id": "scope-client", "client_secret": "scope-secret", "scopes": ["tools.read"]}, + ) + owner: Final = scenario.key(user_id=scenario.user(), object_permission={"mcp_servers": [identity]}) + stranger: Final = scenario.key(user_id=scenario.user(), object_permission={"mcp_servers": [identity]}) + pkce: Final = _Pkce(secrets.token_urlsafe(32)) + code: Final = _authorize_through_gateway(gateway, auth, alias, owner, "scope-client", pkce, "tools.read") + first: Final = _redeem(gateway, alias, owner, "scope-client", code, pkce) + assert first.status_code == 200, first.text + caller: Final = McpCaller(gateway, owner, entry, alias) + read: Final = caller.call(f"{alias}-read_op", {}, identity) + assert read.ok and read.text == "read completed", read.raw + peer.drain() + wire.drain() + denied: Final = ( + call_tool(gateway, owner, identity, f"{alias}-write_op", {}) if entry == "rest" + else caller.rpc("tools/call", {"name": f"{alias}-write_op", "arguments": {}}) + ) + assert denied.status_code == 403, denied.text + challenge: Final = denied.headers["www-authenticate"] + assert 'error="insufficient_scope"' in challenge and 'scope="tools.write"' in challenge + assert f'resource_metadata="{_base(gateway)}/.well-known/oauth-protected-resource/' in challenge + attempts: Final = tuple(request for request in wire.drain() if request.body and json.loads(request.body).get("method") == "tools/call") + assert len(attempts) == 1 and tool_calls(peer.drain()) == (), "denied operation was retried or executed" + second_code: Final = _authorize_through_gateway(gateway, auth, alias, owner, "scope-client", pkce, "tools.write") + second: Final = _redeem(gateway, alias, owner, "scope-client", second_code, pkce) + assert second.status_code == 200, second.text + with auth.lock: + granted: Final = auth.access_tokens[_issued_token(second.json())]["scope"] + assert granted.split() == ["tools.read", "tools.write"] + for operation in ("read", "write"): + result: Final = caller.call(f"{alias}-{operation}_op", {}, identity) + assert result.ok and result.text == f"{operation} completed", result.raw + forbidden: Final = caller.call(f"{alias}-forbid_op", {}, identity) + assert forbidden.status == 200 and not forbidden.ok, forbidden.raw + unauthorized: Final = McpCaller(gateway, stranger, entry, alias).call(f"{alias}-read_op", {}, identity) + assert unauthorized.status == 401, unauthorized.raw + + +@pytest.mark.parametrize("route", ["/mcp", "/toolset/{toolset}/mcp"]) +def test_managed_step_up_advertises_sdk_compatible_metadata_and_bound_permissions(gateway: Gateway, route: str) -> None: + from mcp.shared.auth_utils import check_resource_allowed + + with ( + oauth_server(scopes=("tools.read", "tools.write")) as auth, + _scope_enforcing_peer(auth) as (peer, wire), + gateway.scenario() as scenario, + ): + alias: Final = "scope" + uuid.uuid4().hex[:8] + identity: Final = _register_oauth( + scenario, + peer, + auth, + alias, + auth_type="oauth2", + oauth2_flow="authorization_code", + credentials={"client_id": "scope-client", "client_secret": "scope-secret", "scopes": ["tools.read"]}, + ) + toolset_name: Final = "scope" + uuid.uuid4().hex[:8] + toolset_id: Final = create_toolset( + scenario, ((identity, "read_op"), (identity, "write_op")), toolset_name=toolset_name + ) + owner: Final = scenario.key( + user_id=scenario.user(), object_permission={"mcp_servers": [identity], "mcp_toolsets": [toolset_id]} + ) + pkce: Final = _Pkce(secrets.token_urlsafe(32)) + code: Final = _authorize_through_gateway(gateway, auth, alias, owner, "scope-client", pkce, "tools.read") + first: Final = _redeem(gateway, alias, owner, "scope-client", code, pkce) + assert first.status_code == 200, first.text + path: Final = route.format(toolset=toolset_name) + headers: Final = {**ACCEPT, "Authorization": f"Bearer {owner}"} + initialized: Final = gateway.client.post( + path, headers=headers, json={"jsonrpc": "2.0", "id": 1, "method": "initialize", "params": dict(INITIALIZE)} + ) + assert initialized.status_code == 200, initialized.text + peer.drain() + wire.drain() + response: Final = gateway.client.post( + path, + headers=headers, + json={ + "jsonrpc": "2.0", + "id": 2, + "method": "tools/call", + "params": {"name": f"{alias}-write_op", "arguments": {}}, + }, + ) + assert response.status_code == 403, response.text + challenge: Final = response.headers["www-authenticate"] + metadata_url: Final = challenge.split('resource_metadata="', 1)[1].split('"', 1)[0] + metadata: Final = gateway.client.get(metadata_url) + assert metadata.status_code == 200, metadata.text + assert check_resource_allowed( + requested_resource=_base(gateway) + path, configured_resource=metadata.json()["resource"] + ) + encoded_scope: Final = challenge.split('scope="', 1)[1].split('"', 1)[0] + assert encoded_scope == f"litellm:mcp_scope:{identity.encode().hex()}:{b'tools.write'.hex()}" + assert tool_calls(peer.drain()) == (), "denied operation executed upstream" + attempts: Final = tuple( + request + for request in wire.drain() + if request.body and json.loads(request.body).get("method") == "tools/call" + ) + assert len(attempts) == 1, "denied operation was retried" diff --git a/tests/integration/mcp/test_mcp_toolsets.py b/tests/integration/mcp/test_mcp_toolsets.py index 7ff512a699b..f275bd0e427 100644 --- a/tests/integration/mcp/test_mcp_toolsets.py +++ b/tests/integration/mcp/test_mcp_toolsets.py @@ -524,8 +524,7 @@ def test_owned_proxy_loss_retains_exit_status_and_first_http_failure(gateway: Ga with owned_proxy_process(gateway, tmp_path, {}) as owned: response: Final = owned.gateway.client.get("/health/readiness") assert response.status_code == 200, response.text - owned.process.kill() - owned.process.wait(timeout=3) + owned.kill_supervisor() owned.gateway.client.get("/health/readiness") record: Final = TypeAdapter(dict[str, int | float | str | None]).validate_json( owned.log.with_suffix(".exit.json").read_text() diff --git a/tests/integration/mcp/test_mcp_user_env_vars.py b/tests/integration/mcp/test_mcp_user_env_vars.py index 0170eb77491..75bb9ede182 100644 --- a/tests/integration/mcp/test_mcp_user_env_vars.py +++ b/tests/integration/mcp/test_mcp_user_env_vars.py @@ -1,4 +1,3 @@ -import signal import uuid from collections.abc import Mapping from concurrent.futures import ThreadPoolExecutor @@ -311,8 +310,7 @@ def test_concurrent_users_across_processes_never_leak_and_survive_a_killed_proce outcomes: Final = burst(pool, (gateway, peer, doomed.gateway)) assert all(code in {200, 412} for _, code in outcomes), outcomes assert all(code == 200 for index, code in outcomes if index % 4 != 2), outcomes - doomed.process.send_signal(signal.SIGKILL) - doomed.process.wait(timeout=10) + doomed.kill_supervisor() after_kill: Final = burst(pool, (gateway, peer)) assert all(code in {200, 412} for _, code in after_kill), after_kill assert all(code == 200 for index, code in after_kill if index % 4 != 2), after_kill diff --git a/tests/integration/mcp/test_proxy_mcp_e2e.py b/tests/integration/mcp/test_proxy_mcp_e2e.py index 4cfac6ed784..959d2620f58 100644 --- a/tests/integration/mcp/test_proxy_mcp_e2e.py +++ b/tests/integration/mcp/test_proxy_mcp_e2e.py @@ -41,6 +41,8 @@ from litellm.proxy.proxy_server import ( initialize, ) from tests._master_key import MASTER_KEY +from tests.integration._support.process import stop_child +from tests.integration._support.timing import assert_enforced_within from tests.integration._support.wire import Reply, Wire, wire_server from tests.integration._support.wire import Request as WireRequest @@ -83,7 +85,10 @@ async def drain_logging_worker() -> typing.AsyncIterator[None]: async def test_schema_validation_stops_expensive_work_and_recovers(cancel: bool, concurrency: int) -> None: import psutil - from litellm.proxy._experimental.mcp_server.tool_search import _tool_argument_validation_error + from litellm.proxy._experimental.mcp_server.tool_search import ( + _VALIDATION_TIMEOUT_SECONDS, + _tool_argument_validation_error, + ) existing_children: Final = frozenset(child.pid for child in psutil.Process().children()) warm_count: Final = min(concurrency, 4) @@ -117,7 +122,7 @@ async def test_schema_validation_stops_expensive_work_and_recovers(cancel: bool, assert all(isinstance(result, asyncio.CancelledError) for result in await group) else: assert await group == ["Tool argument validation exceeded its time limit"] * concurrency - assert time.monotonic() - started < 35 + assert_enforced_within(time.monotonic() - started, configured=_VALIDATION_TIMEOUT_SECONDS, slack=5) assert all(not worker.is_running() for worker in workers), "Cancelled validation must terminate worker CPU work" assert await _tool_argument_validation_error({"type": "object"}, {}) is None finally: @@ -391,12 +396,7 @@ def _math_http_server(offset: int) -> typing.Iterator[str]: time.sleep(0.05) yield f"http://{host}:{port}" finally: - process.terminate() - try: - process.wait(timeout=5) - except subprocess.TimeoutExpired: - process.kill() - process.wait(timeout=5) + stop_child(process, grace_seconds=5) @pytest.fixture(scope="module") diff --git a/tests/integration/mcp_coverage.toml b/tests/integration/mcp_coverage.toml index 2357269b59f..daa782eade2 100644 --- a/tests/integration/mcp_coverage.toml +++ b/tests/integration/mcp_coverage.toml @@ -1,5 +1,6 @@ [tool.coverage.run] -branch = true +branch = false +core = "sysmon" parallel = true relative_files = true include = [ diff --git a/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_messages_pre_content_retry_chaos.py b/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_messages_pre_content_retry_chaos.py index b6e848bad75..8944d1990d9 100644 --- a/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_messages_pre_content_retry_chaos.py +++ b/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_messages_pre_content_retry_chaos.py @@ -2,7 +2,6 @@ import asyncio import base64 import json import re -import signal import threading import uuid from collections.abc import Callable, Mapping @@ -364,15 +363,15 @@ async def test_worker_sigkill_mid_burst_leaves_the_sibling_retrying_pre_content_ held_by: Final = MappingProxyType({pid: _open_upstream_connections(pid, wire.url) for pid in workers}) assert sum(held_by.values()) == len(calls), held_by victim_pid, survivor_pid = sorted(workers, key=held_by.__getitem__) - victim: Final = psutil.Process(victim_pid) - victim.suspend() - victim.send_signal(signal.SIGKILL) - release.set() - served: Final = await burst - assert held_by[survivor_pid] >= 10, held_by - assert len(served) == held_by[survivor_pid], (held_by, len(served)) - for item in served: - _assert_completed_on_the_second_attempt(item) + with owned.replacing_worker(victim_pid) as kill: + kill.freeze() + await asyncio.to_thread(kill) + release.set() + served: Final = await burst + assert held_by[survivor_pid] >= 10, held_by + assert len(served) == held_by[survivor_pid], (held_by, len(served)) + for item in served: + _assert_completed_on_the_second_attempt(item) follow_up: Final = _Call(endpoint="messages", stream=True, marker=uuid.uuid4().hex) (answered,) = await _burst(base_url, candidate.key, models, (follow_up,)) _assert_completed_on_the_second_attempt(answered) diff --git a/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_messages_timeout_wire.py b/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_messages_timeout_wire.py index 76c29ad7763..fb0a9d24c0f 100644 --- a/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_messages_timeout_wire.py +++ b/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_messages_timeout_wire.py @@ -5,6 +5,7 @@ from typing import Final import pytest from integration._support.client import JSON_OBJECT, Gateway +from integration._support.timing import assert_cut_short, timed from integration._support.wire import Reply, Request, wire_server _UPSTREAM_STALL_SECONDS: Final = 4.0 @@ -46,13 +47,13 @@ def test_messages_endpoint_honors_configured_timeout_against_stalled_upstream(ga api_key="synthetic-anthropic-key", timeout=_CONFIGURED_TIMEOUT_SECONDS, ) - started: Final = time.monotonic() - response: Final = gateway.request( - "POST", - "/v1/messages", - {"model": model, "max_tokens": 16, "messages": [{"role": "user", "content": prompt}]}, + response: Final = timed( + lambda: gateway.request( + "POST", + "/v1/messages", + {"model": model, "max_tokens": 16, "messages": [{"role": "user", "content": prompt}]}, + ) ) - elapsed: Final = time.monotonic() - started - assert response.status_code == 408, response.text - assert elapsed < _UPSTREAM_STALL_SECONDS, f"timed out only after {elapsed:.2f}s: {response.text}" + assert response.value.status_code == 408, response.value.text + assert_cut_short(response.seconds, injected_delay=_UPSTREAM_STALL_SECONDS) assert len(wire.drain()) == 1 diff --git a/tests/integration/messages_endpoint/providers/anthropic/test_pdf_document_count_tokens_chaos.py b/tests/integration/messages_endpoint/providers/anthropic/test_pdf_document_count_tokens_chaos.py index f196018f117..80ace4c4535 100644 --- a/tests/integration/messages_endpoint/providers/anthropic/test_pdf_document_count_tokens_chaos.py +++ b/tests/integration/messages_endpoint/providers/anthropic/test_pdf_document_count_tokens_chaos.py @@ -1,7 +1,5 @@ import asyncio -import os import re -import signal import threading import uuid from collections.abc import Callable @@ -209,11 +207,13 @@ async def test_worker_sigkill_mid_pdf_count_burst_leaves_the_sibling_pricing_pag ) held_by: Final = MappingProxyType({pid: _held_upstream_connections(pid, wire.url) for pid in workers}) victim: Final = max(workers, key=held_by.__getitem__) - os.kill(victim, signal.SIGKILL) + with owned.replacing_worker(victim) as kill: + await asyncio.to_thread(kill) + release.set() + served: Final = await burst + wire.drain() finally: release.set() - served: Final = await burst - wire.drain() baseline: Final = _baseline(owned.gateway) after: Final = await _fire_counts(owned_url, owned.gateway.key) after_received: Final = wire.drain() diff --git a/tests/integration/observability/test_arize_otel_v2_openinference_chaos.py b/tests/integration/observability/test_arize_otel_v2_openinference_chaos.py index 9ba196cb28e..79b0d9084aa 100644 --- a/tests/integration/observability/test_arize_otel_v2_openinference_chaos.py +++ b/tests/integration/observability/test_arize_otel_v2_openinference_chaos.py @@ -9,7 +9,6 @@ from typing import Final from urllib.parse import urlsplit import httpx -import psutil import pytest from _openinference_support import ( CHAT_TOOLS, @@ -334,13 +333,6 @@ def test_arize_otel_v2_f2_slow_sink_does_not_deadlock(gateway: Gateway, tmp_path _matching_marker_span(rig.destination, marker) -def _has_exited(process: psutil.Process) -> bool: - try: - return process.status() == psutil.STATUS_ZOMBIE - except psutil.NoSuchProcess: - return True - - @pytest.mark.timeout(2 * graceful_stop_seconds() + 120) def test_arize_otel_v2_f3_one_proxy_worker_can_die(gateway: Gateway, tmp_path: Path) -> None: markers: Final = tuple("f3-" + uuid.uuid4().hex for _ in range(8)) @@ -365,11 +357,8 @@ def test_arize_otel_v2_f3_one_proxy_worker_can_die(gateway: Gateway, tmp_path: P destination_handler=sink, fresh_client_connections=True, ) as rig: - children: Final = psutil.Process(rig.owned.process.pid).children(recursive=True) - workers: Final = tuple( - child for child in children if child.is_running() and "resource_tracker" not in " ".join(child.cmdline()) - ) - assert len(workers) >= 2, tuple((worker.pid, worker.name()) for worker in workers) + workers: Final = rig.owned.worker_pids() + assert len(workers) >= 2, workers with ThreadPoolExecutor(max_workers=len(markers)) as executor: try: futures: Final = tuple( @@ -388,12 +377,12 @@ def test_arize_otel_v2_f3_one_proxy_worker_can_die(gateway: Gateway, tmp_path: P seconds=10, ) assert observed >= 2, observed - workers[0].kill() - assert eventually(lambda: _has_exited(workers[0]), bool, seconds=10), workers[0] - release.set() - in_flight: Final = tuple( - (marker, future.result(timeout=60)) for marker, future in zip(markers, futures, strict=True) - ) + with rig.owned.replacing_worker(min(workers)) as kill: + kill() + release.set() + in_flight: Final = tuple( + (marker, future.result(timeout=60)) for marker, future in zip(markers, futures, strict=True) + ) finally: release.set() survivor: Final = "f3-survivor-" + uuid.uuid4().hex diff --git a/tests/integration/observability/test_azure_content_safety_audit.py b/tests/integration/observability/test_azure_content_safety_audit.py index 99d868efe6f..85ced5f719b 100644 --- a/tests/integration/observability/test_azure_content_safety_audit.py +++ b/tests/integration/observability/test_azure_content_safety_audit.py @@ -7,7 +7,6 @@ from contextlib import ExitStack from pathlib import Path from typing import Final -import psutil import pytest import yaml from integration._support.client import Gateway, eventually, gateway_from_environment, object_value @@ -953,13 +952,8 @@ def test_killing_a_worker_mid_burst_leaves_no_duplicate_rows( chaos_rig: tuple[OwnedProxy, Wire, Wire, threading.Event], ) -> None: owned, azure, provider, _ = chaos_rig - port: Final = owned.gateway.client.base_url.port - workers: Final = tuple( - child - for child in psutil.Process(owned.process.pid).children(recursive=False) - if any(connection.laddr.port == port for connection in child.net_connections(kind="tcp")) - ) - assert len(workers) == 2, [worker.pid for worker in workers] + workers: Final = owned.worker_pids() + assert len(workers) == 2, workers with owned.gateway.scenario() as scenario: model: Final = scenario.model( model="openai/gpt-4.1-mini", api_base=provider.url + "/v1", api_key="synthetic-provider-key" @@ -972,11 +966,12 @@ def test_killing_a_worker_mid_burst_leaves_no_duplicate_rows( with ThreadPoolExecutor(max_workers=6) as pool: future_map: Final = tuple(pool.submit(call, identity) for identity in identities) - workers[0].kill() - outcomes: Final = tuple( - future.result() if not future.exception() else (identities[index], -1) - for index, future in enumerate(future_map) - ) + with owned.replacing_worker(min(workers)) as kill: + kill() + outcomes: Final = tuple( + future.result() if not future.exception() else (identities[index], -1) + for index, future in enumerate(future_map) + ) survivors: Final = tuple(status for _, status in outcomes if status != -1) assert survivors and {status for status in survivors} == {200}, outcomes scans: Final = _shield_prompts(azure.drain()) diff --git a/tests/integration/observability/test_azure_content_safety_dispatch_resilience.py b/tests/integration/observability/test_azure_content_safety_dispatch_resilience.py index 5d7b5948d96..d443ab9a2d0 100644 --- a/tests/integration/observability/test_azure_content_safety_dispatch_resilience.py +++ b/tests/integration/observability/test_azure_content_safety_dispatch_resilience.py @@ -2,18 +2,15 @@ import concurrent.futures import socket import threading import uuid -from collections.abc import Iterator from contextlib import ExitStack from dataclasses import dataclass from pathlib import Path from typing import Final import httpx -import psutil import pytest from integration._support.client import Gateway, Scenario, eventually from integration._support.database import read_rows -from integration._support.process import OwnedProxy from integration._support.redis_process import owned_redis from integration._support.wire import Wire, wire_server from integration.observability.azure_dispatch_support import ( @@ -25,7 +22,6 @@ from integration.observability.azure_dispatch_support import ( dispatch_rig, provider_handler, provider_texts, - write_dispatch_config, ) from pydantic import JsonValue @@ -104,16 +100,6 @@ def _safe_request(candidate: Gateway, model: str, prompt: str) -> httpx.Response return None -def _worker_processes(owned: OwnedProxy) -> tuple[psutil.Process, ...]: - descendants: Final = tuple(psutil.Process(owned.process.pid).children(recursive=True)) - return tuple( - process - for process in descendants - if process.is_running() - and any("spawn_main" in argument for argument in process.cmdline()) - ) - - def _operation_specs( phase: str, operation: str, @@ -388,8 +374,8 @@ def test_x3_killing_one_worker_leaves_the_other_serving( _, ): candidate: Final = owned.gateway - workers: Final = _worker_processes(owned) - assert len(workers) == 2, tuple(process.cmdline() for process in workers) + workers: Final = owned.worker_pids() + assert len(workers) == 2, workers with candidate.scenario() as scenario: model: Final = _model(scenario, provider) with concurrent.futures.ThreadPoolExecutor(max_workers=12) as executor: @@ -398,17 +384,17 @@ def test_x3_killing_one_worker_leaves_the_other_serving( for _ in range(4) ) assert entered.wait(timeout=30), "No request reached the slow Azure edge" - killed_worker: Final = workers[0] - survivor: Final = workers[1] - killed_worker.kill() - killed_worker.wait(timeout=10) - assert survivor.is_running(), "The second proxy worker exited with the killed worker" - prompts: Final = tuple(f"worker survivor {uuid.uuid4().hex}" for _ in range(8)) - futures: Final = tuple( - executor.submit(_safe_request, candidate, model, prompt) - for prompt in prompts - ) - responses: Final = tuple(future.result(timeout=70) for future in (*inflight, *futures)) + killed_worker: Final = min(workers) + survivor: Final = max(workers) + with owned.replacing_worker(killed_worker) as kill: + kill() + assert survivor in owned.worker_pids(), "The second proxy worker exited with the killed worker" + prompts: Final = tuple(f"worker survivor {uuid.uuid4().hex}" for _ in range(8)) + futures: Final = tuple( + executor.submit(_safe_request, candidate, model, prompt) + for prompt in prompts + ) + responses: Final = tuple(future.result(timeout=70) for future in (*inflight, *futures)) surviving_responses: Final = tuple(response for response in responses[4:] if response is not None) assert all(response.status_code == 200 for response in surviving_responses), tuple( response.text for response in surviving_responses @@ -470,7 +456,7 @@ def test_x4_proxy_restart_preserves_completed_spend_rows( assert arrived_count == len(interrupted_prompts), ( f"Only {arrived_count} of {len(interrupted_prompts)} interrupted requests reached Azure" ) - first.process.terminate() + first.begin_shutdown() release.set() first_proxy_lifetime.close() interrupted_results: Final = tuple(future.result(timeout=70) for future in interrupted) diff --git a/tests/integration/observability/test_azure_storage_chaos.py b/tests/integration/observability/test_azure_storage_chaos.py index 2235dc8fc56..acc2eb1439d 100644 --- a/tests/integration/observability/test_azure_storage_chaos.py +++ b/tests/integration/observability/test_azure_storage_chaos.py @@ -1,5 +1,3 @@ -import os -import signal import uuid from pathlib import Path from typing import Final @@ -149,13 +147,14 @@ def test_killing_one_worker_keeps_the_other_serving_and_uploading(gateway: Gatew key: Final = scenario.key(models=[openai_model, anthropic_model]) first: Final = mixed_burst(candidate, openai_model, anthropic_model, key, f"{marker}-first", per_surface=2) collect_files(sink, len(first)) - workers: Final = tuple( - process for process in group_members(owned.process.pid) if process.pid != owned.process.pid - ) + workers: Final = owned.worker_pids() assert workers, "no uvicorn workers in the owned proxy process group" - os.kill(workers[0].pid, signal.SIGKILL) - eventually(lambda: _readiness_ok(candidate), lambda ok: ok, seconds=30) - rest: Final = mixed_burst(candidate, openai_model, anthropic_model, key, f"{marker}-rest", per_surface=4) + with owned.replacing_worker(min(workers)) as kill: + kill() + eventually(lambda: _readiness_ok(candidate), lambda ok: ok, seconds=30) + rest: Final = mixed_burst( + candidate, openai_model, anthropic_model, key, f"{marker}-rest", per_surface=4 + ) payloads: Final = eventually( lambda: tuple(sink.payloads().values()), lambda stored: _present_count(stored, rest) == len(rest), diff --git a/tests/integration/observability/test_cache_hit_guardrail_metrics.py b/tests/integration/observability/test_cache_hit_guardrail_metrics.py index 888cfd7ab9d..605130acd38 100644 --- a/tests/integration/observability/test_cache_hit_guardrail_metrics.py +++ b/tests/integration/observability/test_cache_hit_guardrail_metrics.py @@ -1,6 +1,5 @@ import asyncio import json -import subprocess import uuid from collections.abc import Callable, Generator, Mapping from contextlib import contextmanager @@ -15,7 +14,7 @@ import pytest import yaml from integration._support.client import Gateway, Scenario, eventually, object_value, string_value from integration._support.database import read_rows -from integration._support.process import owned_proxy_process +from integration._support.process import OwnedProxy, owned_proxy_process from integration._support.wire import Reply, Request, Wire, wire_server from prometheus_client.parser import text_string_to_metric_families @@ -178,7 +177,7 @@ class Rig: guardrail_name: str policy: Wire provider: Wire - process: subprocess.Popen[bytes] + owned: OwnedProxy @contextmanager @@ -228,7 +227,7 @@ def _rig( identity, policy, provider, - owned.process, + owned, ) diff --git a/tests/integration/observability/test_cache_hit_guardrail_metrics_chaos.py b/tests/integration/observability/test_cache_hit_guardrail_metrics_chaos.py index 7e3b479e5f0..d4b62b36e8b 100644 --- a/tests/integration/observability/test_cache_hit_guardrail_metrics_chaos.py +++ b/tests/integration/observability/test_cache_hit_guardrail_metrics_chaos.py @@ -1,5 +1,4 @@ import json -import signal import threading import uuid from collections.abc import Callable @@ -8,7 +7,6 @@ from pathlib import Path from typing import Final import httpx -import psutil from integration._support.client import Gateway, eventually, object_value, string_value from integration._support.process import owned_proxy_process from integration._support.redis_process import owned_redis @@ -239,14 +237,13 @@ def test_worker_kill_mid_burst_keeps_counting(gateway: Gateway, tmp_path: Path) with _rig(gateway, tmp_path, marker, workers=2) as rig: bodies: Final = _burst_bodies(rig, marker, None) _warm(rig, bodies) - children: Final = psutil.Process(rig.process.pid).children(recursive=True) - assert children, "no uvicorn worker children found" - with ThreadPoolExecutor(max_workers=8) as pool: + workers: Final = eventually(rig.owned.worker_pids, lambda pids: len(pids) == 2) + with rig.owned.replacing_worker(min(workers)) as kill, ThreadPoolExecutor(max_workers=8) as pool: futures: Final = tuple( pool.submit(lambda b: rig.candidate.request("POST", b[0], b[1]), body) for body in bodies ) eventually(lambda: rig.policy.received.qsize() >= 3, bool, seconds=30) - children[0].send_signal(signal.SIGKILL) + kill() statuses: list[int] = [] # mutable-ok: collect per-request outcomes from concurrent futures for future in futures: try: diff --git a/tests/integration/observability/test_grayswan_wire_chaos.py b/tests/integration/observability/test_grayswan_wire_chaos.py index 16800d6c235..b57e9ac9f75 100644 --- a/tests/integration/observability/test_grayswan_wire_chaos.py +++ b/tests/integration/observability/test_grayswan_wire_chaos.py @@ -1,6 +1,4 @@ import json -import os -import signal import threading import time import uuid @@ -8,10 +6,9 @@ from concurrent.futures import ThreadPoolExecutor from pathlib import Path from typing import Final -import psutil import yaml from integration._support.client import Gateway -from integration._support.process import group_members, owned_proxy_process +from integration._support.process import owned_proxy_process from integration._support.wire import Reply, Request, wire_server from pydantic import JsonValue, TypeAdapter from test_grayswan_wire import _PROVIDER_KEY, _REQUEST_MESSAGES, _VENDOR_KEY, _monitor_bodies, _serving_model_probe @@ -206,23 +203,13 @@ def test_worker_kill_mid_burst_survivor_keeps_serving(gateway: Gateway, tmp_path model: Final = scenario.model(model="openai/gpt-4o-mini", api_base=upstream.url, api_key=_PROVIDER_KEY) warm: Final = _fire(candidate, model, "marker-warm", False) assert warm == 200 - members: Final = group_members(owned.process.pid) - candidate_port: Final = candidate.client.base_url.port - workers_listening: Final = tuple( - member - for member in members - if member.pid != owned.process.pid - and any( - connection.laddr.port == candidate_port and connection.status == "LISTEN" - for connection in member.net_connections(kind="inet") + workers: Final = owned.worker_pids() + assert len(workers) == 2, workers + with owned.replacing_worker(min(workers)) as kill: + kill() + statuses: Final = tuple( + _fire(candidate, model, f"marker-kill-{index}", False) for index in range(6) ) - ) - assert len(workers_listening) == 2, [member.pid for member in members] - victim: Final = workers_listening[0] - os.kill(victim.pid, signal.SIGKILL) - psutil.wait_procs((victim,), timeout=10) - assert not psutil.pid_exists(victim.pid), victim.pid - statuses: Final = tuple(_fire(candidate, model, f"marker-kill-{index}", False) for index in range(6)) assert all(status == 200 for status in statuses), statuses bodies: Final = _monitor_bodies(vendor_wire, expected=7) kill_bodies: Final = [ diff --git a/tests/integration/observability/test_guardrail_effects.py b/tests/integration/observability/test_guardrail_effects.py index 4acd9c73ec6..0009e21602b 100644 --- a/tests/integration/observability/test_guardrail_effects.py +++ b/tests/integration/observability/test_guardrail_effects.py @@ -1,8 +1,6 @@ import asyncio import json -import os import re -import signal import socket import threading import uuid @@ -17,13 +15,12 @@ from typing import Final import anthropic import httpx -import psutil import pytest import yaml from integration._support.client import Gateway, eventually, object_value from integration._support.database import read_rows from integration._support.mcp import mcp_peer, register_mcp, tool_names -from integration._support.process import OwnedProxy, group_members, owned_proxy, owned_proxy_process +from integration._support.process import OwnedProxy, owned_proxy, owned_proxy_process from integration._support.wire import Reply, Request, Wire, wire_server from openai import AsyncOpenAI, OpenAI from pydantic import JsonValue @@ -1772,27 +1769,9 @@ def test_responses_pre_call_denial_stream_survives_worker_kill(gateway: Gateway, candidate: Final = owned.gateway with candidate.scenario() as scenario: model: Final = scenario.model(api_base=_dead_api_base()) - members: Final = tuple( - member for member in group_members(owned.process.pid) if member.pid != owned.process.pid - ) - children: Final = tuple(member.pid for member in members) - workers: Final = tuple( - member.pid for member in members if any("spawn_main" in part for part in member.cmdline()) - ) + workers: Final = owned.worker_pids() assert len(workers) >= 2, workers - os.kill(workers[0], signal.SIGKILL) - expected: Final = len(children) - eventually( - lambda: tuple( - member.pid - for member in group_members(owned.process.pid) - if member.pid != owned.process.pid - and member.is_running() - and member.status() != psutil.STATUS_ZOMBIE - ), - lambda pids: len(pids) >= expected and any(pid not in children for pid in pids), - seconds=30, - ) + owned.kill_worker(min(workers)) def burst(index: int) -> httpx.Response: return candidate.request( @@ -2737,14 +2716,10 @@ def test_concurrent_disconnects_during_a_guardrail_outage_each_record_exactly_on def test_disconnect_scans_keep_recording_after_a_worker_is_killed(gateway: Gateway, tmp_path: Path) -> None: with _disconnect_rig(gateway, tmp_path, gated=False, pause=3, workers=2) as rig: - members: Final = tuple( - member for member in group_members(rig.owned.process.pid) if member.pid != rig.owned.process.pid - ) - workers: Final = tuple(member for member in members if any("spawn_main" in part for part in member.cmdline())) - assert len(workers) >= 2, members - workers[0].send_signal(signal.SIGKILL) - psutil.wait_procs((workers[0],), timeout=10) - with ThreadPoolExecutor(max_workers=8) as pool: + workers: Final = rig.owned.worker_pids() + assert len(workers) >= 2, workers + with rig.owned.replacing_worker(min(workers)) as kill, ThreadPoolExecutor(max_workers=8) as pool: + kill() received: Final = tuple(pool.map(lambda index: _chat_httpx(rig, index), range(8))) assert all(rig.secret(index) in line for index, line in enumerate(received)), received assert _post_call_statuses(rig, rig.model, rows=8) == (("success",),) * 8 diff --git a/tests/integration/observability/test_guardrail_stream_scope_chaos.py b/tests/integration/observability/test_guardrail_stream_scope_chaos.py index f8a9d0a4bba..6aa705e52b3 100644 --- a/tests/integration/observability/test_guardrail_stream_scope_chaos.py +++ b/tests/integration/observability/test_guardrail_stream_scope_chaos.py @@ -3,7 +3,6 @@ from __future__ import annotations import json import os import shutil -import signal import socket import subprocess import uuid @@ -18,7 +17,6 @@ from types import MappingProxyType from typing import Final, Literal, TypeAlias, cast import httpx -import psutil import pytest import yaml from integration._support.client import Gateway, JsonValue, eventually @@ -596,19 +594,6 @@ def test_h2_stream_sink_gate_does_not_block_out_of_scope_calls(rig: ChaosRig, tm ) -def _worker_processes(owned: OwnedProxy) -> tuple[psutil.Process, ...]: - return tuple( - child for child in psutil.Process(owned.process.pid).children() if "spawn_main" in _process_command(child) - ) - - -def _process_command(process: psutil.Process) -> str: - try: - return " ".join(process.cmdline()) - except psutil.Error: - return "" - - def _safe_response(future: Future[httpx.Response]) -> httpx.Response | None: try: return future.result(timeout=70) @@ -636,14 +621,14 @@ def test_h3_worker_kill_mid_burst_keeps_remaining_worker_serving(rig: ChaosRig, _rail(name_b, sink_b.url, "non_streaming", default_on=True), ) with _owned_proxy(rig, tmp_path, rails, workers=2) as owned, ThreadPoolExecutor(max_workers=30) as pool: - workers: Final = _worker_processes(owned) - assert len(workers) == 2, tuple(worker.pid for worker in workers) + workers: Final = owned.worker_pids() + assert len(workers) == 2, workers futures: Final[tuple[Future[httpx.Response], ...]] = tuple( pool.submit(_request, owned.gateway, plan, ()) for plan in plans ) assert started.wait(timeout=30), "stream requests did not reach the owned sink" - victim: Final = workers[0] - survivor: Final = workers[1] + victim: Final = min(workers) + survivor: Final = max(workers) survivor_plan: Final = CallPlan( f"audit-h3-survivor-{uuid.uuid4().hex}", f"h3-survivor-{uuid.uuid4().hex}", @@ -651,13 +636,14 @@ def test_h3_worker_kill_mid_burst_keeps_remaining_worker_serving(rig: ChaosRig, False, ) try: - victim.send_signal(signal.SIGKILL) - survivor_response: Final = _request(owned.gateway, survivor_plan, ()) - assert survivor_response.status_code == 200 and survivor_plan.marker in survivor_response.text, ( - survivor_response.status_code, - survivor_response.text, - ) - assert survivor.is_running(), survivor.pid + with owned.replacing_worker(victim) as kill: + kill() + survivor_response: Final = _request(owned.gateway, survivor_plan, ()) + assert survivor_response.status_code == 200 and survivor_plan.marker in survivor_response.text, ( + survivor_response.status_code, + survivor_response.text, + ) + assert survivor in owned.worker_pids(), survivor finally: release.set() responses: Final = tuple(_safe_response(future) for future in futures) diff --git a/tests/integration/observability/test_guardrail_timeout_all_providers.py b/tests/integration/observability/test_guardrail_timeout_all_providers.py index 1f33bd19393..8f6c098a3d4 100644 --- a/tests/integration/observability/test_guardrail_timeout_all_providers.py +++ b/tests/integration/observability/test_guardrail_timeout_all_providers.py @@ -13,6 +13,7 @@ from __future__ import annotations import json import re import socket +import sys import threading import time from collections.abc import Iterator, Mapping @@ -33,11 +34,13 @@ from cryptography.hazmat.primitives import serialization from cryptography.hazmat.primitives.asymmetric import rsa from integration._support.client import Gateway, gateway_from_environment from integration._support.process import owned_proxy_process +from integration._support.timing import assert_enforced_within from integration._support.wire import Reply, Request, wire_server SLOW_SECONDS: Final = 20 FAST_SECONDS: Final = 3 -BOUND_SECONDS: Final = 8 +GUARDRAIL_TIMEOUT_SECONDS: Final = 1 +TIMEOUT_SLACK_SECONDS: Final = 7 PROBE_CONCURRENCY: Final = 8 TOKEN_PATH: Final = "/token" TOKEN_REPLY: Final = json.dumps( @@ -173,6 +176,7 @@ class Sink: port: int seen: list[Seen] = field(default_factory=list) lock: threading.Lock = field(default_factory=threading.Lock) + closing: threading.Event = field(default_factory=threading.Event) server: ThreadingHTTPServer | None = None thread: threading.Thread | None = None @@ -194,7 +198,7 @@ class Sink: ) is_token: Final = self.path.startswith(TOKEN_PATH) if not is_token: - time.sleep(SLOW_SECONDS if self.path.startswith("/slow/") else FAST_SECONDS) + sink.closing.wait(SLOW_SECONDS if self.path.startswith("/slow/") else FAST_SECONDS) payload: Final = TOKEN_REPLY if is_token else b"{}" self.send_response(200) self.send_header("content-type", "application/json") @@ -212,9 +216,14 @@ class Sink: class Server(ThreadingHTTPServer): allow_reuse_address = True - daemon_threads = True + daemon_threads = False + block_on_close = True request_queue_size: Final = socket.SOMAXCONN + def handle_error(self, request: object, client_address: object) -> None: + if not isinstance(sys.exc_info()[1], ConnectionError): + super().handle_error(request, client_address) + self.server = Server(("127.0.0.1", self.port), Handler) self.port = self.server.server_port self.thread = threading.Thread(target=self.server.serve_forever, daemon=True) @@ -222,6 +231,7 @@ class Sink: def stop(self) -> None: assert self.server is not None and self.thread is not None + self.closing.set() self.server.shutdown() self.server.server_close() self.thread.join(timeout=5) @@ -323,7 +333,7 @@ def _rig_config(sink_url: str, root: Path) -> Path: config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) config["litellm_settings"]["cache"] = False config["guardrails"] = [ - _guardrail(name, provider, sink_url, 1, dict(extra), mode) + _guardrail(name, provider, sink_url, GUARDRAIL_TIMEOUT_SECONDS, dict(extra), mode) for name, provider, extra, mode, _ in _provider_values() ] + [ _guardrail("control-generic", "generic_guardrail_api", sink_url, None, {}, "pre_call"), @@ -431,9 +441,7 @@ def test_litellm_params_timeout_bounds_outbound_call( outcome: Final = outcomes[name] calls: Final = rig.sink.calls_for(name) assert calls, f"{name}: sink saw no request for {provider}" - assert outcome.elapsed < BOUND_SECONDS, ( - f"{name}: elapsed {outcome.elapsed:.2f}s, expected under {BOUND_SECONDS}s with timeout=1" - ) + assert_enforced_within(outcome.elapsed, configured=GUARDRAIL_TIMEOUT_SECONDS, slack=TIMEOUT_SLACK_SECONDS) assert isinstance(outcome.response, httpx.Response), f"{name}: client gave up: {outcome.response!r}" assert outcome.response.status_code != 504, outcome.response.text diff --git a/tests/integration/observability/test_langfuse_delivery.py b/tests/integration/observability/test_langfuse_delivery.py index 018c6a35f31..bee754297f0 100644 --- a/tests/integration/observability/test_langfuse_delivery.py +++ b/tests/integration/observability/test_langfuse_delivery.py @@ -1,9 +1,7 @@ import asyncio import base64 import json -import signal import threading -import time import uuid from collections.abc import Iterator, Mapping, Sequence from concurrent.futures import ThreadPoolExecutor @@ -15,7 +13,6 @@ from typing import Final import anthropic import httpx import openai -import psutil import pytest import yaml from _s3_v2_support import _chat_stream_frames, _responses_stream_frames @@ -28,7 +25,8 @@ from integration._support.client import ( string_value, ) from integration._support.database import read_rows, scratch_database -from integration._support.process import group_members, owned_proxy, owned_proxy_process +from integration._support.process import owned_proxy, owned_proxy_process +from integration._support.timing import assert_cut_short, assert_within_factor_of_control, timed from integration._support.wire import Reply, Request, Wire, wire_server from opentelemetry.proto.collector.trace.v1.trace_service_pb2 import ExportTraceServiceRequest from opentelemetry.proto.common.v1.common_pb2 import KeyValue @@ -336,7 +334,10 @@ def test_prompt_fetch_encodes_the_name_retries_a_5xx_once_and_keeps_langfuse_hea flaky_prompt: Final = f"{marker}/what?" encoded_flaky_prompt: Final = f"{marker}%2Fwhat%3F" missing_prompt: Final = marker + "-missing" + steady_prompt: Final = marker + "-steady" + retry_after: Final = 30 + retried_within_control_factor: Final = 3 seen_prompt_gets: Final[list[str]] = [] # mutable-ok: the double counts attempts across requests def upstream(request: Request) -> Reply: @@ -353,8 +354,10 @@ def test_prompt_fetch_encodes_the_name_retries_a_5xx_once_and_keeps_langfuse_hea prior: Final = sum(1 for seen in seen_prompt_gets if seen.startswith(PROMPTS_PATH + encoded_flaky_prompt)) seen_prompt_gets.append(request.target) if prior == 0: - return Reply(status=503, body=b'{"message":"try later"}', headers={"retry-after": "30"}) + return Reply(status=503, body=b'{"message":"try later"}', headers={"retry-after": str(retry_after)}) return _text_prompt(flaky_prompt) + if request.target.startswith(PROMPTS_PATH + steady_prompt): + return _text_prompt(steady_prompt) seen_prompt_gets.append(request.target) return Reply( status=404, @@ -376,15 +379,24 @@ def test_prompt_fetch_encodes_the_name_retries_a_5xx_once_and_keeps_langfuse_hea missing: Final = scenario.model( model="langfuse/gpt-4o-mini", prompt_id=missing_prompt, api_base=provider.url + "/v1", api_key="synthetic" ) - started: Final = time.monotonic() - response: Final = candidate.request( - "POST", - "/v1/chat/completions", - {"model": flaky, "messages": [{"role": "user", "content": marker}], "prompt_variables": {"word": marker}}, + steady: Final = scenario.model( + model="langfuse/gpt-4o-mini", prompt_id=steady_prompt, api_base=provider.url + "/v1", api_key="synthetic" ) - elapsed: Final = time.monotonic() - started + + def cold_prompt_chat(model: str) -> httpx.Response: + return candidate.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": marker}], "prompt_variables": {"word": marker}}, + ) + + control: Final = timed(lambda: cold_prompt_chat(steady)) + assert control.value.status_code == 200, control.value.text + call: Final = timed(lambda: cold_prompt_chat(flaky)) + response: Final = call.value assert response.status_code == 200, response.text - assert elapsed < 5, f"a retried cold prompt miss took {elapsed:.1f}s" + assert_cut_short(call.seconds, injected_delay=retry_after) + assert_within_factor_of_control(call.seconds, control=control.seconds, factor=retried_within_control_factor) attempts: Final = tuple( target for target in seen_prompt_gets if target.startswith(PROMPTS_PATH + encoded_flaky_prompt) ) @@ -1855,12 +1867,8 @@ def test_audit_surviving_worker_keeps_serving_after_kill(gateway: Gateway, tmp_p ): candidate: Final = owned.gateway model: Final = scenario.model(api_base=provider.url + "/v1", api_key=provider_secret) - workers: Final = tuple( - member for member in group_members(owned.process.pid) if member.pid != owned.process.pid - ) + workers: Final = owned.worker_pids() assert workers, "expected worker processes under the owned proxy" - workers[0].send_signal(signal.SIGKILL) - psutil.wait_procs([workers[0]], timeout=5) def fire(index: int) -> tuple[str, httpx.Response]: caller_trace: Final = uuid.uuid4().hex @@ -1878,7 +1886,8 @@ def test_audit_surviving_worker_keeps_serving_after_kill(gateway: Gateway, tmp_p ) return caller_session, response - with ThreadPoolExecutor(max_workers=10) as pool: + with owned.replacing_worker(min(workers)) as kill, ThreadPoolExecutor(max_workers=10) as pool: + kill() answered: Final = tuple(pool.map(fire, range(10))) received: Final[list[Request]] = [] # mutable-ok: drain() consumes the queue, later polls keep earlier ones for caller_session, response in answered: diff --git a/tests/integration/observability/test_langtrace_delivery.py b/tests/integration/observability/test_langtrace_delivery.py index 03e42087dc7..93006cb6c6b 100644 --- a/tests/integration/observability/test_langtrace_delivery.py +++ b/tests/integration/observability/test_langtrace_delivery.py @@ -1,7 +1,7 @@ import asyncio import json import re -import signal +import threading import time import uuid from collections.abc import Callable, Iterator, Sequence @@ -14,12 +14,12 @@ from queue import SimpleQueue from typing import Final import httpx -import psutil import pytest import yaml from anthropic import Anthropic from integration._support.client import Gateway, eventually from integration._support.process import OwnedProxy, owned_proxy, owned_proxy_process +from integration._support.timing import assert_not_delayed_by, timed from integration._support.wire import Reply, Request, Wire, wire_server from openai import AsyncOpenAI, OpenAI from opentelemetry.proto.collector.trace.v1.trace_service_pb2 import ExportTraceServiceRequest @@ -627,15 +627,21 @@ def test_langtrace_sink_outage_mid_burst_recovers_on_the_same_port(gateway: Gate def test_langtrace_slow_sink_does_not_delay_callers_or_duplicate_spans(gateway: Gateway, tmp_path: Path) -> None: + sink_delay: Final = 3 + slowed: Final = threading.Event() + def slow(request: Request) -> Reply: - time.sleep(1) + if slowed.is_set(): + time.sleep(sink_delay) return _accepted(request) with _langtrace_rig(gateway, tmp_path, respond=slow) as rig: - started: Final = time.monotonic() - markers: Final = _burst(rig, 6) - assert time.monotonic() - started < 5 - _assert_each_once(rig.sink, markers, seconds=40) + control: Final = timed(lambda: _burst(rig, 6)) + _assert_each_once(rig.sink, control.value) + slowed.set() + observed: Final = timed(lambda: _burst(rig, 6)) + assert_not_delayed_by((observed.seconds,), control=(control.seconds,), injected_delay=sink_delay) + _assert_each_once(rig.sink, observed.value, seconds=40) def test_langtrace_survives_a_killed_worker(gateway: Gateway, tmp_path: Path) -> None: @@ -659,25 +665,10 @@ def test_langtrace_survives_a_killed_worker(gateway: Gateway, tmp_path: Path) -> _assert_kill_and_recovery(owned, rig) -def _cmdline(process: psutil.Process) -> str: - try: - return " ".join(process.cmdline()) - except psutil.Error: - return "" - - def _assert_kill_and_recovery(owned: OwnedProxy, rig: _Rig) -> None: _assert_delivered(rig, _Surface(_chat_httpx, False), _marker()) - def uvicorn_workers() -> tuple[psutil.Process, ...]: - return tuple(child for child in psutil.Process(owned.process.pid).children() if "spawn_main" in _cmdline(child)) - - workers: Final = uvicorn_workers() + workers: Final = owned.worker_pids() assert len(workers) == 2, workers - workers[0].send_signal(signal.SIGKILL) - eventually( - uvicorn_workers, - lambda value: len(value) == 2 and workers[0].pid not in {child.pid for child in value}, - seconds=30, - ) + owned.kill_worker(min(workers)) _assert_each_once(rig.sink, _burst(rig, 6)) diff --git a/tests/integration/observability/test_logging_only_scope_chaos.py b/tests/integration/observability/test_logging_only_scope_chaos.py index a34b13f048f..afcaab91b5f 100644 --- a/tests/integration/observability/test_logging_only_scope_chaos.py +++ b/tests/integration/observability/test_logging_only_scope_chaos.py @@ -1,6 +1,5 @@ from __future__ import annotations -import signal import socket import threading import time @@ -45,10 +44,13 @@ from _logging_only_scope_support import ( from anthropic import APIConnectionError as AnthropicAPIConnectionError from integration._support.client import Gateway, eventually from integration._support.process import owned_proxy, owned_proxy_process +from integration._support.timing import Timed, assert_not_delayed_by, timed from integration._support.wire import Reply, Request from openai import APIConnectionError as OpenAIAPIConnectionError from pydantic import JsonValue +_EDGE_DELAY_SECONDS: Final = 5 + @pytest.mark.timeout(180) def test_K1_policy_edge_restart_mid_burst_keeps_output_observation_fail_open(gateway: Gateway, tmp_path: Path) -> None: @@ -216,7 +218,7 @@ def test_K2_policy_edge_delay_does_not_delay_concurrent_callers(gateway: Gateway marker: Final = uuid.uuid4().hex def policy(_request: Request) -> Reply: - time.sleep(2) + time.sleep(_EDGE_DELAY_SECONDS) return Reply(body=b'{"action":"NONE"}') with gateway.scenario() as scenario: @@ -227,18 +229,23 @@ def test_K2_policy_edge_delay_does_not_delay_concurrent_callers(gateway: Gateway expected_directions: Final = _directions_for_audit_leg(("request", "response"), "output") control_config: Final = _chaos_control_configuration(tmp_path, identity, model_list) with owned_proxy(gateway, tmp_path, {}, config=control_config, workers=2) as control_proxy: - controls: Final = tuple( - _call_client( - call.client_kind, - call.endpoint, - control_proxy, - call.model, - call.prompt, - call.stream, - f"{call.call_id}-control", + + def run_control(call: ChaosCall) -> Timed[CallerResult]: + return timed( + lambda: _call_client( + call.client_kind, + call.endpoint, + control_proxy, + call.model, + call.prompt, + call.stream, + f"{call.call_id}-control", + ) ) - for call in chaos_calls - ) + + with ThreadPoolExecutor(max_workers=30) as control_pool: + timed_controls: Final = tuple(control_pool.map(run_control, chaos_calls)) + controls: Final = tuple(control.value for control in timed_controls) control_upstream: Final = _drain_upstream(gateway.upstream_url) assert len(control_upstream) == 30, control_upstream assert ( @@ -270,7 +277,11 @@ def test_K2_policy_edge_delay_does_not_delay_concurrent_callers(gateway: Gateway assert tuple(index for index, _, _ in results) == tuple(range(30)), results assert all(result.status == controls[index].status for index, result, _ in results), results assert all(result.text == controls[index].text for index, result, _ in results), results - assert all(duration < 2 for _, _, duration in results), results + assert_not_delayed_by( + (seconds for _, _, seconds in results), + control=(control.seconds for control in timed_controls), + injected_delay=_EDGE_DELAY_SECONDS, + ) eventually( lambda: edge.received.qsize(), lambda count: count == 30 * len(expected_directions), @@ -364,16 +375,7 @@ def test_K3_two_worker_sigkill_checks_post_kill_spend_rows( config: Final = _configuration(tmp_path, identity, edge.url, "output", model_list=model_list) with owned_proxy_process(gateway, tmp_path, {}, config=config, workers=2) as owned: candidate: Final = owned.gateway - root: Final = psutil.Process(owned.process.pid) - workers: Final = eventually( - lambda: tuple( - child - for child in root.children(recursive=True) - if any("spawn_main" in part for part in child.cmdline()) - ), - lambda children: len(children) == 2, - seconds=30, - ) + workers: Final = tuple(sorted(eventually(owned.worker_pids, lambda pids: len(pids) == 2, seconds=30))) def run(call: ChaosCall) -> tuple[int, CallerResult | None, str | None]: try: @@ -390,17 +392,12 @@ def test_K3_two_worker_sigkill_checks_post_kill_spend_rows( except (OpenAIAPIConnectionError, AnthropicAPIConnectionError, httpx.RemoteProtocolError) as error: return call.index, None, str(error) - with ThreadPoolExecutor(max_workers=30) as pool: + with owned.replacing_worker(workers[0]) as kill, ThreadPoolExecutor(max_workers=30) as pool: futures: Final = tuple(pool.submit(run, call) for call in calls) try: assert eventually(lambda: scan_started.is_set(), bool, seconds=30) eventually(lambda: edge.received.qsize(), lambda count: count >= 5, seconds=30) - workers[0].send_signal(signal.SIGKILL) - killed_workers, surviving_workers = psutil.wait_procs((workers[0],), timeout=10) - assert len(killed_workers) == 1 and not surviving_workers, ( - killed_workers, - surviving_workers, - ) + kill() finally: release_scans.set() outcomes: Final = tuple(future.result(timeout=90) for future in futures) diff --git a/tests/integration/observability/test_otel_conversation_id.py b/tests/integration/observability/test_otel_conversation_id.py index e77e5603148..6d0e5ea197d 100644 --- a/tests/integration/observability/test_otel_conversation_id.py +++ b/tests/integration/observability/test_otel_conversation_id.py @@ -1,7 +1,6 @@ import asyncio import base64 import json -import os import re import signal import threading @@ -16,12 +15,11 @@ from typing import Final import anthropic import httpx import openai -import psutil import pytest import yaml from integration._support.client import Gateway, eventually, gateway_from_environment from integration._support.database import read_rows -from integration._support.process import OwnedProxy, owned_proxy_process +from integration._support.process import OwnedProxy, owned_proxy_process, stop_root_process from integration._support.wire import Reply, Request, Wire, wire_server from pydantic import JsonValue @@ -762,29 +760,26 @@ def test_slow_sink_during_a_burst_lands_every_response_exactly_once(rig: Rig) -> @pytest.mark.covers("other.observability.otel.conversation_id_survives_worker_kill") def test_killing_one_of_two_workers_mid_burst_keeps_serving_and_never_duplicates_a_span(two_worker_rig: Rig) -> None: rig: Final = two_worker_rig - root: Final = psutil.Process(rig.process.process.pid) - workers: Final = eventually( - lambda: tuple(child for child in root.children() if "resource_tracker" not in " ".join(child.cmdline())), - lambda found: len(found) == 2, - seconds=30, - ) + workers: Final = eventually(rig.process.worker_pids, lambda found: len(found) == 2, seconds=30) sessions: Final = {index: f"conv-{index}-{uuid.uuid4().hex}" for index in range(24)} markers: Final = tuple(_marker() for _ in range(24)) - def one(index: int) -> tuple[str, str | None]: - if index == 8: - os.kill(workers[0].pid, signal.SIGKILL) - try: - response: Final = rig.chat(markers[index], headers={"x-litellm-session-id": sessions[index]}) - return f"chatcmpl-{markers[index]}", None if response.status_code == 200 else response.text - except httpx.HTTPError as error: - return f"chatcmpl-{markers[index]}", repr(error) + with rig.process.replacing_worker(min(workers)) as kill: - with ThreadPoolExecutor(max_workers=6) as pool: - results: Final = tuple(pool.map(one, range(24))) - assert rig.process.process.poll() is None, "Proxy root exited after a worker was killed" - after: Final = rig.chat(_marker(), headers={"x-litellm-session-id": "conv-after-kill"}) - assert after.status_code == 200, after.text + def one(index: int) -> tuple[str, str | None]: + if index == 8: + kill() + try: + response: Final = rig.chat(markers[index], headers={"x-litellm-session-id": sessions[index]}) + return f"chatcmpl-{markers[index]}", None if response.status_code == 200 else response.text + except httpx.HTTPError as error: + return f"chatcmpl-{markers[index]}", repr(error) + + with ThreadPoolExecutor(max_workers=6) as pool: + results: Final = tuple(pool.map(one, range(24))) + assert rig.process.process.poll() is None, "Proxy root exited after a worker was killed" + after: Final = rig.chat(_marker(), headers={"x-litellm-session-id": "conv-after-kill"}) + assert after.status_code == 200, after.text assert rig.sink.single_span(after.json()["id"]) == "conv-after-kill" failures: Final = tuple(error for _, error in results if error) assert all(error.startswith(("ReadError(", "RemoteProtocolError(", "ConnectError(")) for error in failures), ( @@ -823,8 +818,8 @@ def test_terminating_the_proxy_right_after_a_burst_flushes_every_span_before_exi lambda: _landed(rig, expected), lambda seen: all(len(values) >= 1 for values in seen.values()), seconds=60 ) assert drained == {session: (identity,) for identity, session in expected.items()}, drained - rig.process.process.terminate() - assert rig.process.process.wait(timeout=40) in (0, -signal.SIGTERM) + assert stop_root_process(rig.process.process), "Proxy did not exit within the graceful stop budget" + assert rig.process.process.returncode in (0, -signal.SIGTERM) assert _landed(rig, expected) == drained with pytest.raises(httpx.ConnectError): next(started) diff --git a/tests/integration/observability/test_otel_excluded_services.py b/tests/integration/observability/test_otel_excluded_services.py index 1de071e7103..a90d1836d24 100644 --- a/tests/integration/observability/test_otel_excluded_services.py +++ b/tests/integration/observability/test_otel_excluded_services.py @@ -3,6 +3,8 @@ from __future__ import annotations import time import uuid from collections.abc import Callable, Iterator, Mapping +from dataclasses import dataclass, replace +from itertools import accumulate from pathlib import Path from types import MappingProxyType from typing import Final @@ -17,10 +19,12 @@ from integration._support.client import ( ) from integration._support.otlp_sink import ( AuditConfigWriter, + ConnectSink, Span, SpanSinks, audit_configuration_writer, langfuse_callback_variables, + owned_connect_sink, owned_sinks, recorded_spans, spans_for_trace, @@ -29,6 +33,13 @@ from integration._support.process import owned_proxy, owned_proxy_process from pydantic import JsonValue DB_SYSTEM_KEYS: Final = frozenset({"db.system.name", "db.system"}) +TENANT_EXPORT_INTERVAL_MILLISECONDS: Final = 1000 +TAIL_QUIET_SECONDS: Final = 5 * TENANT_EXPORT_INTERVAL_MILLISECONDS / 1000 + + +@pytest.fixture(autouse=True) +def tenant_export_interval(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("OTEL_BSP_SCHEDULE_DELAY", str(TENANT_EXPORT_INTERVAL_MILLISECONDS)) @pytest.fixture @@ -79,19 +90,43 @@ def _operator_langfuse(audit_sinks: SpanSinks) -> dict[str, str]: } -def _add_callback(gateway: Gateway, team_id: str, callback_vars: Mapping[str, JsonValue]) -> httpx.Response: +def _newrelic_environment(audit_sinks: SpanSinks, newrelic_sink: ConnectSink) -> dict[str, str]: + return { + "LITELLM_OTEL_V2": "1", + **_operator_langfuse(audit_sinks), + "HTTPS_PROXY": newrelic_sink.proxy_url, + "NO_PROXY": "127.0.0.1,localhost", + "REQUESTS_CA_BUNDLE": newrelic_sink.ca_pem, + } + + +def _newrelic_vars(excluded_services: str | None = None) -> dict[str, JsonValue]: + variables: Final[dict[str, JsonValue]] = {"newrelic_api_key": "nr-team-excluded-services"} + return variables if excluded_services is None else {**variables, "excluded_services": excluded_services} + + +def _add_callback( + gateway: Gateway, + team_id: str, + callback_vars: Mapping[str, JsonValue], + callback_name: str = "langfuse_otel", +) -> httpx.Response: return gateway.request( "POST", f"/team/{team_id}/callback", - {"callback_name": "langfuse_otel", "callback_vars": dict(callback_vars)}, + {"callback_name": callback_name, "callback_vars": dict(callback_vars)}, ) -def _drive(candidate: Gateway, langfuse_vars: Mapping[str, JsonValue]) -> httpx.Response: +def _drive( + candidate: Gateway, + callback_vars: Mapping[str, JsonValue], + callback_name: str = "langfuse_otel", +) -> httpx.Response: with candidate.scenario() as scenario: model: Final = scenario.model(model="openai/audit-chat", api_base=f"{candidate.upstream_url}/v1") team_id: Final = scenario.team() - callback: Final = _add_callback(candidate, team_id, langfuse_vars) + callback: Final = _add_callback(candidate, team_id, callback_vars, callback_name) assert callback.status_code == 200, callback.text key: Final = scenario.key(team_id=team_id) response: Final = candidate.request( @@ -125,21 +160,33 @@ def _trace_id(sink_url: str, response: httpx.Response, seconds: float = 40) -> s return found -def _trace_spans(sink_url: str, trace_id: str, seconds: float = 30) -> tuple[Span, ...]: - """The trace's spans once the post-call tail has landed. +@dataclass(frozen=True, slots=True) +class _TraceSample: + at: float + changed_at: float + spans: tuple[Span, ...] - The spend-writer and other post-response spans flush after the request - answers, so absence assertions poll for the whole window instead of - settling at the first glimpse of the root span. - """ - deadline: Final = time.monotonic() + seconds - group: tuple[Span, ...] = () # rebind-ok: drains samples until the post-call tail lands - while time.monotonic() < deadline: - _, spans = recorded_spans(sink_url) - group = spans_for_trace(spans, trace_id) + +def _trace_samples(sink_url: str, trace_id: str) -> Iterator[_TraceSample]: + while True: + now: Final = time.monotonic() + yield _TraceSample(now, now, spans_for_trace(recorded_spans(sink_url)[1], trace_id)) time.sleep(0.5) - assert group, f"trace {trace_id} never reached {sink_url}" - return group + + +def _keep_unchanged_since(previous: _TraceSample, current: _TraceSample) -> _TraceSample: + return current if len(current.spans) != len(previous.spans) else replace(current, changed_at=previous.changed_at) + + +def _trace_spans(sink_url: str, trace_id: str, seconds: float = 30) -> tuple[Span, ...]: + deadline: Final = time.monotonic() + seconds + settled: Final = next( + sample + for sample in accumulate(_trace_samples(sink_url, trace_id), _keep_unchanged_since) + if (sample.spans and sample.at - sample.changed_at >= TAIL_QUIET_SECONDS) or sample.at >= deadline + ) + assert settled.spans, f"trace {trace_id} never reached {sink_url}" + return settled.spans def _trace_spans_when( @@ -231,7 +278,7 @@ def test_excluded_services_drops_db_spans_at_tenant_only( config: Final = _config_with( tmp_path, otel_audit_config, otel={"excluded_services": ["redis", "postgres"]}, extra=_guardrail_block ) - with owned_proxy(gateway, tmp_path, {"LITELLM_OTEL_V2": "1"}, config=config, workers=2) as candidate: + with owned_proxy(gateway, tmp_path, {"LITELLM_OTEL_V2": "1"}, config=config) as candidate: ten_start, _ = recorded_spans(audit_sinks.tenant) op_start, _ = recorded_spans(audit_sinks.operator) traffic: Final = _drive(candidate, langfuse_vars) @@ -266,7 +313,7 @@ def test_without_excluded_services_the_tenant_still_gets_redis_and_postgres_span tmp_path: Path, ) -> None: config: Final = _config_with(tmp_path, otel_audit_config, extra=_guardrail_block) - with owned_proxy(gateway, tmp_path, {"LITELLM_OTEL_V2": "1"}, config=config, workers=2) as candidate: + with owned_proxy(gateway, tmp_path, {"LITELLM_OTEL_V2": "1"}, config=config) as candidate: tenant_start, _ = recorded_spans(audit_sinks.tenant) traffic: Final = _drive(candidate, langfuse_vars) tenant_trace: Final = _trace_id(audit_sinks.tenant, traffic) @@ -278,6 +325,85 @@ def test_without_excluded_services_the_tenant_still_gets_redis_and_postgres_span assert {"redis", "postgresql"} <= systems, f"datastore spans missing at tenant: {systems}" +@pytest.mark.timeout(180) +def test_team_excluded_services_replaces_global_and_keeps_operator_db_spans( + gateway: Gateway, + audit_sinks: SpanSinks, + otel_audit_config: AuditConfigWriter, + tmp_path: Path, +) -> None: + config: Final = _config_with(tmp_path, otel_audit_config, extra=_guardrail_block) + with owned_connect_sink(tmp_path / "newrelic") as newrelic_sink: + overrides: Final = _newrelic_environment(audit_sinks, newrelic_sink) + with owned_proxy(gateway, tmp_path, overrides, config=config, workers=2) as candidate: + operator_start, _ = recorded_spans(audit_sinks.operator) + tenant_start, _ = recorded_spans(newrelic_sink.control_url) + traffic: Final = _drive(candidate, _newrelic_vars("redis,postgres"), callback_name="newrelic") + _await_db_span( + audit_sinks.operator, + None, + "postgres.update LiteLLM_VerificationToken", + seconds=60, + since=operator_start, + ) + operator_trace: Final = _trace_id(audit_sinks.operator, traffic) + tenant_trace: Final = _trace_id(newrelic_sink.control_url, traffic) + assert operator_trace == tenant_trace + operator_trace_spans: Final = _trace_spans(audit_sinks.operator, operator_trace, seconds=15) + assert "redis" in _db_systems(operator_trace_spans), _db_systems(operator_trace_spans) + _, operator_spans = recorded_spans(audit_sinks.operator, operator_start) + assert "postgresql" in _db_systems(operator_spans), _db_systems(operator_spans) + tenant_trace_spans: Final = _trace_spans(newrelic_sink.control_url, tenant_trace, seconds=15) + _, tenant_spans = recorded_spans(newrelic_sink.control_url, tenant_start) + assert _db_systems(tenant_trace_spans) == set(), _db_systems(tenant_trace_spans) + assert _db_systems(tenant_spans) == set(), _db_systems(tenant_spans) + + +@pytest.mark.timeout(180) +def test_team_excluded_services_none_overrides_global( + gateway: Gateway, + audit_sinks: SpanSinks, + otel_audit_config: AuditConfigWriter, + tmp_path: Path, +) -> None: + config: Final = _config_with( + tmp_path, otel_audit_config, otel={"excluded_services": ["redis", "postgres"]}, extra=_guardrail_block + ) + with owned_connect_sink(tmp_path / "newrelic") as newrelic_sink: + overrides: Final = _newrelic_environment(audit_sinks, newrelic_sink) + with owned_proxy(gateway, tmp_path, overrides, config=config, workers=2) as candidate: + tenant_start, _ = recorded_spans(newrelic_sink.control_url) + traffic: Final = _drive(candidate, _newrelic_vars("none"), callback_name="newrelic") + tenant_trace: Final = _trace_id(newrelic_sink.control_url, traffic) + _await_db_span(newrelic_sink.control_url, tenant_trace, "redis", seconds=60) + _await_db_span(newrelic_sink.control_url, None, "postgresql", seconds=60, since=tenant_start) + tenant_trace_spans: Final = _trace_spans(newrelic_sink.control_url, tenant_trace, seconds=15) + _, tenant_spans = recorded_spans(newrelic_sink.control_url, tenant_start) + assert {"redis", "postgresql"} <= _db_systems(tenant_spans), _db_systems(tenant_spans) + assert {"redis", "postgresql"} <= _db_systems(tenant_trace_spans), _db_systems(tenant_trace_spans) + + +@pytest.mark.timeout(180) +def test_team_excluded_services_unset_inherits_global( + gateway: Gateway, + audit_sinks: SpanSinks, + otel_audit_config: AuditConfigWriter, + tmp_path: Path, +) -> None: + config: Final = _config_with( + tmp_path, otel_audit_config, otel={"excluded_services": ["redis", "postgres"]}, extra=_guardrail_block + ) + with owned_connect_sink(tmp_path / "newrelic") as newrelic_sink: + overrides: Final = _newrelic_environment(audit_sinks, newrelic_sink) + with owned_proxy(gateway, tmp_path, overrides, config=config, workers=2) as candidate: + tenant_start, _ = recorded_spans(newrelic_sink.control_url) + traffic: Final = _drive(candidate, _newrelic_vars(), callback_name="newrelic") + tenant_trace: Final = _trace_id(newrelic_sink.control_url, traffic) + _trace_spans(newrelic_sink.control_url, tenant_trace, seconds=15) + _, tenant_spans = recorded_spans(newrelic_sink.control_url, tenant_start) + assert _db_systems(tenant_spans) == set(), _db_systems(tenant_spans) + + @pytest.mark.parametrize("otel", [None, True, "on", "", []], ids=["null", "true", "on", "empty_string", "empty_list"]) @pytest.mark.timeout(180) def test_a_non_mapping_otel_block_still_publishes_the_tenant_fan_out( @@ -294,7 +420,7 @@ def test_a_non_mapping_otel_block_still_publishes_the_tenant_fan_out( config: Final = _config_with(tmp_path, otel_audit_config, extra=with_callback_settings) overrides: Final = {"LITELLM_OTEL_V2": "1", **_operator_langfuse(audit_sinks)} - with owned_proxy(gateway, tmp_path, overrides, config=config, workers=2) as candidate: + with owned_proxy(gateway, tmp_path, overrides, config=config) as candidate: traffic: Final = _drive(candidate, langfuse_vars) tenant_trace: Final = _trace_id(audit_sinks.tenant, traffic) _await_db_span(audit_sinks.tenant, tenant_trace, "redis") @@ -328,7 +454,7 @@ def test_a_bare_excluded_services_env_var_is_ignored( ) -> None: config: Final = _config_with(tmp_path, otel_audit_config) overrides: Final = {"LITELLM_OTEL_V2": "1", name: "redis,postgres"} - with owned_proxy(gateway, tmp_path, overrides, config=config, workers=2) as candidate: + with owned_proxy(gateway, tmp_path, overrides, config=config) as candidate: tenant_start, _ = recorded_spans(audit_sinks.tenant) traffic: Final = _drive(candidate, langfuse_vars) tenant_trace: Final = _trace_id(audit_sinks.tenant, traffic) @@ -360,7 +486,7 @@ def test_the_documented_env_var_wins_over_a_bare_excluded_services( "LITELLM_OTEL_EXCLUDED_SERVICES": "redis", "EXCLUDED_SERVICES": "postgres", } - with owned_proxy(gateway, tmp_path, overrides, config=config, workers=2) as candidate: + with owned_proxy(gateway, tmp_path, overrides, config=config) as candidate: tenant_start, _ = recorded_spans(audit_sinks.tenant) traffic: Final = _drive(candidate, langfuse_vars) tenant_trace: Final = _trace_id(audit_sinks.tenant, traffic) @@ -397,7 +523,6 @@ def test_case_sensitive_otel_settings_read_only_the_exact_env_name( overrides, config=config, remove_environment=("LITELLM_OTEL_EXCLUDED_SERVICES", "litellm_otel_excluded_services"), - workers=2, ) as candidate: tenant_start, _ = recorded_spans(audit_sinks.tenant) traffic: Final = _drive(candidate, langfuse_vars) @@ -431,7 +556,7 @@ def test_env_ignore_empty_keeps_the_default_service_name( ) -> None: config: Final = _config_with(tmp_path, otel_audit_config, otel={"_env_ignore_empty": True}) overrides: Final = {"LITELLM_OTEL_V2": "1", "OTEL_SERVICE_NAME": ""} - with owned_proxy(gateway, tmp_path, overrides, config=config, workers=2) as candidate: + with owned_proxy(gateway, tmp_path, overrides, config=config) as candidate: traffic: Final = _drive(candidate, langfuse_vars) operator_trace: Final = _trace_id(audit_sinks.operator, traffic) operator_spans: Final = _trace_spans_when( @@ -456,7 +581,7 @@ def test_env_parse_none_str_reads_a_null_traces_endpoint_as_unset( ) -> None: config: Final = _config_with(tmp_path, otel_audit_config, otel={"_env_parse_none_str": "null"}) overrides: Final = {"LITELLM_OTEL_V2": "1", "OTEL_TRACES_ENDPOINT": "null"} - with owned_proxy(gateway, tmp_path, overrides, config=config, workers=2) as candidate: + with owned_proxy(gateway, tmp_path, overrides, config=config) as candidate: traffic: Final = _drive(candidate, langfuse_vars) operator_trace: Final = _trace_id(audit_sinks.operator, traffic) operator_spans: Final = _trace_spans_when( @@ -477,7 +602,7 @@ def test_env_excluded_services_drops_only_redis( ) -> None: config: Final = _config_with(tmp_path, otel_audit_config) with owned_proxy( - gateway, tmp_path, {"LITELLM_OTEL_V2": "1", "LITELLM_OTEL_EXCLUDED_SERVICES": "redis"}, config=config, workers=2 + gateway, tmp_path, {"LITELLM_OTEL_V2": "1", "LITELLM_OTEL_EXCLUDED_SERVICES": "redis"}, config=config ) as candidate: start, _ = recorded_spans(audit_sinks.tenant) _drive(candidate, langfuse_vars) @@ -503,7 +628,7 @@ def test_config_excluded_services_wins_over_env( tmp_path, otel_audit_config, otel={"excluded_services": ["postgres"]}, extra=with_langfuse_otel ) with owned_proxy( - gateway, tmp_path, {"LITELLM_OTEL_V2": "1", "LITELLM_OTEL_EXCLUDED_SERVICES": "redis"}, config=config, workers=2 + gateway, tmp_path, {"LITELLM_OTEL_V2": "1", "LITELLM_OTEL_EXCLUDED_SERVICES": "redis"}, config=config ) as candidate: _assert_tenant_keeps_redis_without_postgres(candidate, audit_sinks, langfuse_vars) @@ -523,7 +648,7 @@ def test_excluded_services_applies_with_preset_ordered_first( tmp_path, otel_audit_config, otel={"excluded_services": ["postgres"]}, extra=preset_first ) overrides: Final = {"LITELLM_OTEL_V2": "1", **_operator_langfuse(audit_sinks)} - with owned_proxy(gateway, tmp_path, overrides, config=config, workers=2) as candidate: + with owned_proxy(gateway, tmp_path, overrides, config=config) as candidate: _assert_tenant_keeps_redis_without_postgres(candidate, audit_sinks, langfuse_vars) @@ -536,7 +661,7 @@ def test_bogus_excluded_service_logs_error_and_drops_at_proxy_start( tmp_path: Path, ) -> None: config: Final = _config_with(tmp_path, otel_audit_config, otel={"excluded_services": ["auth", "postgres"]}) - with owned_proxy_process(gateway, tmp_path, {"LITELLM_OTEL_V2": "1"}, config=config, workers=2) as owned: + with owned_proxy_process(gateway, tmp_path, {"LITELLM_OTEL_V2": "1"}, config=config) as owned: assert "'auth' is not a datastore service; ignored" in owned.log.read_text(), owned.log.read_text()[-3000:] _assert_tenant_keeps_redis_without_postgres(owned.gateway, audit_sinks, langfuse_vars) @@ -551,7 +676,7 @@ def test_valid_config_excluded_services_tolerates_bogus_env( ) -> None: config: Final = _config_with(tmp_path, otel_audit_config, otel={"excluded_services": ["postgres"]}) with owned_proxy( - gateway, tmp_path, {"LITELLM_OTEL_V2": "1", "LITELLM_OTEL_EXCLUDED_SERVICES": "auth"}, config=config, workers=2 + gateway, tmp_path, {"LITELLM_OTEL_V2": "1", "LITELLM_OTEL_EXCLUDED_SERVICES": "auth"}, config=config ) as candidate: _assert_tenant_keeps_redis_without_postgres(candidate, audit_sinks, langfuse_vars) @@ -569,7 +694,7 @@ def test_bogus_excluded_services_env_logs_and_drops_with_preset_alongside_otel( config: Final = _config_with(tmp_path, otel_audit_config, extra=with_langfuse_otel) overrides: Final = {"LITELLM_OTEL_V2": "1", "LITELLM_OTEL_EXCLUDED_SERVICES": "auth,postgres"} - with owned_proxy_process(gateway, tmp_path, overrides, config=config, workers=2) as owned: + with owned_proxy_process(gateway, tmp_path, overrides, config=config) as owned: assert "'auth' is not a datastore service; ignored" in owned.log.read_text(), owned.log.read_text()[-3000:] _assert_tenant_keeps_redis_without_postgres(owned.gateway, audit_sinks, langfuse_vars) @@ -591,7 +716,7 @@ def test_bogus_excluded_services_env_logs_and_drops_without_otel_callback( "LITELLM_OTEL_EXCLUDED_SERVICES": "auth,postgres", **_operator_langfuse(audit_sinks), } - with owned_proxy_process(gateway, tmp_path, overrides, config=config, workers=2) as owned: + with owned_proxy_process(gateway, tmp_path, overrides, config=config) as owned: assert "'auth' is not a datastore service; ignored" in owned.log.read_text(), owned.log.read_text()[-3000:] _assert_tenant_keeps_redis_without_postgres(owned.gateway, audit_sinks, langfuse_vars) @@ -604,7 +729,7 @@ def test_postgres_exclusion_covers_spend_flush( tmp_path: Path, ) -> None: config: Final = _config_with(tmp_path, otel_audit_config, otel={"excluded_services": ["postgres"]}) - with owned_proxy(gateway, tmp_path, {"LITELLM_OTEL_V2": "1"}, config=config, workers=2) as candidate: + with owned_proxy(gateway, tmp_path, {"LITELLM_OTEL_V2": "1"}, config=config) as candidate: op_start, _ = recorded_spans(audit_sinks.operator) ten_start, _ = recorded_spans(audit_sinks.tenant) traffic: Final = _drive(candidate, langfuse_vars) diff --git a/tests/integration/observability/test_otel_excluded_services_matrix.py b/tests/integration/observability/test_otel_excluded_services_matrix.py index 7bca7ffc0dc..07646a0b189 100644 --- a/tests/integration/observability/test_otel_excluded_services_matrix.py +++ b/tests/integration/observability/test_otel_excluded_services_matrix.py @@ -1,8 +1,6 @@ import asyncio import json -import os import re -import signal import uuid from collections.abc import Callable, Generator, Iterator, Mapping from concurrent.futures import ThreadPoolExecutor @@ -15,7 +13,6 @@ from typing import Final, Literal, Protocol, cast import anthropic import httpx import openai -import psutil import pytest import yaml from integration._support.client import Gateway, Scenario, eventually, gateway_from_environment, object_value @@ -687,25 +684,22 @@ def test_stalled_tenant_destination_during_a_burst_does_not_block_responses(rig: @pytest.mark.timeout(300) def test_killing_one_of_two_workers_mid_burst_keeps_the_filter_on_the_survivor(rig: Rig) -> None: - root: Final = psutil.Process(rig.owned.process.pid) - workers: Final = eventually( - lambda: tuple(child for child in root.children() if "resource_tracker" not in " ".join(child.cmdline())), - lambda found: len(found) == 2, - seconds=30, - ) + workers: Final = eventually(rig.owned.worker_pids, lambda found: len(found) == 2, seconds=30) cursors: Final = rig.cursors() - def one(index: int) -> Sent | str: - if index == 6: - os.kill(workers[0].pid, signal.SIGKILL) - try: - return rig.raw("chat", _marker(), stream=index % 2 == 0) - except (httpx.HTTPError, AssertionError) as error: - return repr(error) + with rig.owned.replacing_worker(min(workers)) as kill: - with ThreadPoolExecutor(max_workers=6) as pool: - results: Final = tuple(pool.map(one, range(18))) - assert rig.owned.process.poll() is None, "Proxy root exited after a worker was killed" + def one(index: int) -> Sent | str: + if index == 6: + kill() + try: + return rig.raw("chat", _marker(), stream=index % 2 == 0) + except (httpx.HTTPError, AssertionError) as error: + return repr(error) + + with ThreadPoolExecutor(max_workers=6) as pool: + results: Final = tuple(pool.map(one, range(18))) + assert rig.owned.process.poll() is None, "Proxy root exited after a worker was killed" failures: Final = tuple(result for result in results if isinstance(result, str)) assert all(failure.startswith(("ReadError(", "RemoteProtocolError(", "ConnectError(")) for failure in failures), ( failures @@ -918,25 +912,22 @@ def test_unconfigured_tenant_outage_during_a_mixed_burst(unconfigured_rig: Rig) @pytest.mark.timeout(300) def test_unconfigured_killing_one_of_two_workers_keeps_the_fan_out(unconfigured_rig: Rig) -> None: - root: Final = psutil.Process(unconfigured_rig.owned.process.pid) - workers: Final = eventually( - lambda: tuple(child for child in root.children() if "resource_tracker" not in " ".join(child.cmdline())), - lambda found: len(found) == 2, - seconds=30, - ) + workers: Final = eventually(unconfigured_rig.owned.worker_pids, lambda found: len(found) == 2, seconds=30) cursors: Final = unconfigured_rig.cursors() - def one(index: int) -> Sent | str: - if index == 6: - os.kill(workers[0].pid, signal.SIGKILL) - try: - return unconfigured_rig.raw("chat", _marker(), stream=index % 2 == 0) - except (httpx.HTTPError, AssertionError) as error: - return repr(error) + with unconfigured_rig.owned.replacing_worker(min(workers)) as kill: - with ThreadPoolExecutor(max_workers=6) as pool: - results: Final = tuple(pool.map(one, range(18))) - assert unconfigured_rig.owned.process.poll() is None, "Proxy root exited after a worker was killed" + def one(index: int) -> Sent | str: + if index == 6: + kill() + try: + return unconfigured_rig.raw("chat", _marker(), stream=index % 2 == 0) + except (httpx.HTTPError, AssertionError) as error: + return repr(error) + + with ThreadPoolExecutor(max_workers=6) as pool: + results: Final = tuple(pool.map(one, range(18))) + assert unconfigured_rig.owned.process.poll() is None, "Proxy root exited after a worker was killed" failures: Final = tuple(result for result in results if isinstance(result, str)) assert all(failure.startswith(("ReadError(", "RemoteProtocolError(", "ConnectError(")) for failure in failures), ( failures diff --git a/tests/integration/observability/test_passthrough_upstream_error_chaos.py b/tests/integration/observability/test_passthrough_upstream_error_chaos.py index 014942ccfaf..b4a8524d3f8 100644 --- a/tests/integration/observability/test_passthrough_upstream_error_chaos.py +++ b/tests/integration/observability/test_passthrough_upstream_error_chaos.py @@ -1,10 +1,8 @@ import asyncio import json import re -import signal import threading -from collections.abc import Callable, Iterator -from contextlib import contextmanager +from collections.abc import Callable from dataclasses import dataclass from pathlib import Path from typing import Final @@ -15,7 +13,7 @@ import pytest import yaml from integration._support.client import Gateway, eventually, object_value from integration._support.database import read_rows -from integration._support.process import graceful_stop_seconds, owned_proxy_process +from integration._support.process import WorkerHealthcheck, graceful_stop_seconds, owned_proxy_process from integration._support.wire import Reply, Request, wire_server from pydantic import JsonValue @@ -60,16 +58,6 @@ def _held_connections(pid: int, port: int) -> int: ) -@contextmanager -def _suspended_worker(pid: int) -> Iterator[None]: - worker: Final = psutil.Process(pid) - worker.suspend() - try: - yield - finally: - worker.resume() - - def _error_information(call_id: str) -> dict[str, JsonValue]: rows: Final = eventually( lambda: read_rows('SELECT metadata FROM "LiteLLM_SpendLogs" WHERE request_id=%s', (call_id,)), @@ -166,21 +154,19 @@ async def test_passthrough_worker_sigkill_leaves_sibling_serving_and_logging(gat with wire_server(_held_reply(release)) as wire: config["environment_variables"] = {"GEMINI_API_BASE": wire.url, "GEMINI_API_KEY": "scripted"} path.write_text(yaml.safe_dump(config)) - with owned_proxy_process(gateway, tmp_path, {}, config=path, workers=2) as owned: + with owned_proxy_process( + gateway, tmp_path, {}, config=path, workers=2, healthcheck=WorkerHealthcheck.LENIENT + ) as owned: candidate: Final = owned.gateway - workers: Final = eventually( - lambda: tuple(int(pid) for pid in _STARTED_WORKER.findall(owned.log.read_text())), - lambda pids: len(pids) == 2, - seconds=30, - ) + workers: Final = tuple(sorted(eventually(owned.worker_pids, lambda pids: len(pids) == 2, seconds=30))) async with asyncio.TaskGroup() as tasks: try: - with _suspended_worker(workers[0]): + with owned.suspended(workers[0]): seed: Final = tasks.create_task( _fire_burst(str(candidate.client.base_url), candidate.key, 1, tolerate_transport_errors=True) ) await asyncio.to_thread(eventually, lambda: wire.received.qsize(), lambda size: size >= 1, 30) - with _suspended_worker(workers[1]): + with owned.suspended(workers[1]): rest: Final = tasks.create_task( _fire_burst( str(candidate.client.base_url), @@ -194,16 +180,18 @@ async def test_passthrough_worker_sigkill_leaves_sibling_serving_and_logging(gat peer_port: Final = int(wire.url.rsplit(":", 1)[1]) held_by: Final = {pid: _held_connections(pid, peer_port) for pid in workers} assert sum(held_by.values()) == 20 and min(held_by.values()) > 0, held_by - victim: Final = psutil.Process(max(workers, key=held_by.__getitem__)) - victim.suspend() - victim_ports: Final = frozenset( - connection.raddr.port for connection in victim.net_connections(kind="tcp") if connection.raddr - ) - victim.send_signal(signal.SIGKILL) - victim.wait(timeout=5) + with owned.replacing_worker(max(workers, key=held_by.__getitem__)) as kill: + victim: Final = kill.freeze() + victim_ports: Final = frozenset( + connection.raddr.port + for connection in victim.net_connections(kind="tcp") + if connection.raddr + ) + kill() + release.set() + served: Final = (*await seed, *await rest) finally: release.set() - served: Final = (*await seed, *await rest) for item in served: assert item.response.status_code in (200, 404, 500, 502), item.response.status_code follow_up: Final = candidate.request( diff --git a/tests/integration/observability/test_passthrough_upstream_error_visibility.py b/tests/integration/observability/test_passthrough_upstream_error_visibility.py index bc6747844ba..0d86440045b 100644 --- a/tests/integration/observability/test_passthrough_upstream_error_visibility.py +++ b/tests/integration/observability/test_passthrough_upstream_error_visibility.py @@ -126,7 +126,7 @@ async def test_gemini_passthrough_async_client_404_body_reaches_proxy_log_and_sp path: Final = tmp_path / "gemini-async.yaml" with wire_server(respond) as wire: _gemini_config(path, wire.url) - with owned_proxy_process(gateway, tmp_path, {}, config=path, workers=2) as owned: + with owned_proxy_process(gateway, tmp_path, {}, config=path) as owned: candidate: Final = owned.gateway async with httpx.AsyncClient( base_url=str(candidate.client.base_url), timeout=15, trust_env=False @@ -156,7 +156,7 @@ def test_gemini_passthrough_streaming_500_relays_full_body_and_logs_bounded_prev path: Final = tmp_path / "gemini-stream-500.yaml" with wire_server(respond) as wire: _gemini_config(path, wire.url) - with owned_proxy_process(gateway, tmp_path, {}, config=path, workers=2) as owned: + with owned_proxy_process(gateway, tmp_path, {}, config=path) as owned: candidate: Final = owned.gateway with candidate.client.stream( "POST", @@ -188,7 +188,7 @@ def test_gemini_passthrough_success_logs_nothing_and_spend_row_is_success(gatewa path: Final = tmp_path / "gemini-200.yaml" with wire_server(respond) as wire: _gemini_config(path, wire.url) - with owned_proxy_process(gateway, tmp_path, {}, config=path, workers=2) as owned: + with owned_proxy_process(gateway, tmp_path, {}, config=path) as owned: candidate: Final = owned.gateway response: Final = candidate.request( "POST", _GEMINI_MODEL_PATH, _GENERATE_CONTENT, headers={"x-goog-api-key": candidate.key} @@ -208,7 +208,7 @@ def test_gemini_passthrough_streaming_200_relays_every_chunk(gateway: Gateway, t path: Final = tmp_path / "gemini-stream-200.yaml" with wire_server(respond) as wire: _gemini_config(path, wire.url) - with owned_proxy_process(gateway, tmp_path, {}, config=path, workers=2) as owned: + with owned_proxy_process(gateway, tmp_path, {}, config=path) as owned: candidate: Final = owned.gateway with candidate.client.stream( "POST", @@ -239,7 +239,7 @@ def test_gemini_streaming_passthrough_request_tags_reach_spend_log(gateway: Gate config_path: Final = tmp_path / "gemini-stream-tags.yaml" with wire_server(respond) as wire: _gemini_config(config_path, wire.url) - with owned_proxy_process(gateway, tmp_path, {}, config=config_path, workers=2) as owned: + with owned_proxy_process(gateway, tmp_path, {}, config=config_path) as owned: candidate: Final = owned.gateway with candidate.client.stream( "POST", @@ -293,7 +293,7 @@ def test_config_pass_through_route_logs_body_and_strips_query(gateway: Gateway, } ] path.write_text(yaml.safe_dump(config)) - with owned_proxy_process(gateway, tmp_path, {}, config=path, workers=2) as owned: + with owned_proxy_process(gateway, tmp_path, {}, config=path) as owned: candidate: Final = owned.gateway response: Final = candidate.request("POST", "/audit-pt", _GENERATE_CONTENT) assert response.status_code == 403, response.text @@ -324,7 +324,7 @@ def test_openai_passthrough_sdk_error_body_reaches_proxy_log_and_spend_row(gatew path: Final = tmp_path / "openai-404.yaml" with wire_server(respond) as wire: _openai_config(path, wire.url) - with owned_proxy_process(gateway, tmp_path, {}, config=path, workers=2) as owned: + with owned_proxy_process(gateway, tmp_path, {}, config=path) as owned: candidate: Final = owned.gateway with OpenAI( api_key=candidate.key, @@ -350,7 +350,7 @@ async def test_openai_passthrough_async_sdk_error_body_reaches_proxy_log_and_spe path: Final = tmp_path / "openai-async-404.yaml" with wire_server(respond) as wire: _openai_config(path, wire.url) - with owned_proxy_process(gateway, tmp_path, {}, config=path, workers=2) as owned: + with owned_proxy_process(gateway, tmp_path, {}, config=path) as owned: candidate: Final = owned.gateway async with AsyncOpenAI( api_key=candidate.key, @@ -377,7 +377,7 @@ def test_gemini_passthrough_control_characters_cannot_forge_log_lines(gateway: G path: Final = tmp_path / "gemini-forged.yaml" with wire_server(respond) as wire: _gemini_config(path, wire.url) - with owned_proxy_process(gateway, tmp_path, {}, config=path, workers=2) as owned: + with owned_proxy_process(gateway, tmp_path, {}, config=path) as owned: candidate: Final = owned.gateway response: Final = candidate.request( "POST", _GEMINI_MODEL_PATH, _GENERATE_CONTENT, headers={"x-goog-api-key": candidate.key} @@ -401,7 +401,7 @@ def test_gemini_passthrough_empty_error_body_still_logged_and_proxy_serves(gatew path: Final = tmp_path / "gemini-empty.yaml" with wire_server(respond) as wire: _gemini_config(path, wire.url) - with owned_proxy_process(gateway, tmp_path, {}, config=path, workers=2) as owned: + with owned_proxy_process(gateway, tmp_path, {}, config=path) as owned: candidate: Final = owned.gateway response: Final = candidate.request( "POST", _GEMINI_MODEL_PATH, _GENERATE_CONTENT, headers={"x-goog-api-key": candidate.key} @@ -434,7 +434,7 @@ def test_gemini_passthrough_gzip_error_body_decoded_for_log_and_client(gateway: path: Final = tmp_path / "gemini-gzip.yaml" with wire_server(respond) as wire: _gemini_config(path, wire.url) - with owned_proxy_process(gateway, tmp_path, {}, config=path, workers=2) as owned: + with owned_proxy_process(gateway, tmp_path, {}, config=path) as owned: candidate: Final = owned.gateway response: Final = candidate.request( "POST", _GEMINI_MODEL_PATH, _GENERATE_CONTENT, headers={"x-goog-api-key": candidate.key} @@ -462,7 +462,7 @@ def test_gemini_passthrough_streaming_gzip_error_body_decoded_for_log_and_client path: Final = tmp_path / "gemini-stream-gzip.yaml" with wire_server(respond) as wire: _gemini_config(path, wire.url) - with owned_proxy_process(gateway, tmp_path, {}, config=path, workers=2) as owned: + with owned_proxy_process(gateway, tmp_path, {}, config=path) as owned: candidate: Final = owned.gateway with candidate.client.stream( "POST", @@ -490,7 +490,7 @@ def test_gemini_passthrough_error_body_redacted_when_message_logging_off(gateway config["environment_variables"] = {"GEMINI_API_BASE": wire.url, "GEMINI_API_KEY": "scripted"} config["litellm_settings"]["turn_off_message_logging"] = True path.write_text(yaml.safe_dump(config)) - with owned_proxy_process(gateway, tmp_path, {}, config=path, workers=2) as owned: + with owned_proxy_process(gateway, tmp_path, {}, config=path) as owned: candidate: Final = owned.gateway response: Final = candidate.request( "POST", _GEMINI_MODEL_PATH, _GENERATE_CONTENT, headers={"x-goog-api-key": candidate.key} @@ -516,7 +516,7 @@ def test_gemini_passthrough_exact_4096_byte_body_logged_without_marker(gateway: path: Final = tmp_path / "gemini-exact.yaml" with wire_server(respond) as wire: _gemini_config(path, wire.url) - with owned_proxy_process(gateway, tmp_path, {}, config=path, workers=2) as owned: + with owned_proxy_process(gateway, tmp_path, {}, config=path) as owned: candidate: Final = owned.gateway response: Final = candidate.request( "POST", _GEMINI_MODEL_PATH, _GENERATE_CONTENT, headers={"x-goog-api-key": candidate.key} @@ -537,7 +537,7 @@ def test_gemini_passthrough_4097_byte_body_truncated_with_marker(gateway: Gatewa path: Final = tmp_path / "gemini-over.yaml" with wire_server(respond) as wire: _gemini_config(path, wire.url) - with owned_proxy_process(gateway, tmp_path, {}, config=path, workers=2) as owned: + with owned_proxy_process(gateway, tmp_path, {}, config=path) as owned: candidate: Final = owned.gateway response: Final = candidate.request( "POST", _GEMINI_MODEL_PATH, _GENERATE_CONTENT, headers={"x-goog-api-key": candidate.key} @@ -557,7 +557,7 @@ def test_gemini_passthrough_one_byte_stream_chunks_reassembled_and_logged(gatewa path: Final = tmp_path / "gemini-one-byte.yaml" with wire_server(respond) as wire: _gemini_config(path, wire.url) - with owned_proxy_process(gateway, tmp_path, {}, config=path, workers=2) as owned: + with owned_proxy_process(gateway, tmp_path, {}, config=path) as owned: candidate: Final = owned.gateway with candidate.client.stream( "POST", @@ -580,7 +580,7 @@ def test_gemini_passthrough_repeated_errors_each_get_row_and_log_line(gateway: G path: Final = tmp_path / "gemini-twice.yaml" with wire_server(respond) as wire: _gemini_config(path, wire.url) - with owned_proxy_process(gateway, tmp_path, {}, config=path, workers=2) as owned: + with owned_proxy_process(gateway, tmp_path, {}, config=path) as owned: candidate: Final = owned.gateway responses: Final = tuple( candidate.request( @@ -604,7 +604,7 @@ def test_gemini_passthrough_repeated_errors_each_get_row_and_log_line(gateway: G def test_budget_rejected_call_keeps_budget_normalized_error(gateway: Gateway, tmp_path: Path) -> None: path: Final = tmp_path / "budget.yaml" path.write_text(Path("tests/integration/proxy_config.yaml").read_text()) - with owned_proxy_process(gateway, tmp_path, {}, config=path, workers=2) as owned: + with owned_proxy_process(gateway, tmp_path, {}, config=path) as owned: candidate: Final = owned.gateway with candidate.scenario() as scenario: model: Final = scenario.model() diff --git a/tests/integration/observability/test_presidio_streaming_output.py b/tests/integration/observability/test_presidio_streaming_output.py index 680617f1652..d784fecd146 100644 --- a/tests/integration/observability/test_presidio_streaming_output.py +++ b/tests/integration/observability/test_presidio_streaming_output.py @@ -1,6 +1,5 @@ import json import re -import signal import threading import uuid from collections.abc import Callable, Iterator, Mapping @@ -10,16 +9,17 @@ from dataclasses import dataclass from pathlib import Path from typing import Final -import psutil +import httpx import yaml from integration._support.client import Gateway, eventually -from integration._support.process import OwnedProxy, group_members, owned_proxy_process +from integration._support.process import OwnedProxy, owned_proxy_process from integration._support.wire import Reply, Request, Wire, wire_server from openai import OpenAI from pydantic import BaseModel PERSON: Final = "John Smith" MASK: Final = "" +CUT_BY_WORKER_DEATH: Final = 0 GEMINI_MODEL: Final = "gemini-2.5-flash" @@ -170,6 +170,12 @@ class Rig: ) as response: return Received(response.status_code, tuple(response.iter_raw())) + def stream_on_a_worker_that_may_die(self, path: str, body: dict[str, object]) -> Received: + try: + return self.stream(path, body) + except httpx.TransportError: + return Received(CUT_BY_WORKER_DEATH, ()) + def gemini_path(self) -> str: return f"/v1beta/models/{self.gemini}:streamGenerateContent?alt=sse" @@ -567,18 +573,13 @@ def test_native_gemini_keeps_streaming_after_one_worker_is_killed(gateway: Gatew frames: Final = (gemini_frame("alive "), gemini_frame("still.")) provider: Final = gemini_provider(Reply(content_type="text/event-stream", chunks=frames, pause_between_chunks=0.5)) with presidio_rig(gateway, tmp_path, provider) as rig: - workers: Final = eventually( - lambda: tuple( - member for member in group_members(rig.proxy.process.pid) if member.pid != rig.proxy.process.pid - ), - lambda members: len(members) >= 2, - seconds=30, - ) - victim: Final = workers[0] - with ThreadPoolExecutor(max_workers=8) as pool: - futures: Final = tuple(pool.submit(rig.stream, rig.gemini_path(), rig.gemini_body()) for _ in range(8)) - victim.send_signal(signal.SIGKILL) - psutil.wait_procs((victim,), timeout=10) + workers: Final = eventually(rig.proxy.worker_pids, lambda members: len(members) >= 2, seconds=30) + with rig.proxy.replacing_worker(min(workers)) as kill, ThreadPoolExecutor(max_workers=8) as pool: + futures: Final = tuple( + pool.submit(rig.stream_on_a_worker_that_may_die, rig.gemini_path(), rig.gemini_body()) + for _ in range(8) + ) + kill() first_wave: Final = tuple(future.result() for future in futures) survivors: Final = tuple(received for received in first_wave if received.status == 200) assert survivors, [received.text[:200] for received in first_wave] diff --git a/tests/integration/observability/test_prometheus_series_cap_chaos.py b/tests/integration/observability/test_prometheus_series_cap_chaos.py index c2f165ca37e..c62c5b4fb0f 100644 --- a/tests/integration/observability/test_prometheus_series_cap_chaos.py +++ b/tests/integration/observability/test_prometheus_series_cap_chaos.py @@ -6,7 +6,6 @@ from __future__ import annotations import json import re -import signal import threading from collections.abc import Iterator, Mapping, Sequence from concurrent.futures import ThreadPoolExecutor @@ -19,7 +18,6 @@ from types import MappingProxyType from typing import Final import httpx -import psutil import pytest from integration._support.client import eventually, object_value, string_value from integration._support.process import owned_gateway_image, setup_only_proxy_run @@ -255,10 +253,10 @@ def test_killed_worker_is_replaced_by_one_that_reads_the_same_admissions(tmp_pat plans: Final = _plans(extra) with ThreadPoolExecutor(max_workers=1) as pool: burst: Final = pool.submit(_burst, rig, plans, True) - victim: Final = psutil.Process(workers[0]) - victim.suspend() - victim.send_signal(signal.SIGKILL) - served: Final = burst.result() + with rig.proxy.replacing_worker(workers[0]) as kill: + kill.freeze() + kill() + served: Final = burst.result() answered: Final = tuple(item for item in served if item.status == 200) assert answered, [(item.status, item.text[:200]) for item in served] assert all(item.call.answer in item.text for item in answered) diff --git a/tests/integration/observability/test_s3_v2_partition_granularity.py b/tests/integration/observability/test_s3_v2_partition_granularity.py index a362ff55245..382463eb46c 100644 --- a/tests/integration/observability/test_s3_v2_partition_granularity.py +++ b/tests/integration/observability/test_s3_v2_partition_granularity.py @@ -13,7 +13,6 @@ from urllib.parse import quote, unquote import httpx import openai -import psutil import pytest import yaml from _s3_v2_support import ( @@ -31,7 +30,7 @@ from _s3_v2_support import ( from integration._support.client import Gateway, JsonValue, Scenario, eventually, object_value from integration._support.database import read_rows, scratch_database from integration._support.database_relay import database_relay -from integration._support.process import OwnedProxy, group_members, owned_proxy_process +from integration._support.process import OwnedProxy, owned_proxy_process from integration._support.wire import Reply, Request, wire_server FLUSH: Final = {"DEFAULT_S3_FLUSH_INTERVAL_SECONDS": "1"} @@ -1105,14 +1104,6 @@ def test_s3_v2_hour_slow_sink_batches_never_duplicate_an_upload(gateway: Gateway assert _batches_outside_layout(objects, "hour") == () -def _worker_processes(owned: OwnedProxy) -> tuple[int, ...]: - return tuple( - process.pid - for process in group_members(owned.process.pid) - if process.pid != owned.process.pid and "spawn_main" in " ".join(process.cmdline()) - ) - - def test_s3_v2_hour_worker_kill_mid_burst_keeps_the_other_worker_logging(gateway: Gateway, tmp_path: Path) -> None: marker: Final = "s3hkill" + uuid.uuid4().hex[:8] upstream: Final = CountingUpstream() @@ -1124,7 +1115,8 @@ def test_s3_v2_hour_worker_kill_mid_burst_keeps_the_other_worker_logging(gateway owned.gateway.scenario() as scenario, ): openai_model, _, key = _models(scenario, provider.url) - workers: Final = _worker_processes(owned) + workers: Final = owned.worker_pids() + assert len(workers) == 2, workers sent: Final = tuple(f"{marker}-{index}" for index in range(40)) def send(prompt: str) -> tuple[str, bool]: @@ -1143,13 +1135,14 @@ def test_s3_v2_hour_worker_kill_mid_burst_keeps_the_other_worker_logging(gateway return prompt, False return prompt, response.status_code == 200 and response.json()["id"] == prompt - with ThreadPoolExecutor(max_workers=16) as pool: - futures: Final = tuple(pool.submit(send, prompt) for prompt in sent) - eventually(lambda: len(upstream.received()), lambda count: count >= 8, seconds=30) - psutil.Process(workers[0]).kill() - results: Final = tuple(future.result() for future in futures) - later: Final = tuple(f"{marker}-later-{index}" for index in range(8)) - later_results: Final = tuple(send(prompt) for prompt in later) + with owned.replacing_worker(min(workers)) as kill: + with ThreadPoolExecutor(max_workers=16) as pool: + futures: Final = tuple(pool.submit(send, prompt) for prompt in sent) + eventually(lambda: len(upstream.received()), lambda count: count >= 8, seconds=30) + kill() + results: Final = tuple(future.result() for future in futures) + later: Final = tuple(f"{marker}-later-{index}" for index in range(8)) + later_results: Final = tuple(send(prompt) for prompt in later) eventually( lambda: frozenset(str(payload["id"]) for payload in sink.payloads()), lambda landed: frozenset(later) <= landed, @@ -1157,7 +1150,6 @@ def test_s3_v2_hour_worker_kill_mid_burst_keeps_the_other_worker_logging(gateway ) payloads: Final = sink.payloads() objects: Final = sink.objects() - assert len(workers) == 2, workers assert all(ok for _, ok in later_results), "the surviving worker must keep serving after the kill" landed: Final = tuple(str(payload["id"]) for payload in payloads) assert frozenset(landed) <= frozenset((*sent, *later)), "only ids this test sent may land" @@ -1215,7 +1207,7 @@ def test_s3_v2_hour_proxy_restart_mid_burst_keeps_the_layout_without_duplicates( with ThreadPoolExecutor(max_workers=16) as pool: futures: Final = tuple(pool.submit(send, first.gateway, first_key, prompt) for prompt in sent) eventually(lambda: len(upstream.received()), lambda count: count >= 8, seconds=30) - first.process.terminate() + first.begin_shutdown() results: Final = tuple(future.result() for future in futures) first.process.wait(timeout=30) answered: Final = frozenset(prompt for prompt, ok in results if ok) diff --git a/tests/integration/observability/test_s3_v2_streaming_failure_dedupe.py b/tests/integration/observability/test_s3_v2_streaming_failure_dedupe.py index 003181c72d8..d35a644ea3d 100644 --- a/tests/integration/observability/test_s3_v2_streaming_failure_dedupe.py +++ b/tests/integration/observability/test_s3_v2_streaming_failure_dedupe.py @@ -174,7 +174,6 @@ def test_retried_failure_uploads_one_s3_object( tmp_path, {"DEFAULT_S3_FLUSH_INTERVAL_SECONDS": "2"}, config=config, - workers=2, ) as candidate, candidate.scenario() as scenario, ): @@ -216,7 +215,6 @@ def test_fallback_chain_failure_uploads_one_s3_object( tmp_path, {"DEFAULT_S3_FLUSH_INTERVAL_SECONDS": "2"}, config=config, - workers=2, ) as candidate, candidate.scenario() as scenario, ): @@ -268,7 +266,6 @@ def test_streaming_success_uploads_one_s3_object( tmp_path, {"DEFAULT_S3_FLUSH_INTERVAL_SECONDS": "2"}, config=config, - workers=2, ) as candidate, candidate.scenario() as scenario, ): diff --git a/tests/integration/observability/test_s3_v2_upload_fanout.py b/tests/integration/observability/test_s3_v2_upload_fanout.py index 0154cce10df..ead1e5b028b 100644 --- a/tests/integration/observability/test_s3_v2_upload_fanout.py +++ b/tests/integration/observability/test_s3_v2_upload_fanout.py @@ -15,7 +15,7 @@ import yaml from _s3_v2_support import RecordingS3Sink, collect_payloads from _s3_v2_support import s3_config as _recording_s3_config from integration._support.client import Gateway, JsonValue, eventually -from integration._support.process import group_members, owned_proxy, owned_proxy_process +from integration._support.process import owned_proxy, owned_proxy_process from integration._support.wire import Reply, Request, Wire, wire_server BUCKET: Final = "integration-bucket" @@ -814,7 +814,8 @@ def test_s3_v2_worker_kill_mid_burst_keeps_surviving_deliveries(gateway: Gateway ): model: Final = scenario.model(api_base=provider.url + "/v1", api_key="synthetic-provider-key") key: Final = scenario.key(models=[model]) - sent: Final = tuple(f"{marker}-{index}" for index in range(REQUESTS)) + burst: Final = tuple(f"{marker}-burst-{index}" for index in range(REQUESTS)) + after: Final = tuple(f"{marker}-after-{index}" for index in range(REQUESTS // 4)) def send(identity: str) -> tuple[str, bool]: try: @@ -832,23 +833,28 @@ def test_s3_v2_worker_kill_mid_burst_keeps_surviving_deliveries(gateway: Gateway return identity, False return identity, response.status_code == 200 - with ThreadPoolExecutor(max_workers=32) as pool: - futures: Final = tuple(pool.submit(send, identity) for identity in sent) - time.sleep(0.5) - children: Final = tuple( - process for process in group_members(owned.process.pid) if process.pid != owned.process.pid - ) - assert children, "no worker children found to kill" - children[0].kill() - results: Final = tuple(future.result() for future in futures) - survivors: Final = frozenset(identity for identity, ok in results if ok) - assert survivors, "no request survived the worker kill" + workers: Final = owned.worker_pids() + assert workers, "no worker children found to kill" + with owned.replacing_worker(min(workers)) as kill: + with ThreadPoolExecutor(max_workers=32) as pool: + futures: Final = tuple(pool.submit(send, identity) for identity in burst) + time.sleep(0.5) + kill() + results: Final = tuple(future.result() for future in futures) readiness: Final = owned.gateway.client.get("/health/readiness") assert readiness.status_code == 200, readiness.text - payloads: Final = collect_payloads(sink, len(survivors), seconds=90) - landed: Final = frozenset(payload["id"] for payload in payloads) - assert survivors <= landed, "an id whose response succeeded never landed" - assert landed <= frozenset(sent), "an id that was never sent landed" + with ThreadPoolExecutor(max_workers=8) as pool: + after_results: Final = tuple(pool.map(send, after)) + burst_survivors: Final = frozenset(identity for identity, ok in results if ok) + after_survivors: Final = frozenset(identity for identity, ok in after_results if ok) + assert burst_survivors, "no request survived the worker kill" + assert after_survivors == frozenset(after), "the proxy did not serve after the worker was replaced" + landed: Final = eventually( + lambda: frozenset(payload["id"] for payload in sink.payloads()), + lambda ids: after_survivors <= ids and bool(ids & burst_survivors), + seconds=90, + ) + assert landed <= frozenset(burst) | frozenset(after), "an id that was never sent landed" @pytest.mark.covers("other.observability.s3_v2.sigterm_mid_burst_loses_only_inflight_without_duplicates") @@ -895,7 +901,7 @@ def test_s3_v2_sigterm_mid_burst_loses_only_inflight_without_duplicates(gateway: with ThreadPoolExecutor(max_workers=32) as pool: futures: Final = tuple(pool.submit(send, identity) for identity in sent) time.sleep(0.5) - candidate_owned.process.terminate() + candidate_owned.begin_shutdown() results: Final = tuple(future.result() for future in futures) candidate_owned.process.wait(timeout=30) finally: diff --git a/tests/integration/observability/test_signoz_delivery.py b/tests/integration/observability/test_signoz_delivery.py index 56a98b3f17c..5d4b66bbdd6 100644 --- a/tests/integration/observability/test_signoz_delivery.py +++ b/tests/integration/observability/test_signoz_delivery.py @@ -1,7 +1,6 @@ import asyncio import base64 import json -import os import re import signal import threading @@ -17,12 +16,11 @@ from typing import Final import anthropic import httpx import openai -import psutil import pytest import yaml from integration._support.client import Gateway, eventually, gateway_from_environment, object_value, string_value from integration._support.database import read_rows -from integration._support.process import OwnedProxy, owned_proxy_process +from integration._support.process import OwnedProxy, owned_proxy_process, stop_root_process from integration._support.wire import Reply, Request, Wire, wire_server from opentelemetry.proto.collector.trace.v1.trace_service_pb2 import ExportTraceServiceRequest from opentelemetry.proto.common.v1.common_pb2 import AnyValue @@ -915,28 +913,25 @@ def test_slow_sink_during_a_burst_lands_every_response_exactly_once(rig: Rig) -> def test_killing_one_of_two_workers_mid_burst_keeps_serving_and_never_duplicates_a_span(rig: Rig) -> None: - root: Final = psutil.Process(rig.process.process.pid) - workers: Final = eventually( - lambda: tuple(child for child in root.children() if "resource_tracker" not in " ".join(child.cmdline())), - lambda found: len(found) == 2, - seconds=30, - ) + workers: Final = eventually(rig.process.worker_pids, lambda found: len(found) == 2, seconds=30) markers: Final = tuple(_marker() for _ in range(24)) - def one(index: int) -> tuple[str, str | None]: - if index == 8: - os.kill(workers[0].pid, signal.SIGKILL) - try: - response: Final = rig.chat(markers[index]) - return f"chatcmpl-{markers[index]}", None if response.status_code == 200 else response.text - except httpx.HTTPError as error: - return f"chatcmpl-{markers[index]}", repr(error) + with rig.process.replacing_worker(min(workers)) as kill: - with ThreadPoolExecutor(max_workers=6) as pool: - results: Final = tuple(pool.map(one, range(24))) - assert rig.process.process.poll() is None, "Proxy root exited after a worker was killed" - after: Final = rig.chat(_marker()) - assert after.status_code == 200, after.text + def one(index: int) -> tuple[str, str | None]: + if index == 8: + kill() + try: + response: Final = rig.chat(markers[index]) + return f"chatcmpl-{markers[index]}", None if response.status_code == 200 else response.text + except httpx.HTTPError as error: + return f"chatcmpl-{markers[index]}", repr(error) + + with ThreadPoolExecutor(max_workers=6) as pool: + results: Final = tuple(pool.map(one, range(24))) + assert rig.process.process.poll() is None, "Proxy root exited after a worker was killed" + after: Final = rig.chat(_marker()) + assert after.status_code == 200, after.text rig.sink.single_span(_body_id(after)) failures: Final = tuple(error for _, error in results if error) assert all(error.startswith(("ReadError(", "RemoteProtocolError(", "ConnectError(")) for error in failures), ( @@ -969,8 +964,8 @@ def test_terminating_the_proxy_right_after_a_burst_flushes_every_span_before_exi assert all(response.status_code == 200 for response in responses), [response.text for response in responses] identities: Final = tuple(_body_id(response) for response in responses) pending_at_signal: Final = rig.sink.landed(identities) - rig.process.process.terminate() - assert rig.process.process.wait(timeout=40) in (0, -signal.SIGTERM) + assert stop_root_process(rig.process.process), "Proxy did not exit within the graceful stop budget" + assert rig.process.process.returncode in (0, -signal.SIGTERM) with pytest.raises(httpx.ConnectError): next(started) landed: Final = rig.sink.landed(identities) diff --git a/tests/integration/observability/test_straiker_v3_platform.py b/tests/integration/observability/test_straiker_v3_platform.py index 45ae562f5f7..a419d343535 100644 --- a/tests/integration/observability/test_straiker_v3_platform.py +++ b/tests/integration/observability/test_straiker_v3_platform.py @@ -10,8 +10,6 @@ from __future__ import annotations import hashlib import itertools import json -import os -import signal import socket import threading import uuid @@ -30,7 +28,7 @@ import pytest import yaml from integration._support.client import Gateway, eventually, gateway_from_environment, object_value from integration._support.database import read_rows -from integration._support.process import OwnedProxy, owned_proxy, owned_proxy_process +from integration._support.process import OwnedProxy, WorkerHealthcheck, owned_proxy, owned_proxy_process from integration._support.wire import Reply, Request, wire_server V3_KEY: Final = "sk_agt_synthetic_integration_key" @@ -436,7 +434,9 @@ def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]: with gateway_from_environment() as gateway, wire_server(_provider) as provider: config: Final = _rig_config(sink.url, root) with ( - owned_proxy_process(gateway, root, {}, config=config, workers=2) as owned, + owned_proxy_process( + gateway, root, {}, config=config, workers=2, healthcheck=WorkerHealthcheck.LENIENT + ) as owned, owned.gateway.scenario() as scenario, ): chat: Final = scenario.model( @@ -1070,31 +1070,17 @@ def test_burst_with_platform_outage_recovers_without_duplicate_spend(rig: Rig) - # C2: one proxy worker is killed during a burst; the other keeps serving and detect still runs for each call -def _is_live_worker(child: psutil.Process, exclude: int) -> bool: - if child.pid == exclude: - return False - try: - return child.status() != psutil.STATUS_ZOMBIE and "spawn_main" in " ".join(child.cmdline()) - except psutil.Error: - return False - - -def _uvicorn_workers(parent: psutil.Process, *, exclude: int = 0) -> tuple[psutil.Process, ...]: - return tuple(c for c in parent.children() if _is_live_worker(c, exclude)) - - def test_an_unreaped_worker_does_not_hide_the_surviving_proxy_worker(rig: Rig) -> None: - parent: Final = psutil.Process(rig.owned.process.pid) - workers: Final = eventually(lambda: _uvicorn_workers(parent), lambda children: len(children) == 2) + workers: Final = tuple(sorted(eventually(rig.owned.worker_pids, lambda pids: len(pids) == 2))) victim: Final = workers[0] survivor: Final = workers[1] marker: Final = rig.marker() - os.kill(parent.pid, signal.SIGSTOP) - try: - os.kill(victim.pid, signal.SIGKILL) - assert eventually(victim.status, lambda status: status == psutil.STATUS_ZOMBIE) == psutil.STATUS_ZOMBIE - assert tuple(child.pid for child in _uvicorn_workers(parent)) == (survivor.pid,) - assert _uvicorn_workers(parent, exclude=survivor.pid) == () + with rig.owned.replacing_worker(victim) as kill, rig.owned.suspended(rig.owned.process.pid): + kill() + assert eventually(psutil.Process(victim).status, lambda status: status == psutil.STATUS_ZOMBIE) == ( + psutil.STATUS_ZOMBIE + ) + assert rig.owned.worker_pids() == frozenset((survivor,)) with httpx.Client(base_url=rig._base(), timeout=15, trust_env=False) as client: response: Final = client.post( "/v1/chat/completions", @@ -1103,15 +1089,11 @@ def test_an_unreaped_worker_does_not_hide_the_surviving_proxy_worker(rig: Rig) - ) assert response.status_code == 200, response.text assert len(_v3_request_calls(rig, marker)) == 1 - finally: - os.kill(parent.pid, signal.SIGCONT) - assert len(eventually(lambda: _uvicorn_workers(parent, exclude=victim.pid), lambda children: len(children) == 2)) == 2 + assert len(rig.owned.worker_pids() - frozenset((victim,))) == 2 def test_burst_survives_one_worker_kill(rig: Rig) -> None: - parent: Final = psutil.Process(rig.owned.process.pid) - workers: Final = eventually(lambda: _uvicorn_workers(parent), lambda c: len(c) >= 2) - victim: Final = workers[0].pid + workers: Final = tuple(sorted(eventually(rig.owned.worker_pids, lambda pids: len(pids) == 2))) markers: Final = tuple(rig.marker() for _ in range(24)) def fresh_chat(text: str) -> tuple[int, str]: @@ -1126,18 +1108,19 @@ def test_burst_survives_one_worker_kill(rig: Rig) -> None: return 0, repr(error) return response.status_code, response.text - def call(index: int) -> tuple[int, str]: - if index == 6: - os.kill(victim, signal.SIGKILL) - return fresh_chat("kill " + markers[index]) + with rig.owned.replacing_worker(workers[0]) as kill: - with ThreadPoolExecutor(max_workers=4) as pool: - results: Final = tuple(pool.map(call, range(24))) + def call(index: int) -> tuple[int, str]: + if index == 6: + kill() + return fresh_chat("kill " + markers[index]) + + with ThreadPoolExecutor(max_workers=4) as pool: + results: Final = tuple(pool.map(call, range(24))) ok: Final = tuple(i for i, (status, _) in enumerate(results) if status == 200) assert len(ok) >= 20, results for index in ok: assert len(_v3_request_calls(rig, markers[index])) >= 1, markers[index] - eventually(lambda: _uvicorn_workers(parent, exclude=victim), lambda c: len(c) >= 2) after: Final = fresh_chat("after-kill " + rig.marker()) assert after[0] == 200, after diff --git a/tests/integration/providers/_coralbricks.py b/tests/integration/providers/_coralbricks.py new file mode 100644 index 00000000000..0cfbcd94d64 --- /dev/null +++ b/tests/integration/providers/_coralbricks.py @@ -0,0 +1,439 @@ +from __future__ import annotations + +import json +import re +import socket +import uuid +from collections.abc import Callable, Mapping +from dataclasses import dataclass +from pathlib import Path +from typing import Final + +import anthropic +import httpx +import openai +import pytest +from integration._support.client import JSON_OBJECT, Gateway, Scenario, eventually, object_value, string_value +from integration._support.database import read_rows +from integration._support.wire import Reply, Request, Wire +from pydantic import JsonValue + +PROVIDER: Final = "coralbricks" +MODEL: Final = "deepseek-v4.1-flash-fast" +MODELS: Final = (MODEL, "glm-5.3-fast") +API_KEY: Final = "cb_synthetic_integration_key" +ANSWER: Final = "coral ok" +NO_CACHE: Final[dict[str, JsonValue]] = {"no-cache": True} +PROMPT_TOKENS: Final = 4357 +CACHED_TOKENS: Final = 4352 +CACHE_WRITE_TOKENS: Final = 5 +COMPLETION_TOKENS: Final = 20 +REASONING_TOKENS: Final = 12 +COST_MAP: Final = Path(__file__).resolve().parents[3] / "model_prices_and_context_window.json" +LEAK_FIELDS: Final = frozenset( + {"litellm_params", "litellm_logging_obj", "litellm_call_id", "litellm_metadata", "proxy_server_request"} +) +READINESS_TEXT: Final = "model readiness" +MARKER: Final = re.compile(r"marker-([0-9a-f]{32})") +USER_MESSAGES: Final[list[dict[str, JsonValue]]] = [{"role": "user", "content": "Say hello to the reef"}] +SPEND_COLUMNS: Final = ( + "SELECT request_id, model, custom_llm_provider, model_group, call_type, status, spend, prompt_tokens, " + 'completion_tokens, api_base FROM "LiteLLM_SpendLogs"' +) + + +@dataclass(frozen=True, slots=True) +class Rates: + input: float + output: float + cache_write: float + cache_read: float + + +def number(value: JsonValue) -> float: + assert isinstance(value, (int, float)), value + return float(value) + + +def row_of(model: str) -> Mapping[str, JsonValue]: + row: Final = object_value(JSON_OBJECT.validate_json(COST_MAP.read_bytes())[f"{PROVIDER}/{model}"]) + assert row["litellm_provider"] == PROVIDER, row + return row + + +def catalog_rows() -> frozenset[str]: + rows: Final = frozenset( + key for key in JSON_OBJECT.validate_json(COST_MAP.read_bytes()) if key.startswith(f"{PROVIDER}/") + ) + assert {f"{PROVIDER}/{model}" for model in MODELS} <= rows, sorted(rows) + return rows + + +def rates_of(model: str) -> Rates: + row: Final = row_of(model) + return Rates( + input=number(row["input_cost_per_token"]), + output=number(row["output_cost_per_token"]), + cache_write=number(row["cache_creation_input_token_cost"]), + cache_read=number(row["cache_read_input_token_cost"]), + ) + + +def expected_spend(model: str) -> float: + rates: Final = rates_of(model) + return ( + (PROMPT_TOKENS - CACHED_TOKENS - CACHE_WRITE_TOKENS) * rates.input + + CACHED_TOKENS * rates.cache_read + + CACHE_WRITE_TOKENS * rates.cache_write + + COMPLETION_TOKENS * rates.output + ) + + +def chat_usage() -> dict[str, JsonValue]: + return { + "prompt_tokens": PROMPT_TOKENS, + "completion_tokens": COMPLETION_TOKENS, + "total_tokens": PROMPT_TOKENS + COMPLETION_TOKENS, + "prompt_tokens_details": { + "cached_tokens": CACHED_TOKENS, + "cache_write_tokens": CACHE_WRITE_TOKENS, + "billable_cache_write_tokens": CACHE_WRITE_TOKENS, + "cache_write_blocks": 3, + }, + "completion_tokens_details": {"reasoning_tokens": REASONING_TOKENS}, + } + + +def sse(payload: Mapping[str, JsonValue], event: str | None = None) -> bytes: + prefix: Final = f"event: {event}\n" if event is not None else "" + return f"{prefix}data: {json.dumps(payload)}\n\n".encode() + + +def chat_reply(identity: str, model: str, text: str = ANSWER) -> Reply: + body: Final[dict[str, JsonValue]] = { + "id": identity, + "object": "chat.completion", + "created": 1, + "model": model, + "choices": [{"index": 0, "message": {"role": "assistant", "content": text}, "finish_reason": "stop"}], + "usage": chat_usage(), + } + return Reply(body=json.dumps(body).encode()) + + +def chat_stream_reply(identity: str, model: str, text: str = ANSWER, pause: float = 0) -> Reply: + chunk: Final[dict[str, JsonValue]] = {"id": identity, "object": "chat.completion.chunk", "created": 1, "model": model} + frames: Final[tuple[dict[str, JsonValue], ...]] = ( + {**chunk, "choices": [{"index": 0, "delta": {"role": "assistant", "content": ""}}]}, + {**chunk, "choices": [{"index": 0, "delta": {"content": text}}]}, + {**chunk, "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}]}, + {**chunk, "choices": [], "usage": chat_usage()}, + ) + return Reply( + content_type="text/event-stream", + chunks=(*(sse(frame) for frame in frames), b"data: [DONE]\n\n"), + pause_between_chunks=pause, + ) + + +def responses_body(identity: str, model: str, text: str = ANSWER) -> dict[str, JsonValue]: + return { + "id": identity, + "object": "response", + "created_at": 1, + "status": "completed", + "model": model, + "output": [ + { + "id": f"msg_{identity}", + "type": "message", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": text, "annotations": []}], + } + ], + "usage": { + "input_tokens": PROMPT_TOKENS, + "output_tokens": COMPLETION_TOKENS, + "total_tokens": PROMPT_TOKENS + COMPLETION_TOKENS, + "input_tokens_details": { + "cached_tokens": CACHED_TOKENS, + "cache_write_tokens": CACHE_WRITE_TOKENS, + "billable_cache_write_tokens": CACHE_WRITE_TOKENS, + }, + "output_tokens_details": {"reasoning_tokens": REASONING_TOKENS}, + }, + } + + +def responses_reply(identity: str, model: str, stream: bool, text: str = ANSWER, pause: float = 0) -> Reply: + body: Final = responses_body(identity, model, text) + if not stream: + return Reply(body=json.dumps(body).encode()) + events: Final[tuple[dict[str, JsonValue], ...]] = ( + {"type": "response.created", "sequence_number": 0, "response": {**body, "status": "in_progress", "output": []}}, + { + "type": "response.output_text.delta", + "sequence_number": 1, + "item_id": f"msg_{identity}", + "output_index": 0, + "content_index": 0, + "delta": text, + }, + {"type": "response.completed", "sequence_number": 2, "response": body}, + ) + return Reply( + content_type="text/event-stream", + chunks=tuple(sse(event, string_value(event["type"])) for event in events), + pause_between_chunks=pause, + ) + + +def messages_usage(output_tokens: int) -> dict[str, JsonValue]: + return { + "input_tokens": PROMPT_TOKENS - CACHED_TOKENS - CACHE_WRITE_TOKENS, + "cache_read_input_tokens": CACHED_TOKENS, + "cache_creation_input_tokens": CACHE_WRITE_TOKENS, + "output_tokens": output_tokens, + } + + +def messages_reply(identity: str, model: str, text: str = ANSWER) -> Reply: + body: Final[dict[str, JsonValue]] = { + "id": identity, + "type": "message", + "role": "assistant", + "model": model, + "content": [{"type": "text", "text": text}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": messages_usage(COMPLETION_TOKENS), + } + return Reply(body=json.dumps(body).encode()) + + +def messages_stream_reply(identity: str, model: str, text: str = ANSWER, pause: float = 0) -> Reply: + events: Final[tuple[dict[str, JsonValue], ...]] = ( + { + "type": "message_start", + "message": { + "id": identity, + "type": "message", + "role": "assistant", + "model": model, + "content": [], + "stop_reason": None, + "stop_sequence": None, + "usage": messages_usage(1), + }, + }, + {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}, + {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": text}}, + {"type": "content_block_stop", "index": 0}, + { + "type": "message_delta", + "delta": {"stop_reason": "end_turn", "stop_sequence": None}, + "usage": {"output_tokens": COMPLETION_TOKENS}, + }, + {"type": "message_stop"}, + ) + return Reply( + content_type="text/event-stream", + chunks=tuple(sse(event, string_value(event["type"])) for event in events), + pause_between_chunks=pause, + ) + + +def error_reply(status: int, message: str) -> Reply: + body: Final[dict[str, JsonValue]] = {"error": {"message": message, "type": "scripted_error", "code": str(status)}} + return Reply(status=status, body=json.dumps(body).encode()) + + +def is_readiness_probe(request: Request) -> bool: + return bool(request.body) and READINESS_TEXT in request.body.decode(errors="replace") + + +def is_streaming(request: Request) -> bool: + return bool(request.body) and JSON_OBJECT.validate_json(request.body).get("stream") is True + + +def marker_of(request: Request) -> str | None: + found: Final = MARKER.search(request.body.decode(errors="replace")) + return found.group(1) if found is not None else None + + +def route(request: Request, identity: str, model: str, text: str = ANSWER, pause: float = 0) -> Reply: + match request.target: + case "/v1/chat/completions": + return ( + chat_stream_reply(identity, model, text, pause) + if is_streaming(request) + else chat_reply(identity, model, text) + ) + case "/v1/responses": + return responses_reply(identity, model, is_streaming(request), text, pause) + case "/v1/messages": + return ( + messages_stream_reply(identity, model, text, pause) + if is_streaming(request) + else messages_reply(identity, model, text) + ) + case _: + return error_reply(404, f"unknown route {request.target}") + + +def provider(identity: str, model: str = MODEL, stream_pause: float = 0) -> Callable[[Request], Reply]: + def respond(request: Request) -> Reply: + if is_readiness_probe(request): + return chat_reply(f"probe_{uuid.uuid4().hex}", model) + return route(request, identity, model, pause=stream_pause) + + return respond + + +def marker_identity(request: Request, marker: str) -> str: + match request.target: + case "/v1/responses": + return f"resp_{marker}" + case "/v1/messages": + return f"msg_{marker}" + case _: + return f"req_{marker}" + + +def marker_provider(model: str = MODEL, stream_pause: float = 0) -> Callable[[Request], Reply]: + def respond(request: Request) -> Reply: + marker: Final = marker_of(request) + if marker is None: + return chat_reply(f"probe_{uuid.uuid4().hex}", model) + return route(request, marker_identity(request, marker), model, f"answer marker-{marker}", stream_pause) + + return respond + + +def failing_provider(status: int, message: str) -> Callable[[Request], Reply]: + return lambda _request: error_reply(status, message) + + +def gateway_v1(gateway: Gateway) -> str: + return f"{str(gateway.client.base_url).rstrip('/')}/v1" + + +def openai_client(gateway: Gateway) -> openai.OpenAI: + return openai.OpenAI(base_url=gateway_v1(gateway), api_key=gateway.key, max_retries=0) + + +def async_openai_client(gateway: Gateway) -> openai.AsyncOpenAI: + return openai.AsyncOpenAI(base_url=gateway_v1(gateway), api_key=gateway.key, max_retries=0) + + +def anthropic_client(gateway: Gateway) -> anthropic.Anthropic: + return anthropic.Anthropic(base_url=str(gateway.client.base_url).rstrip("/"), api_key=gateway.key, max_retries=0) + + +def async_anthropic_client(gateway: Gateway) -> anthropic.AsyncAnthropic: + return anthropic.AsyncAnthropic( + base_url=str(gateway.client.base_url).rstrip("/"), api_key=gateway.key, max_retries=0 + ) + + +def not_yet_routable(response: object) -> bool: + status: Final = getattr(response, "status_code", None) + text: Final = getattr(response, "text", "") + return status == 400 and "Invalid model name" in text + + +def ready(gateway: Gateway, model: str, seconds: float = 60) -> None: + def probe() -> tuple[int, str]: + try: + response: Final = gateway.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": READINESS_TEXT}], "cache": NO_CACHE}, + ) + except httpx.TransportError as lost: + return 0, repr(lost) + return response.status_code, response.text + + def routable(status: int, text: str) -> bool: + return status != 0 and not (status == 400 and "Invalid model name" in text) + + eventually( + lambda: tuple(probe() for _ in range(2)), + lambda probes: all(routable(status, text) for status, text in probes), + seconds=seconds, + ) + + +def deployment(scenario: Scenario, wire: Wire, model: str = MODEL, **overrides: JsonValue) -> str: + alias: Final = scenario.model(model=f"{PROVIDER}/{model}", api_base=f"{wire.url}/v1", api_key=API_KEY, **overrides) + ready(scenario.gateway, alias) + return alias + + +def raw_deployment(scenario: Scenario, model_name: str, litellm_params: Mapping[str, JsonValue]) -> str: + created: Final = scenario.gateway.post( + "/model/new", {"model_name": model_name, "litellm_params": dict(litellm_params), "model_info": {}} + ) + identity: Final = string_value(object_value(created["model_info"])["id"]) + scenario.cleanups.callback(scenario.delete_model, identity) + return model_name + + +def closed_port() -> int: + with socket.socket() as probe: + probe.bind(("127.0.0.1", 0)) + return int(probe.getsockname()[1]) + + +def provider_calls(wire: Wire) -> tuple[Request, ...]: + return tuple(request for request in wire.drain() if not is_readiness_probe(request)) + + +def only_call(wire: Wire, target: str, model: str = MODEL, key: str = API_KEY) -> dict[str, JsonValue]: + received: Final = provider_calls(wire) + assert [(request.method, request.target) for request in received] == [("POST", target)], [ + (request.method, request.target) for request in received + ] + assert received[0].headers.get("authorization") == f"Bearer {key}", received[0].headers + if target == "/v1/messages": + assert received[0].headers.get("anthropic-version"), received[0].headers + body: Final = JSON_OBJECT.validate_json(received[0].body) + assert body["model"] == model, body + assert not LEAK_FIELDS & body.keys(), sorted(LEAK_FIELDS & body.keys()) + return body + + +def spend_row(identity: str, seconds: float = 70) -> dict[str, JsonValue]: + rows: Final = eventually( + lambda: read_rows(f"{SPEND_COLUMNS} WHERE request_id=%s", (identity,)), + lambda found: len(found) == 1, + seconds=seconds, + ) + return rows[0] + + +def group_rows(model_group: str, count: int, seconds: float = 70) -> list[dict[str, JsonValue]]: + return eventually( + lambda: read_rows(f"{SPEND_COLUMNS} WHERE model_group=%s AND request_id NOT LIKE 'probe_%%'", (model_group,)), + lambda found: len(found) == count, + seconds=seconds, + ) + + +def assert_billed(row: Mapping[str, JsonValue], alias: str, call_type: str, model: str = MODEL) -> None: + assert row["model_group"] == alias, row + assert row["model"] == f"{PROVIDER}/{model}", row + assert row["custom_llm_provider"] == PROVIDER, row + assert row["status"] == "success", row + assert row["call_type"] == call_type, row + assert row["prompt_tokens"] == PROMPT_TOKENS and row["completion_tokens"] == COMPLETION_TOKENS, row + assert number(row["spend"]) == pytest.approx(expected_spend(model)), (row, expected_spend(model)) + + +def assert_failed(row: Mapping[str, JsonValue], alias: str, model: str = MODEL) -> None: + assert row["model_group"] == alias, row + assert row["model"] == f"{PROVIDER}/{model}", row + assert row["custom_llm_provider"] == PROVIDER, row + assert row["status"] == "failure", row + assert number(row["spend"]) == 0, row diff --git a/tests/integration/providers/test_anthropic_thinking_signature_stream_wire.py b/tests/integration/providers/test_anthropic_thinking_signature_stream_wire.py index cf46885b6b5..fc0234410bb 100644 --- a/tests/integration/providers/test_anthropic_thinking_signature_stream_wire.py +++ b/tests/integration/providers/test_anthropic_thinking_signature_stream_wire.py @@ -1,7 +1,6 @@ import asyncio import json import re -import signal import threading import uuid from collections.abc import Mapping, Sequence @@ -763,19 +762,19 @@ async def test_worker_sigkill_mid_burst_leaves_the_sibling_streaming_signed_thin held_by: Final = MappingProxyType({pid: _open_upstream_connections(pid, wire.url) for pid in workers}) assert sum(held_by.values()) == 20, held_by victim_pid, survivor_pid = sorted(workers, key=held_by.__getitem__) - victim: Final = psutil.Process(victim_pid) - victim.suspend() - victim.send_signal(signal.SIGKILL) - release.set() - served: Final = await burst - assert held_by[survivor_pid] >= 10, held_by - completed: Final = tuple(item for item in served if item.status == 200) - assert len(completed) == held_by[survivor_pid], (held_by, [item.status for item in served]) - for item in completed: - if item.call.stream: - _assert_signed_once(deltas_of(chunks_of(item.text)), item.call.marker) - else: - assert answer(item.call.marker) in item.text, item.text + with owned.replacing_worker(victim_pid) as kill: + kill.freeze() + await asyncio.to_thread(kill) + release.set() + served: Final = await burst + assert held_by[survivor_pid] >= 10, held_by + completed: Final = tuple(item for item in served if item.status == 200) + assert len(completed) == held_by[survivor_pid], (held_by, [item.status for item in served]) + for item in completed: + if item.call.stream: + _assert_signed_once(deltas_of(chunks_of(item.text)), item.call.marker) + else: + assert answer(item.call.marker) in item.text, item.text follow_up: Final = _Call(endpoint="chat", stream=True, marker=uuid.uuid4().hex) (answered,) = await _burst(str(candidate.client.base_url), candidate.key, _CONFIG_MODEL, (follow_up,)) assert answered.status == 200, answered.text diff --git a/tests/integration/providers/test_assistants_route_errors_chaos.py b/tests/integration/providers/test_assistants_route_errors_chaos.py index 5e863932050..045c1ed852a 100644 --- a/tests/integration/providers/test_assistants_route_errors_chaos.py +++ b/tests/integration/providers/test_assistants_route_errors_chaos.py @@ -1,39 +1,36 @@ import asyncio -import signal import threading +from contextlib import AsyncExitStack from pathlib import Path from queue import SimpleQueue from types import MappingProxyType from typing import Final from urllib.parse import urlsplit -import psutil import pytest from integration._support.client import Gateway, eventually -from integration._support.process import graceful_stop_seconds, owned_proxy_process +from integration._support.process import WorkerHealthcheck, graceful_stop_seconds, owned_proxy_process from integration.providers._assistants_route_errors_support import ( MAPPED_BY_STATUS, MARKED_ROUTES, ROUTE_BY_NAME, Call, + Held, answer_all, assert_answered_with_its_own_marker, assistants_wire, held_error_peer, held_one_by_one, held_upstream_connections, - live_worker_pids, marker_of, new_marker, - open_upstream_connections, openai_assistants_config, provider_requests, worker_pids, ) _STATUSES: Final = (400, 404, 429, 500, 503) -_MIN_CALLS: Final = 20 -_MAX_CALLS: Final = 60 +_PER_WORKER: Final = 10 _STARTUP_COMPLETE: Final = "Application startup complete." @@ -45,7 +42,9 @@ def _startups(log: Path) -> tuple[int, int]: async def test_worker_sigkill_mid_burst_leaves_the_sibling_answering_each_mapped_status( gateway: Gateway, tmp_path: Path ) -> None: - calls: Final = tuple(Call(MARKED_ROUTES[index % len(MARKED_ROUTES)], new_marker()) for index in range(_MAX_CALLS)) + calls: Final = tuple( + Call(MARKED_ROUTES[index % len(MARKED_ROUTES)], new_marker()) for index in range(2 * _PER_WORKER) + ) follow_up: Final = Call(ROUTE_BY_NAME["get_thread"], new_marker()) status_by_marker: Final = MappingProxyType( { @@ -59,25 +58,30 @@ async def test_worker_sigkill_mid_burst_leaves_the_sibling_answering_each_mapped port: Final = urlsplit(wire.url).port assert port is not None config: Final = openai_assistants_config(tmp_path, port, 60) - with owned_proxy_process(gateway, tmp_path, {}, config=config, workers=2) as owned: + with owned_proxy_process( + gateway, tmp_path, {}, config=config, workers=2, healthcheck=WorkerHealthcheck.LENIENT + ) as owned: eventually(lambda: _startups(owned.log), lambda found: found == (2, 2), seconds=120) - workers: Final = live_worker_pids(owned.log) - assert len(workers) == 2, workers - - def both_workers_hold_enough() -> bool: - held: Final = tuple(open_upstream_connections(pid, port) for pid in workers) - return sum(held) >= _MIN_CALLS and min(held) > 0 - + victim_pid, survivor_pid = sorted(eventually(owned.worker_pids, lambda pids: len(pids) == 2)) try: - async with held_one_by_one(owned.gateway, calls, arrived, both_workers_hold_enough) as held: - held_by: Final = await asyncio.to_thread(held_upstream_connections, workers, port, len(held.calls)) - assert min(held_by.values()) > 0, held_by - victim_pid, survivor_pid = sorted(workers, key=held_by.__getitem__) - victim: Final = psutil.Process(victim_pid) - victim.suspend() - victim.send_signal(signal.SIGKILL) - release.set() - served: Final = await held.answers() + async with AsyncExitStack() as bursts: + with owned.suspended(survivor_pid): + on_victim: Final = await bursts.enter_async_context( + held_one_by_one(owned.gateway, calls[:_PER_WORKER], arrived, lambda: False) + ) + with owned.suspended(victim_pid): + on_survivor: Final = await bursts.enter_async_context( + held_one_by_one(owned.gateway, calls[_PER_WORKER:], arrived, lambda: False) + ) + held: Final = Held(on_victim.calls + on_survivor.calls, on_victim.pending + on_survivor.pending) + held_by: Final = await asyncio.to_thread( + held_upstream_connections, (victim_pid, survivor_pid), port, len(held.calls) + ) + assert held_by == {victim_pid: _PER_WORKER, survivor_pid: _PER_WORKER}, held_by + with owned.replacing_worker(victim_pid) as kill: + await asyncio.to_thread(kill) + release.set() + served: Final = await held.answers() finally: release.set() (answered,) = await answer_all(owned.gateway, (follow_up,)) diff --git a/tests/integration/providers/test_assistants_route_errors_wire.py b/tests/integration/providers/test_assistants_route_errors_wire.py index ec6d58e6670..395a99a8afe 100644 --- a/tests/integration/providers/test_assistants_route_errors_wire.py +++ b/tests/integration/providers/test_assistants_route_errors_wire.py @@ -1,9 +1,9 @@ import asyncio import json import threading +from contextlib import AsyncExitStack from collections.abc import Iterator, Mapping from dataclasses import dataclass -from pathlib import Path from queue import SimpleQueue from types import MappingProxyType from typing import Final @@ -11,8 +11,8 @@ from typing import Final import httpx import openai import pytest -from integration._support.client import Gateway, gateway_from_environment -from integration._support.process import graceful_stop_seconds, owned_proxy_process +from integration._support.client import Gateway, eventually, gateway_from_environment +from integration._support.process import OwnedProxy, WorkerHealthcheck, graceful_stop_seconds, owned_proxy_process from integration._support.wire import Reply, Request from integration.providers._assistants_route_errors_support import ( ASSISTANT_MODEL, @@ -28,6 +28,7 @@ from integration.providers._assistants_route_errors_support import ( STATUS_TABLE, UNMAPPED_500, Call, + Held, Mapped, Route, answer_all, @@ -42,7 +43,6 @@ from integration.providers._assistants_route_errors_support import ( held_error_peer, held_one_by_one, held_upstream_connections, - live_worker_pids, new_marker, openai_assistants_config, owned_config, @@ -71,15 +71,19 @@ _AZURE_CREDENTIALS: Final = ( _PRIVATE_DETAILS: Final = ("10.20.30.40", "/etc/litellm/secrets/db.yaml", "sk-proj-" + "a1B2c3D4" * 5) _URL_MODEL: Final = "http://169.254.169.254/latest" _BURST_STATUSES: Final = (400, 401, 404, 429, 500, 503) +_BURST_PER_WORKER: Final = 24 _SDK_ERRORS: Final = MappingProxyType({404: openai.NotFoundError, 429: openai.RateLimitError}) _BODY_VARIANTS: Final = ("text_502", "empty_503", "hostile_400", "null_message_404", "garbage_200") @dataclass(frozen=True, slots=True) class _Scripted: - gateway: Gateway + owned: OwnedProxy port: int - log: Path + + @property + def gateway(self) -> Gateway: + return self.owned.gateway @pytest.fixture(scope="module") @@ -89,9 +93,11 @@ def scripted(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Scripted]: config: Final = openai_assistants_config(directory, port, _DEPLOYMENT_TIMEOUT_SECONDS) with ( gateway_from_environment() as shared, - owned_proxy_process(shared, directory, {}, config=config, workers=2) as owned, + owned_proxy_process( + shared, directory, {}, config=config, workers=2, healthcheck=WorkerHealthcheck.LENIENT + ) as owned, ): - yield _Scripted(owned.gateway, port, owned.log) + yield _Scripted(owned, port) @pytest.fixture(scope="module") @@ -372,7 +378,9 @@ def test_repeated_failing_request_answers_the_same_each_time(scripted: _Scripted async def test_mixed_status_burst_answers_each_request_with_its_own_status_and_marker(scripted: _Scripted) -> None: - calls: Final = tuple(Call(MARKED_ROUTES[index % len(MARKED_ROUTES)], new_marker()) for index in range(48)) + calls: Final = tuple( + Call(MARKED_ROUTES[index % len(MARKED_ROUTES)], new_marker()) for index in range(2 * _BURST_PER_WORKER) + ) status_by_marker: Final = MappingProxyType( { call.marker: _BURST_STATUSES[(index // len(MARKED_ROUTES)) % len(_BURST_STATUSES)] @@ -381,18 +389,29 @@ async def test_mixed_status_burst_answers_each_request_with_its_own_status_and_m ) release: Final = threading.Event() arrived: Final[SimpleQueue[str]] = SimpleQueue() - workers: Final = live_worker_pids(scripted.log) - assert len(workers) == 2, workers + owned: Final = scripted.owned + first, second = sorted(eventually(owned.worker_pids, lambda pids: len(pids) == 2)) with assistants_wire(held_error_peer(status_by_marker, arrived, release), port=scripted.port) as wire: try: - async with held_one_by_one(scripted.gateway, calls, arrived, lambda: False) as held: - held_by: Final = await asyncio.to_thread(held_upstream_connections, workers, scripted.port, len(calls)) + async with AsyncExitStack() as bursts: + with owned.suspended(second): + on_first: Final = await bursts.enter_async_context( + held_one_by_one(owned.gateway, calls[:_BURST_PER_WORKER], arrived, lambda: False) + ) + with owned.suspended(first): + on_second: Final = await bursts.enter_async_context( + held_one_by_one(owned.gateway, calls[_BURST_PER_WORKER:], arrived, lambda: False) + ) + held: Final = Held(on_first.calls + on_second.calls, on_first.pending + on_second.pending) + held_by: Final = await asyncio.to_thread( + held_upstream_connections, (first, second), scripted.port, len(calls) + ) release.set() answers: Final = await held.answers() finally: release.set() received: Final = wire.drain() - assert sum(held_by.values()) == len(calls) and min(held_by.values()) > 0, held_by + assert held_by == {first: _BURST_PER_WORKER, second: _BURST_PER_WORKER}, held_by assert len(answers) == len(calls) for answer in answers: assert_answered_with_its_own_marker(answer, MAPPED_BY_STATUS[status_by_marker[answer.call.marker]]) diff --git a/tests/integration/providers/test_bedrock_converse_lookaround_regex_chaos.py b/tests/integration/providers/test_bedrock_converse_lookaround_regex_chaos.py index b92360a2193..2031be5f3be 100644 --- a/tests/integration/providers/test_bedrock_converse_lookaround_regex_chaos.py +++ b/tests/integration/providers/test_bedrock_converse_lookaround_regex_chaos.py @@ -1,7 +1,6 @@ import asyncio import json import re -import signal import threading import uuid from collections.abc import Mapping @@ -331,15 +330,15 @@ async def test_worker_sigkill_mid_burst_leaves_the_sibling_cleaning_schemas(gate held_by: Final = MappingProxyType({pid: _open_upstream_connections(pid, wire.url) for pid in workers}) assert sum(held_by.values()) == 20, held_by victim_pid, survivor_pid = sorted(workers, key=held_by.__getitem__) - victim: Final = psutil.Process(victim_pid) - victim.suspend() - victim.send_signal(signal.SIGKILL) - release.set() - served: Final = await burst - assert held_by[survivor_pid] >= 10, held_by - assert len(served) == held_by[survivor_pid], (held_by, len(served)) - for item in served: - _assert_answered_with_its_own_marker(item) + with owned.replacing_worker(victim_pid) as kill: + kill.freeze() + await asyncio.to_thread(kill) + release.set() + served: Final = await burst + assert held_by[survivor_pid] >= 10, held_by + assert len(served) == held_by[survivor_pid], (held_by, len(served)) + for item in served: + _assert_answered_with_its_own_marker(item) follow_up: Final = _Call(model=fleet.kimi_bare, endpoint="chat", stream=False, marker=uuid.uuid4().hex) (answered,) = await _burst(str(candidate.client.base_url), candidate.key, (follow_up,)) _assert_answered_with_its_own_marker(answered) diff --git a/tests/integration/providers/test_bedrock_converse_lookaround_regex_wire.py b/tests/integration/providers/test_bedrock_converse_lookaround_regex_wire.py index 97b64c663bb..428f0015766 100644 --- a/tests/integration/providers/test_bedrock_converse_lookaround_regex_wire.py +++ b/tests/integration/providers/test_bedrock_converse_lookaround_regex_wire.py @@ -10,6 +10,7 @@ import httpx import openai import pytest from integration._support.client import Gateway, Scenario, eventually, object_value, string_value +from integration._support.timing import assert_within_factor_of_control, timed from integration._support.upstream import _aws_event_frame from integration._support.wire import Reply, Request, Wire, wire_server from pydantic import JsonValue, TypeAdapter @@ -765,6 +766,10 @@ def test_the_worst_case_lookaround_input_scans_in_linear_time(gateway: Gateway) "type": "object", "properties": {"collection": {"type": "string", "pattern": pattern}}, } + benign: Final[dict[str, JsonValue]] = { + "type": "object", + "properties": {"collection": {"type": "string", "pattern": "a" * len(pattern)}}, + } liveliness: Final[list[tuple[float, int]]] = [] stop: Final = threading.Event() @@ -774,16 +779,24 @@ def test_the_worst_case_lookaround_input_scans_in_linear_time(gateway: Gateway) with wire_server(_bedrock_peer) as wire, gateway.scenario() as scenario: model: Final = _deployment(scenario, wire, f"bedrock/{_KIMI}") + control: Final = timed( + lambda: gateway.request("POST", _path("chat"), _body("chat", model, (_openai_tool(_TOOL, benign),))) + ) + _assert_tool_call_relayed("chat", control.value) + assert len(wire.drain()) == 1 poller: Final = threading.Thread(target=poll) poller.start() - started: Final = time.perf_counter() - response: Final = gateway.request("POST", _path("chat"), _body("chat", model, (_openai_tool(_TOOL, schema),))) - elapsed: Final = time.perf_counter() - started + worst: Final = timed( + lambda: gateway.request("POST", _path("chat"), _body("chat", model, (_openai_tool(_TOOL, schema),))) + ) stop.set() poller.join() - _assert_tool_call_relayed("chat", response) - assert elapsed < 30, elapsed - assert liveliness and max(latency for latency, _ in liveliness) < 5, liveliness + _assert_tool_call_relayed("chat", worst.value) + assert_within_factor_of_control(worst.seconds, control=control.seconds, factor=10) + assert liveliness, liveliness + assert_within_factor_of_control( + max(latency for latency, _ in liveliness), control=control.seconds, factor=10 + ) assert {status for _, status in liveliness} == {200}, liveliness assert len(wire.drain()) == 1 diff --git a/tests/integration/providers/test_bedrock_converse_missing_content_wire.py b/tests/integration/providers/test_bedrock_converse_missing_content_wire.py index 057596941a2..b0a19bbe817 100644 --- a/tests/integration/providers/test_bedrock_converse_missing_content_wire.py +++ b/tests/integration/providers/test_bedrock_converse_missing_content_wire.py @@ -3,7 +3,6 @@ import base64 import json import os import re -import signal import threading import uuid from collections.abc import Iterable, Iterator, Mapping, Sequence @@ -904,13 +903,14 @@ async def test_c02_worker_sigkill_mid_burst_leaves_the_sibling_serving_content_l held_by: Final = MappingProxyType({pid: _open_peer_connections(pid, wire.url) for pid in workers}) assert sum(held_by.values()) == 20, held_by victim_pid, survivor_pid = sorted(workers, key=held_by.__getitem__) - psutil.Process(victim_pid).send_signal(signal.SIGKILL) - release.set() - served: Final = await burst - assert held_by[survivor_pid] >= 10, held_by - assert len(served) == held_by[survivor_pid], (held_by, len(served)) - for item in served: - assert item.status == 200 and _ANSWER in item.text, (item.call, item.status, item.text) + with owned.replacing_worker(victim_pid) as kill: + await asyncio.to_thread(kill) + release.set() + served: Final = await burst + assert held_by[survivor_pid] >= 10, held_by + assert len(served) == held_by[survivor_pid], (held_by, len(served)) + for item in served: + assert item.status == 200 and _ANSWER in item.text, (item.call, item.status, item.text) eventually( lambda: len(_STARTED_WORKER.findall(owned.log.read_text())), lambda count: count == 3, seconds=60 ) @@ -947,7 +947,7 @@ async def test_c03_proxy_terminated_mid_burst_lands_every_answered_content_less_ _burst(_proxy_url(candidate), candidate.key, _PLAIN, calls, tolerate_transport_errors=True) ) await asyncio.to_thread(eventually, held_indexes.qsize, lambda size: size == 12, 60) - owned.process.terminate() + owned.begin_shutdown() release.set() served: Final = await burst eventually(owned.process.poll, lambda code: code is not None, seconds=60) diff --git a/tests/integration/providers/test_bedrock_converse_stream_event_frames_wire.py b/tests/integration/providers/test_bedrock_converse_stream_event_frames_wire.py index bf4f611107d..df566ea69e4 100644 --- a/tests/integration/providers/test_bedrock_converse_stream_event_frames_wire.py +++ b/tests/integration/providers/test_bedrock_converse_stream_event_frames_wire.py @@ -3,7 +3,6 @@ import base64 import json import os import re -import signal import threading import uuid from collections.abc import Callable, Iterable, Mapping @@ -873,17 +872,18 @@ async def test_c02_worker_sigkill_mid_burst_leaves_the_sibling_rejecting_the_hel held_by: Final = MappingProxyType({pid: _open_peer_connections(pid, wire.url) for pid in workers}) assert sum(held_by.values()) == 20, held_by victim_pid, survivor_pid = sorted(workers, key=held_by.__getitem__) - psutil.Process(victim_pid).send_signal(signal.SIGKILL) - release.set() - served: Final = await burst - assert held_by[survivor_pid] >= 10, held_by - assert len(served) == held_by[survivor_pid], (held_by, len(served)) - for item in served: - assert item.status in (200, 400) and _is_rejection(item, _REJECTION), ( - item.call, - item.status, - item.text, - ) + with owned.replacing_worker(victim_pid) as kill: + await asyncio.to_thread(kill) + release.set() + served: Final = await burst + assert held_by[survivor_pid] >= 10, held_by + assert len(served) == held_by[survivor_pid], (held_by, len(served)) + for item in served: + assert item.status in (200, 400) and _is_rejection(item, _REJECTION), ( + item.call, + item.status, + item.text, + ) eventually( lambda: len(_STARTED_WORKER.findall(owned.log.read_text())), lambda count: count == 3, seconds=60 ) @@ -921,7 +921,7 @@ async def test_c03_proxy_terminated_mid_burst_lands_every_served_rejection_at_mo _burst(_proxy_url(candidate), candidate.key, _PLAIN, calls, tolerate_transport_errors=True) ) await asyncio.to_thread(eventually, held_indexes.qsize, lambda size: size == 12, 60) - owned.process.terminate() + owned.begin_shutdown() release.set() served: Final = await burst eventually(owned.process.poll, lambda code: code is not None, seconds=60) diff --git a/tests/integration/providers/test_bedrock_kimi_k3_cache_point_wire.py b/tests/integration/providers/test_bedrock_kimi_k3_cache_point_wire.py index 9bf526bf3d1..db4e564b8e5 100644 --- a/tests/integration/providers/test_bedrock_kimi_k3_cache_point_wire.py +++ b/tests/integration/providers/test_bedrock_kimi_k3_cache_point_wire.py @@ -3,7 +3,6 @@ import base64 import json import os import re -import signal import socket import threading import time @@ -25,6 +24,7 @@ import yaml from integration._support.client import Gateway, Scenario, eventually, gateway_from_environment from integration._support.database import read_rows from integration._support.process import OwnedProxy, owned_proxy_process +from integration._support.timing import assert_within_factor_of_control, timed from integration._support.upstream import _aws_event_frame from integration._support.wire import Reply, Request, Wire, wire_server from pydantic import JsonValue, TypeAdapter @@ -1021,16 +1021,20 @@ def test_p02_a_plain_kimi_reply_is_priced_from_its_own_row(rig: _Rig) -> None: @pytest.mark.timeout(600) def test_s07_a_5kb_caller_model_under_the_wildcard_is_a_bounded_4xx_and_liveliness_stays_up(rig: _Rig) -> None: + control: Final = timed( + lambda: _send(rig.gateway, "chat", _chat_body(f"bedrock/{'a' * 8}", _prompt(), _SYSTEM_AND_USER)) + ) rig.wire.drain() - started: Final = time.monotonic() - outcome: Final = _send(rig.gateway, "chat", _chat_body(f"bedrock/{'a' * 5120}", _prompt(), _SYSTEM_AND_USER)) - elapsed: Final = time.monotonic() - started + oversized: Final = timed( + lambda: _send(rig.gateway, "chat", _chat_body(f"bedrock/{'a' * 5120}", _prompt(), _SYSTEM_AND_USER)) + ) + outcome: Final = oversized.value liveliness: Final = rig.gateway.client.get("/health/liveliness", timeout=5) assert liveliness.status_code == 200, liveliness.text assert outcome.status in (400, 404), (outcome.status, outcome.raw[:300]) error: Final = _JSON.validate_python(_JSON.validate_python(json.loads(outcome.raw))["error"]) assert "a" * 5120 in str(error["message"]), outcome.raw[:300] - assert elapsed < 10, elapsed + assert_within_factor_of_control(oversized.seconds, control=control.seconds, factor=20) assert rig.wire.drain() == () _failure_row(outcome.call_id) @@ -1494,14 +1498,15 @@ async def test_c02_a_worker_killed_mid_traffic_leaves_the_survivor_and_the_respa held_by: Final = MappingProxyType({pid: _open_peer_connections(pid, wire.url) for pid in workers}) assert sum(held_by.values()) == 20, held_by victim_pid, survivor_pid = sorted(workers, key=held_by.__getitem__) - psutil.Process(victim_pid).send_signal(signal.SIGKILL) - hold.set() - served: Final = await burst - assert len(served) == held_by[survivor_pid], (held_by, len(served)) - for item in served: - _assert_answered(item.outcome) - received: Final = _received(wire) - assert len(received) == 20 and all(item.cache_points == 0 for item in received), received + with owned.replacing_worker(victim_pid) as kill: + await asyncio.to_thread(kill) + hold.set() + served: Final = await burst + assert len(served) == held_by[survivor_pid], (held_by, len(served)) + for item in served: + _assert_answered(item.outcome) + received: Final = _received(wire) + assert len(received) == 20 and all(item.cache_points == 0 for item in received), received respawned: Final = eventually( lambda: tuple(int(pid) for pid in _STARTED_WORKER.findall(owned.log.read_text())), lambda pids: len(pids) == 3, @@ -1570,8 +1575,14 @@ def test_c03_a_restart_re_registers_a_stored_deployment_flag_from_the_database( @pytest.mark.timeout(900) async def test_c04_a_slow_peer_under_ten_concurrent_streams_completes_every_call_once(rig: _Rig) -> None: - rig.wire.drain() calls: Final = tuple(_Call("chat", index, True) for index in range(10)) + alone_started: Final = time.monotonic() + (alone,) = await _burst( + _proxy_url(rig.gateway), rig.gateway.key, "kimi-us", calls[:1], f"{_prompt()} {_SLOW_MARKER}" + ) + alone_seconds: Final = time.monotonic() - alone_started + _assert_answered(alone.outcome) + rig.wire.drain() started: Final = time.monotonic() served: Final = await _burst( _proxy_url(rig.gateway), rig.gateway.key, "kimi-us", calls, f"{_prompt()} {_SLOW_MARKER}" @@ -1582,5 +1593,5 @@ async def test_c04_a_slow_peer_under_ten_concurrent_streams_completes_every_call _assert_answered(item.outcome) received: Final = _received(rig.wire) assert len(received) == 10 and all(item.cache_points == 0 for item in received), received - assert elapsed < 30, elapsed + assert_within_factor_of_control(elapsed, control=alone_seconds, factor=len(calls)) _assert_rows_once(served) diff --git a/tests/integration/providers/test_bedrock_mantle_claude_chat_chaos.py b/tests/integration/providers/test_bedrock_mantle_claude_chat_chaos.py index d8a3306ecd5..9c7903497e3 100644 --- a/tests/integration/providers/test_bedrock_mantle_claude_chat_chaos.py +++ b/tests/integration/providers/test_bedrock_mantle_claude_chat_chaos.py @@ -1,7 +1,6 @@ import asyncio import json import re -import signal import threading import uuid from collections.abc import Callable @@ -322,15 +321,15 @@ async def test_worker_sigkill_mid_burst_leaves_the_sibling_serving_the_native_ro held_by: Final = MappingProxyType({pid: _open_upstream_connections(pid, wire.url) for pid in workers}) assert sum(held_by.values()) == 20, held_by victim_pid, survivor_pid = sorted(workers, key=held_by.__getitem__) - victim: Final = psutil.Process(victim_pid) - victim.suspend() - victim.send_signal(signal.SIGKILL) - release.set() - served: Final = await burst - assert held_by[survivor_pid] >= 10, held_by - assert len(served) == held_by[survivor_pid], (held_by, len(served)) - for item in served: - _assert_answered_with_its_own_marker(item) + with owned.replacing_worker(victim_pid) as kill: + kill.freeze() + await asyncio.to_thread(kill) + release.set() + served: Final = await burst + assert held_by[survivor_pid] >= 10, held_by + assert len(served) == held_by[survivor_pid], (held_by, len(served)) + for item in served: + _assert_answered_with_its_own_marker(item) follow_up: Final = _Call(endpoint="chat", stream=False, marker=uuid.uuid4().hex) (answered,) = await _burst(str(candidate.client.base_url), candidate.key, _CONFIG_MODEL, (follow_up,)) _assert_answered_with_its_own_marker(answered) diff --git a/tests/integration/providers/test_bedrock_mantle_context_overflow_fallbacks_wire.py b/tests/integration/providers/test_bedrock_mantle_context_overflow_fallbacks_wire.py new file mode 100644 index 00000000000..f40fd5295c1 --- /dev/null +++ b/tests/integration/providers/test_bedrock_mantle_context_overflow_fallbacks_wire.py @@ -0,0 +1,531 @@ +import json +from collections.abc import Iterator, Mapping +from dataclasses import dataclass +from pathlib import Path +from typing import Final +from uuid import uuid4 + +import httpx +import pytest +import yaml +from integration._support.client import Gateway, eventually, gateway_from_environment, object_value +from integration._support.process import OwnedProxy, graceful_stop_seconds, owned_proxy_process +from integration._support.wire import Reply, Request, Wire, wire_server +from pydantic import JsonValue, TypeAdapter + +pytestmark = pytest.mark.timeout(int(2 * graceful_stop_seconds() + 120)) + +_BACKEND: Final = "openai.gpt-5.6-luna" +_MANTLE_MODEL: Final = f"bedrock_mantle/{_BACKEND}" +_MANTLE_KEY: Final = "synthetic-mantle-bearer" +_OPENAI_KEY: Final = "synthetic-openai-key" +_RESPONSES_PATH: Final = "/openai/v1/responses" +_GENERIC: Final = "prompt is too long: your prompt exceeds the model's context window" +_UPSTREAM_MESSAGE: Final = ( + "Your input exceeds the context window of this model. Please adjust your input and try again." +) +_OPENAI_OVERFLOW_MESSAGE: Final = ( + "This model's maximum context length is 128000 tokens. However, your messages resulted in 130000 tokens." +) +_OPENAI_OVERFLOW_MARK: Final = "maximum context length is 128000 tokens" +_INVALID_INPUT_MESSAGE: Final = "Invalid 'input': expected a string or array" +_INVALID_PROMPT_MESSAGE: Final = "Invalid prompt: your prompt was flagged as potentially violating our usage policy." +_FALLBACK_TEXT: Final = "fallback answered" +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +_OVERFLOW_ENVELOPE: Final[dict[str, JsonValue]] = { + "error": { + "code": "context_length_exceeded", + "message": _UPSTREAM_MESSAGE, + "param": "input", + "type": "invalid_request_error", + } +} +_OVERFLOW_BODY: Final = json.dumps(_OVERFLOW_ENVELOPE).encode() +_BAD_INPUT_BODY: Final = json.dumps( + {"error": {"code": None, "message": _INVALID_INPUT_MESSAGE, "param": "input", "type": "invalid_request_error"}} +).encode() +_OPENAI_OVERFLOW_ERROR: Final[dict[str, JsonValue]] = { + "message": _OPENAI_OVERFLOW_MESSAGE, + "type": "invalid_request_error", + "param": "messages", + "code": "context_length_exceeded", +} +_OPENAI_OVERFLOW_BODY: Final = json.dumps({"error": _OPENAI_OVERFLOW_ERROR}).encode() + + +def _sse(events: tuple[dict[str, JsonValue], ...]) -> tuple[bytes, ...]: + return tuple(f"event: {event['type']}\ndata: {json.dumps(event)}\n\n".encode() for event in events) + + +def _response_object(identity: str, status: str, model: str) -> dict[str, JsonValue]: + return { + "id": f"resp_{identity}", + "object": "response", + "created_at": 1789788253, + "status": status, + "model": model, + "output": [], + } + + +def _failed_frames(identity: str, error: JsonValue, model: str = _BACKEND) -> tuple[bytes, ...]: + return _sse( + ( + { + "type": "response.created", + "sequence_number": 0, + "response": _response_object(identity, "in_progress", model), + }, + { + "type": "response.failed", + "sequence_number": 1, + "response": {**_response_object(identity, "failed", model), "error": error}, + }, + ) + ) + + +_PASSTHROUGH_FAILED_FRAMES: Final = _failed_frames("passthrough", _OVERFLOW_ENVELOPE["error"]) + + +def _streaming(request: Request) -> bool: + return _JSON_OBJECT.validate_json(request.body).get("stream") is True + + +def _overflow_peer(request: Request) -> Reply: + assert request.method == "POST" and request.target == _RESPONSES_PATH, request.target + if not _streaming(request): + return Reply(status=400, body=_OVERFLOW_BODY) + return Reply(content_type="text/event-stream", chunks=_failed_frames(uuid4().hex, _OVERFLOW_ENVELOPE["error"])) + + +def _passthrough_peer(request: Request) -> Reply: + assert request.method == "POST" and request.target == _RESPONSES_PATH, request.target + if not _streaming(request): + return Reply(status=400, body=_OVERFLOW_BODY) + return Reply(content_type="text/event-stream", chunks=_PASSTHROUGH_FAILED_FRAMES) + + +def _bad_input_peer(request: Request) -> Reply: + assert request.method == "POST" and request.target == _RESPONSES_PATH, request.target + if not _streaming(request): + return Reply(status=400, body=_BAD_INPUT_BODY) + error: Final[dict[str, JsonValue]] = {"code": "invalid_prompt", "message": _INVALID_PROMPT_MESSAGE} + return Reply(content_type="text/event-stream", chunks=_failed_frames(uuid4().hex, error)) + + +_EMPTY_MODEL_LIST: Final = Reply(body=b'{"object":"list","data":[]}') + + +def _openai_overflow_peer(request: Request) -> Reply: + if request.method == "GET": + return _EMPTY_MODEL_LIST + if request.target == "/v1/responses" and _streaming(request): + return Reply( + content_type="text/event-stream", + chunks=_failed_frames(uuid4().hex, _OPENAI_OVERFLOW_ERROR, "gpt-4o-mini"), + ) + assert request.target in ("/v1/chat/completions", "/v1/responses"), request.target + return Reply(status=400, body=_OPENAI_OVERFLOW_BODY) + + +def _chat_reply(identity: str, stream: bool) -> Reply: + if not stream: + return Reply( + body=json.dumps( + { + "id": f"chatcmpl-{identity}", + "object": "chat.completion", + "created": 1789788253, + "model": "gpt-4o-mini", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": _FALLBACK_TEXT}, + "finish_reason": "stop", + } + ], + "usage": {"prompt_tokens": 9, "completion_tokens": 2, "total_tokens": 11}, + } + ).encode() + ) + chunk: Final[dict[str, JsonValue]] = { + "id": f"chatcmpl-{identity}", + "object": "chat.completion.chunk", + "created": 1789788253, + "model": "gpt-4o-mini", + } + frames: Final = ( + { + **chunk, + "choices": [{"index": 0, "delta": {"role": "assistant", "content": _FALLBACK_TEXT}, "finish_reason": None}], + }, + {**chunk, "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}]}, + ) + return Reply( + content_type="text/event-stream", + chunks=tuple(f"data: {json.dumps(frame)}\n\n".encode() for frame in frames) + (b"data: [DONE]\n\n",), + ) + + +def _responses_reply(identity: str, stream: bool) -> Reply: + completed: Final[dict[str, JsonValue]] = { + **_response_object(identity, "completed", "gpt-4o-mini"), + "output": [ + { + "type": "message", + "id": f"msg_{identity}", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": _FALLBACK_TEXT, "annotations": []}], + } + ], + "usage": {"input_tokens": 9, "output_tokens": 2, "total_tokens": 11}, + } + if not stream: + return Reply(body=json.dumps(completed).encode()) + return Reply( + content_type="text/event-stream", + chunks=_sse( + ( + { + "type": "response.created", + "sequence_number": 0, + "response": {**completed, "status": "in_progress", "output": []}, + }, + { + "type": "response.output_text.delta", + "sequence_number": 1, + "item_id": f"msg_{identity}", + "output_index": 0, + "content_index": 0, + "delta": _FALLBACK_TEXT, + }, + {"type": "response.completed", "sequence_number": 2, "response": completed}, + ) + ), + ) + + +def _fallback_peer(request: Request) -> Reply: + if request.method == "GET": + return _EMPTY_MODEL_LIST + assert request.headers["authorization"] == f"Bearer {_OPENAI_KEY}", dict(request.headers) + identity: Final = uuid4().hex + if request.target == "/v1/responses": + return _responses_reply(identity, _streaming(request)) + assert request.target == "/v1/chat/completions", request.target + return _chat_reply(identity, _streaming(request)) + + +@dataclass(frozen=True, slots=True) +class Rig: + gateway: Gateway + owned: OwnedProxy + fallback: Wire + peers: Mapping[str, Wire] + models: Mapping[str, str] + + +def _mantle_params(api_base: str) -> dict[str, JsonValue]: + return {"model": _MANTLE_MODEL, "api_base": api_base, "api_key": _MANTLE_KEY} + + +def _openai_params(api_base: str) -> dict[str, JsonValue]: + return {"model": "openai/gpt-4o-mini", "api_base": api_base + "/v1", "api_key": _OPENAI_KEY} + + +def _write_config( + root: Path, + name: str, + model_list: list[dict[str, JsonValue]], + router_settings: dict[str, JsonValue], + pass_through_target: str | None, +) -> Path: + config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + config["model_list"] = model_list + config["router_settings"] = {"disable_cooldowns": True, "num_retries": 0, **router_settings} + if pass_through_target is not None: + config["general_settings"]["pass_through_endpoints"] = [ + { + "path": "/mantle-passthrough", + "target": pass_through_target, + "headers": {"Authorization": f"Bearer {_MANTLE_KEY}"}, + "auth": True, + } + ] + path: Final = root / f"{name}.yaml" + path.write_text(yaml.safe_dump(config)) + return path + + +@pytest.fixture(scope="module") +def cwf_rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]: + root: Final = tmp_path_factory.mktemp("mantle-cwf") + rig_id: Final = uuid4().hex[:8] + cwf: Final = f"mantle-cwf-{rig_id}" + fallback_name: Final = f"fallback-{rig_id}" + with ( + gateway_from_environment() as gateway, + wire_server(_overflow_peer) as overflow, + wire_server(_fallback_peer) as fallback, + ): + config: Final = _write_config( + root, + "cwf", + [ + {"model_name": cwf, "litellm_params": _mantle_params(overflow.url)}, + {"model_name": fallback_name, "litellm_params": _openai_params(fallback.url)}, + ], + {"context_window_fallbacks": [{cwf: [fallback_name]}]}, + None, + ) + with owned_proxy_process(gateway, root, {}, config=config, workers=2) as owned: + yield Rig(owned.gateway, owned, fallback, {"overflow": overflow}, {"cwf": cwf}) + + +@pytest.fixture(scope="module") +def fb_rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]: + root: Final = tmp_path_factory.mktemp("mantle-fb") + rig_id: Final = uuid4().hex[:8] + names: Final = { + "fb": f"mantle-fb-{rig_id}", + "bad": f"mantle-bad-{rig_id}", + "openai_fb": f"openai-fb-{rig_id}", + } + fallback_name: Final = f"fallback-{rig_id}" + with ( + gateway_from_environment() as gateway, + wire_server(_overflow_peer) as overflow, + wire_server(_passthrough_peer) as passthrough, + wire_server(_bad_input_peer) as bad, + wire_server(_openai_overflow_peer) as openai_overflow, + wire_server(_fallback_peer) as fallback, + ): + config: Final = _write_config( + root, + "fb", + [ + {"model_name": names["fb"], "litellm_params": _mantle_params(overflow.url)}, + {"model_name": names["bad"], "litellm_params": _mantle_params(bad.url)}, + {"model_name": names["openai_fb"], "litellm_params": _openai_params(openai_overflow.url)}, + {"model_name": fallback_name, "litellm_params": _openai_params(fallback.url)}, + ], + {"fallbacks": [{name: [fallback_name]} for name in names.values()]}, + passthrough.url + _RESPONSES_PATH, + ) + with owned_proxy_process(gateway, root, {}, config=config, workers=2) as owned: + yield Rig( + owned.gateway, + owned, + fallback, + {"overflow": overflow, "passthrough": passthrough, "bad": bad, "openai_overflow": openai_overflow}, + names, + ) + + +def _chat(model: str, prompt: str, stream: bool = False) -> dict[str, JsonValue]: + return {"model": model, "messages": [{"role": "user", "content": prompt}], "stream": stream} + + +def _messages(model: str, prompt: str, stream: bool = False) -> dict[str, JsonValue]: + return {"model": model, "max_tokens": 32, "messages": [{"role": "user", "content": prompt}], "stream": stream} + + +def _responses(model: str, prompt: str, stream: bool = False) -> dict[str, JsonValue]: + return {"model": model, "input": prompt, "stream": stream} + + +def _consumed(gateway: Gateway, path: str, body: Mapping[str, JsonValue]) -> tuple[httpx.Response, str]: + with gateway.client.stream("POST", path, json=body, headers={"Authorization": f"Bearer {gateway.key}"}) as response: + text: Final = b"".join(response.iter_bytes()).decode() + return response, text + + +def _events(text: str) -> tuple[dict[str, JsonValue], ...]: + return tuple( + _JSON_OBJECT.validate_json(line.removeprefix("data: ")) + for line in text.splitlines() + if line.startswith("data: ") and line != "data: [DONE]" + ) + + +def _only_error(text: str) -> dict[str, JsonValue]: + errors: Final = tuple(event for event in _events(text) if event.get("type") == "error") + assert len(errors) == 1, text + return object_value(errors[0]["error"]) + + +def _only_failed(text: str) -> dict[str, JsonValue]: + failed: Final = tuple(event for event in _events(text) if event.get("type") == "response.failed") + assert len(failed) == 1, text + return object_value(object_value(failed[0]["response"])["error"]) + + +def _calls_with(wire: Wire, prompt: str) -> int: + return sum(1 for request in wire.drain() if prompt in request.body.decode()) + + +def _prompt(label: str) -> str: + return f"{label} {uuid4().hex}" + + +def _assert_fell_back(rig: Rig, peer: str, prompt: str, response_text: str) -> None: + assert _FALLBACK_TEXT in response_text, response_text + assert _calls_with(rig.peers[peer], prompt) == 1 + assert _calls_with(rig.fallback, prompt) == 1 + + +def _assert_no_fallback(rig: Rig, peer: str, prompt: str) -> None: + assert _calls_with(rig.peers[peer], prompt) == 1 + assert _calls_with(rig.fallback, prompt) == 0 + + +def test_context_window_fallback_fires_for_the_overflow_on_chat_completions(cwf_rig: Rig) -> None: + prompt: Final = _prompt("cwf chat") + response: Final = cwf_rig.gateway.request("POST", "/v1/chat/completions", _chat(cwf_rig.models["cwf"], prompt)) + assert response.status_code == 200, response.text + _assert_fell_back(cwf_rig, "overflow", prompt, response.text) + + +def test_context_window_fallback_fires_for_the_overflow_on_messages(cwf_rig: Rig) -> None: + prompt: Final = _prompt("cwf messages") + response: Final = cwf_rig.gateway.request("POST", "/v1/messages", _messages(cwf_rig.models["cwf"], prompt)) + assert response.status_code == 200, response.text + _assert_fell_back(cwf_rig, "overflow", prompt, response.text) + + +def test_context_window_fallback_fires_for_the_overflow_on_responses(cwf_rig: Rig) -> None: + prompt: Final = _prompt("cwf responses") + response: Final = cwf_rig.gateway.request("POST", "/v1/responses", _responses(cwf_rig.models["cwf"], prompt)) + assert response.status_code == 200, response.text + _assert_fell_back(cwf_rig, "overflow", prompt, response.text) + + +def test_context_window_fallback_does_not_fire_on_streamed_chat_completions(cwf_rig: Rig) -> None: + prompt: Final = _prompt("cwf chat stream") + response, text = _consumed(cwf_rig.gateway, "/v1/chat/completions", _chat(cwf_rig.models["cwf"], prompt, True)) + assert response.status_code == 400, text + assert _GENERIC in text, text + _assert_no_fallback(cwf_rig, "overflow", prompt) + + +def test_context_window_fallback_is_not_consulted_on_streamed_messages(cwf_rig: Rig) -> None: + prompt: Final = _prompt("cwf messages stream") + response, text = _consumed(cwf_rig.gateway, "/v1/messages", _messages(cwf_rig.models["cwf"], prompt, True)) + assert response.status_code == 200, text + error: Final = _only_error(text) + assert error["type"] == "invalid_request_error", text + assert _GENERIC in str(error["message"]), text + _assert_no_fallback(cwf_rig, "overflow", prompt) + + +def test_context_window_fallback_does_not_fire_on_streamed_responses(cwf_rig: Rig) -> None: + prompt: Final = _prompt("cwf responses stream") + response, text = _consumed(cwf_rig.gateway, "/v1/responses", _responses(cwf_rig.models["cwf"], prompt, True)) + assert response.status_code == 200, text + assert _GENERIC in str(_only_failed(text)["message"]), text + _assert_no_fallback(cwf_rig, "overflow", prompt) + + +def test_plain_fallback_fires_for_the_overflow_on_chat_completions(fb_rig: Rig) -> None: + prompt: Final = _prompt("fb chat") + response: Final = fb_rig.gateway.request("POST", "/v1/chat/completions", _chat(fb_rig.models["fb"], prompt)) + assert response.status_code == 200, response.text + _assert_fell_back(fb_rig, "overflow", prompt, response.text) + + +def test_plain_fallback_fires_for_the_overflow_on_messages(fb_rig: Rig) -> None: + prompt: Final = _prompt("fb messages") + response: Final = fb_rig.gateway.request("POST", "/v1/messages", _messages(fb_rig.models["fb"], prompt)) + assert response.status_code == 200, response.text + _assert_fell_back(fb_rig, "overflow", prompt, response.text) + + +def test_plain_fallback_fires_for_the_overflow_on_responses(fb_rig: Rig) -> None: + prompt: Final = _prompt("fb responses") + response: Final = fb_rig.gateway.request("POST", "/v1/responses", _responses(fb_rig.models["fb"], prompt)) + assert response.status_code == 200, response.text + _assert_fell_back(fb_rig, "overflow", prompt, response.text) + + +def test_plain_fallback_does_not_fire_on_streamed_chat_completions(fb_rig: Rig) -> None: + prompt: Final = _prompt("fb chat stream") + response, text = _consumed(fb_rig.gateway, "/v1/chat/completions", _chat(fb_rig.models["fb"], prompt, True)) + assert response.status_code == 400, text + assert _GENERIC in text, text + _assert_no_fallback(fb_rig, "overflow", prompt) + + +def test_streamed_messages_overflow_on_a_plain_fallback_deployment_returns_400_without_falling_back( + fb_rig: Rig, +) -> None: + prompt: Final = _prompt("fb messages stream") + response, text = _consumed(fb_rig.gateway, "/v1/messages", _messages(fb_rig.models["fb"], prompt, True)) + assert response.status_code == 400, text + error: Final = object_value(_JSON_OBJECT.validate_json(text)["error"]) + assert error["type"] == "invalid_request_error", text + assert _GENERIC in str(error["message"]), text + _assert_no_fallback(fb_rig, "overflow", prompt) + + +def test_plain_fallback_does_not_fire_on_streamed_responses(fb_rig: Rig) -> None: + prompt: Final = _prompt("fb responses stream") + response, text = _consumed(fb_rig.gateway, "/v1/responses", _responses(fb_rig.models["fb"], prompt, True)) + assert response.status_code == 200, text + assert _GENERIC in str(_only_failed(text)["message"]), text + _assert_no_fallback(fb_rig, "overflow", prompt) + + +def test_streamed_messages_overflow_on_an_openai_fallback_deployment_returns_the_error_event_without_falling_back( + fb_rig: Rig, +) -> None: + prompt: Final = _prompt("openai fb messages stream") + response, text = _consumed(fb_rig.gateway, "/v1/messages", _messages(fb_rig.models["openai_fb"], prompt, True)) + assert response.status_code == 200, text + error: Final = _only_error(text) + assert error["type"] == "invalid_request_error", text + assert _OPENAI_OVERFLOW_MARK in str(error["message"]), text + _assert_no_fallback(fb_rig, "openai_overflow", prompt) + + +def test_streamed_messages_non_overflow_failure_returns_400_without_falling_back(fb_rig: Rig) -> None: + prompt: Final = _prompt("bad messages stream") + response, text = _consumed(fb_rig.gateway, "/v1/messages", _messages(fb_rig.models["bad"], prompt, True)) + assert response.status_code == 400, text + error: Final = object_value(_JSON_OBJECT.validate_json(text)["error"]) + assert error["type"] == "invalid_request_error", text + assert _INVALID_PROMPT_MESSAGE in str(error["message"]), text + assert _GENERIC not in text, text + _assert_no_fallback(fb_rig, "bad", prompt) + + +def test_pass_through_relays_the_overflow_envelope_verbatim(fb_rig: Rig) -> None: + prompt: Final = _prompt("passthrough") + response: Final = fb_rig.gateway.request("POST", "/mantle-passthrough", {"model": _BACKEND, "input": prompt}) + assert response.status_code == 400, response.text + assert _JSON_OBJECT.validate_json(response.content) == _OVERFLOW_ENVELOPE, response.text + assert _calls_with(fb_rig.peers["passthrough"], prompt) == 1 + + +def test_pass_through_relays_the_failed_stream_frames_verbatim(fb_rig: Rig) -> None: + prompt: Final = _prompt("passthrough stream") + body: Final[dict[str, JsonValue]] = {"model": _BACKEND, "input": prompt, "stream": True} + response, text = _consumed(fb_rig.gateway, "/mantle-passthrough", body) + assert response.status_code == 200, text + assert text.encode() == b"".join(_PASSTHROUGH_FAILED_FRAMES), text + assert _calls_with(fb_rig.peers["passthrough"], prompt) == 1 + + +def test_killing_one_worker_leaves_the_sibling_serving_the_fallback(fb_rig: Rig) -> None: + workers: Final = eventually(fb_rig.owned.worker_pids, lambda pids: len(pids) >= 2, seconds=30) + with fb_rig.owned.replacing_worker(min(workers)) as kill: + kill() + responses: Final = tuple( + fb_rig.gateway.request("POST", "/v1/chat/completions", _chat(fb_rig.models["fb"], _prompt("after kill"))) + for _ in range(3) + ) + for response in responses: + assert response.status_code == 200, response.text + assert _FALLBACK_TEXT in response.text, response.text + assert fb_rig.gateway.request("GET", "/health/liveliness").status_code == 200 diff --git a/tests/integration/providers/test_bedrock_mantle_context_overflow_wire.py b/tests/integration/providers/test_bedrock_mantle_context_overflow_wire.py new file mode 100644 index 00000000000..51fc1aa283b --- /dev/null +++ b/tests/integration/providers/test_bedrock_mantle_context_overflow_wire.py @@ -0,0 +1,795 @@ +import json +from collections.abc import Callable, Mapping +from concurrent.futures import ThreadPoolExecutor +from itertools import product +from typing import Final +from uuid import uuid4 + +import anthropic +import httpx +import openai +import pytest +from integration._support.client import Gateway, eventually, object_value +from integration._support.database import read_rows +from integration._support.wire import Reply, Request, Wire, wire_server +from pydantic import JsonValue, TypeAdapter + +_BACKEND: Final = "openai.gpt-5.6-luna" +_MODEL: Final = f"bedrock_mantle/{_BACKEND}" +_API_KEY: Final = "synthetic-mantle-bearer" +_RESPONSES_PATH: Final = "/openai/v1/responses" +_OPENAI_MODEL: Final = "openai/gpt-4o-mini" +_OPENAI_CHAT_PATH: Final = "/v1/chat/completions" +_OPENAI_RESPONSES_PATH: Final = "/v1/responses" +_OPENAI_API_KEY: Final = "synthetic-openai-key" +_GENERIC: Final = "prompt is too long: your prompt exceeds the model's context window" +_TOO_LONG: Final = "prompt is too long" +_UPSTREAM_MESSAGE: Final = ( + "Your input exceeds the context window of this model. Please adjust your input and try again." +) +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +_JSON_VALUE: Final = TypeAdapter(JsonValue) +_INVALID_INPUT_MESSAGE: Final = "Invalid 'input': expected a string or array" +_INVALID_PROMPT_MESSAGE: Final = "Invalid prompt: your prompt was flagged as potentially violating our usage policy." +_OPENAI_OVERFLOW_MESSAGE: Final = ( + "This model's maximum context length is 128000 tokens. However, your messages resulted in 130000 tokens. " + "Please reduce the length of the messages." +) +_OPENAI_OVERFLOW_MARK: Final = "maximum context length is 128000 tokens" +_LEGACY_PROMPT_TOKENS: Final = 1055489 +_LEGACY_MODEL_MAXIMUM: Final = 1050000 +_LEGACY_MESSAGE: Final = ( + f"prompt tokens ({_LEGACY_PROMPT_TOKENS}) exceed model maximum ({_LEGACY_MODEL_MAXIMUM}) for {_BACKEND}" +) +_HAPPY_TEXT: Final = "mantle overflow audit control" + + +def _envelope(code: str | None, message: str) -> bytes: + return json.dumps( + {"error": {"code": code, "message": message, "param": "input", "type": "invalid_request_error"}} + ).encode() + + +_OVERFLOW_BODY: Final = _envelope("context_length_exceeded", _UPSTREAM_MESSAGE) +_LEGACY_BODY: Final = _envelope("validation_error", _LEGACY_MESSAGE) +_BAD_INPUT_BODY: Final = _envelope(None, _INVALID_INPUT_MESSAGE) +_OPENAI_OVERFLOW_ERROR: Final[dict[str, JsonValue]] = { + "message": _OPENAI_OVERFLOW_MESSAGE, + "type": "invalid_request_error", + "param": "messages", + "code": "context_length_exceeded", +} +_OPENAI_OVERFLOW_BODY: Final = json.dumps({"error": _OPENAI_OVERFLOW_ERROR}).encode() +_OVERFLOW: Final = Reply(status=400, body=_OVERFLOW_BODY) +_OVERFLOW_STREAM_ERROR: Final[dict[str, JsonValue]] = {"code": "context_length_exceeded", "message": _UPSTREAM_MESSAGE} + + +def _response_object(identity: str, status: str, text: str | None, model: str = _BACKEND) -> dict[str, JsonValue]: + output: Final[list[JsonValue]] = ( + [] + if text is None + else [ + { + "type": "message", + "id": f"msg_{identity}", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": text, "annotations": []}], + } + ] + ) + return { + "id": f"resp_{identity}", + "object": "response", + "created_at": 1789788253, + "status": status, + "model": model, + "output": output, + "usage": {"input_tokens": 21, "output_tokens": 4, "total_tokens": 25}, + } + + +def _frames(events: tuple[dict[str, JsonValue], ...]) -> tuple[bytes, ...]: + return tuple(f"event: {event['type']}\ndata: {json.dumps(event)}\n\n".encode() for event in events) + + +def _failed_stream(error: JsonValue, model: str = _BACKEND) -> Reply: + identity: Final = uuid4().hex + return Reply( + content_type="text/event-stream", + chunks=_frames( + ( + { + "type": "response.created", + "sequence_number": 0, + "response": _response_object(identity, "in_progress", None, model), + }, + { + "type": "response.failed", + "sequence_number": 1, + "response": {**_response_object(identity, "failed", None, model), "error": error}, + }, + ) + ), + ) + + +def _happy_reply(stream: bool) -> Reply: + identity: Final = uuid4().hex + completed: Final = _response_object(identity, "completed", _HAPPY_TEXT) + if not stream: + return Reply(body=json.dumps(completed).encode()) + return Reply( + content_type="text/event-stream", + chunks=_frames( + ( + { + "type": "response.created", + "sequence_number": 0, + "response": _response_object(identity, "in_progress", None), + }, + { + "type": "response.output_text.delta", + "sequence_number": 1, + "item_id": f"msg_{identity}", + "output_index": 0, + "content_index": 0, + "delta": _HAPPY_TEXT, + }, + {"type": "response.completed", "sequence_number": 2, "response": completed}, + ) + ), + ) + + +def _mantle_peer(prompt: str, reply: Reply) -> Callable[[Request], Reply]: + def respond(request: Request) -> Reply: + assert request.method == "POST" and request.target == _RESPONSES_PATH, request.target + assert request.headers["authorization"] == f"Bearer {_API_KEY}", dict(request.headers) + body: Final = _JSON_OBJECT.validate_json(request.body) + assert body["model"] == _BACKEND, body + assert prompt in json.dumps(body["input"]), body + return reply + + return respond + + +def _overflowing_mantle_peer(prompt: str) -> Callable[[Request], Reply]: + def respond(request: Request) -> Reply: + assert request.method == "POST" and request.target == _RESPONSES_PATH, request.target + body: Final = _JSON_OBJECT.validate_json(request.body) + assert prompt in json.dumps(body["input"]), body + return _failed_stream(_OVERFLOW_STREAM_ERROR) if body.get("stream") is True else _OVERFLOW + + return respond + + +def _openai_peer(prompt: str) -> Callable[[Request], Reply]: + def respond(request: Request) -> Reply: + if request.method == "GET": + return Reply(body=b'{"object":"list","data":[]}') + assert request.headers["authorization"] == f"Bearer {_OPENAI_API_KEY}", dict(request.headers) + assert prompt in request.body.decode(), request.body + body: Final = _JSON_OBJECT.validate_json(request.body) + if request.target == _OPENAI_RESPONSES_PATH and body.get("stream") is True: + return _failed_stream(_OPENAI_OVERFLOW_ERROR, "gpt-4o-mini") + assert request.target in (_OPENAI_CHAT_PATH, _OPENAI_RESPONSES_PATH), request.target + return Reply(status=400, body=_OPENAI_OVERFLOW_BODY) + + return respond + + +def _chat(model: str, prompt: str, stream: bool = False) -> dict[str, JsonValue]: + return {"model": model, "messages": [{"role": "user", "content": prompt}], "stream": stream} + + +def _messages(model: str, prompt: str, stream: bool = False) -> dict[str, JsonValue]: + return {"model": model, "max_tokens": 32, "messages": [{"role": "user", "content": prompt}], "stream": stream} + + +def _responses(model: str, prompt: str, stream: bool = False) -> dict[str, JsonValue]: + return {"model": model, "input": prompt, "stream": stream} + + +def _consumed(gateway: Gateway, path: str, body: Mapping[str, JsonValue]) -> tuple[httpx.Response, str]: + with gateway.client.stream("POST", path, json=body, headers={"Authorization": f"Bearer {gateway.key}"}) as response: + text: Final = b"".join(response.iter_bytes()).decode() + return response, text + + +def _events(text: str) -> tuple[dict[str, JsonValue], ...]: + return tuple( + _JSON_OBJECT.validate_json(line.removeprefix("data: ")) + for line in text.splitlines() + if line.startswith("data: ") and line != "data: [DONE]" + ) + + +def _error_events(text: str) -> tuple[dict[str, JsonValue], ...]: + return tuple(event for event in _events(text) if event.get("type") == "error") + + +def _only_error(text: str) -> dict[str, JsonValue]: + errors: Final = _error_events(text) + assert len(errors) == 1, text + return object_value(errors[0]["error"]) + + +def _only_failed(text: str) -> dict[str, JsonValue]: + failed: Final = tuple(event for event in _events(text) if event.get("type") == "response.failed") + assert len(failed) == 1, text + return object_value(object_value(failed[0]["response"])["error"]) + + +def _error_object(response: httpx.Response) -> dict[str, JsonValue]: + return object_value(_JSON_OBJECT.validate_json(response.content)["error"]) + + +def _error_message(response: httpx.Response) -> str: + return str(_error_object(response)["message"]) + + +def _only_call(wire: Wire, target: str = _RESPONSES_PATH) -> None: + assert [(request.method, request.target) for request in wire.drain()] == [("POST", target)] + + +def _spend_rows(call_id: str) -> list[dict[str, JsonValue]]: + return read_rows('SELECT status, metadata FROM "LiteLLM_SpendLogs" WHERE request_id=%s', (call_id,)) + + +def _single_row(call_id: str) -> dict[str, JsonValue]: + rows: Final = eventually(lambda: _spend_rows(call_id), lambda values: len(values) == 1, seconds=70) + return rows[0] + + +def _failure_message(call_id: str) -> str: + row: Final = _single_row(call_id) + assert row["status"] == "failure", row + metadata: Final = row["metadata"] + parsed: Final = _JSON_OBJECT.validate_json(metadata) if isinstance(metadata, str) else object_value(metadata) + return str(object_value(parsed["error_information"])["error_message"]) + + +def _proxy_url(gateway: Gateway) -> str: + return str(gateway.client.base_url).rstrip("/") + + +def _openai_client(gateway: Gateway) -> openai.OpenAI: + return openai.OpenAI( + base_url=f"{_proxy_url(gateway)}/v1", + api_key=gateway.key, + max_retries=0, + http_client=httpx.Client(timeout=15, trust_env=False), + ) + + +def _anthropic_client(gateway: Gateway) -> anthropic.Anthropic: + return anthropic.Anthropic( + base_url=_proxy_url(gateway), + api_key=gateway.key, + max_retries=0, + http_client=httpx.Client(timeout=15, trust_env=False), + ) + + +def test_chat_completions_overflow_envelope_returns_400_prompt_too_long_and_logs_failure(gateway: Gateway) -> None: + prompt: Final = f"overflow chat {uuid4().hex}" + with wire_server(_mantle_peer(prompt, _OVERFLOW)) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=_MODEL, api_base=wire.url, api_key=_API_KEY) + response: Final = gateway.request("POST", "/v1/chat/completions", _chat(model, prompt)) + assert response.status_code == 400, response.text + assert _GENERIC in _error_message(response), response.text + _only_call(wire) + assert _GENERIC in _failure_message(response.headers["x-litellm-call-id"]) + + +def test_chat_completions_stream_overflow_envelope_returns_400_prompt_too_long(gateway: Gateway) -> None: + prompt: Final = f"overflow chat stream {uuid4().hex}" + with wire_server(_mantle_peer(prompt, _OVERFLOW)) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=_MODEL, api_base=wire.url, api_key=_API_KEY) + response, text = _consumed(gateway, "/v1/chat/completions", _chat(model, prompt, stream=True)) + assert response.status_code == 400, text + assert _GENERIC in text, text + _only_call(wire) + + +def test_messages_overflow_envelope_returns_400_invalid_request_and_logs_failure(gateway: Gateway) -> None: + prompt: Final = f"overflow messages {uuid4().hex}" + with wire_server(_mantle_peer(prompt, _OVERFLOW)) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=_MODEL, api_base=wire.url, api_key=_API_KEY) + response: Final = gateway.request("POST", "/v1/messages", _messages(model, prompt)) + assert response.status_code == 400, response.text + error: Final = _error_object(response) + assert error["type"] == "invalid_request_error", response.text + assert _GENERIC in str(error["message"]), response.text + _only_call(wire) + assert _GENERIC in _failure_message(response.headers["x-litellm-call-id"]) + + +def test_messages_stream_overflow_envelope_before_the_stream_returns_400_invalid_request(gateway: Gateway) -> None: + prompt: Final = f"overflow messages stream {uuid4().hex}" + with wire_server(_mantle_peer(prompt, _OVERFLOW)) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=_MODEL, api_base=wire.url, api_key=_API_KEY) + response, text = _consumed(gateway, "/v1/messages", _messages(model, prompt, stream=True)) + assert response.status_code == 400, text + error: Final = object_value(_JSON_OBJECT.validate_json(text)["error"]) + assert error["type"] == "invalid_request_error", text + assert _GENERIC in str(error["message"]), text + _only_call(wire) + + +def test_responses_overflow_envelope_returns_400_prompt_too_long_and_logs_failure(gateway: Gateway) -> None: + prompt: Final = f"overflow responses {uuid4().hex}" + with wire_server(_mantle_peer(prompt, _OVERFLOW)) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=_MODEL, api_base=wire.url, api_key=_API_KEY) + response: Final = gateway.request("POST", "/v1/responses", _responses(model, prompt)) + assert response.status_code == 400, response.text + assert _GENERIC in _error_message(response), response.text + _only_call(wire) + assert _GENERIC in _failure_message(response.headers["x-litellm-call-id"]) + + +def test_responses_stream_overflow_envelope_before_the_stream_returns_400_prompt_too_long(gateway: Gateway) -> None: + prompt: Final = f"overflow responses stream {uuid4().hex}" + with wire_server(_mantle_peer(prompt, _OVERFLOW)) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=_MODEL, api_base=wire.url, api_key=_API_KEY) + response, text = _consumed(gateway, "/v1/responses", _responses(model, prompt, stream=True)) + assert response.status_code == 400, text + assert _GENERIC in str(object_value(_JSON_OBJECT.validate_json(text)["error"])["message"]), text + _only_call(wire) + + +def test_openai_sdk_chat_completions_raises_bad_request_saying_prompt_too_long(gateway: Gateway) -> None: + prompt: Final = f"overflow sdk chat {uuid4().hex}" + with wire_server(_mantle_peer(prompt, _OVERFLOW)) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=_MODEL, api_base=wire.url, api_key=_API_KEY) + with pytest.raises(openai.BadRequestError) as caught: + _openai_client(gateway).chat.completions.create(model=model, messages=[{"role": "user", "content": prompt}]) + assert caught.value.status_code == 400 + assert _GENERIC in str(caught.value), str(caught.value) + _only_call(wire) + + +async def test_async_openai_sdk_chat_completions_raises_bad_request_saying_prompt_too_long( + gateway: Gateway, +) -> None: + prompt: Final = f"overflow async sdk chat {uuid4().hex}" + with wire_server(_mantle_peer(prompt, _OVERFLOW)) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=_MODEL, api_base=wire.url, api_key=_API_KEY) + async with openai.AsyncOpenAI( + base_url=f"{_proxy_url(gateway)}/v1", + api_key=gateway.key, + max_retries=0, + http_client=httpx.AsyncClient(timeout=15, trust_env=False), + ) as client: + with pytest.raises(openai.BadRequestError) as caught: + await client.chat.completions.create(model=model, messages=[{"role": "user", "content": prompt}]) + assert caught.value.status_code == 400 + assert _GENERIC in str(caught.value), str(caught.value) + _only_call(wire) + + +def test_openai_sdk_responses_raises_bad_request_saying_prompt_too_long(gateway: Gateway) -> None: + prompt: Final = f"overflow sdk responses {uuid4().hex}" + with wire_server(_mantle_peer(prompt, _OVERFLOW)) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=_MODEL, api_base=wire.url, api_key=_API_KEY) + with pytest.raises(openai.BadRequestError) as caught: + _openai_client(gateway).responses.create(model=model, input=prompt) + assert caught.value.status_code == 400 + assert _GENERIC in str(caught.value), str(caught.value) + _only_call(wire) + + +def test_anthropic_sdk_messages_raises_bad_request_saying_prompt_too_long(gateway: Gateway) -> None: + prompt: Final = f"overflow sdk messages {uuid4().hex}" + with wire_server(_mantle_peer(prompt, _OVERFLOW)) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=_MODEL, api_base=wire.url, api_key=_API_KEY) + with pytest.raises(anthropic.BadRequestError) as caught: + _anthropic_client(gateway).messages.create( + model=model, max_tokens=32, messages=[{"role": "user", "content": prompt}] + ) + assert caught.value.status_code == 400 + assert _GENERIC in str(caught.value), str(caught.value) + _only_call(wire) + + +async def test_async_anthropic_sdk_messages_raises_bad_request_saying_prompt_too_long(gateway: Gateway) -> None: + prompt: Final = f"overflow async sdk messages {uuid4().hex}" + with wire_server(_mantle_peer(prompt, _OVERFLOW)) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=_MODEL, api_base=wire.url, api_key=_API_KEY) + async with anthropic.AsyncAnthropic( + base_url=_proxy_url(gateway), + api_key=gateway.key, + max_retries=0, + http_client=httpx.AsyncClient(timeout=15, trust_env=False), + ) as client: + with pytest.raises(anthropic.BadRequestError) as caught: + await client.messages.create(model=model, max_tokens=32, messages=[{"role": "user", "content": prompt}]) + assert caught.value.status_code == 400 + assert _GENERIC in str(caught.value), str(caught.value) + _only_call(wire) + + +def test_anthropic_sdk_messages_stream_raises_invalid_request_error_saying_prompt_too_long( + gateway: Gateway, +) -> None: + prompt: Final = f"overflow sdk messages stream {uuid4().hex}" + reply: Final = _failed_stream(_OVERFLOW_STREAM_ERROR) + with wire_server(_mantle_peer(prompt, reply)) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=_MODEL, api_base=wire.url, api_key=_API_KEY) + with pytest.raises(anthropic.APIStatusError) as caught: + for _ in _anthropic_client(gateway).messages.create( + model=model, max_tokens=32, messages=[{"role": "user", "content": prompt}], stream=True + ): + pass + body: Final = object_value(_JSON_VALUE.validate_python(caught.value.body)) + error: Final = object_value(body["error"]) + assert error["type"] == "invalid_request_error", body + assert _GENERIC in str(error["message"]), body + _only_call(wire) + + +def test_chat_completions_stream_overflow_in_response_failed_event_returns_400_prompt_too_long( + gateway: Gateway, +) -> None: + prompt: Final = f"failed event chat {uuid4().hex}" + reply: Final = _failed_stream(_OVERFLOW_STREAM_ERROR) + with wire_server(_mantle_peer(prompt, reply)) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=_MODEL, api_base=wire.url, api_key=_API_KEY) + response, text = _consumed(gateway, "/v1/chat/completions", _chat(model, prompt, stream=True)) + assert response.status_code == 400, text + assert _GENERIC in text, text + _only_call(wire) + + +def test_messages_stream_overflow_in_response_failed_event_emits_invalid_request_error_event( + gateway: Gateway, +) -> None: + prompt: Final = f"failed event messages {uuid4().hex}" + reply: Final = _failed_stream(_OVERFLOW_STREAM_ERROR) + with wire_server(_mantle_peer(prompt, reply)) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=_MODEL, api_base=wire.url, api_key=_API_KEY) + response, text = _consumed(gateway, "/v1/messages", _messages(model, prompt, stream=True)) + assert response.status_code == 200, text + assert response.headers["content-type"].startswith("text/event-stream"), dict(response.headers) + error: Final = _only_error(text) + assert error["type"] == "invalid_request_error", text + assert _GENERIC in str(error["message"]), text + assert not any(event["type"] == "content_block_delta" for event in _events(text)), text + _only_call(wire) + + +def test_messages_stream_overflow_with_a_retry_budget_returns_400_invalid_request_and_logs_failure( + gateway: Gateway, +) -> None: + prompt: Final = f"held frames messages stream {uuid4().hex}" + reply: Final = _failed_stream(_OVERFLOW_STREAM_ERROR) + with wire_server(_mantle_peer(prompt, reply)) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=_MODEL, api_base=wire.url, api_key=_API_KEY) + response: Final = gateway.request( + "POST", "/v1/messages", _messages(model, prompt, stream=True), headers={"x-litellm-num-retries": "2"} + ) + assert response.status_code == 400, response.text + assert response.headers["content-type"].startswith("application/json"), dict(response.headers) + error: Final = _error_object(response) + assert error["type"] == "invalid_request_error", response.text + assert _GENERIC in str(error["message"]), response.text + _only_call(wire) + assert _GENERIC in _failure_message(response.headers["x-litellm-call-id"]) + + +def test_responses_stream_overflow_in_response_failed_event_relays_failure_saying_prompt_too_long( + gateway: Gateway, +) -> None: + prompt: Final = f"failed event responses {uuid4().hex}" + reply: Final = _failed_stream(_OVERFLOW_STREAM_ERROR) + with wire_server(_mantle_peer(prompt, reply)) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=_MODEL, api_base=wire.url, api_key=_API_KEY) + response, text = _consumed(gateway, "/v1/responses", _responses(model, prompt, stream=True)) + assert response.status_code == 200, text + assert _GENERIC in str(_only_failed(text)["message"]), text + _only_call(wire) + + +def test_chat_completions_non_overflow_400_keeps_the_upstream_message(gateway: Gateway) -> None: + prompt: Final = f"bad input chat {uuid4().hex}" + reply: Final = Reply(status=400, body=_BAD_INPUT_BODY) + with wire_server(_mantle_peer(prompt, reply)) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=_MODEL, api_base=wire.url, api_key=_API_KEY) + response: Final = gateway.request("POST", "/v1/chat/completions", _chat(model, prompt)) + assert response.status_code == 400, response.text + assert _INVALID_INPUT_MESSAGE in response.text, response.text + assert _TOO_LONG not in response.text, response.text + _only_call(wire) + + +def test_messages_stream_non_overflow_400_before_the_stream_returns_400_with_the_upstream_message( + gateway: Gateway, +) -> None: + prompt: Final = f"bad input messages stream {uuid4().hex}" + reply: Final = Reply(status=400, body=_BAD_INPUT_BODY) + with wire_server(_mantle_peer(prompt, reply)) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=_MODEL, api_base=wire.url, api_key=_API_KEY) + response, text = _consumed(gateway, "/v1/messages", _messages(model, prompt, stream=True)) + assert response.status_code == 400, text + assert _INVALID_INPUT_MESSAGE in text, text + assert _TOO_LONG not in text, text + _only_call(wire) + + +def test_messages_stream_non_overflow_response_failed_event_emits_the_upstream_invalid_request_error_event( + gateway: Gateway, +) -> None: + prompt: Final = f"invalid prompt messages stream {uuid4().hex}" + reply: Final = _failed_stream({"code": "invalid_prompt", "message": _INVALID_PROMPT_MESSAGE}) + with wire_server(_mantle_peer(prompt, reply)) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=_MODEL, api_base=wire.url, api_key=_API_KEY) + response, text = _consumed(gateway, "/v1/messages", _messages(model, prompt, stream=True)) + assert response.status_code == 200, text + assert response.headers["content-type"].startswith("text/event-stream"), dict(response.headers) + error: Final = _only_error(text) + assert error["type"] == "invalid_request_error", text + assert _INVALID_PROMPT_MESSAGE in str(error["message"]), text + assert _TOO_LONG not in text, text + _only_call(wire) + + +@pytest.mark.parametrize( + ("status", "body", "expected"), + ( + (401, b'{"error":{"message":"invalid bearer","type":"authentication_error","code":null}}', 401), + (429, b'{"error":{"message":"slow down","type":"rate_limit_error","code":"rate_limit_exceeded"}}', 429), + (500, b'{"error":{"message":"boom","type":"server_error","code":null}}', 503), + ( + 404, + b'{"error":{"message":"The model `x` does not exist","type":"invalid_request_error","code":"model_not_found"}}', + 404, + ), + ), + ids=("401", "429", "500", "404"), +) +def test_chat_completions_other_upstream_statuses_keep_their_mapping( + gateway: Gateway, status: int, body: bytes, expected: int +) -> None: + prompt: Final = f"status {status} chat {uuid4().hex}" + with wire_server(_mantle_peer(prompt, Reply(status=status, body=body))) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=_MODEL, api_base=wire.url, api_key=_API_KEY) + response: Final = gateway.request("POST", "/v1/chat/completions", _chat(model, prompt)) + assert response.status_code == expected, response.text + assert _TOO_LONG not in response.text, response.text + _only_call(wire) + + +def test_messages_stream_legacy_token_count_envelope_returns_400_with_the_counts(gateway: Gateway) -> None: + prompt: Final = f"legacy messages stream {uuid4().hex}" + with ( + wire_server(_mantle_peer(prompt, Reply(status=400, body=_LEGACY_BODY))) as wire, + gateway.scenario() as scenario, + ): + model: Final = scenario.model(model=_MODEL, api_base=wire.url, api_key=_API_KEY) + response, text = _consumed(gateway, "/v1/messages", _messages(model, prompt, stream=True)) + assert response.status_code == 400, text + error: Final = object_value(_JSON_OBJECT.validate_json(text)["error"]) + assert error["type"] == "invalid_request_error", text + assert f"prompt is too long: {_LEGACY_PROMPT_TOKENS} tokens > {_LEGACY_MODEL_MAXIMUM} maximum" in str( + error["message"] + ), text + _only_call(wire) + + +def test_chat_completions_code_only_overflow_envelope_returns_400_prompt_too_long(gateway: Gateway) -> None: + prompt: Final = f"code only chat {uuid4().hex}" + reply: Final = Reply(status=400, body=b'{"error":{"code":"context_length_exceeded"}}') + with wire_server(_mantle_peer(prompt, reply)) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=_MODEL, api_base=wire.url, api_key=_API_KEY) + response: Final = gateway.request("POST", "/v1/chat/completions", _chat(model, prompt)) + assert response.status_code == 400, response.text + assert _GENERIC in _error_message(response), response.text + _only_call(wire) + + +def test_chat_completions_text_plain_body_naming_the_overflow_code_returns_the_plain_400(gateway: Gateway) -> None: + prompt: Final = f"text plain chat {uuid4().hex}" + reply: Final = Reply(status=400, body=b"request rejected: context_length_exceeded", content_type="text/plain") + with wire_server(_mantle_peer(prompt, reply)) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=_MODEL, api_base=wire.url, api_key=_API_KEY) + response: Final = gateway.request("POST", "/v1/chat/completions", _chat(model, prompt)) + assert response.status_code == 400, response.text + assert _TOO_LONG not in _error_message(response), response.text + assert "request rejected: context_length_exceeded" in _error_message(response), response.text + _only_call(wire) + + +def test_chat_completions_400_merely_quoting_the_overflow_code_keeps_the_upstream_message(gateway: Gateway) -> None: + prompt: Final = f"quoted code chat {uuid4().hex}" + upstream_message: Final = "Invalid input value context_length_exceeded" + reply: Final = Reply(status=400, body=_envelope("invalid_request_error", upstream_message)) + with wire_server(_mantle_peer(prompt, reply)) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=_MODEL, api_base=wire.url, api_key=_API_KEY) + response: Final = gateway.request("POST", "/v1/chat/completions", _chat(model, prompt)) + assert response.status_code == 400, response.text + assert _TOO_LONG not in _error_message(response), response.text + assert upstream_message in _error_message(response), response.text + _only_call(wire) + + +def test_chat_completions_five_kilobyte_overflow_message_returns_400_and_keeps_the_proxy_alive( + gateway: Gateway, +) -> None: + prompt: Final = f"large envelope chat {uuid4().hex}" + reply: Final = Reply(status=400, body=_envelope("context_length_exceeded", "x" * 5120)) + with wire_server(_mantle_peer(prompt, reply)) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=_MODEL, api_base=wire.url, api_key=_API_KEY) + response: Final = gateway.request("POST", "/v1/chat/completions", _chat(model, prompt)) + assert response.status_code == 400, response.text + assert _GENERIC in _error_message(response), response.text + _only_call(wire) + assert gateway.request("GET", "/health/liveliness").status_code == 200 + + +def test_chat_completions_stream_response_failed_without_error_fails_and_keeps_the_proxy_alive( + gateway: Gateway, +) -> None: + prompt: Final = f"null error chat stream {uuid4().hex}" + with wire_server(_mantle_peer(prompt, _failed_stream(None))) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=_MODEL, api_base=wire.url, api_key=_API_KEY) + response, text = _consumed(gateway, "/v1/chat/completions", _chat(model, prompt, stream=True)) + assert response.status_code >= 400, text + assert _TOO_LONG not in text, text + _only_call(wire) + assert gateway.request("GET", "/health/liveliness").status_code == 200 + + +def test_chat_completions_repeated_overflow_logs_one_failure_row_per_call(gateway: Gateway) -> None: + prompt: Final = f"repeated overflow chat {uuid4().hex}" + with wire_server(_mantle_peer(prompt, _OVERFLOW)) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=_MODEL, api_base=wire.url, api_key=_API_KEY) + responses: Final = tuple( + gateway.request("POST", "/v1/chat/completions", _chat(model, prompt)) for _ in range(2) + ) + call_ids: Final = tuple(response.headers["x-litellm-call-id"] for response in responses) + assert len(set(call_ids)) == 2, call_ids + assert [request.target for request in wire.drain()] == [_RESPONSES_PATH, _RESPONSES_PATH] + for response, call_id in zip(responses, call_ids, strict=True): + assert response.status_code == 400, response.text + assert _GENERIC in _failure_message(call_id) + + +def test_chat_completions_prompt_naming_the_error_code_still_succeeds(gateway: Gateway) -> None: + prompt: Final = f"my prompt mentions context_length_exceeded {uuid4().hex}" + with wire_server(_mantle_peer(prompt, _happy_reply(stream=False))) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=_MODEL, api_base=wire.url, api_key=_API_KEY) + response: Final = gateway.request("POST", "/v1/chat/completions", _chat(model, prompt)) + assert response.status_code == 200, response.text + assert _HAPPY_TEXT in response.text, response.text + _only_call(wire) + + +def test_messages_stream_openai_deployment_overflow_in_the_stream_emits_invalid_request_error_event( + gateway: Gateway, +) -> None: + prompt: Final = f"openai overflow messages stream {uuid4().hex}" + with wire_server(_openai_peer(prompt)) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=_OPENAI_MODEL, api_base=f"{wire.url}/v1", api_key=_OPENAI_API_KEY) + response, text = _consumed(gateway, "/v1/messages", _messages(model, prompt, stream=True)) + assert response.status_code == 200, text + error: Final = _only_error(text) + assert error["type"] == "invalid_request_error", text + assert _OPENAI_OVERFLOW_MARK in str(error["message"]), text + _only_call(wire, _OPENAI_RESPONSES_PATH) + + +def test_chat_completions_stream_openai_deployment_overflow_returns_400(gateway: Gateway) -> None: + prompt: Final = f"openai overflow chat stream {uuid4().hex}" + with wire_server(_openai_peer(prompt)) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=_OPENAI_MODEL, api_base=f"{wire.url}/v1", api_key=_OPENAI_API_KEY) + response, text = _consumed(gateway, "/v1/chat/completions", _chat(model, prompt, stream=True)) + assert response.status_code == 400, text + assert _OPENAI_OVERFLOW_MARK in text, text + _only_call(wire, _OPENAI_CHAT_PATH) + + +def test_messages_openai_deployment_overflow_returns_400_invalid_request(gateway: Gateway) -> None: + prompt: Final = f"openai overflow messages {uuid4().hex}" + with wire_server(_openai_peer(prompt)) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=_OPENAI_MODEL, api_base=f"{wire.url}/v1", api_key=_OPENAI_API_KEY) + response: Final = gateway.request("POST", "/v1/messages", _messages(model, prompt)) + assert response.status_code == 400, response.text + error: Final = _error_object(response) + assert error["type"] == "invalid_request_error", response.text + assert _OPENAI_OVERFLOW_MARK in str(error["message"]), response.text + _only_call(wire, _OPENAI_RESPONSES_PATH) + + +_OVERFLOW_MARKER: Final = "chaos-overflow" +_HAPPY_MARKER: Final = "chaos-happy" +_Call = tuple[str, dict[str, JsonValue], bool, str] +_Outcome = tuple[str, bool, bool, int, str, str] +_Builder = Callable[[str, str, bool], dict[str, JsonValue]] +_Cell = tuple[tuple[str, _Builder], str, bool] +_BUILDERS: Final = (("/v1/chat/completions", _chat), ("/v1/messages", _messages), ("/v1/responses", _responses)) + + +def _chaos_peer(request: Request) -> Reply: + body: Final = _JSON_OBJECT.validate_json(request.body) + stream: Final = body.get("stream") is True + if _OVERFLOW_MARKER in json.dumps(body["input"]): + return _failed_stream(_OVERFLOW_STREAM_ERROR) if stream else _OVERFLOW + return _happy_reply(stream) + + +def _burst_bodies(model: str, round_name: str) -> tuple[_Call, ...]: + def call(cell: _Cell) -> _Call: + (path, build), marker, stream = cell + tag: Final = f"chaos-{round_name}-{uuid4().hex}" + prompt: Final = f"{marker} {round_name} {path} stream={stream} {tag}" + return path, build(model, prompt, stream), marker == _OVERFLOW_MARKER, tag + + return tuple(call(cell) for cell in product(_BUILDERS, (_HAPPY_MARKER, _OVERFLOW_MARKER), (False, True))) + + +def _tagged_rows(tag: str) -> list[dict[str, JsonValue]]: + return read_rows( + 'SELECT status FROM "LiteLLM_SpendLogs" WHERE request_tags::jsonb @> %s::jsonb', (json.dumps([tag]),) + ) + + +def _single_tagged_status(tag: str) -> str: + rows: Final = eventually(lambda: _tagged_rows(tag), lambda values: len(values) == 1, seconds=70) + return str(rows[0]["status"]) + + +def _fire(gateway: Gateway, bodies: tuple[_Call, ...]) -> tuple[_Outcome, ...]: + def one(item: _Call) -> _Outcome: + path, body, overflow, tag = item + with gateway.client.stream( + "POST", path, json=body, headers={"Authorization": f"Bearer {gateway.key}", "x-litellm-tags": tag} + ) as response: + text: Final = b"".join(response.iter_bytes()).decode() + return path, overflow, body["stream"] is True, response.status_code, tag, text + + with ThreadPoolExecutor(max_workers=len(bodies)) as pool: + return tuple(pool.map(one, bodies)) + + +def _assert_served(outcomes: tuple[_Outcome, ...]) -> None: + for path, overflow, _, status, tag, text in outcomes: + if not overflow: + assert status == 200, (path, text) + assert _HAPPY_TEXT in text, (path, text) + assert _single_tagged_status(tag) == "success", (path, tag) + continue + assert _GENERIC in text, (path, status, text) + assert status in (200, 400), (path, status, text) + assert _single_tagged_status(tag) == "failure", (path, tag) + + +def test_messages_stream_overflow_logs_one_failure_row(gateway: Gateway) -> None: + with wire_server(_chaos_peer) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=_MODEL, api_base=wire.url, api_key=_API_KEY) + tag: Final = f"chaos-single-{uuid4().hex}" + body: Final = _messages(model, f"{_OVERFLOW_MARKER} single {tag}", True) + ((_, _, _, status, _, text),) = _fire(gateway, (("/v1/messages", body, True, tag),)) + assert status == 200, text + assert _single_tagged_status(tag) == "failure", tag + assert _GENERIC in text, text + + +def test_chaos_mantle_peer_outage_mid_burst_logs_every_call_once_and_keeps_the_proxy_alive(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + with wire_server(_chaos_peer) as wire: + port: Final = int(wire.url.rsplit(":", 1)[1]) + model: Final = scenario.model(model=_MODEL, api_base=wire.url, api_key=_API_KEY) + before: Final = _fire(gateway, _burst_bodies(model, "before")) + assert len(wire.drain()) == len(before) + during: Final = _fire(gateway, _burst_bodies(model, "during")) + assert gateway.request("GET", "/health/liveliness").status_code == 200 + with wire_server(_chaos_peer, port=port) as revived: + after: Final = _fire(gateway, _burst_bodies(model, "after")) + assert len(revived.drain()) == len(after) + _assert_served(before) + _assert_served(after) + for path, _, _, status, tag, text in during: + assert status >= 500 or _error_events(text) or '"response.failed"' in text, (path, status, text) + assert _TOO_LONG not in text, (path, text) + assert _single_tagged_status(tag) == "failure", (path, tag) diff --git a/tests/integration/providers/test_bedrock_mantle_count_tokens_wire.py b/tests/integration/providers/test_bedrock_mantle_count_tokens_wire.py index 933fc49c149..7f59b818cba 100644 --- a/tests/integration/providers/test_bedrock_mantle_count_tokens_wire.py +++ b/tests/integration/providers/test_bedrock_mantle_count_tokens_wire.py @@ -1,7 +1,6 @@ import asyncio import json import re -import signal import socket import threading import uuid @@ -1270,10 +1269,11 @@ def test_worker_sigkill_mid_burst_leaves_the_sibling_counting( shares: Final = {pid: _accepted_client_ports(pid, proxy_url.port or 0) & ports for pid in workers} assert sum(map(len, shares.values())) == len(clients), shares victim: Final = min((pid for pid in workers if shares[pid]), key=lambda pid: len(shares[pid])) - psutil.Process(victim).send_signal(signal.SIGKILL) - release.set() - for port, result in (future.result(timeout=60) for future in futures): - assert result == (None if port in shares[victim] else (200, _MANTLE_COUNT)), (port, result, shares) + with owned.replacing_worker(victim) as kill: + kill() + release.set() + for port, result in (future.result(timeout=60) for future in futures): + assert result == (None if port in shares[victim] else (200, _MANTLE_COUNT)), (port, result, shares) second_wave: Final = _clients(stack, str(proxy_url), 6) assert tuple(_counted_on(client, owned.gateway.key, body) for client in second_wave) == ( (200, _MANTLE_COUNT), diff --git a/tests/integration/providers/test_bedrock_mantle_output_limit_wire.py b/tests/integration/providers/test_bedrock_mantle_output_limit_wire.py index 7891522a26f..4c06a70070e 100644 --- a/tests/integration/providers/test_bedrock_mantle_output_limit_wire.py +++ b/tests/integration/providers/test_bedrock_mantle_output_limit_wire.py @@ -3,7 +3,6 @@ import base64 import json import os import re -import signal import threading import uuid from collections.abc import Callable, Mapping @@ -939,7 +938,7 @@ def _chaos_budget(boots: int, stops: int) -> float: return (boots + stops) * graceful_stop_seconds() + _CHAOS_TRAFFIC_SECONDS -@pytest.mark.timeout(_chaos_budget(boots=2, stops=1)) +@pytest.mark.timeout(_chaos_budget(boots=3, stops=1)) async def test_worker_sigkill_mid_burst_leaves_the_sibling_renaming_the_yaml_limit( gateway: Gateway, tmp_path: Path ) -> None: @@ -969,18 +968,18 @@ async def test_worker_sigkill_mid_burst_leaves_the_sibling_renaming_the_yaml_lim held_by: Final = MappingProxyType({pid: _open_upstream_connections(pid, wire.url) for pid in workers}) assert sum(held_by.values()) == 20, held_by victim_pid, survivor_pid = sorted(workers, key=held_by.__getitem__) - victim: Final = psutil.Process(victim_pid) - victim.suspend() - victim.send_signal(signal.SIGKILL) - release.set() - outcome: Final = await burst - assert held_by[survivor_pid] >= 10, held_by - assert (len(outcome.served), outcome.lost) == (held_by[survivor_pid], held_by[victim_pid]), ( - held_by, - outcome, - ) - for item in outcome.served: - _assert_answered_with_its_own_marker(item) + with owned.replacing_worker(victim_pid) as kill: + kill.freeze() + await asyncio.to_thread(kill) + release.set() + outcome: Final = await burst + assert held_by[survivor_pid] >= 10, held_by + assert (len(outcome.served), outcome.lost) == (held_by[survivor_pid], held_by[victim_pid]), ( + held_by, + outcome, + ) + for item in outcome.served: + _assert_answered_with_its_own_marker(item) follow_up: Final = _Call(endpoint="chat", stream=False, marker=uuid.uuid4().hex) (answered,) = ( await _burst(_proxy_url(candidate), candidate.key, _CONFIG_MODEL, (follow_up,), limit=None) @@ -1011,7 +1010,7 @@ async def test_proxy_restart_mid_burst_loses_only_in_flight_requests_and_renames ) ) await asyncio.to_thread(eventually, held_markers.qsize, lambda size: size == 20, 60) - first.process.terminate() + first.begin_shutdown() release.set() outcome: Final = await burst assert len(outcome.served) + outcome.lost == len(calls), outcome diff --git a/tests/integration/providers/test_bedrock_model_discovery_chaos.py b/tests/integration/providers/test_bedrock_model_discovery_chaos.py index f1358837b7e..34350e5846c 100644 --- a/tests/integration/providers/test_bedrock_model_discovery_chaos.py +++ b/tests/integration/providers/test_bedrock_model_discovery_chaos.py @@ -1,6 +1,5 @@ from __future__ import annotations -import signal import threading from collections.abc import Iterator from concurrent.futures import ThreadPoolExecutor @@ -32,7 +31,7 @@ from integration._support.bedrock_discovery import ( wait_for_replacement_worker, ) from integration._support.client import Gateway, eventually, gateway_from_environment -from integration._support.process import graceful_stop_seconds +from integration._support.process import OwnedProxy, graceful_stop_seconds from integration._support.wire import Reply pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120) @@ -43,6 +42,7 @@ class Rig: gateway: Gateway log: Path plane: ControlPlane + owned: OwnedProxy @pytest.fixture(scope="module") @@ -51,7 +51,7 @@ def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]: directory: Final = tmp_path_factory.mktemp("bedrock_discovery_chaos") with control_plane(directory, hosts) as plane, gateway_from_environment() as parent: with discovery_proxy(parent, directory, plane, check_provider_endpoint=True, workers=WORKERS) as owned: - yield Rig(owned.gateway, owned.log, plane) + yield Rig(owned.gateway, owned.log, plane, owned) def _victim(pids: tuple[int, ...], plane_url: str) -> int | None: @@ -85,13 +85,14 @@ def test_worker_killed_mid_listing_leaves_the_sibling_serving_and_a_replacement_ seconds=LISTING_TIMEOUT_SECONDS, ) assert victim is not None - psutil.Process(victim).send_signal(signal.SIGKILL) - sibling: Final = next(pid for pid in original if pid != victim) - for _ in range(5): - assert rig.gateway.chat(CONTROL_MODEL)["choices"], "the sibling worker stopped answering" - assert psutil.pid_exists(sibling) - severed: Final = in_flight.exception(timeout=LISTING_TIMEOUT_SECONDS * 2) - assert severed is not None or in_flight.result().status_code >= 500 + with rig.owned.replacing_worker(victim) as kill: + kill() + sibling: Final = next(pid for pid in original if pid != victim) + for _ in range(5): + assert rig.gateway.chat(CONTROL_MODEL)["choices"], "the sibling worker stopped answering" + assert psutil.pid_exists(sibling) + severed: Final = in_flight.exception(timeout=LISTING_TIMEOUT_SECONDS * 2) + assert severed is not None or in_flight.result().status_code >= 500 wait_for_replacement_worker(rig.log, original) release.set() assert discovered(rig.gateway, catalog, marker) == catalog.invocable_ids() diff --git a/tests/integration/providers/test_bedrock_runtime_chat_completions_chaos.py b/tests/integration/providers/test_bedrock_runtime_chat_completions_chaos.py index 820d033135f..71b63786c27 100644 --- a/tests/integration/providers/test_bedrock_runtime_chat_completions_chaos.py +++ b/tests/integration/providers/test_bedrock_runtime_chat_completions_chaos.py @@ -7,7 +7,6 @@ import itertools import multiprocessing import os import re -import signal import socket import threading import uuid @@ -29,7 +28,7 @@ import yaml from integration._support.bedrock_runtime_peer import MARKER, marker_of, respond, serve_peer from integration._support.client import Gateway, eventually, gateway_from_environment, object_value from integration._support.database import read_rows -from integration._support.process import owned_proxy_process +from integration._support.process import kill_stray_process, owned_proxy_process from integration._support.wire import Reply, Request, wire_server from pydantic import JsonValue, TypeAdapter @@ -216,7 +215,7 @@ async def _burst_killing_the_peer_once_it_answered( await asyncio.to_thread(eventually, lambda: peer.received.value, lambda count: count == len(calls), 60) first: Final = [await finished for finished in itertools.islice(asyncio.as_completed(tasks), answered)] assert all(item.status == 200 for item in first), [(item.call.marker, item.status) for item in first] - peer.process.kill() + kill_stray_process(peer.process.pid) peer.process.join(timeout=10) return tuple(await asyncio.gather(*tasks)) @@ -254,7 +253,8 @@ def _child_peer(port: int, answer_first: int) -> Iterator[_ChildPeer]: eventually(lambda: _accepts_connections(port), bool, seconds=30) yield _ChildPeer(process=process, received=received, url=f"http://127.0.0.1:{port}") finally: - process.kill() + if process.is_alive(): + kill_stray_process(process.pid) process.join(timeout=10) assert not process.is_alive(), "Owned peer survived cleanup" @@ -428,15 +428,15 @@ async def test_worker_sigkill_mid_burst_leaves_the_sibling_serving(gateway: Gate held_by: Final = MappingProxyType({pid: _open_upstream_connections(pid, wire.url) for pid in workers}) assert sum(held_by.values()) == 20, held_by victim_pid, survivor_pid = sorted(workers, key=held_by.__getitem__) - victim: Final = psutil.Process(victim_pid) - victim.suspend() - victim.send_signal(signal.SIGKILL) - release.set() - served: Final = await burst - assert held_by[survivor_pid] >= 10, held_by - assert len(served) == held_by[survivor_pid], (held_by, len(served)) - for item in served: - _assert_answered_with_its_own_marker(item) + with owned.replacing_worker(victim_pid) as kill: + kill.freeze() + await asyncio.to_thread(kill) + release.set() + served: Final = await burst + assert held_by[survivor_pid] >= 10, held_by + assert len(served) == held_by[survivor_pid], (held_by, len(served)) + for item in served: + _assert_answered_with_its_own_marker(item) follow_up: Final = _Call(endpoint="chat", stream=False, marker=uuid.uuid4().hex) (answered,) = await _burst(str(candidate.client.base_url), candidate.key, _CONFIG_MODEL, (follow_up,)) _assert_answered_with_its_own_marker(answered) diff --git a/tests/integration/providers/test_bedrock_runtime_chat_completions_sad_wire.py b/tests/integration/providers/test_bedrock_runtime_chat_completions_sad_wire.py index 8d6193cd542..0860d79459e 100644 --- a/tests/integration/providers/test_bedrock_runtime_chat_completions_sad_wire.py +++ b/tests/integration/providers/test_bedrock_runtime_chat_completions_sad_wire.py @@ -17,6 +17,7 @@ from integration._support.bedrock_runtime_peer import answer, forwarded_effort, from integration._support.client import Gateway, Scenario, eventually, object_value, string_value from integration._support.database import read_rows from integration._support.process import owned_proxy_process +from integration._support.timing import assert_within_factor_of_control, timed from integration._support.wire import Request, Wire, wire_server from pydantic import JsonValue, TypeAdapter @@ -255,16 +256,21 @@ def test_model_group_info_lists_the_native_supported_params(gateway: Gateway) -> def test_thirty_thousand_digit_version_is_classified_quickly_and_served_by_converse(gateway: Gateway) -> None: marker: Final = uuid.uuid4().hex + control_marker: Final = uuid.uuid4().hex with wire_server(respond) as wire, gateway.scenario() as scenario, ThreadPoolExecutor(max_workers=1) as pool: + control_model: Final = _deployment(scenario, wire) model: Final = _deployment(scenario, wire, model=f"bedrock/{LONG_VERSION_GPT}") + control: Final = timed(lambda: _chat(gateway, control_model, control_marker)) + assert _content(control.value) == answer(control_marker), control.value.text + _only_request(wire, control_marker) liveliness: Final = pool.submit(_timed_liveliness, gateway) - started: Final = time.monotonic() - response: Final = _chat(gateway, model, marker) - elapsed: Final = time.monotonic() - started + served: Final = timed(lambda: _chat(gateway, model, marker)) + response: Final = served.value health_status, health_elapsed = liveliness.result() assert _content(response) == answer(marker), response.text - assert elapsed < 10, elapsed - assert (health_status, health_elapsed < 2) == (200, True), (health_status, health_elapsed) + assert_within_factor_of_control(served.seconds, control=control.seconds, factor=20) + assert health_status == 200, (health_status, health_elapsed) + assert_within_factor_of_control(health_elapsed, control=control.seconds, factor=20) request: Final = _only_request(wire, marker) assert (request.method, target_of(request)) == ("POST", f"/model/{LONG_VERSION_GPT}/converse"), request.target _assert_row(string_value(_payload(response)["id"]), model, "success") @@ -274,12 +280,14 @@ def test_bad_key_on_the_long_version_model_is_refused_before_any_route(gateway: marker: Final = uuid.uuid4().hex control_marker: Final = uuid.uuid4().hex with wire_server(respond) as wire, gateway.scenario() as scenario: + short_model: Final = _deployment(scenario, wire) model: Final = _deployment(scenario, wire, model=f"bedrock/{LONG_VERSION_GPT}") - started: Final = time.monotonic() - refused: Final = _chat(gateway, model, marker, key=BAD_KEY) - elapsed: Final = time.monotonic() - started + short_refusal: Final = timed(lambda: _chat(gateway, short_model, uuid.uuid4().hex, key=BAD_KEY)) + assert short_refusal.value.status_code == 401, short_refusal.value.text + long_refusal: Final = timed(lambda: _chat(gateway, model, marker, key=BAD_KEY)) + refused: Final = long_refusal.value assert refused.status_code == 401, refused.text - assert elapsed < 2, elapsed + assert_within_factor_of_control(long_refusal.seconds, control=short_refusal.seconds, factor=10) assert "Authentication Error" in _error_message(refused), refused.text refused_rows: Final = eventually( lambda: read_rows( diff --git a/tests/integration/providers/test_bedrock_stream_timeout_chaos.py b/tests/integration/providers/test_bedrock_stream_timeout_chaos.py index 2632b81b4dc..496b83e0b2e 100644 --- a/tests/integration/providers/test_bedrock_stream_timeout_chaos.py +++ b/tests/integration/providers/test_bedrock_stream_timeout_chaos.py @@ -1,7 +1,6 @@ import asyncio import json import re -import signal import threading import uuid from collections.abc import Iterator, Mapping @@ -353,14 +352,14 @@ async def test_x3_killing_one_worker_mid_burst_leaves_the_sibling_timing_its_str held_by: Final = {pid: _open_upstream_connections(pid, peer.wire.url) for pid in workers} assert sum(held_by.values()) == 20, held_by victim_pid, survivor_pid = sorted(workers, key=held_by.__getitem__) - victim: Final = psutil.Process(victim_pid) - victim.suspend() - victim.send_signal(signal.SIGKILL) - served: Final = await burst - assert held_by[survivor_pid] >= 10, held_by - assert len(served) == held_by[survivor_pid], (held_by, len(served)) - for item in served: - _assert_timed_out(item, _KILL_TIMEOUT_SECONDS) + with owned.replacing_worker(victim_pid) as kill: + kill.freeze() + await asyncio.to_thread(kill) + served: Final = await burst + assert held_by[survivor_pid] >= 10, held_by + assert len(served) == held_by[survivor_pid], (held_by, len(served)) + for item in served: + _assert_timed_out(item, _KILL_TIMEOUT_SECONDS) peer.hold.clear() _assert_answered(await _one(owned.gateway, _KILL_MODEL, _Call("chat"))) assert len(peer.wire.drain()) == 21 diff --git a/tests/integration/providers/test_cache_control_tool_call_marks_owned_proxy.py b/tests/integration/providers/test_cache_control_tool_call_marks_owned_proxy.py index c945eaff38e..bc0a6f31e74 100644 --- a/tests/integration/providers/test_cache_control_tool_call_marks_owned_proxy.py +++ b/tests/integration/providers/test_cache_control_tool_call_marks_owned_proxy.py @@ -1,7 +1,5 @@ import asyncio -import os import re -import signal import socket import threading from collections import Counter @@ -209,10 +207,12 @@ async def test_worker_sigkill_mid_burst_leaves_the_sibling_serving_capped_reques {pid: _held_upstream_connections(pid, wire.url) for pid in workers} ) victim: Final = max(workers, key=held_by.__getitem__) - os.kill(victim, signal.SIGKILL) + with owned.replacing_worker(victim) as kill: + await asyncio.to_thread(kill) + release.set() + served: Final = await burst finally: release.set() - served: Final = await burst during: Final = wire.drain() after: Final = await _fire(owned_url, owned.gateway.key) after_received: Final = wire.drain() @@ -277,7 +277,7 @@ async def test_proxy_restart_mid_burst_serves_capped_requests_after_the_reboot( ) try: await asyncio.to_thread(eventually, lambda: wire.received.qsize(), lambda size: size >= _BURST, 90) - first.process.terminate() + first.begin_shutdown() finally: release.set() served: Final = await burst diff --git a/tests/integration/providers/test_chatgpt_responses_caller_stream_chaos.py b/tests/integration/providers/test_chatgpt_responses_caller_stream_chaos.py index d1e199c949a..40b87ba1b8b 100644 --- a/tests/integration/providers/test_chatgpt_responses_caller_stream_chaos.py +++ b/tests/integration/providers/test_chatgpt_responses_caller_stream_chaos.py @@ -1,6 +1,5 @@ import asyncio import re -import signal import threading import uuid from collections.abc import Iterator, Mapping, Sequence @@ -304,15 +303,15 @@ async def test_worker_sigkill_mid_burst_leaves_the_sibling_answering_json(gatewa held_by: Final = MappingProxyType({pid: _open_upstream_connections(pid, wire.url) for pid in workers}) assert sum(held_by.values()) == 20, held_by victim_pid, survivor_pid = sorted(workers, key=held_by.__getitem__) - victim: Final = psutil.Process(victim_pid) - victim.suspend() - victim.send_signal(signal.SIGKILL) - release.set() - served: Final = await burst - assert held_by[survivor_pid] >= 10, held_by - assert len(served) == held_by[survivor_pid], (held_by, len(served)) - for item in served: - _assert_answered_in_its_own_shape(item) + with owned.replacing_worker(victim_pid) as kill: + kill.freeze() + await asyncio.to_thread(kill) + release.set() + served: Final = await burst + assert held_by[survivor_pid] >= 10, held_by + assert len(served) == held_by[survivor_pid], (held_by, len(served)) + for item in served: + _assert_answered_in_its_own_shape(item) follow_up: Final = _Call("responses", False, uuid.uuid4().hex) (answered,) = await _burst(candidate, _CONFIG_MODEL, (follow_up,)) _assert_answered_in_its_own_shape(answered) diff --git a/tests/integration/providers/test_cloudflare_clef_omni_media_chaos.py b/tests/integration/providers/test_cloudflare_clef_omni_media_chaos.py new file mode 100644 index 00000000000..a1cb2a84d71 --- /dev/null +++ b/tests/integration/providers/test_cloudflare_clef_omni_media_chaos.py @@ -0,0 +1,293 @@ +import asyncio +import json +import re +import socket +import threading +import uuid +from dataclasses import dataclass +from pathlib import Path +from queue import SimpleQueue +from types import MappingProxyType +from typing import Final +from urllib.parse import unquote, urlsplit + +import httpx +import psutil +import pytest +import yaml +from integration._support.client import Gateway, eventually, object_value, string_value +from integration._support.database import read_rows +from integration._support.process import graceful_stop_seconds, owned_proxy_process +from integration._support.wire import Reply, Request, Wire, wire_server +from pydantic import JsonValue, TypeAdapter + +_API_KEY: Final = "synthetic-cloudflare-key" +_CANONICAL: Final = "cloudflare/@cf/cloudflare/clef-omni" +_CONFIG_MODEL: Final = "clef-omni-media-chaos" +_UPSTREAM_PATH: Final = "/ai/run/@cf/cloudflare/clef-omni" +_WAV: Final = "data:audio/wav;base64,UklGRiQAAABXQVZFZm10IBAAAAABAAEARKwAAIhYAQACABAAZGF0YQAAAAA=" +_MP4: Final = "data:video/mp4;base64,AAAAHGZ0eXBpc29tAAACAGlzb21pc28ybXA0MQ==" +_INPUT_TOKENS: Final = 537 +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +_MARKER: Final = re.compile(r"marker-([0-9a-f]{32})") +_STARTED_WORKER: Final = re.compile(r"Started server process \[(\d+)\]") +_QUESTIONS: Final[dict[str, JsonValue]] = { + "pineapple": {"type": "noul", "instructions": "Does the audio say pineapple?"}, + "red": {"type": "noul", "instructions": "Is the video red?"}, +} +_RED: Final = 0.97 +_CHAOS_TIMEOUT: Final = int(2 * graceful_stop_seconds() + 120) + + +@dataclass(frozen=True, slots=True) +class _Served: + marker: str + call_id: str + status: int + text: str + + +def _marker_of(request: Request) -> str: + found: Final = _MARKER.search(request.body.decode()) + assert found is not None, request.body[:200] + return found.group(1) + + +def _state(marker: str) -> str: + return f"Does the clip say pineapple? marker-{marker}" + + +def _probability(marker: str) -> float: + return round(int(marker[:6], 16) / 0x1000000, 7) + + +def _echo(request: Request) -> Reply: + marker: Final = _marker_of(request) + reply: Final[dict[str, JsonValue]] = { + "result": { + "model": "clef-omni", + "answers": { + "pineapple": {"type": "noul", "noul": _probability(marker)}, + "red": {"type": "noul", "noul": _RED}, + }, + "usage": {"input_tokens": _INPUT_TOKENS, "output_tokens": 0}, + }, + "success": True, + "errors": [], + "messages": [], + } + return Reply(body=json.dumps(reply).encode()) + + +def _body(model: str, marker: str) -> dict[str, JsonValue]: + return { + "model": model, + "state": _state(marker), + "questions": _QUESTIONS, + "audio": [_WAV], + "videos": [_MP4], + "num_retries": 0, + } + + +def _expected_upstream_body(marker: str) -> dict[str, JsonValue]: + return {"model": "clef-omni", "state": _state(marker), "audio": [_WAV], "videos": [_MP4], "questions": _QUESTIONS} + + +def _upstream_body(request: Request) -> dict[str, JsonValue]: + assert request.method == "POST", request.method + assert unquote(request.target) == _UPSTREAM_PATH, request.target + assert request.headers["authorization"] == f"Bearer {_API_KEY}", request.headers + return _JSON_OBJECT.validate_json(request.body) + + +def _assert_each_landed_once_with_media(received: tuple[Request, ...], markers: frozenset[str]) -> None: + bodies: Final = {_marker_of(request): _upstream_body(request) for request in received} + assert len(bodies) == len(received) == len(markers), (len(bodies), len(received), len(markers)) + for marker in markers: + assert bodies[marker] == _expected_upstream_body(marker), marker + + +def _assert_answered_with_its_own_probability(served: _Served) -> None: + assert served.status == 200, served.text + answers: Final = object_value(_JSON_OBJECT.validate_json(served.text)["answers"]) + assert object_value(answers["pineapple"])["noul"] == _probability(served.marker), served.text + + +async def _send(client: httpx.AsyncClient, key: str, model: str, marker: str) -> _Served: + response: Final = await client.post( + "/v1/systemone", json=_body(model, marker), headers={"Authorization": f"Bearer {key}"} + ) + return _Served( + marker=marker, + call_id=response.headers.get("x-litellm-call-id", ""), + status=response.status_code, + text=response.text, + ) + + +async def _burst( + base_url: str, key: str, model: str, markers: tuple[str, ...], *, tolerate_transport_errors: bool = False +) -> tuple[_Served, ...]: + async with httpx.AsyncClient(base_url=base_url, timeout=60, trust_env=False) as client: + results: Final = await asyncio.gather( + *(_send(client, key, model, marker) for marker in markers), return_exceptions=tolerate_transport_errors + ) + for result in results: + assert not isinstance(result, BaseException) or isinstance(result, httpx.TransportError), repr(result) + return tuple(result for result in results if isinstance(result, _Served)) + + +def _markers(count: int) -> tuple[str, ...]: + return tuple(uuid.uuid4().hex for _ in range(count)) + + +def _spend_rows(model: str, expected: int) -> list[dict[str, JsonValue]]: + rows: Final = eventually( + lambda: read_rows('SELECT litellm_call_id, status FROM "LiteLLM_SpendLogs" WHERE model_group=%s', (model,)), + lambda found: len(found) >= expected, + seconds=70, + ) + assert len({row["litellm_call_id"] for row in rows}) == len(rows) == expected, rows + return rows + + +def _assert_each_billed_once( + rows: list[dict[str, JsonValue]], failed: tuple[_Served, ...], served: tuple[_Served, ...] +) -> None: + by_call: Final = {string_value(row["litellm_call_id"]): row for row in rows} + assert len(by_call) == len(failed) + len(served), (len(by_call), len(failed), len(served)) + for item in failed: + assert by_call[item.call_id]["status"] == "failure", (item.call_id, by_call[item.call_id]) + for item in served: + assert by_call[item.call_id]["status"] == "success", (item.call_id, by_call[item.call_id]) + + +def _free_port() -> int: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as probe: + probe.bind(("127.0.0.1", 0)) + return int(probe.getsockname()[1]) + + +def _assert_failed_cleanly(served: _Served) -> None: + assert served.status >= 500, (served.status, served.text) + assert "answers" not in served.text, served.text + assert "error" in served.text.lower(), served.text + assert served.call_id, served + + +@pytest.mark.timeout(_CHAOS_TIMEOUT) +async def test_an_upstream_dropping_connections_mid_burst_then_stopped_and_restarted_lands_each_success_once( + gateway: Gateway, +) -> None: + port: Final = _free_port() + first_wave: Final = _markers(12) + dropped: Final = frozenset(marker for index, marker in enumerate(first_wave) if index % 2 == 1) + while_down: Final = _markers(6) + after_restart: Final = _markers(12) + + def flaky(request: Request) -> Reply: + if _marker_of(request) in dropped: + return Reply(drop_connection=True) + return _echo(request) + + with gateway.scenario() as scenario: + model: Final = scenario.model(model=_CANONICAL, api_base=f"http://127.0.0.1:{port}", api_key=_API_KEY) + base_url: Final = str(gateway.client.base_url) + with wire_server(flaky, port=port) as first: + first_served: Final = await _burst(base_url, gateway.key, model, first_wave) + assert len(first_served) == 12 + for item in first_served: + if item.marker in dropped: + _assert_failed_cleanly(item) + else: + _assert_answered_with_its_own_probability(item) + _assert_each_landed_once_with_media(first.drain(), frozenset(first_wave)) + down: Final = await _burst(base_url, gateway.key, model, while_down) + assert len(down) == 6 + for item in down: + _assert_failed_cleanly(item) + with wire_server(_echo, port=port) as second: + recovered: Final = await _burst(base_url, gateway.key, model, after_restart) + assert len(recovered) == 12 + for item in recovered: + _assert_answered_with_its_own_probability(item) + _assert_each_landed_once_with_media(second.drain(), frozenset(after_restart)) + failed: Final = (*(item for item in first_served if item.marker in dropped), *down) + served: Final = (*(item for item in first_served if item.marker not in dropped), *recovered) + _assert_each_billed_once(_spend_rows(model, 30), failed, served) + + +def _chaos_config(wire: Wire, tmp_path: Path) -> Path: + config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + config["model_list"] = [ + { + "model_name": _CONFIG_MODEL, + "litellm_params": {"model": _CANONICAL, "api_base": wire.url, "api_key": _API_KEY}, + } + ] + path: Final = tmp_path / "clef-omni-media-chaos.yaml" + path.write_text(yaml.safe_dump(config)) + return path + + +def _open_upstream_connections(pid: int, upstream: str) -> int: + port: Final = urlsplit(upstream).port + return sum( + 1 + for connection in psutil.Process(pid).net_connections(kind="tcp") + if connection.status == psutil.CONN_ESTABLISHED and connection.raddr and connection.raddr.port == port + ) + + +@pytest.mark.timeout(_CHAOS_TIMEOUT) +async def test_worker_sigkill_mid_burst_leaves_the_sibling_forwarding_media(gateway: Gateway, tmp_path: Path) -> None: + markers: Final = _markers(20) + release: Final = threading.Event() + held_markers: Final[SimpleQueue[str]] = SimpleQueue() + + def held(request: Request) -> Reply: + if request.method != "POST": + return Reply(body=b"{}") + held_markers.put(_marker_of(request)) + assert release.wait(timeout=60), "The burst was never released" + return _echo(request) + + with wire_server(held) as wire: + path: Final = _chaos_config(wire, tmp_path) + with owned_proxy_process(gateway, tmp_path, {}, config=path, workers=2) as owned: + candidate: Final = owned.gateway + workers: Final = eventually( + lambda: tuple(int(pid) for pid in _STARTED_WORKER.findall(owned.log.read_text())), + lambda pids: len(pids) == 2, + seconds=30, + ) + burst: Final = asyncio.create_task( + _burst( + str(candidate.client.base_url), + candidate.key, + _CONFIG_MODEL, + markers, + tolerate_transport_errors=True, + ) + ) + await asyncio.to_thread(eventually, held_markers.qsize, lambda size: size == 20, 60) + held_by: Final = MappingProxyType({pid: _open_upstream_connections(pid, wire.url) for pid in workers}) + assert sum(held_by.values()) == 20, held_by + victim_pid, survivor_pid = sorted(workers, key=held_by.__getitem__) + with owned.replacing_worker(victim_pid) as kill: + kill.freeze() + await asyncio.to_thread(kill) + release.set() + served: Final = await burst + assert held_by[survivor_pid] >= 10, held_by + assert len(served) == held_by[survivor_pid], (held_by, len(served)) + for item in served: + _assert_answered_with_its_own_probability(item) + follow_up: Final = uuid.uuid4().hex + (answered,) = await _burst(str(candidate.client.base_url), candidate.key, _CONFIG_MODEL, (follow_up,)) + _assert_answered_with_its_own_probability(answered) + received: Final = wire.drain() + posts: Final = tuple(request for request in received if request.method == "POST") + assert len(posts) == len(received), [(request.method, request.target) for request in received] + _assert_each_landed_once_with_media(posts, frozenset((*markers, follow_up))) diff --git a/tests/integration/providers/test_cloudflare_clef_omni_media_wire.py b/tests/integration/providers/test_cloudflare_clef_omni_media_wire.py new file mode 100644 index 00000000000..12a8576c31a --- /dev/null +++ b/tests/integration/providers/test_cloudflare_clef_omni_media_wire.py @@ -0,0 +1,856 @@ +import asyncio +import json +import math +import re +import threading +import time +import uuid +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from pathlib import Path +from queue import SimpleQueue +from typing import Final, Literal +from urllib.parse import unquote + +import httpx +import pytest +from integration._support.client import Gateway, eventually, list_value, object_value, string_value +from integration._support.database import read_rows +from integration._support.timing import assert_within_factor_of_control +from integration._support.wire import Reply, Request, wire_server +from pydantic import JsonValue, TypeAdapter + +import litellm + +Route = Literal["/v1/systemone", "/systemone", "/v1/decisions", "/decisions"] + +_ROUTES: Final[tuple[Route, ...]] = ("/v1/systemone", "/systemone", "/v1/decisions", "/decisions") +_API_KEY: Final = "synthetic-cloudflare-key" +_CANONICAL: Final = "cloudflare/@cf/cloudflare/clef-omni" +_ALIAS: Final = "cloudflare/clef-omni" +_COST_MAP_KEY: Final = "cloudflare/@cf/cloudflare/clef-omni" +_UPSTREAM_PATH: Final = "/ai/run/@cf/cloudflare/clef-omni" +_WAV: Final = "data:audio/wav;base64,UklGRiQAAABXQVZFZm10IBAAAAABAAEARKwAAIhYAQACABAAZGF0YQAAAAA=" +_MP4: Final = "data:video/mp4;base64,AAAAHGZ0eXBpc29tAAACAGlzb21pc28ybXA0MQ==" +_WAV_OBJECT: Final[dict[str, JsonValue]] = {"content_type": "audio/wav", "base64": "UklGRiQAAABXQVZF"} +_MP4_OBJECT: Final[dict[str, JsonValue]] = {"content_type": "video/mp4", "base64": "AAAAHGZ0eXBpc29t"} +_PARAMETERIZED_WAV_OBJECT: Final[dict[str, JsonValue]] = {"content_type": "audio/wav;rate=16000", "base64": "UklGRg=="} +_INPUT_TOKENS: Final = 537 +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +_JSON_ROWS: Final = TypeAdapter(list[dict[str, JsonValue]]) +_MARKER: Final = re.compile(r"marker-([0-9a-f]{32})") +_SYSTEM_ONE_QUESTIONS: Final[dict[str, JsonValue]] = { + "pineapple": {"type": "noul", "instructions": "Does the audio say pineapple?"}, + "red": {"type": "noul", "instructions": "Is the video red?"}, +} +_OPENAI_QUESTIONS: Final[list[JsonValue]] = [ + {"type": "predicate", "name": "pineapple", "instructions": "Does the audio say pineapple?"}, + {"type": "predicate", "name": "red", "instructions": "Is the video red?"}, +] +_RED: Final = 0.97 +_SPEND_QUERY: Final = ( + 'SELECT spend, status, call_type, model_group, prompt_tokens, completion_tokens FROM "LiteLLM_SpendLogs" ' + "WHERE request_id = %s" +) +_WORST_CASE_LENGTH: Final = 1_048_576 + + +@dataclass(frozen=True, slots=True) +class _Call: + route: Route + marker: str + + +@dataclass(frozen=True, slots=True) +class _Served: + call: _Call + status: int + text: str + + +def _marker_of(request: Request) -> str: + found: Final = _MARKER.search(request.body.decode()) + assert found is not None, request.body[:200] + return found.group(1) + + +def _state(marker: str) -> str: + return f"Does the clip say pineapple? marker-{marker}" + + +def _probability(marker: str) -> float: + return round(int(marker[:6], 16) / 0x1000000, 7) + + +def _upstream_reply(marker: str) -> dict[str, JsonValue]: + return { + "result": { + "model": "clef-omni", + "answers": { + "pineapple": {"type": "noul", "noul": _probability(marker)}, + "red": {"type": "noul", "noul": _RED}, + }, + "usage": {"input_tokens": _INPUT_TOKENS, "output_tokens": 0}, + }, + "success": True, + "errors": [], + "messages": [], + } + + +def _openai_decider_reply(marker: str) -> dict[str, JsonValue]: + return { + "model": "gpt-6-luna", + "answers": [ + {"type": "predicate", "name": "pineapple", "probability": _probability(marker)}, + {"type": "predicate", "name": "red", "probability": _RED}, + ], + "usage": {"input_tokens": _INPUT_TOKENS, "output_tokens": 0, "total_tokens": _INPUT_TOKENS}, + } + + +def _echo(request: Request) -> Reply: + marker: Final = _marker_of(request) + if request.target.endswith("/v1/decisions"): + return Reply(body=json.dumps(_openai_decider_reply(marker)).encode()) + if request.target.endswith("/systemone"): + return Reply(body=json.dumps(_upstream_reply(marker)["result"]).encode()) + return Reply(body=json.dumps(_upstream_reply(marker)).encode()) + + +def _speaks_system_one(route: Route) -> bool: + return route.endswith("systemone") + + +def _body(route: Route, model: str, marker: str, **extra: JsonValue) -> dict[str, JsonValue]: + if _speaks_system_one(route): + return {"model": model, "state": _state(marker), "questions": _SYSTEM_ONE_QUESTIONS, **extra} + return {"model": model, "input": _state(marker), "questions": _OPENAI_QUESTIONS, **extra} + + +def _expected_response(route: Route, marker: str) -> dict[str, JsonValue]: + if _speaks_system_one(route): + return { + "model": "clef-omni", + "answers": { + "pineapple": {"type": "noul", "noul": _probability(marker)}, + "red": {"type": "noul", "noul": _RED}, + }, + "usage": {"input_tokens": _INPUT_TOKENS, "output_tokens": 0}, + } + return { + "model": "clef-omni", + "answers": [ + {"type": "predicate", "name": "pineapple", "probability": _probability(marker)}, + {"type": "predicate", "name": "red", "probability": _RED}, + ], + "usage": { + "input_tokens": _INPUT_TOKENS, + "input_tokens_details": {"cached_tokens": 0, "cache_write_tokens": 0}, + "output_tokens": 0, + "output_tokens_details": {"reasoning_tokens": 0}, + "total_tokens": _INPUT_TOKENS, + }, + } + + +def _expected_upstream_body(marker: str, **media: JsonValue) -> dict[str, JsonValue]: + return {"model": "clef-omni", "state": _state(marker), **media, "questions": _SYSTEM_ONE_QUESTIONS} + + +def _upstream_body(request: Request) -> dict[str, JsonValue]: + assert request.method == "POST", request.method + assert unquote(request.target) == _UPSTREAM_PATH, request.target + assert request.headers["authorization"] == f"Bearer {_API_KEY}", request.headers + return _JSON_OBJECT.validate_json(request.body) + + +def _decide(gateway: Gateway, route: Route, model: str, marker: str, **extra: JsonValue) -> httpx.Response: + return gateway.request("POST", route, _body(route, model, marker, **extra)) + + +def _answered_probability(route: Route, text: str) -> JsonValue: + answers: Final = _JSON_OBJECT.validate_json(text)["answers"] + if _speaks_system_one(route): + return object_value(object_value(answers)["pineapple"])["noul"] + return object_value(list_value(answers)[0])["probability"] + + +def _price() -> dict[str, JsonValue]: + cost_map: Final = _JSON_OBJECT.validate_json(Path("model_prices_and_context_window.json").read_bytes()) + priced: Final = cost_map.get(_COST_MAP_KEY) + assert priced is not None, f"{_COST_MAP_KEY} has no cost map entry" + return object_value(priced) + + +def _number(value: JsonValue) -> float: + assert isinstance(value, (int, float)) and not isinstance(value, bool), value + return float(value) + + +def _expected_spend() -> float: + return _INPUT_TOKENS * _number(_price()["input_cost_per_token"]) + + +def _spend_row(call_id: str) -> dict[str, JsonValue]: + rows: Final = eventually(lambda: read_rows(_SPEND_QUERY, (call_id,)), lambda found: len(found) == 1, seconds=70) + return rows[0] + + +def _spend_rows_for(model: str, expected: int) -> list[dict[str, JsonValue]]: + rows: Final = eventually( + lambda: read_rows('SELECT litellm_call_id, status FROM "LiteLLM_SpendLogs" WHERE model_group=%s', (model,)), + lambda found: len(found) >= expected, + seconds=70, + ) + assert len({row["litellm_call_id"] for row in rows}) == len(rows) == expected, rows + return rows + + +def _call_type(route: Route) -> str: + return "asystemone" if route.endswith("systemone") else "adecisions" + + +def _assert_billed(gateway: Gateway, response: httpx.Response, model: str, route: Route) -> None: + expected: Final = _expected_spend() + assert math.isclose(float(response.headers["x-litellm-response-cost"]), expected, rel_tol=1e-9), response.headers + call_id: Final = response.headers["x-litellm-call-id"] + row: Final = _spend_row(call_id) + assert (row["status"], row["call_type"], row["model_group"], row["prompt_tokens"], row["completion_tokens"]) == ( + "success", + _call_type(route), + model, + _INPUT_TOKENS, + 0, + ), row + assert math.isclose(_number(row["spend"]), expected, rel_tol=1e-9), row + listed: Final = gateway.request("GET", "/spend/logs", params={"request_id": call_id}) + assert listed.status_code == 200, listed.text + (reported,) = _JSON_ROWS.validate_json(listed.content) + assert reported["request_id"] == call_id, reported + assert (reported["model_group"], reported["prompt_tokens"], reported["completion_tokens"]) == ( + model, + _INPUT_TOKENS, + 0, + ), reported + assert math.isclose(_number(reported["spend"]), expected, rel_tol=1e-9), reported + + +def _assert_served_with_media(response: httpx.Response, route: Route, marker: str, request: Request) -> None: + assert response.status_code == 200, response.text + assert response.json() == _expected_response(route, marker) + assert _upstream_body(request) == _expected_upstream_body(marker, audio=[_WAV], videos=[_MP4]) + + +@pytest.mark.parametrize("route", _ROUTES) +def test_data_url_clips_reach_clef_omni_on_every_route_and_are_billed(gateway: Gateway, route: Route) -> None: + marker: Final = uuid.uuid4().hex + with wire_server(_echo) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=_CANONICAL, api_base=wire.url, api_key=_API_KEY) + response: Final = _decide(gateway, route, model, marker, audio=[_WAV], videos=[_MP4]) + (request,) = wire.drain() + _assert_served_with_media(response, route, marker, request) + _assert_billed(gateway, response, model, route) + + +async def test_an_async_caller_on_the_unversioned_decisions_route_gets_its_media_forwarded(gateway: Gateway) -> None: + marker: Final = uuid.uuid4().hex + with wire_server(_echo) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=_CANONICAL, api_base=wire.url, api_key=_API_KEY) + async with httpx.AsyncClient(base_url=str(gateway.client.base_url), timeout=60, trust_env=False) as client: + response: Final = await client.post( + "/decisions", + json=_body("/decisions", model, marker, audio=[_WAV], videos=[_MP4]), + headers={"Authorization": f"Bearer {gateway.key}"}, + ) + (request,) = wire.drain() + _assert_served_with_media(response, "/decisions", marker, request) + _assert_billed(gateway, response, model, "/decisions") + + +@pytest.mark.parametrize( + "audio_clip", (_WAV_OBJECT, _PARAMETERIZED_WAV_OBJECT), ids=("plain_content_type", "parameterized_content_type") +) +async def test_object_form_clips_are_forwarded_verbatim(gateway: Gateway, audio_clip: dict[str, JsonValue]) -> None: + marker: Final = uuid.uuid4().hex + with wire_server(_echo) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=_CANONICAL, api_base=wire.url, api_key=_API_KEY) + async with httpx.AsyncClient(base_url=str(gateway.client.base_url), timeout=60, trust_env=False) as client: + response: Final = await client.post( + "/v1/systemone", + json=_body("/v1/systemone", model, marker, audio=[audio_clip], videos=[_MP4_OBJECT]), + headers={"Authorization": f"Bearer {gateway.key}"}, + ) + assert response.status_code == 200, response.text + assert response.json() == _expected_response("/v1/systemone", marker) + (request,) = wire.drain() + assert _upstream_body(request) == _expected_upstream_body(marker, audio=[audio_clip], videos=[_MP4_OBJECT]) + + +def test_both_clef_omni_spellings_forward_media_and_carry_the_cost_map_price(gateway: Gateway) -> None: + price: Final = _price() + with wire_server(_echo) as wire, gateway.scenario() as scenario: + deployments: Final = tuple( + scenario.model(model=spelling, api_base=wire.url, api_key=_API_KEY) for spelling in (_ALIAS, _CANONICAL) + ) + for model in deployments: + marker: Final = uuid.uuid4().hex + response: Final = _decide(gateway, "/v1/systemone", model, marker, audio=[_WAV], videos=[_MP4]) + (request,) = wire.drain() + _assert_served_with_media(response, "/v1/systemone", marker, request) + _assert_billed(gateway, response, model, "/v1/systemone") + listed: Final = list_value(gateway.get("/model/info")["data"]) + for model in deployments: + entry: Final = next(object_value(item) for item in listed if object_value(item)["model_name"] == model) + info: Final = object_value(entry["model_info"]) + assert ( + info.get("input_cost_per_token"), + info.get("output_cost_per_token"), + info.get("mode"), + info.get("max_input_tokens"), + info.get("supports_audio_input"), + info.get("supports_video_input"), + ) == ( + price["input_cost_per_token"], + price["output_cost_per_token"], + price["mode"], + price["max_input_tokens"], + price["supports_audio_input"], + price["supports_video_input"], + ), (model, info) + + +def test_the_maxima_of_four_audio_and_two_video_clips_are_forwarded(gateway: Gateway) -> None: + marker: Final = uuid.uuid4().hex + audio: Final[list[JsonValue]] = [_WAV, _WAV_OBJECT, _WAV, _WAV_OBJECT] + videos: Final[list[JsonValue]] = [_MP4, _MP4_OBJECT] + with wire_server(_echo) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=_CANONICAL, api_base=wire.url, api_key=_API_KEY) + response: Final = _decide(gateway, "/v1/systemone", model, marker, audio=audio, videos=videos) + assert response.status_code == 200, response.text + (request,) = wire.drain() + assert _upstream_body(request) == _expected_upstream_body(marker, audio=audio, videos=videos) + + +def test_missing_empty_and_null_media_send_the_same_text_only_body(gateway: Gateway) -> None: + marker: Final = uuid.uuid4().hex + with wire_server(_echo) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=_CANONICAL, api_base=wire.url, api_key=_API_KEY) + missing: Final = _decide(gateway, "/v1/systemone", model, marker) + empty: Final = _decide(gateway, "/v1/systemone", model, marker, audio=[], videos=[]) + null: Final = _decide(gateway, "/v1/systemone", model, marker, audio=None, videos=None) + for response in (missing, empty, null): + assert response.status_code == 200, response.text + assert response.json() == _expected_response("/v1/systemone", marker) + bodies: Final = tuple(_upstream_body(request) for request in wire.drain()) + assert bodies == (_expected_upstream_body(marker),) * 3, bodies + + +@pytest.mark.parametrize("route", _ROUTES) +def test_null_clips_mean_no_clips_of_that_kind_on_every_route(gateway: Gateway, route: Route) -> None: + marker: Final = uuid.uuid4().hex + with wire_server(_echo) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=_CANONICAL, api_base=wire.url, api_key=_API_KEY) + both_null: Final = _decide(gateway, route, model, marker, audio=None, videos=None) + audio_only: Final = _decide(gateway, route, model, marker, audio=[_WAV], videos=None) + video_only: Final = _decide(gateway, route, model, marker, audio=None, videos=[_MP4]) + for response in (both_null, audio_only, video_only): + assert response.status_code == 200, response.text + assert response.json() == _expected_response(route, marker) + bodies: Final = tuple(_upstream_body(request) for request in wire.drain()) + assert bodies == ( + _expected_upstream_body(marker), + _expected_upstream_body(marker, audio=[_WAV]), + _expected_upstream_body(marker, videos=[_MP4]), + ), bodies + + +_SDK_QUESTIONS: Final = TypeAdapter(dict[str, dict[str, object]]).validate_python(_SYSTEM_ONE_QUESTIONS) + + +async def test_positional_sdk_callers_keep_the_api_key_fifth_and_the_api_base_sixth() -> None: + marker: Final = uuid.uuid4().hex + with wire_server(_echo) as wire: + synchronous: Final = litellm.decisions( + _CANONICAL, None, _SDK_QUESTIONS, None, _API_KEY, wire.url, state=_state(marker), audio=[_WAV] + ) + asynchronous: Final = await litellm.adecisions( + _CANONICAL, None, _SDK_QUESTIONS, None, _API_KEY, wire.url, state=_state(marker), videos=[_MP4] + ) + for response in (synchronous, asynchronous): + assert response.model_dump(mode="json") == _expected_response("/v1/systemone", marker) + bodies: Final = tuple(_upstream_body(request) for request in wire.drain()) + assert bodies == ( + _expected_upstream_body(marker, audio=[_WAV]), + _expected_upstream_body(marker, videos=[_MP4]), + ), bodies + + +_MALFORMED: Final[tuple[tuple[str, dict[str, JsonValue], str], ...]] = ( + ("audio_int", {"audio": 5}, "audio"), + ("audio_list_of_ints", {"audio": [1, 2]}, "audio"), + ("audio_empty_string", {"audio": ""}, "audio"), + ("audio_5kb_garbage", {"audio": ["a" * 5000]}, "pattern"), + ("videos_object_missing_base64", {"videos": [{"content_type": "video/mp4"}]}, "base64"), + ("videos_object_extra_key", {"videos": [{**_MP4_OBJECT, "url": "https://example.invalid/clip.mp4"}]}, "url"), + ("image_data_url_in_audio", {"audio": ["data:image/png;base64,iVBORw0KGgo="]}, "pattern"), + ( + "parameterized_audio_data_url", + {"audio": ["data:audio/webm;codecs=opus;base64,GkXfo0AgQoaBAUL3gQFC8oEEQvOBCA=="]}, + "pattern", + ), + ("five_audio_clips", {"audio": [_WAV] * 5}, "at most 4"), + ("three_videos", {"videos": [_MP4] * 3}, "at most 2"), + ( + "audio_object_with_image_content_type", + {"audio": [{"content_type": "image/png", "base64": "iVBORw0KGgo="}]}, + "pattern", + ), + ("video_object_with_empty_base64", {"videos": [{"content_type": "video/mp4", "base64": ""}]}, "at least 1"), + ( + "audio_object_content_type_with_a_space", + {"audio": [{"content_type": "audio/wav x", "base64": "UklGRg=="}]}, + "pattern", + ), +) + + +@pytest.mark.parametrize(("label", "extra", "fragment"), _MALFORMED, ids=[label for label, _, _ in _MALFORMED]) +def test_malformed_clips_are_refused_before_any_upstream_call( + gateway: Gateway, label: str, extra: dict[str, JsonValue], fragment: str +) -> None: + marker: Final = uuid.uuid4().hex + with wire_server(_echo) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=_CANONICAL, api_base=wire.url, api_key=_API_KEY) + response: Final = _decide(gateway, "/v1/systemone", model, marker, **extra) + assert response.status_code == 400, (label, response.text) + assert "Invalid Decisions request" in response.text and fragment in response.text, (label, response.text) + assert wire.drain() == (), label + + +def _raw_system_one(gateway: Gateway, model: str, marker: str, audio_fragment: str) -> httpx.Response: + raw: Final = ( + f'{{"model": {json.dumps(model)}, "state": {json.dumps(_state(marker))}, ' + f'"questions": {json.dumps(_SYSTEM_ONE_QUESTIONS)}, {audio_fragment}}}' + ) + return gateway.client.post( + "/v1/systemone", + content=raw.encode(), + headers={"Authorization": f"Bearer {gateway.key}", "Content-Type": "application/json"}, + ) + + +def test_a_duplicated_audio_key_is_judged_by_its_last_value(gateway: Gateway) -> None: + marker: Final = uuid.uuid4().hex + image: Final = '"data:image/png;base64,iVBORw0KGgo="' + with wire_server(_echo) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=_CANONICAL, api_base=wire.url, api_key=_API_KEY) + refused: Final = _raw_system_one(gateway, model, marker, f'"audio": [{json.dumps(_WAV)}], "audio": [{image}]') + assert refused.status_code == 400, refused.text + assert "Invalid Decisions request" in refused.text and "pattern" in refused.text, refused.text + assert wire.drain() == () + accepted: Final = _raw_system_one(gateway, model, marker, f'"audio": [{image}], "audio": [{json.dumps(_WAV)}]') + assert accepted.status_code == 200, accepted.text + (request,) = wire.drain() + assert _upstream_body(request) == _expected_upstream_body(marker, audio=[_WAV]) + + +_UNSUPPORTED: Final[tuple[tuple[str, Route], ...]] = ( + ("cloudflare/clef-flash", "/v1/systemone"), + ("cloudflare/clef", "/v1/systemone"), + ("cloudflare/@cf/cloudflare/clef-flash", "/v1/decisions"), + ("openai/gpt-6-luna", "/v1/systemone"), + ("openai/gpt-6-luna", "/v1/decisions"), + ("fireworks_ai/accounts/fireworks/routers/rune", "/v1/systemone"), + ("fireworks_ai/accounts/fireworks/routers/rune", "/v1/decisions"), +) + + +@pytest.mark.parametrize(("spelling", "route"), _UNSUPPORTED, ids=[f"{m}@{r}" for m, r in _UNSUPPORTED]) +def test_media_on_a_model_without_support_is_refused_before_any_upstream_call( + gateway: Gateway, spelling: str, route: Route +) -> None: + marker: Final = uuid.uuid4().hex + with wire_server(_echo) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=spelling, api_base=wire.url, api_key=_API_KEY) + response: Final = _decide(gateway, route, model, marker, audio=[_WAV]) + assert response.status_code == 400, response.text + assert "does not support parameters: ['audio']" in response.text, response.text + assert "drop_params" in response.text, response.text + assert wire.drain() == () + text_only: Final = _decide(gateway, route, model, marker) + assert text_only.status_code == 200, text_only.text + assert len(wire.drain()) == 1 + + +@pytest.mark.parametrize("route", ("/v1/systemone", "/v1/decisions")) +def test_media_on_a_model_without_support_is_dropped_when_the_deployment_drops_params( + gateway: Gateway, route: Route +) -> None: + marker: Final = uuid.uuid4().hex + with wire_server(_echo) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model="cloudflare/clef-flash", api_base=wire.url, api_key=_API_KEY, drop_params=True + ) + text_only: Final = _decide(gateway, route, model, marker) + assert text_only.status_code == 200, text_only.text + (text_only_request,) = wire.drain() + with_clips: Final = _decide(gateway, route, model, marker, audio=[_WAV], videos=[_MP4]) + assert with_clips.status_code == 200, with_clips.text + (dropped_request,) = wire.drain() + text_only_body: Final = _JSON_OBJECT.validate_json(text_only_request.body) + assert "audio" not in text_only_body and "videos" not in text_only_body, text_only_body + assert _JSON_OBJECT.validate_json(dropped_request.body) == text_only_body + assert with_clips.json() == text_only.json() + + +_MALFORMED_CLIPS: Final[dict[str, JsonValue]] = { + "audio": ["https://example.invalid/clip.wav"], + "videos": [{"content_type": "video/mp4"}], +} +_NO_MEDIA_SPELLINGS: Final = ("cloudflare/clef-flash", "cloudflare/clef") + + +def _assert_refused_before_any_upstream_call(response: httpx.Response, fragment: str, wire_received: int) -> None: + assert response.status_code == 400, response.text + assert fragment in response.text, response.text + assert wire_received == 0 + + +@pytest.mark.parametrize("route", ("/v1/systemone", "/v1/decisions")) +def test_malformed_clips_on_a_model_without_support_are_refused_unless_the_deployment_drops_params( + gateway: Gateway, route: Route +) -> None: + marker: Final = uuid.uuid4().hex + with wire_server(_echo) as wire, gateway.scenario() as scenario: + strict: Final = scenario.model(model="cloudflare/clef-flash", api_base=wire.url, api_key=_API_KEY) + dropping: Final = scenario.model( + model="cloudflare/clef-flash", api_base=wire.url, api_key=_API_KEY, drop_params=True + ) + refused: Final = _decide(gateway, route, strict, marker, **_MALFORMED_CLIPS) + _assert_refused_before_any_upstream_call(refused, "Invalid Decisions request", len(wire.drain())) + text_only: Final = _decide(gateway, route, dropping, marker) + assert text_only.status_code == 200, text_only.text + (text_only_request,) = wire.drain() + dropped: Final = _decide(gateway, route, dropping, marker, **_MALFORMED_CLIPS) + assert dropped.status_code == 200, dropped.text + (dropped_request,) = wire.drain() + assert _JSON_OBJECT.validate_json(dropped_request.body) == _JSON_OBJECT.validate_json(text_only_request.body) + assert dropped.json() == text_only.json() + + +def test_a_malformed_clip_on_a_media_model_is_refused_even_when_the_deployment_drops_params( + gateway: Gateway, +) -> None: + marker: Final = uuid.uuid4().hex + with wire_server(_echo) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=_CANONICAL, api_base=wire.url, api_key=_API_KEY, drop_params=True) + refused: Final = _decide(gateway, "/v1/systemone", model, marker, **_MALFORMED_CLIPS) + _assert_refused_before_any_upstream_call(refused, "Invalid Decisions request", len(wire.drain())) + served: Final = _decide(gateway, "/v1/systemone", model, marker, audio=[_WAV], videos=[_MP4]) + (request,) = wire.drain() + _assert_served_with_media(served, "/v1/systemone", marker, request) + + +def _advertised_media_flags(gateway: Gateway, deployments: Mapping[str, str]) -> dict[str, tuple[JsonValue, JsonValue]]: + listed: Final = list_value(gateway.get("/model/info")["data"]) + entries: Final = { + string_value(object_value(item)["model_name"]): object_value(object_value(item)["model_info"]) + for item in listed + } + return { + spelling: (entries[name].get("supports_audio_input"), entries[name].get("supports_video_input")) + for spelling, name in deployments.items() + } + + +def test_model_info_advertises_media_support_only_on_clef_omni_deployments(gateway: Gateway) -> None: + with wire_server(_echo) as wire, gateway.scenario() as scenario: + deployments: Final = { + spelling: scenario.model(model=spelling, api_base=wire.url, api_key=_API_KEY) + for spelling in (_ALIAS, _CANONICAL, *_NO_MEDIA_SPELLINGS) + } + flags: Final = _advertised_media_flags(gateway, deployments) + assert flags[_ALIAS] == (True, True), flags + assert flags[_CANONICAL] == (True, True), flags + advertised_without_media: Final = tuple(any(flags[spelling]) for spelling in _NO_MEDIA_SPELLINGS) + assert advertised_without_media == (False, False), flags + + +def _unique_backend(kind: str) -> str: + return f"@cf/cloudflare/clef-{kind}-{uuid.uuid4().hex[:12]}" + + +def _body_sent_to(request: Request, backend: str) -> dict[str, JsonValue]: + assert request.method == "POST", request.method + assert unquote(request.target) == f"/ai/run/{backend}", request.target + assert request.headers["authorization"] == f"Bearer {_API_KEY}", request.headers + return _JSON_OBJECT.validate_json(request.body) + + +def _expected_backend_body(backend: str, marker: str, **media: JsonValue) -> dict[str, JsonValue]: + return {**_expected_upstream_body(marker, **media), "model": backend.rsplit("/", 1)[-1]} + + +def test_a_deployment_advertising_only_audio_forwards_audio_and_refuses_or_drops_video(gateway: Gateway) -> None: + marker: Final = uuid.uuid4().hex + backend: Final = _unique_backend("audio") + advertised: Final[dict[str, JsonValue]] = {"mode": "evaluation", "supports_audio_input": True} + with wire_server(_echo) as wire, gateway.scenario() as scenario: + strict: Final = scenario.model( + model=f"cloudflare/{backend}", api_base=wire.url, api_key=_API_KEY, model_info=advertised + ) + dropping: Final = scenario.model( + model=f"cloudflare/{backend}", api_base=wire.url, api_key=_API_KEY, drop_params=True, model_info=advertised + ) + with_audio: Final = _decide(gateway, "/v1/systemone", strict, marker, audio=[_WAV]) + assert with_audio.status_code == 200, with_audio.text + (audio_request,) = wire.drain() + assert _body_sent_to(audio_request, backend) == _expected_backend_body(backend, marker, audio=[_WAV]) + with_video: Final = _decide(gateway, "/v1/systemone", strict, marker, videos=[_MP4]) + _assert_refused_before_any_upstream_call( + with_video, "does not support parameters: ['videos']", len(wire.drain()) + ) + assert "drop_params" in with_video.text, with_video.text + dropped: Final = _decide(gateway, "/v1/systemone", dropping, marker, audio=[_WAV], videos=[_MP4]) + assert dropped.status_code == 200, dropped.text + (dropped_request,) = wire.drain() + assert _body_sent_to(dropped_request, backend) == _expected_backend_body(backend, marker, audio=[_WAV]) + + +def test_a_chat_mode_deployment_advertising_media_is_not_a_decisions_media_model(gateway: Gateway) -> None: + marker: Final = uuid.uuid4().hex + backend: Final = _unique_backend("chat") + advertised: Final[dict[str, JsonValue]] = { + "mode": "chat", + "supports_audio_input": True, + "supports_video_input": True, + } + with wire_server(_echo) as wire, gateway.scenario() as scenario: + strict: Final = scenario.model( + model=f"cloudflare/{backend}", api_base=wire.url, api_key=_API_KEY, model_info=advertised + ) + dropping: Final = scenario.model( + model=f"cloudflare/{backend}", api_base=wire.url, api_key=_API_KEY, drop_params=True, model_info=advertised + ) + refused: Final = _decide(gateway, "/v1/systemone", strict, marker, audio=[_WAV], videos=[_MP4]) + _assert_refused_before_any_upstream_call(refused, "does not support parameters", len(wire.drain())) + assert "'audio'" in refused.text and "'videos'" in refused.text, refused.text + assert "drop_params" in refused.text, refused.text + text_only: Final = _decide(gateway, "/v1/systemone", dropping, marker) + assert text_only.status_code == 200, text_only.text + (text_only_request,) = wire.drain() + assert _body_sent_to(text_only_request, backend) == _expected_backend_body(backend, marker) + dropped: Final = _decide(gateway, "/v1/systemone", dropping, marker, audio=[_WAV], videos=[_MP4]) + assert dropped.status_code == 200, dropped.text + (dropped_request,) = wire.drain() + assert _body_sent_to(dropped_request, backend) == _expected_backend_body(backend, marker) + assert dropped.json() == text_only.json() + + +def test_clips_on_an_openai_chat_model_are_refused_or_dropped_never_forwarded(gateway: Gateway) -> None: + marker: Final = uuid.uuid4().hex + with wire_server(_echo) as wire, gateway.scenario() as scenario: + strict: Final = scenario.model(model="openai/gpt-audio", api_base=wire.url, api_key=_API_KEY) + dropping: Final = scenario.model( + model="openai/gpt-audio", api_base=wire.url, api_key=_API_KEY, drop_params=True + ) + refused: Final = _decide(gateway, "/v1/decisions", strict, marker, audio=[_WAV]) + _assert_refused_before_any_upstream_call(refused, "does not support parameters: ['audio']", len(wire.drain())) + assert "drop_params" in refused.text, refused.text + text_only: Final = _decide(gateway, "/v1/decisions", dropping, marker) + assert text_only.status_code == 200, text_only.text + (text_only_request,) = wire.drain() + assert unquote(text_only_request.target) == "/v1/decisions", text_only_request.target + text_only_body: Final = _JSON_OBJECT.validate_json(text_only_request.body) + assert "audio" not in text_only_body and "videos" not in text_only_body, text_only_body + dropped: Final = _decide(gateway, "/v1/decisions", dropping, marker, audio=[_WAV], videos=[_MP4]) + assert dropped.status_code == 200, dropped.text + (dropped_request,) = wire.drain() + assert _JSON_OBJECT.validate_json(dropped_request.body) == text_only_body + assert dropped.json() == text_only.json() + + +_UPSTREAM_REJECTIONS: Final[tuple[tuple[int, str], ...]] = ( + (422, "Each audio entry must be a data URL or an object with content_type and base64"), + (413, "Request entity too large"), +) + + +@pytest.mark.parametrize(("status", "message"), _UPSTREAM_REJECTIONS, ids=("422", "413")) +def test_an_upstream_rejection_of_a_media_request_reaches_the_caller_like_a_text_only_one( + gateway: Gateway, status: int, message: str +) -> None: + marker: Final = uuid.uuid4().hex + rejection: Final = json.dumps( + {"result": {}, "success": False, "errors": [{"code": 5006, "message": message}], "messages": []} + ).encode() + with wire_server(lambda _: Reply(status=status, body=rejection)) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=_CANONICAL, api_base=wire.url, api_key=_API_KEY) + text_only: Final = _decide(gateway, "/v1/systemone", model, marker, num_retries=0) + with_media: Final = _decide(gateway, "/v1/systemone", model, marker, num_retries=0, audio=[_WAV], videos=[_MP4]) + for response in (text_only, with_media): + assert response.status_code >= 400, response.text + assert message in response.text, response.text + assert "answers" not in response.text, response.text + assert with_media.status_code == text_only.status_code, (text_only.status_code, with_media.status_code) + first, second = (_upstream_body(request) for request in wire.drain()) + assert first == _expected_upstream_body(marker) + assert second == _expected_upstream_body(marker, audio=[_WAV], videos=[_MP4]) + + +def test_an_unauthenticated_media_request_is_refused_before_validation(gateway: Gateway) -> None: + marker: Final = uuid.uuid4().hex + with wire_server(_echo) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=_CANONICAL, api_base=wire.url, api_key=_API_KEY) + for route in _ROUTES: + response: Final = gateway.client.post(route, json=_body(route, model, marker, audio=[_WAV], videos=[_MP4])) + assert response.status_code == 401, (route, response.text) + assert "Authentication Error" in response.text, response.text + assert wire.drain() == () + + +def test_unrelated_chat_keeps_working_after_a_run_of_refused_media_requests(gateway: Gateway) -> None: + marker: Final = uuid.uuid4().hex + with wire_server(_echo) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=_CANONICAL, api_base=wire.url, api_key=_API_KEY) + control: Final = scenario.model() + for _, extra, _ in _MALFORMED: + assert _decide(gateway, "/v1/systemone", model, marker, **extra).status_code == 400 + assert wire.drain() == () + assert gateway.chat(control)["choices"], "the control deployment stopped answering" + accepted: Final = _decide(gateway, "/v1/systemone", model, marker, audio=[_WAV]) + assert accepted.status_code == 200, accepted.text + (request,) = wire.drain() + assert _upstream_body(request) == _expected_upstream_body(marker, audio=[_WAV]) + + +def test_repeated_and_clip_changed_requests_each_reach_the_upstream_and_are_each_billed(gateway: Gateway) -> None: + marker: Final = uuid.uuid4().hex + with wire_server(_echo) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=_CANONICAL, api_base=wire.url, api_key=_API_KEY) + first: Final = _decide(gateway, "/v1/systemone", model, marker, audio=[_WAV], videos=[_MP4]) + second: Final = _decide(gateway, "/v1/systemone", model, marker, audio=[_WAV], videos=[_MP4]) + swapped: Final = _decide(gateway, "/v1/systemone", model, marker, audio=[_WAV_OBJECT], videos=[_MP4]) + responses: Final = (first, second, swapped) + for response in responses: + assert response.status_code == 200, response.text + assert response.json() == _expected_response("/v1/systemone", marker) + call_ids: Final = tuple(response.headers["x-litellm-call-id"] for response in responses) + assert len(set(call_ids)) == 3, call_ids + bodies: Final = tuple(_upstream_body(request) for request in wire.drain()) + assert bodies == ( + _expected_upstream_body(marker, audio=[_WAV], videos=[_MP4]), + _expected_upstream_body(marker, audio=[_WAV], videos=[_MP4]), + _expected_upstream_body(marker, audio=[_WAV_OBJECT], videos=[_MP4]), + ), bodies + for call_id in call_ids: + assert _spend_row(call_id)["status"] == "success", call_id + + +async def _send(client: httpx.AsyncClient, key: str, model: str, call: _Call) -> _Served: + response: Final = await client.post( + call.route, + json=_body(call.route, model, call.marker, audio=[_WAV], videos=[_MP4]), + headers={"Authorization": f"Bearer {key}"}, + ) + return _Served(call=call, status=response.status_code, text=response.text) + + +async def _burst(base_url: str, key: str, model: str, calls: Sequence[_Call]) -> tuple[_Served, ...]: + async with httpx.AsyncClient(base_url=base_url, timeout=60, trust_env=False) as client: + return tuple(await asyncio.gather(*(_send(client, key, model, call) for call in calls))) + + +def _assert_answered_with_its_own_probability(served: _Served) -> None: + assert served.status == 200, served.text + assert _answered_probability(served.call.route, served.text) == _probability(served.call.marker), served.text + + +def _assert_each_landed_once_with_media(received: tuple[Request, ...], calls: Sequence[_Call]) -> None: + bodies: Final = {_marker_of(request): _upstream_body(request) for request in received} + assert len(bodies) == len(received) == len(calls), (len(bodies), len(received), len(calls)) + for call in calls: + assert bodies[call.marker] == _expected_upstream_body(call.marker, audio=[_WAV], videos=[_MP4]), call + + +async def test_a_concurrent_burst_across_the_four_routes_lands_each_request_once_with_its_media( + gateway: Gateway, +) -> None: + calls: Final = tuple(_Call(route=_ROUTES[index % len(_ROUTES)], marker=uuid.uuid4().hex) for index in range(20)) + with wire_server(_echo) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=_CANONICAL, api_base=wire.url, api_key=_API_KEY) + served: Final = await _burst(str(gateway.client.base_url), gateway.key, model, calls) + assert len(served) == 20 + for item in served: + _assert_answered_with_its_own_probability(item) + _assert_each_landed_once_with_media(wire.drain(), calls) + rows: Final = _spend_rows_for(model, 20) + assert [row["status"] for row in rows] == ["success"] * 20, rows + + +def test_openapi_documents_the_media_fields_on_every_decisions_route(gateway: Gateway) -> None: + paths: Final = object_value(gateway.get("/openapi.json")["paths"]) + for route in _ROUTES: + post: Final = object_value(object_value(paths[route])["post"]) + content: Final = object_value(object_value(post["requestBody"])["content"]) + schema: Final = object_value(object_value(content["application/json"])["schema"]) + properties: Final = object_value(schema["properties"]) + assert "audio" in properties and "videos" in properties, (route, sorted(properties)) + assert object_value(properties["audio"])["maxItems"] == 4, (route, properties["audio"]) + assert object_value(properties["videos"])["maxItems"] == 2, (route, properties["videos"]) + + +def _probe_liveliness(base_url: str, stop: threading.Event, latencies: SimpleQueue[float]) -> None: + with httpx.Client(base_url=base_url, timeout=10, trust_env=False) as client: + while not stop.is_set(): + started: Final = time.perf_counter() + response: Final = client.get("/health/liveliness") + latencies.put(time.perf_counter() - started) + assert response.status_code == 200, response.text + + +def _timed_refusal_under_liveliness_probes( + gateway: Gateway, model: str, clip: str +) -> tuple[httpx.Response, float, tuple[float, ...]]: + stop: Final = threading.Event() + latencies: Final[SimpleQueue[float]] = SimpleQueue() + probe: Final = threading.Thread( + target=_probe_liveliness, args=(str(gateway.client.base_url), stop, latencies), daemon=True + ) + probe.start() + started: Final = time.perf_counter() + with httpx.Client(base_url=str(gateway.client.base_url), timeout=60, trust_env=False) as client: + response: Final = client.post( + "/v1/systemone", + json=_body("/v1/systemone", model, uuid.uuid4().hex, audio=[clip]), + headers={"Authorization": f"Bearer {gateway.key}"}, + ) + elapsed: Final = time.perf_counter() - started + stop.set() + probe.join(timeout=15) + return response, elapsed, tuple(latencies.get_nowait() for _ in range(latencies.qsize())) + + +def test_the_data_url_pattern_worst_case_is_refused_in_bounded_time_while_liveliness_answers( + gateway: Gateway, +) -> None: + with wire_server(_echo) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=_CANONICAL, api_base=wire.url, api_key=_API_KEY) + tenth_response, tenth_elapsed, _ = _timed_refusal_under_liveliness_probes( + gateway, model, "data:audio/" + "a" * (_WORST_CASE_LENGTH // 10) + ) + response, elapsed, observed = _timed_refusal_under_liveliness_probes( + gateway, model, "data:audio/" + "a" * _WORST_CASE_LENGTH + ) + timing: Final = { + "tenth_seconds": round(tenth_elapsed, 3), + "classifier_seconds": round(elapsed, 3), + "slowest_liveliness_seconds": round(max(observed), 3), + } + assert tenth_response.status_code == 400, tenth_response.text[:300] + assert_within_factor_of_control(elapsed, control=tenth_elapsed, factor=25) + assert observed and max(observed) < max(1.0, elapsed / 2), timing + assert response.status_code == 400, (timing, response.text[:300]) + assert "pattern" in response.text, response.text[:300] + assert wire.drain() == (), timing diff --git a/tests/integration/providers/test_coralbricks_provider_chaos.py b/tests/integration/providers/test_coralbricks_provider_chaos.py new file mode 100644 index 00000000000..dc4014fdb75 --- /dev/null +++ b/tests/integration/providers/test_coralbricks_provider_chaos.py @@ -0,0 +1,396 @@ +from __future__ import annotations + +import asyncio +import re +import threading +import uuid +from collections.abc import Callable, Iterator +from contextlib import contextmanager +from dataclasses import dataclass +from pathlib import Path +from queue import SimpleQueue +from types import MappingProxyType +from typing import Final, Literal +from urllib.parse import urlsplit + +import httpx +import psutil +import pytest +import yaml +from integration._support.client import JSON_OBJECT, Gateway, eventually, string_value +from integration._support.database import read_rows +from integration._support.process import graceful_stop_seconds, owned_proxy_process +from integration._support.responses_vendor import same_response +from integration._support.wire import Reply, Request, Wire, wire_server +from integration.providers._coralbricks import ( + API_KEY, + MARKER, + MODEL, + NO_CACHE, + PROVIDER, + SPEND_COLUMNS, + deployment, + marker_of, + marker_provider, + ready, +) +from pydantic import JsonValue + +Endpoint = Literal["chat", "messages", "responses"] +_STARTED_WORKER: Final = re.compile(r"Started server process \[(\d+)\]") + + +@dataclass(frozen=True, slots=True) +class Call: + endpoint: Endpoint + stream: bool + marker: str + + +@dataclass(frozen=True, slots=True) +class Served: + call: Call + status: int + text: str + call_id: str + + +def path_of(endpoint: Endpoint) -> str: + match endpoint: + case "chat": + return "/v1/chat/completions" + case "messages": + return "/v1/messages" + case "responses": + return "/v1/responses" + + +def body_of(model: str, call: Call) -> dict[str, JsonValue]: + question: Final = f"Say hello marker-{call.marker}" + common: Final[dict[str, JsonValue]] = {"model": model, "stream": call.stream, "cache": NO_CACHE} + match call.endpoint: + case "chat": + return {**common, "messages": [{"role": "user", "content": question}]} + case "messages": + return {**common, "max_tokens": 64, "messages": [{"role": "user", "content": question}]} + case "responses": + return {**common, "input": question} + + +def calls_of(count: int, endpoints: tuple[Endpoint, ...], stream: Callable[[int], bool]) -> tuple[Call, ...]: + return tuple( + Call(endpoint=endpoints[index % len(endpoints)], stream=stream(index), marker=uuid.uuid4().hex) + for index in range(count) + ) + + +async def send(client: httpx.AsyncClient, key: str, model: str, call: Call) -> Served: + async with client.stream( + "POST", + path_of(call.endpoint), + json=body_of(model, call), + headers={"Authorization": f"Bearer {key}", "anthropic-version": "2023-06-01"}, + ) as response: + raw: Final = await response.aread() + return Served( + call=call, status=response.status_code, text=raw.decode(), call_id=response.headers.get("x-litellm-call-id", "") + ) + + +async def burst( + base_url: str, key: str, model: str, calls: tuple[Call, ...], *, tolerate_transport_errors: bool = False +) -> tuple[Served, ...]: + async with httpx.AsyncClient(base_url=base_url, timeout=90, trust_env=False) as client: + results: Final = await asyncio.gather( + *(send(client, key, model, call) for call in calls), return_exceptions=tolerate_transport_errors + ) + for result in results: + assert not isinstance(result, BaseException) or isinstance(result, httpx.TransportError), repr(result) + return tuple(result for result in results if isinstance(result, Served)) + + +def assert_answered_with_its_own_marker(served: Served) -> None: + assert served.status == 200, served.text + assert set(MARKER.findall(served.text)) == {served.call.marker}, served.text + + +def expected_row_id(call: Call) -> str: + match call.endpoint: + case "chat": + return f"req_{call.marker}" + case "messages": + return f"msg_{call.marker}" + case "responses": + return f"resp_{call.marker}" + + +def landed_once(model_group: str, served: tuple[Served, ...]) -> None: + expected: Final = tuple(expected_row_id(item.call) for item in served) + rows: Final = eventually( + lambda: read_rows( + f"{SPEND_COLUMNS} WHERE model_group=%s AND status='success' AND request_id NOT LIKE 'probe_%%'", + (model_group,), + ), + lambda found: len(found) >= len(expected), + seconds=90, + ) + identities: Final = [string_value(row["request_id"]) for row in rows] + for identity in expected: + matches: Final = [found for found in identities if same_response(found, identity)] + assert len(matches) == 1, (identity, matches) + assert len(rows) == len(expected), (len(rows), len(expected)) + assert {row["model"] for row in rows} == {f"{PROVIDER}/{MODEL}"}, rows + assert {row["custom_llm_provider"] for row in rows} == {PROVIDER}, rows + + +def failed_once(model_group: str, served: tuple[Served, ...]) -> None: + call_ids: Final = tuple(item.call_id for item in served) + assert all(call_ids), [item.call_id for item in served] + rows: Final = eventually( + lambda: read_rows(f"{SPEND_COLUMNS} WHERE model_group=%s AND status='failure'", (model_group,)), + lambda found: {string_value(row["request_id"]) for row in found} >= set(call_ids), + seconds=90, + ) + identities: Final = [string_value(row["request_id"]) for row in rows] + assert sorted(identities) == sorted(call_ids), (sorted(identities), sorted(call_ids)) + + +def assert_no_bleed(received: tuple[Request, ...], markers: frozenset[str]) -> None: + forwarded: Final = [marker_of(request) for request in received if marker_of(request) is not None] + assert sorted(marker for marker in forwarded if marker is not None) == sorted(markers), forwarded + + +@contextmanager +def released_on_exit(release: threading.Event) -> Iterator[threading.Event]: + try: + yield release + finally: + release.set() + + +def held_provider(release: threading.Event, held: SimpleQueue[str]) -> Callable[[Request], Reply]: + respond: Final = marker_provider() + + def hold(request: Request) -> Reply: + marker: Final = marker_of(request) + if marker is None: + return respond(request) + held.put(marker) + assert release.wait(timeout=90), "The burst was never released" + return respond(request) + + return hold + + +def health_of(gateway: Gateway, model: str) -> tuple[int, int, int, str]: + response: Final = gateway.request("GET", "/health", params={"model": model}) + health: Final = JSON_OBJECT.validate_json(response.content) + healthy_count: Final = health["healthy_count"] + unhealthy_count: Final = health["unhealthy_count"] + assert isinstance(healthy_count, int) and isinstance(unhealthy_count, int), response.text + return response.status_code, healthy_count, unhealthy_count, response.text + + +def assert_unhealthy(gateway: Gateway, model: str) -> None: + observed: Final = eventually( + lambda: health_of(gateway, model), lambda found: found[:3] == (503, 0, 1), seconds=30, return_last_on_timeout=True + ) + assert observed[:3] == (503, 0, 1), observed[3] + + +def assert_healthy(gateway: Gateway, model: str) -> None: + observed: Final = eventually( + lambda: health_of(gateway, model), lambda found: found[:3] == (200, 1, 0), seconds=30, return_last_on_timeout=True + ) + assert observed[:3] == (200, 1, 0), observed[3] + + +@pytest.mark.timeout(400) +async def test_provider_outage_mid_burst_fails_fast_reports_unhealthy_and_recovers(gateway: Gateway) -> None: + before: Final = calls_of(15, ("chat", "messages", "responses"), lambda index: index % 2 == 0) + during: Final = calls_of(15, ("chat", "messages", "responses"), lambda index: index % 2 == 1) + after: Final = calls_of(15, ("chat", "messages", "responses"), lambda index: index % 2 == 0) + release: Final = threading.Event() + held: Final[SimpleQueue[str]] = SimpleQueue() + with gateway.scenario() as scenario: + with wire_server(held_provider(release, held)) as wire, released_on_exit(release): + port: Final = int(urlsplit(wire.url).port or 0) + alias: Final = deployment(scenario, wire) + wire.drain() + first: Final = asyncio.create_task(burst(str(gateway.client.base_url), gateway.key, alias, before)) + await asyncio.to_thread(eventually, held.qsize, lambda size: size == 15, 60) + release.set() + served_before: Final = await first + assert len(served_before) == 15 + for item in served_before: + assert_answered_with_its_own_marker(item) + assert_no_bleed(wire.drain(), frozenset(call.marker for call in before)) + served_during: Final = await burst(str(gateway.client.base_url), gateway.key, alias, during) + assert len(served_during) == 15 + for item in served_during: + assert item.status == 500, item.text + assert "CoralbricksException" in item.text, item.text + assert "Connection error" in item.text or "Cannot connect to host" in item.text, item.text + assert "marker-" not in item.text, item.text + assert_unhealthy(gateway, alias) + with wire_server(marker_provider(), port=port) as restarted: + served_after: Final = await burst(str(gateway.client.base_url), gateway.key, alias, after) + assert len(served_after) == 15 + for item in served_after: + assert_answered_with_its_own_marker(item) + assert_no_bleed(restarted.drain(), frozenset(call.marker for call in after)) + assert_healthy(gateway, alias) + landed_once(alias, (*served_before, *served_after)) + failed_once(alias, served_during) + + +@pytest.mark.timeout(300) +async def test_slow_provider_streams_are_forwarded_once_under_concurrency(gateway: Gateway) -> None: + calls: Final = calls_of(20, ("chat", "messages", "responses"), lambda _: True) + with wire_server(marker_provider(stream_pause=0.3)) as wire, gateway.scenario() as scenario: + alias: Final = deployment(scenario, wire) + wire.drain() + served: Final = await burst(str(gateway.client.base_url), gateway.key, alias, calls) + assert len(served) == 20 + for item in served: + assert_answered_with_its_own_marker(item) + if item.call.endpoint == "chat": + assert item.text.rstrip().endswith("data: [DONE]"), item.text + assert_no_bleed(wire.drain(), frozenset(call.marker for call in calls)) + landed_once(alias, served) + + +def chaos_model() -> str: + return f"coralbricks-chaos-{uuid.uuid4().hex[:12]}" + + +def chaos_config(wire: Wire, directory: Path, model_name: str) -> Path: + base: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + config: Final = { + **base, + "model_list": [ + { + "model_name": model_name, + "litellm_params": {"model": f"{PROVIDER}/{MODEL}", "api_base": f"{wire.url}/v1", "api_key": API_KEY}, + } + ], + } + path: Final = directory / "coralbricks-chaos.yaml" + path.write_text(yaml.safe_dump(config)) + return path + + +def open_upstream_connections(pid: int, upstream: str) -> int: + port: Final = urlsplit(upstream).port + return sum( + 1 + for connection in psutil.Process(pid).net_connections(kind="tcp") + if connection.status == psutil.CONN_ESTABLISHED and connection.raddr and connection.raddr.port == port + ) + + +def worker_pids(log: Path) -> tuple[int, ...]: + return eventually( + lambda: tuple(int(pid) for pid in _STARTED_WORKER.findall(log.read_text())), + lambda pids: len(pids) == 2, + seconds=60, + ) + + +@pytest.mark.timeout(2 * graceful_stop_seconds() + 240) +async def test_worker_sigkill_mid_burst_leaves_the_sibling_serving_coralbricks(gateway: Gateway, tmp_path: Path) -> None: + calls: Final = calls_of(20, ("chat", "messages", "responses"), lambda _: False) + release: Final = threading.Event() + held: Final[SimpleQueue[str]] = SimpleQueue() + alias: Final = chaos_model() + with wire_server(held_provider(release, held)) as wire, released_on_exit(release): + config: Final = chaos_config(wire, tmp_path, alias) + with owned_proxy_process(gateway, tmp_path, {}, config=config, workers=2) as owned: + candidate: Final = owned.gateway + workers: Final = worker_pids(owned.log) + ready(candidate, alias, seconds=120) + wire.drain() + task: Final = asyncio.create_task( + burst(str(candidate.client.base_url), candidate.key, alias, calls, tolerate_transport_errors=True) + ) + await asyncio.to_thread(eventually, held.qsize, lambda size: size == 20, 60) + held_by: Final = MappingProxyType({pid: open_upstream_connections(pid, wire.url) for pid in workers}) + assert sum(held_by.values()) == 20, held_by + victim_pid, survivor_pid = sorted(workers, key=held_by.__getitem__) + with owned.replacing_worker(victim_pid) as kill: + kill.freeze() + await asyncio.to_thread(kill) + release.set() + served: Final = await task + assert held_by[survivor_pid] >= 10, held_by + assert len(served) == held_by[survivor_pid], (held_by, len(served)) + for item in served: + assert_answered_with_its_own_marker(item) + follow_up: Final = Call(endpoint="chat", stream=False, marker=uuid.uuid4().hex) + (answered,) = await burst(str(candidate.client.base_url), candidate.key, alias, (follow_up,)) + assert_answered_with_its_own_marker(answered) + assert_no_bleed(wire.drain(), frozenset(call.marker for call in (*calls, follow_up))) + landed_once(alias, (*served, answered)) + eventually( + lambda: owned.log.read_text().count("Application startup complete."), + lambda started: started == 3, + seconds=graceful_stop_seconds(), + ) + + +@pytest.mark.timeout(480) +async def test_proxy_sigterm_mid_burst_then_reboot_serves_coralbricks_again(gateway: Gateway, tmp_path: Path) -> None: + calls: Final = calls_of(20, ("chat", "messages", "responses"), lambda index: index % 2 == 0) + release: Final = threading.Event() + held: Final[SimpleQueue[str]] = SimpleQueue() + alias: Final = chaos_model() + with wire_server(held_provider(release, held)) as wire, released_on_exit(release): + config: Final = chaos_config(wire, tmp_path, alias) + with owned_proxy_process(gateway, tmp_path, {}, config=config, workers=2) as owned: + worker_pids(owned.log) + ready(owned.gateway, alias, seconds=120) + wire.drain() + task: Final = asyncio.create_task( + burst( + str(owned.gateway.client.base_url), + owned.gateway.key, + alias, + calls, + tolerate_transport_errors=True, + ) + ) + await asyncio.to_thread(eventually, held.qsize, lambda size: size == 20, 60) + owned.begin_shutdown() + release.set() + served: Final = await task + answered: Final = tuple(item for item in served if item.status == 200) + for item in answered: + assert_answered_with_its_own_marker(item) + await asyncio.to_thread(eventually, owned.process.poll, lambda code: code is not None, 90) + assert_no_bleed(wire.drain(), frozenset(call.marker for call in calls)) + with owned_proxy_process(gateway, tmp_path, {}, config=config, workers=2) as rebooted: + worker_pids(rebooted.log) + ready(rebooted.gateway, alias, seconds=120) + wire.drain() + follow_up: Final = Call(endpoint="responses", stream=True, marker=uuid.uuid4().hex) + (recovered,) = await burst(str(rebooted.gateway.client.base_url), rebooted.gateway.key, alias, (follow_up,)) + assert_answered_with_its_own_marker(recovered) + assert_no_bleed(wire.drain(), frozenset({follow_up.marker})) + rows: Final = eventually( + lambda: read_rows( + f"{SPEND_COLUMNS} WHERE model_group=%s AND status='success' AND request_id NOT LIKE 'probe_%%'", + (alias,), + ), + lambda found: any(same_response(string_value(row["request_id"]), expected_row_id(follow_up)) for row in found), + seconds=90, + ) + identities: Final = [string_value(row["request_id"]) for row in rows] + assert len(identities) == len(set(identities)), identities + allowed: Final = (*(expected_row_id(call) for call in calls), expected_row_id(follow_up)) + for identity in identities: + assert any(same_response(identity, candidate) for candidate in allowed), identity + billed: Final = { + item.call.marker: [found for found in identities if same_response(found, expected_row_id(item.call))] + for item in answered + } + assert all(len(found) == 1 for found in billed.values()), billed diff --git a/tests/integration/providers/test_coralbricks_provider_env_wire.py b/tests/integration/providers/test_coralbricks_provider_env_wire.py new file mode 100644 index 00000000000..6d8cc766b62 --- /dev/null +++ b/tests/integration/providers/test_coralbricks_provider_env_wire.py @@ -0,0 +1,190 @@ +from __future__ import annotations + +import uuid +from collections.abc import Iterator +from dataclasses import dataclass +from pathlib import Path +from typing import Final + +import pytest +import yaml +from integration._support.client import JSON_OBJECT, Gateway, eventually, gateway_from_environment, string_value +from integration._support.process import owned_proxy_process +from integration._support.responses_vendor import same_response +from integration._support.wire import Request, Wire, wire_server +from integration.providers._coralbricks import ( + API_KEY, + MODEL, + NO_CACHE, + PROVIDER, + assert_billed, + is_readiness_probe, + marker_provider, + not_yet_routable, + raw_deployment, + ready, + spend_row, +) +from pydantic import JsonValue + +ENV_KEY: Final = "cb_env_integration_key" +CONFIG_MODEL: Final = "coralbricks-env-config" + + +@dataclass(frozen=True, slots=True) +class EnvRig: + gateway: Gateway + env_wire: Wire + + +def env_config(env_wire: Wire, directory: Path) -> Path: + base: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + config: Final = { + **base, + "model_list": [{"model_name": CONFIG_MODEL, "litellm_params": {"model": f"{PROVIDER}/{MODEL}"}}], + "environment_variables": { + **base.get("environment_variables", {}), + "CORALBRICKS_API_KEY": ENV_KEY, + "CORALBRICKS_API_BASE": f"{env_wire.url}/v1", + }, + } + path: Final = directory / "coralbricks-env.yaml" + path.write_text(yaml.safe_dump(config)) + return path + + +@pytest.fixture(scope="module") +def env_rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[EnvRig]: + directory: Final = tmp_path_factory.mktemp("coralbricks-env") + with gateway_from_environment() as rig_gateway, wire_server(marker_provider()) as env_wire: + config: Final = env_config(env_wire, directory) + with owned_proxy_process( + rig_gateway, directory, {}, remove_environment=("CORALBRICKS_API_KEY", "CORALBRICKS_API_BASE"), + config=config, workers=2, + ) as owned: + ready(owned.gateway, CONFIG_MODEL, seconds=120) + yield EnvRig(owned.gateway, env_wire) + + +def calls(wire: Wire) -> tuple[Request, ...]: + return tuple(request for request in wire.drain() if not is_readiness_probe(request)) + + +def body_for(target: str, model: str, marker: str) -> dict[str, JsonValue]: + question: Final = f"Say hello marker-{marker}" + match target: + case "/v1/responses": + return {"model": model, "input": question, "cache": NO_CACHE} + case "/v1/messages": + return {"model": model, "max_tokens": 64, "messages": [{"role": "user", "content": question}], "cache": NO_CACHE} + case _: + return {"model": model, "messages": [{"role": "user", "content": question}], "cache": NO_CACHE} + + +def only_call_with_key(wire: Wire, target: str, key: str) -> None: + received: Final = calls(wire) + assert [(request.method, request.target) for request in received] == [("POST", target)], [ + (request.method, request.target) for request in received + ] + assert received[0].headers.get("authorization") == f"Bearer {key}", received[0].headers + assert JSON_OBJECT.validate_json(received[0].body)["model"] == MODEL, received[0].body + + +def row_id(target: str, marker: str, served: dict[str, JsonValue]) -> str: + match target: + case "/v1/responses": + assert same_response(string_value(served["id"]), f"resp_{marker}"), served + return string_value(served["id"]) + case "/v1/messages": + assert served["id"] == f"msg_{marker}", served + return f"msg_{marker}" + case _: + assert served["id"] == f"req_{marker}", served + return f"req_{marker}" + + +@pytest.mark.parametrize( + ("target", "call_type"), + (("/v1/chat/completions", "acompletion"), ("/v1/responses", "aresponses"), ("/v1/messages", "anthropic_messages")), +) +def test_config_model_without_credentials_uses_the_env_key_and_base(env_rig: EnvRig, target: str, call_type: str) -> None: + marker: Final = uuid.uuid4().hex + env_rig.env_wire.drain() + served: Final = env_rig.gateway.post(target, body_for(target, CONFIG_MODEL, marker)) + only_call_with_key(env_rig.env_wire, target, ENV_KEY) + assert_billed(spend_row(row_id(target, marker, served)), CONFIG_MODEL, call_type) + + +def test_explicit_credentials_beat_the_env_credentials(env_rig: EnvRig) -> None: + marker: Final = uuid.uuid4().hex + with wire_server(marker_provider()) as explicit_wire, env_rig.gateway.scenario() as scenario: + alias: Final = raw_deployment( + scenario, + f"integration-{uuid.uuid4().hex}", + {"model": f"{PROVIDER}/{MODEL}", "api_base": f"{explicit_wire.url}/v1", "api_key": API_KEY}, + ) + ready(env_rig.gateway, alias) + env_rig.env_wire.drain() + explicit_wire.drain() + served: Final = env_rig.gateway.post("/v1/chat/completions", body_for("/v1/chat/completions", alias, marker)) + only_call_with_key(explicit_wire, "/v1/chat/completions", API_KEY) + assert calls(env_rig.env_wire) == (), "the env base was called although the deployment names its own" + assert_billed(spend_row(row_id("/v1/chat/completions", marker, served)), alias, "acompletion") + + +@pytest.mark.parametrize("api_key", ("missing", None, ""), ids=("missing", "null", "empty")) +def test_deployment_without_a_key_uses_the_env_key(env_rig: EnvRig, api_key: str | None) -> None: + marker: Final = uuid.uuid4().hex + with wire_server(marker_provider()) as explicit_wire, env_rig.gateway.scenario() as scenario: + alias: Final = raw_deployment( + scenario, + f"integration-{uuid.uuid4().hex}", + { + "model": f"{PROVIDER}/{MODEL}", + "api_base": f"{explicit_wire.url}/v1", + **({} if api_key == "missing" else {"api_key": api_key}), + }, + ) + ready(env_rig.gateway, alias) + explicit_wire.drain() + served: Final = env_rig.gateway.post("/v1/chat/completions", body_for("/v1/chat/completions", alias, marker)) + only_call_with_key(explicit_wire, "/v1/chat/completions", ENV_KEY) + assert_billed(spend_row(row_id("/v1/chat/completions", marker, served)), alias, "acompletion") + + +@pytest.mark.parametrize("api_base", ("missing", None, ""), ids=("missing", "null", "empty")) +def test_deployment_without_a_base_uses_the_env_base(env_rig: EnvRig, api_base: str | None) -> None: + marker: Final = uuid.uuid4().hex + with env_rig.gateway.scenario() as scenario: + alias: Final = raw_deployment( + scenario, + f"integration-{uuid.uuid4().hex}", + { + "model": f"{PROVIDER}/{MODEL}", + "api_key": API_KEY, + **({} if api_base == "missing" else {"api_base": api_base}), + }, + ) + ready(env_rig.gateway, alias) + env_rig.env_wire.drain() + served: Final = env_rig.gateway.post("/v1/chat/completions", body_for("/v1/chat/completions", alias, marker)) + only_call_with_key(env_rig.env_wire, "/v1/chat/completions", API_KEY) + assert_billed(spend_row(row_id("/v1/chat/completions", marker, served)), alias, "acompletion") + + +def test_env_credentials_are_not_echoed_by_model_info(env_rig: EnvRig) -> None: + entries: Final = env_rig.gateway.get("/model/info")["data"] + assert isinstance(entries, list) + mine: Final = [JSON_OBJECT.validate_python(entry) for entry in entries] + config_rows: Final = [entry for entry in mine if entry["model_name"] == CONFIG_MODEL] + assert len(config_rows) == 1, [entry["model_name"] for entry in mine] + assert ENV_KEY not in str(config_rows[0]), config_rows[0] + response: Final = eventually( + lambda: env_rig.gateway.request("GET", "/health", params={"model": CONFIG_MODEL}), + lambda candidate: not not_yet_routable(candidate), + seconds=30, + ) + assert response.status_code == 200, response.text + assert ENV_KEY not in response.text, response.text + health: Final = JSON_OBJECT.validate_json(response.content) + assert health["healthy_count"] == 1 and health["unhealthy_count"] == 0, health diff --git a/tests/integration/providers/test_coralbricks_provider_wire.py b/tests/integration/providers/test_coralbricks_provider_wire.py new file mode 100644 index 00000000000..d8231df5151 --- /dev/null +++ b/tests/integration/providers/test_coralbricks_provider_wire.py @@ -0,0 +1,612 @@ +from __future__ import annotations + +import uuid +from itertools import count +from typing import Final + +import pytest +from integration._support.client import JSON_OBJECT, Gateway, eventually, object_value, string_value +from integration._support.responses_vendor import same_response +from integration._support.wire import Reply, Request, wire_server +from integration.providers._coralbricks import ( + ANSWER, + API_KEY, + COMPLETION_TOKENS, + MODEL, + MODELS, + NO_CACHE, + PROMPT_TOKENS, + PROVIDER, + USER_MESSAGES, + anthropic_client, + assert_billed, + assert_failed, + catalog_rows, + async_anthropic_client, + async_openai_client, + closed_port, + deployment, + expected_spend, + failing_provider, + group_rows, + is_readiness_probe, + not_yet_routable, + number, + only_call, + openai_client, + provider, + provider_calls, + rates_of, + raw_deployment, + row_of, + ready, + spend_row, +) +from pydantic import JsonValue + +DEPLOYMENT_RETRIES: Final = 2 + + +@pytest.mark.parametrize("model", MODELS) +def test_openai_sdk_chat_reaches_coralbricks_and_bills_the_cost_map_row(gateway: Gateway, model: str) -> None: + identity: Final = f"req_{uuid.uuid4().hex}" + with wire_server(provider(identity, model)) as wire, gateway.scenario() as scenario: + alias: Final = deployment(scenario, wire, model) + response: Final = openai_client(gateway).chat.completions.create( + model=alias, + messages=[{"role": "user", "content": "Say hello to the reef"}], + extra_body={"cache": NO_CACHE}, + ) + assert response.id == identity + assert response.choices[0].message.content == ANSWER + sent: Final = only_call(wire, "/v1/chat/completions", model) + assert sent["messages"] == USER_MESSAGES, sent + assert_billed(spend_row(identity), alias, "acompletion", model) + + +async def test_async_openai_sdk_chat_reaches_coralbricks(gateway: Gateway) -> None: + identity: Final = f"req_{uuid.uuid4().hex}" + with wire_server(provider(identity)) as wire, gateway.scenario() as scenario: + alias: Final = deployment(scenario, wire) + response: Final = await async_openai_client(gateway).chat.completions.create( + model=alias, + messages=[{"role": "user", "content": "Say hello to the reef"}], + extra_body={"cache": NO_CACHE}, + ) + assert response.id == identity + assert response.choices[0].message.content == ANSWER + assert only_call(wire, "/v1/chat/completions")["messages"] == USER_MESSAGES + assert_billed(spend_row(identity), alias, "acompletion") + + +def test_openai_sdk_chat_stream_is_forwarded_and_billed_from_the_usage_frame(gateway: Gateway) -> None: + identity: Final = f"req_{uuid.uuid4().hex}" + with wire_server(provider(identity)) as wire, gateway.scenario() as scenario: + alias: Final = deployment(scenario, wire) + chunks: Final = list( + openai_client(gateway).chat.completions.create( + model=alias, + messages=[{"role": "user", "content": "Say hello to the reef"}], + stream=True, + stream_options={"include_usage": True}, + extra_body={"cache": NO_CACHE}, + ) + ) + assert {chunk.id for chunk in chunks} == {identity} + assert "".join(chunk.choices[0].delta.content or "" for chunk in chunks if chunk.choices) == ANSWER + usage: Final = [chunk.usage for chunk in chunks if chunk.usage is not None] + assert len(usage) == 1 and usage[0].prompt_tokens == PROMPT_TOKENS, usage + sent: Final = only_call(wire, "/v1/chat/completions") + assert sent["stream"] is True, sent + assert_billed(spend_row(identity), alias, "acompletion") + + +async def test_async_openai_sdk_chat_stream_is_forwarded_and_billed(gateway: Gateway) -> None: + identity: Final = f"req_{uuid.uuid4().hex}" + with wire_server(provider(identity)) as wire, gateway.scenario() as scenario: + alias: Final = deployment(scenario, wire) + stream: Final = await async_openai_client(gateway).chat.completions.create( + model=alias, + messages=[{"role": "user", "content": "Say hello to the reef"}], + stream=True, + stream_options={"include_usage": True}, + extra_body={"cache": NO_CACHE}, + ) + chunks: Final = [chunk async for chunk in stream] + assert {chunk.id for chunk in chunks} == {identity} + assert "".join(chunk.choices[0].delta.content or "" for chunk in chunks if chunk.choices) == ANSWER + assert only_call(wire, "/v1/chat/completions")["stream"] is True + assert_billed(spend_row(identity), alias, "acompletion") + + +def test_raw_chat_request_carries_the_call_id_and_lands_once(gateway: Gateway) -> None: + identity: Final = f"req_{uuid.uuid4().hex}" + with wire_server(provider(identity)) as wire, gateway.scenario() as scenario: + alias: Final = deployment(scenario, wire) + response: Final = gateway.request( + "POST", "/v1/chat/completions", {"model": alias, "messages": USER_MESSAGES, "cache": NO_CACHE} + ) + assert response.status_code == 200, response.text + assert "x-litellm-call-id" in response.headers, response.headers + body: Final = JSON_OBJECT.validate_json(response.content) + assert body["id"] == identity, response.text + assert only_call(wire, "/v1/chat/completions")["messages"] == USER_MESSAGES + assert_billed(spend_row(identity), alias, "acompletion") + + +def test_openai_sdk_responses_go_natively_to_coralbricks_responses(gateway: Gateway) -> None: + identity: Final = f"resp_{uuid.uuid4().hex}" + with wire_server(provider(identity)) as wire, gateway.scenario() as scenario: + alias: Final = deployment(scenario, wire) + response: Final = openai_client(gateway).responses.create( + model=alias, input="Say hello to the reef", extra_body={"cache": NO_CACHE} + ) + assert response.output_text == ANSWER + assert same_response(response.id, identity), response.id + sent: Final = only_call(wire, "/v1/responses") + assert sent["input"] == "Say hello to the reef", sent + assert_billed(spend_row(response.id), alias, "aresponses") + + +async def test_async_openai_sdk_responses_go_natively_to_coralbricks_responses(gateway: Gateway) -> None: + identity: Final = f"resp_{uuid.uuid4().hex}" + with wire_server(provider(identity)) as wire, gateway.scenario() as scenario: + alias: Final = deployment(scenario, wire) + response: Final = await async_openai_client(gateway).responses.create( + model=alias, input="Say hello to the reef", extra_body={"cache": NO_CACHE} + ) + assert response.output_text == ANSWER + assert same_response(response.id, identity), response.id + assert only_call(wire, "/v1/responses")["input"] == "Say hello to the reef" + assert_billed(spend_row(response.id), alias, "aresponses") + + +def test_openai_sdk_responses_stream_is_forwarded_natively_and_billed_once(gateway: Gateway) -> None: + identity: Final = f"resp_{uuid.uuid4().hex}" + with wire_server(provider(identity)) as wire, gateway.scenario() as scenario: + alias: Final = deployment(scenario, wire) + events: Final = list( + openai_client(gateway).responses.create( + model=alias, input="Say hello to the reef", stream=True, extra_body={"cache": NO_CACHE} + ) + ) + assert [event.type for event in events] == [ + "response.created", + "response.output_text.delta", + "response.completed", + ], [event.type for event in events] + completed: Final = events[-1] + assert completed.type == "response.completed" + assert completed.response.output_text == ANSWER + sent: Final = only_call(wire, "/v1/responses") + assert sent["stream"] is True and sent["input"] == "Say hello to the reef", sent + rows: Final = group_rows(alias, 1) + assert same_response(completed.response.id, string_value(rows[0]["request_id"])), rows + assert_billed(rows[0], alias, "aresponses") + + +async def test_async_openai_sdk_responses_stream_is_forwarded_natively_and_billed_once(gateway: Gateway) -> None: + identity: Final = f"resp_{uuid.uuid4().hex}" + with wire_server(provider(identity)) as wire, gateway.scenario() as scenario: + alias: Final = deployment(scenario, wire) + stream: Final = await async_openai_client(gateway).responses.create( + model=alias, input="Say hello to the reef", stream=True, extra_body={"cache": NO_CACHE} + ) + events: Final = [event async for event in stream] + assert [event.type for event in events] == [ + "response.created", + "response.output_text.delta", + "response.completed", + ], [event.type for event in events] + completed: Final = events[-1] + assert completed.type == "response.completed" + assert completed.response.output_text == ANSWER + assert only_call(wire, "/v1/responses")["stream"] is True + rows: Final = group_rows(alias, 1) + assert same_response(completed.response.id, string_value(rows[0]["request_id"])), rows + assert_billed(rows[0], alias, "aresponses") + + +def test_raw_responses_request_lands_once(gateway: Gateway) -> None: + identity: Final = f"resp_{uuid.uuid4().hex}" + with wire_server(provider(identity)) as wire, gateway.scenario() as scenario: + alias: Final = deployment(scenario, wire) + response: Final = gateway.request( + "POST", "/v1/responses", {"model": alias, "input": "Say hello to the reef", "cache": NO_CACHE} + ) + assert response.status_code == 200, response.text + body: Final = JSON_OBJECT.validate_json(response.content) + assert same_response(string_value(body["id"]), identity), response.text + assert only_call(wire, "/v1/responses")["input"] == "Say hello to the reef" + assert_billed(spend_row(string_value(body["id"])), alias, "aresponses") + + +def test_anthropic_sdk_messages_go_natively_to_coralbricks_messages(gateway: Gateway) -> None: + identity: Final = f"msg_{uuid.uuid4().hex}" + with wire_server(provider(identity)) as wire, gateway.scenario() as scenario: + alias: Final = deployment(scenario, wire) + response: Final = anthropic_client(gateway).messages.create( + model=alias, + max_tokens=64, + messages=[{"role": "user", "content": "Say hello to the reef"}], + extra_body={"cache": NO_CACHE}, + ) + assert response.id == identity + assert response.content[0].type == "text" and response.content[0].text == ANSWER + assert only_call(wire, "/v1/messages")["messages"] == USER_MESSAGES + assert_billed(spend_row(identity), alias, "anthropic_messages") + + +async def test_async_anthropic_sdk_messages_go_natively_to_coralbricks_messages(gateway: Gateway) -> None: + identity: Final = f"msg_{uuid.uuid4().hex}" + with wire_server(provider(identity)) as wire, gateway.scenario() as scenario: + alias: Final = deployment(scenario, wire) + response: Final = await async_anthropic_client(gateway).messages.create( + model=alias, + max_tokens=64, + messages=[{"role": "user", "content": "Say hello to the reef"}], + extra_body={"cache": NO_CACHE}, + ) + assert response.id == identity + assert response.content[0].type == "text" and response.content[0].text == ANSWER + assert only_call(wire, "/v1/messages")["messages"] == USER_MESSAGES + assert_billed(spend_row(identity), alias, "anthropic_messages") + + +def test_anthropic_sdk_messages_stream_is_forwarded_natively_and_billed_once(gateway: Gateway) -> None: + identity: Final = f"msg_{uuid.uuid4().hex}" + with wire_server(provider(identity)) as wire, gateway.scenario() as scenario: + alias: Final = deployment(scenario, wire) + with anthropic_client(gateway).messages.stream( + model=alias, + max_tokens=64, + messages=[{"role": "user", "content": "Say hello to the reef"}], + extra_body={"cache": NO_CACHE}, + ) as stream: + text: Final = "".join(stream.text_stream) + final: Final = stream.get_final_message() + assert text == ANSWER + assert final.id == identity and final.usage.output_tokens == COMPLETION_TOKENS, final + assert only_call(wire, "/v1/messages")["stream"] is True + assert_billed(spend_row(identity), alias, "anthropic_messages") + + +async def test_async_anthropic_sdk_messages_stream_is_forwarded_natively_and_billed_once(gateway: Gateway) -> None: + identity: Final = f"msg_{uuid.uuid4().hex}" + with wire_server(provider(identity)) as wire, gateway.scenario() as scenario: + alias: Final = deployment(scenario, wire) + async with async_anthropic_client(gateway).messages.stream( + model=alias, + max_tokens=64, + messages=[{"role": "user", "content": "Say hello to the reef"}], + extra_body={"cache": NO_CACHE}, + ) as stream: + text: Final = "".join([piece async for piece in stream.text_stream]) + final: Final = await stream.get_final_message() + assert text == ANSWER + assert final.id == identity, final + assert only_call(wire, "/v1/messages")["stream"] is True + assert_billed(spend_row(identity), alias, "anthropic_messages") + + +def test_raw_messages_request_lands_once(gateway: Gateway) -> None: + identity: Final = f"msg_{uuid.uuid4().hex}" + with wire_server(provider(identity)) as wire, gateway.scenario() as scenario: + alias: Final = deployment(scenario, wire) + response: Final = gateway.request( + "POST", + "/v1/messages", + {"model": alias, "max_tokens": 64, "messages": USER_MESSAGES, "cache": NO_CACHE}, + ) + assert response.status_code == 200, response.text + assert JSON_OBJECT.validate_json(response.content)["id"] == identity, response.text + assert only_call(wire, "/v1/messages")["messages"] == USER_MESSAGES + assert_billed(spend_row(identity), alias, "anthropic_messages") + + +def request_body(target: str, alias: str) -> dict[str, JsonValue]: + match target: + case "/v1/responses": + return {"model": alias, "input": "Say hello to the reef", "cache": NO_CACHE} + case "/v1/messages": + return {"model": alias, "max_tokens": 64, "messages": USER_MESSAGES, "cache": NO_CACHE} + case _: + return {"model": alias, "messages": USER_MESSAGES, "cache": NO_CACHE} + + +@pytest.mark.parametrize( + ("target", "status"), + (("/v1/chat/completions", 401), ("/v1/chat/completions", 404), ("/v1/responses", 401), ("/v1/messages", 401)), +) +def test_provider_client_error_reaches_the_caller_after_one_attempt(gateway: Gateway, target: str, status: int) -> None: + message: Final = f"scripted coralbricks {status} {uuid.uuid4().hex}" + with wire_server(failing_provider(status, message)) as wire, gateway.scenario() as scenario: + alias: Final = scenario.model(model=f"{PROVIDER}/{MODEL}", api_base=f"{wire.url}/v1", api_key=API_KEY) + probe: Final = eventually( + lambda: gateway.request("POST", "/v1/chat/completions", {"model": alias, "messages": USER_MESSAGES}), + lambda response: not not_yet_routable(response), + seconds=60, + ) + assert probe.status_code == status, probe.text + wire.drain() + response: Final = gateway.request("POST", target, request_body(target, alias)) + assert response.status_code == status, response.text + assert message in response.text, response.text + attempts: Final = [(request.method, request.target) for request in wire.drain()] + assert attempts == [("POST", target)], attempts + assert_failed(spend_row(response.headers["x-litellm-call-id"]), alias) + assert gateway.request("GET", "/health/liveliness").status_code == 200 + + +@pytest.mark.parametrize("status", (429, 500)) +def test_provider_server_error_is_retried_then_reaches_the_caller(gateway: Gateway, status: int) -> None: + message: Final = f"scripted coralbricks {status} {uuid.uuid4().hex}" + with wire_server(failing_provider(status, message)) as wire, gateway.scenario() as scenario: + alias: Final = scenario.model( + model=f"{PROVIDER}/{MODEL}", api_base=f"{wire.url}/v1", api_key=API_KEY, num_retries=DEPLOYMENT_RETRIES + ) + eventually( + lambda: gateway.request("POST", "/v1/chat/completions", {"model": alias, "messages": USER_MESSAGES}), + lambda response: not not_yet_routable(response), + seconds=60, + ) + wire.drain() + response: Final = gateway.request( + "POST", "/v1/chat/completions", {"model": alias, "messages": USER_MESSAGES, "cache": NO_CACHE} + ) + assert response.status_code == status, response.text + assert message in response.text, response.text + attempts: Final = [(request.method, request.target) for request in wire.drain()] + assert attempts == [("POST", "/v1/chat/completions")] * (1 + DEPLOYMENT_RETRIES), attempts + assert_failed(spend_row(response.headers["x-litellm-call-id"]), alias) + + +def test_unreachable_api_base_reports_the_connection_error_and_leaves_the_proxy_serving(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + control: Final = scenario.model() + alias: Final = scenario.model( + model=f"{PROVIDER}/{MODEL}", api_base=f"http://127.0.0.1:{closed_port()}/v1", api_key=API_KEY + ) + response: Final = eventually( + lambda: gateway.request( + "POST", "/v1/chat/completions", {"model": alias, "messages": USER_MESSAGES, "cache": NO_CACHE} + ), + lambda candidate: not not_yet_routable(candidate), + seconds=60, + ) + assert response.status_code == 500, response.text + assert "Connection error" in response.text, response.text + assert_failed(spend_row(response.headers["x-litellm-call-id"]), alias) + assert gateway.request("GET", "/health/liveliness").status_code == 200 + ready(gateway, control) + assert gateway.chat(control, text="still serving")["object"] == "chat.completion" + + +@pytest.mark.parametrize("api_key", ("missing", None, ""), ids=("missing", "null", "empty")) +def test_deployment_without_a_key_and_no_env_key_fails_before_the_provider( + gateway: Gateway, api_key: str | None +) -> None: + with wire_server(provider("unused")) as wire, gateway.scenario() as scenario: + litellm_params: Final[dict[str, JsonValue]] = { + "model": f"{PROVIDER}/{MODEL}", + "api_base": f"{wire.url}/v1", + **({} if api_key == "missing" else {"api_key": api_key}), + } + alias: Final = raw_deployment(scenario, f"integration-{uuid.uuid4().hex}", litellm_params) + response: Final = eventually( + lambda: gateway.request( + "POST", "/v1/chat/completions", {"model": alias, "messages": USER_MESSAGES, "cache": NO_CACHE} + ), + lambda candidate: not not_yet_routable(candidate), + seconds=60, + ) + assert response.status_code == 500, response.text + assert "AuthenticationError" in response.text and "CORALBRICKS_API_KEY" in response.text, response.text + assert_failed(spend_row(response.headers["x-litellm-call-id"]), alias) + assert wire.drain() == (), "the provider was called without a key" + assert gateway.request("GET", "/health/liveliness").status_code == 200 + + +@pytest.mark.parametrize( + "model", + (7, ["coralbricks/deepseek-v4.1-flash-fast"], f"{PROVIDER}/" + "g" * 5120), + ids=("int", "list", "5kb"), +) +def test_hostile_model_field_is_rejected_without_crashing(gateway: Gateway, model: JsonValue) -> None: + for _ in range(2): + response: Final = gateway.request( + "POST", "/v1/chat/completions", {"model": model, "messages": USER_MESSAGES, "cache": NO_CACHE} + ) + assert response.status_code in (400, 422), response.text + assert gateway.request("GET", "/health/liveliness").status_code == 200 + + +def test_identical_uncached_chats_are_each_forwarded_and_billed(gateway: Gateway) -> None: + prefix: Final = f"req_{uuid.uuid4().hex}" + calls: Final = count(1) + + def respond(request: Request) -> Reply: + if is_readiness_probe(request): + return provider(f"{prefix}-probe")(request) + return provider(f"{prefix}-{next(calls)}")(request) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + alias: Final = deployment(scenario, wire) + wire.drain() + first: Final = gateway.post( + "/v1/chat/completions", {"model": alias, "messages": USER_MESSAGES, "cache": NO_CACHE} + ) + second: Final = gateway.post( + "/v1/chat/completions", {"model": alias, "messages": USER_MESSAGES, "cache": NO_CACHE} + ) + assert first["id"] != second["id"], (first["id"], second["id"]) + assert [(request.method, request.target) for request in provider_calls(wire)] == [ + ("POST", "/v1/chat/completions") + ] * 2 + rows: Final = group_rows(alias, 2) + assert {string_value(row["request_id"]) for row in rows} == {string_value(first["id"]), string_value(second["id"])} + for row in rows: + assert_billed(row, alias, "acompletion") + + +def test_model_info_reports_the_coralbricks_cost_map_row(gateway: Gateway) -> None: + rates: Final = rates_of(MODEL) + with wire_server(provider("unused")) as wire, gateway.scenario() as scenario: + alias: Final = scenario.model(model=f"{PROVIDER}/{MODEL}", api_base=f"{wire.url}/v1", api_key=API_KEY) + entries: Final = gateway.get("/model/info")["data"] + assert isinstance(entries, list) + mine: Final = [object_value(entry) for entry in entries if object_value(entry)["model_name"] == alias] + assert len(mine) == 1, [object_value(entry)["model_name"] for entry in entries] + litellm_params: Final = object_value(mine[0]["litellm_params"]) + info: Final = object_value(mine[0]["model_info"]) + assert litellm_params["model"] == f"{PROVIDER}/{MODEL}", litellm_params + assert info["litellm_provider"] == PROVIDER, info + assert number(info["input_cost_per_token"]) == pytest.approx(rates.input), info + assert number(info["output_cost_per_token"]) == pytest.approx(rates.output), info + assert number(info["cache_creation_input_token_cost"]) == pytest.approx(rates.cache_write), info + assert info["max_input_tokens"] == row_of(MODEL)["max_input_tokens"] and info["mode"] == "chat", info + assert info["supports_function_calling"] is True and info["supports_prompt_caching"] is True, info + + +def test_models_listing_and_model_group_info_carry_the_coralbricks_deployment(gateway: Gateway) -> None: + rates: Final = rates_of(MODEL) + with wire_server(provider("unused")) as wire, gateway.scenario() as scenario: + alias: Final = scenario.model(model=f"{PROVIDER}/{MODEL}", api_base=f"{wire.url}/v1", api_key=API_KEY) + listed: Final = gateway.get("/v1/models")["data"] + assert isinstance(listed, list) + assert alias in {object_value(entry)["id"] for entry in listed} + groups: Final = gateway.get("/model_group/info")["data"] + assert isinstance(groups, list) + mine: Final = [object_value(group) for group in groups if object_value(group)["model_group"] == alias] + assert len(mine) == 1, [object_value(group)["model_group"] for group in groups] + assert mine[0]["providers"] == [PROVIDER], mine[0] + assert number(mine[0]["input_cost_per_token"]) == pytest.approx(rates.input), mine[0] + assert number(mine[0]["output_cost_per_token"]) == pytest.approx(rates.output), mine[0] + + +def test_wildcard_deployment_lists_every_coralbricks_catalog_row(gateway: Gateway) -> None: + with wire_server(provider("unused")) as wire, gateway.scenario() as scenario: + created: Final = gateway.post( + "/model/new", + { + "model_name": f"{PROVIDER}/*", + "litellm_params": {"model": f"{PROVIDER}/*", "api_base": f"{wire.url}/v1", "api_key": API_KEY}, + }, + ) + scenario.cleanups.callback(scenario.delete_model, string_value(object_value(created["model_info"])["id"])) + catalog: Final = catalog_rows() + listed: Final = gateway.get("/v1/models")["data"] + assert isinstance(listed, list) + listed_names: Final = {object_value(entry)["id"] for entry in listed} + assert catalog <= listed_names, sorted(listed_names) + entries: Final = gateway.get("/model/info")["data"] + assert isinstance(entries, list) + expanded: Final = { + object_value(entry)["model_name"]: object_value(entry) + for entry in entries + if object_value(entry)["model_name"] in catalog + } + assert set(expanded) == catalog, [object_value(entry)["model_name"] for entry in entries] + for name, entry in expanded.items(): + assert object_value(entry["litellm_params"])["model"] == name, entry + assert object_value(entry["model_info"])["litellm_provider"] == PROVIDER, entry + + +def test_health_probe_reaches_coralbricks_with_the_deployment_key(gateway: Gateway) -> None: + identity: Final = f"req_{uuid.uuid4().hex}" + with wire_server(provider(identity)) as wire, gateway.scenario() as scenario: + alias: Final = deployment(scenario, wire) + wire.drain() + health: Final = gateway.get("/health", {"model": alias}) + assert health["healthy_count"] == 1 and health["unhealthy_count"] == 0, health + healthy: Final = health["healthy_endpoints"] + assert isinstance(healthy, list) and object_value(healthy[0])["model"] == f"{PROVIDER}/{MODEL}", health + probes: Final = provider_calls(wire) + assert [(request.method, request.target) for request in probes] == [("POST", "/v1/chat/completions")], probes + assert probes[0].headers.get("authorization") == f"Bearer {API_KEY}", probes[0].headers + assert JSON_OBJECT.validate_json(probes[0].body)["model"] == MODEL + + +def test_add_model_form_fields_name_the_coralbricks_credentials(gateway: Gateway) -> None: + response: Final = gateway.request("GET", "/public/providers/fields") + assert response.status_code == 200, response.text + providers: Final = response.json() + assert isinstance(providers, list), response.text + mine: Final = [object_value(entry) for entry in providers if object_value(entry)["litellm_provider"] == PROVIDER] + assert len(mine) == 1, [object_value(entry)["litellm_provider"] for entry in providers] + assert mine[0]["provider_display_name"] == "CoralBricks", mine[0] + fields: Final = mine[0]["credential_fields"] + assert isinstance(fields, list), mine[0] + by_key: Final = {string_value(object_value(field)["key"]): object_value(field) for field in fields} + assert set(by_key) == {"api_key", "api_base"}, sorted(by_key) + assert by_key["api_key"]["required"] is True and by_key["api_base"]["required"] is False, by_key + assert string_value(by_key["api_base"]["placeholder"]).endswith("/v1"), by_key["api_base"] + + +def test_key_scoped_to_the_coralbricks_deployment_is_charged_its_spend(gateway: Gateway) -> None: + identity: Final = f"req_{uuid.uuid4().hex}" + with wire_server(provider(identity)) as wire, gateway.scenario() as scenario: + alias: Final = deployment(scenario, wire) + token: Final = scenario.key(models=[alias], max_budget=1.0) + wire.drain() + response: Final = gateway.post( + "/v1/chat/completions", {"model": alias, "messages": USER_MESSAGES, "cache": NO_CACHE}, key=token + ) + assert response["id"] == identity, response + only_call(wire, "/v1/chat/completions") + assert_billed(spend_row(identity), alias, "acompletion") + info: Final = eventually( + lambda: object_value(gateway.get("/key/info", {"key": token})["info"]), + lambda found: number(found["spend"]) > 0, + seconds=70, + ) + assert number(info["spend"]) == pytest.approx(expected_spend(MODEL)), info + + +def test_wildcard_coralbricks_deployment_routes_the_requested_model(gateway: Gateway) -> None: + identity: Final = f"req_{uuid.uuid4().hex}" + with wire_server(provider(identity)) as wire, gateway.scenario() as scenario: + prefix: Final = f"integration-{uuid.uuid4().hex}" + raw_deployment( + scenario, + f"{prefix}/*", + {"model": f"{PROVIDER}/*", "api_base": f"{wire.url}/v1", "api_key": API_KEY}, + ) + requested: Final = f"{prefix}/{MODEL}" + response: Final = eventually( + lambda: gateway.request( + "POST", "/v1/chat/completions", {"model": requested, "messages": USER_MESSAGES, "cache": NO_CACHE} + ), + lambda candidate: not not_yet_routable(candidate), + seconds=60, + ) + assert response.status_code == 200, response.text + assert JSON_OBJECT.validate_json(response.content)["id"] == identity, response.text + only_call(wire, "/v1/chat/completions") + row: Final = spend_row(identity) + assert row["model"] == f"{PROVIDER}/{MODEL}" and row["custom_llm_provider"] == PROVIDER, row + assert number(row["spend"]) == pytest.approx(expected_spend(MODEL)), row + + +def test_fallback_from_a_refused_primary_lands_on_coralbricks(gateway: Gateway) -> None: + identity: Final = f"req_{uuid.uuid4().hex}" + with wire_server(provider(identity)) as wire, gateway.scenario() as scenario: + alias: Final = deployment(scenario, wire) + primary: Final = scenario.model( + model=f"{PROVIDER}/{MODEL}", api_base=f"http://127.0.0.1:{closed_port()}/v1", api_key=API_KEY + ) + wire.drain() + response: Final = eventually( + lambda: gateway.request( + "POST", + "/v1/chat/completions", + {"model": primary, "messages": USER_MESSAGES, "cache": NO_CACHE, "fallbacks": [alias]}, + ), + lambda candidate: not not_yet_routable(candidate), + seconds=60, + ) + assert response.status_code == 200, response.text + assert response.headers.get("x-litellm-attempted-fallbacks") == "1", response.headers + assert JSON_OBJECT.validate_json(response.content)["id"] == identity, response.text + only_call(wire, "/v1/chat/completions") + assert_billed(spend_row(identity), alias, "acompletion") diff --git a/tests/integration/providers/test_databricks_decisions_wire.py b/tests/integration/providers/test_databricks_decisions_wire.py index 24db0153bcd..d3b17044207 100644 --- a/tests/integration/providers/test_databricks_decisions_wire.py +++ b/tests/integration/providers/test_databricks_decisions_wire.py @@ -10,6 +10,7 @@ import pytest from integration._support.client import Gateway, Scenario, eventually, object_value, string_value from integration._support.database import read_rows from integration._support.process import graceful_stop_seconds, owned_proxy_process +from integration._support.timing import assert_within_factor_of_control from integration._support.upstream import ScenarioHandle, delete_scenario, register_scenario from integration.cost_calculation.cost_tracking_case import JsonResponse from pydantic import JsonValue @@ -111,10 +112,28 @@ def _upstream_calls(gateway: Gateway, handle: ScenarioHandle) -> list[dict[str, return _calls_to(_observed(gateway), handle) -def _timed_decide(gateway: Gateway, model: str) -> tuple[httpx.Response, float]: - started: Final = time.perf_counter() - response: Final = _decide(gateway, model) - return response, time.perf_counter() - started +def _timed_decide_under_liveliness_probes( + gateway: Gateway, model: str +) -> tuple[httpx.Response, float, tuple[tuple[int, float], ...]]: + liveliness: Final[list[tuple[int, float]]] = [] + stop: Final = threading.Event() + + def poll() -> None: + while not stop.is_set(): + started: Final = time.perf_counter() + probe: Final = gateway.client.get("/health/liveliness") + liveliness.append((probe.status_code, time.perf_counter() - started)) + + poller: Final = threading.Thread(target=poll) + poller.start() + try: + started: Final = time.perf_counter() + response: Final = _decide(gateway, model) + elapsed: Final = time.perf_counter() - started + finally: + stop.set() + poller.join() + return response, elapsed, tuple(liveliness) def _spend_row(call_id: str) -> dict[str, JsonValue]: @@ -623,23 +642,12 @@ def test_a_megabyte_model_name_is_refused_in_linear_time_while_liveliness_stays_ "/v1/systemone", json={"model": model, "state": _STATE, "questions": _QUESTIONS} ) assert anonymous.status_code == 401, anonymous.text[:300] - tenth_response, tenth_elapsed = _timed_decide(gateway, f"{prefix}/{_TENTH_MEGABYTE_NAME}") - liveliness: Final[list[tuple[int, float]]] = [] - stop: Final = threading.Event() - - def poll() -> None: - while not stop.is_set(): - started: Final = time.perf_counter() - probe: Final = gateway.client.get("/health/liveliness") - liveliness.append((probe.status_code, time.perf_counter() - started)) - - poller: Final = threading.Thread(target=poll) - poller.start() - response, elapsed = _timed_decide(gateway, model) - stop.set() - poller.join() + tenth_response, tenth_elapsed, _ = _timed_decide_under_liveliness_probes( + gateway, f"{prefix}/{_TENTH_MEGABYTE_NAME}" + ) + response, elapsed, liveliness = _timed_decide_under_liveliness_probes(gateway, model) assert tenth_response.status_code == 400, tenth_response.text[:300] - assert elapsed < 25 * tenth_elapsed, (elapsed, tenth_elapsed) + assert_within_factor_of_control(elapsed, control=tenth_elapsed, factor=25) assert liveliness and all(status == 200 for status, _ in liveliness), liveliness assert max(seconds for _, seconds in liveliness) < max(1.0, elapsed / 2), (elapsed, liveliness) assert response.status_code == 400, response.text[:300] diff --git a/tests/integration/providers/test_decisions_chaos.py b/tests/integration/providers/test_decisions_chaos.py index 63b4a482e92..d868e46e6a5 100644 --- a/tests/integration/providers/test_decisions_chaos.py +++ b/tests/integration/providers/test_decisions_chaos.py @@ -1,7 +1,6 @@ import asyncio import json import re -import signal import threading import uuid from dataclasses import dataclass @@ -215,12 +214,12 @@ async def test_worker_sigkill_mid_burst_leaves_the_sibling_serving_the_default_m held_by: Final = MappingProxyType({pid: _open_upstream_connections(pid, wire.url) for pid in workers}) assert sum(held_by.values()) == 20, held_by victim_pid, survivor_pid = sorted(workers, key=held_by.__getitem__) - victim: Final = psutil.Process(victim_pid) - victim.suspend() - victim.send_signal(signal.SIGKILL) - release.set() - served: Final = await burst - assert len(served) == held_by[survivor_pid], (held_by, len(served)) + with owned.replacing_worker(victim_pid) as kill: + kill.freeze() + await asyncio.to_thread(kill) + release.set() + served: Final = await burst + assert len(served) == held_by[survivor_pid], (held_by, len(served)) follow_up: Final = _Call(route="/systemone", marker=f"ok-{uuid.uuid4().hex}", fail=False) (answered,) = await _burst(base_url, candidate.key, None, (follow_up,)) await asyncio.to_thread( diff --git a/tests/integration/providers/test_decisions_wire.py b/tests/integration/providers/test_decisions_wire.py index 2862632d132..6d4e9ac8a30 100644 --- a/tests/integration/providers/test_decisions_wire.py +++ b/tests/integration/providers/test_decisions_wire.py @@ -99,6 +99,15 @@ _PROVIDERS: Final = ( _Provider( "hosted_vllm", "hosted_vllm/Qwen/Qwen3-0.6B", "/v1/systemone", "Qwen/Qwen3-0.6B", None, False, None ), + _Provider( + "fireworks_ai", + "fireworks_ai/accounts/fireworks/routers/rune", + "/v1/systemone", + "accounts/fireworks/routers/rune", + _API_KEY, + False, + "fireworks_ai/accounts/fireworks/routers/rune", + ), _Provider( "cloudflare", "cloudflare/clef", diff --git a/tests/integration/providers/test_device_code_login_guard_boot.py b/tests/integration/providers/test_device_code_login_guard_boot.py index 687a8f1d313..34ab7f1597f 100644 --- a/tests/integration/providers/test_device_code_login_guard_boot.py +++ b/tests/integration/providers/test_device_code_login_guard_boot.py @@ -2,7 +2,6 @@ import asyncio import hashlib import itertools import re -import signal import threading import uuid from collections.abc import Mapping @@ -422,12 +421,13 @@ async def test_worker_sigkill_mid_burst_leaves_the_sibling_refusing_logins_and_s owned, key, workers, control_port, held_markers, (), tasks ) victim_pid, survivor_pid = sorted(workers, key=held_by.__getitem__) - victim: Final = psutil.Process(victim_pid) - victim.suspend() - victim.send_signal(signal.SIGKILL) + with owned.replacing_worker(victim_pid) as kill: + kill.freeze() + await asyncio.to_thread(kill) + release.set() + answers: Final = await asyncio.gather(*bursts) finally: release.set() - answers: Final = await asyncio.gather(*bursts) served: Final = tuple(item for item in itertools.chain.from_iterable(answers) if item is not None) assert len(served) == held_by[survivor_pid], (held_by, len(served)) for item in served: diff --git a/tests/integration/providers/test_device_code_login_guard_wire.py b/tests/integration/providers/test_device_code_login_guard_wire.py index 072de4f3782..c88bd0b9273 100644 --- a/tests/integration/providers/test_device_code_login_guard_wire.py +++ b/tests/integration/providers/test_device_code_login_guard_wire.py @@ -24,7 +24,9 @@ from integration._support import responses_vendor as rv from integration._support.client import eventually, gateway_from_environment from integration._support.database import read_rows from integration._support.process import OwnedProxy, owned_proxy_process +from integration._support.timing import assert_cut_short, timed from integration._support.wire import Request +from litellm.llms.chatgpt.authenticator import DEVICE_CODE_COOLDOWN_SECONDS from pydantic import JsonValue pytestmark: Final = pytest.mark.timeout(300) @@ -713,16 +715,18 @@ def test_a_token_mounted_after_a_refusal_serves_without_a_restart_and_removing_i _COOLDOWN_SECONDS_LEFT: Final = 45 +_REFUSAL_MARGIN_SECONDS: Final = 15 def test_a_recent_device_code_request_does_not_hold_the_worker_for_the_rest_of_the_cooldown(scene: _Scene) -> None: - dl.write_chatgpt(scene.rig.chatgpt_dir, {"device_code_requested_at": time.time() - (300 - _COOLDOWN_SECONDS_LEFT)}) - started: Final = time.monotonic() - served: Final = _serve(scene.call(_Cell("chatgpt", "chat", False, "httpx"))) - elapsed: Final = time.monotonic() - started + requested_at: Final = time.time() - (DEVICE_CODE_COOLDOWN_SECONDS - _COOLDOWN_SECONDS_LEFT) + dl.write_chatgpt(scene.rig.chatgpt_dir, {"device_code_requested_at": requested_at}) + remaining: Final = requested_at + DEVICE_CODE_COOLDOWN_SECONDS - time.time() + refusal: Final = timed(lambda: _serve(scene.call(_Cell("chatgpt", "chat", False, "httpx")))) + served: Final = refusal.value assert served.status == 400, served.text assert dl.CHATGPT_REFUSAL in _error_message(served.text), served.text - assert elapsed < _COOLDOWN_SECONDS_LEFT - 15, elapsed + assert_cut_short(refusal.seconds, injected_delay=remaining - _REFUSAL_MARGIN_SECONDS) assert scene.rig.peers.auth_connections() == () diff --git a/tests/integration/providers/test_gemini_embedding_file_block_chaos.py b/tests/integration/providers/test_gemini_embedding_file_block_chaos.py index cee0758b433..17483147107 100644 --- a/tests/integration/providers/test_gemini_embedding_file_block_chaos.py +++ b/tests/integration/providers/test_gemini_embedding_file_block_chaos.py @@ -1,8 +1,6 @@ from __future__ import annotations import json -import os -import signal import threading import uuid import zlib @@ -192,15 +190,17 @@ def test_block_embeddings_survive_a_worker_kill(gateway: Gateway, tmp_path: Path victim: Final = min(workers) assert outcome(candidate, name, 0).startswith("ok:") - def attempt_around_the_kill(index: int) -> str: - if index == 2: - os.kill(victim, signal.SIGKILL) - return outcome(candidate, name, index) + with owned.replacing_worker(victim) as kill: - with ThreadPoolExecutor(max_workers=BURST) as pool: - outcomes: Final = tuple(pool.map(attempt_around_the_kill, range(1, BURST + 1))) - assert outcomes.count("ok") == 0 and any(item.startswith("ok:") for item in outcomes), outcomes - assert all(item.startswith("ok:") or item.startswith("transport:") for item in outcomes), outcomes + def attempt_around_the_kill(index: int) -> str: + if index == 2: + kill() + return outcome(candidate, name, index) + + with ThreadPoolExecutor(max_workers=BURST) as pool: + outcomes: Final = tuple(pool.map(attempt_around_the_kill, range(1, BURST + 1))) + assert outcomes.count("ok") == 0 and any(item.startswith("ok:") for item in outcomes), outcomes + assert all(item.startswith("ok:") or item.startswith("transport:") for item in outcomes), outcomes respawned: Final = eventually( lambda: worker_pids(owned.process.pid), lambda pids: len(pids) == 2 and victim not in pids, diff --git a/tests/integration/providers/test_gemini_omitted_temperature.py b/tests/integration/providers/test_gemini_omitted_temperature.py new file mode 100644 index 00000000000..fa8a9c890c6 --- /dev/null +++ b/tests/integration/providers/test_gemini_omitted_temperature.py @@ -0,0 +1,83 @@ +"""A Gemini 3+ request that omits temperature reaches the provider without one: streamed on each endpoint, and in +a burst of mixed-endpoint requests spread across the proxy's workers.""" + +import json +from collections.abc import Mapping +from concurrent.futures import ThreadPoolExecutor +from typing import Final + +import pytest +from integration._support.client import JSON_OBJECT, Gateway, object_value +from integration._support.provider import SharedProvider +from integration._support.wire import Reply, Request +from integration.translation.case import TranslationTestCase +from integration.translation.chat_completions.bases.gemini import ( + GEMINI_3_8_FLASH_TEST_CASE as CHAT_COMPLETIONS_TEST_CASE, +) +from integration.translation.messages.bases.gemini import GEMINI_3_8_FLASH_TEST_CASE as MESSAGES_TEST_CASE +from integration.translation.responses.bases.gemini import GEMINI_3_8_FLASH_TEST_CASE as RESPONSES_TEST_CASE +from pydantic import JsonValue + +GENERATE_TARGET: Final = "/models/gemini-3.8-flash:generateContent" +STREAM_TARGET: Final = "/models/gemini-3.8-flash:streamGenerateContent?alt=sse" +EXPECTED_GENERATION_CONFIG: Final = {"max_output_tokens": 1024} +ENDPOINT_CASES: Final = (CHAT_COMPLETIONS_TEST_CASE, MESSAGES_TEST_CASE, RESPONSES_TEST_CASE) +BURST: Final = ENDPOINT_CASES * 4 +STREAM_USAGE: Final = { + "promptTokenCount": 10, + "candidatesTokenCount": 2, + "totalTokenCount": 12, + "promptTokensDetails": [{"modality": "TEXT", "tokenCount": 10}], +} + + +def _frame(payload: Mapping[str, JsonValue]) -> bytes: + return f"data: {json.dumps(payload)}\n\n".encode() + + +def _streamed_reply() -> Reply: + opening: Final = { + "candidates": [{"content": {"role": "model", "parts": [{"text": "Hello."}]}, "index": 0}], + "modelVersion": "gemini-3.8-flash", + "responseId": "stream-gemini-3-8-flash", + } + closing: Final = { + "candidates": [{"content": {"role": "model", "parts": [{"text": ""}]}, "finishReason": "STOP", "index": 0}], + "usageMetadata": STREAM_USAGE, + "modelVersion": "gemini-3.8-flash", + "responseId": "stream-gemini-3-8-flash", + } + return Reply(content_type="text/event-stream", chunks=(_frame(opening), _frame(closing))) + + +def _generation_config(request: Request) -> dict[str, JsonValue]: + return object_value(JSON_OBJECT.validate_json(request.body)["generationConfig"]) + + +@pytest.mark.parametrize("case", ENDPOINT_CASES, ids=lambda case: case.litellm_endpoint) +def test_streamed_request_without_temperature_forwards_none( + case: TranslationTestCase, gateway: Gateway, provider: SharedProvider +) -> None: + provider.expect(_streamed_reply()) + response: Final = gateway.request("POST", case.litellm_endpoint, {**case.litellm_request, "stream": True}) + assert response.status_code == 200, response.text + assert "Hello." in response.text, response.text + received: Final = provider.received() + assert [(request.method, request.target) for request in received] == [("POST", STREAM_TARGET)] + assert _generation_config(received[0]) == EXPECTED_GENERATION_CONFIG + + +def test_concurrent_mixed_endpoint_burst_without_temperature_forwards_none( + gateway: Gateway, provider: SharedProvider +) -> None: + reply_body: Final = json.dumps(CHAT_COMPLETIONS_TEST_CASE.mock_provider_response).encode() + provider.expect(*(Reply(body=reply_body) for _ in BURST)) + with ThreadPoolExecutor(max_workers=len(BURST)) as pool: + responses: Final = tuple( + pool.map(lambda case: gateway.request("POST", case.litellm_endpoint, case.litellm_request), BURST) + ) + assert [response.status_code for response in responses] == [200] * len(BURST), [r.text for r in responses] + assert all("Hello." in response.text for response in responses), [r.text for r in responses] + received: Final = provider.received() + assert [(request.method, request.target) for request in received] == [("POST", GENERATE_TARGET)] * len(BURST) + assert [_generation_config(request) for request in received] == [EXPECTED_GENERATION_CONFIG] * len(BURST) diff --git a/tests/integration/providers/test_gemini_thinking_replay_chaos.py b/tests/integration/providers/test_gemini_thinking_replay_chaos.py index 57058eb5fbe..f8bb52ef60a 100644 --- a/tests/integration/providers/test_gemini_thinking_replay_chaos.py +++ b/tests/integration/providers/test_gemini_thinking_replay_chaos.py @@ -1,7 +1,6 @@ import asyncio import json import re -import signal import threading import uuid from collections.abc import Callable, Mapping @@ -327,15 +326,15 @@ async def test_worker_sigkill_mid_burst_leaves_the_sibling_replaying_thinking_si held_by: Final = MappingProxyType({pid: _open_upstream_connections(pid, wire.url) for pid in workers}) assert sum(held_by.values()) == 20, held_by victim_pid, survivor_pid = sorted(workers, key=held_by.__getitem__) - victim: Final = psutil.Process(victim_pid) - victim.suspend() - victim.send_signal(signal.SIGKILL) - release.set() - served: Final = await burst - assert held_by[survivor_pid] >= 10, held_by - assert len(served) == held_by[survivor_pid], (held_by, len(served)) - for item in served: - _assert_answered_with_its_own_marker(item) + with owned.replacing_worker(victim_pid) as kill: + kill.freeze() + await asyncio.to_thread(kill) + release.set() + served: Final = await burst + assert held_by[survivor_pid] >= 10, held_by + assert len(served) == held_by[survivor_pid], (held_by, len(served)) + for item in served: + _assert_answered_with_its_own_marker(item) follow_up: Final = _Call(endpoint="messages", stream=False, marker=uuid.uuid4().hex) (answered,) = await _burst(str(candidate.client.base_url), candidate.key, _CONFIG_MODEL, (follow_up,)) _assert_answered_with_its_own_marker(answered) diff --git a/tests/integration/providers/test_gemini_tts_voice_mapping_owned_proxy.py b/tests/integration/providers/test_gemini_tts_voice_mapping_owned_proxy.py index 8a977427741..3a5a9bf5d06 100644 --- a/tests/integration/providers/test_gemini_tts_voice_mapping_owned_proxy.py +++ b/tests/integration/providers/test_gemini_tts_voice_mapping_owned_proxy.py @@ -8,9 +8,7 @@ so these rows boot their own proxy with ``GEMINI_API_BASE`` pointed at the scrip from __future__ import annotations import asyncio -import os import re -import signal import threading from collections.abc import Callable, Iterator from dataclasses import dataclass @@ -306,11 +304,13 @@ async def test_worker_sigkill_mid_speech_burst_leaves_the_sibling_serving(gatewa ) held_by: Final = MappingProxyType({pid: _held_connections(pid, wire.url) for pid in workers}) victim: Final = max(workers, key=held_by.__getitem__) - os.kill(victim, signal.SIGKILL) + with proxy.replacing_worker(victim) as kill: + await asyncio.to_thread(kill) + release.set() + served: Final = await burst + during: Final = wire.drain() finally: release.set() - served: Final = await burst - during: Final = wire.drain() after: Final = await _fire(url, proxy.gateway.key, model) after_received: Final = wire.drain() assert sum(held_by.values()) == BURST, held_by @@ -335,7 +335,7 @@ async def test_proxy_restart_mid_speech_burst_serves_the_next_burst_after_the_re burst: Final = tasks.create_task(_fire(url, first.gateway.key, model, tolerate_transport_errors=True)) try: await asyncio.to_thread(eventually, lambda: wire.received.qsize(), lambda size: size >= BURST, 90) - first.process.terminate() + first.begin_shutdown() finally: release.set() served: Final = await burst diff --git a/tests/integration/providers/test_hosted_vllm_reasoning_content_chaos.py b/tests/integration/providers/test_hosted_vllm_reasoning_content_chaos.py index 84a5d71e8b8..eabbba425df 100644 --- a/tests/integration/providers/test_hosted_vllm_reasoning_content_chaos.py +++ b/tests/integration/providers/test_hosted_vllm_reasoning_content_chaos.py @@ -1,7 +1,6 @@ import asyncio import json import re -import signal import threading import uuid from collections.abc import Callable @@ -374,15 +373,15 @@ async def test_worker_sigkill_mid_burst_leaves_the_sibling_forwarding_reasoning( held_by: Final = MappingProxyType({pid: _open_upstream_connections(pid, wire.url) for pid in workers}) assert sum(held_by.values()) == 20, held_by victim_pid, survivor_pid = sorted(workers, key=held_by.__getitem__) - victim: Final = psutil.Process(victim_pid) - victim.suspend() - victim.send_signal(signal.SIGKILL) - release.set() - served: Final = await burst - assert held_by[survivor_pid] >= 10, held_by - assert len(served) == held_by[survivor_pid], (held_by, len(served)) - for item in served: - _assert_answered_with_its_own_marker(item) + with owned.replacing_worker(victim_pid) as kill: + kill.freeze() + await asyncio.to_thread(kill) + release.set() + served: Final = await burst + assert held_by[survivor_pid] >= 10, held_by + assert len(served) == held_by[survivor_pid], (held_by, len(served)) + for item in served: + _assert_answered_with_its_own_marker(item) follow_up: Final = _Call(endpoint="chat", stream=False, marker=uuid.uuid4().hex) (answered,) = await _burst(str(candidate.client.base_url), candidate.key, _CONFIG_MODEL, (follow_up,)) _assert_answered_with_its_own_marker(answered) diff --git a/tests/integration/providers/test_ollama_prompt_tools_chaos.py b/tests/integration/providers/test_ollama_prompt_tools_chaos.py index 243c85b2a64..37529f9f692 100644 --- a/tests/integration/providers/test_ollama_prompt_tools_chaos.py +++ b/tests/integration/providers/test_ollama_prompt_tools_chaos.py @@ -1,7 +1,6 @@ import asyncio import json import re -import signal import threading import uuid from collections.abc import Callable @@ -389,15 +388,15 @@ async def test_worker_sigkill_mid_burst_leaves_the_sibling_instructing_ollama(ga held_by: Final = MappingProxyType({pid: _open_upstream_connections(pid, wire.url) for pid in workers}) assert sum(held_by.values()) == 20, held_by victim_pid, survivor_pid = sorted(workers, key=held_by.__getitem__) - victim: Final = psutil.Process(victim_pid) - victim.suspend() - victim.send_signal(signal.SIGKILL) - release.set() - served: Final = await burst - assert held_by[survivor_pid] >= 10, held_by - assert len(served) == held_by[survivor_pid], (held_by, len(served)) - for item in served: - _assert_answered_with_its_own_marker(item) + with owned.replacing_worker(victim_pid) as kill: + kill.freeze() + await asyncio.to_thread(kill) + release.set() + served: Final = await burst + assert held_by[survivor_pid] >= 10, held_by + assert len(served) == held_by[survivor_pid], (held_by, len(served)) + for item in served: + _assert_answered_with_its_own_marker(item) follow_up: Final = _Call(endpoint="chat", stream=False, marker=uuid.uuid4().hex) (answered,) = await _burst(str(candidate.client.base_url), candidate.key, _CONFIG_MODEL, (follow_up,)) _assert_answered_with_its_own_marker(answered) diff --git a/tests/integration/providers/test_openai_dialect_prompt_cache_breakpoint_owned_proxy.py b/tests/integration/providers/test_openai_dialect_prompt_cache_breakpoint_owned_proxy.py index cc7ab06da90..920002316bc 100644 --- a/tests/integration/providers/test_openai_dialect_prompt_cache_breakpoint_owned_proxy.py +++ b/tests/integration/providers/test_openai_dialect_prompt_cache_breakpoint_owned_proxy.py @@ -1,6 +1,5 @@ import asyncio import re -import signal import threading from collections.abc import Callable, Iterator from concurrent.futures import ThreadPoolExecutor @@ -304,12 +303,12 @@ async def test_f2_a_worker_killed_mid_burst_leaves_the_survivor_marking_requests held_by: Final = MappingProxyType({pid: _open_peer_connections(pid, wire.url) for pid in workers}) assert sum(held_by.values()) == _HELD_BURST, held_by victim_pid, survivor_pid = sorted(workers, key=held_by.__getitem__) - victim: Final = psutil.Process(victim_pid) - victim.suspend() - victim.send_signal(signal.SIGKILL) - release.set() - served: Final = await burst - assert len(served) == held_by[survivor_pid], (held_by, len(served)) + with owned.replacing_worker(victim_pid) as kill: + kill.freeze() + await asyncio.to_thread(kill) + release.set() + served: Final = await burst + assert len(served) == held_by[survivor_pid], (held_by, len(served)) follow_up: Final = fresh_marker() (answered,) = await _held_burst(url, candidate.key, (follow_up,)) _assert_served_and_landed(wire, (*served, answered), expected_received=_HELD_BURST + 1) @@ -365,7 +364,7 @@ async def test_f3_a_graceful_restart_mid_burst_drains_the_held_requests_and_keep release.clear() burst: Final = asyncio.create_task(_held_burst(url, candidate.key, markers)) await asyncio.to_thread(eventually, held.qsize, lambda size: size == _HELD_BURST, 60) - owned.process.terminate() + owned.begin_shutdown() release.set() served: Final = await burst assert len(served) == _HELD_BURST, len(served) diff --git a/tests/integration/providers/test_openai_realtime_transcription_url.py b/tests/integration/providers/test_openai_realtime_transcription_url.py index eec0dd37baa..81c450d8f70 100644 --- a/tests/integration/providers/test_openai_realtime_transcription_url.py +++ b/tests/integration/providers/test_openai_realtime_transcription_url.py @@ -30,7 +30,7 @@ from websockets.exceptions import ConnectionClosed, InvalidStatus from tests.integration._support.client import JSON_OBJECT, Gateway, Scenario, eventually, object_value, string_value from tests.integration._support.database import read_rows -from tests.integration._support.process import owned_proxy_process, owned_upstream, stop_root_process +from tests.integration._support.process import OwnedProxy, owned_proxy_process, owned_upstream, stop_root_process from tests.integration._support.upstream import ScenarioHandle, delete_scenario, register_scenario from tests.integration.cost_calculation.cost_tracking_case import RealtimeResponse @@ -532,7 +532,7 @@ async def _first_frames(sockets: tuple[ClientConnection, ...]) -> tuple[dict[str return tuple([JSON_OBJECT.validate_json(await socket.recv()) for socket in sockets]) -async def _sessions_through_worker_kill(ws_base: str, model: str, key: str, root: psutil.Process) -> KillOutcome: +async def _sessions_through_worker_kill(ws_base: str, model: str, key: str, owned: OwnedProxy) -> KillOutcome: query: Final = f"model={model}&{TRANSCRIPTION_QUERY}" headers: Final = {"Authorization": f"Bearer {key}"} sockets: Final = tuple( @@ -544,13 +544,15 @@ async def _sessions_through_worker_kill(ws_base: str, model: str, key: str, root try: created: Final = await asyncio.wait_for(_first_frames(sockets), 60) assert [event.get("type") for event in created] == ["session.created"] * OPEN_SESSIONS, created - workers: Final = _workers(root) - assert len(workers) == WORKERS, [process.pid for process in workers] + workers: Final = tuple(sorted(owned.worker_pids())) + assert len(workers) == WORKERS, workers victim: Final = workers[0] - victim.kill() - await asyncio.to_thread(victim.wait, 10) - served: Final = await asyncio.wait_for(asyncio.gather(*(_one_turn_or_close(socket) for socket in sockets)), 60) - return KillOutcome(served.count(True), served.count(False), victim.pid) + with owned.replacing_worker(victim) as kill: + await asyncio.to_thread(kill) + served: Final = await asyncio.wait_for( + asyncio.gather(*(_one_turn_or_close(socket) for socket in sockets)), 60 + ) + return KillOutcome(served.count(True), served.count(False), victim) finally: await asyncio.gather(*(socket.close() for socket in sockets)) @@ -607,7 +609,7 @@ def test_worker_kill_then_proxy_restart_keep_transcription_sessions_serving( handle.scenario_id, ) root: Final = psutil.Process(owned.process.pid) - outcome: Final = asyncio.run(_sessions_through_worker_kill(_ws_base(owned_url), model, key, root)) + outcome: Final = asyncio.run(_sessions_through_worker_kill(_ws_base(owned_url), model, key, owned)) record_property( "worker_kill", {"served": outcome.served, "closed": outcome.closed, "killed": outcome.killed_pid} ) diff --git a/tests/integration/providers/test_realtime_transcription_guardrail_session_update.py b/tests/integration/providers/test_realtime_transcription_guardrail_session_update.py index e39f7e360f6..cda00b5f3fa 100644 --- a/tests/integration/providers/test_realtime_transcription_guardrail_session_update.py +++ b/tests/integration/providers/test_realtime_transcription_guardrail_session_update.py @@ -2067,18 +2067,18 @@ async def _open_sessions(ws_base: str, model: str, key: str) -> tuple[ClientConn return sockets -async def _sessions_through_worker_kill(ws_base: str, model: str, key: str, root: psutil.Process) -> KillOutcome: +async def _sessions_through_worker_kill(ws_base: str, model: str, key: str, owned: OwnedProxy) -> KillOutcome: sockets: Final = await _open_sessions(ws_base, model, key) try: - workers: Final = _workers(root) - assert len(workers) == WORKERS, [process.pid for process in workers] + workers: Final = tuple(sorted(owned.worker_pids())) + assert len(workers) == WORKERS, workers victim: Final = workers[0] - victim.kill() - await asyncio.to_thread(victim.wait, 10) - served: Final = await asyncio.wait_for( - asyncio.gather(*(_one_transcript_or_close(socket) for socket in sockets)), 60 - ) - return KillOutcome(served.count(True), served.count(False), victim.pid) + with owned.replacing_worker(victim) as kill: + await asyncio.to_thread(kill) + served: Final = await asyncio.wait_for( + asyncio.gather(*(_one_transcript_or_close(socket) for socket in sockets)), 60 + ) + return KillOutcome(served.count(True), served.count(False), victim) finally: await asyncio.gather(*(socket.close() for socket in sockets)) @@ -2138,7 +2138,7 @@ def test_worker_kill_then_proxy_restart_keep_transcription_sessions_uninjected( blocked_after_restart.scenario_id, ) root: Final = psutil.Process(owned.process.pid) - outcome: Final = asyncio.run(_sessions_through_worker_kill(_ws_base(owned_url), model, key, root)) + outcome: Final = asyncio.run(_sessions_through_worker_kill(_ws_base(owned_url), model, key, owned)) record_property( "worker_kill", {"served": outcome.served, "closed": outcome.closed, "killed": outcome.killed_pid} ) diff --git a/tests/integration/providers/test_realtime_unguarded_voice_turns.py b/tests/integration/providers/test_realtime_unguarded_voice_turns.py index 396b266f45f..4dd7fd9aeef 100644 --- a/tests/integration/providers/test_realtime_unguarded_voice_turns.py +++ b/tests/integration/providers/test_realtime_unguarded_voice_turns.py @@ -853,13 +853,6 @@ def _spawned_worker(child: psutil.Process) -> bool: return False -def _exited(process: psutil.Process) -> bool: - try: - return process.status() == psutil.STATUS_ZOMBIE - except psutil.NoSuchProcess: - return True - - def _workers(root: psutil.Process) -> tuple[psutil.Process, ...]: return tuple( sorted((child for child in root.children() if _spawned_worker(child)), key=lambda process: process.pid) @@ -913,16 +906,16 @@ async def _sessions_on_every_worker( async def _sessions_through_worker_kill( - url: str, key: str, workers: tuple[psutil.Process, ...] + url: str, key: str, owned: OwnedProxy, workers: tuple[psutil.Process, ...] ) -> tuple[tuple[Session, bool], ...]: opened: Final = await _sessions_on_every_worker(url, key, workers, ()) sockets: Final = tuple(session.socket for session in opened) try: on_killed_worker: Final = tuple(session.pid == workers[0].pid for session in opened) - workers[0].kill() - await asyncio.to_thread(eventually, lambda: _exited(workers[0]), bool, graceful_stop_seconds()) - turns: Final = (_turn_or_close(socket, label) for socket, label in zip(sockets, _labels(len(sockets)))) - sessions: Final = await asyncio.wait_for(asyncio.gather(*turns), 90) + with owned.replacing_worker(workers[0].pid) as kill: + await asyncio.to_thread(kill) + turns: Final = (_turn_or_close(socket, label) for socket, label in zip(sockets, _labels(len(sockets)))) + sessions: Final = await asyncio.wait_for(asyncio.gather(*turns), 90) return tuple(zip(sessions, on_killed_worker)) finally: await asyncio.gather(*(socket.close() for socket in sockets)) @@ -941,7 +934,7 @@ def test_worker_kill_leaves_the_surviving_sessions_with_only_their_own_frames(ga workers: Final = _workers(root) assert len(workers) == WORKERS, [process.pid for process in workers] outcomes: Final = asyncio.run( - _sessions_through_worker_kill(_session_url(_owned_url(owned), f"model={model}"), key, workers) + _sessions_through_worker_kill(_session_url(_owned_url(owned), f"model={model}"), key, owned, workers) ) served: Final = tuple(label for label, (_, killed) in zip(_labels(len(outcomes)), outcomes) if not killed) assert [session.types for session, _ in outcomes] == [ diff --git a/tests/integration/providers/test_responses_bridge_incomplete.py b/tests/integration/providers/test_responses_bridge_incomplete.py index 5d3877c954e..e970c806330 100644 --- a/tests/integration/providers/test_responses_bridge_incomplete.py +++ b/tests/integration/providers/test_responses_bridge_incomplete.py @@ -2,9 +2,11 @@ import json import uuid from typing import Final +import httpx import pytest -from integration._support.client import Gateway +from integration._support.client import Gateway, list_value, object_value from integration._support.wire import Reply, Request, wire_server +from pydantic import JsonValue @pytest.mark.covers("other.provider_wire.responses_bridge.max_output_tokens_incomplete_maps_to_length") @@ -572,3 +574,726 @@ def test_chat_over_responses_deployment_merges_function_call_followed_by_message } ], response.text assert [(request.method, request.target) for request in wire.drain()] == [("POST", "/responses")] + + +_WEATHER_TOOL: Final[dict[str, JsonValue]] = { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get the weather for a city.", + "parameters": { + "type": "object", + "properties": {"city": {"type": "string"}}, + "required": ["city"], + }, + }, +} + + +def _message_item(item_id: str, text: str, phase: str | None = None) -> dict[str, JsonValue]: + return { + "type": "message", + "id": item_id, + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": text, "annotations": []}], + **({} if phase is None else {"phase": phase}), + } + + +_WEATHER_CALL: Final[dict[str, JsonValue]] = { + "type": "function_call", + "id": "fc_1", + "call_id": "call_1", + "name": "get_weather", + "arguments": '{"city":"Paris"}', + "status": "completed", +} + +_TWO_MESSAGES_THEN_CALL: Final = ( + _message_item("msg_first", "First message."), + _message_item("msg_second", "Second message."), + _WEATHER_CALL, +) + + +def test_chat_over_responses_deployment_folds_two_messages_and_function_call_into_one_choice( + gateway: Gateway, +) -> None: + identity: Final = "responses-bridge-two-messages-" + uuid.uuid4().hex + + def respond(request: Request) -> Reply: + if request.method == "GET" and request.target == "/v1/models": + return Reply(body=b'{"object":"list","data":[]}') + assert request.method == "POST" and request.target == "/responses", request.target + return Reply( + body=json.dumps( + { + "id": "resp_two_messages", + "object": "response", + "created_at": 1789788253, + "status": "completed", + "model": "gpt-6-sol", + "output": list(_TWO_MESSAGES_THEN_CALL), + "usage": {"input_tokens": 10, "output_tokens": 5, "total_tokens": 15}, + } + ).encode() + ) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model="openai/responses/gpt-6-sol", api_base=wire.url, api_key="synthetic-openai-key" + ) + response: Final = gateway.request( + "POST", + "/v1/chat/completions", + { + "model": model, + "messages": [{"role": "user", "content": f"What is the weather in Paris? {identity}"}], + "tools": [_WEATHER_TOOL], + "cache": {"no-cache": True}, + }, + ) + assert response.status_code == 200, response.text + body: Final = response.json() + assert body["choices"] == [ + { + "finish_reason": "tool_calls", + "index": 0, + "message": { + "role": "assistant", + "content": "First message.Second message.", + "tool_calls": [ + { + "id": "fc_1", + "type": "function", + "function": {"name": "get_weather", "arguments": '{"city":"Paris"}'}, + "index": 0, + } + ], + }, + } + ], response.text + assert [(request.method, request.target) for request in wire.drain()] == [("POST", "/responses")] + + +_COMMENTARY_THEN_FINAL_ANSWER: Final = ( + _message_item("msg_commentary", "Let me answer that.", phase="commentary"), + _message_item("msg_final", '{"city": "Paris", "summary": "Sunny, 21C"}', phase="final_answer"), +) + + +def test_chat_over_responses_deployment_folds_commentary_and_final_answer_into_one_choice( + gateway: Gateway, +) -> None: + identity: Final = "responses-bridge-commentary-final-" + uuid.uuid4().hex + + def respond(request: Request) -> Reply: + if request.method == "GET" and request.target == "/v1/models": + return Reply(body=b'{"object":"list","data":[]}') + assert request.method == "POST" and request.target == "/responses", request.target + return Reply( + body=json.dumps( + { + "id": "resp_commentary_final", + "object": "response", + "created_at": 1789788253, + "status": "completed", + "model": "gpt-6-sol", + "output": list(_COMMENTARY_THEN_FINAL_ANSWER), + "usage": {"input_tokens": 10, "output_tokens": 5, "total_tokens": 15}, + } + ).encode() + ) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model="openai/responses/gpt-6-sol", api_base=wire.url, api_key="synthetic-openai-key" + ) + response: Final = gateway.request( + "POST", + "/v1/chat/completions", + { + "model": model, + "messages": [{"role": "user", "content": f"Weather in Paris as JSON. {identity}"}], + "tools": [_WEATHER_TOOL], + "cache": {"no-cache": True}, + }, + ) + assert response.status_code == 200, response.text + body: Final = response.json() + assert body["choices"] == [ + { + "finish_reason": "stop", + "index": 0, + "message": { + "role": "assistant", + "content": 'Let me answer that.{"city": "Paris", "summary": "Sunny, 21C"}', + }, + } + ], response.text + assert [(request.method, request.target) for request in wire.drain()] == [("POST", "/responses")] + + +_REASONING_ITEM: Final[dict[str, JsonValue]] = { + "type": "reasoning", + "id": "rs_fold", + "summary": [{"type": "summary_text", "text": "Checking the forecast."}], +} + +_TIME_TOOL: Final[dict[str, JsonValue]] = { + "type": "function", + "function": { + "name": "get_time", + "description": "Get the local time for a city.", + "parameters": { + "type": "object", + "properties": {"city": {"type": "string"}}, + "required": ["city"], + }, + }, +} + +_TIME_CALL: Final[dict[str, JsonValue]] = { + "type": "function_call", + "id": "fc_2", + "call_id": "call_2", + "name": "get_time", + "arguments": '{"city":"Paris","tz":"CET"}', + "status": "completed", +} + +_REASONING_THEN_TWO_MESSAGES_THEN_CALL: Final = ( + _REASONING_ITEM, + _message_item("msg_first", "First message."), + _message_item("msg_second", "Second message."), + _WEATHER_CALL, +) + +_MESSAGE_THEN_TWO_CALLS: Final = ( + _message_item("msg_only", "Checking both."), + _WEATHER_CALL, + _TIME_CALL, +) + +_TWO_CALLS_CHAT_SHAPE: Final = [ + { + "id": "fc_1", + "type": "function", + "function": {"name": "get_weather", "arguments": '{"city":"Paris"}'}, + "index": 0, + }, + { + "id": "fc_2", + "type": "function", + "function": {"name": "get_time", "arguments": '{"city":"Paris","tz":"CET"}'}, + "index": 1, + }, +] + + +def _responses_body(response_id: str, output: tuple[dict[str, JsonValue], ...]) -> bytes: + return json.dumps( + { + "id": response_id, + "object": "response", + "created_at": 1789788253, + "status": "completed", + "model": "gpt-6-sol", + "output": list(output), + "usage": {"input_tokens": 10, "output_tokens": 5, "total_tokens": 15}, + } + ).encode() + + +def test_chat_over_responses_deployment_folds_reasoning_two_messages_and_function_call_into_one_choice( + gateway: Gateway, +) -> None: + identity: Final = "responses-bridge-reasoning-two-messages-" + uuid.uuid4().hex + + def respond(request: Request) -> Reply: + if request.method == "GET" and request.target == "/v1/models": + return Reply(body=b'{"object":"list","data":[]}') + assert request.method == "POST" and request.target == "/responses", request.target + return Reply(body=_responses_body("resp_reasoning_two_messages", _REASONING_THEN_TWO_MESSAGES_THEN_CALL)) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model="openai/responses/gpt-6-sol", api_base=wire.url, api_key="synthetic-openai-key" + ) + response: Final = gateway.request( + "POST", + "/v1/chat/completions", + { + "model": model, + "messages": [{"role": "user", "content": f"What is the weather in Paris? {identity}"}], + "tools": [_WEATHER_TOOL], + "cache": {"no-cache": True}, + }, + ) + assert response.status_code == 200, response.text + body: Final = response.json() + assert body["choices"] == [ + { + "finish_reason": "tool_calls", + "index": 0, + "message": { + "role": "assistant", + "content": "First message.Second message.", + "reasoning_content": "Checking the forecast.", + "reasoning_items": [_REASONING_ITEM], + "tool_calls": _TWO_CALLS_CHAT_SHAPE[:1], + }, + } + ], response.text + assert [(request.method, request.target) for request in wire.drain()] == [("POST", "/responses")] + + +def test_chat_over_responses_deployment_folds_message_and_two_function_calls_into_one_choice( + gateway: Gateway, +) -> None: + identity: Final = "responses-bridge-two-calls-" + uuid.uuid4().hex + + def respond(request: Request) -> Reply: + if request.method == "GET" and request.target == "/v1/models": + return Reply(body=b'{"object":"list","data":[]}') + assert request.method == "POST" and request.target == "/responses", request.target + return Reply(body=_responses_body("resp_two_calls", _MESSAGE_THEN_TWO_CALLS)) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model="openai/responses/gpt-6-sol", api_base=wire.url, api_key="synthetic-openai-key" + ) + response: Final = gateway.request( + "POST", + "/v1/chat/completions", + { + "model": model, + "messages": [{"role": "user", "content": f"Weather and time in Paris? {identity}"}], + "tools": [_WEATHER_TOOL, _TIME_TOOL], + "cache": {"no-cache": True}, + }, + ) + assert response.status_code == 200, response.text + body: Final = response.json() + assert body["choices"] == [ + { + "finish_reason": "tool_calls", + "index": 0, + "message": { + "role": "assistant", + "content": "Checking both.", + "tool_calls": _TWO_CALLS_CHAT_SHAPE, + }, + } + ], response.text + assert [(request.method, request.target) for request in wire.drain()] == [("POST", "/responses")] + + +def _responses_frame(event: str, payload: dict[str, JsonValue]) -> bytes: + return f"event: {event}\ndata: {json.dumps(payload, separators=(',', ':'))}\n\n".encode() + + +def _message_item_frames(output_index: int, item: dict[str, JsonValue]) -> tuple[bytes, ...]: + item_id: Final = item["id"] + text: Final = object_value(list_value(item["content"])[0])["text"] + return ( + _responses_frame( + "response.output_item.added", + { + "type": "response.output_item.added", + "output_index": output_index, + "item": {**item, "status": "in_progress", "content": []}, + }, + ), + _responses_frame( + "response.output_text.delta", + { + "type": "response.output_text.delta", + "item_id": item_id, + "output_index": output_index, + "content_index": 0, + "delta": text, + }, + ), + _responses_frame( + "response.output_item.done", + {"type": "response.output_item.done", "output_index": output_index, "item": item}, + ), + ) + + +def _function_call_frames(output_index: int, item: dict[str, JsonValue]) -> tuple[bytes, ...]: + return ( + _responses_frame( + "response.output_item.added", + { + "type": "response.output_item.added", + "output_index": output_index, + "item": {**item, "status": "in_progress", "arguments": ""}, + }, + ), + _responses_frame( + "response.function_call_arguments.delta", + { + "type": "response.function_call_arguments.delta", + "item_id": item["id"], + "output_index": output_index, + "delta": item["arguments"], + }, + ), + _responses_frame( + "response.function_call_arguments.done", + { + "type": "response.function_call_arguments.done", + "item_id": item["id"], + "output_index": output_index, + "arguments": item["arguments"], + }, + ), + _responses_frame( + "response.output_item.done", + {"type": "response.output_item.done", "output_index": output_index, "item": item}, + ), + ) + + +def _reasoning_item_frames(output_index: int, item: dict[str, JsonValue]) -> tuple[bytes, ...]: + text: Final = object_value(list_value(item["summary"])[0])["text"] + return ( + _responses_frame( + "response.output_item.added", + { + "type": "response.output_item.added", + "output_index": output_index, + "item": {**item, "summary": []}, + }, + ), + _responses_frame( + "response.reasoning_summary_text.delta", + { + "type": "response.reasoning_summary_text.delta", + "item_id": item["id"], + "output_index": output_index, + "summary_index": 0, + "delta": text, + }, + ), + _responses_frame( + "response.output_item.done", + {"type": "response.output_item.done", "output_index": output_index, "item": item}, + ), + ) + + +def _sse_payloads(lines: list[str]) -> list[dict[str, JsonValue]]: + return [json.loads(line[6:]) for line in lines if line.startswith("data: ") and line != "data: [DONE]"] + + +def test_chat_stream_over_responses_deployment_folds_two_messages_and_function_call_into_one_choice( + gateway: Gateway, +) -> None: + identity: Final = "responses-bridge-two-messages-stream-" + uuid.uuid4().hex + response_payload: Final[dict[str, JsonValue]] = { + "id": "resp_two_messages_stream", + "object": "response", + "created_at": 1789788253, + "model": "gpt-6-sol", + "status": "in_progress", + "output": [], + } + + def respond(request: Request) -> Reply: + if request.method == "GET" and request.target == "/v1/models": + return Reply(body=b'{"object":"list","data":[]}') + assert request.method == "POST" and request.target == "/responses", request.target + assert json.loads(request.body)["stream"] is True, request.body + return Reply( + content_type="text/event-stream", + chunks=( + _responses_frame("response.created", {"type": "response.created", "response": response_payload}), + *_message_item_frames(0, _TWO_MESSAGES_THEN_CALL[0]), + *_message_item_frames(1, _TWO_MESSAGES_THEN_CALL[1]), + *_function_call_frames(2, _TWO_MESSAGES_THEN_CALL[2]), + _responses_frame( + "response.completed", + { + "type": "response.completed", + "response": { + **response_payload, + "status": "completed", + "output": list(_TWO_MESSAGES_THEN_CALL), + "usage": {"input_tokens": 10, "output_tokens": 5, "total_tokens": 15}, + }, + }, + ), + ), + ) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model="openai/responses/gpt-6-sol", api_base=wire.url, api_key="synthetic-openai-key" + ) + response: Final = gateway.request( + "POST", + "/v1/chat/completions", + { + "model": model, + "messages": [{"role": "user", "content": f"What is the weather in Paris? {identity}"}], + "tools": [_WEATHER_TOOL], + "stream": True, + "cache": {"no-cache": True}, + }, + ) + assert response.status_code == 200, response.text + chunks: Final = _sse_payloads(list(response.iter_lines())) + choices: Final = [object_value(choice) for chunk in chunks for choice in list_value(chunk["choices"])] + assert [choice["index"] for choice in choices] == [0] * len(choices), chunks + deltas: Final = [object_value(choice["delta"]) for choice in choices] + assert ( + "".join(str(delta["content"]) for delta in deltas if delta.get("content")) + == "First message.Second message." + ) + tool_calls: Final = [ + object_value(call) + for delta in deltas + if delta.get("tool_calls") + for call in list_value(delta["tool_calls"]) + ] + assert [ + object_value(call["function"])["name"] for call in tool_calls if object_value(call["function"]).get("name") + ] == ["get_weather"], chunks + assert ( + "".join(str(object_value(call["function"]).get("arguments") or "") for call in tool_calls) + == '{"city":"Paris"}' + ), chunks + assert [choice["finish_reason"] for choice in choices if choice.get("finish_reason")] == ["tool_calls"], chunks + assert [(request.method, request.target) for request in wire.drain()] == [("POST", "/responses")] + + +def test_chat_stream_over_responses_deployment_folds_commentary_and_final_answer_into_one_choice( + gateway: Gateway, +) -> None: + identity: Final = "responses-bridge-commentary-final-stream-" + uuid.uuid4().hex + response_payload: Final[dict[str, JsonValue]] = { + "id": "resp_commentary_final_stream", + "object": "response", + "created_at": 1789788253, + "model": "gpt-6-sol", + "status": "in_progress", + "output": [], + } + + def respond(request: Request) -> Reply: + if request.method == "GET" and request.target == "/v1/models": + return Reply(body=b'{"object":"list","data":[]}') + assert request.method == "POST" and request.target == "/responses", request.target + assert json.loads(request.body)["stream"] is True, request.body + return Reply( + content_type="text/event-stream", + chunks=( + _responses_frame("response.created", {"type": "response.created", "response": response_payload}), + *_message_item_frames(0, _COMMENTARY_THEN_FINAL_ANSWER[0]), + *_message_item_frames(1, _COMMENTARY_THEN_FINAL_ANSWER[1]), + _responses_frame( + "response.completed", + { + "type": "response.completed", + "response": { + **response_payload, + "status": "completed", + "output": list(_COMMENTARY_THEN_FINAL_ANSWER), + "usage": {"input_tokens": 10, "output_tokens": 5, "total_tokens": 15}, + }, + }, + ), + ), + ) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model="openai/responses/gpt-6-sol", api_base=wire.url, api_key="synthetic-openai-key" + ) + response: Final = gateway.request( + "POST", + "/v1/chat/completions", + { + "model": model, + "messages": [{"role": "user", "content": f"Weather in Paris as JSON. {identity}"}], + "tools": [_WEATHER_TOOL], + "stream": True, + "cache": {"no-cache": True}, + }, + ) + assert response.status_code == 200, response.text + chunks: Final = _sse_payloads(list(response.iter_lines())) + choices: Final = [object_value(choice) for chunk in chunks for choice in list_value(chunk["choices"])] + assert [choice["index"] for choice in choices] == [0] * len(choices), chunks + deltas: Final = [object_value(choice["delta"]) for choice in choices] + assert ( + "".join(str(delta["content"]) for delta in deltas if delta.get("content")) + == 'Let me answer that.{"city": "Paris", "summary": "Sunny, 21C"}' + ), chunks + assert not [delta for delta in deltas if delta.get("tool_calls")], chunks + assert [choice["finish_reason"] for choice in choices if choice.get("finish_reason")] == ["stop"], chunks + assert [(request.method, request.target) for request in wire.drain()] == [("POST", "/responses")] + + +def _stream_reply(response_id: str, output: tuple[dict[str, JsonValue], ...], frames: tuple[bytes, ...]) -> Reply: + response_payload: Final[dict[str, JsonValue]] = { + "id": response_id, + "object": "response", + "created_at": 1789788253, + "model": "gpt-6-sol", + "status": "in_progress", + "output": [], + } + return Reply( + content_type="text/event-stream", + chunks=( + _responses_frame("response.created", {"type": "response.created", "response": response_payload}), + *frames, + _responses_frame( + "response.completed", + { + "type": "response.completed", + "response": { + **response_payload, + "status": "completed", + "output": list(output), + "usage": {"input_tokens": 10, "output_tokens": 5, "total_tokens": 15}, + }, + }, + ), + ), + ) + + +def _stream_choices(response: httpx.Response) -> list[dict[str, JsonValue]]: + chunks: Final = _sse_payloads(list(response.iter_lines())) + return [object_value(choice) for chunk in chunks for choice in list_value(chunk["choices"])] + + +def _stream_tool_calls(deltas: list[dict[str, JsonValue]]) -> list[dict[str, JsonValue]]: + return [ + object_value(call) for delta in deltas if delta.get("tool_calls") for call in list_value(delta["tool_calls"]) + ] + + +def test_chat_stream_over_responses_deployment_folds_reasoning_two_messages_and_function_call_into_one_choice( + gateway: Gateway, +) -> None: + identity: Final = "responses-bridge-reasoning-two-messages-stream-" + uuid.uuid4().hex + + def respond(request: Request) -> Reply: + if request.method == "GET" and request.target == "/v1/models": + return Reply(body=b'{"object":"list","data":[]}') + assert request.method == "POST" and request.target == "/responses", request.target + assert json.loads(request.body)["stream"] is True, request.body + return _stream_reply( + "resp_reasoning_two_messages_stream", + _REASONING_THEN_TWO_MESSAGES_THEN_CALL, + ( + *_reasoning_item_frames(0, _REASONING_THEN_TWO_MESSAGES_THEN_CALL[0]), + *_message_item_frames(1, _REASONING_THEN_TWO_MESSAGES_THEN_CALL[1]), + *_message_item_frames(2, _REASONING_THEN_TWO_MESSAGES_THEN_CALL[2]), + *_function_call_frames(3, _REASONING_THEN_TWO_MESSAGES_THEN_CALL[3]), + ), + ) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model="openai/responses/gpt-6-sol", api_base=wire.url, api_key="synthetic-openai-key" + ) + response: Final = gateway.request( + "POST", + "/v1/chat/completions", + { + "model": model, + "messages": [{"role": "user", "content": f"What is the weather in Paris? {identity}"}], + "tools": [_WEATHER_TOOL], + "stream": True, + "cache": {"no-cache": True}, + }, + ) + assert response.status_code == 200, response.text + choices: Final = _stream_choices(response) + assert [choice["index"] for choice in choices] == [0] * len(choices), choices + deltas: Final = [object_value(choice["delta"]) for choice in choices] + assert ( + "".join(str(delta["reasoning_content"]) for delta in deltas if delta.get("reasoning_content")) + == "Checking the forecast." + ), choices + assert ( + "".join(str(delta["content"]) for delta in deltas if delta.get("content")) + == "First message.Second message." + ), choices + tool_calls: Final = _stream_tool_calls(deltas) + assert [ + object_value(call["function"])["name"] for call in tool_calls if object_value(call["function"]).get("name") + ] == ["get_weather"], choices + assert ( + "".join(str(object_value(call["function"]).get("arguments") or "") for call in tool_calls) + == '{"city":"Paris"}' + ), choices + assert [choice["finish_reason"] for choice in choices if choice.get("finish_reason")] == ["tool_calls"], choices + assert [(request.method, request.target) for request in wire.drain()] == [("POST", "/responses")] + + +def test_chat_stream_over_responses_deployment_folds_message_and_two_function_calls_into_one_choice( + gateway: Gateway, +) -> None: + identity: Final = "responses-bridge-two-calls-stream-" + uuid.uuid4().hex + + def respond(request: Request) -> Reply: + if request.method == "GET" and request.target == "/v1/models": + return Reply(body=b'{"object":"list","data":[]}') + assert request.method == "POST" and request.target == "/responses", request.target + assert json.loads(request.body)["stream"] is True, request.body + return _stream_reply( + "resp_two_calls_stream", + _MESSAGE_THEN_TWO_CALLS, + ( + *_message_item_frames(0, _MESSAGE_THEN_TWO_CALLS[0]), + *_function_call_frames(1, _MESSAGE_THEN_TWO_CALLS[1]), + *_function_call_frames(2, _MESSAGE_THEN_TWO_CALLS[2]), + ), + ) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model="openai/responses/gpt-6-sol", api_base=wire.url, api_key="synthetic-openai-key" + ) + response: Final = gateway.request( + "POST", + "/v1/chat/completions", + { + "model": model, + "messages": [{"role": "user", "content": f"Weather and time in Paris? {identity}"}], + "tools": [_WEATHER_TOOL, _TIME_TOOL], + "stream": True, + "cache": {"no-cache": True}, + }, + ) + assert response.status_code == 200, response.text + choices: Final = _stream_choices(response) + assert [choice["index"] for choice in choices] == [0] * len(choices), choices + deltas: Final = [object_value(choice["delta"]) for choice in choices] + assert "".join(str(delta["content"]) for delta in deltas if delta.get("content")) == "Checking both.", choices + tool_calls: Final = _stream_tool_calls(deltas) + assert [ + (call["index"], object_value(call["function"])["name"]) + for call in tool_calls + if object_value(call["function"]).get("name") + ] == [(0, "get_weather"), (1, "get_time")], choices + arguments_by_index: Final = { + index: "".join( + str(object_value(call["function"]).get("arguments") or "") + for call in tool_calls + if call["index"] == index + ) + for index in (0, 1) + } + assert arguments_by_index == {0: '{"city":"Paris"}', 1: '{"city":"Paris","tz":"CET"}'}, choices + assert [choice["finish_reason"] for choice in choices if choice.get("finish_reason")] == ["tool_calls"], choices + assert [(request.method, request.target) for request in wire.drain()] == [("POST", "/responses")] diff --git a/tests/integration/providers/test_responses_bridge_prompt_cache_breakpoint_chaos.py b/tests/integration/providers/test_responses_bridge_prompt_cache_breakpoint_chaos.py index 10d5cfd9646..e03af85d01a 100644 --- a/tests/integration/providers/test_responses_bridge_prompt_cache_breakpoint_chaos.py +++ b/tests/integration/providers/test_responses_bridge_prompt_cache_breakpoint_chaos.py @@ -2,7 +2,6 @@ from __future__ import annotations import asyncio import re -import signal import threading import uuid from contextlib import ExitStack @@ -126,11 +125,11 @@ async def test_worker_and_peer_outages_preserve_markers_and_recover( survivor_pids: Final = tuple(pid for pid in workers if pid != victim_pid) assert active_by_worker[victim_pid] > 0 and len(survivor_pids) == 1, active_by_worker (survivor_pid,) = survivor_pids - victim: Final = psutil.Process(victim_pid) - victim.suspend() - victim.send_signal(signal.SIGKILL) - release.set() - remaining_served: Final = await asyncio.gather(*burst_tasks[len(early_calls) :]) + with owned.replacing_worker(victim_pid) as kill: + kill.freeze() + await asyncio.to_thread(kill) + release.set() + remaining_served: Final = await asyncio.gather(*burst_tasks[len(early_calls) :]) served: Final = (*early_served, *remaining_served) successful: Final = tuple(item for item in served if item.status == 200) connection_errors: Final = tuple(item for item in served if item.status == 0) diff --git a/tests/integration/providers/test_responses_bridge_prompt_cache_breakpoint_wire.py b/tests/integration/providers/test_responses_bridge_prompt_cache_breakpoint_wire.py index 9fad3a2b191..055eec3f6f6 100644 --- a/tests/integration/providers/test_responses_bridge_prompt_cache_breakpoint_wire.py +++ b/tests/integration/providers/test_responses_bridge_prompt_cache_breakpoint_wire.py @@ -367,6 +367,142 @@ def test_anthropic_sdk_system_string_gets_the_injected_marker(bridge: _Bridge) - bridge.spend.landed(bridge.injecting_off, raw.headers["x-litellm-call-id"], None) +_EPHEMERAL: Final[anthropic.types.CacheControlEphemeralParam] = {"type": "ephemeral"} +_OPTED_IN: Final[Mapping[str, JsonValue]] = {"prompt_cache_options": {"mode": "explicit"}} + + +def _anthropic_message( + client: anthropic.Anthropic, + model: str, + messages: Sequence[anthropic.types.MessageParam], + *, + system: Sequence[anthropic.types.TextBlockParam] = (), + stream: bool = False, +) -> anthropic.types.Message: + if not stream: + return client.messages.create( + model=model, max_tokens=64, system=system, messages=messages, extra_body=dict(pcb.NO_CACHE) + ) + with client.messages.stream( + model=model, max_tokens=64, system=system, messages=messages, extra_body=dict(pcb.NO_CACHE) + ) as events: + return events.get_final_message() + + +def _marked_system() -> list[anthropic.types.TextBlockParam]: + return [ + {"type": "text", "text": "Answer briefly"}, + {"type": "text", "text": "Stable rules", "cache_control": _EPHEMERAL}, + ] + + +def _marked_user(marker: str) -> anthropic.types.MessageParam: + return { + "role": "user", + "content": [ + {"type": "text", "text": pcb.prompt(marker), "cache_control": _EPHEMERAL}, + {"type": "text", "text": "and nothing else"}, + ], + } + + +@pytest.mark.parametrize("stream", [False, True], ids=["unary", "stream"]) +def test_anthropic_sdk_cache_control_reaches_the_responses_wire_as_a_prompt_cache_breakpoint_on_the_same_block( + bridge: _Bridge, stream: bool +) -> None: + marker: Final = uuid.uuid4().hex + with bridge.gateway.scenario() as scenario, _anthropic(bridge) as client: + model: Final = scenario.model(model=pcb.MODEL, api_base=bridge.api_base, **_OPTED_IN) + message: Final = _anthropic_message( + client, model, [_marked_user(marker)], system=_marked_system(), stream=stream + ) + _anthropic_text(message, marker) + body: Final = _wire_body(pcb.posted(bridge.wire, marker), stream=stream) + assert "instructions" not in body, body + assert body["input"] == [ + { + "type": "message", + "role": "developer", + "content": [ + {"type": "input_text", "text": "Answer briefly"}, + {"type": "input_text", "text": "Stable rules", "prompt_cache_breakpoint": pcb.EXPLICIT}, + ], + }, + { + "type": "message", + "role": "user", + "content": [ + {"type": "input_text", "text": pcb.prompt(marker), "prompt_cache_breakpoint": pcb.EXPLICIT}, + {"type": "input_text", "text": "and nothing else"}, + ], + }, + ], body + + +def test_anthropic_sdk_cache_control_sends_no_mark_to_a_deployment_that_does_not_read_breakpoints( + bridge: _Bridge, +) -> None: + marker: Final = uuid.uuid4().hex + with _anthropic(bridge) as client: + message: Final = _anthropic_message(client, bridge.off, [_marked_user(marker)], system=_marked_system()) + _anthropic_text(message, marker) + request: Final = pcb.posted(bridge.wire, marker) + body: Final = _wire_body(request) + assert body["instructions"] == "Answer briefly\nStable rules", body + assert body["input"] == [ + { + "type": "message", + "role": "user", + "content": [ + {"type": "input_text", "text": pcb.prompt(marker)}, + {"type": "input_text", "text": "and nothing else"}, + ], + } + ], body + assert b"cache_control" not in request.body and b"prompt_cache_breakpoint" not in request.body, request.body + + +def test_anthropic_sdk_cache_control_on_a_tool_result_marks_its_function_call_output(bridge: _Bridge) -> None: + marker: Final = uuid.uuid4().hex + messages: Final[list[anthropic.types.MessageParam]] = [ + {"role": "user", "content": pcb.prompt(marker)}, + {"role": "assistant", "content": [{"type": "tool_use", "id": "call_1", "name": "lookup", "input": {}}]}, + { + "role": "user", + "content": [ + {"type": "tool_result", "tool_use_id": "call_1", "content": "found it", "cache_control": _EPHEMERAL} + ], + }, + ] + with bridge.gateway.scenario() as scenario, _anthropic(bridge) as client: + model: Final = scenario.model(model=pcb.MODEL, api_base=bridge.api_base, **_OPTED_IN) + message: Final = _anthropic_message(client, model, messages) + _anthropic_text(message, marker) + request: Final = pcb.posted(bridge.wire, marker) + _wire_body(request) + assert pcb.function_output(pcb.input_items(request), "call_1") == [ + {"type": "input_text", "text": "found it", "prompt_cache_breakpoint": pcb.EXPLICIT} + ], "OpenAI caches up to a breakpoint on a function_call_output part (prompt caching guide, read 2026-10-12)" + + +def test_anthropic_sdk_cache_control_on_an_assistant_turn_is_dropped(bridge: _Bridge) -> None: + marker: Final = uuid.uuid4().hex + messages: Final[list[anthropic.types.MessageParam]] = [ + {"role": "user", "content": "first question"}, + {"role": "assistant", "content": [{"type": "text", "text": "earlier answer", "cache_control": _EPHEMERAL}]}, + {"role": "user", "content": pcb.prompt(marker)}, + ] + with bridge.gateway.scenario() as scenario, _anthropic(bridge) as client: + model: Final = scenario.model(model=pcb.MODEL, api_base=bridge.api_base, **_OPTED_IN) + message: Final = _anthropic_message(client, model, messages) + _anthropic_text(message, marker) + request: Final = pcb.posted(bridge.wire, marker) + _wire_body(request) + assert pcb.content_of(pcb.input_items(request), "assistant") == [ + {"type": "output_text", "text": "earlier answer"} + ], "OpenAI accepts a breakpoint on output_text but writes no cache for it (live gpt-5.6, 2026-10-12)" + + def test_native_responses_request_never_enters_the_bridge(bridge: _Bridge) -> None: marker: Final = uuid.uuid4().hex block: Final[dict[str, JsonValue]] = {"type": "input_text", "text": pcb.prompt(marker)} diff --git a/tests/integration/providers/test_responses_bridge_prompt_cache_breakpoint_wire_chaos.py b/tests/integration/providers/test_responses_bridge_prompt_cache_breakpoint_wire_chaos.py index 6f9a05d8dfa..ee73d780170 100644 --- a/tests/integration/providers/test_responses_bridge_prompt_cache_breakpoint_wire_chaos.py +++ b/tests/integration/providers/test_responses_bridge_prompt_cache_breakpoint_wire_chaos.py @@ -1,6 +1,5 @@ import asyncio import dataclasses -import signal import threading import uuid from collections import Counter @@ -13,7 +12,6 @@ from typing import Final, Literal, TypeAlias from urllib.parse import urlsplit import httpx -import psutil import pytest import yaml from integration._support import prompt_cache_breakpoint as pcb @@ -371,16 +369,16 @@ async def test_worker_sigkill_mid_burst_leaves_the_sibling_answering( held_by: Final = MappingProxyType({pid: pcb.open_upstream_connections(pid, wire.url) for pid in workers}) assert sum(held_by.values()) == 20, held_by victim_pid, survivor_pid = sorted(workers, key=held_by.__getitem__) - victim: Final = psutil.Process(victim_pid) - victim.suspend() - victim.send_signal(signal.SIGKILL) - held.release.set() - served: Final = await burst - assert held_by[survivor_pid] >= 10, held_by - assert len(served) == held_by[survivor_pid], (held_by, len(served)) - for item in served: - _assert_answered_in_its_own_shape(item) - _marked_once(pcb.drained_posts(wire), calls, pcb.EXPLICIT) + with owned.replacing_worker(victim_pid) as kill: + kill.freeze() + await asyncio.to_thread(kill) + held.release.set() + served: Final = await burst + assert held_by[survivor_pid] >= 10, held_by + assert len(served) == held_by[survivor_pid], (held_by, len(served)) + for item in served: + _assert_answered_in_its_own_shape(item) + _marked_once(pcb.drained_posts(wire), calls, pcb.EXPLICIT) follow_up: Final = uuid.uuid4().hex call_id: Final = _completion(_chat(candidate, _GLOBAL_UNSET, follow_up, "yes"), follow_up) pcb.assert_marker(_user_block_on_wire(wire, follow_up), None) @@ -397,7 +395,7 @@ async def test_proxy_restart_mid_burst_never_lands_a_served_call_twice( with owned_proxy_process(gateway, tmp_path, {}, config=config, workers=2) as first: _worker_pids(first) burst: Final = await _hold_burst(held, first.gateway, _GLOBAL_UNSET, calls) - first.process.terminate() + first.begin_shutdown() held.release.set() served: Final = await burst for item in served: diff --git a/tests/integration/providers/test_responses_minted_reasoning_replay_chaos.py b/tests/integration/providers/test_responses_minted_reasoning_replay_chaos.py index d62b67e08d6..f003e5ffa63 100644 --- a/tests/integration/providers/test_responses_minted_reasoning_replay_chaos.py +++ b/tests/integration/providers/test_responses_minted_reasoning_replay_chaos.py @@ -1,7 +1,6 @@ import asyncio import json import re -import signal import socket import threading import time @@ -351,15 +350,15 @@ async def test_worker_sigkill_mid_burst_leaves_the_sibling_dropping_the_minted_i held_by: Final = MappingProxyType({pid: _open_upstream_connections(pid, wire.url) for pid in workers}) assert sum(held_by.values()) == 20, held_by victim_pid, survivor_pid = sorted(workers, key=held_by.__getitem__) - victim: Final = psutil.Process(victim_pid) - victim.suspend() - victim.send_signal(signal.SIGKILL) - release.set() - served: Final = await burst - assert held_by[survivor_pid] >= 10, held_by - assert len(served) == held_by[survivor_pid], (held_by, len(served)) - for item in served: - _assert_answered_with_its_own_marker(item) + with owned.replacing_worker(victim_pid) as kill: + kill.freeze() + await asyncio.to_thread(kill) + release.set() + served: Final = await burst + assert held_by[survivor_pid] >= 10, held_by + assert len(served) == held_by[survivor_pid], (held_by, len(served)) + for item in served: + _assert_answered_with_its_own_marker(item) follow_up: Final = _Call("responses", False, uuid.uuid4().hex) (answered,) = await _burst(str(candidate.client.base_url), candidate.key, models, (follow_up,)) _assert_answered_with_its_own_marker(answered) diff --git a/tests/integration/providers/test_responses_minted_reasoning_replay_wire.py b/tests/integration/providers/test_responses_minted_reasoning_replay_wire.py index 036f96916a6..2253443dfc0 100644 --- a/tests/integration/providers/test_responses_minted_reasoning_replay_wire.py +++ b/tests/integration/providers/test_responses_minted_reasoning_replay_wire.py @@ -17,6 +17,7 @@ from integration._support import claude_code as cc from integration._support import responses_vendor as rv from integration._support.client import Gateway, Scenario, eventually from integration._support.database import read_rows +from integration._support.timing import assert_within_factor_of_control, timed from integration._support.wire import Request, Wire, wire_server from pydantic import JsonValue, TypeAdapter @@ -26,6 +27,7 @@ _CLAUDE: Final = cc.OPUS _OPENAI_KEY: Final = "synthetic-openai-key" _AZURE_KEY: Final = "synthetic-azure-key" _CACHE_BUST: Final[Mapping[str, JsonValue]] = MappingProxyType({"cache": {"no-cache": True}}) +_RESPONSIVE_FACTOR: Final = 10 _JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) _ITEMS: Final = TypeAdapter(list[dict[str, JsonValue]]) @@ -564,6 +566,11 @@ def test_a_megabyte_of_minted_thinking_is_dropped_while_the_proxy_stays_responsi "id": f"rs_{uuid.uuid4()}", "encrypted_content": encrypted, } + single: Final[dict[str, JsonValue]] = { + **minted, + "id": f"rs_{uuid.uuid4()}", + "encrypted_content": json.dumps([block]), + } history: Final = rv.agents_sdk_history(marker, minted) latencies: Final[deque[float]] = deque() done: Final = threading.Event() @@ -577,16 +584,22 @@ def test_a_megabyte_of_minted_thinking_is_dropped_while_the_proxy_stays_responsi with wire_server(rv.ResponsesVendor().respond) as wire, gateway.scenario() as scenario: model: Final = _OPENAI.register(scenario, wire) + same_size: Final[list[dict[str, JsonValue]]] = [ + {"role": "user", "content": "t" * len(encrypted)}, + *rv.agents_sdk_history(marker, single), + ] + control: Final = timed(lambda: _raw(gateway, "/v1/responses", {"model": model, "input": same_size})) + assert control.value.status_code == 200, control.value.text[:500] + _only_request(wire) prober: Final = threading.Thread(target=probe) prober.start() - started: Final = time.monotonic() - response: Final = _raw(gateway, "/v1/responses", {"model": model, "input": history}) - elapsed: Final = time.monotonic() - started + dropped: Final = timed(lambda: _raw(gateway, "/v1/responses", {"model": model, "input": history})) done.set() prober.join(timeout=35) - assert response.status_code == 200, response.text[:500] - assert elapsed < 20, elapsed - assert latencies and max(latencies) < 5, (max(latencies), len(latencies)) + assert dropped.value.status_code == 200, dropped.value.text[:500] + assert_within_factor_of_control(dropped.seconds, control=control.seconds, factor=_RESPONSIVE_FACTOR) + assert latencies, "no liveliness probe completed while the megabyte request was in flight" + assert_within_factor_of_control(max(latencies), control=control.seconds, factor=_RESPONSIVE_FACTOR) _, body = _only_request(wire) assert body["input"] == rv.without(history, (minted,)) diff --git a/tests/integration/providers/test_responses_websocket_session_limit.py b/tests/integration/providers/test_responses_websocket_session_limit.py index 32b01d8387a..f2b3a48c6d1 100644 --- a/tests/integration/providers/test_responses_websocket_session_limit.py +++ b/tests/integration/providers/test_responses_websocket_session_limit.py @@ -24,6 +24,7 @@ import websockets from integration._support.client import Gateway, eventually, gateway_from_environment from integration._support.database import read_rows, scratch_database from integration._support.process import OwnedProxy, graceful_stop_seconds, owned_proxy, owned_proxy_process +from integration._support.timing import assert_cut_short, assert_enforced_within from integration._support.tls import server_context, write_self_signed_cert from integration._support.upstream import ( JsonResponse, @@ -44,6 +45,8 @@ STALL_PROVIDER_MODEL: Final = "ws-stall-peer-model" DEAF_PROVIDER_MODEL: Final = "ws-deaf-peer-model" DEAF_RESPONSE_DELAY_SECONDS: Final = 10 DEAF_READ_PAUSE_SECONDS: Final = 20 +MID_TURN_SECONDS: Final = 40 +DEAF_UNTIL_SECONDS: Final = MID_TURN_SECONDS + DEAF_RESPONSE_DELAY_SECONDS + DEAF_READ_PAUSE_SECONDS PEER_TEXT: Final = "responses websocket peer" TERMINAL: Final = frozenset({"response.completed", "response.failed", "error"}) JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) @@ -52,6 +55,8 @@ LIMIT_CLOSE_REASON: Final = "Session duration limit reached" RELOAD_INTERVAL_SECONDS: Final = 3 WORKER_SYNC_SECONDS: Final = RELOAD_INTERVAL_SECONDS + 7 CAP_SECONDS: Final = 60 +CAP_CLOSE_SLACK_SECONDS: Final = 15 +DEAF_CLOSE_SLACK_SECONDS: Final = 3 RESTORED_IDLE_SECONDS: Final = 75 CAPPED_PATHS: Final = ("/v1/responses", "/responses") CAPPED_SOCKETS_MINIMUM: Final = 8 @@ -725,7 +730,7 @@ async def _mid_session( additional_headers={"Authorization": f"Bearer {key}"}, open_timeout=10, ) as connection: - await asyncio.sleep(max(0, started + 40 - time.monotonic())) + await asyncio.sleep(max(0, started + MID_TURN_SECONDS - time.monotonic())) created, turn_error = await _stall_turn(connection, stall_model) remaining: Final = max(0, 75 - (time.monotonic() - started)) close: Final = CloseResult(await _wait_for_close(connection, remaining), time.monotonic() - started) @@ -982,19 +987,23 @@ async def _kill_workload(owned: OwnedProxy, key: str, model: str) -> KillResults ) held_by: Final = MappingProxyType({pid: holders.count(pid) for pid in sorted(workers)}) victim: Final = max(sorted(workers), key=held_by.__getitem__) - psutil.Process(victim).kill() victims: Final = tuple( connection for connection, holder in zip(connections, holders, strict=True) if holder == victim ) survivors: Final = tuple( connection for connection, holder in zip(connections, holders, strict=True) if holder != victim ) - victim_closes, survivor_turns = await asyncio.gather( - asyncio.gather(*tuple(_wait_for_close(connection, 30) for connection in victims)), - asyncio.gather( - *tuple(_turn_result(connection, _create(model, f"survivor-{uuid.uuid4().hex}")) for connection in survivors) - ), - ) + with owned.replacing_worker(victim) as kill: + await asyncio.to_thread(kill) + victim_closes, survivor_turns = await asyncio.gather( + asyncio.gather(*tuple(_wait_for_close(connection, 30) for connection in victims)), + asyncio.gather( + *tuple( + _turn_result(connection, _create(model, f"survivor-{uuid.uuid4().hex}")) + for connection in survivors + ) + ), + ) health: Final = await asyncio.to_thread( eventually, lambda: _health_status(candidate), @@ -1327,7 +1336,9 @@ def test_session_cap_closes_never_started_socket(cap_results: CapResults) -> Non assert cap_results.idle.outcome.closed, cap_results.idle assert cap_results.idle.outcome.close_code == 1000, cap_results.idle assert cap_results.idle.outcome.close_reason == "Session duration limit reached", cap_results.idle - assert 59 <= cap_results.idle.elapsed <= 75, cap_results.idle + assert_enforced_within( + cap_results.idle.elapsed, configured=CAP_SECONDS, slack=CAP_CLOSE_SLACK_SECONDS, not_before=CAP_SECONDS - 1 + ) def test_session_cap_closes_active_socket_and_provider(cap_results: CapResults) -> None: @@ -1336,7 +1347,12 @@ def test_session_cap_closes_active_socket_and_provider(cap_results: CapResults) assert cap_results.active.close.outcome.closed, cap_results.active assert cap_results.active.close.outcome.close_code == 1000, cap_results.active assert cap_results.active.close.outcome.close_reason == "Session duration limit reached", cap_results.active - assert 59 <= cap_results.active.close.elapsed <= 75, cap_results.active + assert_enforced_within( + cap_results.active.close.elapsed, + configured=CAP_SECONDS, + slack=CAP_CLOSE_SLACK_SECONDS, + not_before=CAP_SECONDS - 1, + ) assert cap_results.active.provider_closed, cap_results.active assert all(snapshot.closed for snapshot in cap_results.provider), cap_results.provider @@ -1347,7 +1363,9 @@ def test_session_cap_closes_mid_response_and_allows_new_session(cap_results: Cap assert cap_results.mid.close.outcome.closed, cap_results.mid assert cap_results.mid.close.outcome.close_code == 1000, cap_results.mid assert cap_results.mid.close.outcome.close_reason == "Session duration limit reached", cap_results.mid - assert 59 <= cap_results.mid.close.elapsed <= 75, cap_results.mid + assert_enforced_within( + cap_results.mid.close.elapsed, configured=CAP_SECONDS, slack=CAP_CLOSE_SLACK_SECONDS, not_before=CAP_SECONDS - 1 + ) assert cap_results.mid.provider_closed, cap_results.mid assert cap_results.mid.fresh_error is None, cap_results.mid assert cap_results.mid.fresh_completed, cap_results.mid @@ -1359,7 +1377,13 @@ def test_session_cap_closes_client_promptly_when_provider_ignores_close(cap_resu assert cap_results.deaf.close.outcome.closed, cap_results.deaf assert cap_results.deaf.close.outcome.close_code == 1000, cap_results.deaf assert cap_results.deaf.close.outcome.close_reason == "Session duration limit reached", cap_results.deaf - assert 59 <= cap_results.deaf.close.elapsed <= 63, cap_results.deaf + assert_enforced_within( + cap_results.deaf.close.elapsed, + configured=CAP_SECONDS, + slack=DEAF_CLOSE_SLACK_SECONDS, + not_before=CAP_SECONDS - 1, + ) + assert_cut_short(cap_results.deaf.close.elapsed, injected_delay=DEAF_UNTIL_SECONDS) assert cap_results.deaf.provider_closed, cap_results.deaf assert cap_results.deaf.fresh_error is None, cap_results.deaf assert cap_results.deaf.fresh_completed, cap_results.deaf @@ -1420,7 +1444,10 @@ def test_db_override_caps_new_sessions_on_every_worker_without_restart(override_ ) assert all(socket.close.outcome.closed and socket.close.outcome.close_code == 1000 for socket in capped), capped assert all(socket.close.outcome.close_reason == LIMIT_CLOSE_REASON for socket in capped), capped - assert all(CAP_SECONDS - 1 <= socket.close.elapsed <= CAP_SECONDS + 15 for socket in capped), capped + for socket in capped: + assert_enforced_within( + socket.close.elapsed, configured=CAP_SECONDS, slack=CAP_CLOSE_SLACK_SECONDS, not_before=CAP_SECONDS - 1 + ) assert frozenset(socket.worker_pid for socket in capped) == override_results.workers, capped diff --git a/tests/integration/providers/test_sagemaker_chat_credential_sources.py b/tests/integration/providers/test_sagemaker_chat_credential_sources.py index daaec6a7583..b01454acfcf 100644 --- a/tests/integration/providers/test_sagemaker_chat_credential_sources.py +++ b/tests/integration/providers/test_sagemaker_chat_credential_sources.py @@ -1,5 +1,4 @@ import os -import signal from collections.abc import Mapping from pathlib import Path from types import MappingProxyType @@ -186,13 +185,6 @@ def _worker_pids(log: str) -> tuple[int, ...]: ) -def _dead(process: psutil.Process) -> bool: - try: - return process.status() == psutil.STATUS_ZOMBIE - except psutil.NoSuchProcess: - return True - - @pytest.mark.timeout(_OWNED_PROXY_CELL_SECONDS) def test_every_aws_credential_source_keeps_its_credentials_out_of_the_body_when_a_worker_is_killed( gateway: Gateway, tmp_path: Path @@ -212,11 +204,10 @@ def test_every_aws_credential_source_keeps_its_credentials_out_of_the_body_when_ _assert_each_source_signed_its_own_clean_body( container, {source: _ask(owned, source) for source in _SIGNERS} ) - victim: Final = psutil.Process(workers[0]) - victim.send_signal(signal.SIGKILL) - eventually(lambda: _dead(victim), bool, seconds=graceful_stop_seconds()) - _assert_each_source_signed_its_own_clean_body( - container, {source: _ask(owned, source) for source in _SIGNERS} - ) - assert psutil.Process(workers[1]).is_running(), workers + with owned.replacing_worker(workers[0]) as kill: + kill() + _assert_each_source_signed_its_own_clean_body( + container, {source: _ask(owned, source) for source in _SIGNERS} + ) + assert psutil.Process(workers[1]).is_running(), workers assert owned.process.poll() is None, owned.log.read_text() diff --git a/tests/integration/providers/test_sagemaker_embedding_chaos.py b/tests/integration/providers/test_sagemaker_embedding_chaos.py new file mode 100644 index 00000000000..680ad3bfd2d --- /dev/null +++ b/tests/integration/providers/test_sagemaker_embedding_chaos.py @@ -0,0 +1,302 @@ +import asyncio +import re +import socket +import threading +import uuid +from collections.abc import Iterator +from dataclasses import dataclass +from pathlib import Path +from queue import SimpleQueue +from types import MappingProxyType +from typing import Final +from urllib.parse import urlsplit + +import httpx +import psutil +import pytest +import yaml +from integration._support.client import Gateway, eventually, gateway_from_environment, string_value +from integration._support.database import read_rows +from integration._support.process import owned_proxy, owned_proxy_process +from integration._support.sagemaker_embedding import ( + COMPONENT, + COMPONENT_HEADER, + JSON_OBJECT, + SERVED_MODEL, + json_reply, + openai_deployment, + openai_payload, + respond, +) +from integration._support.wire import Reply, Request, wire_server +from pydantic import JsonValue + +pytestmark = pytest.mark.timeout(300) + +_MODEL: Final = "sm-openai-chaos" +_ALIASES: Final = ( + "/v1/embeddings", + "/embeddings", + f"/openai/deployments/{_MODEL}/embeddings", + f"/engines/{_MODEL}/embeddings", +) +_STARTED_WORKER: Final = re.compile(r"Started server process \[(\d+)\]") +_ROWS_BY_CALL: Final = ( + "SELECT litellm_call_id, status FROM \"LiteLLM_SpendLogs\" WHERE litellm_call_id = ANY(string_to_array(%s, ','))" +) + + +@dataclass(frozen=True, slots=True) +class _Served: + text: str + status: int + body: str + call_id: str + + +def _free_port() -> int: + with socket.socket() as probe: + probe.bind(("127.0.0.1", 0)) + return probe.getsockname()[1] + + +def _config(directory: Path) -> Path: + config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + config["model_list"] = [openai_deployment(_MODEL, "openai", model_info={"mode": "embedding"})] + path: Final = directory / "sagemaker-embedding-chaos.yaml" + path.write_text(yaml.safe_dump(config)) + return path + + +def _runtime(peer_url: str) -> dict[str, str]: + return {"AWS_ENDPOINT_URL_SAGEMAKER_RUNTIME": peer_url, "AWS_MAX_ATTEMPTS": "1"} + + +@pytest.fixture(scope="module") +def rig() -> Iterator[Gateway]: + with gateway_from_environment() as gateway: + yield gateway + + +@pytest.fixture(scope="module") +def peer_port() -> int: + return _free_port() + + +@pytest.fixture(scope="module") +def proxy(rig: Gateway, peer_port: int, tmp_path_factory: pytest.TempPathFactory) -> Iterator[Gateway]: + directory: Final = tmp_path_factory.mktemp("sagemaker-embedding-chaos") + runtime: Final = _runtime(f"http://127.0.0.1:{peer_port}") + with owned_proxy(rig, directory, runtime, config=_config(directory), workers=2) as owned: + yield owned + + +def _texts(count: int) -> tuple[str, ...]: + return tuple(f"integration embedding {uuid.uuid4().hex}" for _ in range(count)) + + +async def _send(client: httpx.AsyncClient, key: str, index: int, text: str) -> _Served: + path: Final = _ALIASES[index % len(_ALIASES)] + body: Final[dict[str, JsonValue]] = { + "input": text, + "cache": {"no-cache": True}, + **({"model": _MODEL} if path.endswith("/v1/embeddings") or path == "/embeddings" else {}), + } + response: Final = await client.post(path, json=body, headers={"Authorization": f"Bearer {key}"}) + return _Served(text, response.status_code, response.text, response.headers.get("x-litellm-call-id", "")) + + +async def _burst( + base_url: str, key: str, texts: tuple[str, ...], *, tolerate_transport_errors: bool = False +) -> tuple[_Served, ...]: + async with httpx.AsyncClient(base_url=base_url, timeout=90, trust_env=False) as client: + results: Final = await asyncio.gather( + *(_send(client, key, index, text) for index, text in enumerate(texts)), + return_exceptions=tolerate_transport_errors, + ) + for result in results: + assert not isinstance(result, BaseException) or isinstance(result, httpx.TransportError), repr(result) + return tuple(result for result in results if isinstance(result, _Served)) + + +def _assert_embedded(served: _Served) -> None: + assert served.status == 200, served.body + payload: Final = JSON_OBJECT.validate_json(served.body) + assert payload["data"] == [{"object": "embedding", "index": 0, "embedding": [0.25, -0.5, 1.0]}], served.body + + +def _assert_peer_saw(received: tuple[Request, ...], texts: tuple[str, ...]) -> None: + assert {request.headers.get(COMPONENT_HEADER) for request in received} == {COMPONENT}, received + bodies: Final = [JSON_OBJECT.validate_json(request.body) for request in received] + assert sorted(string_value(body["input"]) for body in bodies) == sorted(texts), bodies + assert {body["model"] for body in bodies} == {SERVED_MODEL}, bodies + + +def _spend_status_by_call(served: tuple[_Served, ...]) -> dict[str, JsonValue]: + wanted: Final = sorted(item.call_id for item in served) + rows: Final = eventually( + lambda: read_rows(_ROWS_BY_CALL, (",".join(wanted),)), lambda found: len(found) >= len(wanted), seconds=70 + ) + assert sorted(string_value(row["litellm_call_id"]) for row in rows) == wanted, rows + return {string_value(row["litellm_call_id"]): row["status"] for row in rows} + + +def _unhealthy_count(gateway: Gateway) -> JsonValue: + response: Final = gateway.request("GET", "/health", params={"model": _MODEL}) + return JSON_OBJECT.validate_json(response.content).get("unhealthy_count") + + +async def test_concurrent_burst_across_every_alias_reaches_the_component_and_is_logged_once_per_call( + proxy: Gateway, peer_port: int +) -> None: + texts: Final = _texts(24) + with wire_server(respond, port=peer_port) as wire: + served: Final = await _burst(str(proxy.client.base_url), proxy.key, texts) + assert len(served) == 24 + for item in served: + _assert_embedded(item) + _assert_peer_saw(wire.drain(), texts) + assert _spend_status_by_call(served) == {item.call_id: "success" for item in served} + + +async def test_peer_outage_fails_each_call_once_reports_unhealthy_and_recovers_on_restart( + proxy: Gateway, peer_port: int +) -> None: + before: Final = _texts(12) + during: Final = _texts(4) + after: Final = _texts(4) + base_url: Final = str(proxy.client.base_url) + with wire_server(respond, port=peer_port) as wire: + served: Final = await _burst(base_url, proxy.key, before) + assert len(served) == 12 + for item in served: + _assert_embedded(item) + _assert_peer_saw(wire.drain(), before) + refused: Final = await _burst(base_url, proxy.key, during) + assert len(refused) == 4 + for item in refused: + assert item.status == 503, item.body + assert "Could not connect to the endpoint URL" in item.body, item.body + assert await asyncio.to_thread(eventually, lambda: _unhealthy_count(proxy), lambda count: count == 1, 60) + with wire_server(respond, port=peer_port) as restarted: + recovered: Final = await _burst(base_url, proxy.key, after) + assert len(recovered) == 4 + for item in recovered: + _assert_embedded(item) + healthy: Final = [request for request in restarted.drain() if "test from litellm" not in request.body.decode()] + _assert_peer_saw(tuple(healthy), after) + assert _spend_status_by_call((*served, *refused, *recovered)) == { + **{item.call_id: "success" for item in served}, + **{item.call_id: "failure" for item in refused}, + **{item.call_id: "success" for item in recovered}, + } + + +def _slow(request: Request) -> Reply: + body: Final = json_reply(openai_payload(JSON_OBJECT.validate_json(request.body))).body + return Reply(chunks=(body[:8], body[8:]), pause_between_chunks=0.5) + + +async def test_slow_peer_answers_every_caller_whole_without_duplicate_rows(proxy: Gateway, peer_port: int) -> None: + texts: Final = _texts(10) + with wire_server(_slow, port=peer_port) as wire: + served: Final = await _burst(str(proxy.client.base_url), proxy.key, texts) + assert len(served) == 10 + for item in served: + _assert_embedded(item) + _assert_peer_saw(wire.drain(), texts) + assert _spend_status_by_call(served) == {item.call_id: "success" for item in served} + + +def _live_workers(log: Path) -> tuple[int, ...]: + started: Final = (int(pid) for pid in _STARTED_WORKER.findall(log.read_text())) + return tuple(pid for pid in started if psutil.pid_exists(pid)) + + +def _open_peer_connections(pid: int, peer_url: str) -> int: + port: Final = urlsplit(peer_url).port + return sum( + 1 + for connection in psutil.Process(pid).net_connections(kind="tcp") + if connection.status == psutil.CONN_ESTABLISHED and connection.raddr and connection.raddr.port == port + ) + + +@dataclass(frozen=True, slots=True) +class _HeldPeer: + release: threading.Event + held: SimpleQueue[str] + + def respond(self, request: Request) -> Reply: + self.held.put(string_value(JSON_OBJECT.validate_json(request.body)["input"])) + assert self.release.wait(timeout=120), "The burst was never released" + return respond(request) + + +@pytest.mark.timeout(600) +async def test_worker_sigkill_mid_burst_leaves_the_sibling_serving_the_component(rig: Gateway, tmp_path: Path) -> None: + texts: Final = _texts(8) + peer: Final = _HeldPeer(threading.Event(), SimpleQueue()) + with wire_server(peer.respond) as wire: + with owned_proxy_process(rig, tmp_path, _runtime(wire.url), config=_config(tmp_path), workers=2) as owned: + candidate: Final = owned.gateway + workers: Final = eventually(lambda: _live_workers(owned.log), lambda pids: len(pids) == 2, seconds=30) + burst: Final = asyncio.create_task( + _burst(str(candidate.client.base_url), candidate.key, texts, tolerate_transport_errors=True) + ) + await asyncio.to_thread(eventually, peer.held.qsize, lambda size: size == len(texts), 120) + held_by: Final = MappingProxyType({pid: _open_peer_connections(pid, wire.url) for pid in workers}) + assert sum(held_by.values()) == len(texts), held_by + victim_pid, survivor_pid = sorted(workers, key=held_by.__getitem__) + with owned.replacing_worker(victim_pid) as kill: + kill.freeze() + await asyncio.to_thread(kill) + peer.release.set() + served: Final = await burst + assert len(served) == held_by[survivor_pid], (held_by, len(served)) + for item in served: + _assert_embedded(item) + (follow_up,) = _texts(1) + (answered,) = await _burst(str(candidate.client.base_url), candidate.key, (follow_up,)) + _assert_embedded(answered) + _assert_peer_saw(wire.drain(), (*texts, follow_up)) + assert _spend_status_by_call((*served, answered)) == { + item.call_id: "success" for item in (*served, answered) + } + + +@pytest.mark.timeout(600) +async def test_proxy_sigterm_mid_burst_finishes_in_flight_calls_and_a_relaunch_serves_again( + rig: Gateway, tmp_path: Path +) -> None: + texts: Final = _texts(8) + peer: Final = _HeldPeer(threading.Event(), SimpleQueue()) + with wire_server(peer.respond) as wire: + first: Final = tmp_path / "first" + first.mkdir() + with owned_proxy_process(rig, first, _runtime(wire.url), config=_config(first), workers=2) as owned: + burst: Final = asyncio.create_task( + _burst(str(owned.gateway.client.base_url), owned.gateway.key, texts, tolerate_transport_errors=True) + ) + await asyncio.to_thread(eventually, peer.held.qsize, lambda size: size == len(texts), 120) + owned.begin_shutdown() + peer.release.set() + served: Final = await burst + exit_code: Final = await asyncio.to_thread( + eventually, owned.process.poll, lambda code: code is not None, 120 + ) + assert exit_code is not None + for item in served: + _assert_embedded(item) + _assert_peer_saw(wire.drain(), texts) + assert len(served) == len(texts), (len(served), len(texts)) + second: Final = tmp_path / "second" + second.mkdir() + with owned_proxy_process(rig, second, _runtime(wire.url), config=_config(second), workers=2) as relaunched: + (follow_up,) = _texts(1) + (answered,) = await _burst(str(relaunched.gateway.client.base_url), relaunched.gateway.key, (follow_up,)) + _assert_embedded(answered) + _assert_peer_saw(wire.drain(), (follow_up,)) + assert _spend_status_by_call((*served, answered)) == { + item.call_id: "success" for item in (*served, answered) + } diff --git a/tests/integration/providers/test_sagemaker_embedding_wire.py b/tests/integration/providers/test_sagemaker_embedding_wire.py new file mode 100644 index 00000000000..0ed38b1a768 --- /dev/null +++ b/tests/integration/providers/test_sagemaker_embedding_wire.py @@ -0,0 +1,742 @@ +import json +import uuid +from collections.abc import Iterator, Mapping +from pathlib import Path +from typing import Final + +import httpx +import openai +import pytest +import yaml +from botocore.exceptions import ConnectionClosedError +from integration._support.client import Gateway, eventually, gateway_from_environment +from integration._support.database import read_rows +from integration._support.process import owned_proxy +from integration._support.sagemaker_embedding import ( + ACCESS_KEY, + COMPONENT, + COMPONENT_HEADER, + EXPECTED_SHAPE, + IC_NOT_ALLOWED, + IC_REQUIRED, + JSON_OBJECT, + NON_STRING_MODEL_ID, + NOT_JSON_BODY, + PASS_IN_HINT, + REMOVE_HINT, + SERVED_MODEL, + deployment, + embedding_value, + endpoint, + invocations_target, + openai_deployment, + openai_route, + plain_route, + respond, +) +from integration._support.wire import Request, Wire, wire_server +from pydantic import JsonValue + +pytestmark = pytest.mark.timeout(240) + +_NO_CACHE: Final[dict[str, JsonValue]] = {"no-cache": True} +_SPEND_ROW: Final = 'SELECT status, call_type, model_group FROM "LiteLLM_SpendLogs" WHERE litellm_call_id=%s' + + +def _model_list(health_endpoint: str) -> list[JsonValue]: + return [ + openai_deployment("sm-openai", "openai"), + deployment( + "sm-openai-explicit", + f"openai/{endpoint('openai')}", + custom_llm_provider="sagemaker", + model_id=COMPONENT, + extra_body={"model": SERVED_MODEL}, + ), + deployment("sm-openai-bare", openai_route("openai")), + deployment("sm-openai-no-model-id", openai_route("openai"), extra_body={"model": SERVED_MODEL}), + openai_deployment("sm-openai-custom-key", "openai", input_type="query"), + deployment( + "sm-openai-pinned", + openai_route("openai"), + model_id=COMPONENT, + extra_body={"model": SERVED_MODEL, "encoding_format": "float"}, + ), + openai_deployment("sm-openai-model-id-int", "openai", model_id=7), + openai_deployment("sm-openai-model-id-list", "openai", model_id=[COMPONENT]), + openai_deployment("sm-openai-model-id-empty", "openai", model_id=""), + openai_deployment("sm-openai-model-id-null", "openai", model_id=None), + openai_deployment("sm-openai-model-id-long", "openai", model_id="c" * 5120), + openai_deployment("sm-openai-extra-body-string", "openai", extra_body="not-an-object"), + openai_deployment("sm-openai-extra-body-empty", "openai", extra_body={}), + openai_deployment("sm-openai-nousage", "nousage"), + openai_deployment("sm-openai-nomodel", "nomodel"), + openai_deployment("sm-openai-noindex", "noindex"), + openai_deployment("sm-openai-promptonly", "promptonly"), + openai_deployment("sm-openai-datadict", "datadict"), + openai_deployment("sm-openai-hfshape", "hfshape"), + openai_deployment("sm-openai-notjson", "notjson"), + openai_deployment("sm-openai-err400", "err400"), + openai_deployment("sm-openai-err424", "err424"), + openai_deployment("sm-openai-err429", "err429"), + openai_deployment("sm-openai-err503", "err503"), + openai_deployment("sm-openai-err500", "err500"), + openai_deployment("sm-openai-drop", "drop"), + deployment("sm-openai-unknown", "sagemaker/openai/unknown-embed-endpoint", model_id=COMPONENT), + openai_deployment("sm-openai-ic", "icrequired"), + deployment("sm-openai-ic-missing", openai_route("icrequired"), extra_body={"model": SERVED_MODEL}), + deployment("sm-plain-noic", plain_route("noic"), model_id=COMPONENT), + deployment("sm-plain-noic-control", plain_route("noic")), + deployment( + "sm-openai-health", + f"sagemaker/openai/{health_endpoint}", + model_id=COMPONENT, + extra_body={"model": SERVED_MODEL}, + model_info={"mode": "embedding"}, + ), + deployment("sm-plain", plain_route("hf"), model_id=COMPONENT), + deployment("sm-plain-control", plain_route("hf")), + deployment("sm-plain-drop-params", plain_route("hf"), model_id=COMPONENT, drop_params=True), + deployment("sm-plain-dict", plain_route("hfdict"), model_id=COMPONENT), + deployment("sm-plain-err424", plain_route("err424"), model_id=COMPONENT), + deployment("sm-voyage", plain_route("voyage"), model_id=COMPONENT), + deployment("sm-cohere", plain_route("cohere"), model_id=COMPONENT), + ] + + +def _config(directory: Path, health_endpoint: str) -> Path: + config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + config["model_list"] = _model_list(health_endpoint) + path: Final = directory / "sagemaker-embedding.yaml" + path.write_text(yaml.safe_dump(config)) + return path + + +@pytest.fixture(scope="module") +def rig() -> Iterator[Gateway]: + with gateway_from_environment() as gateway: + yield gateway + + +@pytest.fixture(scope="module") +def wire() -> Iterator[Wire]: + with wire_server(respond) as served: + yield served + + +@pytest.fixture(autouse=True) +def _drained_wire(wire: Wire) -> None: + wire.drain() + + +@pytest.fixture(scope="module") +def health_endpoint() -> str: + return endpoint(f"openai-health-{uuid.uuid4().hex[:8]}") + + +@pytest.fixture(scope="module") +def proxy( + rig: Gateway, wire: Wire, health_endpoint: str, tmp_path_factory: pytest.TempPathFactory +) -> Iterator[Gateway]: + directory: Final = tmp_path_factory.mktemp("sagemaker-embedding") + overrides: Final = {"AWS_ENDPOINT_URL_SAGEMAKER_RUNTIME": wire.url, "AWS_MAX_ATTEMPTS": "2"} + with owned_proxy(rig, directory, overrides, config=_config(directory, health_endpoint), workers=2) as owned: + yield owned + + +def _text() -> str: + return f"integration embedding {uuid.uuid4().hex}" + + +def _embed(proxy: Gateway, model: str, text: JsonValue, **fields: JsonValue) -> httpx.Response: + return proxy.request("POST", "/v1/embeddings", {"model": model, "input": text, "cache": _NO_CACHE, **fields}) + + +def _spend_row(response: httpx.Response) -> dict[str, JsonValue]: + call_id: Final = response.headers["x-litellm-call-id"] + rows: Final = eventually(lambda: read_rows(_SPEND_ROW, (call_id,)), lambda found: len(found) == 1, seconds=70) + return rows[0] + + +def _usage(payload: Mapping[str, JsonValue]) -> tuple[JsonValue, ...]: + usage: Final = payload["usage"] + assert isinstance(usage, dict), payload + return tuple(usage[key] for key in ("prompt_tokens", "completion_tokens", "total_tokens")) + + +def _only_request(wire: Wire) -> Request: + received: Final = wire.drain() + assert len(received) == 1, [(request.method, request.target) for request in received] + return received[0] + + +def _body(request: Request) -> dict[str, JsonValue]: + return JSON_OBJECT.validate_json(request.body) + + +def _forwarded_body(sent: httpx.Request, text: str) -> dict[str, JsonValue]: + asked: Final = JSON_OBJECT.validate_json(sent.content) + encoding: Final = {"encoding_format": asked["encoding_format"]} if "encoding_format" in asked else {} + return {"input": [text], **encoding, "model": SERVED_MODEL} + + +def _signed_headers(request: Request) -> tuple[str, ...]: + authorization: Final = request.headers["authorization"] + assert authorization.startswith(f"AWS4-HMAC-SHA256 Credential={ACCESS_KEY}/"), authorization + signed: Final = next(part for part in authorization.split(", ") if part.startswith("SignedHeaders=")) + return tuple(signed.removeprefix("SignedHeaders=").split(";")) + + +def _sdk_base_url(proxy: Gateway) -> str: + return f"{str(proxy.client.base_url).rstrip('/')}/v1" + + +def _sdk(proxy: Gateway) -> openai.OpenAI: + return openai.OpenAI(base_url=_sdk_base_url(proxy), api_key=proxy.key, max_retries=0, timeout=60) + + +def _async_sdk(proxy: Gateway) -> openai.AsyncOpenAI: + return openai.AsyncOpenAI(base_url=_sdk_base_url(proxy), api_key=proxy.key, max_retries=0, timeout=60) + + +def test_openai_route_signs_model_id_as_the_inference_component_header_and_sends_the_openai_body( + proxy: Gateway, wire: Wire +) -> None: + texts: Final = [_text(), _text()] + response: Final = _embed(proxy, "sm-openai", texts) + assert response.status_code == 200, response.text + payload: Final = JSON_OBJECT.validate_json(response.content) + assert payload["data"] == [ + {"object": "embedding", "index": 0, "embedding": [0.25, -0.5, 1.0]}, + {"object": "embedding", "index": 1, "embedding": [0.25, -0.5, 2.0]}, + ], response.text + assert payload["model"] == "sm-openai", response.text + assert _usage(payload) == (7, 0, 7), response.text + request: Final = _only_request(wire) + assert (request.method, request.target) == ("POST", invocations_target(endpoint("openai"))), request + assert request.headers[COMPONENT_HEADER] == COMPONENT, dict(request.headers) + assert request.headers["x-amzn-sagemaker-custom-attributes"] == "accept_eula=true", dict(request.headers) + assert request.headers["content-type"] == "application/json", dict(request.headers) + assert COMPONENT_HEADER in _signed_headers(request), request.headers["authorization"] + assert _body(request) == {"input": texts, "model": SERVED_MODEL}, request.body + row: Final = _spend_row(response) + assert row == {"status": "success", "call_type": "aembedding", "model_group": "sm-openai"}, row + + +@pytest.mark.parametrize( + ("path", "model_in_body"), + [ + ("/v1/embeddings", True), + ("/embeddings", True), + ("/openai/deployments/sm-openai/embeddings", False), + ("/engines/sm-openai/embeddings", False), + ], +) +def test_every_embeddings_alias_reaches_the_inference_component( + proxy: Gateway, wire: Wire, path: str, model_in_body: bool +) -> None: + text: Final = _text() + body: Final[dict[str, JsonValue]] = { + "input": text, + "cache": _NO_CACHE, + **({"model": "sm-openai"} if model_in_body else {}), + } + response: Final = proxy.request("POST", path, body) + assert response.status_code == 200, response.text + payload: Final = JSON_OBJECT.validate_json(response.content) + assert payload["data"] == [{"object": "embedding", "index": 0, "embedding": [0.25, -0.5, 1.0]}], response.text + request: Final = _only_request(wire) + assert request.headers[COMPONENT_HEADER] == COMPONENT, dict(request.headers) + assert _body(request) == {"input": text, "model": SERVED_MODEL}, request.body + assert _spend_row(response)["status"] == "success" + + +def test_openai_sdk_sync_client_reads_the_vectors_in_the_encoding_it_asked_for(proxy: Gateway, wire: Wire) -> None: + text: Final = _text() + raw: Final = _sdk(proxy).embeddings.with_raw_response.create( + model="sm-openai", input=[text], extra_body={"cache": _NO_CACHE} + ) + created: Final = raw.parse() + assert [row.embedding for row in created.data] == [[0.25, -0.5, 1.0]], created + assert created.model == "sm-openai", created + assert (created.usage.prompt_tokens, created.usage.total_tokens) == (7, 7), created + request: Final = _only_request(wire) + assert request.headers[COMPONENT_HEADER] == COMPONENT, dict(request.headers) + assert _body(request) == _forwarded_body(raw.http_request, text), request.body + call_id: Final = raw.headers["x-litellm-call-id"] + rows: Final = eventually(lambda: read_rows(_SPEND_ROW, (call_id,)), lambda found: len(found) == 1, seconds=70) + assert rows[0]["status"] == "success", rows + + +async def test_openai_sdk_async_client_reads_the_vectors_in_the_encoding_it_asked_for( + proxy: Gateway, wire: Wire +) -> None: + text: Final = _text() + raw: Final = await _async_sdk(proxy).embeddings.with_raw_response.create( + model="sm-openai", input=[text], extra_body={"cache": _NO_CACHE} + ) + created: Final = raw.parse() + assert [row.embedding for row in created.data] == [[0.25, -0.5, 1.0]], created + assert created.model == "sm-openai", created + request: Final = _only_request(wire) + assert request.headers[COMPONENT_HEADER] == COMPONENT, dict(request.headers) + assert _body(request) == _forwarded_body(raw.http_request, text), request.body + call_id: Final = raw.headers["x-litellm-call-id"] + rows: Final = eventually(lambda: read_rows(_SPEND_ROW, (call_id,)), lambda found: len(found) == 1, seconds=70) + assert rows[0]["status"] == "success", rows + + +def test_explicit_custom_llm_provider_spelling_takes_the_openai_route(proxy: Gateway, wire: Wire) -> None: + text: Final = _text() + response: Final = _embed(proxy, "sm-openai-explicit", [text]) + assert response.status_code == 200, response.text + payload: Final = JSON_OBJECT.validate_json(response.content) + assert payload["model"] == "sm-openai-explicit", response.text + request: Final = _only_request(wire) + assert request.target == f"/endpoints/{endpoint('openai')}/invocations", request + assert request.headers[COMPONENT_HEADER] == COMPONENT, dict(request.headers) + assert _body(request) == {"input": [text], "model": SERVED_MODEL}, request.body + assert _spend_row(response)["status"] == "success" + + +def test_plain_route_signs_model_id_as_the_header_and_keeps_the_hf_body(proxy: Gateway, wire: Wire) -> None: + text: Final = _text() + response: Final = _embed(proxy, "sm-plain", [text]) + assert response.status_code == 200, response.text + payload: Final = JSON_OBJECT.validate_json(response.content) + assert payload["data"] == [{"object": "embedding", "index": 0, "embedding": [0.25, -0.5, 1.0]}], response.text + assert payload["model"] == "sm-plain", response.text + assert _usage(payload) == (3, 0, 3), response.text + request: Final = _only_request(wire) + assert request.target == f"/endpoints/{endpoint('hf')}/invocations", request + assert request.headers[COMPONENT_HEADER] == COMPONENT, dict(request.headers) + assert COMPONENT_HEADER in _signed_headers(request), request.headers["authorization"] + assert _body(request) == {"inputs": [text]}, request.body + assert _spend_row(response)["status"] == "success" + + +def test_plain_route_without_model_id_sends_no_inference_component_header(proxy: Gateway, wire: Wire) -> None: + text: Final = _text() + response: Final = _embed(proxy, "sm-plain-control", [text]) + assert response.status_code == 200, response.text + request: Final = _only_request(wire) + assert COMPONENT_HEADER not in request.headers, dict(request.headers) + assert COMPONENT_HEADER not in _signed_headers(request), request.headers["authorization"] + assert _body(request) == {"inputs": [text]}, request.body + assert _spend_row(response)["status"] == "success" + + +def test_plain_route_reads_the_wrapped_hf_embedding_key(proxy: Gateway, wire: Wire) -> None: + texts: Final = [_text(), _text()] + response: Final = _embed(proxy, "sm-plain-dict", texts) + assert response.status_code == 200, response.text + payload: Final = JSON_OBJECT.validate_json(response.content) + assert payload["data"] == [ + {"object": "embedding", "index": 0, "embedding": [0.25, -0.5, 1.0]}, + {"object": "embedding", "index": 1, "embedding": [0.25, -0.5, 2.0]}, + ], response.text + request: Final = _only_request(wire) + assert request.headers[COMPONENT_HEADER] == COMPONENT, dict(request.headers) + assert _body(request) == {"inputs": texts}, request.body + + +def test_voyage_endpoint_keeps_the_voyage_body_and_signs_model_id_as_the_header(proxy: Gateway, wire: Wire) -> None: + text: Final = _text() + response: Final = _embed(proxy, "sm-voyage", [text]) + assert response.status_code == 200, response.text + payload: Final = JSON_OBJECT.validate_json(response.content) + assert payload["data"] == [{"object": "embedding", "index": 0, "embedding": [0.25, -0.5, 1.0]}], response.text + assert payload["model"] == "sm-voyage", response.text + assert _usage(payload) == (9, 0, 9), response.text + request: Final = _only_request(wire) + assert request.target == f"/endpoints/{endpoint('voyage')}/invocations", request + assert request.headers[COMPONENT_HEADER] == COMPONENT, dict(request.headers) + assert _body(request) == {"input": [text], "model": endpoint("voyage")}, request.body + assert _spend_row(response)["status"] == "success" + + +def test_cohere_endpoint_keeps_the_cohere_body_and_signs_model_id_as_the_header(proxy: Gateway, wire: Wire) -> None: + text: Final = _text() + response: Final = _embed(proxy, "sm-cohere", [text]) + assert response.status_code == 200, response.text + payload: Final = JSON_OBJECT.validate_json(response.content) + assert payload["data"] == [{"object": "embedding", "index": 0, "embedding": [0.25, -0.5, 1.0]}], response.text + assert _usage(payload) == (6, 0, 6), response.text + request: Final = _only_request(wire) + assert request.target == f"/endpoints/{endpoint('cohere')}/invocations", request + assert request.headers[COMPONENT_HEADER] == COMPONENT, dict(request.headers) + assert _body(request) == {"texts": [text], "input_type": "search_document"}, request.body + assert _spend_row(response)["status"] == "success" + + +def test_openai_route_forwards_dimensions_encoding_format_and_user_in_the_body(proxy: Gateway, wire: Wire) -> None: + text: Final = _text() + response: Final = _embed( + proxy, "sm-openai", [text], dimensions=3, encoding_format="float", user="integration-embedding-user" + ) + assert response.status_code == 200, response.text + payload: Final = JSON_OBJECT.validate_json(response.content) + assert payload["data"] == [{"object": "embedding", "index": 0, "embedding": [0.25, -0.5, 1.0]}], response.text + request: Final = _only_request(wire) + assert request.headers[COMPONENT_HEADER] == COMPONENT, dict(request.headers) + assert _body(request) == { + "input": [text], + "dimensions": 3, + "encoding_format": "float", + "user": "integration-embedding-user", + "model": SERVED_MODEL, + }, request.body + + +def test_openai_route_without_model_id_or_extra_body_sends_the_bare_openai_body(proxy: Gateway, wire: Wire) -> None: + text: Final = _text() + response: Final = _embed(proxy, "sm-openai-bare", [text]) + assert response.status_code == 200, response.text + request: Final = _only_request(wire) + assert COMPONENT_HEADER not in request.headers, dict(request.headers) + assert _body(request) == {"input": [text]}, request.body + + +def test_openai_route_sends_a_custom_deployment_key_at_the_top_level_of_the_body(proxy: Gateway, wire: Wire) -> None: + text: Final = _text() + response: Final = _embed(proxy, "sm-openai-custom-key", [text]) + assert response.status_code == 200, response.text + request: Final = _only_request(wire) + assert request.headers[COMPONENT_HEADER] == COMPONENT, dict(request.headers) + assert _body(request) == {"input": [text], "input_type": "query", "model": SERVED_MODEL}, request.body + + +def test_deployment_extra_body_wins_over_the_callers_encoding_format(proxy: Gateway, wire: Wire) -> None: + text: Final = _text() + created: Final = _sdk(proxy).embeddings.create( + model="sm-openai-pinned", input=[text], extra_body={"cache": _NO_CACHE} + ) + assert [row.embedding for row in created.data] == [[0.25, -0.5, 1.0]], created + request: Final = _only_request(wire) + assert _body(request) == {"input": [text], "encoding_format": "float", "model": SERVED_MODEL}, request.body + + +def test_openai_route_returns_the_base64_string_a_raw_caller_asked_for_unchanged(proxy: Gateway, wire: Wire) -> None: + text: Final = _text() + response: Final = _embed(proxy, "sm-openai", [text], encoding_format="base64") + assert response.status_code == 200, response.text + payload: Final = JSON_OBJECT.validate_json(response.content) + assert payload["data"] == [{"object": "embedding", "index": 0, "embedding": embedding_value(0, "base64")}], ( + response.text + ) + request: Final = _only_request(wire) + assert _body(request) == {"input": [text], "encoding_format": "base64", "model": SERVED_MODEL}, request.body + + +def test_openai_route_reports_zero_usage_when_the_container_sends_none(proxy: Gateway, wire: Wire) -> None: + response: Final = _embed(proxy, "sm-openai-nousage", [_text()]) + assert response.status_code == 200, response.text + payload: Final = JSON_OBJECT.validate_json(response.content) + assert _usage(payload) == (0, 0, 0), response.text + assert _only_request(wire).headers[COMPONENT_HEADER] == COMPONENT + assert _spend_row(response)["status"] == "success" + + +def test_openai_route_names_the_endpoint_when_the_container_sends_no_model(proxy: Gateway, wire: Wire) -> None: + response: Final = _embed(proxy, "sm-openai-nomodel", [_text()]) + assert response.status_code == 200, response.text + payload: Final = JSON_OBJECT.validate_json(response.content) + assert payload["model"] == "sm-openai-nomodel", response.text + assert _usage(payload) == (7, 0, 7), response.text + assert _only_request(wire).headers[COMPONENT_HEADER] == COMPONENT + + +def test_openai_route_numbers_rows_the_container_left_unindexed(proxy: Gateway, wire: Wire) -> None: + response: Final = _embed(proxy, "sm-openai-noindex", [_text(), _text()]) + assert response.status_code == 200, response.text + payload: Final = JSON_OBJECT.validate_json(response.content) + assert payload["data"] == [ + {"object": "embedding", "index": 0, "embedding": [0.25, -0.5, 1.0]}, + {"object": "embedding", "index": 1, "embedding": [0.25, -0.5, 2.0]}, + ], response.text + assert _only_request(wire).headers[COMPONENT_HEADER] == COMPONENT + + +def test_openai_route_takes_total_tokens_from_prompt_tokens_when_the_container_omits_it( + proxy: Gateway, wire: Wire +) -> None: + response: Final = _embed(proxy, "sm-openai-promptonly", [_text()]) + assert response.status_code == 200, response.text + payload: Final = JSON_OBJECT.validate_json(response.content) + assert _usage(payload) == (5, 0, 5), response.text + assert _only_request(wire).headers[COMPONENT_HEADER] == COMPONENT + + +@pytest.mark.parametrize("model", ["sm-openai-datadict", "sm-openai-hfshape"]) +def test_openai_route_rejects_a_container_body_that_is_not_an_openai_embeddings_object( + proxy: Gateway, wire: Wire, model: str +) -> None: + response: Final = _embed(proxy, model, [_text()]) + assert response.status_code == 503, response.text + assert EXPECTED_SHAPE in response.text, response.text + assert _only_request(wire).headers[COMPONENT_HEADER] == COMPONENT + assert _spend_row(response)["status"] == "failure" + + +def test_openai_route_reports_a_non_json_container_body_as_a_server_error(proxy: Gateway, wire: Wire) -> None: + response: Final = _embed(proxy, "sm-openai-notjson", [_text()]) + with pytest.raises(json.JSONDecodeError) as decode_error: + json.loads(NOT_JSON_BODY) + assert response.status_code == 500, response.text + assert str(decode_error.value) in response.text, response.text + assert _only_request(wire).headers[COMPONENT_HEADER] == COMPONENT + assert _spend_row(response)["status"] == "failure" + + +def test_container_validation_error_reaches_the_caller_as_400(proxy: Gateway, wire: Wire) -> None: + response: Final = _embed(proxy, "sm-openai-err400", [_text()]) + assert response.status_code == 400, response.text + assert "1 validation error detected" in response.text, response.text + assert _only_request(wire).headers[COMPONENT_HEADER] == COMPONENT + assert _spend_row(response)["status"] == "failure" + + +def test_container_model_error_reaches_the_caller_with_the_original_message(proxy: Gateway, wire: Wire) -> None: + response: Final = _embed(proxy, "sm-openai-err424", [_text()]) + assert response.status_code == 424, response.text + assert "unexpected field inputs" in response.text, response.text + assert _only_request(wire).headers[COMPONENT_HEADER] == COMPONENT + assert _spend_row(response)["status"] == "failure" + + +def test_plain_route_body_rejected_by_an_openai_container_reaches_the_caller(proxy: Gateway, wire: Wire) -> None: + text: Final = _text() + response: Final = _embed(proxy, "sm-plain-err424", [text]) + assert response.status_code == 424, response.text + assert "unexpected field inputs" in response.text, response.text + request: Final = _only_request(wire) + assert request.headers[COMPONENT_HEADER] == COMPONENT, dict(request.headers) + assert _body(request) == {"inputs": [text]}, request.body + assert _spend_row(response)["status"] == "failure" + + +@pytest.mark.parametrize( + ("model", "status", "message"), + [ + ("sm-openai-err429", 429, "Model is not ready for inference yet"), + ("sm-openai-err503", 503, "The endpoint is scaling, try again later"), + ("sm-openai-err500", 503, "An internal failure occurred"), + ], +) +def test_container_outage_errors_reach_the_caller_and_every_attempt_carries_the_header( + proxy: Gateway, wire: Wire, model: str, status: int, message: str +) -> None: + text: Final = _text() + response: Final = _embed(proxy, model, [text]) + assert response.status_code == status, response.text + assert message in response.text, response.text + attempts: Final = wire.drain() + assert len(attempts) >= 1, attempts + assert {attempt.headers[COMPONENT_HEADER] for attempt in attempts} == {COMPONENT}, attempts + assert {attempt.body for attempt in attempts} == {json.dumps({"input": [text], "model": SERVED_MODEL}).encode()} + assert _spend_row(response)["status"] == "failure" + + +def test_dropped_connection_reaches_the_caller_as_a_server_error(proxy: Gateway, wire: Wire) -> None: + response: Final = _embed(proxy, "sm-openai-drop", [_text()]) + assert response.status_code == 503, response.text + assert ConnectionClosedError.fmt.partition(":")[0] in response.text, response.text + attempts: Final = wire.drain() + assert len(attempts) >= 1, attempts + assert {attempt.headers[COMPONENT_HEADER] for attempt in attempts} == {COMPONENT}, attempts + assert _spend_row(response)["status"] == "failure" + + +def test_inference_component_endpoint_answers_once_the_header_unlocks_it(proxy: Gateway, wire: Wire) -> None: + text: Final = _text() + response: Final = _embed(proxy, "sm-openai-ic", [text]) + assert response.status_code == 200, response.text + payload: Final = JSON_OBJECT.validate_json(response.content) + assert payload["data"] == [{"object": "embedding", "index": 0, "embedding": [0.25, -0.5, 1.0]}], response.text + request: Final = _only_request(wire) + assert request.headers[COMPONENT_HEADER] == COMPONENT, dict(request.headers) + assert _body(request) == {"input": [text], "model": SERVED_MODEL}, request.body + assert _spend_row(response)["status"] == "success" + + +def test_inference_component_endpoint_without_model_id_tells_the_caller_how_to_pass_it( + proxy: Gateway, wire: Wire +) -> None: + text: Final = _text() + response: Final = _embed(proxy, "sm-openai-ic-missing", [text]) + assert response.status_code == 400, response.text + assert IC_REQUIRED in response.text, response.text + assert PASS_IN_HINT in response.text, response.text + request: Final = _only_request(wire) + assert COMPONENT_HEADER not in request.headers, dict(request.headers) + assert _body(request) == {"input": [text], "model": SERVED_MODEL}, request.body + assert _spend_row(response)["status"] == "failure" + + +def test_endpoint_without_inference_components_tells_the_deployment_to_drop_model_id( + proxy: Gateway, wire: Wire +) -> None: + text: Final = _text() + response: Final = _embed(proxy, "sm-plain-noic", [text]) + assert response.status_code == 400, response.text + assert IC_NOT_ALLOWED in response.text, response.text + assert REMOVE_HINT in response.text, response.text + request: Final = _only_request(wire) + assert request.headers[COMPONENT_HEADER] == COMPONENT, dict(request.headers) + assert _body(request) == {"inputs": [text]}, request.body + assert _spend_row(response)["status"] == "failure" + + +def test_endpoint_without_inference_components_serves_a_deployment_without_model_id(proxy: Gateway, wire: Wire) -> None: + text: Final = _text() + response: Final = _embed(proxy, "sm-plain-noic-control", [text]) + assert response.status_code == 200, response.text + request: Final = _only_request(wire) + assert COMPONENT_HEADER not in request.headers, dict(request.headers) + assert _body(request) == {"inputs": [text]}, request.body + assert _spend_row(response)["status"] == "success" + + +def test_unknown_endpoint_reaches_the_caller_as_400(proxy: Gateway, wire: Wire) -> None: + response: Final = _embed(proxy, "sm-openai-unknown", [_text()]) + assert response.status_code == 400, response.text + assert "Endpoint unknown-embed-endpoint of account 000000000000 not found." in response.text, response.text + request: Final = _only_request(wire) + assert request.target == "/endpoints/unknown-embed-endpoint/invocations", request + assert request.headers[COMPONENT_HEADER] == COMPONENT, dict(request.headers) + assert _spend_row(response)["status"] == "failure" + + +def test_plain_route_rejects_an_unsupported_openai_param_without_drop_params(proxy: Gateway, wire: Wire) -> None: + response: Final = _embed(proxy, "sm-plain", [_text()], dimensions=3) + assert response.status_code == 400, response.text + assert "does not support parameters: {'dimensions': 3}" in response.text, response.text + assert "drop_params" in response.text, response.text + assert wire.drain() == (), response.text + assert _spend_row(response)["status"] == "failure" + + +def test_plain_route_drops_an_unsupported_openai_param_under_drop_params(proxy: Gateway, wire: Wire) -> None: + text: Final = _text() + response: Final = _embed(proxy, "sm-plain-drop-params", [text], dimensions=3) + assert response.status_code == 200, response.text + request: Final = _only_request(wire) + assert request.headers[COMPONENT_HEADER] == COMPONENT, dict(request.headers) + assert _body(request) == {"inputs": [text]}, request.body + + +@pytest.mark.parametrize("model", ["sm-openai-model-id-int", "sm-openai-model-id-list"]) +def test_a_non_string_model_id_fails_before_any_request_leaves_the_proxy( + proxy: Gateway, wire: Wire, model: str +) -> None: + response: Final = _embed(proxy, model, [_text()]) + assert response.status_code == 400, response.text + assert NON_STRING_MODEL_ID in response.text, response.text + assert wire.drain() == (), response.text + assert _spend_row(response)["status"] == "failure" + + +@pytest.mark.parametrize("model", ["sm-openai-model-id-empty", "sm-openai-model-id-null"]) +def test_an_empty_or_null_model_id_sends_no_inference_component_header(proxy: Gateway, wire: Wire, model: str) -> None: + text: Final = _text() + response: Final = _embed(proxy, model, [text]) + assert response.status_code == 200, response.text + request: Final = _only_request(wire) + assert COMPONENT_HEADER not in request.headers, dict(request.headers) + assert _body(request) == {"input": [text], "model": SERVED_MODEL}, request.body + + +def test_a_long_model_id_is_signed_and_forwarded_whole(proxy: Gateway, wire: Wire) -> None: + text: Final = _text() + response: Final = _embed(proxy, "sm-openai-model-id-long", [text]) + assert response.status_code == 200, response.text + request: Final = _only_request(wire) + assert request.headers[COMPONENT_HEADER] == "c" * 5120, len(request.headers[COMPONENT_HEADER]) + assert COMPONENT_HEADER in _signed_headers(request), request.headers["authorization"] + assert _body(request) == {"input": [text], "model": SERVED_MODEL}, request.body + + +@pytest.mark.parametrize("model", ["sm-openai-extra-body-string", "sm-openai-extra-body-empty"]) +def test_an_extra_body_that_is_not_a_populated_object_adds_nothing_to_the_body( + proxy: Gateway, wire: Wire, model: str +) -> None: + text: Final = _text() + response: Final = _embed(proxy, model, [text]) + assert response.status_code == 200, response.text + request: Final = _only_request(wire) + assert request.headers[COMPONENT_HEADER] == COMPONENT, dict(request.headers) + assert _body(request) == {"input": [text]}, request.body + + +def test_a_caller_sent_model_id_overrides_the_deployments_inference_component(proxy: Gateway, wire: Wire) -> None: + text: Final = _text() + response: Final = _embed(proxy, "sm-openai", [text], model_id="caller-chosen-component") + assert response.status_code == 200, response.text + request: Final = _only_request(wire) + assert request.headers[COMPONENT_HEADER] == "caller-chosen-component", dict(request.headers) + assert _body(request) == {"input": [text], "model": SERVED_MODEL}, request.body + + +def test_a_caller_sent_model_id_is_forwarded_when_the_deployment_has_none(proxy: Gateway, wire: Wire) -> None: + text: Final = _text() + response: Final = _embed(proxy, "sm-openai-bare", [text], model_id="caller-chosen-component") + assert response.status_code == 200, response.text + request: Final = _only_request(wire) + assert request.headers[COMPONENT_HEADER] == "caller-chosen-component", dict(request.headers) + assert _body(request) == {"input": [text]}, request.body + + +def test_a_caller_sent_non_string_model_id_fails_before_any_request_leaves_the_proxy( + proxy: Gateway, wire: Wire +) -> None: + response: Final = _embed(proxy, "sm-openai-bare", [_text()], model_id=["caller-chosen-component"]) + assert response.status_code == 400, response.text + assert NON_STRING_MODEL_ID in response.text, response.text + assert wire.drain() == (), response.text + + +def test_a_caller_sent_extra_body_replaces_the_deployments_extra_body(proxy: Gateway, wire: Wire) -> None: + text: Final = _text() + response: Final = _embed(proxy, "sm-openai", [text], extra_body={"truncate": "END"}) + assert response.status_code == 200, response.text + request: Final = _only_request(wire) + assert request.headers[COMPONENT_HEADER] == COMPONENT, dict(request.headers) + assert _body(request) == {"input": [text], "truncate": "END"}, request.body + + +def test_a_failed_call_does_not_poison_the_next_one(proxy: Gateway, wire: Wire) -> None: + failed: Final = _embed(proxy, "sm-openai-err503", [_text()]) + assert failed.status_code == 503, failed.text + wire.drain() + text: Final = _text() + response: Final = _embed(proxy, "sm-openai", [text]) + assert response.status_code == 200, response.text + request: Final = _only_request(wire) + assert request.headers[COMPONENT_HEADER] == COMPONENT, dict(request.headers) + assert _body(request) == {"input": [text], "model": SERVED_MODEL}, request.body + assert _spend_row(response)["status"] == "success" + + +def test_two_identical_uncached_requests_each_reach_the_container_and_log_once(proxy: Gateway, wire: Wire) -> None: + text: Final = _text() + first: Final = _embed(proxy, "sm-openai", [text]) + second: Final = _embed(proxy, "sm-openai", [text]) + assert (first.status_code, second.status_code) == (200, 200), (first.text, second.text) + assert first.headers["x-litellm-call-id"] != second.headers["x-litellm-call-id"], (first.headers, second.headers) + requests: Final = wire.drain() + assert [_body(request) for request in requests] == [{"input": [text], "model": SERVED_MODEL}] * 2, requests + assert {request.headers[COMPONENT_HEADER] for request in requests} == {COMPONENT}, requests + assert (_spend_row(first)["status"], _spend_row(second)["status"]) == ("success", "success") + + +def test_health_check_reaches_the_container_with_the_inference_component_header( + proxy: Gateway, wire: Wire, health_endpoint: str +) -> None: + report: Final = proxy.get("/health", {"model": "sm-openai-health"}) + assert (report["healthy_count"], report["unhealthy_count"]) == (1, 0), report + request: Final = _only_request(wire) + assert request.target == f"/endpoints/{health_endpoint}/invocations", request + assert request.headers[COMPONENT_HEADER] == COMPONENT, dict(request.headers) + body: Final = _body(request) + assert body["model"] == SERVED_MODEL, request.body + assert isinstance(body["input"], list) and len(body["input"]) == 1, request.body diff --git a/tests/integration/providers/test_strands_decider_agentcore_wire.py b/tests/integration/providers/test_strands_decider_agentcore_wire.py index c7213a66bc5..9b7a9dbbfd0 100644 --- a/tests/integration/providers/test_strands_decider_agentcore_wire.py +++ b/tests/integration/providers/test_strands_decider_agentcore_wire.py @@ -3,7 +3,6 @@ import json import os import re import threading -import time import uuid from collections.abc import Callable, Iterator, Mapping, Sequence from concurrent.futures import ThreadPoolExecutor @@ -21,6 +20,7 @@ from integration._support.database import read_rows from integration._support.forward_proxy import ForwardProxy, tunnelling_forward_proxy from integration._support.process import owned_proxy from integration._support.sigv4 import encoded_path, signature +from integration._support.timing import Timed, assert_within_factor_of_control, timed from integration._support.tls import server_context, write_self_signed_cert from integration._support.wire import Reply, Request, Wire, wire_server from pydantic import JsonValue, TypeAdapter @@ -167,6 +167,8 @@ _load_yaml: Final[Callable[[str], object]] = yaml.safe_load _JSON: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue) _OWNED_PROXY_CELL_SECONDS: Final = 2 * max(30.0, float(os.environ.get("INTEGRATION_PROXY_READY_SECONDS") or 70)) + 120 pytestmark: Final = pytest.mark.timeout(_OWNED_PROXY_CELL_SECONDS) +_LINEAR_FACTOR: Final = 20 +_CLASSIFIER_SAMPLES: Final = 5 def _host(region: str) -> str: @@ -795,6 +797,10 @@ def test_a_plain_api_base_forwards_amazon_shaped_extra_headers_unchanged(gateway assert json.loads(call.body) == _UPSTREAM_BODY, call.body +def _fastest(operation: Callable[[], str]) -> Timed[str]: + return min((timed(operation) for _ in range(_CLASSIFIER_SAMPLES)), key=lambda sample: sample.seconds) + + def _worst_case_api_base() -> str: return "arn:aws" + "-a" * 131072 + "!" @@ -803,31 +809,28 @@ def test_a_worst_case_runtime_arn_pattern_input_is_classified_in_linear_time( gateway: Gateway, record_property: Callable[[str, object], None] ) -> None: worst: Final = _worst_case_api_base() - started: Final = time.perf_counter() - url: Final = StrandsDeciderDecisionsConfig().get_complete_url(worst, _BODY_MODEL) - classified: Final = time.perf_counter() - started - record_property("classifier_seconds", f"{classified:.6f}") - assert url == f"{worst}/v1/systemone" - assert classified < 0.5, classified + matching: Final = f"{worst.removesuffix('!')}:bedrock-agentcore:{_US_EAST}:123456789012:runtime/decider" + config: Final = StrandsDeciderDecisionsConfig() + control: Final = _fastest(lambda: config.get_complete_url(matching, _BODY_MODEL)) + classified: Final = _fastest(lambda: config.get_complete_url(worst, _BODY_MODEL)) + record_property("classifier_seconds", f"{classified.seconds:.6f}") + assert control.value.endswith("/invocations"), control.value[-200:] + assert classified.value == f"{worst}/v1/systemone" + assert_within_factor_of_control(classified.seconds, control=control.seconds, factor=_LINEAR_FACTOR) with wire_server(_answering) as wire, gateway.scenario() as scenario: model: Final = _plain_model(scenario, wire, configurable_clientside_auth_params=["api_base"]) + answered: Final = timed(lambda: _decide(gateway, model)) + _assert_answered(answered.value, model) + assert len(wire.drain()) == 1 with ThreadPoolExecutor(max_workers=2) as pool: - timed: Final = pool.submit(_timed_decide, gateway, model, worst) - liveliness_started: Final = time.perf_counter() - liveliness: Final = gateway.request("GET", "/health/liveliness") - liveliness_seconds: Final = time.perf_counter() - liveliness_started - response, decide_seconds = timed.result() - record_property("decide_seconds", f"{decide_seconds:.6f}") - record_property("liveliness_seconds", f"{liveliness_seconds:.6f}") - assert liveliness.status_code == 200, liveliness.text - assert 500 <= response.status_code < 600, response.text - assert "litellm.APIConnectionError" in response.text, response.text[:2000] - assert decide_seconds < 5, decide_seconds - assert liveliness_seconds < 2, liveliness_seconds + pending: Final = pool.submit(timed, lambda: _decide(gateway, model, api_base=worst)) + liveliness: Final = timed(lambda: gateway.request("GET", "/health/liveliness")) + decided: Final = pending.result() + record_property("decide_seconds", f"{decided.seconds:.6f}") + record_property("liveliness_seconds", f"{liveliness.seconds:.6f}") + assert liveliness.value.status_code == 200, liveliness.value.text + assert 500 <= decided.value.status_code < 600, decided.value.text + assert "litellm.APIConnectionError" in decided.value.text, decided.value.text[:2000] + assert_within_factor_of_control(decided.seconds, control=answered.seconds, factor=_LINEAR_FACTOR) + assert_within_factor_of_control(liveliness.seconds, control=answered.seconds, factor=_LINEAR_FACTOR) assert wire.drain() == () - - -def _timed_decide(gateway: Gateway, model: str, api_base: str) -> tuple[httpx.Response, float]: - started: Final = time.perf_counter() - response: Final = _decide(gateway, model, api_base=api_base) - return response, time.perf_counter() - started diff --git a/tests/integration/providers/test_validated_reply_boundaries_under_faults.py b/tests/integration/providers/test_validated_reply_boundaries_under_faults.py index 0a358e685f6..ae658554668 100644 --- a/tests/integration/providers/test_validated_reply_boundaries_under_faults.py +++ b/tests/integration/providers/test_validated_reply_boundaries_under_faults.py @@ -1,7 +1,6 @@ import base64 import json import re -import signal import threading import uuid from collections import Counter @@ -449,11 +448,11 @@ def test_a_worker_killed_mid_burst_leaves_the_sibling_validating_provider_payloa eventually(held.qsize, lambda size: size == len(calls), seconds=_GATE_SECONDS) held_by: Final = {pid: _open_upstream_connections(pid, wire.url) for pid in workers} victim_pid, survivor_pid = sorted(workers, key=held_by.__getitem__) - victim: Final = psutil.Process(victim_pid) - victim.suspend() - victim.send_signal(signal.SIGKILL) - release.set() - outcomes: Final = tuple(future.result(timeout=60) for future in futures) + with owned.replacing_worker(victim_pid) as kill: + kill.freeze() + kill() + release.set() + outcomes: Final = tuple(future.result(timeout=60) for future in futures) assert sum(held_by.values()) == len(calls), held_by served: Final = tuple( (call, response) for call, response in zip(calls, outcomes, strict=True) if response is not None diff --git a/tests/integration/providers/test_vertex_anthropic_inline_tools_beta_wire.py b/tests/integration/providers/test_vertex_anthropic_inline_tools_beta_wire.py index 5061f0b34a4..a6b039dc8f7 100644 --- a/tests/integration/providers/test_vertex_anthropic_inline_tools_beta_wire.py +++ b/tests/integration/providers/test_vertex_anthropic_inline_tools_beta_wire.py @@ -1,7 +1,6 @@ import asyncio import json import re -import signal import threading import uuid from collections.abc import Callable, Mapping @@ -564,11 +563,12 @@ def test_worker_sigkill_mid_burst_leaves_the_replacement_forwarding_inline_tools shares: Final = {pid: _accepted_client_ports(pid, proxy_url.port or 0) & frozenset(ports) for pid in workers} assert sum(map(len, shares.values())) == len(clients), shares victim: Final = min((pid for pid in workers if shares[pid]), key=lambda pid: len(shares[pid])) - psutil.Process(victim).send_signal(signal.SIGKILL) - release.set() - results: Final = tuple(future.result(timeout=60) for future in futures) - for port, result in zip(ports, results, strict=True): - assert result == (None if port in shares[victim] else 200), (port, result, shares) + with owned.replacing_worker(victim) as kill: + kill() + release.set() + results: Final = tuple(future.result(timeout=60) for future in futures) + for port, result in zip(ports, results, strict=True): + assert result == (None if port in shares[victim] else 200), (port, result, shares) eventually(lambda: len(_STARTED_WORKER.findall(owned.log.read_text())), lambda started: started >= 3, 120) second_wave: Final = tuple(_marker("after kill") for _ in range(6)) for marker in second_wave: diff --git a/tests/integration/providers/test_vertex_gemini_fragmented_stream_wire.py b/tests/integration/providers/test_vertex_gemini_fragmented_stream_wire.py index 449c5c9c105..95a5f881ba1 100644 --- a/tests/integration/providers/test_vertex_gemini_fragmented_stream_wire.py +++ b/tests/integration/providers/test_vertex_gemini_fragmented_stream_wire.py @@ -1,12 +1,14 @@ import json -import time +from collections.abc import Callable +from functools import partial from typing import Final import pytest from cryptography.hazmat.primitives import serialization from cryptography.hazmat.primitives.asymmetric import rsa -from integration._support.client import Gateway -from integration._support.wire import Reply, Request, wire_server +from integration._support.client import Gateway, Scenario +from integration._support.timing import assert_within_factor_of_control, timed +from integration._support.wire import Reply, Request, Wire, wire_server from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter _BACKEND: Final = "gemini-3.7-flash" @@ -17,6 +19,7 @@ _PROMPT: Final = "Write a very long numbered list." _PART_COUNT: Final = 8000 _LINES_PER_FRAGMENT: Final = 64 _STREAM_BUDGET_SECONDS: Final = 10.0 +_FRAGMENTED_FACTOR: Final = 10 _JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) @@ -64,8 +67,8 @@ def _expected_text() -> str: return "".join(f"{index}. item\n" for index in range(_PART_COUNT)) -def _gemini_response_fragments() -> tuple[bytes, ...]: - document: Final = json.dumps( +def _gemini_document() -> str: + return json.dumps( { "candidates": [ { @@ -81,18 +84,18 @@ def _gemini_response_fragments() -> tuple[bytes, ...]: }, indent=2, ) + + +def _gemini_response_fragments(document: str, lines_per_fragment: int) -> tuple[bytes, ...]: lines: Final = document.split("\n") fragments: Final = tuple( - "\n".join(lines[start : start + _LINES_PER_FRAGMENT]).encode() + b"\n" - for start in range(0, len(lines), _LINES_PER_FRAGMENT) + "\n".join(lines[start : start + lines_per_fragment]).encode() + b"\n" + for start in range(0, len(lines), lines_per_fragment) ) return (b"data: " + fragments[0], *fragments[1:], b"\n") -@pytest.mark.covers("providers.vertex_gemini.fragmented_stream_json_is_parsed_once_and_stays_live") -def test_vertex_gemini_stream_split_across_many_fragments_completes_without_stalling(gateway: Gateway) -> None: - fragments: Final = _gemini_response_fragments() - +def _scripted_vertex(fragments: tuple[bytes, ...]) -> Callable[[Request], Reply]: def respond(request: Request) -> Reply: assert request.method == "POST" assert request.target == f"{_MODEL_PATH}:streamGenerateContent?alt=sse" @@ -102,37 +105,60 @@ def test_vertex_gemini_stream_split_across_many_fragments_completes_without_stal assert body["generationConfig"] == {"temperature": 0.0} return Reply(content_type="text/event-stream", chunks=fragments) - with wire_server(respond) as wire, gateway.scenario() as scenario: - model: Final = scenario.model( - model=f"vertex_ai/{_BACKEND}", - api_base=f"{wire.url}{_MODEL_PATH}", - api_key=None, - vertex_project=_PROJECT, - vertex_location=_LOCATION, - vertex_credentials=_service_account_json(gateway.upstream_url.rstrip("/")), - ) - started: Final = time.monotonic() - with gateway.client.stream( - "POST", - "/v1/chat/completions", - json={ - "model": model, - "messages": [{"role": "user", "content": _PROMPT}], - "stream": True, - "temperature": 0.0, - }, - headers={"Authorization": f"Bearer {gateway.key}"}, - timeout=_STREAM_BUDGET_SECONDS, - ) as response: - assert response.status_code == 200, response.read() - lines: Final = tuple(line for line in response.iter_lines() if line.startswith("data: ")) - elapsed: Final = time.monotonic() - started - assert elapsed < _STREAM_BUDGET_SECONDS, f"stream took {elapsed:.1f}s for {len(fragments)} fragments" - assert lines[-1] == "data: [DONE]", lines[-3:] - chunks: Final = tuple(_Chunk.model_validate_json(line.removeprefix("data: ")) for line in lines[:-1]) - choices: Final = tuple(choice for chunk in chunks for choice in chunk.choices) - assert "".join(choice.delta.content or "" for choice in choices) == _expected_text() - assert tuple(choice.finish_reason for choice in choices if choice.finish_reason) == ("stop",) + return respond + + +def _vertex_model(gateway: Gateway, scenario: Scenario, wire: Wire) -> str: + return scenario.model( + model=f"vertex_ai/{_BACKEND}", + api_base=f"{wire.url}{_MODEL_PATH}", + api_key=None, + vertex_project=_PROJECT, + vertex_location=_LOCATION, + vertex_credentials=_service_account_json(gateway.upstream_url.rstrip("/")), + ) + + +def _streamed_lines(gateway: Gateway, model: str) -> tuple[str, ...]: + with gateway.client.stream( + "POST", + "/v1/chat/completions", + json={ + "model": model, + "messages": [{"role": "user", "content": _PROMPT}], + "stream": True, + "temperature": 0.0, + }, + headers={"Authorization": f"Bearer {gateway.key}"}, + timeout=_STREAM_BUDGET_SECONDS, + ) as response: + assert response.status_code == 200, response.read() + return tuple(line for line in response.iter_lines() if line.startswith("data: ")) + + +def _assert_streamed_the_whole_list(lines: tuple[str, ...]) -> None: + assert lines[-1] == "data: [DONE]", lines[-3:] + chunks: Final = tuple(_Chunk.model_validate_json(line.removeprefix("data: ")) for line in lines[:-1]) + choices: Final = tuple(choice for chunk in chunks for choice in chunk.choices) + assert "".join(choice.delta.content or "" for choice in choices) == _expected_text() + assert tuple(choice.finish_reason for choice in choices if choice.finish_reason) == ("stop",) + + +@pytest.mark.covers("providers.vertex_gemini.fragmented_stream_json_is_parsed_once_and_stays_live") +def test_vertex_gemini_stream_split_across_many_fragments_completes_without_stalling(gateway: Gateway) -> None: + document: Final = _gemini_document() + whole: Final = _gemini_response_fragments(document, document.count("\n") + 1) + fragments: Final = _gemini_response_fragments(document, _LINES_PER_FRAGMENT) + with ( + wire_server(_scripted_vertex(whole)) as control_wire, + wire_server(_scripted_vertex(fragments)) as wire, + gateway.scenario() as scenario, + ): + control: Final = timed(partial(_streamed_lines, gateway, _vertex_model(gateway, scenario, control_wire))) + _assert_streamed_the_whole_list(control.value) + streamed: Final = timed(partial(_streamed_lines, gateway, _vertex_model(gateway, scenario, wire))) + assert_within_factor_of_control(streamed.seconds, control=control.seconds, factor=_FRAGMENTED_FACTOR) + _assert_streamed_the_whole_list(streamed.value) assert [(request.method, request.target) for request in wire.drain()] == [ ("POST", f"{_MODEL_PATH}:streamGenerateContent?alt=sse") ] diff --git a/tests/integration/providers/test_vertex_partner_count_tokens_wire.py b/tests/integration/providers/test_vertex_partner_count_tokens_wire.py index 5df8dd6e32e..558ee3a004e 100644 --- a/tests/integration/providers/test_vertex_partner_count_tokens_wire.py +++ b/tests/integration/providers/test_vertex_partner_count_tokens_wire.py @@ -1,6 +1,5 @@ import json import re -import signal import socket import threading import uuid @@ -663,15 +662,16 @@ def test_worker_sigkill_mid_burst_leaves_the_sibling_counting(gateway: Gateway, shares: Final = {pid: _accepted_client_ports(pid, proxy_url.port or 0) & frozenset(ports) for pid in workers} assert sum(map(len, shares.values())) == len(clients), shares victim: Final = min((pid for pid in workers if shares[pid]), key=lambda pid: len(shares[pid])) - psutil.Process(victim).send_signal(signal.SIGKILL) - release.set() - results: Final = tuple(future.result(timeout=60) for future in futures) - for port, result in zip(ports, results, strict=True): - assert result == (None if port in shares[victim] else (200, _PEER_COUNT)), (port, result, shares) - second_wave: Final = _clients(stack, str(proxy_url), 6) - assert tuple(_counted_on(client, gateway.key, _full(_OWNED_MODEL)) for client in second_wave) == ( - (200, _PEER_COUNT), - ) * len(second_wave) + with owned.replacing_worker(victim) as kill: + kill() + release.set() + results: Final = tuple(future.result(timeout=60) for future in futures) + for port, result in zip(ports, results, strict=True): + assert result == (None if port in shares[victim] else (200, _PEER_COUNT)), (port, result, shares) + second_wave: Final = _clients(stack, str(proxy_url), 6) + assert tuple(_counted_on(client, gateway.key, _full(_OWNED_MODEL)) for client in second_wave) == ( + (200, _PEER_COUNT), + ) * len(second_wave) assert len(_counted_bodies(wire)) == len(clients) + len(second_wave) eventually(lambda: len(_STARTED_WORKER.findall(owned.log.read_text())), lambda started: started >= 3, 120) assert owned.process.poll() is None diff --git a/tests/integration/providers/test_voyage_mongodb_host_wire.py b/tests/integration/providers/test_voyage_mongodb_host_wire.py index 1a6ed907c9b..926bd5a7c69 100644 --- a/tests/integration/providers/test_voyage_mongodb_host_wire.py +++ b/tests/integration/providers/test_voyage_mongodb_host_wire.py @@ -2,7 +2,6 @@ import asyncio import json import os import re -import signal import socket import uuid from collections.abc import Callable, Iterator @@ -11,7 +10,6 @@ from pathlib import Path from typing import Final import httpx -import psutil import pytest import yaml from integration._support.client import JSON_OBJECT, Gateway, eventually, gateway_from_environment, object_value @@ -675,12 +673,12 @@ async def test_worker_sigkill_mid_traffic_leaves_the_sibling_routing_by_key(gate _burst(str(candidate.client.base_url), candidate.key, 20, tolerate_transport_errors=True) ) await asyncio.to_thread(eventually, lambda: forward.received.qsize(), lambda size: size >= 4, 60) - victim: Final = psutil.Process(workers[0]) - victim.suspend() - victim.send_signal(signal.SIGKILL) - served: Final = await burst - for probe in served: - assert probe.response.status_code == REFUSED_STATUS, probe.response.text + with owned.replacing_worker(workers[0]) as kill: + kill.freeze() + await asyncio.to_thread(kill) + served: Final = await burst + for probe in served: + assert probe.response.status_code == REFUSED_STATUS, probe.response.text assert set(forward.targets()) <= {MONGODB_HOST, VOYAGE_HOST} follow_up: Final = _rerank(candidate, MONGODB_RERANK) _assert_dialed(forward, follow_up, MONGODB_HOST) diff --git a/tests/integration/proxy_config.yaml b/tests/integration/proxy_config.yaml index a4928a31943..3136ecb5a09 100644 --- a/tests/integration/proxy_config.yaml +++ b/tests/integration/proxy_config.yaml @@ -349,6 +349,11 @@ model_list: model: openrouter/typesafe/jev-1.13 api_base: http://127.0.0.1:8191 api_key: synthetic-openrouter-key + - model_name: fireworks_ai/accounts/fireworks/routers/rune + litellm_params: + model: fireworks_ai/accounts/fireworks/routers/rune + api_base: http://127.0.0.1:8191 + api_key: synthetic-fireworks-key - model_name: openrouter/anthropic/claude-sonnet-4 litellm_params: model: openrouter/anthropic/claude-sonnet-4 diff --git a/tests/integration/routing/test_audit_any_sweep_chaos.py b/tests/integration/routing/test_audit_any_sweep_chaos.py index 1e46299e662..fe95d528fd5 100644 --- a/tests/integration/routing/test_audit_any_sweep_chaos.py +++ b/tests/integration/routing/test_audit_any_sweep_chaos.py @@ -1,8 +1,6 @@ from __future__ import annotations import json -import re -import signal from collections.abc import Mapping from concurrent.futures import ThreadPoolExecutor from pathlib import Path @@ -10,7 +8,6 @@ from types import MappingProxyType from typing import Final import httpx -import psutil import pytest from integration._support.client import Gateway, eventually, gateway_from_environment from integration._support.process import graceful_stop_seconds, owned_proxy_process @@ -139,13 +136,6 @@ def _fire( return kind, 200, response_id if isinstance(response_id, str) else None -_WORKER_PID: Final = re.compile(r"Started server process \[(\d+)\]") - - -def _worker_pids(log: Path) -> tuple[int, ...]: - return tuple(int(match) for match in _WORKER_PID.findall(log.read_text())) - - @pytest.mark.timeout(2 * graceful_stop_seconds() + 120) def test_proxy_survives_worker_kill_mid_burst(tmp_path: Path) -> None: with gateway_from_environment() as upstream_gateway: @@ -181,16 +171,17 @@ def test_proxy_survives_worker_kill_mid_burst(tmp_path: Path) -> None: "responses-stream": responses_stream_model, } ) - workers: Final = eventually(lambda: _worker_pids(owned.log), lambda pids: len(pids) == 2, seconds=30) + workers: Final = eventually(owned.worker_pids, lambda pids: len(pids) == 2, seconds=30) records: dict[int, tuple[str, int | None, str | None]] = {} with ThreadPoolExecutor(max_workers=30) as pool: futures: Final = [ (index, pool.submit(_fire, owned.gateway, models, key, _burst_kind(index), str(index))) for index in range(30) ] - psutil.Process(workers[0]).send_signal(signal.SIGKILL) - for index, future in futures: - records[index] = future.result() + with owned.replacing_worker(min(workers)) as kill: + kill() + for index, future in futures: + records[index] = future.result() proxy_stamped: Final = [ record[2] for record in records.values() diff --git a/tests/integration/routing/test_complexity_router_databricks_classifier.py b/tests/integration/routing/test_complexity_router_databricks_classifier.py index 419b786bfc1..af4b76e888f 100644 --- a/tests/integration/routing/test_complexity_router_databricks_classifier.py +++ b/tests/integration/routing/test_complexity_router_databricks_classifier.py @@ -2,7 +2,6 @@ import asyncio import json import os import re -import signal import threading import time import uuid @@ -1008,18 +1007,18 @@ async def test_worker_sigkill_mid_burst_leaves_the_sibling_classifying_through_t held_by: Final = MappingProxyType({pid: _open_upstream_connections(pid, judge.url) for pid in workers}) assert sum(held_by.values()) == 20, held_by victim_pid, survivor_pid = sorted(workers, key=held_by.__getitem__) - victim: Final = psutil.Process(victim_pid) - victim.suspend() - victim.send_signal(signal.SIGKILL) - release.set() - served: Final = await burst - assert len(served) == held_by[survivor_pid], (held_by, len(served)) - assert [item.cause for item in served] == ["jev_classifier"] * len(served), served - follow_up_prompt: Final = _prompt() - (answered,) = await _burst( - base_url, candidate.key, router, (("/v1/chat/completions", False),), (follow_up_prompt,) - ) - assert (answered.cause, answered.text) == ("jev_classifier", _COMPLEX_TEXT), answered + with owned.replacing_worker(victim_pid) as kill: + kill.freeze() + await asyncio.to_thread(kill) + release.set() + served: Final = await burst + assert len(served) == held_by[survivor_pid], (held_by, len(served)) + assert [item.cause for item in served] == ["jev_classifier"] * len(served), served + follow_up_prompt: Final = _prompt() + (answered,) = await _burst( + base_url, candidate.key, router, (("/v1/chat/completions", False),), (follow_up_prompt,) + ) + assert (answered.cause, answered.text) == ("jev_classifier", _COMPLEX_TEXT), answered await asyncio.to_thread( eventually, lambda: _worker_startups(owned.log), lambda found: len(found[0]) == 3 and found[1] == 3, 180 ) diff --git a/tests/integration/routing/test_fallback_hop_foreign_encrypted_reasoning.py b/tests/integration/routing/test_fallback_hop_foreign_encrypted_reasoning.py index bea36ee6bf8..32c7b99d5f1 100644 --- a/tests/integration/routing/test_fallback_hop_foreign_encrypted_reasoning.py +++ b/tests/integration/routing/test_fallback_hop_foreign_encrypted_reasoning.py @@ -19,8 +19,6 @@ serving the group with one deployment until the next resync. import asyncio import json -import re -import signal import time import uuid from collections.abc import AsyncIterator, Callable, Iterator, Mapping, Sequence @@ -34,7 +32,6 @@ from typing import Final import anthropic import httpx import openai -import psutil import pytest from integration._support.client import ( JSON_OBJECT, @@ -1249,13 +1246,6 @@ def test_order_one_outage_mid_burst_hops_stripped_and_recovers_after_restart(gat assert [json.loads(request.body)["input"] for request in restarted.drain()] == [history(recovered)] -WORKER_PID: Final = re.compile(r"Started server process \[(\d+)\]") - - -def worker_pids(log: Path) -> tuple[int, ...]: - return tuple(int(match) for match in WORKER_PID.findall(log.read_text())) - - def assert_completed_shots_hopped_stripped( gateway: Gateway, groups: Mapping[str, OrderedGroup], fired: Sequence[Fired] ) -> None: @@ -1283,17 +1273,18 @@ def test_worker_killed_mid_burst_leaves_the_survivor_hopping_stripped(tmp_path: owned.gateway.post("/config/update", {"router_settings": {"num_retries": 0}}) with owned.gateway.scenario() as scenario: groups: Final = burst_groups(scenario) - workers: Final = eventually(lambda: worker_pids(owned.log), lambda pids: len(pids) == 2, seconds=60) + workers: Final = eventually(owned.worker_pids, lambda pids: len(pids) == 2, seconds=60) shots: Final = tuple((BURST_KINDS[index % len(BURST_KINDS)], marker()) for index in range(30)) - with lanes(owned.gateway, [group.target for group in groups.values()], 30) as pinned: - with ThreadPoolExecutor(max_workers=30) as pool: - futures: Final = [ - pool.submit(fire, lane, groups, kind, mark) - for lane, (kind, mark) in zip(pinned, shots, strict=True) - ] - psutil.Process(workers[0]).send_signal(signal.SIGKILL) - fired: Final = tuple(future.result() for future in futures) - assert_completed_shots_hopped_stripped(owned.gateway, groups, fired) - probes: Final = tuple(fire(owned.gateway, groups, kind, marker()) for kind in BURST_KINDS) - assert all(probe.status == 200 for probe in probes), probes - assert_burst_hopped_stripped(owned.gateway, groups, probes) + with owned.replacing_worker(min(workers)) as kill: + with lanes(owned.gateway, [group.target for group in groups.values()], 30) as pinned: + with ThreadPoolExecutor(max_workers=30) as pool: + futures: Final = [ + pool.submit(fire, lane, groups, kind, mark) + for lane, (kind, mark) in zip(pinned, shots, strict=True) + ] + kill() + fired: Final = tuple(future.result() for future in futures) + assert_completed_shots_hopped_stripped(owned.gateway, groups, fired) + probes: Final = tuple(fire(owned.gateway, groups, kind, marker()) for kind in BURST_KINDS) + assert all(probe.status == 200 for probe in probes), probes + assert_burst_hopped_stripped(owned.gateway, groups, probes) diff --git a/tests/integration/routing/test_payment_required_cooldown.py b/tests/integration/routing/test_payment_required_cooldown.py index a7b80392636..12190efee26 100644 --- a/tests/integration/routing/test_payment_required_cooldown.py +++ b/tests/integration/routing/test_payment_required_cooldown.py @@ -2,7 +2,6 @@ from __future__ import annotations import json import re -import signal import uuid from collections.abc import Callable, Iterator, Mapping, Sequence from concurrent.futures import ThreadPoolExecutor @@ -14,7 +13,6 @@ from typing import Final import anthropic import httpx import openai -import psutil import pytest import yaml from integration._support.client import GATEWAY_LIMITS, Gateway, Scenario, eventually, gateway_from_environment @@ -54,7 +52,6 @@ COOLED_DOWN: Final = "litellm_deployment_cooled_down_total" PRIMARY_GROUP: Final = new_group() FALLBACK_GROUP: Final = new_group() MARKER: Final = re.compile(rb"pr402-[0-9a-f]{32}(?:-\d+)?") -WORKER_PID: Final = re.compile(r"Started server process \[(\d+)\]") @dataclass(frozen=True, slots=True) @@ -506,10 +503,6 @@ def burst_call(group: str, index: int) -> Call: return Call(marker, "/v1/responses", {"model": group, "input": marker, "stream": stream}) -def worker_pids(log: Path) -> tuple[int, ...]: - return tuple(int(pid) for pid in WORKER_PID.findall(log.read_text())) - - def ready_workers(log: Path) -> int: return log.read_text().count("Application startup complete.") @@ -552,12 +545,7 @@ def test_chaos_burst_cooldown_holds_across_both_workers_and_a_worker_kill(tmp_pa assert_served(warm, chat_identity) assert markers_received(refusing) == () assert markers_received(healthy) == markers(warm) - victim: Final = next(pid for pid in sorted(worker_pids(owned.log)) if psutil.pid_exists(pid)) - ready_before: Final = ready_workers(owned.log) - psutil.Process(victim).send_signal(signal.SIGKILL) - eventually( - lambda: ready_workers(owned.log), lambda ready: ready > ready_before, seconds=graceful_stop_seconds() - ) + owned.kill_worker(min(owned.worker_pids())) after_kill: Final = tuple(raw_chat(patient, group, new_marker()) for _ in range(20)) assert_served(after_kill, chat_identity) assert markers_received(refusing) == () diff --git a/tests/integration/routing/test_per_model_rate_limit_fallbacks.py b/tests/integration/routing/test_per_model_rate_limit_fallbacks.py index 428a485f397..67abc211900 100644 --- a/tests/integration/routing/test_per_model_rate_limit_fallbacks.py +++ b/tests/integration/routing/test_per_model_rate_limit_fallbacks.py @@ -2,7 +2,6 @@ from __future__ import annotations import json import re -import signal import uuid from collections import Counter from collections.abc import Iterator, Mapping, Sequence @@ -16,20 +15,18 @@ from typing import Final import anthropic import httpx import openai -import psutil import pytest import yaml from integration._support.client import ( JSON_OBJECT, Gateway, Scenario, - eventually, gateway_from_environment, object_value, string_value, ) from integration._support.openai_wire import chat_reply, responses_reply -from integration._support.process import graceful_stop_seconds, owned_proxy_process +from integration._support.process import OwnedProxy, graceful_stop_seconds, owned_proxy_process from integration._support.wire import Reply, Request, Wire, wire_server from integration.routing.test_payment_required_mapping import OPENAI_MODEL, is_model_info_probe, model_list_reply from pydantic import JsonValue, TypeAdapter @@ -39,10 +36,7 @@ pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120) SETTING: Final = "disable_fallbacks_on_per_model_rate_limits" FRESH_CONNECTION: Final = MappingProxyType({"Connection": "close"}) MARKER: Final = re.compile(rb"pmrl-[0-9a-f]{32}") -WORKER_STARTED: Final = re.compile(r"Started server process \[(\d+)\]") -STARTUP_COMPLETE: Final = "Application startup complete." CONFIG_LIST: Final = TypeAdapter(list[dict[str, JsonValue]]) -WORKER_PIDS: Final = TypeAdapter(tuple[int, ...]) BURST: Final = 8 TOKEN_LIMIT: Final = 100 INPUT_TOKEN_LIMIT: Final = 120 @@ -75,10 +69,14 @@ class Groups: @dataclass(frozen=True, slots=True) class Rig: gateway: Gateway - log: Path + owned: OwnedProxy groups: Groups wires: Mapping[str, Wire] + @property + def log(self) -> Path: + return self.owned.log + def traffic(self) -> dict[str, tuple[str, ...]]: return {group: markers_on(wire) for group, wire in self.wires.items()} @@ -221,7 +219,7 @@ def enforcing(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]: rig_gateway, directory, {}, config=enforcing_config(directory, groups, wires), workers=2 ) ) - yield Rig(owned.gateway, owned.log, groups, wires) + yield Rig(owned.gateway, owned, groups, wires) def chat(gateway: Gateway, key: str, group: str, content: str, **extra: JsonValue) -> httpx.Response: @@ -326,11 +324,6 @@ def served_once_in_a_burst(rig: Rig, markers: Sequence[str], responses: Sequence return served -def worker_startups(log: Path) -> tuple[tuple[int, ...], int]: - text: Final = log.read_text(errors="replace") - return WORKER_PIDS.validate_python(WORKER_STARTED.findall(text)), text.count(STARTUP_COMPLETE) - - def listed_setting(gateway: Gateway) -> tuple[dict[str, JsonValue], ...]: response: Final = gateway.request("GET", "/config/list", params={"config_type": "general_settings"}) assert response.status_code == 200, response.text @@ -604,16 +597,7 @@ def test_yaml_owned_setting_refuses_a_database_write_and_keeps_enforcing(enforci def test_worker_sigkill_and_replacement_keeps_refusing_per_model_overflow_without_fallback(enforcing: Rig) -> None: enforcing.traffic() - pids, ready = eventually( - lambda: worker_startups(enforcing.log), lambda found: len(found[0]) >= 2 and found[1] >= 2, seconds=60 - ) - victim: Final = next(pid for pid in pids if psutil.pid_exists(pid)) - psutil.Process(victim).send_signal(signal.SIGKILL) - eventually( - lambda: worker_startups(enforcing.log), - lambda found: len(found[0]) > len(pids) and found[1] > ready, - seconds=graceful_stop_seconds(), - ) + enforcing.owned.kill_worker(min(enforcing.owned.worker_pids())) primary: Final = enforcing.groups.primary markers: Final = tuple(new_marker() for _ in range(BURST)) with enforcing.gateway.scenario() as scenario: diff --git a/tests/integration/routing/test_per_model_rate_limit_fallbacks_lifecycle.py b/tests/integration/routing/test_per_model_rate_limit_fallbacks_lifecycle.py index 4361843b2c9..ea47b590f9b 100644 --- a/tests/integration/routing/test_per_model_rate_limit_fallbacks_lifecycle.py +++ b/tests/integration/routing/test_per_model_rate_limit_fallbacks_lifecycle.py @@ -67,7 +67,7 @@ def fallback_rig( owned: Final = stack.enter_context( owned_proxy_process(rig_gateway, directory, environment, config=config, workers=2) ) - yield Rig(owned.gateway, owned.log, groups, wires) + yield Rig(owned.gateway, owned, groups, wires) @pytest.fixture(scope="module") @@ -185,7 +185,7 @@ def test_database_setting_survives_a_restart_and_refuses_per_model_overflow(tmp_ ) listed: Final = listed_setting(restarted.gateway) assert [(entry["field_value"], entry["stored_in_db"]) for entry in listed] == [(True, True)], listed - rig: Final = Rig(restarted.gateway, restarted.log, groups, wires) + rig: Final = Rig(restarted.gateway, restarted, groups, wires) rig.traffic() key: Final = string_value(rig.gateway.post("/key/generate", {"model_rpm_limit": {groups.primary: 1}})["key"]) first, second = new_marker(), new_marker() diff --git a/tests/integration/routing/test_priority_scheduler_queue_cleanup.py b/tests/integration/routing/test_priority_scheduler_queue_cleanup.py index 3c38acadaaf..bcfb28f790a 100644 --- a/tests/integration/routing/test_priority_scheduler_queue_cleanup.py +++ b/tests/integration/routing/test_priority_scheduler_queue_cleanup.py @@ -4,7 +4,6 @@ import asyncio import http.client import json import re -import signal import threading import time import uuid @@ -27,6 +26,7 @@ from integration._support.database import read_rows from integration._support.openai_wire import answering_model_discovery, chat_reply, openai_error, responses_reply from integration._support.process import graceful_stop_seconds, owned_proxy_process from integration._support.redis_process import OwnedRedis, owned_redis +from integration._support.timing import assert_cut_short from integration._support.wire import Reply, Request, Wire, wire_server from pydantic import JsonValue, TypeAdapter from redis import Redis @@ -624,7 +624,7 @@ def test_fresh_queue_during_a_cooldown_is_refused_at_once(pair: Pair) -> None: client.completions.create(model=group, prompt=new_marker(), extra_body={"priority": 1}) assert refused.value.status_code == 429, refused.value.message assert NO_DEPLOYMENTS in refused.value.message - assert time.monotonic() - started < PROMPT_SECONDS + assert_cut_short(time.monotonic() - started, injected_delay=ROUTER_TIMEOUT_SECONDS) assert queue_entries(pair.cache, group) == [] @@ -825,18 +825,18 @@ def test_sibling_worker_keeps_serving_prioritized_requests_after_a_worker_is_kil held_by: Final = {pid: established_upstream_connections(pid, wire) for pid in workers} assert sum(held_by.values()) == 20, held_by victim_pid, survivor_pid = sorted(workers, key=held_by.__getitem__) - victim: Final = psutil.Process(victim_pid) - victim.suspend() - victim.send_signal(signal.SIGKILL) - upstream.release.set() - outcomes: Final = pending.result(timeout=90) - answered: Final = tuple(outcome for outcome in outcomes if isinstance(outcome[2], Answer)) - failed: Final = tuple(outcome for outcome in outcomes if not isinstance(outcome[2], Answer)) - assert len(answered) == held_by[survivor_pid], (held_by, len(answered)) - assert len(failed) == held_by[victim_pid], (held_by, len(failed)) - assert_all_served(answered) - follow_up: Final = new_marker() - served_eventually(owned.gateway, KILL_GROUP, follow_up) + with owned.replacing_worker(victim_pid) as kill: + kill.freeze() + kill() + upstream.release.set() + outcomes: Final = pending.result(timeout=90) + answered: Final = tuple(outcome for outcome in outcomes if isinstance(outcome[2], Answer)) + failed: Final = tuple(outcome for outcome in outcomes if not isinstance(outcome[2], Answer)) + assert len(answered) == held_by[survivor_pid], (held_by, len(answered)) + assert len(failed) == held_by[victim_pid], (held_by, len(failed)) + assert_all_served(answered) + follow_up: Final = new_marker() + served_eventually(owned.gateway, KILL_GROUP, follow_up) eventually( lambda: len(STARTED_WORKER.findall(owned.log.read_text())), lambda count: count == 3, diff --git a/tests/integration/run.py b/tests/integration/run.py index 97c79d1845d..43e606f46a7 100644 --- a/tests/integration/run.py +++ b/tests/integration/run.py @@ -1,13 +1,15 @@ from __future__ import annotations import argparse +import importlib.util import json import os import subprocess import sys +from collections.abc import Mapping from dataclasses import dataclass from pathlib import Path -from types import MappingProxyType +from types import MappingProxyType, ModuleType from typing import Final GROUPS: Final = MappingProxyType( @@ -32,10 +34,15 @@ GITHUB_FILES: Final = frozenset( "tests/integration/database/image/test_offline_image_migration.py", "tests/integration/database/image/test_ui_image_serves_offline.py", "tests/integration/mcp/test_interactions.py", + "tests/integration/sdk/test_redis_cluster_iam_auth.py", } ) +PEER_GROUPS: Final = frozenset({"management", "mcp"}) +COST_MAP_GROUPS: Final = frozenset({"cost"}) + + @dataclass(frozen=True, slots=True) class Selection: nodes: tuple[str, ...] @@ -59,6 +66,28 @@ def uncollected(nodes: tuple[str, ...], collected: frozenset[str]) -> tuple[str, return tuple(node for node in nodes if file_of(node) not in collected_files) +def worker_settings(group: str, workers: int) -> Mapping[str, str]: + if workers < 2: + return MappingProxyType({}) + return MappingProxyType( + { + "INTEGRATION_WORKER_ISOLATION": "1", + "INTEGRATION_WORKER_PEER": "1" if group in PEER_GROUPS else "", + "INTEGRATION_WORKER_COST_MAP": "1" if group in COST_MAP_GROUPS else "", + } + ) + + +def _budget() -> ModuleType: + specification: Final = importlib.util.spec_from_file_location( + "integration_budget", Path(__file__).with_name("_support") / "budget.py" + ) + assert specification is not None and specification.loader is not None + module: Final = importlib.util.module_from_spec(specification) + specification.loader.exec_module(module) + return module + + def main() -> int: parser: Final = argparse.ArgumentParser() parser.add_argument("group", choices=tuple(GROUPS)) @@ -96,12 +125,14 @@ def main() -> int: "LITELLM_LOCAL_MODEL_COST_MAP": "True", "LITELLM_MODE": "PRODUCTION", "PYTHON_DOTENV_DISABLED": "1", - **( - {"INTEGRATION_WORKER_ISOLATION": "1"} - if options.workers > 1 and options.group in {"management", "accounting", "providers", "extensions"} - else {} - ), + **worker_settings(options.group, options.workers), } + budget: Final = _budget() + proxy_root: Final = Path(os.environ.get("INTEGRATION_PROXY_ROOT") or root) + imported: Final = float( + os.environ.get(budget.IMPORT_SECONDS_VARIABLE) or budget.measure_import_seconds(proxy_root, environment) + ) + (output / "boot-calibration.json").write_text(budget.calibration(imported) + "\n") result: Final = subprocess.call( [ sys.executable, @@ -115,7 +146,7 @@ def main() -> int: "no:pytest-retry", "-p", "no:rerunfailures", - "--timeout=90", + f"--timeout={int(budget.case_timeout_for(imported))}", "--durations=15", "--tb=short", f"--hypothesis-seed={options.seed}", @@ -130,7 +161,7 @@ def main() -> int: *(("--dist=loadgroup",) if options.workers > 1 else ()), ], cwd=root, - env=environment, + env={**environment, budget.IMPORT_SECONDS_VARIABLE: str(imported)}, ) if result != 0: return result diff --git a/tests/integration/sdk/test_owned_redis_ports.py b/tests/integration/sdk/test_owned_redis_ports.py new file mode 100644 index 00000000000..0212329b7b7 --- /dev/null +++ b/tests/integration/sdk/test_owned_redis_ports.py @@ -0,0 +1,56 @@ +from __future__ import annotations + +import re +from itertools import chain +from pathlib import Path +from typing import Final + +import pytest +from integration._support.redis_process import PORT_ATTEMPTS, free_loopback_port, owned_redis +from redis import Redis + + +def _redis_pid(host: str, port: int) -> int: + with Redis(host=host, port=port) as client: + return int(client.info("server")["process_id"]) + + +def test_owned_redis_moves_off_a_port_another_owned_redis_already_serves(tmp_path: Path) -> None: + with owned_redis(tmp_path / "first") as first: + candidates: Final = chain((first.port,), iter(free_loopback_port, None)) + with owned_redis(tmp_path / "second", lambda: next(candidates)) as second: + assert second.port != first.port + assert first.serves() + assert second.serves() + assert second.server_pid != first.server_pid + assert _redis_pid(second.host, second.port) == second.server_pid + assert _redis_pid(first.host, first.port) == first.server_pid + assert not second.serves() + assert first.serves() + + +def test_owned_redis_restart_refuses_a_port_another_redis_took_while_it_was_down(tmp_path: Path) -> None: + with owned_redis(tmp_path / "first") as first: + first.stop() + with owned_redis(tmp_path / "squatter", lambda: first.port) as squatter: + with pytest.raises( + AssertionError, match=f"Another Redis already serves {re.escape(first.host)}:{first.port}" + ): + first.start() + assert first.process is None + assert _redis_pid(squatter.host, squatter.port) == squatter.server_pid + + +def test_owned_redis_gives_up_with_the_last_cause_after_a_bounded_number_of_ports(tmp_path: Path) -> None: + with owned_redis(tmp_path / "first") as first: + ports: Final = iter((first.port,) * PORT_ATTEMPTS) + with pytest.raises( + AssertionError, + match=( + f"after {PORT_ATTEMPTS} attempts: Another Redis already serves {re.escape(first.host)}:{first.port}" + ), + ): + with owned_redis(tmp_path / "second", lambda: next(ports)): + pytest.fail("Owned Redis unexpectedly reused a port already serving Redis") + assert next(ports, None) is None + assert first.serves() diff --git a/tests/integration/sdk/test_redis_cluster_iam_auth.py b/tests/integration/sdk/test_redis_cluster_iam_auth.py index c649d349dae..055b73bb415 100644 --- a/tests/integration/sdk/test_redis_cluster_iam_auth.py +++ b/tests/integration/sdk/test_redis_cluster_iam_auth.py @@ -11,6 +11,7 @@ from unittest.mock import MagicMock, patch import pytest import redis +from tests.integration._support.process import stop_child from litellm._redis import _get_redis_env_kwarg_mapping, get_redis_client from litellm._redis_credential_provider import _token_cache @@ -91,12 +92,7 @@ def authenticated_redis_cluster(tmp_path: Path, unused_tcp_port_factory: Callabl admin.execute_command("ACL", "SETUSER", "default", "resetpass", ">local-fixture-token") yield port finally: - process.terminate() - try: - process.wait(timeout=5) - except subprocess.TimeoutExpired: - process.kill() - process.wait() + stop_child(process, grace_seconds=5) if process.stdout is not None: process.stdout.close() diff --git a/tests/integration/sdk/test_sagemaker_embedding_sdk.py b/tests/integration/sdk/test_sagemaker_embedding_sdk.py new file mode 100644 index 00000000000..6b70a27fbec --- /dev/null +++ b/tests/integration/sdk/test_sagemaker_embedding_sdk.py @@ -0,0 +1,171 @@ +import uuid +from collections.abc import Iterator +from typing import Final + +import pytest +from integration._support.sagemaker_embedding import ( + ACCESS_KEY, + COMPONENT, + COMPONENT_HEADER, + JSON_OBJECT, + NON_STRING_MODEL_ID, + REGION, + SECRET_KEY, + SERVED_MODEL, + endpoint, + invocations_target, + litellm_params, + openai_route, + plain_route, + respond, +) +from integration._support.wire import Request, Wire, wire_server +from pydantic import JsonValue + +import litellm + +_AWS: Final[dict[str, str]] = { + "aws_access_key_id": ACCESS_KEY, + "aws_secret_access_key": SECRET_KEY, + "aws_region_name": REGION, +} + + +@pytest.fixture(scope="module") +def wire() -> Iterator[Wire]: + with wire_server(respond) as served: + yield served + + +@pytest.fixture(autouse=True) +def _runtime_at_the_peer(wire: Wire, monkeypatch: pytest.MonkeyPatch) -> None: + wire.drain() + monkeypatch.setenv("AWS_ENDPOINT_URL_SAGEMAKER_RUNTIME", wire.url) + monkeypatch.setenv("AWS_MAX_ATTEMPTS", "1") + + +def _text() -> str: + return f"integration embedding {uuid.uuid4().hex}" + + +def _only_request(wire: Wire) -> Request: + received: Final = wire.drain() + assert len(received) == 1, [(request.method, request.target) for request in received] + return received[0] + + +def _signed_headers(request: Request) -> tuple[str, ...]: + authorization: Final = request.headers["authorization"] + assert authorization.startswith(f"AWS4-HMAC-SHA256 Credential={ACCESS_KEY}/"), authorization + signed: Final = next(part for part in authorization.split(", ") if part.startswith("SignedHeaders=")) + return tuple(signed.removeprefix("SignedHeaders=").split(";")) + + +def _assert_component_request(request: Request, scenario: str, body: dict[str, JsonValue]) -> None: + assert (request.method, request.target) == ("POST", invocations_target(endpoint(scenario))), request + assert request.headers[COMPONENT_HEADER] == COMPONENT, dict(request.headers) + assert COMPONENT_HEADER in _signed_headers(request), request.headers["authorization"] + assert JSON_OBJECT.validate_json(request.body) == body, request.body + + +def _vectors(response: litellm.EmbeddingResponse) -> list[JsonValue]: + return [row["embedding"] for row in response.data] + + +def test_embedding_signs_model_id_as_the_inference_component_and_reads_the_served_model(wire: Wire) -> None: + text: Final = _text() + response: Final = litellm.embedding( + model=openai_route("openai"), input=[text], model_id=COMPONENT, extra_body={"model": SERVED_MODEL}, **_AWS + ) + assert response.model == SERVED_MODEL, response + assert _vectors(response) == [[0.25, -0.5, 1.0]], response + assert (response.usage.prompt_tokens, response.usage.total_tokens) == (7, 7), response.usage + _assert_component_request(_only_request(wire), "openai", {"input": [text], "model": SERVED_MODEL}) + + +async def test_aembedding_signs_model_id_as_the_inference_component_and_reads_the_served_model(wire: Wire) -> None: + texts: Final = [_text(), _text()] + response: Final = await litellm.aembedding( + model=openai_route("openai"), input=texts, model_id=COMPONENT, extra_body={"model": SERVED_MODEL}, **_AWS + ) + assert response.model == SERVED_MODEL, response + assert _vectors(response) == [[0.25, -0.5, 1.0], [0.25, -0.5, 2.0]], response + _assert_component_request(_only_request(wire), "openai", {"input": texts, "model": SERVED_MODEL}) + + +def test_embedding_names_the_endpoint_when_the_container_reports_no_model(wire: Wire) -> None: + text: Final = _text() + response: Final = litellm.embedding(model=openai_route("nomodel"), input=[text], model_id=COMPONENT, **_AWS) + assert response.model == endpoint("nomodel"), response + assert _vectors(response) == [[0.25, -0.5, 1.0]], response + _assert_component_request(_only_request(wire), "nomodel", {"input": [text]}) + + +def test_plain_route_keeps_the_hf_body_and_signs_the_component_header(wire: Wire) -> None: + text: Final = _text() + response: Final = litellm.embedding(model=plain_route("hf"), input=[text], model_id=COMPONENT, **_AWS) + assert response.model == endpoint("hf"), response + assert _vectors(response) == [[0.25, -0.5, 1.0]], response + assert response.usage.prompt_tokens == 3, response.usage + _assert_component_request(_only_request(wire), "hf", {"inputs": [text]}) + + +def test_plain_route_without_model_id_sends_no_component_header(wire: Wire) -> None: + text: Final = _text() + response: Final = litellm.embedding(model=plain_route("hf"), input=[text], **_AWS) + assert _vectors(response) == [[0.25, -0.5, 1.0]], response + request: Final = _only_request(wire) + assert (request.method, request.target) == ("POST", invocations_target(endpoint("hf"))), request + assert COMPONENT_HEADER not in request.headers, dict(request.headers) + assert COMPONENT_HEADER not in _signed_headers(request), request.headers["authorization"] + assert JSON_OBJECT.validate_json(request.body) == {"inputs": [text]}, request.body + + +def test_explicit_provider_spelling_signs_model_id_as_the_inference_component(wire: Wire) -> None: + text: Final = _text() + response: Final = litellm.embedding( + model=f"openai/{endpoint('openai')}", + custom_llm_provider="sagemaker", + input=[text], + model_id=COMPONENT, + extra_body={"model": SERVED_MODEL}, + **_AWS, + ) + assert response.model == SERVED_MODEL, response + assert _vectors(response) == [[0.25, -0.5, 1.0]], response + _assert_component_request(_only_request(wire), "openai", {"input": [text], "model": SERVED_MODEL}) + + +async def test_router_aembedding_forwards_the_deployment_model_id_as_the_component(wire: Wire) -> None: + text: Final = _text() + router: Final = litellm.Router( + model_list=[ + { + "model_name": "sm-openai-router", + "litellm_params": litellm_params( + openai_route("openai"), model_id=COMPONENT, extra_body={"model": SERVED_MODEL} + ), + } + ], + num_retries=0, + ) + response: Final = await router.aembedding(model="sm-openai-router", input=[text]) + assert response.model == SERVED_MODEL, response + assert _vectors(response) == [[0.25, -0.5, 1.0]], response + _assert_component_request(_only_request(wire), "openai", {"input": [text], "model": SERVED_MODEL}) + + +def test_embedding_rejects_a_non_string_model_id_before_calling_aws(wire: Wire) -> None: + with pytest.raises(litellm.BadRequestError) as caught: + litellm.embedding(model=openai_route("openai"), input=[_text()], model_id=[COMPONENT], **_AWS) + assert caught.value.status_code == 400, caught.value + assert NON_STRING_MODEL_ID in str(caught.value), caught.value + assert wire.drain() == (), "AWS was called with a list-valued model_id" + + +def test_embedding_surfaces_the_container_model_error_with_its_status(wire: Wire) -> None: + with pytest.raises(litellm.BadRequestError) as caught: + litellm.embedding(model=openai_route("err424"), input=[_text()], model_id=COMPONENT, **_AWS) + assert caught.value.status_code == 424, caught.value + assert "unexpected field inputs" in str(caught.value), caught.value + assert _only_request(wire).headers[COMPONENT_HEADER] == COMPONENT diff --git a/tests/integration/security/conftest.py b/tests/integration/security/conftest.py deleted file mode 100644 index b9f33c71f63..00000000000 --- a/tests/integration/security/conftest.py +++ /dev/null @@ -1,26 +0,0 @@ -from __future__ import annotations - -import os -from collections.abc import Iterator -from pathlib import Path -from typing import Final - -import pytest - -from tests.integration._support.workers import worker_services - - -@pytest.fixture(scope="session", autouse=True) -def isolated_security_worker(worker_id: str, tmp_path_factory: pytest.TempPathFactory) -> Iterator[None]: - if worker_id == "master": - yield - return - directory: Final = tmp_path_factory.mktemp("security-worker") - results: Final = Path(os.environ.get("INTEGRATION_RESULTS_DIR", str(directory))) / worker_id - results.mkdir(parents=True, exist_ok=True) - with pytest.MonkeyPatch.context() as environment: - environment.setenv("INTEGRATION_RESULTS_DIR", str(results)) - with worker_services(directory) as services: - for name, value in services.items(): - environment.setenv(name, value) - yield diff --git a/tests/integration/spend/test_background_interaction_settlement.py b/tests/integration/spend/test_background_interaction_settlement.py index 11444396697..75cb2912076 100644 --- a/tests/integration/spend/test_background_interaction_settlement.py +++ b/tests/integration/spend/test_background_interaction_settlement.py @@ -761,12 +761,12 @@ def test_killed_workers_leave_their_polls_to_the_respawned_workers(rig: Rig) -> for item in created: _state(rig, item, InteractionState(status="in_progress")) rates: Final = _rates(resumer.gateway, model) - killed: Final = _worker_pids(resumer.process.pid) + killed: Final = resumer.worker_pids() assert len(killed) == 2, killed - victims: Final = tuple(psutil.Process(pid) for pid in killed) - for victim in victims: - victim.kill() - psutil.wait_procs(victims, timeout=15) + first, second = sorted(killed) + with resumer.replacing_worker(first) as kill_first, resumer.replacing_worker(second) as kill_second: + kill_first() + kill_second() for item in created: _state(rig, item, _completed()) for item in created: diff --git a/tests/integration/spend/test_daily_activity_key_alias_probes.py b/tests/integration/spend/test_daily_activity_key_alias_probes.py index 8d9b8616435..db22d2402fd 100644 --- a/tests/integration/spend/test_daily_activity_key_alias_probes.py +++ b/tests/integration/spend/test_daily_activity_key_alias_probes.py @@ -1,4 +1,3 @@ -import time import uuid from collections.abc import Callable, Iterator from concurrent.futures import ThreadPoolExecutor @@ -36,10 +35,12 @@ from integration._support.daily_activity import ( ) from integration._support.database import read_rows, scratch_database from integration._support.process import OwnedProxy, owned_proxy_process +from integration._support.timing import assert_enforced_within, timed +from litellm.constants import SPEND_LOG_KEY_METADATA_QUERY_TIMEOUT_MS from pydantic import JsonValue DAY_OUTSIDE_THE_WINDOW: Final = "2026-02-10" -GIVES_UP_WITHIN_SECONDS: Final = 10 +LOCKED_LOOKUP_SLACK_SECONDS: Final = 5 CONCURRENT_READS: Final = 20 CACHED_MISS_CLEARS_WITHIN_SECONDS: Final = 45 ALIAS_OF_ONE_SPEND_LOG: Final = ( @@ -450,16 +451,18 @@ def test_alias_lookup_gives_up_while_spend_logs_are_locked_and_answers_once_they spend_logs_of_key(api_key, (named_row(0, alias),), database_url=database_url), ): with locked_table(SPEND_LOGS_TABLE, database_url=database_url): - started: Final = time.monotonic() - locked: Final = activity_of_key(owned.gateway, AGGREGATED_USER_ACTIVITY, api_key) - waited: Final = time.monotonic() - started + locked: Final = timed(lambda: activity_of_key(owned.gateway, AGGREGATED_USER_ACTIVITY, api_key)) unlocked: Final = eventually( lambda: activity_of_key(owned.gateway, AGGREGATED_USER_ACTIVITY, api_key), _names_the_key(api_key, alias), seconds=CACHED_MISS_CLEARS_WITHIN_SECONDS, ) - assert waited < GIVES_UP_WITHIN_SECONDS, waited - assert_key_reported(locked, api_key, DAY, key_metadata(user=owner, email=email), seeded_metrics(1)) + assert_enforced_within( + locked.seconds, + configured=SPEND_LOG_KEY_METADATA_QUERY_TIMEOUT_MS / 1000, + slack=LOCKED_LOOKUP_SLACK_SECONDS, + ) + assert_key_reported(locked.value, api_key, DAY, key_metadata(user=owner, email=email), seeded_metrics(1)) assert_key_reported(unlocked, api_key, DAY, key_metadata(alias=alias, user=owner, email=email), seeded_metrics(1)) diff --git a/tests/integration/spend/test_daily_activity_key_metadata_query_timeout.py b/tests/integration/spend/test_daily_activity_key_metadata_query_timeout.py index 747570a01ec..780ee43fb3f 100644 --- a/tests/integration/spend/test_daily_activity_key_metadata_query_timeout.py +++ b/tests/integration/spend/test_daily_activity_key_metadata_query_timeout.py @@ -1,7 +1,5 @@ import itertools import json -import os -import signal import threading import uuid from bisect import bisect_left @@ -1025,14 +1023,6 @@ def _running_children(owned: OwnedProxy) -> tuple[int, ...]: ) -def _worker_pids(owned: OwnedProxy) -> tuple[int, ...]: - return tuple( - member.pid - for member in group_members(owned.process.pid) - if member.pid != owned.process.pid and any("spawn_main" in part for part in member.cmdline()) - ) - - @pytest.mark.timeout(600) def test_usage_page_keeps_serving_when_a_worker_dies_mid_outage(gateway: Gateway, tmp_path: Path) -> None: with ( @@ -1042,23 +1032,25 @@ def test_usage_page_keeps_serving_when_a_worker_dies_mid_outage(gateway: Gateway ): spender: Final = _spender(owned.gateway, database_url, "worker-death") children: Final = _running_children(owned) - workers: Final = _worker_pids(owned) + workers: Final = owned.worker_pids() assert len(workers) == WORKERS, workers + victim: Final = min(workers) with _locked_spend_logs(database_url): timed_out: Final = _fresh_read(owned, spender.digest) assert timed_out.status_code == 200, timed_out.text assert timed_out.elapsed >= FAILED_LOOKUP_FLOOR, timed_out.elapsed assert _aliases(timed_out, spender.digest) == (None,), timed_out.text - os.kill(workers[0], signal.SIGKILL) - after_kill: Final = _fresh_read(owned, spender.digest) - assert after_kill.status_code == 200, after_kill.text - assert _aliases(after_kill, spender.digest) == (None,), after_kill.text + with owned.replacing_worker(victim) as kill: + kill() + after_kill: Final = _fresh_read(owned, spender.digest) + assert after_kill.status_code == 200, after_kill.text + assert _aliases(after_kill, spender.digest) == (None,), after_kill.text respawned: Final = eventually( lambda: _running_children(owned), lambda pids: len(pids) >= len(children) and any(pid not in children for pid in pids), seconds=30, ) - assert workers[0] not in respawned, respawned + assert victim not in respawned, respawned recovered: Final = eventually( lambda: _fresh_read(owned, spender.digest), lambda response: response.status_code == 200 and _aliases(response, spender.digest) == (spender.alias,), diff --git a/tests/integration/spend/test_daily_activity_key_owner_faults.py b/tests/integration/spend/test_daily_activity_key_owner_faults.py index 2ab56622ec4..1dba97d6ae1 100644 --- a/tests/integration/spend/test_daily_activity_key_owner_faults.py +++ b/tests/integration/spend/test_daily_activity_key_owner_faults.py @@ -1,6 +1,3 @@ -import os -import signal -import time import uuid from collections.abc import Iterator from contextlib import contextmanager @@ -32,12 +29,14 @@ from integration._support.daily_activity import ( ) from integration._support.database import scratch_database from integration._support.process import OwnedProxy, group_members, owned_proxy_process +from integration._support.timing import assert_enforced_within, timed +from litellm.constants import SPEND_LOG_KEY_METADATA_QUERY_TIMEOUT_MS USER_ACTIVITY: Final = "/user/daily/activity" TEAM_ACTIVITY: Final = "/team/daily/activity" AGGREGATED_TEAM_ACTIVITY: Final = "/team/daily/activity/aggregated" KEYS_OF_ONE_TEAM: Final = 300 -GIVES_UP_WITHIN_SECONDS: Final = 10 +LOCKED_LOOKUP_SLACK_SECONDS: Final = 5 READS_AFTER_THE_WORKER_IS_REPLACED: Final = 6 @@ -214,12 +213,14 @@ def test_owner_lookup_gives_up_while_daily_user_spend_is_locked_and_answers_once rows: Final = (user_row(owner, api_key, DAY), seeded_row(TEAM_SPEND, "team_id", team, api_key, DAY)) with daily_rows(rows, database_url=database_url): with locked_table(USER_SPEND, database_url=database_url): - started: Final = time.monotonic() - locked: Final = activity_of_key(owned.gateway, TEAM_ACTIVITY, api_key, team_ids=team) - waited: Final = time.monotonic() - started + locked: Final = timed(lambda: activity_of_key(owned.gateway, TEAM_ACTIVITY, api_key, team_ids=team)) unlocked: Final = activity_of_key(owned.gateway, TEAM_ACTIVITY, api_key, team_ids=team) - assert waited < GIVES_UP_WITHIN_SECONDS, waited - assert_key_reported(locked, api_key, DAY, key_metadata(), seeded_metrics(1)) + assert_enforced_within( + locked.seconds, + configured=SPEND_LOG_KEY_METADATA_QUERY_TIMEOUT_MS / 1000, + slack=LOCKED_LOOKUP_SLACK_SECONDS, + ) + assert_key_reported(locked.value, api_key, DAY, key_metadata(), seeded_metrics(1)) assert_key_reported(unlocked, api_key, DAY, key_metadata(user=owner, email=email), seeded_metrics(1)) @@ -234,12 +235,11 @@ def test_owner_is_reported_while_a_worker_is_killed_and_after_it_is_replaced(gat member for member in group_members(owned.process.pid) if member.pid != owned.process.pid ) children: Final = tuple(member.pid for member in members) - workers: Final = tuple( - member.pid for member in members if any("spawn_main" in part for part in member.cmdline()) - ) + workers: Final = owned.worker_pids() assert len(workers) >= 2, workers - os.kill(workers[0], signal.SIGKILL) - during: Final = _read_on_a_new_connection(owned.gateway, api_key) + with owned.replacing_worker(min(workers)) as kill: + kill() + during: Final = _read_on_a_new_connection(owned.gateway, api_key) eventually( lambda: _running_children(owned), lambda pids: len(pids) >= len(children) and any(pid not in children for pid in pids), diff --git a/tests/integration/spend/test_daily_tag_spend_retention.py b/tests/integration/spend/test_daily_tag_spend_retention.py index fcb624cb188..0bdb981c878 100644 --- a/tests/integration/spend/test_daily_tag_spend_retention.py +++ b/tests/integration/spend/test_daily_tag_spend_retention.py @@ -1,12 +1,10 @@ import json import os -import signal import uuid from datetime import datetime, timedelta, timezone from pathlib import Path from typing import Final -import psutil import psycopg import pytest import yaml @@ -14,7 +12,7 @@ from pydantic import JsonValue, TypeAdapter from tests.integration._support.client import Gateway, eventually, string_value from tests.integration._support.database import read_rows -from tests.integration._support.process import OwnedProxy, owned_proxy, owned_proxy_process +from tests.integration._support.process import owned_proxy, owned_proxy_process CLEANUP_EVERY_MINUTE: Final = "* * * * *" RETENTION_SETTING: Final = "maximum_daily_tag_spend_retention_period" @@ -82,15 +80,6 @@ def _store_retention_setting(value: JsonValue) -> None: ) -def _listening_workers(owned: OwnedProxy) -> tuple[psutil.Process, ...]: - port: Final = owned.gateway.client.base_url.port - return tuple( - child - for child in psutil.Process(owned.process.pid).children(recursive=True) - if any(conn.status == psutil.CONN_LISTEN and conn.laddr.port == port for conn in child.net_connections("inet")) - ) - - def _listed_retention_value(gateway: Gateway) -> JsonValue: listed: Final = _SETTINGS.validate_json( gateway.request("GET", "/config/list", params={"config_type": "general_settings"}).content @@ -216,13 +205,11 @@ def test_daily_tag_spend_cleanup_completes_after_one_of_two_workers_is_killed(ga with owned_proxy_process(gateway, tmp_path, {}, config=config, workers=2) as owned: with owned.gateway.scenario() as scenario: model: Final = scenario.model() - workers: Final = eventually( - lambda: _listening_workers(owned), lambda found: len(found) == 2, seconds=30 - ) - workers[0].send_signal(signal.SIGKILL) - eventually(lambda: workers[0].is_running(), lambda alive: not alive, seconds=10) - ids: Final = tuple(_completion_id(owned.gateway, model) for _ in range(6)) - assert len(set(ids)) == 6 and all(identity.startswith("chatcmpl-") for identity in ids), ids + workers: Final = eventually(owned.worker_pids, lambda found: len(found) == 2, seconds=30) + with owned.replacing_worker(min(workers)) as kill: + kill() + ids: Final = tuple(_completion_id(owned.gateway, model) for _ in range(6)) + assert len(set(ids)) == 6 and all(identity.startswith("chatcmpl-") for identity in ids), ids remaining: Final = eventually( lambda: _remaining_days(tag), lambda days: expired not in days, seconds=150 ) diff --git a/tests/integration/spend/test_expired_ui_session_failure_rows.py b/tests/integration/spend/test_expired_ui_session_failure_rows.py new file mode 100644 index 00000000000..0acdc86709e --- /dev/null +++ b/tests/integration/spend/test_expired_ui_session_failure_rows.py @@ -0,0 +1,127 @@ +import uuid +from collections.abc import Iterator +from dataclasses import dataclass +from typing import Final + +import jwt +import pytest +from integration._support.client import Gateway, eventually, gateway_from_environment, held, string_value +from integration._support.database import read_rows +from integration._support.process import OwnedProxy, graceful_stop_seconds, owned_proxy_process + +from litellm.proxy._types import hash_token + +pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 240) + +SESSION_SECONDS: Final = 3 +DASHBOARD_INFO_ROUTES: Final = ( + "/key/list", + "/v2/user/info", + "/team/list", + "/organization/list", + "/v2/model/info", + "/get/user_banner", +) +SPEND_LOG_FLUSH_SECONDS: Final = 25 + + +@dataclass(frozen=True, slots=True) +class ExpiredKeys: + session: str + virtual: str + + +@pytest.fixture(scope="module") +def owned(tmp_path_factory: pytest.TempPathFactory) -> Iterator[OwnedProxy]: + with ( + gateway_from_environment() as base, + owned_proxy_process( + base, + tmp_path_factory.mktemp("expired-ui-session"), + {"LITELLM_UI_SESSION_DURATION": f"{SESSION_SECONDS}s", "LITELLM_DISABLE_NO_REDIS_WARNING": "true"}, + workers=2, + ) as proxy, + ): + yield proxy + + +def _session_key(gateway: Gateway) -> str: + response: Final = gateway.client.post("/v2/login", json={"username": "admin", "password": gateway.key}) + assert response.status_code == 200, response.text + claims: Final = jwt.decode(string_value(response.json()["token"]), options={"verify_signature": False}) + key: Final = string_value(claims["key"]) + rows: Final = read_rows('SELECT team_id FROM "LiteLLM_VerificationToken" WHERE token = %s', (hash_token(key),)) + assert [row["team_id"] for row in rows] == ["litellm-dashboard"], rows + return key + + +def _rejected_as_expired(gateway: Gateway, key: str, route: str) -> None: + response: Final = eventually( + lambda: gateway.request("GET", route, key=key), + lambda reply: reply.status_code == 401, + seconds=30, + ) + assert "Expired Key" in response.text, response.text + + +def _failure_rows(key: str) -> list[str]: + rows: Final = read_rows( + 'SELECT call_type FROM "LiteLLM_SpendLogs" WHERE api_key = %s AND status = %s', + (hash_token(key), "failure"), + ) + return sorted(string_value(row["call_type"]) for row in rows) + + +@pytest.fixture +def expired(owned: OwnedProxy) -> Iterator[ExpiredKeys]: + with owned.gateway.scenario() as scenario: + keys: Final = ExpiredKeys( + session=_session_key(owned.gateway), + virtual=scenario.key(duration="1s", key_alias=f"integration-expired-{uuid.uuid4().hex[:12]}"), + ) + _rejected_as_expired(owned.gateway, keys.session, "/key/info") + _rejected_as_expired(owned.gateway, keys.virtual, "/key/info") + yield keys + + +def test_an_expired_dashboard_session_key_writes_no_failure_rows_for_dashboard_info_calls( + owned: OwnedProxy, expired: ExpiredKeys +) -> None: + for route in DASHBOARD_INFO_ROUTES: + _rejected_as_expired(owned.gateway, expired.session, route) + _rejected_as_expired(owned.gateway, expired.virtual, route) + + virtual_rows: Final = eventually( + lambda: _failure_rows(expired.virtual), + lambda rows: len(rows) == len(DASHBOARD_INFO_ROUTES) + 1, + seconds=70, + ) + assert virtual_rows == sorted(("/key/info", *DASHBOARD_INFO_ROUTES)) + assert ( + held( + lambda: _failure_rows(expired.session), + lambda rows: rows == [], + holding=SPEND_LOG_FLUSH_SECONDS, + seconds=SPEND_LOG_FLUSH_SECONDS + 5, + ) + == [] + ) + + +def test_an_expired_dashboard_session_key_still_logs_a_failed_llm_call(owned: OwnedProxy, expired: ExpiredKeys) -> None: + response: Final = owned.gateway.request( + "POST", + "/v1/chat/completions", + {"model": "integration-unrouted", "messages": [{"role": "user", "content": "expired session"}]}, + key=expired.session, + ) + assert response.status_code == 401, response.text + assert "Expired Key" in response.text, response.text + assert ( + eventually( + lambda: _failure_rows(expired.session), + lambda rows: "acompletion" in rows, + seconds=70, + ).count("acompletion") + == 1 + ) diff --git a/tests/integration/spend/test_normalized_error_long_message.py b/tests/integration/spend/test_normalized_error_long_message.py index accf160b03f..407b40d36c3 100644 --- a/tests/integration/spend/test_normalized_error_long_message.py +++ b/tests/integration/spend/test_normalized_error_long_message.py @@ -1,7 +1,5 @@ import asyncio import json -import os -import signal import threading import time import uuid @@ -16,18 +14,20 @@ from typing import Final import anthropic import httpx import openai -import psutil import pytest from integration._support.client import Gateway, eventually, object_value, string_value from integration._support.database import read_rows -from integration._support.process import OwnedProxy, owned_proxy_process +from integration._support.process import owned_proxy_process +from integration._support.timing import assert_not_delayed_by, assert_within_factor_of_control, timed from integration._support.wire import Reply, Request, wire_server from pydantic import JsonValue CRAFTED_MODEL: Final = ("exceeded " * 32_000)[:288_000] HOSTILE_5KB_MODEL: Final = ("exceeded budget " * 400)[:5_000] +BENIGN_288KB_MODEL: Final = "x" * len(CRAFTED_MODEL) +CRAFTED_SLOWDOWN_FACTOR: Final = 10 FAST_SECONDS: Final = 10.0 -LIVELINESS_MAX_SECONDS: Final = 5.0 +LIVELINESS_STALL_BUDGET_SECONDS: Final = 5.0 ROW_SECONDS: Final = 70 CHAT: Final = "/v1/chat/completions" MESSAGES: Final = "/v1/messages" @@ -129,15 +129,19 @@ def test_crafted_288kb_model_through_async_openai_sdk_fails_fast(gateway: Gatewa def test_crafted_288kb_model_through_anthropic_sdk_fails_fast(gateway: Gateway) -> None: with gateway.scenario() as scenario: client: Final = anthropic.Anthropic(base_url=str(gateway.client.base_url), api_key=scenario.key(), timeout=120) - started: Final = time.perf_counter() - with pytest.raises(anthropic.BadRequestError) as raised: - client.messages.create( - model=CRAFTED_MODEL, max_tokens=8, messages=[{"role": "user", "content": f"audit {uuid.uuid4().hex}"}] - ) - seconds: Final = time.perf_counter() - started - assert seconds < FAST_SECONDS, f"crafted 288 KB model took {seconds:.2f}s" - assert "Invalid model name passed in" in str(raised.value), str(raised.value)[:300] - info: Final = _error_information(raised.value.response.headers["x-litellm-call-id"]) + + def refused(model: str) -> anthropic.BadRequestError: + with pytest.raises(anthropic.BadRequestError) as raised: + client.messages.create( + model=model, max_tokens=8, messages=[{"role": "user", "content": f"audit {uuid.uuid4().hex}"}] + ) + return raised.value + + control: Final = timed(lambda: refused(BENIGN_288KB_MODEL)) + crafted: Final = timed(lambda: refused(CRAFTED_MODEL)) + assert_within_factor_of_control(crafted.seconds, control=control.seconds, factor=CRAFTED_SLOWDOWN_FACTOR) + assert "Invalid model name passed in" in str(crafted.value), str(crafted.value)[:300] + info: Final = _error_information(crafted.value.response.headers["x-litellm-call-id"]) assert info["normalized_error"] == "400_INVALID_REQUEST", info @@ -151,35 +155,31 @@ def _poll_liveliness(client: httpx.Client, stop: threading.Event) -> list[float] return latencies -def _cmdline(process: psutil.Process) -> str: - try: - return " ".join(process.cmdline()) - except psutil.Error: - return "" - - -def _worker_pids(owned: OwnedProxy) -> tuple[int, ...]: - return tuple(child.pid for child in psutil.Process(owned.process.pid).children() if "spawn_main" in _cmdline(child)) +def _liveliness_while_posting( + client: httpx.Client, model: str, key: str +) -> tuple[tuple[_Timed, ...], tuple[float, ...]]: + stop: Final = threading.Event() + with ThreadPoolExecutor(max_workers=3) as pool: + liveliness: Final = pool.submit(_poll_liveliness, client, stop) + posts: Final = tuple( + pool.submit(_timed_post, client, CHAT, _body(CHAT, model, uuid.uuid4().hex), key) for _ in range(2) + ) + results: Final = tuple(future.result() for future in posts) + stop.set() + return results, tuple(liveliness.result()) def test_two_concurrent_crafted_requests_do_not_stall_liveliness_on_a_two_worker_proxy( gateway: Gateway, tmp_path: Path ) -> None: with owned_proxy_process(gateway, tmp_path, {}, workers=2) as owned, _patient_client(owned.gateway) as client: - assert len(_worker_pids(owned)) == 2, _worker_pids(owned) - stop: Final = threading.Event() - with ThreadPoolExecutor(max_workers=3) as pool: - liveliness: Final = pool.submit(_poll_liveliness, client, stop) - crafted: Final = tuple( - pool.submit(_timed_post, client, CHAT, _body(CHAT, CRAFTED_MODEL, uuid.uuid4().hex), gateway.key) - for _ in range(2) - ) - results: Final = tuple(future.result() for future in crafted) - stop.set() - latencies: Final = liveliness.result() - for timed in results: - _assert_crafted_failure(timed) - assert latencies and max(latencies) < LIVELINESS_MAX_SECONDS, f"liveliness max {max(latencies):.2f}s" + assert len(owned.worker_pids()) == 2, owned.worker_pids() + _benign, control = _liveliness_while_posting(client, BENIGN_288KB_MODEL, gateway.key) + results, latencies = _liveliness_while_posting(client, CRAFTED_MODEL, gateway.key) + for posted in results: + _assert_crafted_failure(posted) + assert latencies and control, (latencies, control) + assert_not_delayed_by(latencies, control=control, injected_delay=LIVELINESS_STALL_BUDGET_SECONDS) def _completion(request: Request) -> Reply: @@ -403,25 +403,28 @@ def test_killing_one_worker_mid_burst_leaves_the_other_serving_crafted_and_happy ): model: Final = scenario.model(api_base=wire.url + "/v1", num_retries=0) key: Final = scenario.key(models=[model]) - workers: Final = _worker_pids(owned) + workers: Final = owned.worker_pids() assert len(workers) == 2, workers + victim: Final = min(workers) - def kill_one_worker() -> None: - os.kill(workers[0], signal.SIGKILL) - - results: Final = _burst(client, model, key, scenario.key(), kill_one_worker) - dropped: Final = tuple(result for result in results if result.status is None) - assert len(dropped) < len(results), results - crafted: Final = tuple(result for result in results if result.label.startswith("crafted") and result.status) - assert crafted and all(result.status == 400 and result.seconds < FAST_SECONDS for result in crafted), crafted - late: Final = tuple(result for result in results if result.label.startswith("late")) - assert all(result.status == 200 for result in late), late - _assert_rows_land_exactly_once(results, "late") - survivor: Final = tuple(pid for pid in _worker_pids(owned) if pid != workers[0]) - assert survivor, "no worker left serving" - after: Final = owned.gateway.chat(model, key=key, text=f"after kill {uuid.uuid4().hex}") - assert isinstance(after["id"], str) and after["id"].startswith("chatcmpl-"), after - assert client.get("/health/liveliness").status_code == 200 + with owned.replacing_worker(victim) as kill: + results: Final = _burst(client, model, key, scenario.key(), kill) + dropped: Final = tuple(result for result in results if result.status is None) + assert len(dropped) < len(results), results + crafted: Final = tuple( + result for result in results if result.label.startswith("crafted") and result.status + ) + assert crafted and all( + result.status == 400 and result.seconds < FAST_SECONDS for result in crafted + ), crafted + late: Final = tuple(result for result in results if result.label.startswith("late")) + assert all(result.status == 200 for result in late), late + _assert_rows_land_exactly_once(results, "late") + survivor: Final = tuple(pid for pid in owned.worker_pids() if pid != victim) + assert survivor, "no worker left serving" + after: Final = owned.gateway.chat(model, key=key, text=f"after kill {uuid.uuid4().hex}") + assert isinstance(after["id"], str) and after["id"].startswith("chatcmpl-"), after + assert client.get("/health/liveliness").status_code == 200 def test_upstream_returning_503_mid_burst_logs_every_failure_with_its_own_cluster_key( diff --git a/tests/integration/spend/test_shutdown_flush.py b/tests/integration/spend/test_shutdown_flush.py index 5441f70d59e..e96511511d9 100644 --- a/tests/integration/spend/test_shutdown_flush.py +++ b/tests/integration/spend/test_shutdown_flush.py @@ -1,6 +1,5 @@ import json import os -import signal import threading import uuid from collections.abc import Callable, Iterator @@ -163,7 +162,7 @@ class _Shutdown: def terminate_once(self, blocked: Callable[[], bool], release: Callable[[], None]) -> None: eventually(blocked, lambda state: state, seconds=60) - self.owned.process.send_signal(signal.SIGTERM) + self.owned.begin_shutdown() eventually(lambda: self.logged(CANCEL_LOG_LINE), lambda seen: seen, seconds=60) release() self.owned.process.wait(timeout=120) diff --git a/tests/integration/spend/test_spend_log_tool_payload_content.py b/tests/integration/spend/test_spend_log_tool_payload_content.py index 22ed6c74756..672c0a38407 100644 --- a/tests/integration/spend/test_spend_log_tool_payload_content.py +++ b/tests/integration/spend/test_spend_log_tool_payload_content.py @@ -325,7 +325,7 @@ def test_stored_chat_response_keeps_logprob_tokens(gateway: Gateway, tmp_path: P with ( wire_server(_answering_model_listing(respond)) as wire, - owned_proxy(gateway, tmp_path, {}, config=_prompt_storage_config(tmp_path), workers=2) as isolated, + owned_proxy(gateway, tmp_path, {}, config=_prompt_storage_config(tmp_path)) as isolated, isolated.scenario() as scenario, ): model: Final = scenario.model( @@ -418,7 +418,7 @@ def test_stored_messages_keep_tool_use_input(gateway: Gateway, tmp_path: Path) - with ( wire_server(_answering_model_listing(respond)) as wire, - owned_proxy(gateway, tmp_path, {}, config=_prompt_storage_config(tmp_path), workers=2) as isolated, + owned_proxy(gateway, tmp_path, {}, config=_prompt_storage_config(tmp_path)) as isolated, isolated.scenario() as scenario, ): model: Final = scenario.model( @@ -650,7 +650,7 @@ def test_chat_sdk_keeps_logprob_tokens(gateway: Gateway, tmp_path: Path, client_ with ( wire_server(_answering_model_listing(respond)) as wire, - owned_proxy(gateway, tmp_path, {}, config=_prompt_storage_config(tmp_path), workers=2) as isolated, + owned_proxy(gateway, tmp_path, {}, config=_prompt_storage_config(tmp_path)) as isolated, isolated.scenario() as scenario, ): model: Final = scenario.model( @@ -702,7 +702,7 @@ def test_chat_sdk_stream_include_usage_masks_request_fields(gateway: Gateway, tm with ( wire_server(_answering_model_listing(respond)) as wire, - owned_proxy(gateway, tmp_path, {}, config=_prompt_storage_config(tmp_path), workers=2) as isolated, + owned_proxy(gateway, tmp_path, {}, config=_prompt_storage_config(tmp_path)) as isolated, isolated.scenario() as scenario, ): model: Final = scenario.model( @@ -747,7 +747,7 @@ def test_chat_history_keeps_string_tool_arguments_and_tool_content(gateway: Gate with ( wire_server(_answering_model_listing(respond)) as wire, - owned_proxy(gateway, tmp_path, {}, config=_prompt_storage_config(tmp_path), workers=2) as isolated, + owned_proxy(gateway, tmp_path, {}, config=_prompt_storage_config(tmp_path)) as isolated, isolated.scenario() as scenario, ): model: Final = scenario.model(model="openai/gpt-4o-mini", api_base=wire.url, api_key="synthetic-openai-key") @@ -833,7 +833,7 @@ def test_messages_sdk_keeps_tool_use_input(gateway: Gateway, tmp_path: Path, cli with ( wire_server(_answering_model_listing(respond)) as wire, - owned_proxy(gateway, tmp_path, {}, config=_prompt_storage_config(tmp_path), workers=2) as isolated, + owned_proxy(gateway, tmp_path, {}, config=_prompt_storage_config(tmp_path)) as isolated, isolated.scenario() as scenario, ): model: Final = scenario.model(model=ANTHROPIC_MODEL, api_base=wire.url, api_key="synthetic-anthropic-key") @@ -879,7 +879,7 @@ def test_messages_stream_keeps_tool_use_input(gateway: Gateway, tmp_path: Path) with ( wire_server(_answering_model_listing(respond)) as wire, - owned_proxy(gateway, tmp_path, {}, config=_prompt_storage_config(tmp_path), workers=2) as isolated, + owned_proxy(gateway, tmp_path, {}, config=_prompt_storage_config(tmp_path)) as isolated, isolated.scenario() as scenario, ): model: Final = scenario.model(model=ANTHROPIC_MODEL, api_base=wire.url, api_key="synthetic-anthropic-key") @@ -925,7 +925,7 @@ def test_responses_stream_keeps_function_call_arguments(gateway: Gateway, tmp_pa with ( wire_server(_answering_model_listing(respond)) as wire, - owned_proxy(gateway, tmp_path, {}, config=_prompt_storage_config(tmp_path), workers=2) as isolated, + owned_proxy(gateway, tmp_path, {}, config=_prompt_storage_config(tmp_path)) as isolated, isolated.scenario() as scenario, ): model: Final = scenario.model(model=ANTHROPIC_MODEL, api_base=wire.url, api_key="synthetic-anthropic-key") @@ -1006,7 +1006,7 @@ def test_native_responses_keeps_logprob_tokens(gateway: Gateway, tmp_path: Path) with ( wire_server(_answering_model_listing(respond)) as wire, - owned_proxy(gateway, tmp_path, {}, config=_prompt_storage_config(tmp_path), workers=2) as isolated, + owned_proxy(gateway, tmp_path, {}, config=_prompt_storage_config(tmp_path)) as isolated, isolated.scenario() as scenario, ): model: Final = scenario.model(model="openai/gpt-4o-mini", api_base=wire.url, api_key="synthetic-openai-key") @@ -1052,7 +1052,7 @@ def test_malformed_tool_blocks_keep_only_recognized_content(gateway: Gateway, tm with ( wire_server(_answering_model_listing(respond)) as wire, - owned_proxy(gateway, tmp_path, {}, config=_prompt_storage_config(tmp_path), workers=2) as isolated, + owned_proxy(gateway, tmp_path, {}, config=_prompt_storage_config(tmp_path)) as isolated, isolated.scenario() as scenario, ): model: Final = scenario.model(model="openai/gpt-4o-mini", api_base=wire.url, api_key="synthetic-openai-key") @@ -1097,7 +1097,7 @@ def test_messages_tool_input_handles_mixed_values_and_truncation(gateway: Gatewa with ( wire_server(_answering_model_listing(respond)) as wire, - owned_proxy(gateway, tmp_path, {}, config=_prompt_storage_config(tmp_path), workers=2) as isolated, + owned_proxy(gateway, tmp_path, {}, config=_prompt_storage_config(tmp_path)) as isolated, isolated.scenario() as scenario, ): model: Final = scenario.model(model=ANTHROPIC_MODEL, api_base=wire.url, api_key="synthetic-anthropic-key") @@ -1143,7 +1143,7 @@ def test_messages_without_auth_create_no_spend_row(gateway: Gateway, tmp_path: P with ( wire_server(respond) as wire, - owned_proxy(gateway, tmp_path, {}, config=_prompt_storage_config(tmp_path), workers=2) as isolated, + owned_proxy(gateway, tmp_path, {}, config=_prompt_storage_config(tmp_path)) as isolated, ): response: Final = isolated.client.post( "/v1/messages", @@ -1182,7 +1182,7 @@ def test_messages_upstream_error_keeps_tool_input(gateway: Gateway, tmp_path: Pa with ( wire_server(_answering_model_listing(respond)) as wire, - owned_proxy(gateway, tmp_path, {}, config=_prompt_storage_config(tmp_path), workers=2) as isolated, + owned_proxy(gateway, tmp_path, {}, config=_prompt_storage_config(tmp_path)) as isolated, isolated.scenario() as scenario, ): model: Final = scenario.model(model=ANTHROPIC_MODEL, api_base=wire.url, api_key="synthetic-anthropic-key") @@ -1235,7 +1235,6 @@ def test_store_prompts_off_keeps_chat_and_messages_representation_equal(gateway: tmp_path, {}, config=_prompt_storage_config(tmp_path, store_prompts=False), - workers=2, ) as isolated, isolated.scenario() as scenario, ): @@ -1283,7 +1282,7 @@ def test_identical_messages_requests_have_distinct_spend_rows(gateway: Gateway, with ( wire_server(_answering_model_listing(respond)) as wire, - owned_proxy(gateway, tmp_path, {}, config=_prompt_storage_config(tmp_path), workers=2) as isolated, + owned_proxy(gateway, tmp_path, {}, config=_prompt_storage_config(tmp_path)) as isolated, isolated.scenario() as scenario, ): model: Final = scenario.model(model=ANTHROPIC_MODEL, api_base=wire.url, api_key="synthetic-anthropic-key") diff --git a/tests/integration/spend/test_vertex_image_regional_uplift.py b/tests/integration/spend/test_vertex_image_regional_uplift.py index a06c8842a45..9aedb81c3e0 100644 --- a/tests/integration/spend/test_vertex_image_regional_uplift.py +++ b/tests/integration/spend/test_vertex_image_regional_uplift.py @@ -3,7 +3,6 @@ import base64 import json import os import re -import signal import threading import time import uuid @@ -625,27 +624,28 @@ async def test_worker_sigkill_mid_burst_leaves_the_sibling_billing_the_uplift(ga held_by: Final = {pid: _open_upstream_connections(pid, wire.url) for pid in workers} assert sum(held_by.values()) == len(call_ids), held_by victim_pid, survivor_pid = sorted(workers, key=held_by.__getitem__) - psutil.Process(victim_pid).send_signal(signal.SIGKILL) - release.set() - served: Final = await burst - succeeded: Final = tuple(item.call_id for item in served if item.status == 200) - _artifact( - tmp_path, - "worker-sigkill.json", - { - "held_by_killed_worker": held_by[victim_pid], - "held_by_surviving_worker": held_by[survivor_pid], - "served": len(served), - "succeeded": len(succeeded), - "transport_errors": len(call_ids) - len(served), - }, - ) - assert len(succeeded) == held_by[survivor_pid], ( - held_by, - [(item.status, item.text[:80]) for item in served], - ) - follow_up_id, follow_up = _generate(owned.gateway, _CONFIG_MODEL) - _assert_images(follow_up, follow_up_id, _REGIONAL_TOKEN_COST) + with owned.replacing_worker(victim_pid) as kill: + await asyncio.to_thread(kill) + release.set() + served: Final = await burst + succeeded: Final = tuple(item.call_id for item in served if item.status == 200) + _artifact( + tmp_path, + "worker-sigkill.json", + { + "held_by_killed_worker": held_by[victim_pid], + "held_by_surviving_worker": held_by[survivor_pid], + "served": len(served), + "succeeded": len(succeeded), + "transport_errors": len(call_ids) - len(served), + }, + ) + assert len(succeeded) == held_by[survivor_pid], ( + held_by, + [(item.status, item.text[:80]) for item in served], + ) + follow_up_id, follow_up = _generate(owned.gateway, _CONFIG_MODEL) + _assert_images(follow_up, follow_up_id, _REGIONAL_TOKEN_COST) rows: Final = await asyncio.to_thread(_spend_rows, succeeded, seconds=120) assert rows == {call_id: pytest.approx(_REGIONAL_TOKEN_COST) for call_id in succeeded} await asyncio.to_thread( @@ -670,7 +670,7 @@ async def test_proxy_restart_mid_burst_bills_each_landed_call_once(gateway: Gate ) async with _gated_burst(base_url, gateway.key, call_ids, release) as burst: await asyncio.to_thread(eventually, held.qsize, lambda size: size == len(call_ids), 120) - owned.process.terminate() + owned.begin_shutdown() release.set() served: Final = await burst await asyncio.to_thread(owned.process.wait, 240) diff --git a/tests/integration/streaming/test_decision_model_buffering.py b/tests/integration/streaming/test_decision_model_buffering.py new file mode 100644 index 00000000000..a2594fd2b9b --- /dev/null +++ b/tests/integration/streaming/test_decision_model_buffering.py @@ -0,0 +1,239 @@ +import json +import uuid +from pathlib import Path +from typing import Final, Literal + +import pytest +import yaml +from integration._support.client import JSON_OBJECT, Gateway, list_value, object_value, string_value +from integration._support.process import owned_proxy +from integration._support.upstream import delete_scenario, register_scenario +from integration.cost_calculation.cost_tracking_case import JsonResponse, SseResponse +from pydantic import JsonValue + + +@pytest.mark.parametrize("endpoint", ["chat/completions", "messages", "responses"]) +@pytest.mark.parametrize( + ("buffering", "probability", "judge_status", "fallback"), + [ + (True, 1.0, 200, "fail_closed"), + (True, 0.0, 200, "fail_closed"), + (False, 1.0, 200, "fail_closed"), + (None, 1.0, 200, "fail_closed"), + (True, 0.0, 500, "fail_closed"), + (True, 0.0, 500, "fail_open"), + ], + ids=["buffered-block", "buffered-allow", "unbuffered-block", "default-block", "outage-closed", "outage-open"], +) +def test_decision_model_streaming_buffering( + gateway: Gateway, + tmp_path: Path, + buffering: bool | None, + probability: float, + endpoint: str, + judge_status: int, + fallback: Literal["fail_closed", "fail_open"], +) -> None: + identity: Final = uuid.uuid4().hex + frames: Final = _frames(identity, endpoint) + chat: Final = register_scenario( + identity, + SseResponse(content_type="text/event-stream", frames=frames), + control_url=gateway.upstream_url, + ) + try: + decider: Final = register_scenario( + identity + "-decider", + JsonResponse( + content_type="application/json", + status=judge_status, + body={ + "model": "scripted-decider", + "answers": [{"type": "predicate", "name": "bananas", "probability": probability}], + "usage": {"input_tokens": 1, "output_tokens": 1, "total_tokens": 2}, + }, + ), + control_url=gateway.upstream_url, + ) + try: + config: Final = { + **JSON_OBJECT.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())), + "model_list": [ + { + "model_name": "buffering-chat", + "litellm_params": { + "model": "anthropic/scripted-chat" if endpoint == "messages" else "openai/scripted-chat", + "api_key": "synthetic", + "api_base": chat.api_base(), + }, + }, + { + "model_name": "buffering-decider", + "litellm_params": { + "model": "openai/scripted-decider", + "api_key": "synthetic", + "api_base": decider.api_base(), + }, + }, + ], + "guardrails": [ + { + "guardrail_name": "banana-policy", + "litellm_params": { + "guardrail": "decision_model", + "mode": "post_call", + "default_on": True, + "decision_model": "buffering-decider", + "unreachable_fallback": fallback, + **({"streaming_buffer_until_moderated": buffering} if buffering is not None else {}), + "checks": [ + { + "name": "bananas", + "instructions": "Does the text mention bananas?", + "action": "block", + "threshold": 0.5, + } + ], + }, + } + ], + "general_settings": {"master_key": gateway.key}, + } + config_path: Final = tmp_path / "decision-buffering.yaml" + config_path.write_text(yaml.safe_dump(config)) + with owned_proxy(gateway, tmp_path, {}, config=config_path, workers=2) as candidate: + response: Final = candidate.request( + "POST", + f"/v1/{endpoint}", + { + "model": "buffering-chat", + **( + {"input": "Tell me about bananas"} + if endpoint == "responses" + else {"messages": [{"role": "user", "content": "Tell me about bananas"}]} + ), + **({"max_tokens": 100} if endpoint == "messages" else {}), + "stream": True, + "cache": {"no-cache": True, "no-store": True}, + }, + ) + events: Final = tuple( + JSON_OBJECT.validate_json(line.removeprefix("data: ")) + for line in response.text.splitlines() + if line.startswith("data: {") + ) + content: Final = "".join(_event_text(event, endpoint) for event in events) + if judge_status != 200 and fallback == "fail_closed": + assert "Decision model guardrail failed to respond" in response.text, response.text + assert content == "", response.text + elif probability == 0: + assert content == "Bananas are fruit.", response.text + assert '"error":' not in response.text, response.text + else: + assert "Violated decision model guardrail policy" in response.text, response.text + if buffering: + assert content == "", response.text + else: + assert content and "Bananas are fruit.".startswith(content), response.text + finally: + delete_scenario(decider) + finally: + delete_scenario(chat) + + +def _event_text(event: dict[str, JsonValue], endpoint: str) -> str: + if endpoint == "responses": + return string_value(event.get("delta", "")) if event.get("type") == "response.output_text.delta" else "" + if endpoint == "messages": + return string_value(object_value(event.get("delta", {})).get("text", "")) + choices: Final = list_value(event.get("choices", [{}])) + return string_value(object_value(object_value(choices[0]).get("delta", {})).get("content", "")) + + +def _frames(identity: str, endpoint: str) -> tuple[str, ...]: + pieces: Final = ("Ban", "anas", " are", " fruit", ".") + if endpoint == "messages": + events: Final = ( + { + "type": "message_start", + "message": { + "id": identity, + "type": "message", + "role": "assistant", + "content": [], + "model": "scripted-chat", + "stop_reason": None, + "usage": {"input_tokens": 1}, + }, + }, + {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}, + *( + {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": text}} + for text in pieces + ), + {"type": "content_block_stop", "index": 0}, + {"type": "message_delta", "delta": {"stop_reason": "end_turn"}, "usage": {"output_tokens": 5}}, + {"type": "message_stop"}, + ) + return tuple(f"event: {event['type']}\ndata: {json.dumps(event)}" for event in events) + if endpoint == "responses": + message: Final = { + "id": "msg_" + identity, + "type": "message", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": "".join(pieces), "annotations": []}], + } + response: Final = { + "id": "resp_" + identity, + "object": "response", + "created_at": 1, + "status": "completed", + "model": "scripted-chat", + "output": [message], + "usage": {"input_tokens": 1, "output_tokens": 5, "total_tokens": 6}, + } + response_events: Final = ( + { + "type": "response.created", + "sequence_number": 0, + "response": {**response, "status": "in_progress", "output": []}, + }, + { + "type": "response.output_item.added", + "sequence_number": 1, + "output_index": 0, + "item": {**message, "content": []}, + }, + *( + { + "type": "response.output_text.delta", + "sequence_number": index + 2, + "item_id": message["id"], + "output_index": 0, + "content_index": 0, + "delta": text, + } + for index, text in enumerate(pieces) + ), + {"type": "response.output_item.done", "sequence_number": 7, "output_index": 0, "item": message}, + {"type": "response.completed", "sequence_number": 8, "response": response}, + ) + return tuple(f"event: {event['type']}\ndata: {json.dumps(event)}" for event in response_events) + return ( + *( + f"data: { + json.dumps( + { + 'id': identity, + 'object': 'chat.completion.chunk', + 'created': 1, + 'model': 'scripted-chat', + 'choices': [{'index': 0, 'delta': {'content': text}, 'finish_reason': None}], + } + ) + }" + for text in pieces + ), + "data: [DONE]", + ) diff --git a/tests/integration/streaming/test_decision_model_pipeline_buffering.py b/tests/integration/streaming/test_decision_model_pipeline_buffering.py new file mode 100644 index 00000000000..f76e7d76d4a --- /dev/null +++ b/tests/integration/streaming/test_decision_model_pipeline_buffering.py @@ -0,0 +1,244 @@ +from __future__ import annotations + +import json +import threading +import uuid +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path +from typing import Final + +import pytest +import yaml +from integration._support.client import JSON_OBJECT, Gateway, list_value, object_value +from integration._support.process import owned_proxy +from integration._support.upstream import delete_scenario, register_scenario +from integration._support.wire import Reply, Request, wire_server + +from tests.integration.cost_calculation.cost_tracking_case import SseResponse + + +def _frame(identity: str, content: str, finish: str | None = None) -> str: + return "data: " + json.dumps( + { + "id": identity, + "object": "chat.completion.chunk", + "created": 1, + "model": "gpt-4o-mini", + "choices": [ + { + "index": 0, + "delta": {"content": content}, + "finish_reason": finish, + } + ], + } + ) + + +def _content( + line: str, + content_received: threading.Event, + done_received: threading.Event, +) -> str: + if not line.startswith("data:"): + return "" + payload: Final = line.removeprefix("data:").strip() + if payload == "[DONE]": + done_received.set() + return "" + event: Final = JSON_OBJECT.validate_json(payload) + assert "error" not in event, event + pieces: Final = tuple( + object_value(object_value(choice)["delta"]).get("content", "") + for choice in list_value(event.get("choices", [])) + ) + assert all(isinstance(piece, str) for piece in pieces), event + text: Final = "".join(piece for piece in pieces if isinstance(piece, str)) + if text: + content_received.set() + return text + + +def _read_stream( + candidate: Gateway, + model: str, + content_received: threading.Event, + done_received: threading.Event, +) -> str: + with candidate.client.stream( + "POST", + "/v1/chat/completions", + headers={"Authorization": f"Bearer {candidate.key}"}, + json={ + "model": model, + "messages": [{"role": "user", "content": "synthetic control"}], + "stream": True, + }, + timeout=30, + ) as response: + assert response.status_code == 200, response.read() + return "".join(_content(line, content_received, done_received) for line in response.iter_lines()) + + +@pytest.mark.parametrize("buffering", [None, True, False], ids=["omitted", "enabled", "disabled"]) +def test_detect_only_decision_pipeline_respects_buffering( + gateway: Gateway, + tmp_path: Path, + buffering: bool | None, +) -> None: + identity: Final = f"decision-buffer-{uuid.uuid4().hex}" + chat_model: Final = f"{identity}-chat" + judge_model: Final = f"{identity}-judge" + secret: Final = f"LIT9453_SECRET_{identity}" + judge_entered: Final = threading.Event() + release_judge: Final = threading.Event() + content_received: Final = threading.Event() + done_received: Final = threading.Event() + + def judge(request: Request) -> Reply: + assert request.method == "POST", request + assert request.target == "/v1/decisions", request + body: Final = JSON_OBJECT.validate_json(request.body) + assert secret in json.dumps(body.get("input")), body + judge_entered.set() + assert release_judge.wait(timeout=20), "Judge gate was never released" + return Reply( + body=json.dumps( + { + "model": "gpt-6-luna", + "answers": [ + { + "type": "predicate", + "name": "unsafe_output", + "probability": 0.99, + } + ], + "usage": { + "input_tokens": 10, + "output_tokens": 1, + "total_tokens": 11, + }, + } + ).encode() + ) + + with gateway.scenario() as scenario: + upstream: Final = register_scenario( + identity, + SseResponse( + content_type="text/event-stream", + frames=( + _frame(identity, secret), + _frame(identity, "", "stop"), + "data: [DONE]", + ), + frame_delay_ms=20, + ), + control_url=gateway.upstream_url, + ) + scenario.cleanups.callback(delete_scenario, upstream) + + with wire_server(judge) as judge_server: + base: Final = JSON_OBJECT.validate_python( + yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + ) + config: Final = { + **base, + "model_list": [ + { + "model_name": chat_model, + "litellm_params": { + "model": "openai/gpt-4o-mini", + "api_base": f"{upstream.api_base()}/v1", + "api_key": "synthetic-chat-key", + }, + }, + { + "model_name": judge_model, + "litellm_params": { + "model": "openai/gpt-6-luna", + "api_base": f"{judge_server.url}/v1", + "api_key": "synthetic-judge-key", + }, + }, + ], + "guardrails": [ + { + "guardrail_name": identity, + "litellm_params": { + "guardrail": "decision_model", + "mode": "post_call", + "default_on": False, + "decision_model": judge_model, + "timeout": 25, + "streaming_end_of_stream_only": True, + **({"streaming_buffer_until_moderated": buffering} if buffering is not None else {}), + "checks": [ + { + "name": "unsafe_output", + "instructions": ("Does the response contain the synthetic secret?"), + "action": "block", + "threshold": 0.5, + } + ], + }, + } + ], + "policies": { + identity: { + "guardrails": {"add": [identity]}, + "pipeline": { + "mode": "post_call", + "steps": [ + { + "guardrail": identity, + "on_pass": "allow", + "on_fail": "next", + "on_error": "next", + } + ], + }, + } + }, + "policy_attachments": [{"policy": identity, "models": [chat_model]}], + } + config_path: Final = tmp_path / "decision-pipeline-buffering.yaml" + config_path.write_text(yaml.safe_dump(config)) + + with ( + owned_proxy( + gateway, + tmp_path, + {}, + config=config_path, + workers=2, + ) as candidate, + ThreadPoolExecutor(max_workers=1) as reader, + ): + result: Final = reader.submit( + _read_stream, + candidate, + chat_model, + content_received, + done_received, + ) + try: + assert judge_entered.wait(timeout=15), "The completed chat response never reached the judge" + if buffering is False: + assert content_received.wait(timeout=5), ( + "Pipeline did not release explicitly unbuffered content" + ) + else: + assert not content_received.wait(timeout=1), ( + "Pipeline released content while its judge was held" + ) + assert not result.done(), "Client stream ended before the judge verdict" + finally: + release_judge.set() + + assert result.result(timeout=15) == secret + assert done_received.is_set(), "Missing terminal SSE [DONE]" + requests: Final = judge_server.drain() + assert len(requests) == 1, requests + observed: Final = JSON_OBJECT.validate_json(requests[0].body) + assert secret in json.dumps(observed.get("input")), observed diff --git a/tests/integration/streaming/test_interception_converted_stream_options.py b/tests/integration/streaming/test_interception_converted_stream_options.py index 46a4211b121..74fad7deaf1 100644 --- a/tests/integration/streaming/test_interception_converted_stream_options.py +++ b/tests/integration/streaming/test_interception_converted_stream_options.py @@ -1,6 +1,5 @@ import asyncio import json -import signal import threading import uuid from collections.abc import Callable, Generator, Iterator, Mapping, Sequence @@ -882,18 +881,21 @@ async def test_worker_sigkill_mid_burst_leaves_the_sibling_restreaming_with_usag held_by: Final = MappingProxyType({pid: _open_upstream_connections(pid, owned.wire.url) for pid in workers}) assert sum(held_by.values()) == 20, held_by victim_pid, survivor_pid = sorted(workers, key=held_by.__getitem__) - victim: Final = psutil.Process(victim_pid) - victim.suspend() - victim.send_signal(signal.SIGKILL) - release.set() - served: Final = await burst - assert held_by[survivor_pid] >= 10, held_by - assert len(served) == held_by[survivor_pid], (held_by, len(served)) - identities: Final = frozenset( - _assert_restreamed(item.stream, iv.answer(_BURST_DEPLOYMENT), usage=iv.FOLLOWUP_USAGE) for item in served - ) - (follow_up,) = await _burst(owned.base_url, owned.gateway.key, (_chat_request(deployment),)) - recovered: Final = _assert_restreamed(follow_up.stream, iv.answer(_BURST_DEPLOYMENT), usage=iv.FOLLOWUP_USAGE) + with owned.proxy.replacing_worker(victim_pid) as kill: + kill.freeze() + await asyncio.to_thread(kill) + release.set() + served: Final = await burst + assert held_by[survivor_pid] >= 10, held_by + assert len(served) == held_by[survivor_pid], (held_by, len(served)) + identities: Final = frozenset( + _assert_restreamed(item.stream, iv.answer(_BURST_DEPLOYMENT), usage=iv.FOLLOWUP_USAGE) + for item in served + ) + (follow_up,) = await _burst(owned.base_url, owned.gateway.key, (_chat_request(deployment),)) + recovered: Final = _assert_restreamed( + follow_up.stream, iv.answer(_BURST_DEPLOYMENT), usage=iv.FOLLOWUP_USAGE + ) rows: Final = _spend_rows(_BURST_MODEL, len(identities) + 1) landed: Final = tuple(string_value(row["request_id"]) for row in rows if row["status"] == "success") assert sorted(landed) == sorted((*identities, recovered)), rows diff --git a/tests/integration/streaming/test_responses_bridge_stream_chaos.py b/tests/integration/streaming/test_responses_bridge_stream_chaos.py index 6b2ea47a849..745be2e34b1 100644 --- a/tests/integration/streaming/test_responses_bridge_stream_chaos.py +++ b/tests/integration/streaming/test_responses_bridge_stream_chaos.py @@ -1,6 +1,5 @@ import asyncio import re -import signal import socket import threading import uuid @@ -322,18 +321,18 @@ async def test_worker_sigkill_mid_burst_leaves_the_sibling_answering_the_bridged held_by: Final = MappingProxyType({pid: _open_upstream_connections(pid, wire.url) for pid in workers}) assert sum(held_by.values()) == 20, held_by victim_pid, survivor_pid = sorted(workers, key=held_by.__getitem__) - victim: Final = psutil.Process(victim_pid) - victim.suspend() - victim.send_signal(signal.SIGKILL) - held.release.set() - served: Final = await burst - assert held_by[survivor_pid] >= 10, held_by - assert len(served) == held_by[survivor_pid], (held_by, len(served)) - for item in served: - _assert_answered_in_its_own_shape(item) - follow_up: Final = _Call("chat_healthy", uuid.uuid4().hex) - (answered,) = await _burst(rig.gateway, (follow_up,)) - _assert_answered_in_its_own_shape(answered) + with rig.proxy.replacing_worker(victim_pid) as kill: + kill.freeze() + await asyncio.to_thread(kill) + held.release.set() + served: Final = await burst + assert held_by[survivor_pid] >= 10, held_by + assert len(served) == held_by[survivor_pid], (held_by, len(served)) + for item in served: + _assert_answered_in_its_own_shape(item) + follow_up: Final = _Call("chat_healthy", uuid.uuid4().hex) + (answered,) = await _burst(rig.gateway, (follow_up,)) + _assert_answered_in_its_own_shape(answered) _assert_forwarded(wire.drain(), (*calls, follow_up)) _assert_each_lands_once((*served, answered)) @@ -349,7 +348,7 @@ async def test_proxy_sigterm_mid_burst_drains_the_spend_log_queue_and_the_restar with wire_server(held.respond, port=port) as wire: with owned_proxy_process(gateway, tmp_path, {}, config=config, workers=2) as owned: burst: Final = await _held_burst(owned.gateway, calls, held) - owned.process.terminate() + owned.begin_shutdown() held.release.set() served: Final = await burst await asyncio.to_thread( diff --git a/tests/integration/test_oci_proxy_integration.py b/tests/integration/test_oci_proxy_integration.py index 06ed4486941..d059a55b1f2 100644 --- a/tests/integration/test_oci_proxy_integration.py +++ b/tests/integration/test_oci_proxy_integration.py @@ -42,6 +42,7 @@ from typing import Iterator import httpx import pytest from tests._master_key import MASTER_KEY +from tests.integration._support.process import stop_child # --------------------------------------------------------------------------- # Skip gate @@ -147,12 +148,7 @@ def _serve(config_path: str) -> Iterator[str]: _wait_for_health(base_url, proc, time.monotonic() + STARTUP_TIMEOUT_S) yield base_url finally: - proc.terminate() - try: - proc.wait(timeout=10) - except subprocess.TimeoutExpired: - proc.kill() - proc.wait(timeout=5) + stop_child(proc, grace_seconds=10) @pytest.fixture(scope="module") diff --git a/tests/integration/translation/chat_completions/bases/gemini.py b/tests/integration/translation/chat_completions/bases/gemini.py index e9fbcb6ccda..1ffde32c3be 100644 --- a/tests/integration/translation/chat_completions/bases/gemini.py +++ b/tests/integration/translation/chat_completions/bases/gemini.py @@ -21,7 +21,7 @@ GEMINI_3_5_FLASH_TEST_CASE: Final = TranslationTestCase( expected_provider_headers={"content-type": "application/json", "x-goog-api-key": "synthetic-gemini-key"}, expected_provider_request={ "contents": [{"parts": [{"text": "Say hello."}], "role": "user"}], - "generationConfig": {"max_output_tokens": 1024, "temperature": 1.0}, + "generationConfig": {"max_output_tokens": 1024}, "system_instruction": {"parts": [{"text": "You are a terse assistant."}]}, }, mock_provider_response={ @@ -97,7 +97,7 @@ GEMINI_3_8_FLASH_TEST_CASE: Final = TranslationTestCase( expected_provider_headers={"content-type": "application/json", "x-goog-api-key": "synthetic-gemini-key"}, expected_provider_request={ "contents": [{"parts": [{"text": "Say hello."}], "role": "user"}], - "generationConfig": {"max_output_tokens": 1024, "temperature": 1.0}, + "generationConfig": {"max_output_tokens": 1024}, "system_instruction": {"parts": [{"text": "You are a terse assistant."}]}, }, mock_provider_response={ @@ -173,7 +173,7 @@ GEMINI_3_1_PRO_PREVIEW_TEST_CASE: Final = TranslationTestCase( expected_provider_headers={"content-type": "application/json", "x-goog-api-key": "synthetic-gemini-key"}, expected_provider_request={ "contents": [{"parts": [{"text": "Say hello."}], "role": "user"}], - "generationConfig": {"max_output_tokens": 1024, "temperature": 1.0}, + "generationConfig": {"max_output_tokens": 1024}, "system_instruction": {"parts": [{"text": "You are a terse assistant."}]}, }, mock_provider_response={ diff --git a/tests/integration/translation/chat_completions/bases/vertex_ai.py b/tests/integration/translation/chat_completions/bases/vertex_ai.py index 49018605bd1..236cefb5c72 100644 --- a/tests/integration/translation/chat_completions/bases/vertex_ai.py +++ b/tests/integration/translation/chat_completions/bases/vertex_ai.py @@ -498,7 +498,7 @@ GEMINI_3_5_FLASH_TEST_CASE: Final = TranslationTestCase( expected_provider_request={ "contents": [{"role": "user", "parts": [{"text": "Say hello."}]}], "system_instruction": {"parts": [{"text": "You are a terse assistant."}]}, - "generationConfig": {"max_output_tokens": 1024, "temperature": 1.0}, + "generationConfig": {"max_output_tokens": 1024}, }, mock_provider_response={ "candidates": [ @@ -574,7 +574,7 @@ GEMINI_3_8_FLASH_TEST_CASE: Final = TranslationTestCase( expected_provider_request={ "contents": [{"role": "user", "parts": [{"text": "Say hello."}]}], "system_instruction": {"parts": [{"text": "You are a terse assistant."}]}, - "generationConfig": {"max_output_tokens": 1024, "temperature": 1.0}, + "generationConfig": {"max_output_tokens": 1024}, }, mock_provider_response={ "candidates": [ @@ -650,7 +650,7 @@ GEMINI_3_1_PRO_PREVIEW_TEST_CASE: Final = TranslationTestCase( expected_provider_request={ "contents": [{"role": "user", "parts": [{"text": "Say hello."}]}], "system_instruction": {"parts": [{"text": "You are a terse assistant."}]}, - "generationConfig": {"max_output_tokens": 1024, "temperature": 1.0}, + "generationConfig": {"max_output_tokens": 1024}, }, mock_provider_response={ "candidates": [ diff --git a/tests/integration/translation/chat_completions/sampling/__init__.py b/tests/integration/translation/chat_completions/sampling/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/integration/translation/chat_completions/sampling/test_chat_completions_sampling_gemini.py b/tests/integration/translation/chat_completions/sampling/test_chat_completions_sampling_gemini.py new file mode 100644 index 00000000000..1707fc71238 --- /dev/null +++ b/tests/integration/translation/chat_completions/sampling/test_chat_completions_sampling_gemini.py @@ -0,0 +1,76 @@ +from dataclasses import replace +from typing import Final + +import pytest +from integration._support.client import Gateway, object_value +from integration._support.provider import SharedProvider +from integration.translation.case import TranslationTestCase +from integration.translation.chat_completions.bases.gemini import GEMINI_3_8_FLASH_TEST_CASE +from integration.translation.chat_completions.bases.vertex_ai import ( + GEMINI_3_8_FLASH_TEST_CASE as VERTEX_AI_GEMINI_3_8_FLASH_TEST_CASE, +) +from integration.translation.runner import assert_translation + +GEMINI_GENERATION_CONFIG: Final = object_value(GEMINI_3_8_FLASH_TEST_CASE.expected_provider_request["generationConfig"]) +VERTEX_AI_GENERATION_CONFIG: Final = object_value( + VERTEX_AI_GEMINI_3_8_FLASH_TEST_CASE.expected_provider_request["generationConfig"] +) + +GEMINI_3_8_FLASH_TEMPERATURE_TEST_CASE: Final = replace( + GEMINI_3_8_FLASH_TEST_CASE, + scenario="temperature", + litellm_request={**GEMINI_3_8_FLASH_TEST_CASE.litellm_request, "temperature": 0.4}, + expected_provider_request={ + **GEMINI_3_8_FLASH_TEST_CASE.expected_provider_request, + "generationConfig": {**GEMINI_GENERATION_CONFIG, "temperature": 0.4}, + }, +) + +GEMINI_3_8_FLASH_ZERO_TEMPERATURE_TEST_CASE: Final = replace( + GEMINI_3_8_FLASH_TEST_CASE, + scenario="zero_temperature", + litellm_request={**GEMINI_3_8_FLASH_TEST_CASE.litellm_request, "temperature": 0}, + expected_provider_request={ + **GEMINI_3_8_FLASH_TEST_CASE.expected_provider_request, + "generationConfig": {**GEMINI_GENERATION_CONFIG, "temperature": 0}, + }, +) + +GEMINI_3_8_FLASH_REASONING_EFFORT_TEST_CASE: Final = replace( + GEMINI_3_8_FLASH_TEST_CASE, + scenario="reasoning_effort", + litellm_request={**GEMINI_3_8_FLASH_TEST_CASE.litellm_request, "reasoning_effort": "medium"}, + expected_provider_request={ + **GEMINI_3_8_FLASH_TEST_CASE.expected_provider_request, + "generationConfig": { + **GEMINI_GENERATION_CONFIG, + "thinkingConfig": {"thinkingLevel": "medium", "includeThoughts": True}, + }, + }, +) + +VERTEX_AI_GEMINI_3_8_FLASH_TEMPERATURE_TEST_CASE: Final = replace( + VERTEX_AI_GEMINI_3_8_FLASH_TEST_CASE, + scenario="temperature", + litellm_request={**VERTEX_AI_GEMINI_3_8_FLASH_TEST_CASE.litellm_request, "temperature": 0.4}, + expected_provider_request={ + **VERTEX_AI_GEMINI_3_8_FLASH_TEST_CASE.expected_provider_request, + "generationConfig": {**VERTEX_AI_GENERATION_CONFIG, "temperature": 0.4}, + }, +) + + +@pytest.mark.parametrize( + "case", + [ + GEMINI_3_8_FLASH_TEMPERATURE_TEST_CASE, + GEMINI_3_8_FLASH_ZERO_TEMPERATURE_TEST_CASE, + GEMINI_3_8_FLASH_REASONING_EFFORT_TEST_CASE, + VERTEX_AI_GEMINI_3_8_FLASH_TEMPERATURE_TEST_CASE, + ], + ids=lambda case: case.id, +) +def test_chat_completions_sampling_gemini( + case: TranslationTestCase, gateway: Gateway, provider: SharedProvider +) -> None: + assert_translation(case, gateway, provider) diff --git a/tests/integration/translation/decisions/bases/fireworks_ai.py b/tests/integration/translation/decisions/bases/fireworks_ai.py new file mode 100644 index 00000000000..672dc53d186 --- /dev/null +++ b/tests/integration/translation/decisions/bases/fireworks_ai.py @@ -0,0 +1,174 @@ +from typing import Final + +from integration.translation.case import TranslationTestCase + +"""Provider request and reply shape from Fireworks System One (POST /v1/systemone, private preview). Mock reply captured live on 2026-10-12. +""" +RUNE_TEST_CASE: Final = TranslationTestCase( + scenario="basic", + litellm_endpoint="/v1/decisions", + litellm_request={ + "model": "fireworks_ai/accounts/fireworks/routers/rune", + "input": "Ticket (billing): The export job hangs at 99% and never finishes", + "questions": [ + {"type": "predicate", "name": "defect", "instructions": "Is this a defect?"}, + { + "type": "choice", + "name": "severity", + "instructions": "How severe is it?", + "choices": [ + {"value": "low", "description": "cosmetic"}, + {"value": "high", "description": "blocks users"}, + ], + }, + { + "type": "score", + "name": "confidence", + "instructions": "How sure are you?", + "levels": [{"label": "unsure"}, {"label": "sure"}], + }, + ], + "cache": {"no-cache": True}, + }, + expected_provider_endpoint="/v1/systemone", + expected_provider_headers={"authorization": "Bearer synthetic-fireworks-key", "content-type": "application/json"}, + expected_provider_request={ + "model": "accounts/fireworks/routers/rune", + "state": "Ticket (billing): The export job hangs at 99% and never finishes", + "questions": { + "defect": {"type": "noul", "instructions": "Is this a defect?"}, + "severity": { + "type": "choice", + "instructions": "How severe is it?", + "criteria": {"low": "cosmetic", "high": "blocks users"}, + }, + "confidence": {"type": "score", "instructions": "How sure are you?", "criteria": ["unsure", "sure"]}, + }, + }, + mock_provider_response={ + "model": "accounts/fireworks/routers/rune", + "answers": { + "defect": {"type": "noul", "noul": 0.9902915235185259}, + "severity": { + "type": "choice", + "choice": "high", + "confidence": 0.6709670742068738, + "probabilities": {"high": 0.8354835371034369, "low": 0.1645164628965632}, + }, + "confidence": { + "type": "score", + "score": 0.7549149868676283, + "confidence": 0.5098299737352565, + "legend": {"0": "unsure", "1": "sure"}, + "probabilities": {"0": 0.24508501313237172, "1": 0.7549149868676283}, + }, + }, + "usage": {"input_tokens": 384, "cached_input_tokens": 0, "output_tokens": 0}, + }, + expected_litellm_response={ + "model": "accounts/fireworks/routers/rune", + "answers": [ + {"type": "predicate", "name": "defect", "probability": 0.9902915235185259}, + { + "type": "choice", + "name": "severity", + "choice": "high", + "probabilities": [ + {"value": "low", "probability": 0.1645164628965632}, + {"value": "high", "probability": 0.8354835371034369}, + ], + "confidence": 0.6709670742068738, + }, + { + "type": "score", + "name": "confidence", + "score": 0.7549149868676283, + "probabilities": [ + {"value": 0, "label": "unsure", "probability": 0.24508501313237172}, + {"value": 1, "label": "sure", "probability": 0.7549149868676283}, + ], + "confidence": 0.5098299737352565, + }, + ], + "usage": { + "input_tokens": 384, + "input_tokens_details": {"cached_tokens": 0, "cache_write_tokens": 0}, + "output_tokens": 0, + "output_tokens_details": {"reasoning_tokens": 0}, + "total_tokens": 384, + }, + }, +) +RUNE_SYSTEMONE_TEST_CASE: Final = TranslationTestCase( + scenario="systemone", + litellm_endpoint="/v1/systemone", + litellm_request={ + "model": "fireworks_ai/accounts/fireworks/routers/rune", + "state": "Ticket (billing): The export job hangs at 99% and never finishes", + "questions": { + "defect": {"type": "noul", "instructions": "Is this a defect?"}, + "severity": { + "type": "choice", + "instructions": "How severe is it?", + "criteria": {"low": "cosmetic", "high": "blocks users"}, + }, + "confidence": {"type": "score", "instructions": "How sure are you?", "criteria": ["unsure", "sure"]}, + }, + "cache": {"no-cache": True}, + }, + expected_provider_endpoint="/v1/systemone", + expected_provider_headers={"authorization": "Bearer synthetic-fireworks-key", "content-type": "application/json"}, + expected_provider_request={ + "model": "accounts/fireworks/routers/rune", + "state": "Ticket (billing): The export job hangs at 99% and never finishes", + "questions": { + "defect": {"type": "noul", "instructions": "Is this a defect?"}, + "severity": { + "type": "choice", + "instructions": "How severe is it?", + "criteria": {"low": "cosmetic", "high": "blocks users"}, + }, + "confidence": {"type": "score", "instructions": "How sure are you?", "criteria": ["unsure", "sure"]}, + }, + }, + mock_provider_response={ + "model": "accounts/fireworks/routers/rune", + "answers": { + "defect": {"type": "noul", "noul": 0.9902915235185259}, + "severity": { + "type": "choice", + "choice": "high", + "confidence": 0.6709670742068738, + "probabilities": {"high": 0.8354835371034369, "low": 0.1645164628965632}, + }, + "confidence": { + "type": "score", + "score": 0.7549149868676283, + "confidence": 0.5098299737352565, + "legend": {"0": "unsure", "1": "sure"}, + "probabilities": {"0": 0.24508501313237172, "1": 0.7549149868676283}, + }, + }, + "usage": {"input_tokens": 384, "cached_input_tokens": 0, "output_tokens": 0}, + }, + expected_litellm_response={ + "model": "accounts/fireworks/routers/rune", + "answers": { + "defect": {"type": "noul", "noul": 0.9902915235185259}, + "severity": { + "type": "choice", + "choice": "high", + "confidence": 0.6709670742068738, + "probabilities": {"high": 0.8354835371034369, "low": 0.1645164628965632}, + }, + "confidence": { + "type": "score", + "score": 0.7549149868676283, + "confidence": 0.5098299737352565, + "legend": {"0": "unsure", "1": "sure"}, + "probabilities": {"0": 0.24508501313237172, "1": 0.7549149868676283}, + }, + }, + "usage": {"input_tokens": 384, "cached_input_tokens": 0, "output_tokens": 0}, + }, +) diff --git a/tests/integration/translation/decisions/basic/test_decisions_basic_fireworks_ai.py b/tests/integration/translation/decisions/basic/test_decisions_basic_fireworks_ai.py new file mode 100644 index 00000000000..ae767820d2c --- /dev/null +++ b/tests/integration/translation/decisions/basic/test_decisions_basic_fireworks_ai.py @@ -0,0 +1,11 @@ +import pytest +from integration._support.client import Gateway +from integration._support.provider import SharedProvider +from integration.translation.case import TranslationTestCase +from integration.translation.decisions.bases.fireworks_ai import RUNE_SYSTEMONE_TEST_CASE, RUNE_TEST_CASE +from integration.translation.runner import assert_translation + + +@pytest.mark.parametrize("case", [RUNE_TEST_CASE, RUNE_SYSTEMONE_TEST_CASE], ids=lambda case: case.id) +def test_decisions_basic_fireworks_ai(case: TranslationTestCase, gateway: Gateway, provider: SharedProvider) -> None: + assert_translation(case, gateway, provider) diff --git a/tests/integration/translation/messages/bases/gemini.py b/tests/integration/translation/messages/bases/gemini.py index 217a9960b6e..b464f9a686e 100644 --- a/tests/integration/translation/messages/bases/gemini.py +++ b/tests/integration/translation/messages/bases/gemini.py @@ -18,7 +18,7 @@ GEMINI_3_5_FLASH_TEST_CASE: Final = TranslationTestCase( expected_provider_headers={"content-type": "application/json", "x-goog-api-key": "synthetic-gemini-key"}, expected_provider_request={ "contents": [{"parts": [{"text": "Say hello."}], "role": "user"}], - "generationConfig": {"max_output_tokens": 1024, "temperature": 1.0}, + "generationConfig": {"max_output_tokens": 1024}, "system_instruction": {"parts": [{"text": "You are a terse assistant."}]}, }, mock_provider_response={ @@ -72,7 +72,7 @@ GEMINI_3_8_FLASH_TEST_CASE: Final = TranslationTestCase( expected_provider_headers={"content-type": "application/json", "x-goog-api-key": "synthetic-gemini-key"}, expected_provider_request={ "contents": [{"parts": [{"text": "Say hello."}], "role": "user"}], - "generationConfig": {"max_output_tokens": 1024, "temperature": 1.0}, + "generationConfig": {"max_output_tokens": 1024}, "system_instruction": {"parts": [{"text": "You are a terse assistant."}]}, }, mock_provider_response={ @@ -126,7 +126,7 @@ GEMINI_3_1_PRO_PREVIEW_TEST_CASE: Final = TranslationTestCase( expected_provider_headers={"content-type": "application/json", "x-goog-api-key": "synthetic-gemini-key"}, expected_provider_request={ "contents": [{"parts": [{"text": "Say hello."}], "role": "user"}], - "generationConfig": {"max_output_tokens": 1024, "temperature": 1.0}, + "generationConfig": {"max_output_tokens": 1024}, "system_instruction": {"parts": [{"text": "You are a terse assistant."}]}, }, mock_provider_response={ diff --git a/tests/integration/translation/messages/bases/vertex_ai.py b/tests/integration/translation/messages/bases/vertex_ai.py index 97d48bdc871..be1765f5d26 100644 --- a/tests/integration/translation/messages/bases/vertex_ai.py +++ b/tests/integration/translation/messages/bases/vertex_ai.py @@ -379,7 +379,7 @@ GEMINI_3_5_FLASH_TEST_CASE: Final = TranslationTestCase( expected_provider_request={ "contents": [{"role": "user", "parts": [{"text": "Say hello."}]}], "system_instruction": {"parts": [{"text": "You are a terse assistant."}]}, - "generationConfig": {"max_output_tokens": 1024, "temperature": 1.0}, + "generationConfig": {"max_output_tokens": 1024}, }, mock_provider_response={ "candidates": [ @@ -433,7 +433,7 @@ GEMINI_3_8_FLASH_TEST_CASE: Final = TranslationTestCase( expected_provider_request={ "contents": [{"role": "user", "parts": [{"text": "Say hello."}]}], "system_instruction": {"parts": [{"text": "You are a terse assistant."}]}, - "generationConfig": {"max_output_tokens": 1024, "temperature": 1.0}, + "generationConfig": {"max_output_tokens": 1024}, }, mock_provider_response={ "candidates": [ @@ -487,7 +487,7 @@ GEMINI_3_1_PRO_PREVIEW_TEST_CASE: Final = TranslationTestCase( expected_provider_request={ "contents": [{"role": "user", "parts": [{"text": "Say hello."}]}], "system_instruction": {"parts": [{"text": "You are a terse assistant."}]}, - "generationConfig": {"max_output_tokens": 1024, "temperature": 1.0}, + "generationConfig": {"max_output_tokens": 1024}, }, mock_provider_response={ "candidates": [ diff --git a/tests/integration/translation/messages/sampling/__init__.py b/tests/integration/translation/messages/sampling/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/integration/translation/messages/sampling/test_messages_sampling_gemini.py b/tests/integration/translation/messages/sampling/test_messages_sampling_gemini.py new file mode 100644 index 00000000000..e2b54774b1b --- /dev/null +++ b/tests/integration/translation/messages/sampling/test_messages_sampling_gemini.py @@ -0,0 +1,26 @@ +from dataclasses import replace +from typing import Final + +import pytest +from integration._support.client import Gateway, object_value +from integration._support.provider import SharedProvider +from integration.translation.case import TranslationTestCase +from integration.translation.messages.bases.gemini import GEMINI_3_8_FLASH_TEST_CASE +from integration.translation.runner import assert_translation + +GENERATION_CONFIG: Final = object_value(GEMINI_3_8_FLASH_TEST_CASE.expected_provider_request["generationConfig"]) + +GEMINI_3_8_FLASH_TEMPERATURE_TEST_CASE: Final = replace( + GEMINI_3_8_FLASH_TEST_CASE, + scenario="temperature", + litellm_request={**GEMINI_3_8_FLASH_TEST_CASE.litellm_request, "temperature": 0.4}, + expected_provider_request={ + **GEMINI_3_8_FLASH_TEST_CASE.expected_provider_request, + "generationConfig": {**GENERATION_CONFIG, "temperature": 0.4}, + }, +) + + +@pytest.mark.parametrize("case", [GEMINI_3_8_FLASH_TEMPERATURE_TEST_CASE], ids=lambda case: case.id) +def test_messages_sampling_gemini(case: TranslationTestCase, gateway: Gateway, provider: SharedProvider) -> None: + assert_translation(case, gateway, provider) diff --git a/tests/integration/translation/responses/bases/gemini.py b/tests/integration/translation/responses/bases/gemini.py index 3b72eb4ff00..34591a31843 100644 --- a/tests/integration/translation/responses/bases/gemini.py +++ b/tests/integration/translation/responses/bases/gemini.py @@ -19,7 +19,7 @@ GEMINI_3_5_FLASH_TEST_CASE: Final = TranslationTestCase( expected_provider_headers={"content-type": "application/json", "x-goog-api-key": "synthetic-gemini-key"}, expected_provider_request={ "contents": [{"parts": [{"text": "Say hello."}], "role": "user"}], - "generationConfig": {"max_output_tokens": 1024, "temperature": 1.0}, + "generationConfig": {"max_output_tokens": 1024}, "system_instruction": {"parts": [{"text": "You are a terse assistant."}]}, }, mock_provider_response={ @@ -110,7 +110,7 @@ GEMINI_3_8_FLASH_TEST_CASE: Final = TranslationTestCase( expected_provider_headers={"content-type": "application/json", "x-goog-api-key": "synthetic-gemini-key"}, expected_provider_request={ "contents": [{"parts": [{"text": "Say hello."}], "role": "user"}], - "generationConfig": {"max_output_tokens": 1024, "temperature": 1.0}, + "generationConfig": {"max_output_tokens": 1024}, "system_instruction": {"parts": [{"text": "You are a terse assistant."}]}, }, mock_provider_response={ @@ -201,7 +201,7 @@ GEMINI_3_1_PRO_PREVIEW_TEST_CASE: Final = TranslationTestCase( expected_provider_headers={"content-type": "application/json", "x-goog-api-key": "synthetic-gemini-key"}, expected_provider_request={ "contents": [{"parts": [{"text": "Say hello."}], "role": "user"}], - "generationConfig": {"max_output_tokens": 1024, "temperature": 1.0}, + "generationConfig": {"max_output_tokens": 1024}, "system_instruction": {"parts": [{"text": "You are a terse assistant."}]}, }, mock_provider_response={ diff --git a/tests/integration/translation/responses/bases/vertex_ai.py b/tests/integration/translation/responses/bases/vertex_ai.py index eb0bf46e9db..de2ac6fdd68 100644 --- a/tests/integration/translation/responses/bases/vertex_ai.py +++ b/tests/integration/translation/responses/bases/vertex_ai.py @@ -592,7 +592,7 @@ GEMINI_3_5_FLASH_TEST_CASE: Final = TranslationTestCase( expected_provider_request={ "contents": [{"role": "user", "parts": [{"text": "Say hello."}]}], "system_instruction": {"parts": [{"text": "You are a terse assistant."}]}, - "generationConfig": {"max_output_tokens": 1024, "temperature": 1.0}, + "generationConfig": {"max_output_tokens": 1024}, }, mock_provider_response={ "candidates": [ @@ -684,7 +684,7 @@ GEMINI_3_8_FLASH_TEST_CASE: Final = TranslationTestCase( expected_provider_request={ "contents": [{"role": "user", "parts": [{"text": "Say hello."}]}], "system_instruction": {"parts": [{"text": "You are a terse assistant."}]}, - "generationConfig": {"max_output_tokens": 1024, "temperature": 1.0}, + "generationConfig": {"max_output_tokens": 1024}, }, mock_provider_response={ "candidates": [ @@ -776,7 +776,7 @@ GEMINI_3_1_PRO_PREVIEW_TEST_CASE: Final = TranslationTestCase( expected_provider_request={ "contents": [{"role": "user", "parts": [{"text": "Say hello."}]}], "system_instruction": {"parts": [{"text": "You are a terse assistant."}]}, - "generationConfig": {"max_output_tokens": 1024, "temperature": 1.0}, + "generationConfig": {"max_output_tokens": 1024}, }, mock_provider_response={ "candidates": [ diff --git a/tests/integration/translation/responses/sampling/__init__.py b/tests/integration/translation/responses/sampling/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/integration/translation/responses/sampling/test_responses_sampling_gemini.py b/tests/integration/translation/responses/sampling/test_responses_sampling_gemini.py new file mode 100644 index 00000000000..6e82132aecb --- /dev/null +++ b/tests/integration/translation/responses/sampling/test_responses_sampling_gemini.py @@ -0,0 +1,27 @@ +from dataclasses import replace +from typing import Final + +import pytest +from integration._support.client import Gateway, object_value +from integration._support.provider import SharedProvider +from integration.translation.case import TranslationTestCase +from integration.translation.responses.bases.gemini import GEMINI_3_8_FLASH_TEST_CASE +from integration.translation.runner import assert_translation + +GENERATION_CONFIG: Final = object_value(GEMINI_3_8_FLASH_TEST_CASE.expected_provider_request["generationConfig"]) + +GEMINI_3_8_FLASH_TEMPERATURE_TEST_CASE: Final = replace( + GEMINI_3_8_FLASH_TEST_CASE, + scenario="temperature", + litellm_request={**GEMINI_3_8_FLASH_TEST_CASE.litellm_request, "temperature": 0.4}, + expected_provider_request={ + **GEMINI_3_8_FLASH_TEST_CASE.expected_provider_request, + "generationConfig": {**GENERATION_CONFIG, "temperature": 0.4}, + }, + expected_litellm_response={**GEMINI_3_8_FLASH_TEST_CASE.expected_litellm_response, "temperature": 0.4}, +) + + +@pytest.mark.parametrize("case", [GEMINI_3_8_FLASH_TEMPERATURE_TEST_CASE], ids=lambda case: case.id) +def test_responses_sampling_gemini(case: TranslationTestCase, gateway: Gateway, provider: SharedProvider) -> None: + assert_translation(case, gateway, provider) diff --git a/tests/unit/caching/test_caching.py b/tests/unit/caching/test_caching.py index 0cad6e350e3..c28455f3aca 100644 --- a/tests/unit/caching/test_caching.py +++ b/tests/unit/caching/test_caching.py @@ -1548,3 +1548,23 @@ async def test_redis_caching_ttl_sadd( ) redis_client.expire.assert_awaited_once_with("cache-key", timedelta(seconds=120)) + + +def test_cache_key_separates_model_ids_unless_a_model_group_owns_the_key() -> None: + cache: Final = Cache(type=LiteLLMCacheType.LOCAL) + variants: Final[tuple[dict[str, object], ...]] = ( + {"model_id": "component-a"}, + {"model_id": "component-b"}, + {}, + {"model_id": "component-a"}, + {"model_id": "component-a", "metadata": {"model_group": "embed"}}, + {"model_id": "component-b", "metadata": {"model_group": "embed"}}, + ) + key_a, key_b, bare, key_a_again, grouped_a, grouped_b = ( + cache.get_cache_key(model="sagemaker/openai/embed", input=["hello"], **extra) for extra in variants + ) + + assert key_a != key_b + assert key_a != bare + assert key_a == key_a_again + assert grouped_a == grouped_b diff --git a/tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py b/tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py index 00653a5d127..8072919ed87 100644 --- a/tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py +++ b/tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py @@ -5138,4 +5138,132 @@ def test_transform_request_drop_params_in_litellm_params_gates_the_prompt_cache_ litellm_logging_obj=Mock(), ) - assert "prompt_cache_breakpoint" not in result["input"][0]["content"][0] \ No newline at end of file + assert "prompt_cache_breakpoint" not in result["input"][0]["content"][0] + + +def test_convert_response_output_folds_two_messages_and_function_call_into_one_choice() -> None: + first: Final = ResponseOutputMessage.model_validate( + { + "id": "msg_first", + "type": "message", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": "First message.", "annotations": []}], + } + ) + second_annotation: Final = { + "type": "url_citation", + "start_index": 0, + "end_index": 6, + "title": "Second", + "url": "https://example.com/second", + } + second: Final = ResponseOutputMessage.model_validate( + { + "id": "msg_second", + "type": "message", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": "Second message.", "annotations": [second_annotation]}], + } + ) + function_call: Final = ResponseFunctionToolCall( + id="fc_1", + type="function_call", + status="completed", + arguments='{"city":"Paris"}', + call_id="call_1", + name="get_weather", + ) + + choices: Final = LiteLLMResponsesTransformationHandler._convert_response_output_to_choices( # pyright: ignore[reportPrivateUsage] # exercises the converter directly + (first, second, function_call) + ) + + assert len(choices) == 1 + assert choices[0].index == 0 + assert choices[0].finish_reason == "tool_calls" + assert choices[0].message.content == "First message.Second message." + assert choices[0].message.annotations == [ + {**second_annotation, "start_index": len("First message."), "end_index": len("First message.") + 6} + ] + tool_calls: Final = choices[0].message.tool_calls + assert tool_calls is not None + assert [call.model_dump(exclude_none=True) for call in tool_calls] == [ + { + "id": "fc_1", + "type": "function", + "function": {"name": "get_weather", "arguments": '{"city":"Paris"}'}, + "index": 0, + } + ] + + +def test_convert_response_output_folds_commentary_and_final_answer_into_one_choice() -> None: + commentary: Final = ResponseOutputMessage.model_validate( + { + "id": "msg_commentary", + "type": "message", + "role": "assistant", + "status": "completed", + "phase": "commentary", + "content": [{"type": "output_text", "text": "Let me answer that.", "annotations": []}], + } + ) + final_answer: Final = ResponseOutputMessage.model_validate( + { + "id": "msg_final", + "type": "message", + "role": "assistant", + "status": "completed", + "phase": "final_answer", + "content": [{"type": "output_text", "text": '{"city": "Paris"}', "annotations": []}], + } + ) + + choices: Final = LiteLLMResponsesTransformationHandler._convert_response_output_to_choices( # pyright: ignore[reportPrivateUsage] # exercises the converter directly + (commentary, final_answer) + ) + + assert len(choices) == 1 + assert choices[0].index == 0 + assert choices[0].finish_reason == "stop" + assert choices[0].message.content == 'Let me answer that.{"city": "Paris"}' + assert choices[0].message.tool_calls is None + + +def test_convert_response_output_numbers_two_function_calls_zero_and_one() -> None: + message: Final = ResponseOutputMessage.model_validate( + { + "id": "msg_only", + "type": "message", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": "Checking both.", "annotations": []}], + } + ) + weather: Final = ResponseFunctionToolCall( + id="fc_1", + type="function_call", + status="completed", + arguments='{"city":"Paris"}', + call_id="call_1", + name="get_weather", + ) + time: Final = ResponseFunctionToolCall( + id="fc_2", + type="function_call", + status="completed", + arguments='{"city":"Paris"}', + call_id="call_2", + name="get_time", + ) + + choices: Final = LiteLLMResponsesTransformationHandler._convert_response_output_to_choices( # pyright: ignore[reportPrivateUsage] # the fold is the unit under test + (message, weather, time) + ) + + assert len(choices) == 1 + tool_calls: Final = choices[0].message.tool_calls + assert tool_calls is not None + assert [(call.index, call.function.name) for call in tool_calls] == [(0, "get_weather"), (1, "get_time")] diff --git a/tests/unit/decisions/test_main.py b/tests/unit/decisions/test_main.py index e1b2028fa30..e36d5b8f9b1 100644 --- a/tests/unit/decisions/test_main.py +++ b/tests/unit/decisions/test_main.py @@ -2,9 +2,10 @@ from __future__ import annotations import asyncio import json +import re import threading -from concurrent.futures import wait from collections.abc import Mapping +from concurrent.futures import wait from types import MappingProxyType from typing import Final @@ -16,8 +17,8 @@ from botocore.awsrequest import AWSRequest from botocore.credentials import Credentials import litellm -from litellm.cost_calculator import get_response_cost_from_hidden_params from litellm.constants import AWS_SIGNING_MAX_THREADS +from litellm.cost_calculator import get_response_cost_from_hidden_params from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER from litellm.llms.bedrock.base_aws_llm import AWS_SIGNING_EXECUTOR @@ -167,6 +168,25 @@ def _httpx_transport(monkeypatch: pytest.MonkeyPatch) -> None: litellm.in_memory_llm_clients_cache.flush_cache() +_GPT_6_LUNA_DECISIONS_INPUT_RATE: Final = 1e-07 +_MICROSOFT_DECISION_1_DECISIONS_INPUT_RATE: Final = 7e-08 +_DECISIONS_PRICING_FIELDS: Final[dict[str, dict[str, object]]] = { + "gpt-6-luna": { + "input_cost_per_token_decisions": _GPT_6_LUNA_DECISIONS_INPUT_RATE, + "cache_read_input_token_cost_decisions": _GPT_6_LUNA_DECISIONS_INPUT_RATE, + "cache_creation_input_token_cost_decisions": _GPT_6_LUNA_DECISIONS_INPUT_RATE, + "output_cost_per_token_decisions": 0.0, + }, + "azure_ai/Microsoft-Decision-1": {"input_cost_per_token_decisions": _MICROSOFT_DECISION_1_DECISIONS_INPUT_RATE}, +} + + +@pytest.fixture +def decisions_pricing_fields(monkeypatch: pytest.MonkeyPatch) -> None: + for key, fields in _DECISIONS_PRICING_FIELDS.items(): + monkeypatch.setitem(litellm.model_cost, key, {**litellm.model_cost[key], **fields}) + + @pytest.mark.asyncio @pytest.mark.parametrize(("provider", "model", "url", "upstream_model"), _PROVIDERS) async def test_adecisions_sends_the_provider_wire_contract( @@ -322,6 +342,79 @@ def test_decisions_cost_uses_litellm_token_pricing() -> None: assert cost == pytest.approx(expected_cost) +@pytest.mark.asyncio +async def test_openai_decisions_uses_evaluation_pricing_and_explicit_base_model( + respx_mock: respx.MockRouter, + decisions_pricing_fields: None, +) -> None: + route: Final = respx_mock.post("https://api.openai.com/v1/decisions").respond(json=_OPENAI_RESPONSE) + response: Final = await litellm.adecisions( + model="openai/gpt-6-luna", + input="The package arrived broken.", + questions=[{"type": "predicate", "name": "is_defect", "instructions": "Is this a defect?"}], + api_key="caller-key", + api_base="https://api.openai.com/v1", + ) + chat_cost: Mapping[str, object] = litellm.model_cost["gpt-6-luna"] + expected_chat_cost: Final = ( + (_INPUT_TOKENS - _CACHED_TOKENS - _CACHE_WRITE_TOKENS) * float(chat_cost["input_cost_per_token"]) + + _CACHED_TOKENS * float(chat_cost["cache_read_input_token_cost"]) + + _CACHE_WRITE_TOKENS * float(chat_cost["cache_creation_input_token_cost"]) + + _OUTPUT_TOKENS * float(chat_cost["output_cost_per_token"]) + ) + + assert route.called + assert response._hidden_params["model"] == "openai/gpt-6-luna" + assert response._hidden_params["provider_response_model"] == "openai/gpt-6-luna" + assert litellm.completion_cost(completion_response=response) == pytest.approx( + _INPUT_TOKENS * _GPT_6_LUNA_DECISIONS_INPUT_RATE + ) + assert litellm.completion_cost(completion_response=response, call_type="acompletion") == pytest.approx( + expected_chat_cost + ) + + +@pytest.mark.asyncio +async def test_systemone_uses_evaluation_pricing(respx_mock: respx.MockRouter, decisions_pricing_fields: None) -> None: + route: Final = respx_mock.post("https://api.openai.com/v1/decisions").respond(json=_OPENAI_RESPONSE) + response: Final = await litellm.asystemone( + model="openai/gpt-6-luna", + state="The package arrived broken.", + questions=_QUESTIONS, + api_key="caller-key", + api_base="https://api.openai.com/v1", + ) + + assert route.called + assert response._hidden_params["model"] == "openai/gpt-6-luna" + assert response._hidden_params["provider_response_model"] == "openai/gpt-6-luna" + assert litellm.completion_cost(completion_response=response) == pytest.approx( + _INPUT_TOKENS * _GPT_6_LUNA_DECISIONS_INPUT_RATE + ) + + +@pytest.mark.asyncio +async def test_systemone_falls_back_to_plain_pricing_without_decisions_row( + respx_mock: respx.MockRouter, +) -> None: + route: Final = respx_mock.post("https://decisions.example/v1/decisions").respond(json=_RESPONSE) + response: Final = await litellm.asystemone( + model="perplexity/sonar", + state="review", + questions={"is_defect": {"type": "noul", "instructions": "Is this a defect?"}}, + api_key="caller-key", + api_base="https://decisions.example/v1", + ) + plain_cost: Mapping[str, object] = litellm.model_cost["perplexity/sonar"] + + assert route.called + assert response._hidden_params["provider_response_model"] == "perplexity/sonar" + assert litellm.completion_cost(completion_response=response) == pytest.approx( + _INPUT_TOKENS * float(plain_cost["input_cost_per_token"]) + + _OUTPUT_TOKENS * float(plain_cost["output_cost_per_token"]) + ) + + @pytest.mark.parametrize( "response", ( @@ -1309,6 +1402,7 @@ async def test_openai_decisions_translate_systemone_to_the_openai_wire_contract_ url: str, monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter, + decisions_pricing_fields: None, ) -> None: monkeypatch.setattr(litellm, "api_base", None) monkeypatch.delenv("OPENAI_BASE_URL", raising=False) @@ -1348,13 +1442,7 @@ async def test_openai_decisions_translate_systemone_to_the_openai_wire_contract_ } assert response.answers == {"is_defect": NoulAnswer(type="noul", noul=0.9)} assert response.hidden_params["custom_llm_provider"] == "openai" - luna_cost: Final = litellm.model_cost["gpt-6-luna"] - expected_cost: Final = ( - (_INPUT_TOKENS - _CACHED_TOKENS - _CACHE_WRITE_TOKENS) * float(luna_cost["input_cost_per_token"]) - + _CACHED_TOKENS * float(luna_cost["cache_read_input_token_cost"]) - + _CACHE_WRITE_TOKENS * float(luna_cost["cache_creation_input_token_cost"]) - + _OUTPUT_TOKENS * float(luna_cost["output_cost_per_token"]) - ) + expected_cost: Final = _INPUT_TOKENS * _GPT_6_LUNA_DECISIONS_INPUT_RATE assert expected_cost > 0 assert litellm.completion_cost(completion_response=response) == pytest.approx(expected_cost) @@ -1693,6 +1781,28 @@ async def test_azure_ai_decision_posts_deployment_to_foundry_systemone_route( assert response._hidden_params["model"] == "azure_ai/decision-1" +@pytest.mark.asyncio +async def test_azure_ai_systemone_uses_its_decisions_pricing_fields( + respx_mock: respx.MockRouter, decisions_pricing_fields: None +) -> None: + route: Final = respx_mock.post("https://res.services.ai.azure.com/providers/microsoft/v1/systemone").respond( + json={**_RESPONSE, "model": "microsoft-decision-1"} + ) + response: Final = await litellm.asystemone( + model="azure_ai/Microsoft-Decision-1", + state="review", + questions={"is_defect": {"type": "noul", "instructions": "Is this a defect?"}}, + api_base="https://res.services.ai.azure.com", + api_key="foundry-key", + ) + + assert route.called + assert response._hidden_params["provider_response_model"] == "azure_ai/Microsoft-Decision-1" + assert litellm.completion_cost(completion_response=response) == pytest.approx( + _INPUT_TOKENS * _MICROSOFT_DECISION_1_DECISIONS_INPUT_RATE + ) + + @pytest.mark.asyncio async def test_azure_ai_decision_reads_foundry_env_and_keeps_gateway_path_prefix( monkeypatch: pytest.MonkeyPatch, @@ -1730,3 +1840,389 @@ async def test_azure_ai_decision_requires_api_base_before_http( ) assert len(respx_mock.calls) == 0 + + +_CLEF_OMNI_UPSTREAM: Final = "https://example.com/ai/run/@cf/cloudflare/clef-omni" +_CLEF_OMNI_RESPONSE: Final[Mapping[str, object]] = { + "success": True, + "result": { + "model": "clef-omni", + "answers": {"heard": {"type": "noul", "noul": 0.9}}, + "usage": {"input_tokens": 100, "output_tokens": 0}, + }, +} +_MEDIA_REQUESTS: Final[tuple[tuple[str, Mapping[str, object]], ...]] = ( + ( + "decisions", + MappingProxyType( + { + "input": "Evaluate the clips", + "questions": [{"type": "predicate", "name": "heard", "instructions": "Is there speech?"}], + } + ), + ), + ( + "systemone", + MappingProxyType( + { + "state": "Evaluate the clips", + "questions": {"heard": {"type": "noul", "instructions": "Is there speech?"}}, + } + ), + ), +) + + +@pytest.mark.parametrize("async_mode", (False, True), ids=("sync", "async")) +@pytest.mark.parametrize(("api", "request_kwargs"), _MEDIA_REQUESTS, ids=_PREDICATE_REQUEST_IDS) +@pytest.mark.parametrize("model", ("cloudflare/clef-omni", "cloudflare/@cf/cloudflare/clef-omni")) +@pytest.mark.asyncio +async def test_clef_omni_media_reaches_upstream_in_both_formats( + respx_mock: respx.MockRouter, async_mode: bool, api: str, request_kwargs: Mapping[str, object], model: str +) -> None: + audio: Final = [ + "data:audio/wav;base64,YQ==", + {"content_type": "audio/mpeg", "base64": "Yg=="}, + {"content_type": "audio/webm;codecs=opus", "base64": "ZQ=="}, + ] + videos: Final = ["data:video/mp4;base64,Yw==", {"content_type": "video/webm;codecs=vp9", "base64": "ZA=="}] + route: Final = respx_mock.post(_CLEF_OMNI_UPSTREAM).respond(200, json=dict(_CLEF_OMNI_RESPONSE)) + call: Final = getattr(litellm, f"a{api}" if async_mode else api) + + response: Final = ( + await call( + model=model, api_base="https://example.com", api_key="test", audio=audio, videos=videos, **request_kwargs + ) + if async_mode + else call( + model=model, api_base="https://example.com", api_key="test", audio=audio, videos=videos, **request_kwargs + ) + ) + + assert json.loads(route.calls[0].request.content) == { + "model": "clef-omni", + "state": "Evaluate the clips", + "questions": {"heard": {"type": "noul", "instructions": "Is there speech?"}}, + "audio": audio, + "videos": videos, + } + assert response.usage is not None + assert response.usage.input_tokens == 100 + assert len(route.calls) == 1 + + +@pytest.mark.parametrize( + "model", + ( + "cloudflare/clef", + "cloudflare/clef-flash", + "typesafe/jev-latest", + "openai/decision", + "fireworks_ai/accounts/fireworks/routers/rune", + ), +) +@pytest.mark.parametrize( + "media", + ( + {"audio": ["data:audio/wav;base64,YQ=="]}, + {"videos": ["data:video/mp4;base64,YQ=="]}, + {"audio": ["data:audio/wav;base64,YQ=="], "videos": ["data:video/mp4;base64,YQ=="]}, + ), + ids=("audio", "videos", "both"), +) +@pytest.mark.asyncio +async def test_clips_for_a_model_without_media_support_are_refused_before_http( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, model: str, media: Mapping[str, object] +) -> None: + monkeypatch.setattr(litellm, "drop_params", False) + named_params: Final = str(list(media)) + + with pytest.raises( + litellm.UnsupportedParamsError, match=rf"does not support parameters: {re.escape(named_params)}, .*drop_params" + ): + await litellm.asystemone( + model=model, + state="test", + questions=_QUESTIONS, + api_base="https://example.com", + api_key="test", + **media, + ) + assert len(respx_mock.calls) == 0 + + +@pytest.mark.parametrize( + "clips", + ( + {"audio": ["data:audio/wav;base64,YQ=="], "videos": ["data:video/mp4;base64,YQ=="]}, + {"audio": ["https://example.com/clip.wav"], "videos": [{"content_type": "video/mp4"}]}, + ), + ids=("well_formed", "malformed"), +) +@pytest.mark.parametrize(("api", "request_kwargs"), _PREDICATE_REQUESTS, ids=_PREDICATE_REQUEST_IDS) +@pytest.mark.asyncio +async def test_clips_are_dropped_from_the_wire_for_a_model_without_media_support_under_drop_params( + api: str, + request_kwargs: Mapping[str, object], + clips: Mapping[str, object], + respx_mock: respx.MockRouter, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "drop_params", False) + route: Final = respx_mock.post("https://api.perplexity.ai/v1/decisions").respond(json=_RESPONSE) + call: Final = getattr(litellm, f"a{api}") + + await call(model="perplexity/pplx-decider-v1-27b", api_key="caller-key", **request_kwargs) + await call(model="perplexity/pplx-decider-v1-27b", api_key="caller-key", drop_params=True, **request_kwargs, **clips) + + assert len(route.calls) == 2 + text_only_body: Final = json.loads(route.calls[0].request.content) + assert "audio" not in text_only_body and "videos" not in text_only_body + assert json.loads(route.calls[1].request.content) == text_only_body + + +@pytest.mark.parametrize( + "media", + ( + {"audio": ["https://example.com/clip.wav"]}, + {"videos": ["https://example.com/clip.mp4"]}, + {"audio": [{"content_type": "audio/wav"}]}, + {"videos": [{"content_type": "video/mp4"}]}, + {"audio": [{"content_type": "video/mp4", "base64": "YQ=="}]}, + {"videos": ["data:audio/wav;base64,YQ=="]}, + {"audio": ["data:audio/webm;codecs=opus;base64,YQ=="]}, + {"videos": ["data:video/webm;codecs=vp9;base64,YQ=="]}, + {"audio": ["data:audio/wav;base64,YQ=="] * 5}, + {"videos": ["data:video/mp4;base64,YQ=="] * 3}, + ), +) +@pytest.mark.asyncio +async def test_systemone_validates_media_before_http(respx_mock: respx.MockRouter, media: Mapping[str, object]) -> None: + with pytest.raises(litellm.BadRequestError, match="Invalid Decisions request"): + await litellm.asystemone( + model="cloudflare/clef-omni", + state="test", + questions=_QUESTIONS, + api_base="https://example.com", + api_key="test", + **media, + ) + assert len(respx_mock.calls) == 0 + + +@pytest.mark.parametrize( + "media", ({"audio": [], "videos": []}, {"audio": None, "videos": None}, {"audio": None, "videos": []}), ids=repr +) +@pytest.mark.asyncio +async def test_empty_or_null_media_keeps_text_only_request_unchanged( + respx_mock: respx.MockRouter, media: Mapping[str, object] +) -> None: + route: Final = respx_mock.post("https://example.com/v1/systemone").respond(200, json=dict(_RESPONSE)) + await litellm.asystemone( + model="typesafe/jev-latest", + state="test", + questions=_QUESTIONS, + api_base="https://example.com", + api_key="test", + **media, + ) + assert json.loads(route.calls[0].request.content) == { + "model": "jev-latest", + "state": "test", + "questions": _QUESTIONS, + } + + +@pytest.mark.asyncio +async def test_a_malformed_clip_for_a_model_that_takes_media_is_refused_even_under_drop_params( + respx_mock: respx.MockRouter, +) -> None: + with pytest.raises(litellm.BadRequestError, match="Invalid Decisions request"): + await litellm.asystemone( + model="cloudflare/clef-omni", + state="test", + questions=_QUESTIONS, + api_base="https://example.com", + api_key="test", + audio=["https://example.com/clip.wav"], + drop_params=True, + ) + assert len(respx_mock.calls) == 0 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("mode", ("evaluation", "decisions")) +async def test_media_support_is_read_per_kind_from_the_cost_map( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, local_model_cost_map: None, mode: str +) -> None: + monkeypatch.setattr(litellm, "drop_params", False) + monkeypatch.setitem( + litellm.model_cost, + "typesafe/jev-hears", + {"litellm_provider": "typesafe", "mode": mode, "supports_audio_input": True}, + ) + route: Final = respx_mock.post("https://example.com/v1/systemone").respond(200, json=dict(_RESPONSE)) + audio: Final = ["data:audio/wav;base64,YQ=="] + + await litellm.asystemone( + model="typesafe/jev-hears", + state="test", + questions=_QUESTIONS, + api_base="https://example.com", + api_key="test", + audio=audio, + ) + with pytest.raises(litellm.UnsupportedParamsError, match=re.escape("does not support parameters: ['videos']")): + await litellm.asystemone( + model="typesafe/jev-hears", + state="test", + questions=_QUESTIONS, + api_base="https://example.com", + api_key="test", + audio=audio, + videos=["data:video/mp4;base64,YQ=="], + ) + + assert len(route.calls) == 1 + assert json.loads(route.calls[0].request.content) == { + "model": "jev-hears", + "state": "test", + "questions": _QUESTIONS, + "audio": audio, + } + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("model", "url", "cost_map_entry"), + ( + ( + "typesafe/jev-chatty", + "https://example.com/v1/systemone", + {"litellm_provider": "typesafe", "mode": "chat", "supports_audio_input": True, "supports_video_input": True}, + ), + ("openai/gpt-audio", "https://example.com/v1/decisions", None), + ), +) +async def test_a_chat_models_media_flags_do_not_count_for_decisions( + respx_mock: respx.MockRouter, + monkeypatch: pytest.MonkeyPatch, + local_model_cost_map: None, + model: str, + url: str, + cost_map_entry: Mapping[str, object] | None, +) -> None: + monkeypatch.setattr(litellm, "drop_params", False) + if cost_map_entry is not None: + monkeypatch.setitem(litellm.model_cost, model, cost_map_entry) + assert litellm.supports_audio_input(model) is True + route: Final = respx_mock.post(url).respond(200, json=dict(_RESPONSE if "systemone" in url else _OPENAI_RESPONSE)) + questions: Final = {"is_defect": {"type": "noul", "instructions": "Is this a defect?"}} + audio: Final = ["data:audio/wav;base64,YQ=="] + videos: Final = ["data:video/mp4;base64,YQ=="] + + with pytest.raises(litellm.UnsupportedParamsError, match=re.escape("does not support parameters: ['audio', 'videos']")): + await litellm.asystemone( + model=model, + state="test", + questions=questions, + api_base="https://example.com", + api_key="test", + audio=audio, + videos=videos, + ) + assert not route.calls + + await litellm.asystemone( + model=model, + state="test", + questions=questions, + api_base="https://example.com", + api_key="test", + audio=audio, + videos=videos, + drop_params=True, + ) + sent: Final = json.loads(route.calls[0].request.content) + assert "audio" not in sent + assert "videos" not in sent + assert sent["model"] == model.split("/", 1)[1] + + +def test_positional_arguments_keep_the_order_the_sdk_had_before_media(respx_mock: respx.MockRouter) -> None: + decisions_route: Final = respx_mock.post("https://example.com/v1/decisions").respond(json=_OPENAI_RESPONSE) + systemone_route: Final = respx_mock.post("https://example.com/v1/systemone").respond(json=_RESPONSE) + openai_questions: Final = [{"type": "predicate", "name": "is_defect", "instructions": "Is this a defect?"}] + systemone_questions: Final = {"is_defect": {"type": "noul", "instructions": "Is this a defect?"}} + + litellm.decisions("openai/gpt-6-luna", "review", openai_questions, None, "positional-key", "https://example.com") + litellm.systemone("typesafe/jev-1.13", "review", systemone_questions, "positional-key", "https://example.com") + + assert decisions_route.calls[0].request.headers["authorization"] == "Bearer positional-key" + assert systemone_route.calls[0].request.headers["authorization"] == "Bearer positional-key" + for sent in (decisions_route.calls[0].request.content, systemone_route.calls[0].request.content): + assert "audio" not in json.loads(sent) + + +_FIREWORKS_URL: Final = "https://api.fireworks.ai/inference/v1/systemone" +_FIREWORKS_MODEL: Final = "fireworks_ai/accounts/fireworks/routers/rune" +_FIREWORKS_RESPONSE: Final[Mapping[str, object]] = { + "model": "accounts/fireworks/routers/rune", + "answers": {"is_defect": {"type": "noul", "noul": 0.9902915235185259}}, + "usage": {"input_tokens": 300, "cached_input_tokens": 200, "output_tokens": 0}, +} + + +@pytest.mark.asyncio +async def test_fireworks_ai_posts_the_jev_body_to_systemone_with_the_account_model_name( + monkeypatch: pytest.MonkeyPatch, + respx_mock: respx.MockRouter, +) -> None: + monkeypatch.delenv("FIREWORKS_API_BASE", raising=False) + monkeypatch.delenv("FIREWORKS_API_KEY", raising=False) + monkeypatch.setenv("FIREWORKS_AI_API_KEY", "fireworks-key") + route: Final = respx_mock.post(_FIREWORKS_URL).respond(json=_FIREWORKS_RESPONSE) + + response: Final = await litellm.asystemone( + model=_FIREWORKS_MODEL, + state="The package arrived broken.", + questions={"is_defect": {"type": "noul", "instructions": "Is this a defect?"}}, + ) + + assert route.called + assert respx_mock.calls[0].request.headers["authorization"] == "Bearer fireworks-key" + assert json.loads(respx_mock.calls[0].request.content) == { + "model": "accounts/fireworks/routers/rune", + "state": "The package arrived broken.", + "questions": {"is_defect": {"type": "noul", "instructions": "Is this a defect?"}}, + } + assert response.model == _FIREWORKS_RESPONSE["model"] + assert response.answers["is_defect"] == NoulAnswer(type="noul", noul=0.9902915235185259) + + +@pytest.mark.asyncio +async def test_fireworks_ai_bills_cached_input_tokens_at_the_cache_read_rate( + monkeypatch: pytest.MonkeyPatch, + respx_mock: respx.MockRouter, +) -> None: + monkeypatch.delenv("FIREWORKS_API_BASE", raising=False) + respx_mock.post(_FIREWORKS_URL).respond(json=_FIREWORKS_RESPONSE) + prices: Final = litellm.model_cost[_FIREWORKS_MODEL] + uncached: Final = 100 * float(prices["input_cost_per_token"]) + cached: Final = 200 * float(prices["cache_read_input_token_cost"]) + recording_logger: Final = _RecordingLogger() + monkeypatch.setattr(litellm, "callbacks", [recording_logger]) + + response: Final = await litellm.adecisions( + model=_FIREWORKS_MODEL, + input="The package arrived broken.", + questions=[{"type": "predicate", "name": "is_defect", "instructions": "Is this a defect?"}], + api_key="caller-key", + ) + await _drain_logging_worker() + + assert cached > 0 + assert response.usage.input_tokens_details.cached_tokens == 200 + assert recording_logger.standard_logging_object is not None + assert recording_logger.standard_logging_object["response_cost"] == pytest.approx(uncached + cached) + assert recording_logger.standard_logging_object["response_cost"] < 300 * float(prices["input_cost_per_token"]) diff --git a/tests/unit/enterprise/proxy/hooks/test_managed_files.py b/tests/unit/enterprise/proxy/hooks/test_managed_files.py index ed80bd9b928..0b22d07ee00 100644 --- a/tests/unit/enterprise/proxy/hooks/test_managed_files.py +++ b/tests/unit/enterprise/proxy/hooks/test_managed_files.py @@ -5264,3 +5264,113 @@ async def test_list_user_batches_target_model_names_filter_rejected_with_400(): assert exc.value.type == "invalid_request_error" assert exc.value.param == "target_model_names" assert exc.value.message == "Filtering by 'target_model_names' is not supported when using managed batches." + + +def _tenant_a_input_file_id() -> str: + raw: Final = ( + "litellm_proxy:application/octet-stream;unified_id,7d1c4b0e-5a7f-4a43-9a0e-3f8e2a1b6c11;" + "target_model_names,gpt-4o-mini;llm_output_file_id,file-tenant-a;llm_output_file_model_id,model-a" + ) + return base64.urlsafe_b64encode(raw.encode()).decode().rstrip("=") + + +def _tenant_a_input_file_row(unified_file_id: str) -> LiteLLM_ManagedFileTable: + return LiteLLM_ManagedFileTable( + unified_file_id=unified_file_id, + file_object=None, + model_mappings={"model-a": "file-tenant-a"}, + flat_model_file_ids=["file-tenant-a"], + created_by="user_a", + team_id="team_a", + ) + + +_INPUT_FILE_CREATE_CALLS: Final = ( + ("acreate_batch", "input_file_id"), + ("acreate_fine_tuning_job", "training_file"), + ("acreate_fine_tuning_job", "validation_file"), +) + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("call_type", "file_field"), _INPUT_FILE_CREATE_CALLS) +async def test_create_from_another_tenants_managed_file_is_denied(call_type: str, file_field: str) -> None: + from litellm.proxy._types import UserAPIKeyAuth + + unified_file_id: Final = _tenant_a_input_file_id() + managed_files, _ = _managed_files_with_fake_prisma(_tenant_a_input_file_row(unified_file_id)) + + with pytest.raises(HTTPException) as exc_info: + await managed_files.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(user_id="user_b", team_id="team_b", parent_otel_span=MagicMock()), + cache=MagicMock(), + data={"model": "gpt-4o-mini", "training_file": "file-own-training", file_field: unified_file_id}, + call_type=call_type, + ) + + assert exc_info.value.status_code == 403, f"{call_type}.{file_field} let tenant B use tenant A's file" + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("call_type", "file_field"), _INPUT_FILE_CREATE_CALLS) +async def test_create_from_managed_file_id_with_no_db_row_is_404(call_type: str, file_field: str) -> None: + from litellm.proxy._types import UserAPIKeyAuth + + managed_files, _ = _managed_files_with_fake_prisma() + + with pytest.raises(HTTPException) as exc_info: + await managed_files.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(user_id="user_a", team_id="team_a", parent_otel_span=MagicMock()), + cache=MagicMock(), + data={"model": "gpt-4o-mini", file_field: _tenant_a_input_file_id()}, + call_type=call_type, + ) + + assert exc_info.value.status_code == 404, f"{call_type}.{file_field} dispatched a managed id with no owner row" + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("call_type", "file_field"), _INPUT_FILE_CREATE_CALLS) +@pytest.mark.parametrize( + "caller_kwargs", + [ + {"user_id": "user_a", "team_id": "team_a"}, + {"user_id": "teammate_of_a", "team_id": "team_a"}, + {"user_id": "admin", "team_id": "team_b", "user_role": "proxy_admin"}, + ], + ids=["owner", "same_team", "proxy_admin"], +) +async def test_create_from_accessible_managed_file_is_allowed( + call_type: str, file_field: str, caller_kwargs: dict[str, str] +) -> None: + from litellm.proxy._types import UserAPIKeyAuth + + unified_file_id: Final = _tenant_a_input_file_id() + managed_files, _ = _managed_files_with_fake_prisma(_tenant_a_input_file_row(unified_file_id)) + + result: Final = await managed_files.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(**caller_kwargs, parent_otel_span=MagicMock()), + cache=MagicMock(), + data={"model": "gpt-4o-mini", file_field: unified_file_id}, + call_type=call_type, + ) + + assert isinstance(result, dict) and result[file_field] == unified_file_id, result + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("call_type", "file_field"), _INPUT_FILE_CREATE_CALLS) +async def test_create_from_raw_provider_file_id_skips_managed_ownership_lookup(call_type: str, file_field: str) -> None: + from litellm.proxy._types import UserAPIKeyAuth + + managed_files, table = _managed_files_with_fake_prisma() + + result: Final = await managed_files.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(user_id="user_b", team_id="team_b", parent_otel_span=MagicMock()), + cache=MagicMock(), + data={"model": "gpt-4o-mini", file_field: "file-raw-provider-id"}, + call_type=call_type, + ) + + assert isinstance(result, dict) and result[file_field] == "file-raw-provider-id", result + assert not any("unified_file_id" in where for where in table.find_first_calls), table.find_first_calls diff --git a/tests/unit/experimental_mcp_client/test_mcp_client.py b/tests/unit/experimental_mcp_client/test_mcp_client.py index f949d2d3e58..f92f7f74eb6 100644 --- a/tests/unit/experimental_mcp_client/test_mcp_client.py +++ b/tests/unit/experimental_mcp_client/test_mcp_client.py @@ -2757,6 +2757,35 @@ async def test_optional_discovery_allows_exhaustion_at_page_cap(method: str, mon ) +def test_standalone_client_request_does_not_require_proxy_dependencies(): + import textwrap + + from tests.test_litellm_rust.support.child_interpreter import run_child_interpreter + + script = textwrap.dedent(""" + import asyncio + import sys + for name in ("fastapi", "orjson", "email_validator", "annotated_doc"): + sys.modules[name] = None + import httpx2 + from litellm.experimental_mcp_client.client import MCPClient + + async def request(): + upstream = MCPClient(server_url="https://upstream.example/mcp") + transport = httpx2.MockTransport(lambda request: httpx2.Response(200)) + async with upstream._create_httpx_client_factory(transport=transport)() as client: + response = await client.post(upstream.server_url, json={"method": "tools/call"}) + assert response.status_code == 200 + assert all(sys.modules[name] is None for name in ("fastapi", "orjson", "email_validator", "annotated_doc")) + + asyncio.run(request()) + print("standalone request succeeded") + """) + result = run_child_interpreter(script, timeout=60) + assert result.returncode == 0, result.stderr + assert result.stdout.strip() == "standalone request succeeded" + + def test_client_import_before_proxy_credentials_succeeds_in_fresh_process(): import subprocess @@ -4942,3 +4971,50 @@ async def test_persistent_session_ends_after_a_broken_stream_so_the_next_call_ge await asyncio.wait_for(session.run(broken_stream), 5) await asyncio.wait_for(session.wait_closed(), 5) assert session.closed, "a dead transport must end the session instead of being reused" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("session_started_in_request", [True, False]) +async def test_persistent_session_releases_the_current_requests_scope_response(session_started_in_request) -> None: + from mcp.server.context import ServerRequestContext + from litellm.proxy._experimental.mcp_server.mcp_context import active_mcp_request_ctx_var + from litellm.proxy._experimental.mcp_server.scope_challenge import SCOPE_RESPONSE_KEY, OAuthScopeResponse + + app = _stateful_upstream() + async with app.router.lifespan_context(app): + client = _StatefulUpstreamClient(app, server_url="http://upstream/mcp", transport_type=MCPTransport.http) + first = OAuthScopeResponse(AsyncMock(), "https://gateway.example", "/mcp") + second = OAuthScopeResponse(AsyncMock(), "https://gateway.example", "/mcp") + context = ServerRequestContext( + session=MagicMock(), + lifespan_context={}, + protocol_version="2025-11-25", + method="tools/call", + request=Request({"type": "http", SCOPE_RESPONSE_KEY: first}), + ) + earlier_session = None if session_started_in_request else client.open_persistent_session() + token = active_mcp_request_ctx_var.set(context) + session = earlier_session if earlier_session is not None else client.open_persistent_session() + try: + selected = await client.call_tool( + CallToolRequestParams(name="select_project", arguments={"name": "demo"}), persistent_session=session + ) + assert not selected.is_error and not first.pending + active_mcp_request_ctx_var.set( + ServerRequestContext( + session=context.session, + lifespan_context={}, + protocol_version="2025-11-25", + method="tools/call", + request=Request({"type": "http", SCOPE_RESPONSE_KEY: second}), + ) + ) + created = await client.call_tool( + CallToolRequestParams(name="create_feature", arguments={"title": "feature"}), persistent_session=session + ) + assert not created.is_error and created.content[0].text == "demo/feature" + assert not second.pending + finally: + active_mcp_request_ctx_var.reset(token) + session.close() + await asyncio.wait_for(session.wait_closed(), 5) diff --git a/tests/unit/integration_support/test_budget.py b/tests/unit/integration_support/test_budget.py new file mode 100644 index 00000000000..e0e14fab1af --- /dev/null +++ b/tests/unit/integration_support/test_budget.py @@ -0,0 +1,58 @@ +from __future__ import annotations + +import json +from typing import Final + +import pytest + +from tests.integration._support.budget import ( + IMPORT_SECONDS_VARIABLE, + READINESS_OVERRIDE_VARIABLE, + calibration, + case_timeout_for, + import_seconds, + readiness_for, + readiness_seconds, +) + + +@pytest.mark.parametrize("workers", (1, 2, 4)) +def test_a_box_that_imports_twice_as_slowly_gets_twice_the_readiness_deadline(workers: int) -> None: + assert readiness_for(14.0, workers) == pytest.approx(2 * readiness_for(7.0, workers)) + + +def test_a_multi_worker_proxy_gets_the_two_imports_its_boot_pays_for() -> None: + assert readiness_for(7.0, 2) == pytest.approx(2 * readiness_for(7.0, 1)) + assert readiness_for(7.0, 4) == readiness_for(7.0, 2) + + +@pytest.mark.parametrize("imported", (0.5, 7.0, 40.0)) +def test_a_case_timeout_covers_a_previous_teardown_a_two_worker_boot_and_the_call(imported: float) -> None: + assert case_timeout_for(imported) > readiness_for(imported, 2) + readiness_for(imported, 1) + + +def test_the_calibration_record_matches_the_budgets_it_reports() -> None: + record: Final = json.loads(calibration(7.0)) + assert record == { + "import_seconds": 7.0, + "readiness_seconds_1_worker": readiness_for(7.0, 1), + "readiness_seconds_2_workers": readiness_for(7.0, 2), + "case_timeout_seconds": case_timeout_for(7.0), + } + + +def test_the_exported_calibration_is_reused_instead_of_measured_again(monkeypatch: pytest.MonkeyPatch) -> None: + import_seconds.cache_clear() + monkeypatch.setenv(IMPORT_SECONDS_VARIABLE, "3.5") + monkeypatch.delenv(READINESS_OVERRIDE_VARIABLE, raising=False) + try: + assert import_seconds() == 3.5 + assert readiness_seconds(2) == readiness_for(3.5, 2) + finally: + import_seconds.cache_clear() + + +def test_an_explicit_readiness_override_wins_over_the_calibration(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv(READINESS_OVERRIDE_VARIABLE, "42") + assert readiness_seconds(1) == 42.0 + assert readiness_seconds(2) == 42.0 diff --git a/tests/unit/integration_support/test_process.py b/tests/unit/integration_support/test_process.py index b65c9fa7ec6..b2271a1ed47 100644 --- a/tests/unit/integration_support/test_process.py +++ b/tests/unit/integration_support/test_process.py @@ -5,8 +5,13 @@ import errno import importlib import os import socket +import subprocess +import sys +import threading import time from collections.abc import Iterator +from contextlib import contextmanager +from functools import reduce from pathlib import Path from types import ModuleType from typing import Final @@ -18,8 +23,18 @@ import pytest import yaml from pydantic import TypeAdapter -from tests.integration._support.client import Gateway -from tests.integration._support.process import worker_configuration +from tests.integration._support.client import Gateway, eventually +from tests.integration._support.process import ( + OwnedProxy, + WorkerHealthcheck, + WorkerKill, + healthcheck_seconds, + kill_stray_process, + owned_proxy_process, + replacement_pid, + stop_process_group, + worker_configuration, +) TESTS_DIR: Final = Path(__file__).resolve().parents[2] UNRELATED_CRASH: Final = "Traceback (most recent call last):\nModuleNotFoundError: No module named 'litellm'\n" @@ -224,6 +239,45 @@ def test_only_live_readable_workers_can_satisfy_readiness( process_module._wait_until_ready(launch) +def test_boot_phases_mark_when_prisma_ran_the_port_bound_and_each_worker_answered( + process_module: ModuleType, +) -> None: + probe: Final = process_module._Probe + timeline: Final = ( + probe(0.5, listening=False, prisma=False, ready_pid=None), + probe(4.0, listening=False, prisma=True, ready_pid=None), + probe(5.0, listening=False, prisma=False, ready_pid=None), + probe(5.5, listening=False, prisma=True, ready_pid=None), + probe(6.5, listening=False, prisma=False, ready_pid=None), + probe(7.0, listening=True, prisma=False, ready_pid=None), + probe(10.0, listening=True, prisma=False, ready_pid=41), + probe(11.0, listening=True, prisma=False, ready_pid=42), + ) + phases: Final = reduce(process_module._observe, timeline, process_module.BootPhases()) + assert phases == process_module.BootPhases( + prisma_started=4.0, + prisma_finished=6.5, + listening=7.0, + first_ready=10.0, + ready_pids=frozenset((41, 42)), + ) + + +def test_boot_phases_leave_prisma_unset_when_no_prisma_child_was_seen(process_module: ModuleType) -> None: + probe: Final = process_module._Probe + timeline: Final = ( + probe(1.0, listening=False, prisma=False, ready_pid=None), + probe(2.0, listening=True, prisma=False, ready_pid=7), + ) + phases: Final = reduce(process_module._observe, timeline, process_module.BootPhases()) + assert (phases.prisma_started, phases.prisma_finished, phases.listening, phases.first_ready) == ( + None, + None, + 2.0, + 2.0, + ) + + @pytest.mark.parametrize( ("writer", "explicit_reader", "expected_reader"), ( @@ -278,3 +332,159 @@ def test_owned_runtime_configuration_preserves_custom_worker_startup_hooks( ) if hook ) + + +_SUPERVISOR: Final = """import multiprocessing +import time + +if __name__ == "__main__": + worker = multiprocessing.get_context("spawn").Process(target=time.sleep, args=(60,)) + worker.start() + print(worker.pid, flush=True) + deadline = time.monotonic() + 60 + while time.monotonic() < deadline: + multiprocessing.active_children() + time.sleep(0.05) +""" + + +@contextmanager +def _supervised(tmp_path: Path, healthcheck: WorkerHealthcheck) -> Iterator[tuple[OwnedProxy, int]]: + script: Final = tmp_path / "supervisor.py" + script.write_text(_SUPERVISOR) + with subprocess.Popen( + (sys.executable, "-P", str(script)), stdout=subprocess.PIPE, start_new_session=True + ) as supervisor: + try: + assert supervisor.stdout is not None + worker: Final = int(supervisor.stdout.readline()) + with httpx.Client(base_url="http://127.0.0.1:9") as client: + owned: Final = OwnedProxy( + Gateway(client, "sk-unit", "http://127.0.0.1:9"), supervisor, tmp_path / "proxy.log", healthcheck + ) + eventually(owned.worker_pids, lambda pids: pids == frozenset((worker,)), seconds=30) + yield owned, worker + finally: + stop_process_group(supervisor, 5) + + +def test_worker_pids_are_the_supervisors_spawned_workers_and_not_the_supervisor(tmp_path: Path) -> None: + with _supervised(tmp_path, WorkerHealthcheck.STRICT) as (owned, worker): + assert owned.worker_pids() == frozenset((worker,)) + assert owned.process.pid not in owned.worker_pids() + + +def test_only_a_serving_worker_of_this_proxy_can_be_killed(tmp_path: Path) -> None: + with _supervised(tmp_path, WorkerHealthcheck.STRICT) as (owned, worker): + with pytest.raises(AssertionError, match="not a serving worker"): + owned.kill_worker(owned.process.pid) + assert owned.process.poll() is None + assert psutil.Process(worker).is_running() + + +def _leave_without_killing(owned: OwnedProxy, worker: int) -> None: + with owned.replacing_worker(worker): + pass + + +def test_a_replacing_block_that_never_kills_its_worker_fails_and_leaves_it_serving(tmp_path: Path) -> None: + with _supervised(tmp_path, WorkerHealthcheck.STRICT) as (owned, worker): + with pytest.raises(AssertionError, match="without killing the worker"): + _leave_without_killing(owned, worker) + assert owned.worker_pids() == frozenset((worker,)) + + +def _kill_of(owned: OwnedProxy, worker: int) -> WorkerKill: + return WorkerKill(psutil.Process(worker), psutil.Process(owned.process.pid), 5, threading.Event()) + + +def test_a_worker_kill_returns_once_the_supervisor_reaped_the_victim_and_cannot_fire_twice(tmp_path: Path) -> None: + with _supervised(tmp_path, WorkerHealthcheck.STRICT) as (owned, worker): + kill: Final = _kill_of(owned, worker) + kill() + assert not psutil.pid_exists(worker) + assert kill.fired.is_set() + with pytest.raises(AssertionError, match="already killed"): + kill() + + +def test_a_frozen_victim_stays_stopped_until_the_kill(tmp_path: Path) -> None: + with _supervised(tmp_path, WorkerHealthcheck.STRICT) as (owned, worker): + kill: Final = _kill_of(owned, worker) + assert kill.freeze().status() == psutil.STATUS_STOPPED + kill() + assert not psutil.pid_exists(worker) + + +def test_a_kill_under_a_suspended_supervisor_returns_with_the_victim_left_unreaped(tmp_path: Path) -> None: + with _supervised(tmp_path, WorkerHealthcheck.LENIENT) as (owned, worker): + kill: Final = _kill_of(owned, worker) + with owned.suspended(owned.process.pid): + kill() + assert psutil.Process(worker).status() == psutil.STATUS_ZOMBIE + eventually(lambda: psutil.pid_exists(worker), lambda exists: not exists, seconds=5) + + +def test_a_kill_fails_when_the_supervisor_never_reaps_the_victim(tmp_path: Path) -> None: + with _supervised(tmp_path, WorkerHealthcheck.STRICT) as (owned, worker): + unreaping: Final = psutil.Process(owned.process.pid) + kill: Final = WorkerKill(psutil.Process(worker), psutil.Process(os.getpid()), 0.5, threading.Event()) + unreaping.suspend() + try: + with pytest.raises(AssertionError, match="did not reap"): + kill() + finally: + unreaping.resume() + assert not kill.fired.is_set() + + +def test_a_strict_healthcheck_proxy_refuses_suspension_before_stopping_anything(tmp_path: Path) -> None: + with _supervised(tmp_path, WorkerHealthcheck.STRICT) as (owned, worker): + with pytest.raises(AssertionError, match="LENIENT"): + with owned.suspended(worker): + pass + assert psutil.Process(worker).status() != psutil.STATUS_STOPPED + + +def test_a_lenient_healthcheck_proxy_suspends_for_the_block_and_resumes_after(tmp_path: Path) -> None: + with _supervised(tmp_path, WorkerHealthcheck.LENIENT) as (owned, worker): + with owned.suspended(worker) as suspended: + assert suspended.status() == psutil.STATUS_STOPPED + assert psutil.Process(worker).status() != psutil.STATUS_STOPPED + assert owned.worker_pids() == frozenset((worker,)) + + +def test_a_worker_killed_while_suspended_fails_the_block(tmp_path: Path) -> None: + with _supervised(tmp_path, WorkerHealthcheck.LENIENT) as (owned, worker): + with pytest.raises(AssertionError, match="died while suspended"): + with owned.suspended(worker): + kill_stray_process(worker) + + +def test_the_worker_healthcheck_cannot_be_overridden_through_extra_arguments(tmp_path: Path) -> None: + with pytest.raises(AssertionError, match="healthcheck"): + with owned_proxy_process( + MagicMock(spec=Gateway), tmp_path, {}, extra_arguments=("--timeout_worker_healthcheck", "60") + ): + pass + assert not tuple(tmp_path.glob("owned-proxy-*.log")) + + +def test_a_lenient_healthcheck_outlasts_a_strict_one() -> None: + assert healthcheck_seconds(WorkerHealthcheck.LENIENT) > healthcheck_seconds(WorkerHealthcheck.STRICT) + + +@pytest.mark.parametrize( + ("answered", "serving", "replacement"), + ( + (30, frozenset((20, 30)), 30), + (20, frozenset((20, 30)), None), + (10, frozenset((20, 30)), None), + (40, frozenset((20, 30)), None), + (None, frozenset((20, 30)), None), + ), +) +def test_only_a_new_serving_worker_answering_readiness_counts_as_the_replacement( + answered: int | None, serving: frozenset[int], replacement: int | None +) -> None: + assert replacement_pid(frozenset((10, 20)), answered, serving) == replacement diff --git a/tests/unit/integration_support/test_timing.py b/tests/unit/integration_support/test_timing.py new file mode 100644 index 00000000000..7e098e879b6 --- /dev/null +++ b/tests/unit/integration_support/test_timing.py @@ -0,0 +1,57 @@ +from __future__ import annotations + +import time +from typing import Final + +import pytest + +from tests.integration._support.timing import ( + assert_cut_short, + assert_enforced_within, + assert_not_delayed_by, + assert_within_factor_of_control, + timed, +) + + +def test_timed_returns_the_value_and_at_least_the_time_the_operation_took() -> None: + measured: Final = timed(lambda: time.sleep(0.05) or "done") + assert measured.value == "done" + assert measured.seconds >= 0.05 + + +def test_a_burst_slowed_by_load_but_not_by_the_injected_delay_passes() -> None: + assert_not_delayed_by((1.2, 2.4), control=(1.1, 1.9), injected_delay=2) + + +def test_a_burst_that_waited_out_the_injected_delay_fails() -> None: + with pytest.raises(AssertionError, match="injected delay"): + assert_not_delayed_by((1.2, 4.0), control=(1.1, 1.9), injected_delay=2) + + +def test_a_call_cut_short_before_the_injected_stall_passes_and_one_that_waited_fails() -> None: + assert_cut_short(1.5, injected_delay=20) + with pytest.raises(AssertionError, match="no sooner than"): + assert_cut_short(20.0, injected_delay=20) + + +def test_an_enforced_limit_passes_inside_its_slack_and_fails_on_either_side() -> None: + assert_enforced_within(59.5, configured=60, slack=15, not_before=59) + assert_enforced_within(75.0, configured=60, slack=15, not_before=59) + with pytest.raises(AssertionError, match="configured 60s plus 15s"): + assert_enforced_within(300.0, configured=60, slack=15, not_before=59) + with pytest.raises(AssertionError, match="between 59s"): + assert_enforced_within(58.0, configured=60, slack=15, not_before=59) + + +def test_an_enforced_limit_without_a_floor_only_caps_the_elapsed_time() -> None: + assert_enforced_within(0.0, configured=1, slack=7) + with pytest.raises(AssertionError): + assert_enforced_within(8.5, configured=1, slack=7) + + +def test_a_bound_relative_to_the_control_scales_with_a_slower_machine() -> None: + assert_within_factor_of_control(0.9, control=0.1, factor=10) + assert_within_factor_of_control(9.0, control=1.0, factor=10) + with pytest.raises(AssertionError, match="over 10x"): + assert_within_factor_of_control(1.1, control=0.1, factor=10) diff --git a/tests/unit/integration_support/test_workers.py b/tests/unit/integration_support/test_workers.py index 457b953e481..30e55a91b11 100644 --- a/tests/unit/integration_support/test_workers.py +++ b/tests/unit/integration_support/test_workers.py @@ -1,11 +1,22 @@ +import threading +from contextlib import ExitStack from dataclasses import dataclass +from pathlib import Path +from types import MappingProxyType from typing import Final import httpx import pytest from tests.integration._support.client import Gateway -from tests.integration._support.workers import readiness_failure +from tests.integration._support.workers import ( + GatewayProfile, + WorkerStack, + gateway_profile, + readiness_failure, + start_worker_stacks, +) +from tests.integration.run import worker_settings HEALTHY: Final = { "status": "healthy", @@ -49,3 +60,74 @@ def test_readiness_rejects_a_proxy_without_working_shared_redis(field: str, valu primary: Final = Gateway(client, "owned-key", "http://127.0.0.1:8190") failure: Final = readiness_failure(client, primary.upstream_url, (primary,), Subscribers(1)) assert failure == "proxy 0 Redis read/write health" + + +@pytest.mark.parametrize( + ("group", "peer", "cost_map"), + ( + ("management", True, False), + ("mcp", True, False), + ("cost", False, True), + ("security", False, False), + ("extensions", False, False), + ), +) +def test_each_group_gets_the_worker_gateway_its_shared_proxy_used_to_provide( + group: str, peer: bool, cost_map: bool +) -> None: + assert gateway_profile(worker_settings(group, 4)) == GatewayProfile( + peer=peer, replica=False, cost_map=cost_map, coverage_file=None + ) + + +def test_replica_runs_and_coverage_runs_reach_every_worker_gateway() -> None: + profile: Final = gateway_profile( + { + **worker_settings("management", 2), + "INTEGRATION_PROXY_READ_REPLICA_URL": "postgresql://reader@127.0.0.1/circle_test", + "INTEGRATION_COVERAGE_FILE": "/results/coverage/data", + } + ) + assert profile == GatewayProfile( + peer=True, replica=True, cost_map=False, coverage_file=Path("/results/coverage/data") + ) + + +def _stack(name: str, closed: threading.Event) -> WorkerStack: + resources: Final = ExitStack() + resources.callback(closed.set) + return WorkerStack(resources, MappingProxyType({"INTEGRATION_RESULTS_DIR": name})) + + +def test_worker_stacks_start_together_and_return_in_worker_order(tmp_path: Path) -> None: + directories: Final = tuple(tmp_path / f"gw{index}" for index in range(3)) + closed: Final = MappingProxyType({directory: threading.Event() for directory in directories}) + together: Final = threading.Barrier(len(directories), timeout=10) + + def build(directory: Path) -> WorkerStack: + together.wait() + return _stack(directory.name, closed[directory]) + + with ExitStack() as resources: + environments: Final = start_worker_stacks(directories, build, resources, len(directories)) + assert [environment["INTEGRATION_RESULTS_DIR"] for environment in environments] == ["gw0", "gw1", "gw2"] + assert not any(event.is_set() for event in closed.values()) + assert all(event.is_set() for event in closed.values()) + + +def test_a_failed_worker_stack_stops_every_sibling_that_started(tmp_path: Path) -> None: + directories: Final = tuple(tmp_path / f"gw{index}" for index in range(3)) + closed: Final = MappingProxyType({directory: threading.Event() for directory in directories}) + together: Final = threading.Barrier(len(directories), timeout=10) + + def build(directory: Path) -> WorkerStack: + together.wait() + if directory.name == "gw1": + raise RuntimeError("gw1 gateway never became ready") + return _stack(directory.name, closed[directory]) + + with ExitStack() as resources: + with pytest.raises(RuntimeError, match="gw1 gateway never became ready"): + start_worker_stacks(directories, build, resources, len(directories)) + assert closed[directories[0]].is_set() + assert closed[directories[2]].is_set() diff --git a/tests/unit/integrations/otel/test_otel_v2_destinations.py b/tests/unit/integrations/otel/test_otel_v2_destinations.py index 5c9cab88784..068b2aa2e2e 100644 --- a/tests/unit/integrations/otel/test_otel_v2_destinations.py +++ b/tests/unit/integrations/otel/test_otel_v2_destinations.py @@ -36,6 +36,7 @@ from litellm.integrations.otel.model.config import ( ExporterSpec, OpenTelemetryV2Config, is_otel_v2_enabled, + parse_excluded_services, ) from litellm.integrations.otel.model.destination import OtelDestination from litellm.integrations.otel.model.payloads import ( @@ -580,6 +581,101 @@ class TestFanOut: "chat gpt-4", } + def test_team_exclusions_replace_empty_global_and_operator_keeps_all_database_spans(self): + team, inherited, operator = InMemorySpanExporter(), InMemorySpanExporter(), InMemorySpanExporter() + exporters: Final = {"http://team.local": team, "http://inherited.local": inherited} + provider: Final = TracerProvider() + provider.add_span_processor(SimpleSpanProcessor(operator)) + provider.add_span_processor( + TenantFanOutSpanProcessor( + processor_factory=lambda destination: SimpleSpanProcessor(exporters[destination.endpoint]) + ) + ) + tracer: Final = get_tracer(provider, "litellm") + + def run(): + set_request_destinations( + ( + OtelDestination( + endpoint="http://team.local", excluded_db_systems=frozenset({"redis", "postgresql"}) + ), + OtelDestination(endpoint="http://inherited.local"), + ) + ) + with tracer.start_as_current_span("redis read") as span: + span.set_attribute("db.system.name", "redis") + with tracer.start_as_current_span("postgres write") as span: + span.set_attribute("db.system", "postgresql") + with tracer.start_as_current_span("chat gpt-4"): + pass + + in_fresh_context(run) + + assert {span.name for span in team.get_finished_spans()} == {"chat gpt-4"} + assert {span.name for span in inherited.get_finished_spans()} == { + "redis read", + "postgres write", + "chat gpt-4", + } + assert {span.name for span in operator.get_finished_spans()} == { + "redis read", + "postgres write", + "chat gpt-4", + } + + def test_none_and_partial_team_exclusions_replace_the_global_set(self): + none, redis_only, inherited, operator = (InMemorySpanExporter() for _ in range(4)) + exporters: Final = { + "http://none.local": none, + "http://redis.local": redis_only, + "http://inherited.local": inherited, + } + provider: Final = TracerProvider() + provider.add_span_processor(SimpleSpanProcessor(operator)) + provider.add_span_processor( + TenantFanOutSpanProcessor( + processor_factory=lambda destination: SimpleSpanProcessor(exporters[destination.endpoint]), + excluded_db_systems=frozenset({"redis", "postgresql"}), + ) + ) + tracer: Final = get_tracer(provider, "litellm") + + def run(): + set_request_destinations( + ( + OtelDestination(endpoint="http://none.local", excluded_db_systems=frozenset()), + OtelDestination(endpoint="http://redis.local", excluded_db_systems=frozenset({"redis"})), + OtelDestination(endpoint="http://inherited.local"), + ) + ) + with tracer.start_as_current_span("redis read") as span: + span.set_attribute("db.system.name", "redis") + with tracer.start_as_current_span("postgres write") as span: + span.set_attribute("db.system", "postgresql") + with tracer.start_as_current_span("chat gpt-4"): + pass + + in_fresh_context(run) + + assert {span.name for span in none.get_finished_spans()} == { + "redis read", + "postgres write", + "chat gpt-4", + } + assert {span.name for span in redis_only.get_finished_spans()} == {"postgres write", "chat gpt-4"} + assert {span.name for span in inherited.get_finished_spans()} == {"chat gpt-4"} + assert {span.name for span in operator.get_finished_spans()} == { + "redis read", + "postgres write", + "chat gpt-4", + } + + def test_excluded_system_policy_does_not_change_exporter_cache_identity(self): + inherited: Final = OtelDestination(endpoint="http://same.local") + overridden: Final = OtelDestination(endpoint="http://same.local", excluded_db_systems=frozenset()) + + assert inherited.cache_key() == overridden.cache_key() + def test_a_team_naming_two_backends_gets_the_trace_at_both(self): """The fan-out rides one provider, so it cannot skip a destination on the grounds that some other backend owns it: nothing else would deliver it.""" @@ -4727,6 +4823,83 @@ class TestCaptureMessageContent: assert destinations["langfuse_otel"].capture_message_content == setting assert destinations["arize"].capture_message_content is None, "the setting stays on its own callback" + @pytest.mark.usefixtures("allow_test_hosts") + @pytest.mark.parametrize( + ("setting", "expected"), + [ + ("redis,postgres", frozenset({"redis", "postgresql"})), + ("none", frozenset()), + (None, None), + ], + ) + def test_excluded_services_stays_on_its_destination_and_normalizes_aliases( + self, monkeypatch: pytest.MonkeyPatch, setting: str | None, expected: frozenset[str] | None + ) -> None: + monkeypatch.setenv("LITELLM_OTEL_V2", "true") + is_otel_v2_enabled.cache_clear() + callback_vars: Final = { + "langfuse_public_key": "pk-team", + "langfuse_secret_key": "sk-team", + "langfuse_host": "http://team.local", + **({"excluded_services": setting} if setting is not None else {}), + } + auth: Final = UserAPIKeyAuth( + team_metadata={ + "logging": [ + {"callback_name": "langfuse_otel", "callback_type": "success", "callback_vars": callback_vars}, + { + "callback_name": "arize", + "callback_type": "success", + "callback_vars": {"arize_api_key": "k", "arize_space_id": "s"}, + }, + ] + } + ) + + destinations: Final = {d.callback_name: d for d in resolve_tenant_otel_destinations(auth)} + + assert destinations["langfuse_otel"].excluded_db_systems == expected + assert destinations["arize"].excluded_db_systems is None + + @pytest.mark.usefixtures("allow_test_hosts") + def test_key_excluded_services_is_parsed_onto_its_destination(self, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("LITELLM_OTEL_V2", "true") + is_otel_v2_enabled.cache_clear() + auth: Final = UserAPIKeyAuth( + metadata={ + "logging": [ + { + "callback_name": "langfuse_otel", + "callback_type": "success", + "callback_vars": { + "langfuse_public_key": "pk-key", + "langfuse_secret_key": "sk-key", + "langfuse_host": "http://key.local", + "excluded_services": "redis,postgres", + }, + } + ] + } + ) + + (destination,) = resolve_tenant_otel_destinations(auth) + + assert destination.excluded_db_systems == frozenset({"redis", "postgresql"}) + + @pytest.mark.parametrize( + ("value", "expected"), + [ + (None, None), + (17, None), + ("none", frozenset()), + ("redis,postgres", frozenset({"redis", "postgresql"})), + ], + ) + def test_parse_excluded_services_returns_only_valid_string_overrides( + self, value: object, expected: frozenset[str] | None + ) -> None: + assert parse_excluded_services(value) == expected + @pytest.mark.usefixtures("allow_test_hosts") @pytest.mark.parametrize("setting", ["event_only", "span_and_event"]) def test_a_stored_event_mode_is_skipped_like_any_invalid_entry( @@ -4805,6 +4978,24 @@ class TestCaptureMessageContent: assert flattened.callback_vars == {"langfuse_public_key": "pk"} + def test_a_team_wide_callback_vars_map_cannot_carry_excluded_services(self) -> None: + with pytest.raises(ValueError, match="Invalid callback variable: excluded_services"): + TeamCallbackMetadata( + success_callback=["langfuse_otel", "signoz"], callback_vars={"excluded_services": "redis"} + ) + + def test_flattening_excluded_services_leaves_it_on_its_own_entry(self) -> None: + flattened: Final = convert_key_logging_metadata_to_callback( + AddTeamCallback( + callback_name="langfuse_otel", + callback_type="success", + callback_vars={"langfuse_public_key": "pk", "excluded_services": "redis"}, + ), + None, + ) + + assert flattened.callback_vars == {"langfuse_public_key": "pk"} + class TestArizeProjectRouting: """Arize rejects an export that names no project (400: set x-project-name, diff --git a/tests/unit/integrations/test_anthropic_cache_control_hook.py b/tests/unit/integrations/test_anthropic_cache_control_hook.py index f596c2adea1..5c2b9412d16 100644 --- a/tests/unit/integrations/test_anthropic_cache_control_hook.py +++ b/tests/unit/integrations/test_anthropic_cache_control_hook.py @@ -3070,7 +3070,7 @@ class TestOpenAIPromptCacheBreakpoint: ], ) def test_dialect_resolution(self, model, provider, expected): - assert AnthropicCacheControlHook._targets_openai_prompt_cache_breakpoint(model, provider) is expected + assert AnthropicCacheControlHook.targets_openai_prompt_cache_breakpoint(model, provider) is expected def test_count_covers_both_marker_kinds(self): message = { @@ -3476,8 +3476,8 @@ class TestChatPathProviderStamp: from unittest.mock import patch with patch.object(AnthropicCacheControlHook, "_resolve_provider") as resolve: - assert AnthropicCacheControlHook._targets_openai_prompt_cache_breakpoint("gpt-4.1", None) is False - assert AnthropicCacheControlHook._targets_openai_prompt_cache_breakpoint("my-custom-model", None) is False + assert AnthropicCacheControlHook.targets_openai_prompt_cache_breakpoint("gpt-4.1", None) is False + assert AnthropicCacheControlHook.targets_openai_prompt_cache_breakpoint("my-custom-model", None) is False resolve.assert_not_called() def test_litellm_proxy_target_keeps_anthropic_style_markers(self): @@ -3813,7 +3813,7 @@ class TestHostedOpenAIDialectFlag: def test_flagged_non_openai_deployment_is_eligible(self, monkeypatch): self._register(monkeypatch, self.MANTLE_MODEL, "bedrock_mantle") assert ( - AnthropicCacheControlHook._targets_openai_prompt_cache_breakpoint(self.MANTLE_MODEL, "bedrock_mantle") + AnthropicCacheControlHook.targets_openai_prompt_cache_breakpoint(self.MANTLE_MODEL, "bedrock_mantle") is True ) @@ -3821,7 +3821,7 @@ class TestHostedOpenAIDialectFlag: """The api_base check exists to sniff for api.openai.com, which a Bedrock host never is.""" self._register(monkeypatch, self.MANTLE_MODEL, "bedrock_mantle") assert ( - AnthropicCacheControlHook._targets_openai_prompt_cache_breakpoint( + AnthropicCacheControlHook.targets_openai_prompt_cache_breakpoint( self.MANTLE_MODEL, "bedrock_mantle", api_base="https://bedrock-runtime.us-east-1.amazonaws.com", @@ -3832,7 +3832,7 @@ class TestHostedOpenAIDialectFlag: def test_flag_set_false_keeps_the_deployment_ineligible(self, monkeypatch): self._register(monkeypatch, self.MANTLE_MODEL, "bedrock_mantle", flag=False) assert ( - AnthropicCacheControlHook._targets_openai_prompt_cache_breakpoint(self.MANTLE_MODEL, "bedrock_mantle") + AnthropicCacheControlHook.targets_openai_prompt_cache_breakpoint(self.MANTLE_MODEL, "bedrock_mantle") is False ) @@ -3840,14 +3840,14 @@ class TestHostedOpenAIDialectFlag: def test_an_odd_typed_flag_keeps_the_deployment_ineligible(self, monkeypatch, flag): self._register(monkeypatch, self.MANTLE_MODEL, "bedrock_mantle", flag=flag) assert ( - AnthropicCacheControlHook._targets_openai_prompt_cache_breakpoint(self.MANTLE_MODEL, "bedrock_mantle") + AnthropicCacheControlHook.targets_openai_prompt_cache_breakpoint(self.MANTLE_MODEL, "bedrock_mantle") is False ) def test_unflagged_non_openai_deployment_stays_ineligible(self, monkeypatch): self._register(monkeypatch, self.MANTLE_MODEL, "bedrock_mantle", flag=None) assert ( - AnthropicCacheControlHook._targets_openai_prompt_cache_breakpoint(self.MANTLE_MODEL, "bedrock_mantle") + AnthropicCacheControlHook.targets_openai_prompt_cache_breakpoint(self.MANTLE_MODEL, "bedrock_mantle") is False ) @@ -3855,7 +3855,7 @@ class TestHostedOpenAIDialectFlag: """A flagged entry does not license a different provider serving the same model string.""" self._register(monkeypatch, self.MANTLE_MODEL, "bedrock_mantle") assert ( - AnthropicCacheControlHook._targets_openai_prompt_cache_breakpoint(self.MANTLE_MODEL, "azure") is False + AnthropicCacheControlHook.targets_openai_prompt_cache_breakpoint(self.MANTLE_MODEL, "azure") is False ) def test_openai_entries_still_go_through_the_api_base_check(self, monkeypatch): @@ -3863,7 +3863,7 @@ class TestHostedOpenAIDialectFlag: assert litellm.model_cost["gpt-5.6"]["supports_prompt_cache_breakpoint"] is True assert litellm.model_cost["gpt-5.6"]["litellm_provider"] == "openai" assert ( - AnthropicCacheControlHook._targets_openai_prompt_cache_breakpoint( + AnthropicCacheControlHook.targets_openai_prompt_cache_breakpoint( "gpt-5.6", "openai", api_base="https://some-compatible-host.example.com" ) is False @@ -3871,8 +3871,8 @@ class TestHostedOpenAIDialectFlag: def test_azure_hosted_gpt_5_6_remains_ineligible(self): """Regression guard: the openai entry's flag must not leak to another provider.""" - assert AnthropicCacheControlHook._targets_openai_prompt_cache_breakpoint("gpt-5.6", "azure") is False - assert AnthropicCacheControlHook._targets_openai_prompt_cache_breakpoint("azure/gpt-5.6", None) is False + assert AnthropicCacheControlHook.targets_openai_prompt_cache_breakpoint("gpt-5.6", "azure") is False + assert AnthropicCacheControlHook.targets_openai_prompt_cache_breakpoint("azure/gpt-5.6", None) is False REGIONAL_MODEL = "bedrock_mantle/us-east-1/openai.gpt-5.6-sol" @@ -3880,15 +3880,15 @@ class TestHostedOpenAIDialectFlag: """``bedrock_mantle//`` is a documented routing form the map keys without the region.""" self._register(monkeypatch, self.MANTLE_MODEL, "bedrock_mantle") assert ( - AnthropicCacheControlHook._targets_openai_prompt_cache_breakpoint(self.REGIONAL_MODEL, "bedrock_mantle") + AnthropicCacheControlHook.targets_openai_prompt_cache_breakpoint(self.REGIONAL_MODEL, "bedrock_mantle") is True ) - assert AnthropicCacheControlHook._targets_openai_prompt_cache_breakpoint(self.REGIONAL_MODEL, None) is True + assert AnthropicCacheControlHook.targets_openai_prompt_cache_breakpoint(self.REGIONAL_MODEL, None) is True def test_region_prefixed_deployment_honors_a_flag_set_false(self, monkeypatch): self._register(monkeypatch, self.MANTLE_MODEL, "bedrock_mantle", flag=False) assert ( - AnthropicCacheControlHook._targets_openai_prompt_cache_breakpoint(self.REGIONAL_MODEL, "bedrock_mantle") + AnthropicCacheControlHook.targets_openai_prompt_cache_breakpoint(self.REGIONAL_MODEL, "bedrock_mantle") is False ) @@ -3897,31 +3897,31 @@ class TestHostedOpenAIDialectFlag: self._register(monkeypatch, self.MANTLE_MODEL, "bedrock_mantle") gov_model = "bedrock_mantle/us-gov-west-1/openai.gpt-5.6-sol" self._register(monkeypatch, gov_model, "bedrock_mantle", flag=None) - assert AnthropicCacheControlHook._targets_openai_prompt_cache_breakpoint(gov_model, "bedrock_mantle") is False + assert AnthropicCacheControlHook.targets_openai_prompt_cache_breakpoint(gov_model, "bedrock_mantle") is False def test_region_free_entry_does_not_license_another_provider(self, monkeypatch): """The candidate keys are built for the request's provider, so a flagged Mantle row stays Mantle's.""" self._register(monkeypatch, self.MANTLE_MODEL, "bedrock_mantle") assert ( - AnthropicCacheControlHook._targets_openai_prompt_cache_breakpoint("us-east-1/openai.gpt-5.6-sol", "azure") + AnthropicCacheControlHook.targets_openai_prompt_cache_breakpoint("us-east-1/openai.gpt-5.6-sol", "azure") is False ) def test_unmapped_deployment_name_stays_ineligible(self): - assert AnthropicCacheControlHook._targets_openai_prompt_cache_breakpoint("azure/my-deployment", None) is False - assert AnthropicCacheControlHook._targets_openai_prompt_cache_breakpoint("my-deployment", "azure") is False + assert AnthropicCacheControlHook.targets_openai_prompt_cache_breakpoint("azure/my-deployment", None) is False + assert AnthropicCacheControlHook.targets_openai_prompt_cache_breakpoint("my-deployment", "azure") is False def test_bare_name_colliding_with_an_openai_row_reads_its_own_provider_entry(self, monkeypatch): """The Responses layer hands the hook a bare deployment name plus its provider. The openai row keyed by that bare name neither answers for the deployment nor stops the lookup of the provider's own entry.""" self._register(monkeypatch, "gpt-collide", "openai") self._register(monkeypatch, "azure_ai/gpt-collide", "azure_ai") - assert AnthropicCacheControlHook._targets_openai_prompt_cache_breakpoint("gpt-collide", "azure_ai") is True + assert AnthropicCacheControlHook.targets_openai_prompt_cache_breakpoint("gpt-collide", "azure_ai") is True def test_bare_name_colliding_with_an_openai_row_stays_ineligible_without_its_own_flag(self, monkeypatch): self._register(monkeypatch, "gpt-collide", "openai") self._register(monkeypatch, "azure_ai/gpt-collide", "azure_ai", flag=None) - assert AnthropicCacheControlHook._targets_openai_prompt_cache_breakpoint("gpt-collide", "azure_ai") is False + assert AnthropicCacheControlHook.targets_openai_prompt_cache_breakpoint("gpt-collide", "azure_ai") is False def test_the_support_check_reads_the_served_providers_row_for_a_provider_stripped_name(self, monkeypatch): """The chat-to-Responses bridge holds the provider-stripped deployment name and the provider that @@ -3962,7 +3962,7 @@ class TestBedrockMantleGptShipsTheOpenAIDialect: def test_shipped_entry_makes_the_deployment_eligible(self): assert supports_openai_prompt_cache_breakpoint(self.MANTLE_MODEL) is True assert ( - AnthropicCacheControlHook._targets_openai_prompt_cache_breakpoint(self.MANTLE_MODEL, "bedrock_mantle") + AnthropicCacheControlHook.targets_openai_prompt_cache_breakpoint(self.MANTLE_MODEL, "bedrock_mantle") is True ) diff --git a/tests/unit/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py b/tests/unit/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py index dbd9cb470e4..be39e83f989 100644 --- a/tests/unit/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py +++ b/tests/unit/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py @@ -3996,3 +3996,34 @@ def test_get_batch_cost_rates_claude_haiku_5_5_prompt_length_tiers( ) assert (rates.input, rates.output, rates.cache_read, rates.cache_creation) == expected + + +@pytest.mark.parametrize(("speed", "multiplier"), (("fast", 6.0), ("standard", 1.0), (None, 1.0))) +def test_token_type_cost_breakdown_bills_cache_tokens_at_the_provider_speed_multiplier( + monkeypatch: pytest.MonkeyPatch, speed: str | None, multiplier: float +) -> None: + monkeypatch.setitem( + litellm.model_cost, + "fast-cache-model", + { + "input_cost_per_token": 1e-6, + "output_cost_per_token": 2e-6, + "cache_read_input_token_cost": 1e-7, + "cache_creation_input_token_cost": 1.25e-6, + "litellm_provider": "anthropic", + "mode": "chat", + "provider_specific_entry": {"fast": 6.0}, + }, + ) + usage = Usage( + prompt_tokens=800, + completion_tokens=50, + total_tokens=850, + prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=400, cache_creation_tokens=300), + speed=speed, + ) + + breakdown = get_token_type_cost_breakdown(model="fast-cache-model", custom_llm_provider="anthropic", usage=usage) + + assert breakdown.cache_read_cost == pytest.approx(multiplier * 400 * 1e-7) + assert breakdown.cache_creation_cost == pytest.approx(multiplier * 300 * 1.25e-6) diff --git a/tests/unit/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py b/tests/unit/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py index 5a01b28e8ac..320a4eff371 100644 --- a/tests/unit/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py +++ b/tests/unit/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py @@ -591,3 +591,66 @@ _BEDROCK_MANTLE_WEB_SEARCH_MODELS = ( ) _BEDROCK_MANTLE_WEB_SEARCH_RATE = 0.012 + + +@pytest.mark.parametrize( + ("model_key", "web_search_requests", "medium_cost_per_query", "expected_cost"), + [ + ("azure_ai/claude-haiku-4-5", 2, 0.01, 0.02), + ("azure_ai/claude-sonnet-4-5", 3, 0.01, 0.03), + ("azure_ai/claude-haiku-4-5", 1, 0.01, 0.01), + ("azure_ai/claude-haiku-4-5", 2, 0.0, 0.0), + ("4d696a77-1451-438a-85be-9750c7fb4eb7", 2, 0.01, 0.02), + ("4d696a77-1451-438a-85be-9750c7fb4eb7", 3, 0.01, 0.03), + ("4d696a77-1451-438a-85be-9750c7fb4eb7", 2, 0.0, 0.0), + ], +) +def test_azure_ai_claude_web_search_billed_per_server_tool_use_request( + model_key: str, web_search_requests: int, medium_cost_per_query: float, expected_cost: float +) -> None: + from litellm.llms import get_cost_for_web_search_request + from litellm.types.utils import ModelInfo, SearchContextCostPerQuery, ServerToolUse, Usage + + model_info = ModelInfo( + key=model_key, + litellm_provider="azure_ai", + mode="chat", + max_tokens=None, + max_input_tokens=None, + max_output_tokens=None, + input_cost_per_token=1e-6, + output_cost_per_token=5e-6, + search_context_cost_per_query=SearchContextCostPerQuery(search_context_size_medium=medium_cost_per_query), + ) + usage = Usage( + prompt_tokens=1840, + completion_tokens=412, + total_tokens=2252, + server_tool_use=ServerToolUse(web_search_requests=web_search_requests), + ) + + cost = get_cost_for_web_search_request(custom_llm_provider="azure_ai", usage=usage, model_info=model_info) + + assert cost == pytest.approx(expected_cost), ( + f"azure_ai claude must bill {web_search_requests} searches at {medium_cost_per_query} each, got {cost}" + ) + + +def test_azure_ai_non_claude_web_search_keeps_generic_fallback() -> None: + from litellm.llms import get_cost_for_web_search_request + from litellm.types.utils import ModelInfo, SearchContextCostPerQuery, Usage + + model_info = ModelInfo( + key="azure_ai/gpt-5.4-mini-2026-03-17", + litellm_provider="azure_ai", + mode="chat", + max_tokens=None, + max_input_tokens=None, + max_output_tokens=None, + input_cost_per_token=1e-6, + output_cost_per_token=5e-6, + search_context_cost_per_query=SearchContextCostPerQuery(search_context_size_medium=0.01), + ) + usage = Usage(prompt_tokens=1840, completion_tokens=412, total_tokens=2252) + + assert get_cost_for_web_search_request(custom_llm_provider="azure_ai", usage=usage, model_info=model_info) is None diff --git a/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py index 1d19264fb9e..c040b161ba4 100644 --- a/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py +++ b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py @@ -12,6 +12,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( TOOL_RESULT_IMAGE_BOUNDARY, TOOL_RESULT_IMAGE_PLACEHOLDER, add_system_prompt_to_messages, + cache_control_as_prompt_cache_breakpoint, encrypted_content_from_signature, encrypted_reasoning_signature, get_file_ids_from_messages, @@ -2263,3 +2264,19 @@ def test_get_semantic_cache_prompt_from_messages_matches_get_str_from_messages_w messages: list[dict[str, object]], ) -> None: assert get_semantic_cache_prompt_from_messages(messages) == get_str_from_messages(messages) + + +def test_cache_control_is_replaced_by_an_explicit_prompt_cache_breakpoint_on_the_same_block() -> None: + marked: Final = {"type": "text", "text": "rules", "cache_control": {"type": "ephemeral", "ttl": "1h"}} + + assert cache_control_as_prompt_cache_breakpoint(marked) == { + "type": "text", + "text": "rules", + "prompt_cache_breakpoint": {"mode": "explicit"}, + } + + +def test_a_block_without_cache_control_is_returned_untouched() -> None: + unmarked: Final = {"type": "text", "text": "rules"} + + assert cache_control_as_prompt_cache_breakpoint(unmarked) is unmarked diff --git a/tests/unit/litellm_core_utils/test_cost_discount.py b/tests/unit/litellm_core_utils/test_cost_discount.py new file mode 100644 index 00000000000..79abaa2b563 --- /dev/null +++ b/tests/unit/litellm_core_utils/test_cost_discount.py @@ -0,0 +1,180 @@ +import pytest + +from litellm.litellm_core_utils.cost_discount import ( + parse_cost_discount_key, + resolve_cost_discount, +) + + +def test_parse_cost_discount_key_bare_provider(): + parsed = parse_cost_discount_key("vertex_ai") + assert parsed.provider == "vertex_ai" + assert parsed.model_pattern is None + + +def test_parse_cost_discount_key_splits_on_first_slash(): + parsed = parse_cost_discount_key("vertex_ai/claude-*") + assert parsed.provider == "vertex_ai" + assert parsed.model_pattern == "claude-*" + + +def test_parse_cost_discount_key_pattern_containing_slash_stays_in_pattern(): + parsed = parse_cost_discount_key("vertex_ai/a/b") + assert parsed.provider == "vertex_ai" + assert parsed.model_pattern == "a/b" + + +def test_parse_cost_discount_key_empty_pattern_after_slash(): + parsed = parse_cost_discount_key("vertex_ai/") + assert parsed.provider == "vertex_ai" + assert parsed.model_pattern == "" + + +def test_resolve_cost_discount_exact_key_beats_glob_and_bare(): + config = {"vertex_ai/claude-sonnet-4-5": 0.3, "vertex_ai/claude-*": 0.2, "vertex_ai": 0.05} + assert resolve_cost_discount(config, "vertex_ai", "claude-sonnet-4-5") == 0.3 + + +def test_resolve_cost_discount_glob_beats_bare_provider(): + config = {"vertex_ai/claude-*": 0.2, "vertex_ai": 0.05} + assert resolve_cost_discount(config, "vertex_ai", "claude-sonnet-4-5") == 0.2 + + +def test_resolve_cost_discount_longest_literal_glob_wins(): + config = {"vertex_ai/claude-*": 0.2, "vertex_ai/claude-sonnet-*": 0.25, "vertex_ai": 0.05} + assert resolve_cost_discount(config, "vertex_ai", "claude-sonnet-4-5") == 0.25 + + +def test_resolve_cost_discount_strips_provider_prefix_from_model(): + config = {"vertex_ai/claude-*": 0.2} + assert resolve_cost_discount(config, "vertex_ai", "vertex_ai/claude-sonnet-4-5") == 0.2 + + +@pytest.mark.parametrize("model", ["openai/openai/gpt-4", "openai/openai/openai/gpt-4"]) +def test_resolve_cost_discount_collapses_repeated_provider_prefix(model: str): + config = {"openai": 0.05, "openai/gpt-*": 0.20} + assert resolve_cost_discount(config, "openai", model) == 0.20 + + +def test_resolve_cost_discount_prefixed_pattern_beats_stripped_glob(): + config = {"openrouter": 0.05, "openrouter/openrouter/aurora-*": 0.30, "openrouter/aurora-*": 0.20} + assert resolve_cost_discount(config, "openrouter", "openrouter/openrouter/aurora-alpha") == 0.30 + + +def test_resolve_cost_discount_fully_stripped_name_still_matches(): + config = {"openrouter": 0.05, "openrouter/aurora-*": 0.20} + assert resolve_cost_discount(config, "openrouter", "openrouter/openrouter/aurora-alpha") == 0.20 + + +def test_resolve_cost_discount_exact_prefixed_name_beats_stripped_exact(): + config = {"openrouter/openrouter/auto": 0.40, "openrouter/auto": 0.10} + assert resolve_cost_discount(config, "openrouter", "openrouter/openrouter/auto") == 0.40 + assert resolve_cost_discount(config, "openrouter", "openrouter/auto") == 0.10 + + +def test_resolve_cost_discount_longest_literal_across_candidate_names(): + config = {"openai/*": 0.10, "openai/gpt-*": 0.20} + assert resolve_cost_discount(config, "openai", "openai/openai/gpt-4") == 0.20 + + +def test_resolve_cost_discount_glob_skips_prefixed_candidate_names(): + config = {"openai/o*": 0.30, "openai/gpt-*": 0.20} + assert resolve_cost_discount(config, "openai", "openai/openai/gpt-4") == 0.20 + + +def test_resolve_cost_discount_prefixless_glob_only_sees_fully_stripped_name(): + config = {"openai": 0.05, "openai/o*": 0.30} + assert resolve_cost_discount(config, "openai", "openai/openai/gpt-4") == 0.05 + + +def test_resolve_cost_discount_slash_glob_sees_prefixed_candidate_names(): + config = {"openrouter": 0.05, "openrouter/*/auto": 0.25} + assert resolve_cost_discount(config, "openrouter", "openrouter/openrouter/auto") == 0.25 + + +def test_resolve_cost_discount_non_matching_model_falls_back_to_bare(): + config = {"vertex_ai/claude-*": 0.2, "vertex_ai": 0.05} + assert resolve_cost_discount(config, "vertex_ai", "gemini-3-pro-preview") == 0.05 + + +def test_resolve_cost_discount_non_matching_model_no_bare_returns_none(): + config = {"vertex_ai/claude-*": 0.2} + assert resolve_cost_discount(config, "vertex_ai", "gemini-3-pro-preview") is None + + +def test_resolve_cost_discount_other_provider_untouched(): + config = {"vertex_ai/claude-*": 0.2, "vertex_ai": 0.05} + assert resolve_cost_discount(config, "openai", "claude-sonnet-4-5") is None + + +def test_resolve_cost_discount_none_provider_returns_none(): + config = {"vertex_ai": 0.05} + assert resolve_cost_discount(config, None, "claude-sonnet-4-5") is None + + +def test_resolve_cost_discount_empty_provider_returns_none(): + config = {"vertex_ai": 0.05} + assert resolve_cost_discount(config, "", "claude-sonnet-4-5") is None + + +def test_resolve_cost_discount_none_model_only_bare_matches(): + config = {"vertex_ai/claude-*": 0.2, "vertex_ai": 0.05} + assert resolve_cost_discount(config, "vertex_ai", None) == 0.05 + assert resolve_cost_discount({"vertex_ai/claude-*": 0.2}, "vertex_ai", None) is None + + +def test_resolve_cost_discount_char_class_counts_as_wildcard(): + config = {"vertex_ai/claude-sonnet-*": 0.3, "vertex_ai/claude-[abcdefghijklmnopqrstuvwxyz]*": 0.1} + assert resolve_cost_discount(config, "vertex_ai", "claude-sonnet-4-5") == 0.3 + + +def test_resolve_cost_discount_unclosed_bracket_matches_literally(): + config = {"vertex_ai/weird[": 0.2} + assert resolve_cost_discount(config, "vertex_ai", "weird[") == 0.2 + assert resolve_cost_discount(config, "vertex_ai", "weirdx") is None + + +def test_resolve_cost_discount_unclosed_negated_class_counts_as_literal(): + config = {"vertex_ai/gemini-?*": 0.2, "vertex_ai/gemini-[!]*": 0.3} + assert resolve_cost_discount(config, "vertex_ai", "gemini-[!]") == 0.3 + + +def test_resolve_cost_discount_char_class_with_leading_bracket(): + config = {"vertex_ai/a[]x]*": 0.1, "vertex_ai/a]*": 0.2} + assert resolve_cost_discount(config, "vertex_ai", "a]q") == 0.2 + assert resolve_cost_discount(config, "vertex_ai", "axq") == 0.1 + + +def test_resolve_cost_discount_glob_crosses_slash_in_model(): + config = {"bedrock/*anthropic.claude-*": 0.15, "bedrock": 0.05} + assert resolve_cost_discount(config, "bedrock", "bedrock/us-east-1/anthropic.claude-v2:1") == 0.15 + + +def test_resolve_cost_discount_exact_pattern_key_does_not_count_as_exact(): + config = {"vertex_ai/claude-*": 0.2, "vertex_ai": 0.05} + assert resolve_cost_discount(config, "vertex_ai", "vertex_ai") == 0.05 + + +def test_resolve_cost_discount_region_stripped_name_matches_pattern(): + config = {"bedrock": 0.05, "bedrock/claude-*": 0.2} + assert resolve_cost_discount(config, "bedrock", "bedrock/us-gov-west-1/claude-x", "us-gov-west-1") == 0.2 + + +def test_resolve_cost_discount_region_stripped_exact_beats_glob(): + config = {"bedrock/claude-x": 0.3, "bedrock/claude-*": 0.2} + assert resolve_cost_discount(config, "bedrock", "bedrock/us-gov-west-1/claude-x", "us-gov-west-1") == 0.3 + + +def test_resolve_cost_discount_region_qualified_pattern_wins_longest_literal(): + config = {"bedrock": 0.05, "bedrock/us-gov-west-1/*": 0.4, "bedrock/claude-*": 0.2} + assert resolve_cost_discount(config, "bedrock", "bedrock/us-gov-west-1/claude-x", "us-gov-west-1") == 0.4 + + +def test_resolve_cost_discount_prefixless_glob_cannot_see_region_prefix(): + config = {"bedrock": 0.05, "bedrock/us*": 0.9} + assert resolve_cost_discount(config, "bedrock", "bedrock/us-gov-west-1/claude-x", "us-gov-west-1") == 0.05 + + +def test_resolve_cost_discount_no_region_name_keeps_qualified_name_hidden(): + config = {"bedrock": 0.05, "bedrock/claude-*": 0.2} + assert resolve_cost_discount(config, "bedrock", "bedrock/us-gov-west-1/claude-x") == 0.05 diff --git a/tests/unit/litellm_core_utils/test_decisions_pricing.py b/tests/unit/litellm_core_utils/test_decisions_pricing.py new file mode 100644 index 00000000000..e9814b980d1 --- /dev/null +++ b/tests/unit/litellm_core_utils/test_decisions_pricing.py @@ -0,0 +1,122 @@ +from typing import Final, cast + +import pytest + +import litellm +from litellm import cost_calculator +from litellm.litellm_core_utils.decisions_pricing import ( + apply_decisions_pricing, + decisions_pricing_fields, + is_decisions_call_type, +) +from litellm.types.utils import CallTypesLiteral, PromptTokensDetailsWrapper, Usage + +_ROW: Final[dict[str, object]] = { + "input_cost_per_token": 1e-07, + "input_cost_per_token_above_272k_tokens": 2e-07, + "output_cost_per_token": 5e-07, + "output_cost_per_token_decisions": 0.0, + "cache_read_input_token_cost": 1e-08, + "cache_read_input_token_cost_decisions": 1e-07, + "litellm_provider": "openai", + "mode": "chat", +} + + +def test_is_decisions_call_type_covers_decisions_and_systemone_only() -> None: + assert all( + is_decisions_call_type(call_type) for call_type in ("decisions", "adecisions", "systemone", "asystemone") + ) + assert not any(is_decisions_call_type(call_type) for call_type in ("completion", "acompletion", "aresponses")) + + +def test_decisions_pricing_fields_maps_suffixed_fields_onto_plain_names() -> None: + assert decisions_pricing_fields(_ROW) == { + "output_cost_per_token": 0.0, + "cache_read_input_token_cost": 1e-07, + } + assert decisions_pricing_fields({"input_cost_per_token": 1e-07, "mode": "evaluation"}) == {} + + +def test_apply_decisions_pricing_overrides_plain_fields_and_keeps_the_rest() -> None: + priced: Final = apply_decisions_pricing(_ROW) + + assert priced == { + "input_cost_per_token": 1e-07, + "input_cost_per_token_above_272k_tokens": 2e-07, + "output_cost_per_token": 0.0, + "cache_read_input_token_cost": 1e-07, + "litellm_provider": "openai", + "mode": "chat", + } + assert _ROW["output_cost_per_token"] == 5e-07 + + +def test_apply_decisions_pricing_leaves_rows_without_decisions_fields_unchanged() -> None: + plain_row: Final[dict[str, object]] = { + "input_cost_per_token": 4e-08, + "output_cost_per_token": 0.0, + "mode": "evaluation", + } + + assert apply_decisions_pricing(plain_row) == plain_row + + +_LUNA_USAGE: Final = Usage( + prompt_tokens=100, + completion_tokens=10, + total_tokens=110, + prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=50), +) + + +@pytest.fixture +def luna_with_decisions_fields(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setitem(litellm.model_cost, "gpt-6-luna", {**litellm.model_cost["gpt-6-luna"], **_ROW}) + + +@pytest.mark.parametrize("call_type", ["decisions", "adecisions", "systemone", "asystemone"]) +def test_cost_per_token_bills_decisions_call_types_from_decisions_fields( + luna_with_decisions_fields: None, call_type: str +) -> None: + prompt_cost, completion_cost = cost_calculator.cost_per_token( + model="gpt-6-luna", + custom_llm_provider="openai", + usage_object=_LUNA_USAGE, + call_type=cast(CallTypesLiteral, call_type), + ) + + assert prompt_cost == pytest.approx(50 * 1e-07 + 50 * 1e-07) + assert completion_cost == 0.0 + + +def test_cost_per_token_keeps_chat_on_plain_fields(luna_with_decisions_fields: None) -> None: + prompt_cost, completion_cost = cost_calculator.cost_per_token( + model="gpt-6-luna", + custom_llm_provider="openai", + usage_object=_LUNA_USAGE, + call_type="completion", + ) + + assert prompt_cost == pytest.approx(50 * 1e-07 + 50 * 1e-08) + assert completion_cost == pytest.approx(10 * 5e-07) + + +def test_cost_per_token_falls_back_to_plain_fields_without_decisions_fields(monkeypatch: pytest.MonkeyPatch) -> None: + plain_row: Final[dict[str, object]] = { + "input_cost_per_token": 4e-08, + "output_cost_per_token": 0.0, + "litellm_provider": "perplexity", + "mode": "evaluation", + } + monkeypatch.setitem(litellm.model_cost, "perplexity/pplx-decider-v1-27b", plain_row) + + prompt_cost, completion_cost = cost_calculator.cost_per_token( + model="pplx-decider-v1-27b", + custom_llm_provider="perplexity", + usage_object=_LUNA_USAGE, + call_type="adecisions", + ) + + assert prompt_cost == pytest.approx(100 * 4e-08) + assert completion_cost == 0.0 diff --git a/tests/unit/litellm_core_utils/test_exception_mapping_utils.py b/tests/unit/litellm_core_utils/test_exception_mapping_utils.py index fa8b437a64b..06aa167d0ba 100644 --- a/tests/unit/litellm_core_utils/test_exception_mapping_utils.py +++ b/tests/unit/litellm_core_utils/test_exception_mapping_utils.py @@ -1247,6 +1247,139 @@ def test_bedrock_mantle_context_overflow_maps_to_context_window_exceeded(): assert "prompt is too long: 1055489 tokens > 1050000 maximum" in excinfo.value.message +def test_bedrock_mantle_openai_envelope_context_overflow_maps_to_context_window_exceeded(): + from litellm.llms.base_llm.chat.transformation import BaseLLMException + + original_exception = BaseLLMException( + status_code=400, + message=( + '{"error":{"code":"context_length_exceeded",' + '"message":"Your input exceeds the context window of this model. ' + 'Please adjust your input and try again.",' + '"param":"input","type":"invalid_request_error"}}' + ), + ) + + with pytest.raises(litellm.ContextWindowExceededError) as excinfo: + exception_type( + model="openai.gpt-5.6-sol", + original_exception=original_exception, + custom_llm_provider="bedrock_mantle", + ) + + assert excinfo.value.status_code == 400 + assert "prompt is too long" in excinfo.value.message + + +def test_bedrock_mantle_streamed_context_overflow_event_maps_to_context_window_exceeded(): + from litellm.responses.streaming_iterator import _map_stream_error_to_exception + + mapped_exception = _map_stream_error_to_exception( + { + "code": "context_length_exceeded", + "message": "Your input exceeds the context window of this model. Please adjust your input and try again.", + }, + model="openai.gpt-5.6-luna", + custom_llm_provider="bedrock_mantle", + ) + + assert isinstance(mapped_exception, litellm.ContextWindowExceededError) + assert mapped_exception.status_code == 400 + assert "prompt is too long" in mapped_exception.message + + +@pytest.mark.parametrize("original_status_code", [500, 503]) +def test_bedrock_mantle_in_stream_context_overflow_answers_400_whatever_the_event_status_was(original_status_code): + original_exception = OpenAIError( + status_code=original_status_code, + message="prompt tokens (229378) exceed model maximum (131072) for google.gemma-4-e2b", + ) + + with pytest.raises(litellm.ContextWindowExceededError) as excinfo: + exception_type( + model="google.gemma-4-e2b", + original_exception=original_exception, + custom_llm_provider="bedrock_mantle", + ) + + assert excinfo.value.status_code == 400 + assert excinfo.value.response.status_code == 400 + + +def test_bedrock_mantle_context_overflow_keeps_the_provider_400_response_body(): + from litellm.llms.base_llm.chat.transformation import BaseLLMException + + body = { + "error": { + "code": "context_length_exceeded", + "message": "Your input exceeds the context window of this model.", + "param": "input", + "type": "invalid_request_error", + } + } + provider_response = httpx.Response( + status_code=400, + json=body, + request=httpx.Request("POST", "https://bedrock-mantle.us-east-1.amazonaws.com/v1/chat/completions"), + ) + original_exception = BaseLLMException(status_code=400, message=json.dumps(body), response=provider_response) + + with pytest.raises(litellm.ContextWindowExceededError) as excinfo: + exception_type( + model="openai.gpt-5.6-sol", + original_exception=original_exception, + custom_llm_provider="bedrock_mantle", + ) + + assert excinfo.value.response is provider_response + assert excinfo.value.response.json() == body + assert excinfo.value.status_code == 400 + + +@pytest.mark.parametrize( + "error_str", + [ + '{"error":{"code":"context_length_exceeded","message":"Your input exceeds the context window."}}', + "Error code: 400 - {'error': {'message': 'Your input exceeds the context window.', 'type': None, " + "'code': 'context_length_exceeded'}}", + '{"message":"BedrockException - {\\"error\\":{\\"code\\":\\"context_length_exceeded\\"}}"}', + ], +) +def test_bedrock_mantle_overflow_code_maps_in_every_serialization_of_the_envelope(error_str): + from litellm.llms.base_llm.chat.transformation import BaseLLMException + + with pytest.raises(litellm.ContextWindowExceededError): + exception_type( + model="openai.gpt-5.6-sol", + original_exception=BaseLLMException(status_code=400, message=error_str), + custom_llm_provider="bedrock_mantle", + ) + + +def test_bedrock_mantle_error_merely_quoting_the_overflow_code_stays_a_plain_400(): + from litellm.llms.base_llm.chat.transformation import BaseLLMException + + original_exception = BaseLLMException( + status_code=400, + message=( + '{"error":{"code":"invalid_request_error",' + '"message":"Invalid input value context_length_exceeded",' + '"param":"input","type":"invalid_request_error"}}' + ), + ) + + with pytest.raises(litellm.BadRequestError) as excinfo: + exception_type( + model="openai.gpt-5.6-sol", + original_exception=original_exception, + custom_llm_provider="bedrock_mantle", + ) + + assert not isinstance(excinfo.value, litellm.ContextWindowExceededError) + assert excinfo.value.status_code == 400 + assert "Invalid input value context_length_exceeded" in excinfo.value.message + + def test_branchless_provider_transport_error_maps_to_api_connection_error(): from litellm.llms.base_llm.chat.transformation import BaseLLMException diff --git a/tests/unit/litellm_core_utils/test_health_check_helpers.py b/tests/unit/litellm_core_utils/test_health_check_helpers.py index 123be0bd427..4f62c6fb096 100644 --- a/tests/unit/litellm_core_utils/test_health_check_helpers.py +++ b/tests/unit/litellm_core_utils/test_health_check_helpers.py @@ -5,7 +5,7 @@ import socket import struct import zlib from types import MappingProxyType -from typing import Final +from typing import Final, Literal from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -726,9 +726,11 @@ async def test_ahealth_check_probes_hosted_vllm_with_a_choice_question( @pytest.mark.asyncio +@pytest.mark.parametrize("mode", ("evaluation", "decisions")) async def test_ahealth_check_keeps_the_noul_probe_for_other_decisions_providers( monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter, + mode: Literal["evaluation", "decisions"], ) -> None: monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) litellm.in_memory_llm_clients_cache.flush_cache() @@ -745,7 +747,7 @@ async def test_ahealth_check_keeps_the_noul_probe_for_other_decisions_providers( "model": "typesafe/jev-1.13", "api_key": "sk-test", }, - mode="evaluation", + mode=mode, ) assert "error" not in result, result diff --git a/tests/unit/litellm_core_utils/test_litellm_logging.py b/tests/unit/litellm_core_utils/test_litellm_logging.py index 7dbb7cc3924..870378e93f6 100644 --- a/tests/unit/litellm_core_utils/test_litellm_logging.py +++ b/tests/unit/litellm_core_utils/test_litellm_logging.py @@ -4551,6 +4551,103 @@ def test_process_hidden_params_recalculates_cost_after_failure_handler_zero(): assert slo.get("response_cost", 0) > 0 +@pytest.mark.parametrize("reported_cost", [0.0, 0.37]) +@pytest.mark.parametrize("cache_hit", [False, True]) +def test_provider_reported_cost_reaches_metadata_and_spend(reported_cost: float, cache_hit: bool) -> None: + started: Final = datetime_standard_logging(2026, 1, 1) + logging_obj: Final = LitellmLogging( + model="native-test", + messages=[{"role": "user", "content": "test"}], + stream=False, + call_type="anthropic_messages", + start_time=started, + litellm_call_id="reported-cost", + function_id="reported-cost", + ) + logging_obj.update_from_kwargs(kwargs={}, optional_params={}) + logging_obj._provider_reported_cost = reported_cost + logging_obj.model_call_details.update({"cache_hit": cache_hit, "litellm_params": {}}) + response: Final = ModelResponse( + id="reported-cost", + model="native-test", + choices=[{"message": {"role": "assistant", "content": "ok"}}], + ) + expected: Final = 0.0 if cache_hit else reported_cost + assert logging_obj.response_cost_calculator(response) == expected + logging_obj._process_hidden_params_and_response_cost(response, started, started) + assert logging_obj.model_call_details["response_cost"] == expected + assert logging_obj.model_call_details["standard_logging_object"]["response_cost"] == expected + + +@pytest.mark.parametrize("pricing_policy", ["reported", "custom", "base"]) +@pytest.mark.parametrize("reported_cost", [None, 0.0, 0.37]) +@pytest.mark.parametrize("additional_cost", [0.0, 0.12]) +@pytest.mark.parametrize("cache_hit", [False, True]) +def test_reported_cost_respects_configured_pricing_and_additional_charges( + monkeypatch: pytest.MonkeyPatch, + pricing_policy: Literal["reported", "custom", "base"], + reported_cost: float | None, + additional_cost: float, + cache_hit: bool, +) -> None: + pricing_model: Final = "anthropic/reported-cost-pricing-test" + input_rate: Final = 0.002 + output_rate: Final = 0.004 + prompt_tokens: Final = 10 + completion_tokens: Final = 5 + monkeypatch.setitem( + litellm.model_cost, + pricing_model, + {"input_cost_per_token": input_rate, "output_cost_per_token": output_rate, "litellm_provider": "anthropic"}, + ) + started: Final = datetime_standard_logging(2026, 1, 1) + logging_obj: Final = LitellmLogging( + model=pricing_model, + messages=[{"role": "user", "content": "test"}], + stream=False, + call_type="anthropic_messages", + start_time=started, + litellm_call_id="reported-cost-pricing", + function_id="reported-cost-pricing", + ) + logging_obj.update_from_kwargs( + kwargs={}, + optional_params={"provider_reported_cost": 0.0, "_provider_reported_cost": 0.0}, + litellm_params={ + **({"base_model": pricing_model} if pricing_policy == "base" else {}), + "metadata": { + "model_info": { + "id": pricing_model, + **( + {"input_cost_per_token": input_rate, "output_cost_per_token": output_rate} + if pricing_policy == "custom" + else {} + ), + } + } + } + ) + logging_obj.model_call_details.update( + {"additional_response_cost": additional_cost, "cache_hit": cache_hit} + ) + logging_obj._provider_reported_cost = reported_cost + response: Final = ModelResponse( + model=pricing_model, + choices=[{"message": {"role": "assistant", "content": "ok"}}], + usage=Usage(prompt_tokens=prompt_tokens, completion_tokens=completion_tokens), + ) + base_cost: Final = ( + prompt_tokens * input_rate + completion_tokens * output_rate + if pricing_policy != "reported" or reported_cost is None + else reported_cost + ) + expected: Final = 0.0 if cache_hit else base_cost + additional_cost + assert logging_obj.response_cost_calculator(response) == pytest.approx(expected) + logging_obj._process_hidden_params_and_response_cost(response, started, started) + assert logging_obj.model_call_details["response_cost"] == pytest.approx(expected) + assert logging_obj.model_call_details["standard_logging_object"]["response_cost"] == pytest.approx(expected) + + def test_process_hidden_params_preserves_zero_cost_in_hidden_params(): """Pass-through handlers often set response_cost on result._hidden_params (including 0).""" from datetime import datetime diff --git a/tests/unit/llms/anthropic/pass_through/adapters/test_streaming_iterator_mid_stream_error.py b/tests/unit/llms/anthropic/pass_through/adapters/test_streaming_iterator_mid_stream_error.py index 2f1152174c2..785e1f1096b 100644 --- a/tests/unit/llms/anthropic/pass_through/adapters/test_streaming_iterator_mid_stream_error.py +++ b/tests/unit/llms/anthropic/pass_through/adapters/test_streaming_iterator_mid_stream_error.py @@ -26,6 +26,8 @@ from unittest.mock import MagicMock import pytest +import litellm + sys.path.insert(0, os.path.abspath("../../../../..")) from litellm.exceptions import MidStreamFallbackError @@ -148,6 +150,34 @@ def test_error_event_preserves_midstream_fallback_error(): assert "internalServerException" in payload["error"]["message"] +def test_error_event_keeps_status_of_mapped_litellm_exception(): + exc = litellm.ContextWindowExceededError( + message="prompt is too long: your prompt exceeds the model's context window", + model="openai.gpt-5.6-luna", + llm_provider="bedrock_mantle", + ) + name, payload = _parse_sse(_mid_stream_error_sse_event(exc)) + assert name == "error" + assert payload["error"]["type"] == "invalid_request_error" + assert "prompt is too long" in payload["error"]["message"] + + +@pytest.mark.parametrize( + "exc", + [ + litellm.BadRequestError( + message="temperature must be in the range [0.0, 2.0]", model="m", llm_provider="gemini" + ), + litellm.AuthenticationError(message="API key not valid", llm_provider="gemini", model="m"), + litellm.NotFoundError(message="model is not found", model="m", llm_provider="gemini"), + ], +) +def test_error_event_reports_other_provider_4xx_as_retriable_500(exc): + name, payload = _parse_sse(_mid_stream_error_sse_event(exc)) + assert name == "error" + assert payload["error"]["type"] == "api_error" + + class _AsyncFailureRecorder(CustomLogger): def __init__(self): super().__init__() diff --git a/tests/unit/llms/anthropic/pass_through/responses_adapters/test_responses_adapters_handler.py b/tests/unit/llms/anthropic/pass_through/responses_adapters/test_responses_adapters_handler.py index eb2d62a8d71..df565ef25b5 100644 --- a/tests/unit/llms/anthropic/pass_through/responses_adapters/test_responses_adapters_handler.py +++ b/tests/unit/llms/anthropic/pass_through/responses_adapters/test_responses_adapters_handler.py @@ -200,3 +200,76 @@ async def test_streaming_hands_the_logging_object_the_message_id_the_caller_is_s message_start = next(e for e in events if e["type"] == "message_start") assert message_start["message"]["id"].startswith("msg_") assert logging_obj.streamed_anthropic_message_id == message_start["message"]["id"] + + +OPENAI_API: Final = {"custom_llm_provider": "openai", "api_base": "https://api.openai.com/v1"} +OTHER_HOST: Final = {"custom_llm_provider": "openai", "api_base": "https://gateway.example/v1"} +EPHEMERAL: Final = {"type": "ephemeral"} + + +def _request_with_client_cache_marks(cache_control: dict[str, str] | None) -> dict[str, object]: + mark: Final = {} if cache_control is None else {"cache_control": cache_control} + return { + "max_tokens": 1024, + "system": [{"type": "text", "text": "Be concise."}, {"type": "text", "text": "Be helpful.", **mark}], + "messages": [ + {"role": "user", "content": [{"type": "text", "text": "look", **mark}, {"type": "text", "text": "now"}]} + ], + } + + +@pytest.mark.parametrize( + "extra_kwargs", + [ + pytest.param(OPENAI_API, id="api.openai.com"), + pytest.param({**OTHER_HOST, "prompt_cache_options": {"mode": "implicit"}}, id="other-host-opted-in"), + ], +) +def test_client_cache_control_reaches_gpt_5_6_as_a_prompt_cache_breakpoint_on_the_same_block( + extra_kwargs: dict[str, object], +) -> None: + responses_kwargs: Final = _build_responses_kwargs( + model="openai/gpt-5.6", + extra_kwargs=extra_kwargs, + **_request_with_client_cache_marks({**EPHEMERAL, "ttl": "1h"}), + ) + + assert "instructions" not in responses_kwargs + assert responses_kwargs["input"] == [ + { + "type": "message", + "role": "developer", + "content": [ + {"type": "input_text", "text": "Be concise."}, + {"type": "input_text", "text": "Be helpful.", "prompt_cache_breakpoint": {"mode": "explicit"}}, + ], + }, + { + "type": "message", + "role": "user", + "content": [ + {"type": "input_text", "text": "look", "prompt_cache_breakpoint": {"mode": "explicit"}}, + {"type": "input_text", "text": "now"}, + ], + }, + ] + + +@pytest.mark.parametrize( + ("model", "extra_kwargs"), + [ + pytest.param("openai/gpt-5.5", OPENAI_API, id="model-older-than-gpt-5.6"), + pytest.param("openai/gpt-5.6", OTHER_HOST, id="other-host-not-opted-in"), + ], +) +def test_client_cache_control_changes_nothing_for_a_deployment_that_does_not_read_prompt_cache_breakpoint( + model: str, extra_kwargs: dict[str, object] +) -> None: + marked: Final = _build_responses_kwargs( + model=model, extra_kwargs=extra_kwargs, **_request_with_client_cache_marks(EPHEMERAL) + ) + unmarked: Final = _build_responses_kwargs( + model=model, extra_kwargs=extra_kwargs, **_request_with_client_cache_marks(None) + ) + + assert marked == unmarked diff --git a/tests/unit/llms/anthropic/pass_through/responses_adapters/test_responses_adapters_transformation.py b/tests/unit/llms/anthropic/pass_through/responses_adapters/test_responses_adapters_transformation.py index 9b6b44c3d05..1532b38bfae 100644 --- a/tests/unit/llms/anthropic/pass_through/responses_adapters/test_responses_adapters_transformation.py +++ b/tests/unit/llms/anthropic/pass_through/responses_adapters/test_responses_adapters_transformation.py @@ -2156,7 +2156,7 @@ class TestPromptCacheBreakpointToResponses: } ] - def test_assistant_and_tool_result_blocks_drop_breakpoint(self): + def test_assistant_blocks_drop_breakpoint_and_a_tool_result_keeps_it_on_its_last_output_part(self): items = _ADAPTER.translate_messages_to_responses_input( [ {"role": "user", "content": [{"type": "text", "text": "q"}]}, @@ -2180,8 +2180,14 @@ class TestPromptCacheBreakpointToResponses: }, ] ) - assert len(items) == 4 - assert not _contains_key(items, "prompt_cache_breakpoint") + *earlier_items, tool_output = items + assert len(earlier_items) == 3 + assert not _contains_key(earlier_items, "prompt_cache_breakpoint") + assert tool_output == { + "type": "function_call_output", + "call_id": "toolu_01", + "output": [{"type": "input_text", "text": "r", "prompt_cache_breakpoint": self.EXPLICIT}], + } def test_prompt_cache_options_forwarded_to_responses_kwargs(self): from litellm.llms.anthropic.pass_through.responses_adapters.handler import ( diff --git a/tests/unit/llms/base_llm/decisions/test_base_decisions_transformation.py b/tests/unit/llms/base_llm/decisions/test_base_decisions_transformation.py index a452a621d2e..454fa393ea8 100644 --- a/tests/unit/llms/base_llm/decisions/test_base_decisions_transformation.py +++ b/tests/unit/llms/base_llm/decisions/test_base_decisions_transformation.py @@ -178,3 +178,23 @@ def test_missing_or_mismatched_systemone_answers_are_refusals_left_out_of_the_re assert parsed.answers[1:] == (DecisionsIRRefusal(), DecisionsIRRefusal(), DecisionsIRRefusal()) assert tuple(ir_to_systemone_response(parsed, ir).answers) == ("damaged",) + + +def test_each_media_kind_is_gated_by_its_own_support_flag() -> None: + audio: Final = ["data:audio/wav;base64,YQ=="] + videos: Final = ["data:video/mp4;base64,YQ=="] + audio_only: Final = systemone_request_to_ir(_SYSTEMONE_BODY.validate_python({**_SYSTEMONE_REQUEST, "audio": audio})) + video_only: Final = systemone_request_to_ir(_SYSTEMONE_BODY.validate_python({**_SYSTEMONE_REQUEST, "videos": videos})) + + assert ir_to_systemone_request("clef-omni", audio_only, supports_audio=True) == { + "model": "clef-omni", + **_SYSTEMONE_REQUEST, + "audio": audio, + } + assert ir_to_systemone_request("clef-omni", video_only, supports_video=True) == { + "model": "clef-omni", + **_SYSTEMONE_REQUEST, + "videos": videos, + } + assert isinstance(ir_to_systemone_request("clef-omni", audio_only, supports_video=True), UnsupportedDecisionsRequest) + assert isinstance(ir_to_systemone_request("clef-omni", video_only, supports_audio=True), UnsupportedDecisionsRequest) diff --git a/tests/unit/llms/bedrock/chat/test_converse_transformation.py b/tests/unit/llms/bedrock/chat/test_converse_transformation.py index c68f8a59333..469d06f7205 100644 --- a/tests/unit/llms/bedrock/chat/test_converse_transformation.py +++ b/tests/unit/llms/bedrock/chat/test_converse_transformation.py @@ -5347,12 +5347,12 @@ def test_parallel_tool_calls_flag_decoupled_from_ttl_pricing(monkeypatch): The disable_parallel_tool_use gate must read supports_parallel_tool_use_config, not the 1h-TTL pricing field: a model carrying only the former still gets the flag. """ - from litellm.llms.bedrock.common_utils import is_claude_4_5_on_bedrock + from litellm.llms.bedrock.common_utils import supports_bedrock_prompt_cache_ttl config = AmazonConverseConfig() model = "anthropic.claude-parallel-tool-use-only" monkeypatch.setitem(litellm.model_cost, model, {"supports_parallel_tool_use_config": True}) - assert is_claude_4_5_on_bedrock(model) is False + assert supports_bedrock_prompt_cache_ttl(model) is False messages = [{"role": "user", "content": "What's the weather in SF and NYC?"}] optional_params = config.map_openai_params( @@ -5924,7 +5924,7 @@ def test_cache_control_injection_tool_config_honors_ttl_for_supported_model(monk cache_control behavior, instead of always emitting a bare {"type": "default"} cachePoint with no ttl. - Forces the bundled local cost map so `is_claude_4_5_on_bedrock` (which + Forces the bundled local cost map so `supports_bedrock_prompt_cache_ttl` (which reads `cache_creation_input_token_cost_above_1hr` from litellm.model_cost) sees this branch's pricing data rather than the network-fetched `main` copy, which lacks it until merge. diff --git a/tests/unit/llms/bedrock/test_bedrock_common_utils.py b/tests/unit/llms/bedrock/test_bedrock_common_utils.py index 3a2c851807c..d10ce9d1fae 100644 --- a/tests/unit/llms/bedrock/test_bedrock_common_utils.py +++ b/tests/unit/llms/bedrock/test_bedrock_common_utils.py @@ -473,7 +473,7 @@ def test_capability_lookups_fall_back_to_base_model_when_regional_entry_lacks_fi import litellm from litellm.llms.bedrock.common_utils import ( bedrock_converse_supports_parallel_tool_use_config, - is_claude_4_5_on_bedrock, + supports_bedrock_prompt_cache_ttl, ) base = "anthropic.claude-fallback-test" @@ -488,10 +488,30 @@ def test_capability_lookups_fall_back_to_base_model_when_regional_entry_lacks_fi }, ) - assert is_claude_4_5_on_bedrock(regional) is True + assert supports_bedrock_prompt_cache_ttl(regional) is True assert bedrock_converse_supports_parallel_tool_use_config(regional) is True +def test_deprecated_is_claude_4_5_on_bedrock_warns_and_matches_new_name(monkeypatch): + import litellm + from litellm.llms.bedrock.common_utils import ( + is_claude_4_5_on_bedrock, + supports_bedrock_prompt_cache_ttl, + ) + + with_ttl = "anthropic.claude-ttl-supported" + without_ttl = "anthropic.claude-ttl-unsupported" + monkeypatch.setitem(litellm.model_cost, with_ttl, {"cache_creation_input_token_cost_above_1hr": 1e-05}) + monkeypatch.setitem(litellm.model_cost, without_ttl, {"input_cost_per_token": 1e-06}) + + with pytest.warns(UserWarning, match="supports_bedrock_prompt_cache_ttl"): + assert is_claude_4_5_on_bedrock(with_ttl) is True + with pytest.warns(UserWarning, match="supports_bedrock_prompt_cache_ttl"): + assert is_claude_4_5_on_bedrock(without_ttl) is False + assert supports_bedrock_prompt_cache_ttl(with_ttl) is True + assert supports_bedrock_prompt_cache_ttl(without_ttl) is False + + @pytest.mark.parametrize( ("entry", "expected"), [ diff --git a/tests/unit/llms/databricks/chat/test_databricks_chat_transformation.py b/tests/unit/llms/databricks/chat/test_databricks_chat_transformation.py index c6e0ecd2b41..fe76a667f3e 100644 --- a/tests/unit/llms/databricks/chat/test_databricks_chat_transformation.py +++ b/tests/unit/llms/databricks/chat/test_databricks_chat_transformation.py @@ -2216,3 +2216,42 @@ def test_completion_with_prompt_caching_anthropic_model( assert response.usage.total_tokens == 1666 assert response.usage.cache_read_input_tokens == 0 assert response.usage.cache_creation_input_tokens == 1545 + + +def test_get_optional_params_keeps_service_tier_for_databricks() -> None: + optional_params: Final = litellm.utils.get_optional_params( + model="databricks-claude-opus-5", + custom_llm_provider="databricks", + service_tier="priority", + drop_params=True, + ) + + assert optional_params["service_tier"] == "priority" + + +@pytest.mark.parametrize("served_tier", ["priority", "default", None]) +def test_transform_response_carries_served_service_tier(served_tier: str | None) -> None: + raw_response: Final = MagicMock() + raw_response.json.return_value = { + "id": "chatcmpl-1", + "object": "chat.completion", + "created": 0, + "model": "dbrx", + **({"service_tier": served_tier} if served_tier is not None else {}), + "choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, + } + + result: Final = DatabricksConfig().transform_response( + model="databricks/dbrx", + raw_response=raw_response, + model_response=litellm.ModelResponse(), + logging_obj=MagicMock(), + request_data={}, + messages=[{"role": "user", "content": "hi"}], + optional_params={}, + litellm_params={}, + encoding=None, + ) + + assert result.to_dict().get("service_tier") == served_tier diff --git a/tests/unit/llms/openai_like/test_coralbricks_provider.py b/tests/unit/llms/openai_like/test_coralbricks_provider.py new file mode 100644 index 00000000000..6187c7aa466 --- /dev/null +++ b/tests/unit/llms/openai_like/test_coralbricks_provider.py @@ -0,0 +1,539 @@ +import json +from dataclasses import dataclass +from pathlib import Path +from typing import Final + +import pytest +import respx + +import litellm +from litellm.caching.llm_caching_handler import LLMClientCache + +CORALBRICKS_MODEL: Final = "coralbricks/deepseek-v4.1-flash-fast" +CORALBRICKS_MODEL_ID: Final = "deepseek-v4.1-flash-fast" + + +def test_coralbricks_provider_resolution(monkeypatch: pytest.MonkeyPatch): + from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider + + monkeypatch.setenv("CORALBRICKS_API_KEY", "coralbricks-test-key") + + model, provider, api_key, api_base = get_llm_provider( + model="coralbricks/glm-5.3-fast", + custom_llm_provider=None, + api_base=None, + api_key=None, + ) + + assert model == "glm-5.3-fast" + assert provider == "coralbricks" + assert api_key == "coralbricks-test-key" + assert api_base == "https://inference.coralbricks.ai/v1" + + +def test_coralbricks_provider_keeps_explicit_credentials(monkeypatch: pytest.MonkeyPatch): + from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider + + monkeypatch.setenv("CORALBRICKS_API_KEY", "coralbricks-env-key") + + _, provider, api_key, api_base = get_llm_provider( + model="coralbricks/glm-5.3-fast", + custom_llm_provider=None, + api_base="https://coralbricks.internal.example/v1", + api_key="coralbricks-explicit-key", + ) + + assert provider == "coralbricks" + assert api_key == "coralbricks-explicit-key" + assert api_base == "https://coralbricks.internal.example/v1" + + +def test_coralbricks_is_available_in_add_model_form(): + fields_path: Final = ( + Path(litellm.__file__).parent / "proxy" / "public_endpoints" / "provider_create_fields.json" + ) + providers: Final = json.loads(fields_path.read_text()) + coralbricks: Final = next( + provider for provider in providers if provider["litellm_provider"] == "coralbricks" + ) + + assert coralbricks["provider"] == "CORALBRICKS" + assert coralbricks["provider_display_name"] == "CoralBricks" + assert coralbricks["default_model_placeholder"] == "coralbricks/glm-5.3-fast" + assert {field["key"]: field["required"] for field in coralbricks["credential_fields"]} == { + "api_base": False, + "api_key": True, + } + + +def test_coralbricks_endpoint_support_matches_providers_json(): + matrix_path: Final = Path(litellm.__file__).parent / "provider_endpoints_support_backup.json" + endpoints: Final = json.loads(matrix_path.read_text())["providers"]["coralbricks"]["endpoints"] + providers_json: Final = json.loads( + (Path(litellm.__file__).parent / "llms" / "openai_like" / "providers.json").read_text() + ) + declared: Final = set(providers_json["coralbricks"]["supported_endpoints"]) + + supported_from_matrix: Final = { + "/v1/chat/completions": endpoints["chat_completions"], + "/v1/responses": endpoints["responses"], + "/v1/messages": endpoints["messages"], + } + assert {endpoint for endpoint, supported in supported_from_matrix.items() if supported} == declared + + +def test_coralbricks_chat_completion_request(): + with respx.mock() as upstream: + route: Final = upstream.post("https://inference.coralbricks.ai/v1/chat/completions").respond( + 200, + json={ + "id": "chatcmpl_coralbricks", + "object": "chat.completion", + "created": 1_789_550_000, + "model": CORALBRICKS_MODEL_ID, + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "Hello from CoralBricks"}, + "finish_reason": "stop", + } + ], + "usage": {"prompt_tokens": 4, "completion_tokens": 3, "total_tokens": 7}, + }, + ) + response: Final = litellm.completion( + model=CORALBRICKS_MODEL, + messages=[{"role": "user", "content": "Say hello"}], + api_key="coralbricks-test-key", + ) + + request: Final = route.calls.last.request + body: Final = json.loads(request.content) + assert route.call_count == 1 + assert str(request.url) == "https://inference.coralbricks.ai/v1/chat/completions" + assert request.headers["authorization"] == "Bearer coralbricks-test-key" + assert body["model"] == CORALBRICKS_MODEL_ID + assert body["messages"] == [{"role": "user", "content": "Say hello"}] + assert response.choices[0].message.content == "Hello from CoralBricks" + + +def test_coralbricks_responses_request(): + with respx.mock() as upstream: + route: Final = upstream.post("https://inference.coralbricks.ai/v1/responses").respond( + 200, + json={ + "id": "resp_coralbricks", + "object": "response", + "created_at": 1_789_550_000, + "model": CORALBRICKS_MODEL_ID, + "status": "completed", + "output": [ + { + "id": "msg_coralbricks", + "type": "message", + "role": "assistant", + "status": "completed", + "content": [ + {"type": "output_text", "text": "Hello from CoralBricks", "annotations": []} + ], + } + ], + "usage": {"input_tokens": 4, "output_tokens": 3, "total_tokens": 7}, + }, + ) + response: Final = litellm.responses( + model=CORALBRICKS_MODEL, + input="Say hello", + api_key="coralbricks-test-key", + ) + + request: Final = route.calls.last.request + body: Final = json.loads(request.content) + assert route.call_count == 1 + assert str(request.url) == "https://inference.coralbricks.ai/v1/responses" + assert request.headers["authorization"] == "Bearer coralbricks-test-key" + assert body["model"] == CORALBRICKS_MODEL_ID + assert body["input"] == "Say hello" + assert response.output[0].content[0].text == "Hello from CoralBricks" + + +@pytest.mark.asyncio +async def test_coralbricks_anthropic_messages_request(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setattr(litellm, "in_memory_llm_clients_cache", LLMClientCache()) + with respx.mock() as upstream: + route: Final = upstream.post("https://inference.coralbricks.ai/v1/messages").respond( + 200, + json={ + "id": "msg_coralbricks", + "type": "message", + "role": "assistant", + "model": CORALBRICKS_MODEL_ID, + "content": [{"type": "text", "text": "Hello from CoralBricks"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 4, "output_tokens": 3}, + }, + ) + response: Final = await litellm.anthropic.messages.acreate( + model=CORALBRICKS_MODEL, + messages=[{"role": "user", "content": "Say hello"}], + max_tokens=32, + api_key="coralbricks-test-key", + ) + + request: Final = route.calls.last.request + body: Final = json.loads(request.content) + assert route.call_count == 1 + assert str(request.url) == "https://inference.coralbricks.ai/v1/messages" + assert request.headers["authorization"] == "Bearer coralbricks-test-key" + assert request.headers["anthropic-version"] == "2023-06-01" + assert body["model"] == CORALBRICKS_MODEL_ID + assert body["messages"] == [{"role": "user", "content": "Say hello"}] + assert response["content"][0]["text"] == "Hello from CoralBricks" + + +@dataclass(frozen=True) +class CoralBricksRates: + input: float + output: float + cache_write: float + cache_read: float + + +def _coralbricks_rates() -> CoralBricksRates: + row: Final = litellm.model_cost[CORALBRICKS_MODEL] + assert row["litellm_provider"] == "coralbricks" + return CoralBricksRates( + input=float(row["input_cost_per_token"]), + output=float(row["output_cost_per_token"]), + cache_write=float(row["cache_creation_input_token_cost"]), + cache_read=float(row["cache_read_input_token_cost"]), + ) + + +def _expected_cost( + rates: CoralBricksRates, + uncached_input_tokens: int, + cached_read_tokens: int, + cache_write_tokens: int, + billed_output_tokens: int, +) -> float: + return ( + uncached_input_tokens * rates.input + + cached_read_tokens * rates.cache_read + + cache_write_tokens * rates.cache_write + + billed_output_tokens * rates.output + ) + + +def test_coralbricks_chat_completion_cost_includes_cache_write_and_free_cached_read(): + rates: Final = _coralbricks_rates() + with respx.mock() as upstream: + upstream.post("https://inference.coralbricks.ai/v1/chat/completions").respond( + 200, + json={ + "id": "chatcmpl_coralbricks", + "object": "chat.completion", + "created": 1_789_550_000, + "model": CORALBRICKS_MODEL_ID, + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "ok"}, + "finish_reason": "stop", + } + ], + "usage": { + "prompt_tokens": 4357, + "completion_tokens": 20, + "total_tokens": 4377, + "prompt_tokens_details": { + "cached_tokens": 4352, + "cache_write_tokens": 5, + "billable_cache_write_tokens": 5, + "cache_write_blocks": 3, + }, + "completion_tokens_details": {"reasoning_tokens": 12}, + }, + }, + ) + response: Final = litellm.completion( + model=CORALBRICKS_MODEL, + messages=[{"role": "user", "content": "Say hello"}], + api_key="coralbricks-test-key", + ) + + cost: Final = litellm.completion_cost(completion_response=response, model=CORALBRICKS_MODEL) + assert cost == pytest.approx( + _expected_cost( + rates, + uncached_input_tokens=4357 - 4352 - 5, + cached_read_tokens=4352, + cache_write_tokens=5, + billed_output_tokens=20, + ) + ) + + +def test_coralbricks_responses_cost_includes_cache_write_and_free_cached_read(): + rates: Final = _coralbricks_rates() + with respx.mock() as upstream: + upstream.post("https://inference.coralbricks.ai/v1/responses").respond( + 200, + json={ + "id": "resp_coralbricks", + "object": "response", + "created_at": 1_789_550_000, + "model": CORALBRICKS_MODEL_ID, + "status": "completed", + "output": [ + { + "id": "msg_coralbricks", + "type": "message", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": "ok", "annotations": []}], + } + ], + "usage": { + "input_tokens": 4357, + "output_tokens": 20, + "total_tokens": 4377, + "input_tokens_details": { + "cached_tokens": 4352, + "cache_write_tokens": 5, + "billable_cache_write_tokens": 5, + }, + "output_tokens_details": {"reasoning_tokens": 12}, + }, + }, + ) + response: Final = litellm.responses( + model=CORALBRICKS_MODEL, + input="Say hello", + api_key="coralbricks-test-key", + ) + + cost: Final = litellm.completion_cost(completion_response=response, model=CORALBRICKS_MODEL) + assert cost == pytest.approx( + _expected_cost( + rates, + uncached_input_tokens=4357 - 4352 - 5, + cached_read_tokens=4352, + cache_write_tokens=5, + billed_output_tokens=20, + ) + ) + + +@pytest.mark.asyncio +async def test_coralbricks_anthropic_messages_cost_includes_cache_write_and_free_cached_read( + monkeypatch: pytest.MonkeyPatch, +): + rates: Final = _coralbricks_rates() + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setattr(litellm, "in_memory_llm_clients_cache", LLMClientCache()) + with respx.mock() as upstream: + upstream.post("https://inference.coralbricks.ai/v1/messages").respond( + 200, + json={ + "id": "msg_coralbricks", + "type": "message", + "role": "assistant", + "model": CORALBRICKS_MODEL_ID, + "content": [{"type": "text", "text": "ok"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": { + "input_tokens": 0, + "cache_read_input_tokens": 4352, + "cache_creation_input_tokens": 5, + "output_tokens": 20, + }, + }, + ) + response: Final = await litellm.anthropic.messages.acreate( + model=CORALBRICKS_MODEL, + messages=[{"role": "user", "content": "Say hello"}], + max_tokens=32, + api_key="coralbricks-test-key", + ) + + cost: Final = litellm.completion_cost(completion_response=response, model=CORALBRICKS_MODEL) + assert cost == pytest.approx( + _expected_cost( + rates, + uncached_input_tokens=0, + cached_read_tokens=4352, + cache_write_tokens=5, + billed_output_tokens=20, + ) + ) + + +def test_coralbricks_rows_are_listed_for_wildcard_expansion(): + catalog: Final = {name for name, info in litellm.model_cost.items() if info.get("litellm_provider") == "coralbricks"} + assert CORALBRICKS_MODEL in catalog + assert set(litellm.get_valid_models(custom_llm_provider="coralbricks")) == catalog + + +def test_coralbricks_cost_map_backup_mirrors_main(): + main_map: Final = json.loads( + (Path(litellm.__file__).parent.parent / "model_prices_and_context_window.json").read_text() + ) + backup_map: Final = json.loads( + (Path(litellm.__file__).parent / "model_prices_and_context_window_backup.json").read_text() + ) + main_keys: Final = {key for key in main_map if key.startswith("coralbricks/")} + backup_keys: Final = {key for key in backup_map if key.startswith("coralbricks/")} + assert main_keys and main_keys == backup_keys + for key in main_keys: + assert backup_map[key] == main_map[key] + + +def test_coralbricks_cost_map_rows_declare_same_endpoints_as_providers_json(): + providers_json: Final = json.loads( + (Path(litellm.__file__).parent / "llms" / "openai_like" / "providers.json").read_text() + ) + declared: Final = providers_json["coralbricks"]["supported_endpoints"] + cost_map_paths: Final = ( + Path(litellm.__file__).parent.parent / "model_prices_and_context_window.json", + Path(litellm.__file__).parent / "model_prices_and_context_window_backup.json", + ) + for cost_map_path in cost_map_paths: + cost_map: Final = json.loads(cost_map_path.read_text()) + coralbricks_rows: Final = { + key: row for key, row in cost_map.items() if key.startswith("coralbricks/") + } + assert coralbricks_rows + for row in coralbricks_rows.values(): + assert row["supported_endpoints"] == declared + + +def _billed_chat_cost(model_id: str) -> float: + with respx.mock() as upstream: + upstream.post("https://inference.coralbricks.ai/v1/chat/completions").respond( + 200, + json={ + "id": "chatcmpl_coralbricks", + "object": "chat.completion", + "created": 1_789_550_000, + "model": model_id, + "choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}], + "usage": { + "prompt_tokens": 4400, + "completion_tokens": 20, + "total_tokens": 4420, + "prompt_tokens_details": { + "cached_tokens": 4352, + "cache_write_tokens": 5, + "billable_cache_write_tokens": 5, + }, + }, + }, + ) + response: Final = litellm.completion( + model=f"coralbricks/{model_id}", + messages=[{"role": "user", "content": "Say hello"}], + api_key="coralbricks-test-key", + ) + return litellm.completion_cost(completion_response=response, model=f"coralbricks/{model_id}") + + +@pytest.mark.parametrize( + ("deprecated_slug", "current_slug"), + ( + ("glm-5.3-fp4", "glm-5.3-fast"), + ("glm-5.3", "glm-5.3-fast"), + ("deepseek-v4.1-flash-fast-fp4", "deepseek-v4.1-flash-fast"), + ("deepseek-v4.1-flash", "deepseek-v4.1-flash-fast"), + ), +) +def test_coralbricks_deprecated_slug_bills_like_its_current_slug(deprecated_slug: str, current_slug: str): + current_cost: Final = _billed_chat_cost(current_slug) + assert current_cost > 0 + assert _billed_chat_cost(deprecated_slug) == pytest.approx(current_cost) + + +def test_coralbricks_streaming_chat_cost_includes_cache_write_and_free_cached_read(): + rates: Final = _coralbricks_rates() + usage_chunk: Final = { + "id": "chatcmpl_coralbricks", + "object": "chat.completion.chunk", + "created": 1_789_550_000, + "model": CORALBRICKS_MODEL_ID, + "choices": [], + "usage": { + "prompt_tokens": 4357, + "completion_tokens": 20, + "total_tokens": 4377, + "prompt_tokens_details": { + "cached_tokens": 4352, + "cache_write_tokens": 5, + "billable_cache_write_tokens": 5, + "cache_write_blocks": 3, + }, + "completion_tokens_details": {"reasoning_tokens": 12}, + }, + } + stream_body: Final = ( + "data: " + + json.dumps( + { + "id": "chatcmpl_coralbricks", + "object": "chat.completion.chunk", + "created": 1_789_550_000, + "model": CORALBRICKS_MODEL_ID, + "choices": [{"index": 0, "delta": {"role": "assistant", "content": "ok"}}], + } + ) + + "\n\n" + + "data: " + + json.dumps( + { + "id": "chatcmpl_coralbricks", + "object": "chat.completion.chunk", + "created": 1_789_550_000, + "model": CORALBRICKS_MODEL_ID, + "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}], + } + ) + + "\n\n" + + "data: " + + json.dumps(usage_chunk) + + "\n\n" + + "data: [DONE]\n\n" + ) + with respx.mock() as upstream: + route: Final = upstream.post("https://inference.coralbricks.ai/v1/chat/completions").respond( + 200, headers={"content-type": "text/event-stream"}, content=stream_body + ) + response: Final = litellm.completion( + model=CORALBRICKS_MODEL, + messages=[{"role": "user", "content": "Say hello"}], + api_key="coralbricks-test-key", + stream=True, + stream_options={"include_usage": True}, + ) + chunks: Final = list(response) + + request: Final = route.calls.last.request + body: Final = json.loads(request.content) + assert route.call_count == 1 + assert body["stream"] is True + + built: Final = litellm.stream_chunk_builder( + chunks, messages=[{"role": "user", "content": "Say hello"}] + ) + assert built.usage.prompt_tokens_details.cached_tokens == 4352 + cost: Final = litellm.completion_cost(completion_response=built, model=CORALBRICKS_MODEL) + assert cost == pytest.approx( + _expected_cost( + rates, + uncached_input_tokens=0, + cached_read_tokens=4352, + cache_write_tokens=5, + billed_output_tokens=20, + ) + ) diff --git a/tests/unit/llms/sagemaker/test_sagemaker_common_utils.py b/tests/unit/llms/sagemaker/test_sagemaker_common_utils.py index 51a2ccb5640..17fac935679 100644 --- a/tests/unit/llms/sagemaker/test_sagemaker_common_utils.py +++ b/tests/unit/llms/sagemaker/test_sagemaker_common_utils.py @@ -4,7 +4,9 @@ from unittest.mock import AsyncMock, MagicMock, patch import httpx import pytest -from litellm.llms.sagemaker.common_utils import AWSEventStreamDecoder +from botocore.exceptions import ClientError, ConnectionClosedError + +from litellm.llms.sagemaker.common_utils import AWSEventStreamDecoder, client_error_details from litellm.llms.sagemaker.completion.transformation import SagemakerConfig @@ -32,9 +34,7 @@ def test_sagemaker_response_stream_shape_lazy_loads_once(): import litellm.llms.sagemaker.common_utils as mod sentinel = MagicMock() - with patch.object( - mod, "_load_sagemaker_response_stream_shape", return_value=sentinel - ) as mock_load: + with patch.object(mod, "_load_sagemaker_response_stream_shape", return_value=sentinel) as mock_load: assert mod.get_sagemaker_response_stream_shape() is sentinel assert mod.get_sagemaker_response_stream_shape() is sentinel mock_load.assert_called_once() @@ -94,9 +94,7 @@ def test_sagemaker_response_stream_shape_is_structure_shape(): from litellm.llms.sagemaker.common_utils import get_sagemaker_response_stream_shape shape = get_sagemaker_response_stream_shape() - assert ( - shape is not None - ), "get_sagemaker_response_stream_shape() is None — botocore may not be installed" + assert shape is not None, "get_sagemaker_response_stream_shape() is None — botocore may not be installed" assert isinstance(shape, StructureShape) assert shape.name == "InvokeEndpointWithResponseStreamOutput" @@ -274,3 +272,38 @@ class TestSagemakerTransform: # The function should properly map max_tokens if max_completion_tokens is not provided assert result == {"temperature": 0.7, "max_new_tokens": 200} + + +NULL_MESSAGE_ERROR = ClientError( + {"Error": {"Code": "InternalFailure", "Message": None}, "ResponseMetadata": {"HTTPStatusCode": 500}}, + "InvokeEndpoint", +) +DROPPED_CONNECTION_ERROR = ConnectionClosedError(endpoint_url="https://runtime.sagemaker.test/x") +HTTPX_STATUS_ERROR = httpx.HTTPStatusError( + "boom", request=httpx.Request("POST", "https://runtime.sagemaker.test/x"), response=httpx.Response(400) +) +PLAIN_ERROR = RuntimeError("plain failure") + + +@pytest.mark.parametrize( + ("error", "expected_status", "expected_message"), + [ + ( + ClientError( + { + "Error": {"Code": "ValidationError", "Message": "bad input"}, + "ResponseMetadata": {"HTTPStatusCode": 400}, + }, + "InvokeEndpoint", + ), + 400, + "bad input", + ), + (NULL_MESSAGE_ERROR, 500, str(NULL_MESSAGE_ERROR)), + (DROPPED_CONNECTION_ERROR, 500, str(DROPPED_CONNECTION_ERROR)), + (HTTPX_STATUS_ERROR, 500, str(HTTPX_STATUS_ERROR)), + (PLAIN_ERROR, 500, str(PLAIN_ERROR)), + ], +) +def test_client_error_details_reads_every_error_shape(error, expected_status, expected_message): + assert client_error_details(error) == (expected_status, expected_message) diff --git a/tests/unit/llms/sagemaker/test_sagemaker_embedding_openai.py b/tests/unit/llms/sagemaker/test_sagemaker_embedding_openai.py new file mode 100644 index 00000000000..aacbddf7d10 --- /dev/null +++ b/tests/unit/llms/sagemaker/test_sagemaker_embedding_openai.py @@ -0,0 +1,378 @@ +""" +Test cases for OpenAI-compatible SageMaker embedding endpoints and inference components + +Covers the `sagemaker/openai/` route (request and response in the OpenAI +embeddings shape) and `model_id` forwarded to `invoke_endpoint` as `InferenceComponentName` +for every SageMaker embedding config. +""" + +import asyncio +import json +from collections.abc import Mapping +from unittest.mock import MagicMock, patch + +import httpx +import pytest +from botocore.exceptions import ClientError, ConnectionClosedError + +import litellm +from litellm import Router +from litellm.llms.sagemaker.common_utils import SagemakerError +from litellm.llms.sagemaker.completion.handler import SagemakerLLM +from litellm.llms.sagemaker.embedding.openai_transformation import ( + SagemakerOpenAIEmbeddingConfig, +) +from litellm.llms.sagemaker.embedding.transformation import SagemakerEmbeddingConfig +from litellm.llms.voyage.embedding.transformation import VoyageEmbeddingConfig +from litellm.types.utils import EmbeddingResponse + +OPENAI_RESPONSE = { + "object": "list", + "data": [ + {"object": "embedding", "index": 0, "embedding": [0.1, 0.2, 0.3]}, + {"object": "embedding", "index": 1, "embedding": [0.4, 0.5, 0.6]}, + ], + "model": "served-embedding-model", + "usage": {"prompt_tokens": 11, "total_tokens": 11}, +} + + +def fake_sagemaker_session(response_body: object) -> tuple[MagicMock, MagicMock]: + client = MagicMock() + client.invoke_endpoint.return_value = { + "Body": MagicMock(read=MagicMock(return_value=json.dumps(response_body).encode())) + } + session = MagicMock() + session.client.return_value = client + return session, client + + +def invoke_kwargs(client: MagicMock) -> Mapping[str, object]: + client.invoke_endpoint.assert_called_once() + return client.invoke_endpoint.call_args.kwargs + + +class TestSagemakerEmbeddingFactoryOpenAIRoute: + def test_openai_route_selects_openai_config(self): + config = SagemakerEmbeddingConfig.get_model_config("openai/my-embed-endpoint") + + assert isinstance(config, SagemakerOpenAIEmbeddingConfig) + assert SagemakerEmbeddingConfig().get_supported_openai_params("openai/my-embed-endpoint") == [ + "dimensions", + "encoding_format", + "user", + ] + + def test_openai_route_wins_over_model_name_substrings(self): + assert isinstance( + SagemakerEmbeddingConfig.get_model_config("openai/voyage-endpoint"), + SagemakerOpenAIEmbeddingConfig, + ) + assert isinstance( + SagemakerEmbeddingConfig.get_model_config("openai/cohere-endpoint"), + SagemakerOpenAIEmbeddingConfig, + ) + + def test_plain_endpoint_names_keep_their_config(self): + assert type(SagemakerEmbeddingConfig.get_model_config("my-embed-endpoint")) is SagemakerEmbeddingConfig + assert isinstance(SagemakerEmbeddingConfig.get_model_config("voyage-endpoint"), VoyageEmbeddingConfig) + + def test_endpoint_name_strips_only_the_route(self): + assert SagemakerEmbeddingConfig.get_endpoint_name("openai/my-embed-endpoint") == "my-embed-endpoint" + assert SagemakerEmbeddingConfig.get_endpoint_name("my-embed-endpoint") == "my-embed-endpoint" + assert SagemakerEmbeddingConfig.get_endpoint_name("my-openai-endpoint") == "my-openai-endpoint" + + +class TestSagemakerOpenAIEmbeddingConfig: + def setup_method(self): + self.config = SagemakerOpenAIEmbeddingConfig() + + def test_request_is_openai_shaped_with_extra_body_flattened(self): + result = self.config.transform_embedding_request( + model="my-embed-endpoint", + input=["hello", "world"], + optional_params={"truncate": True, "extra_body": {"model": "served-embedding-model"}}, + headers={}, + ) + + assert result == {"input": ["hello", "world"], "truncate": True, "model": "served-embedding-model"} + + def test_request_without_extra_body_sends_input_only(self): + result = self.config.transform_embedding_request( + model="my-embed-endpoint", input=["hello"], optional_params={}, headers={} + ) + + assert result == {"input": ["hello"]} + + def test_openai_params_pass_through(self): + mapped = self.config.map_openai_params( + non_default_params={"dimensions": 256, "encoding_format": "float", "user": "u1", "stream": True}, + optional_params={}, + model="my-embed-endpoint", + drop_params=False, + ) + + assert mapped == {"dimensions": 256, "encoding_format": "float", "user": "u1"} + + def test_response_reads_data_rows_and_usage(self): + raw = httpx.Response(200, json=OPENAI_RESPONSE) + + result = self.config.transform_embedding_response( + model="my-embed-endpoint", + raw_response=raw, + model_response=EmbeddingResponse(), + logging_obj=MagicMock(), + ) + + assert result.object == "list" + assert result.data == [ + {"object": "embedding", "index": 0, "embedding": [0.1, 0.2, 0.3]}, + {"object": "embedding", "index": 1, "embedding": [0.4, 0.5, 0.6]}, + ] + assert result.model == "served-embedding-model" + assert result.usage is not None + assert (result.usage.prompt_tokens, result.usage.total_tokens) == (11, 11) + + def test_response_without_usage_or_model_falls_back(self): + raw = httpx.Response(200, json={"data": [{"embedding": [0.9]}]}) + + result = self.config.transform_embedding_response( + model="my-embed-endpoint", + raw_response=raw, + model_response=EmbeddingResponse(), + logging_obj=MagicMock(), + ) + + assert result.data == [{"object": "embedding", "index": 0, "embedding": [0.9]}] + assert result.model == "my-embed-endpoint" + assert result.usage is not None + assert (result.usage.prompt_tokens, result.usage.total_tokens) == (0, 0) + + def test_response_keeps_base64_embedding_strings(self): + raw = httpx.Response( + 200, json={"data": [{"index": 0, "embedding": "AACAPwAAAEA="}], "usage": {"prompt_tokens": 2}} + ) + + result = self.config.transform_embedding_response( + model="my-embed-endpoint", + raw_response=raw, + model_response=EmbeddingResponse(), + logging_obj=MagicMock(), + ) + + assert result.data == [{"object": "embedding", "index": 0, "embedding": "AACAPwAAAEA="}] + assert result.usage is not None + assert (result.usage.prompt_tokens, result.usage.total_tokens) == (2, 2) + + @pytest.mark.parametrize( + "body", + [ + {"embedding": [[0.1, 0.2]]}, + [[0.1, 0.2]], + {"data": [{"index": 0}]}, + {"data": "not-a-list"}, + ], + ) + def test_response_rejects_non_openai_shapes(self, body): + raw = httpx.Response(200, json=body) + + with pytest.raises(SagemakerError) as exc_info: + self.config.transform_embedding_response( + model="my-embed-endpoint", + raw_response=raw, + model_response=EmbeddingResponse(), + logging_obj=MagicMock(), + ) + + assert exc_info.value.status_code == 500 + assert "'data'" in exc_info.value.message + + +class TestSagemakerEmbeddingInferenceComponents: + def setup_method(self): + self.sagemaker_llm = SagemakerLLM() + + def embed_with(self, session: MagicMock, model: str, optional_params: Mapping[str, object]) -> EmbeddingResponse: + with patch("boto3.Session", return_value=session): + return self.sagemaker_llm.embedding( + model=model, + input=["hello", "good morning"], + model_response=EmbeddingResponse(), + print_verbose=print, + encoding=None, + logging_obj=MagicMock(), + optional_params={"aws_region_name": "us-east-1", **optional_params}, + ) + + def run_embedding( + self, model: str, optional_params: Mapping[str, object], response_body: object + ) -> tuple[EmbeddingResponse, Mapping[str, object]]: + session, client = fake_sagemaker_session(response_body) + result = self.embed_with(session, model, optional_params) + return result, invoke_kwargs(client) + + def test_openai_route_forwards_component_and_keeps_it_out_of_the_body(self): + result, kwargs = self.run_embedding( + model="openai/my-embed-endpoint", + optional_params={"model_id": "my-component", "extra_body": {"model": "served-embedding-model"}}, + response_body=OPENAI_RESPONSE, + ) + + assert kwargs["EndpointName"] == "my-embed-endpoint" + assert kwargs["InferenceComponentName"] == "my-component" + assert json.loads(kwargs["Body"]) == {"input": ["hello", "good morning"], "model": "served-embedding-model"} + assert [row["embedding"] for row in result.data] == [[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]] + assert result.usage is not None + assert result.usage.prompt_tokens == 11 + + def test_hf_endpoint_forwards_component_and_keeps_its_body_shape(self): + result, kwargs = self.run_embedding( + model="my-embed-endpoint", + optional_params={"model_id": "my-component", "normalize": True}, + response_body={"embedding": [[0.1, 0.2], [0.3, 0.4]]}, + ) + + assert kwargs["EndpointName"] == "my-embed-endpoint" + assert kwargs["InferenceComponentName"] == "my-component" + assert json.loads(kwargs["Body"]) == {"inputs": ["hello", "good morning"], "normalize": True} + assert [row["embedding"] for row in result.data] == [[0.1, 0.2], [0.3, 0.4]] + + def test_no_component_without_model_id(self): + _, kwargs = self.run_embedding( + model="my-embed-endpoint", + optional_params={}, + response_body={"embedding": [[0.1, 0.2], [0.3, 0.4]]}, + ) + + assert "InferenceComponentName" not in kwargs + assert json.loads(kwargs["Body"]) == {"inputs": ["hello", "good morning"]} + + @pytest.mark.parametrize( + ("sagemaker_message", "expected_hint"), + [ + ( + "Inference Component Name header is not allowed for endpoints to which you dont plan to deploy" + " inference components. Please remove the Inference Component Name header and try again.", + "remove `model_id` from this deployment, the endpoint has no inference components", + ), + ( + "Inference Component Name header is required for this endpoint.", + "pass in via `litellm.embedding(..., model_id={InferenceComponentName})`", + ), + ('Received server error (500) from model with message "boom".', ""), + ], + ) + def test_component_errors_keep_status_and_name_model_id(self, sagemaker_message: str, expected_hint: str): + session, client = fake_sagemaker_session(OPENAI_RESPONSE) + client.invoke_endpoint.side_effect = ClientError( + { + "Error": {"Code": "ValidationError", "Message": sagemaker_message}, + "ResponseMetadata": {"HTTPStatusCode": 400}, + }, + "InvokeEndpoint", + ) + with pytest.raises(SagemakerError) as raised: + self.embed_with(session, "my-embed-endpoint", {"model_id": "stray-component"}) + + assert raised.value.status_code == 400 + assert raised.value.message == ( + sagemaker_message + "\n " + expected_hint if expected_hint else sagemaker_message + ) + + def test_null_error_message_surfaces_the_error_code(self): + session, client = fake_sagemaker_session(OPENAI_RESPONSE) + error = ClientError( + { + "Error": {"Code": "InternalFailure", "Message": None}, + "ResponseMetadata": {"HTTPStatusCode": 500}, + }, + "InvokeEndpoint", + ) + client.invoke_endpoint.side_effect = error + with pytest.raises(SagemakerError) as raised: + self.embed_with(session, "openai/my-embed-endpoint", {"model_id": "x" * 2048}) + + assert raised.value.status_code == 500 + assert raised.value.message == str(error) + + def test_dropped_connection_surfaces_as_a_server_error(self): + session, client = fake_sagemaker_session(OPENAI_RESPONSE) + error = ConnectionClosedError(endpoint_url="https://runtime.sagemaker.test/x") + client.invoke_endpoint.side_effect = error + with pytest.raises(SagemakerError) as raised: + self.embed_with(session, "openai/my-embed-endpoint", {"model_id": "my-component"}) + + assert raised.value.status_code == 500 + assert raised.value.message == str(error) + + def test_non_string_model_id_is_rejected_before_the_call(self): + session, client = fake_sagemaker_session(OPENAI_RESPONSE) + with pytest.raises(SagemakerError) as raised: + self.embed_with(session, "openai/my-embed-endpoint", {"model_id": ["my-component"]}) + + assert raised.value.status_code == 400 + assert raised.value.message == "model_id must be the inference component name as a string" + client.invoke_endpoint.assert_not_called() + + +SDK_KWARGS = { + "model": "sagemaker/openai/my-embed-endpoint", + "input": ["hello", "good morning"], + "model_id": "my-component", + "extra_body": {"model": "served-embedding-model"}, + "aws_region_name": "us-east-1", + "aws_access_key_id": "test-key", + "aws_secret_access_key": "test-secret", +} + +EXPECTED_BODY = {"input": ["hello", "good morning"], "model": "served-embedding-model"} + + +class TestSagemakerOpenAIEmbeddingSdkCalls: + def test_sdk_sync(self): + session, client = fake_sagemaker_session(OPENAI_RESPONSE) + with patch("boto3.Session", return_value=session): + result = litellm.embedding(**SDK_KWARGS) + + kwargs = invoke_kwargs(client) + assert (kwargs["EndpointName"], kwargs["InferenceComponentName"]) == ("my-embed-endpoint", "my-component") + assert json.loads(kwargs["Body"]) == EXPECTED_BODY + assert len(result.data) == 2 + assert result.usage is not None + assert result.usage.total_tokens == 11 + + def test_sdk_async(self): + session, client = fake_sagemaker_session(OPENAI_RESPONSE) + with patch("boto3.Session", return_value=session): + result = asyncio.run(litellm.aembedding(**SDK_KWARGS)) + + kwargs = invoke_kwargs(client) + assert (kwargs["EndpointName"], kwargs["InferenceComponentName"]) == ("my-embed-endpoint", "my-component") + assert json.loads(kwargs["Body"]) == EXPECTED_BODY + assert len(result.data) == 2 + + def test_router_deployment_with_model_id(self): + router = Router( + model_list=[ + { + "model_name": "my-embeddings", + "litellm_params": { + "model": "sagemaker/openai/my-embed-endpoint", + "model_id": "my-component", + "extra_body": {"model": "served-embedding-model"}, + "aws_region_name": "us-east-1", + "aws_access_key_id": "test-key", + "aws_secret_access_key": "test-secret", + }, + "model_info": {"mode": "embedding"}, + } + ] + ) + session, client = fake_sagemaker_session(OPENAI_RESPONSE) + with patch("boto3.Session", return_value=session): + result = asyncio.run(router.aembedding(model="my-embeddings", input=["hello", "good morning"])) + + kwargs = invoke_kwargs(client) + assert (kwargs["EndpointName"], kwargs["InferenceComponentName"]) == ("my-embed-endpoint", "my-component") + assert json.loads(kwargs["Body"]) == EXPECTED_BODY + assert [row["index"] for row in result.data] == [0, 1] diff --git a/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py b/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py index fb92cb39f33..f98c78b2ae8 100644 --- a/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py +++ b/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py @@ -2,7 +2,7 @@ import asyncio, importlib, os, uuid import json import re from copy import deepcopy -from typing import Final, List, cast, get_args +from typing import Final, List, Literal, cast, get_args from unittest.mock import MagicMock, Mock, patch import httpx @@ -812,6 +812,102 @@ def test_vertex_ai_url_context_tool_use_tokens_billed_as_input_tokens(): assert usage.total_tokens - usage.prompt_tokens - usage.completion_tokens == 0 +def test_url_context_tool_use_tokens_are_priced_as_input(): + """ + URL context tool-use tokens are billed at the input token rate, so the input cost must + cover every token counted in usage.prompt_tokens, not just the reported text modality. + See https://github.com/BerriAI/litellm/issues/45991 + """ + v = VertexGeminiConfig() + completion_response = { + "candidates": [ + { + "urlContextMetadata": {"urlMetadata": []}, + "groundingMetadata": { + "groundingChunks": [{"web": {"uri": "https://example.com", "title": "Example"}}] + }, + } + ], + "usageMetadata": UsageMetadata( + promptTokenCount=23, + candidatesTokenCount=10, + toolUsePromptTokenCount=14115, + totalTokenCount=14148, + promptTokensDetails=[{"modality": "TEXT", "tokenCount": 23}], + toolUsePromptTokensDetails=[{"modality": "TEXT", "tokenCount": 14115}], + ), + } + + usage = v.calculate_usage(completion_response=completion_response) + + assert usage.prompt_tokens == 23 + 14115 + assert usage.prompt_tokens_details.text_tokens == usage.prompt_tokens + input_rate: Final = litellm.get_model_info("gemini-2.5-flash")["input_cost_per_token"] + prompt_cost, _ = litellm.cost_per_token(model="gemini/gemini-2.5-flash", usage_object=usage) + assert prompt_cost == usage.prompt_tokens * input_rate + + +def test_search_grounding_tool_use_tokens_stay_unpriced(): + """ + Grounding with Google Search tokens are not billed at the input token rate (a separate + per-query search fee applies), so neither prompt_tokens nor the input cost may grow with + toolUsePromptTokenCount. + """ + v = VertexGeminiConfig() + completion_response = { + "candidates": [{"groundingMetadata": {"webSearchQueries": ["latest nobel physics"]}}], + "usageMetadata": UsageMetadata( + promptTokenCount=23, + candidatesTokenCount=10, + toolUsePromptTokenCount=14115, + totalTokenCount=14148, + promptTokensDetails=[{"modality": "TEXT", "tokenCount": 23}], + toolUsePromptTokensDetails=[{"modality": "TEXT", "tokenCount": 14115}], + ), + } + + usage = v.calculate_usage(completion_response=completion_response) + + assert usage.prompt_tokens == 23 + assert usage.prompt_tokens_details.text_tokens == 23 + input_rate: Final = litellm.get_model_info("gemini-2.5-flash")["input_cost_per_token"] + prompt_cost, _ = litellm.cost_per_token(model="gemini/gemini-2.5-flash", usage_object=usage) + assert prompt_cost == 23 * input_rate + + +def test_missing_prompt_tokens_details_still_prices_the_full_prompt(): + """ + Gemini omits promptTokensDetails on some tool-use responses, leaving text_tokens unset. + That None must be preserved so input billing falls back to the full prompt_tokens, and + must not be turned into a concrete count that prices only part of the prompt. + """ + v = VertexGeminiConfig() + completion_response = { + "candidates": [ + { + "urlContextMetadata": {"urlMetadata": []}, + "groundingMetadata": { + "groundingChunks": [{"web": {"uri": "https://example.com", "title": "Example"}}] + }, + } + ], + "usageMetadata": UsageMetadata( + promptTokenCount=23, + candidatesTokenCount=10, + toolUsePromptTokenCount=14115, + totalTokenCount=14148, + ), + } + + usage = v.calculate_usage(completion_response=completion_response) + + assert usage.prompt_tokens == 23 + 14115 + assert usage.prompt_tokens_details.text_tokens is None + input_rate: Final = litellm.get_model_info("gemini-2.5-flash")["input_cost_per_token"] + prompt_cost, _ = litellm.cost_per_token(model="gemini/gemini-2.5-flash", usage_object=usage) + assert prompt_cost == usage.prompt_tokens * input_rate + + def test_streaming_chunk_includes_reasoning_tokens(): from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( ModelResponseIterator, @@ -3013,7 +3109,7 @@ def test_is_gemini_3_or_newer(): ], ) def test_gemini_3_reasoning_effort_maps_to_thinking_level(model: str): - """Test that reasoning_effort maps to thinkingLevel and default temperature=1.0""" + """Test that reasoning_effort maps to thinkingLevel and no temperature is injected""" from litellm.llms.gemini.chat.transformation import GoogleAIStudioGeminiConfig from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( VertexGeminiConfig, @@ -3038,7 +3134,7 @@ def test_gemini_3_reasoning_effort_maps_to_thinking_level(model: str): "thinkingLevel": effort, "includeThoughts": True, } - assert mapped["temperature"] == 1.0 + assert "temperature" not in mapped assert "thinkingBudget" not in mapped["thinkingConfig"] @@ -3473,28 +3569,58 @@ def test_reasoning_effort_dict_format_gemini_3(): assert "thinkingConfig" not in result -def test_temperature_default_for_gemini_3(): - """Test that temperature defaults to 1.0 for Gemini 3+ models when not specified""" +@pytest.mark.parametrize("model", ["gemini-3-pro-preview", "gemini-flash-latest"]) +@pytest.mark.parametrize("config_cls", ["vertex", "studio"]) +def test_gemini_3_omits_temperature_when_not_specified(model: str, config_cls: Literal["vertex", "studio"]) -> None: + """Gemini 3+ requests must not gain a temperature the caller did not send""" from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( VertexGeminiConfig, ) - v = VertexGeminiConfig() - model = "gemini-3-pro-preview" - optional_params = {} - - # No temperature specified - non_default_params = {} - result = v.map_openai_params( - non_default_params=non_default_params, - optional_params=optional_params, + cfg: Final = VertexGeminiConfig() if config_cls == "vertex" else GoogleAIStudioGeminiConfig() + result: Final = cfg.map_openai_params( + non_default_params={}, + optional_params={}, model=model, drop_params=False, ) - # Should default to 1.0 - assert "temperature" in result - assert result["temperature"] == 1.0 + assert "temperature" not in result + + +@pytest.mark.parametrize("config_cls", ["vertex", "studio"]) +def test_gemini_3_explicit_temperature_is_still_forwarded(config_cls: Literal["vertex", "studio"]) -> None: + """An explicitly supplied temperature is still forwarded""" + from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( + VertexGeminiConfig, + ) + + cfg: Final = VertexGeminiConfig() if config_cls == "vertex" else GoogleAIStudioGeminiConfig() + result: Final = cfg.map_openai_params( + non_default_params={"temperature": 0.4}, + optional_params={}, + model="gemini-3-pro-preview", + drop_params=False, + ) + + assert result["temperature"] == 0.4 + + +@pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai"]) +def test_get_optional_params_gemini_3_has_no_default_temperature( + custom_llm_provider: Literal["gemini", "vertex_ai"], +) -> None: + """Neither the gemini nor the vertex_ai route injects a temperature""" + from litellm.utils import get_optional_params + + optional_params: Final = get_optional_params( + model="gemini-3-pro-preview", + custom_llm_provider=custom_llm_provider, + reasoning_effort="medium", + ) + + assert "temperature" not in optional_params + assert optional_params["thinkingConfig"]["thinkingLevel"] == "medium" def test_media_resolution_from_detail_parameter(): @@ -3744,8 +3870,7 @@ def test_gemini_3_image_models_no_thinking_config(): # Should NOT have thinkingConfig automatically added assert "thinkingConfig" not in result - # But should still get temperature=1.0 for Gemini 3 - assert result["temperature"] == 1.0 + assert "temperature" not in result def test_gemini_3_text_models_get_thinking_config(): @@ -3773,7 +3898,7 @@ def test_gemini_3_text_models_get_thinking_config(): # Should NOT have thinkingConfig automatically added when user provides no reasoning_effort assert "thinkingConfig" not in result - assert result["temperature"] == 1.0 + assert "temperature" not in result def test_gemini_image_models_excluded_from_thinking(): diff --git a/tests/unit/proxy/_experimental/mcp_server/faults/test_list_outcomes.py b/tests/unit/proxy/_experimental/mcp_server/faults/test_list_outcomes.py index b41e786881d..fe2254fbd46 100644 --- a/tests/unit/proxy/_experimental/mcp_server/faults/test_list_outcomes.py +++ b/tests/unit/proxy/_experimental/mcp_server/faults/test_list_outcomes.py @@ -267,3 +267,32 @@ def test_pure_non_auth_response_still_classifies_upstream_error(): fault = classify_list_exception(exc) assert fault.tag == "upstream_error" assert fault.status_code == 502 + + +@pytest.mark.parametrize("status, challenge, expected", [ + (403, 'Bearer error="insufficient_scope", scope="tools.write"', "tools.write"), + (403, 'bEaReR error=insufficient_scope, scope="tools.read tools.write"', "tools.read tools.write"), + (403, 'Bearer error="insufficient_scope"', ""), + (403, 'Bearer error="insufficient_scope", scope=', None), + (403, 'Bearer error="insufficient_scope", scope', None), + (401, 'Bearer error="insufficient_scope", scope="tools.write"', None), + (403, 'Bearer error="invalid_token", scope="tools.write"', None), + (403, '', None), + (403, 'Basic error="insufficient_scope", scope="tools.write"', None), + (403, 'Bearer error_description="error=insufficient_scope", scope="tools.write"', None), + (403, 'Bearer error="insufficient_scope", error="invalid_token"', None), + (403, 'Bearer error="insufficient_scope", scope="tools.write\r\ninjected"', None), + (403, 'Bearer error="insufficient_scope", scope="tools.write", Basic realm="other"', None), +]) +@pytest.mark.parametrize("client", ["httpx", "httpx2"]) +def test_managed_scope_step_up_only_accepts_unambiguous_bearer_challenges(status, challenge, expected, client): + import httpx2 + + from litellm.proxy._experimental.mcp_server.faults.list_outcomes import upstream_insufficient_scope + + transport = httpx if client == "httpx" else httpx2 + response = transport.Response(status, headers={"www-authenticate": challenge}, request=transport.Request("POST", "https://peer.test/mcp")) + cause = transport.HTTPStatusError("upstream rejected", request=response.request, response=response) + wrapped = RuntimeError("transport failure") + wrapped.__cause__ = cause + assert upstream_insufficient_scope(wrapped) == expected diff --git a/tests/unit/proxy/_experimental/mcp_server/outbound_credentials/test_authz_code_refresher.py b/tests/unit/proxy/_experimental/mcp_server/outbound_credentials/test_authz_code_refresher.py index 350e7f1f910..3068e5a0116 100644 --- a/tests/unit/proxy/_experimental/mcp_server/outbound_credentials/test_authz_code_refresher.py +++ b/tests/unit/proxy/_experimental/mcp_server/outbound_credentials/test_authz_code_refresher.py @@ -210,16 +210,17 @@ async def test_unrotated_refresh_token_is_carried_forward(): @pytest.mark.asyncio -async def test_unrecorded_scope_is_carried_forward(): +@pytest.mark.parametrize("prior_scopes", [(), ("read", "write")]) +async def test_unrecorded_scope_is_carried_forward(prior_scopes): persisted = [] refresher = _refresher(body={"access_token": "new-at"}, persist_sink=persisted) token = await refresher.refresh( - "a", "s", OAuthToken("old", refresh_token="rt", scopes=("read", "write")) + "a", "s", OAuthToken("old", refresh_token="rt", scopes=prior_scopes) ) assert token is not None # response omitted "scope" -> the user's recorded grant is preserved, not dropped - assert token.scopes == ("read", "write") - assert persisted[0][5] == ("read", "write") + assert token.scopes == prior_scopes + assert persisted[0][5] == prior_scopes @pytest.mark.asyncio @@ -257,18 +258,19 @@ async def test_refresh_omits_resource_when_unset(): @pytest.mark.asyncio -async def test_returned_scope_overrides_prior_when_present(): +@pytest.mark.parametrize("scope, expected", [("read", ("read",)), ("", ()), (" ", ())]) +async def test_returned_scope_overrides_prior_when_present(scope, expected): persisted = [] refresher = _refresher( - body={"access_token": "new-at", "scope": "read"}, + body={"access_token": "new-at", "scope": scope}, persist_sink=persisted, ) token = await refresher.refresh( "a", "s", OAuthToken("old", refresh_token="rt", scopes=("read", "write")) ) assert token is not None - assert token.scopes == ("read",) # a present scope replaces the prior grant - assert persisted[0][5] == ("read",) + assert token.scopes == expected + assert persisted[0][5] == expected @pytest.mark.asyncio @@ -366,7 +368,7 @@ async def test_cimd_refresh_on_fresh_replica_preserves_user_and_client_identity( {}, ) ] - assert persisted == [("alice", "srv", "new-at", "old-rt", 3600, None)] + assert persisted == [("alice", "srv", "new-at", "old-rt", 3600, ())] assert server.client_id is None diff --git a/tests/unit/proxy/_experimental/mcp_server/outbound_credentials/test_per_user_oauth_store.py b/tests/unit/proxy/_experimental/mcp_server/outbound_credentials/test_per_user_oauth_store.py index e7308884060..5e2bb0f3722 100644 --- a/tests/unit/proxy/_experimental/mcp_server/outbound_credentials/test_per_user_oauth_store.py +++ b/tests/unit/proxy/_experimental/mcp_server/outbound_credentials/test_per_user_oauth_store.py @@ -26,6 +26,63 @@ class _RecordingStore: self.invalidations.append((user_id, server_id)) +@pytest.fixture +def credential_database(monkeypatch): + from unittest.mock import AsyncMock, MagicMock + + from litellm.proxy import proxy_server + + prisma = MagicMock() + table = prisma.db.litellm_mcpusercredentials + table.upsert = AsyncMock() + monkeypatch.setattr(proxy_server, "prisma_client", prisma) + monkeypatch.setenv("LITELLM_SALT_KEY", "credential-persistence-test-salt") + return table + + +@pytest.mark.asyncio +@pytest.mark.parametrize("scopes, expected", [(None, None), ((), []), (("read",), ["read"])]) +async def test_credential_persistence_preserves_known_and_unknown_scopes(credential_database, scopes, expected): + from litellm.proxy._experimental.mcp_server.db import decode_oauth_payload + from litellm.proxy._experimental.mcp_server.outbound_credentials.per_user_oauth_store import _persist_credential + + await _persist_credential("alice", "srv", "access", "refresh", 60, scopes) + row = credential_database.upsert.call_args.kwargs["data"]["create"] + payload = decode_oauth_payload(row["credential_b64"]) + assert payload is not None and payload.get("scopes") == expected + assert row["user_id"] == "alice" and row["server_id"] == "srv" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "prior_scopes, response_scope, expected", + [((), None, []), (("read", "write"), "", []), (("read", "write"), None, ["read", "write"])], +) +async def test_native_refresh_persists_actual_grant_through_storage_adapter( + credential_database, prior_scopes, response_scope, expected, +): + from litellm.proxy._experimental.mcp_server.db import decode_oauth_payload + from litellm.proxy._experimental.mcp_server.outbound_credentials.authz_code_refresher import AuthorizationCodeRefresher + from litellm.proxy._experimental.mcp_server.outbound_credentials.per_user_oauth_store import _persist_credential + from litellm.types.mcp_server.mcp_server_manager import MCPServer + + server = MCPServer( + server_id="srv", name="srv", transport="http", auth_type="oauth2", + token_url="https://idp.example/token", client_id="client", client_secret="secret", + ) + + async def token_endpoint(url, form, headers): + assert url == server.token_url and form["grant_type"] == "refresh_token" + return {"access_token": "rotated", **({"scope": response_scope} if response_scope is not None else {})} + + refresher = AuthorizationCodeRefresher(lambda _: server, token_endpoint, _persist_credential) + refreshed = await refresher.refresh("alice", "srv", OAuthToken("old", refresh_token="refresh", scopes=prior_scopes)) + assert refreshed is not None and refreshed.scopes == tuple(expected) + payload = decode_oauth_payload(credential_database.upsert.call_args.kwargs["data"]["create"]["credential_b64"]) + assert payload is not None and payload.get("scopes") == expected + assert payload["access_token"] == "rotated" + + class _BlockingStore: def __init__(self, access_token: str) -> None: self._access_token = access_token diff --git a/tests/unit/proxy/_experimental/mcp_server/test_db_credentials.py b/tests/unit/proxy/_experimental/mcp_server/test_db_credentials.py index 10993280e89..6b7c4350426 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_db_credentials.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_db_credentials.py @@ -1856,3 +1856,15 @@ async def test_saved_cimd_grant_refreshes_from_encrypted_storage_on_a_fresh_repl assert parse_qs(upstream.calls[0].request.content.decode())["client_id"] == [identity] assert "client_secret" not in parse_qs(upstream.calls[0].request.content.decode()) assert server.client_id is None and server.client_id_metadata_document_supported is capability + + +@pytest.mark.asyncio +@pytest.mark.parametrize("provider_scope, expected", [(None, ["a"]), ("b", ["b"]), ("", [])]) +async def test_silent_refresh_preserves_omitted_scope_but_accepts_explicit_grant(monkeypatch, provider_scope, expected): + import litellm.proxy._experimental.mcp_server.db as db_mod + + await _run_refresh(monkeypatch, _refresh_server(), { + "access_token": "new-access", + **({"scope": provider_scope} if provider_scope is not None else {}), + }) + assert db_mod.store_user_oauth_credential.await_args.kwargs["scopes"] == expected diff --git a/tests/unit/proxy/_experimental/mcp_server/test_discoverable_endpoints.py b/tests/unit/proxy/_experimental/mcp_server/test_discoverable_endpoints.py index 4f22200145e..2590052f329 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_discoverable_endpoints.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_discoverable_endpoints.py @@ -11981,6 +11981,7 @@ def jwt_oauth_identity(monkeypatch: pytest.MonkeyPatch) -> tuple["JWTHandler", " monkeypatch.setattr(proxy_server, "user_api_key_cache", cache) prisma: Final = MagicMock() prisma.db.litellm_teammembership.find_unique = AsyncMock(return_value=None) + prisma.db.litellm_mcpusercredentials.find_unique = AsyncMock(return_value=None) monkeypatch.setattr(proxy_server, "prisma_client", prisma) return handler, signing_key @@ -13732,7 +13733,7 @@ async def test_token_route_preserves_saved_cimd_grant_after_origin_change( "matching", "unicode_grant", "foreign_refresh", "missing_refresh", "missing_grant", "database_missing", "database_outage", "static_client", "anonymous", ]) -async def test_saved_cimd_refresh_identity_is_bound_to_the_callers_stored_grant( +async def test_saved_refresh_metadata_is_bound_to_the_callers_stored_grant( monkeypatch: pytest.MonkeyPatch, state: str, ) -> None: from types import SimpleNamespace @@ -13766,12 +13767,210 @@ async def test_saved_cimd_refresh_identity_is_bound_to_the_callers_stored_grant( table.find_unique.side_effect = RuntimeError("database unavailable") if state == "static_client": server.client_id = "configured-client" - resolved: Final = await endpoints._saved_cimd_refresh_client_id( + resolved: Final = await endpoints._saved_refresh_credential( server, None if state == "anonymous" else "alice", "foreign-refresh" if state == "foreign_refresh" else grant, ) - assert resolved == (identity if state in ("matching", "unicode_grant") else None) - if state in ("anonymous", "static_client", "database_missing"): + if state in ("matching", "unicode_grant", "static_client"): + assert resolved is not None and resolved["cimd_client_id"] == identity + assert resolved["refresh_token"] == grant + else: + assert resolved is None + if state in ("anonymous", "database_missing"): table.find_unique.assert_not_awaited() else: table.find_unique.assert_awaited_once_with(where={"user_id_server_id": {"user_id": "alice", "server_id": server.server_id}}) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("prior_scopes, expected", [(["tools.read"], "tools.read tools.write"), (None, "tools.write"), ([], "tools.write")]) +async def test_managed_authorize_retains_the_callers_previous_grant(jwt_oauth_identity, monkeypatch, prior_scopes, expected): + from types import SimpleNamespace + from urllib.parse import parse_qs, urlsplit + + from litellm.proxy import proxy_server + from litellm.proxy._experimental.mcp_server import db, discoverable_endpoints as endpoints + from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager + + _, signing_key = jwt_oauth_identity + monkeypatch.setenv("LITELLM_SALT_KEY", "scope-grant-test-salt") + server = _cimd_oauth_server().model_copy(update={"client_id": "scope-client", "scopes": ["tools.read"]}) + global_mcp_server_manager.registry[server.server_id] = server + credential = db.encrypt_value_helper(json.dumps({"type": "oauth2", "access_token": "old-access", "scopes": prior_scopes})) + table = proxy_server.prisma_client.db.litellm_mcpusercredentials + table.find_unique = AsyncMock(return_value=SimpleNamespace(credential_b64=credential)) + request = _token_request({"Authorization": f"Bearer {_oauth_identity_jwt(signing_key, scope='litellm_proxy_admin')}"}) + response = await endpoints.authorize_with_server( + request, server, "scope-client", "http://127.0.0.1:9/cb", scope="tools.write", + code_challenge="challenge", code_challenge_method="S256", + ) + assert parse_qs(urlsplit(response.headers["location"]).query)["scope"] == [expected] + assert table.find_unique.call_args.kwargs["where"] == { + "user_id_server_id": {"user_id": "jwt-owner", "server_id": server.server_id} + } + + +@pytest.mark.asyncio +@pytest.mark.parametrize("returned_scope, expected", [(None, ["tools.read", "tools.write"]), ("tools.read", ["tools.read"]), ("", [])]) +async def test_managed_code_exchange_preserves_requested_scope_only_when_omitted( + jwt_oauth_identity, monkeypatch, respx_mock, returned_scope, expected, +): + from http.cookies import SimpleCookie + from urllib.parse import parse_qs, urlsplit + + from litellm.proxy import proxy_server + from litellm.proxy._experimental.mcp_server import db, discoverable_endpoints as endpoints + from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager + + _, signing_key = jwt_oauth_identity + monkeypatch.setenv("LITELLM_SALT_KEY", "requested-scope-test-salt") + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + server = _cimd_oauth_server().model_copy(update={"client_id": "scope-client"}) + global_mcp_server_manager.registry[server.server_id] = server + table = proxy_server.prisma_client.db.litellm_mcpusercredentials + table.find_unique = AsyncMock(return_value=None) + table.upsert = AsyncMock() + request = _token_request({"Authorization": f"Bearer {_oauth_identity_jwt(signing_key, scope='litellm_proxy_admin')}"}) + started = await endpoints.authorize_with_server( + request, server, "scope-client", "http://127.0.0.1:9/cb", scope="tools.read tools.write", + code_challenge="challenge", code_challenge_method="S256", + ) + state = parse_qs(urlsplit(started.headers["location"]).query)["state"][0] + cookies = SimpleCookie() + cookies.load(started.headers["set-cookie"]) + callback_request = _token_request({"Cookie": "; ".join(f"{name}={value.value}" for name, value in cookies.items())}) + returned = await endpoints.callback(callback_request, code="provider-code", state=state, error=None, iss=None) + code = parse_qs(urlsplit(returned.headers["location"]).query)["code"][0] + upstream = respx_mock.post(server.token_url).respond(200, json={ + "access_token": "new-access", "token_type": "Bearer", "refresh_token": "new-refresh", + **({"scope": returned_scope} if returned_scope is not None else {}), + }) + exchanged = await endpoints.exchange_token_with_server( + request, server, "authorization_code", client_id="scope-client", code=code, code_verifier="verifier", redirect_uri="http://127.0.0.1:9/cb", client_secret=None, + ) + assert parse_qs(upstream.calls[0].request.content.decode())["code"] == ["provider-code"] + saved = db.decode_oauth_payload(table.upsert.call_args.kwargs["data"]["create"]["credential_b64"]) + assert saved is not None and saved.get("scopes") == expected + assert json.loads(exchanged.body)["scope"] == " ".join(expected or ()) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("advertised, expected", [(True, "tools.read offline_access"), (False, "tools.read")]) +async def test_managed_authorize_requests_only_selected_scope_and_supported_refresh_access(monkeypatch, advertised, expected): + from urllib.parse import parse_qs, urlsplit + + from mcp.shared.auth import OAuthMetadata + + from litellm.proxy._experimental.mcp_server import discoverable_endpoints as endpoints + + monkeypatch.setenv("LITELLM_SALT_KEY", "metadata-scope-test-salt") + metadata = OAuthMetadata( + issuer="https://idp.example.com", authorization_endpoint="https://idp.example.com/authorize", + token_endpoint="https://idp.example.com/token", response_types_supported=["code"], + scopes_supported=["tools.read", "admin"] + (["offline_access"] if advertised else []), + ) + server = _cimd_oauth_server().model_copy(update={"authorization_server_metadata": metadata.model_dump(mode="json")}) + response = await endpoints.authorize_with_server( + _cimd_request(), server, "scope-client", "http://127.0.0.1:9/cb", scope="tools.read", + code_challenge="challenge", code_challenge_method="S256", + ) + assert parse_qs(urlsplit(response.headers["location"]).query)["scope"] == [expected] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("mismatch", ["server", "issuer", "user", "mode", "tampered", "invalid_payload"]) +async def test_managed_scope_code_cannot_cross_oauth_boundaries(jwt_oauth_identity, monkeypatch, respx_mock, mismatch): + from litellm.proxy._experimental.mcp_server import discoverable_endpoints as endpoints + + _, signing_key = jwt_oauth_identity + monkeypatch.setenv("LITELLM_SALT_KEY", "scope-boundary-test-salt") + server = _cimd_oauth_server().model_copy(update={"client_id": "scope-client"}) + payload = endpoints._ManagedAuthorizationCode( + upstream_code="provider-code", mcp_server_id="other-server" if mismatch == "server" else server.server_id, + issuer="https://other-issuer.test" if mismatch == "issuer" else server.issuer, + requested_scope="tools.read tools.write", scope_user_id="other-user" if mismatch == "user" else "jwt-owner", + ) + code = endpoints._MANAGED_AUTH_CODE_PREFIX + ("invalid" if mismatch == "tampered" else endpoints.encrypt_value_helper("{}" if mismatch == "invalid_payload" else payload.model_dump_json())) + request = _token_request({"Authorization": f"Bearer {_oauth_identity_jwt(signing_key, scope='litellm_proxy_admin')}"}) + if mismatch == "mode": + server.auth_type = MCPAuth.true_passthrough + with pytest.raises(HTTPException) as raised: + await endpoints.exchange_token_with_server( + request, server, "authorization_code", code=code, redirect_uri="http://127.0.0.1:9/cb", + client_id="scope-client", client_secret=None, code_verifier="verifier", + ) + assert raised.value.status_code == 400 + assert respx_mock.calls.call_count == 0 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("scope, requested_scope, expected", [ + (None, None, ["tools.read", "tools.write"]), ("tools.read", None, ["tools.read"]), ("", None, []), + (None, "tools.write", ["tools.write"]), ("tools.read", "tools.write", ["tools.read"]), ("", "tools.write", []), +]) +async def test_managed_refresh_retains_only_the_actual_grant(jwt_oauth_identity, monkeypatch, respx_mock, scope, requested_scope, expected): + from types import SimpleNamespace + + from litellm.proxy import proxy_server + from litellm.proxy._experimental.mcp_server import db, discoverable_endpoints as endpoints + from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager + + _, signing_key = jwt_oauth_identity + monkeypatch.setenv("LITELLM_SALT_KEY", "refresh-scope-test-salt") + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + server = _cimd_oauth_server().model_copy(update={"client_id": "scope-client"}) + global_mcp_server_manager.registry[server.server_id] = server + credential = db.encrypt_value_helper(json.dumps({ + "type": "oauth2", "access_token": "old-access", "refresh_token": "owner-refresh", "scopes": ["tools.read", "tools.write"], + })) + table = proxy_server.prisma_client.db.litellm_mcpusercredentials + table.find_unique = AsyncMock(return_value=SimpleNamespace(credential_b64=credential)) + table.upsert = AsyncMock() + upstream = respx_mock.post(server.token_url).respond(200, json={ + "access_token": "new-access", "token_type": "Bearer", "refresh_token": "new-refresh", + **({"scope": scope} if scope is not None else {}), + }) + request = _token_request({"Authorization": f"Bearer {_oauth_identity_jwt(signing_key, scope='litellm_proxy_admin')}"}) + await endpoints.exchange_token_with_server( + request, server, "refresh_token", code=None, redirect_uri=None, client_id="scope-client", client_secret=None, + code_verifier=None, refresh_token="owner-refresh", scope=requested_scope, + ) + saved = db.decode_oauth_payload(table.upsert.call_args.kwargs["data"]["create"]["credential_b64"]) + assert saved is not None and saved.get("scopes") == expected + assert upstream.call_count == 1 + from urllib.parse import parse_qs + sent = parse_qs(upstream.calls.last.request.content.decode()) + assert sent.get("scope") == ([requested_scope] if requested_scope else None) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("per_server_relay", [False, True]) +async def test_managed_resource_metadata_excludes_refresh_permission(monkeypatch, per_server_relay): + from litellm.proxy._experimental.mcp_server import discoverable_endpoints as endpoints + from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager + + server = _cimd_oauth_server().model_copy(update={ + "scopes": ["tools.read", "offline_access"], "per_server_oauth_discovery": per_server_relay, + }) + global_mcp_server_manager.registry[server.server_id] = server + result = await endpoints._build_oauth_protected_resource_response(_cimd_request(), server.server_id, True) + assert result["scopes_supported"] == ["tools.read"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("granted, required, expected", [ + (["tools.read"], ("tools.write",), "absent"), + (["tools.read", "tools.write"], ("tools.write",), "present"), + (None, ("tools.write",), "absent"), + (["tools.read"], (), "present"), +]) +async def test_connect_requires_the_requested_vendor_grant(monkeypatch, granted, required, expected): + from litellm.proxy._experimental.mcp_server import db, discoverable_endpoints as endpoints + + credential = {**_stored_grant(), "scopes": granted} + read = AsyncMock(return_value=credential) + prisma = MagicMock() + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", prisma) + monkeypatch.setattr(db, "get_user_oauth_credential", read) + assert await endpoints._vendor_credential_state("owner", "server", required) == expected + read.assert_awaited_once_with(prisma, "owner", "server") diff --git a/tests/unit/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py b/tests/unit/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py index 76d071b1852..c79b2efde5e 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py @@ -991,7 +991,7 @@ SCOPED_RESOURCE = "https://llm.example.com/mcp/github" _MANAGER_PATCH = "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager" -def _scoped_authorize(client_id, resource, session_user_id="u1"): +def _scoped_authorize(client_id, resource, session_user_id="u1", **kwargs): return aggregate_authorize( request=_request(query=f"client_id={client_id}"), client_id=client_id, @@ -1002,6 +1002,7 @@ def _scoped_authorize(client_id, resource, session_user_id="u1"): response_type="code", session_user_id=session_user_id, resource=resource, + **kwargs, ) @@ -1914,7 +1915,9 @@ def _redis_that(async_increment, get=None): cache.redis_cache.async_increment = async_increment cache.redis_cache.check_and_fix_namespace = MagicMock(side_effect=lambda key: key) cache.redis_cache.init_async_client.return_value.get = get or AsyncMock(return_value=None) - cache.redis_cache.async_get_cache = AsyncMock(side_effect=AssertionError("peek must read the client, not the wrapper")) + cache.redis_cache.async_get_cache = AsyncMock( + side_effect=AssertionError("peek must read the client, not the wrapper") + ) cache.async_increment_cache = AsyncMock(side_effect=AssertionError("must not fall back to in-memory")) return cache @@ -2068,7 +2071,11 @@ async def test_refresh_of_a_rotated_token_answers_503_before_minting_while_redis (await _redeem_native(await _native_code(client_id, cache=issued), client_id, _Minter(), cache=issued)).body ) rotated = json.loads( - (await _refresh_native(payload["refresh_token"], client_id, _Minter(), _redis_that(AsyncMock(return_value=1)))).body + ( + await _refresh_native( + payload["refresh_token"], client_id, _Minter(), _redis_that(AsyncMock(return_value=1)) + ) + ).body )["refresh_token"] minter = _Minter() @@ -2578,3 +2585,287 @@ async def test_token_exchange_relays_a_mint_refusal(failure, status, error): response = await _exchange_native(client_id, _Minter(failure), _Exchanger()) assert response.status_code == status assert json.loads(response.body)["error"] == error + + +@pytest.mark.asyncio +async def test_connect_flow_retains_step_up_scope_until_the_vendor_grants_it(): + from unittest.mock import AsyncMock, patch + + client_id = (await _register([REDIRECT_URI]))["client_id"] + server = _scoped_mcp_server() + with patch(_MANAGER_PATCH) as manager: + manager.get_mcp_server_by_name.return_value = server + response = _scoped_authorize(client_id, SCOPED_RESOURCE, scope="tools.write") + vendor = AsyncMock(return_value="absent") + described = await _describe_page(response, scoped_server=server, vendor=vendor) + assert json.loads(described.body)["requested_scopes"] == ["tools.write"] + assert json.loads(described.body)["connected"] is False + vendor.assert_awaited_once_with("u1", server.server_id, ("tools.write",)) + premature = await _complete_page(response, scoped_server=server, vendor=vendor) + assert premature.status_code == 400 + vendor.return_value = "present" + complete = await _complete_page(response, scoped_server=server, vendor=vendor) + assert complete.status_code == 303 + assert parse_qs(urlparse(complete.headers["location"]).query)["code"][0].startswith(GATEWAY_AUTH_CODE_PREFIX) + + +@pytest.mark.asyncio +async def test_aggregate_step_up_requires_the_upstream_grant_without_narrowing_gateway_access(): + from unittest.mock import AsyncMock + + client_id = (await _register([REDIRECT_URI]))["client_id"] + server = _scoped_mcp_server() + encoded = f"litellm:mcp_scope:{server.server_id.encode().hex()}:{b'tools.write'.hex()}" + response = _scoped_authorize(client_id, "https://llm.example.com/mcp", scope=encoded) + vendor = AsyncMock(return_value="absent") + described = await _describe_page(response, scoped_server=server, vendor=vendor) + assert json.loads(described.body)["connected"] is False + assert json.loads(described.body)["requested_scopes"] == ["tools.write"] + premature = await _complete_page(response, scoped_server=server, vendor=vendor) + assert premature.status_code == 400 + vendor.return_value = "present" + ready = await _describe_page(response, scoped_server=server, vendor=vendor) + assert json.loads(ready.body)["state"] == "unscoped" + assert json.loads(ready.body)["server_id"] is None + completed = await _complete_page(response, scoped_server=server, vendor=vendor) + code = parse_qs(urlparse(completed.headers["location"]).query)["code"][0] + tokens = await _redeem(code, client_id) + assert tokens.status_code == 200 + assert _opened_principal(json.loads(tokens.body)).resource_server_id is None + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "scope", + [ + "litellm:mcp_scope:", + "litellm:mcp_scope:aa:bb", + "litellm:mcp_scope:6769746875622d6964:20", + "litellm:mcp_scope:6964:5A", + ], +) +async def test_aggregate_step_up_rejects_malformed_permissions(scope): + client_id = (await _register([REDIRECT_URI]))["client_id"] + response = _scoped_authorize(client_id, "https://llm.example.com/mcp", scope=scope) + assert response.status_code == 400 + assert json.loads(response.body)["error"] == "invalid_scope" + + +@pytest.mark.asyncio +async def test_step_up_validates_every_target_before_credential_reads_and_preserves_each_grant(): + from unittest.mock import AsyncMock, patch + + from litellm.proxy._experimental.mcp_server.oauth_utils import encode_upstream_scope + + client_id = (await _register([REDIRECT_URI]))["client_id"] + first, second = _scoped_mcp_server("first"), _scoped_mcp_server("second") + scopes = " ".join( + [encode_upstream_scope(first.server_id, "first.write"), encode_upstream_scope(second.server_id, "second.write")] + ) + response = _scoped_authorize(client_id, "https://llm.example.com", scope=scopes) + handle, cookies = _flow_cookie_from(response) + vendor = AsyncMock(return_value="absent") + reachable = AsyncMock(side_effect=[True, False]) + with patch(_MANAGER_PATCH) as manager: + manager.get_mcp_server_by_id.side_effect = lambda identity: { + first.server_id: first, + second.server_id: second, + }.get(identity) + refused = await describe_connect_flow(_request(cookies=cookies), handle, "u1", vendor, reachable) + assert refused.status_code == 400 + vendor.assert_not_awaited() + reachable.side_effect = None + reachable.return_value = True + vendor.side_effect = ["present", "absent"] + pending = await describe_connect_flow(_request(cookies=cookies), handle, "u1", vendor, reachable) + assert json.loads(pending.body)["server_id"] == second.server_id + assert json.loads(pending.body)["requested_scopes"] == ["second.write"] + assert [call.args for call in vendor.await_args_list] == [ + ("u1", first.server_id, ("first.write",)), + ("u1", second.server_id, ("second.write",)), + ] + vendor.reset_mock() + wrong_user = await describe_connect_flow(_request(cookies=cookies), handle, "u2", vendor, reachable) + assert wrong_user.status_code == 403 + vendor.assert_not_awaited() + vendor.side_effect = None + vendor.return_value = "present" + complete = await complete_connect_flow( + _request(cookies=cookies, method="POST"), + handle, + "u1", + DualCache(), + lookup_vendor_credential=vendor, + lookup_server_reachability=reachable, + ) + assert complete.status_code == 303 + assert vendor.await_count == 2 + + +@pytest.mark.asyncio +async def test_named_resource_rejects_another_upstreams_permission_before_consent(): + from unittest.mock import patch + + from litellm.proxy._experimental.mcp_server.oauth_utils import encode_upstream_scope + + client_id = (await _register([REDIRECT_URI]))["client_id"] + with patch(_MANAGER_PATCH) as manager: + manager.get_mcp_server_by_name.return_value = _scoped_mcp_server() + response = _scoped_authorize(client_id, SCOPED_RESOURCE, scope=encode_upstream_scope("other-id", "write")) + assert response.status_code == 400 + assert json.loads(response.body)["error"] == "invalid_scope" + + +@pytest.mark.parametrize( + "identity, permission", + [ + ("server:id/one", "tools.write"), + ("snowman-☃", "https://permissions.example/read"), + ("id", "!#$%&'()*+,-./:;<=>?@[]^_`{|}~"), + ], +) +def test_upstream_permission_scope_round_trip_preserves_identity_and_permission(identity, permission): + from litellm.proxy._experimental.mcp_server.oauth_utils import decode_upstream_scope, encode_upstream_scope + + encoded = encode_upstream_scope(identity, permission) + decoded = decode_upstream_scope(encoded) + assert decoded is not None + assert (decoded.server_id, decoded.permission) == (identity, permission) + assert " " not in encoded + + +@pytest.mark.asyncio +@pytest.mark.parametrize("state, expected", [("absent", 400), ("unavailable", 503)]) +async def test_upstream_step_up_cannot_finish_with_an_unproven_grant(state, expected): + from unittest.mock import AsyncMock + + from litellm.proxy._experimental.mcp_server.oauth_utils import encode_upstream_scope + + client_id = (await _register([REDIRECT_URI]))["client_id"] + server = _scoped_mcp_server() + response = _scoped_authorize( + client_id, "https://llm.example.com/mcp", scope=encode_upstream_scope(server.server_id, "tools.write") + ) + vendor = AsyncMock(return_value=state) + completed = await _complete_page(response, scoped_server=server, vendor=vendor) + assert completed.status_code == expected + assert "location" not in completed.headers + vendor.return_value = "present" + retried = await _complete_page(response, scoped_server=server, vendor=vendor) + assert retried.status_code == 303 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("route", ["/mcp", "/toolset/analysis/mcp"]) +async def test_sdk_completes_aggregate_step_up_discovery_consent_and_pkce(route): + from unittest.mock import AsyncMock + + import httpx2 + from mcp.client.auth import OAuthClientProvider + from mcp.shared.auth import AuthorizationCodeResult, OAuthClientMetadata + + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + _build_aggregate_authorization_server_response, + oauth_protected_resource_root, + ) + from litellm.proxy._experimental.mcp_server.exceptions import MCPUpstreamAuthError + + server = _scoped_mcp_server() + storage = AsyncMock() + storage.get_tokens.return_value = None + storage.get_client_info.return_value = None + callback = AsyncMock() + vendor = AsyncMock(return_value="absent") + cache = DualCache() + + async def consent(url): + params = {key: values[0] for key, values in parse_qs(urlparse(url).query).items()} + response = aggregate_authorize(request=_request(), session_user_id="u1", **params) + described = await _describe_page(response, scoped_server=server, vendor=vendor) + assert json.loads(described.body)["requested_scopes"] == ["tools.write"] + assert json.loads(described.body)["connected"] is False + premature = await _complete_page(response, scoped_server=server, vendor=vendor, cache=cache) + assert premature.status_code == 400 + vendor.return_value = "present" + completed = await _complete_page(response, scoped_server=server, vendor=vendor, cache=cache) + assert completed.status_code == 303 + returned = parse_qs(urlparse(completed.headers["location"]).query) + callback.return_value = AuthorizationCodeResult(code=returned["code"][0], state=returned["state"][0]) + + provider = OAuthClientProvider( + server_url=f"https://llm.example.com{route}", + client_metadata=OAuthClientMetadata(redirect_uris=[REDIRECT_URI], token_endpoint_auth_method="none"), + storage=storage, + redirect_handler=consent, + callback_handler=callback, + ) + flow = provider.async_auth_flow(httpx2.Request("POST", f"https://llm.example.com{route}")) + initial = await anext(flow) + challenge = MCPUpstreamAuthError( + status_code=403, + www_authenticate=None, + server_name=server.server_name, + server_id=server.server_id, + required_scope="tools.write", + ).to_http_exception(base_url="https://llm.example.com", request_path=route) + request = await flow.asend(httpx2.Response(challenge.status_code, headers=challenge.headers, request=initial)) + assert request.url.path == "/.well-known/oauth-protected-resource" + request = await flow.asend(httpx2.Response(200, json=oauth_protected_resource_root(_request()), request=request)) + assert request.url.path == "/.well-known/oauth-authorization-server/mcp" + request = await flow.asend( + httpx2.Response(200, json=_build_aggregate_authorization_server_response(_request(), False), request=request) + ) + assert request.url.path == "/register" + registered = await register_aggregate_client( + _request("/register", method="POST"), json.loads(request.content), token_exchange_available=False + ) + request = await flow.asend(httpx2.Response(registered.status_code, content=registered.body, request=request)) + assert request.url.path == "/token" + params = {key: values[0] for key, values in parse_qs(request.content.decode()).items()} + tokens = await _redeem(cache=cache, **params) + assert tokens.status_code == 200 + payload = json.loads(tokens.body) + assert _opened_principal(payload).resource_server_id is None + retry = await flow.asend(httpx2.Response(200, json=payload, request=request)) + assert retry.url == initial.url + assert retry.headers["authorization"] == f"Bearer {payload['access_token']}" + with pytest.raises(StopAsyncIteration): + await flow.asend(httpx2.Response(200, json={"result": "write completed"}, request=retry)) + callback.assert_awaited_once() + storage.set_tokens.assert_awaited_once() + + +@pytest.mark.parametrize( + "scope", + ["tools.write", "litellm:mcp_scope::7772697465", "litellm:mcp_scope:6964:5A", "litellm:mcp_scope:6964:7772 697465"], +) +def test_upstream_permission_rejects_noncanonical_tokens(scope): + from litellm.proxy._experimental.mcp_server.oauth_utils import decode_upstream_scope + + assert decode_upstream_scope(scope) is None + + +@pytest.mark.asyncio +async def test_named_step_up_finish_describes_and_retains_its_single_server_ceiling(): + from unittest.mock import AsyncMock, patch + + from litellm.proxy._experimental.mcp_server.oauth_utils import encode_upstream_scope + + client_id = (await _register([REDIRECT_URI]))["client_id"] + server = _scoped_mcp_server() + with patch(_MANAGER_PATCH) as manager: + manager.get_mcp_server_by_name.return_value = server + response = _scoped_authorize( + client_id, SCOPED_RESOURCE, scope=encode_upstream_scope(server.server_id, "tools.write") + ) + vendor = AsyncMock(return_value="present") + ready = await _describe_page(response, scoped_server=server, vendor=vendor) + payload = json.loads(ready.body) + assert payload["state"] == "interactive" + assert payload["server_id"] == server.server_id + assert payload["connected"] is True + completed = await _complete_page(response, scoped_server=server, vendor=vendor) + code = parse_qs(urlparse(completed.headers["location"]).query)["code"][0] + tokens = await _redeem(code, client_id) + assert tokens.status_code == 200 + assert _opened_principal(json.loads(tokens.body)).resource_server_id == server.server_id diff --git a/tests/unit/proxy/_experimental/mcp_server/test_mcp_max_concurrent_requests.py b/tests/unit/proxy/_experimental/mcp_server/test_mcp_max_concurrent_requests.py index 24b0f450d0b..8497d7bb7c4 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_mcp_max_concurrent_requests.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_mcp_max_concurrent_requests.py @@ -166,7 +166,11 @@ async def test_persistent_tool_call_with_limit_one_completes_without_nested_admi oauth2_headers: Final = ( {"Authorization": "Bearer subject-token"} if auth_type == MCPAuth.oauth2_token_exchange else None ) - context_token: Final = active_mcp_request_ctx_var.set(cast(ServerRequestContext, object())) + context_token: Final = active_mcp_request_ctx_var.set( + ServerRequestContext( + session=AsyncMock(), lifespan_context={}, protocol_version="2025-11-25", method="tools/call", request=None + ) + ) manager.track_gateway_session(gateway_session_id) try: with patch.object(manager, "create_mcp_client", new=AsyncMock(return_value=client)): @@ -209,7 +213,11 @@ async def test_persistent_session_admission_uses_updated_server_limit() -> None: submitted=second_submitted, release=release_second, ) - context_token: Final = active_mcp_request_ctx_var.set(cast(ServerRequestContext, object())) + context_token: Final = active_mcp_request_ctx_var.set( + ServerRequestContext( + session=AsyncMock(), lifespan_context={}, protocol_version="2025-11-25", method="tools/call", request=None + ) + ) manager.track_gateway_session(gateway_session_id) async def dispatch() -> CallToolResult: @@ -275,7 +283,11 @@ async def test_queued_persistent_call_survives_header_change_while_waiting_for_a second_client: Final = _AdmissionProbeClient( server.server_id, "second-fingerprint", tracker, submitted=second_submitted ) - context_token: Final = active_mcp_request_ctx_var.set(cast(ServerRequestContext, object())) + context_token: Final = active_mcp_request_ctx_var.set( + ServerRequestContext( + session=AsyncMock(), lifespan_context={}, protocol_version="2025-11-25", method="tools/call", request=None + ) + ) manager.track_gateway_session(gateway_session_id) async def hold_slot() -> None: diff --git a/tests/unit/proxy/_experimental/mcp_server/test_mcp_oauth_passthrough_tools.py b/tests/unit/proxy/_experimental/mcp_server/test_mcp_oauth_passthrough_tools.py index 7290f0de86d..2d72d2299c6 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_mcp_oauth_passthrough_tools.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_mcp_oauth_passthrough_tools.py @@ -1,6 +1,8 @@ from litellm.proxy._experimental.mcp_server import operations as mcp_operations + """Unit tests for MCP OAuth passthrough tool-fetch behavior.""" +import asyncio import logging import sys from unittest.mock import AsyncMock, MagicMock @@ -462,7 +464,6 @@ async def test_fetch_tools_logs_upstream_request_details_on_500(caplog): assert "upstream-token-0123456789" not in caplog.text - @pytest.mark.asyncio async def test_client_creation_failure_logs_sanitized_exchange(monkeypatch, caplog): manager = MCPServerManager() @@ -498,3 +499,243 @@ def test_passthrough_admission_recognizes_only_matching_authorization(oauth_head server = MCPServer(server_id="catalog", name="catalog", alias="catalog", transport=MCPTransport.http) assert _client_has_passthrough_authorization(server, oauth_headers, server_headers) is authorized + + +@pytest.mark.parametrize("base_url, request_path, scope, expected", [ + ("https://gateway.test/proxy", "/alias/mcp", "tools.write", 'Bearer resource_metadata="https://gateway.test/proxy/.well-known/oauth-protected-resource/alias/mcp", error="insufficient_scope", scope="tools.write"'), + ("https://gateway.test/proxy", "/mcp/alias", "", 'Bearer resource_metadata="https://gateway.test/proxy/.well-known/oauth-protected-resource/mcp/alias", error="insufficient_scope"'), + (None, None, "tools.write", 'Bearer error="insufficient_scope", scope="tools.write"'), +]) +def test_managed_step_up_challenge_uses_only_the_gateway_metadata(base_url, request_path, scope, expected): + error = MCPUpstreamAuthError(status_code=403, www_authenticate=None, server_name="alias", required_scope=scope) + response = error.to_http_exception(base_url=base_url, request_path=request_path) + assert response.status_code == 403 + assert response.headers == {"www-authenticate": expected} + + +@pytest.mark.asyncio +async def test_scope_challenge_replaces_pending_stream_headers_once(): + from litellm.proxy._experimental.mcp_server.scope_challenge import OAuthScopeResponse + + send = AsyncMock() + response = OAuthScopeResponse(send, "https://gateway.example", "/docs/mcp") + await response({"type": "http.response.start", "status": 200, "headers": []}) + body = asyncio.create_task(response({"type": "http.response.body", "body": b"event: message\n"})) + await asyncio.sleep(0) + send.assert_not_awaited() + response.deny(MCPUpstreamAuthError(403, None, "docs", required_scope="tools.write")) + await asyncio.wait_for(body, 1) + await response({"type": "http.response.body", "body": b"private tool error"}) + assert send.await_count == 2 + start, result = [call.args[0] for call in send.await_args_list] + assert start["status"] == 403 + challenge = dict(start["headers"])[b"www-authenticate"].decode() + assert challenge == ('Bearer resource_metadata="https://gateway.example/.well-known/oauth-protected-resource/docs/mcp", ' + 'error="insufficient_scope", scope="tools.write"') + assert result["body"] == b'{"detail":"Forbidden"}' + + +@pytest.mark.asyncio +async def test_scope_authorization_releases_progress_without_waiting_for_tool_result(): + import httpx2 + from mcp.server.context import ServerRequestContext + from starlette.requests import Request + + from litellm.proxy._experimental.mcp_server.mcp_context import active_mcp_request_ctx_var + from litellm.proxy._experimental.mcp_server.scope_challenge import ( + SCOPE_RESPONSE_KEY, + OAuthScopeResponse, + ) + + from litellm.proxy._experimental.mcp_server.legacy_callbacks import get_scope_response, record_upstream_tool_authorization + + send = AsyncMock() + gate = OAuthScopeResponse(send, "https://gateway.example", "/mcp/docs") + request = Request({"type": "http", SCOPE_RESPONSE_KEY: gate}) + context = ServerRequestContext(session=MagicMock(), lifespan_context={}, protocol_version="2025-11-25", method="tools/call", request=request) + token = active_mcp_request_ctx_var.set(context) + try: + await gate({"type": "http.response.start", "status": 200, "headers": []}) + progress = asyncio.create_task(gate({"type": "http.response.body", "body": b"progress", "more_body": True})) + await asyncio.sleep(0) + for method, status in [("initialize", 200), ("tools/list", 200), ("tools/call", 403)]: + await record_upstream_tool_authorization(httpx2.Response(status, request=httpx2.Request("POST", "https://upstream.example/mcp", json={"method": method})), pending=get_scope_response()) + assert gate.pending + unread = httpx2.Request("POST", "https://upstream.example/mcp", content=iter([b'{"method":"tools/call"}'])) + await record_upstream_tool_authorization(httpx2.Response(200, request=unread), pending=get_scope_response()) + assert gate.pending + # Use the actual HTTP response hook contract, before reading the tool response body. + from litellm.experimental_mcp_client.client import MCPClient + + upstream = MCPClient(server_url="https://upstream.example/mcp") + async with upstream._create_httpx_client_factory(transport=httpx2.MockTransport(lambda request: httpx2.Response(200, content=b"tool result")))() as client: + async with client.stream("POST", "https://upstream.example/mcp", json={"method": "tools/call"}): + await asyncio.wait_for(progress, 1) + assert send.await_args_list[-1].args[0]["body"] == b"progress" + gate.deny(MCPUpstreamAuthError(403, None, "docs", required_scope="tools.write")) + await gate({"type": "http.response.body", "body": b"result"}) + assert send.await_args_list[0].args[0]["status"] == 200 + assert send.await_args_list[-1].args[0]["body"] == b"result" + finally: + active_mcp_request_ctx_var.reset(token) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("failure", [ValueError, asyncio.CancelledError]) +async def test_scope_middleware_releases_waiting_response_on_validation_or_cancellation(failure): + from mcp.server.context import ServerRequestContext + from starlette.requests import Request + + from litellm.proxy._experimental.mcp_server.scope_challenge import SCOPE_RESPONSE_KEY, OAuthScopeResponse, finish_scope_response + + send = AsyncMock() + gate = OAuthScopeResponse(send, "https://gateway.example", "/mcp/docs") + context = ServerRequestContext(session=MagicMock(), lifespan_context={}, protocol_version="2025-11-25", method="tools/call", request=Request({"type": "http", SCOPE_RESPONSE_KEY: gate})) + with pytest.raises(failure): + await finish_scope_response(context, AsyncMock(side_effect=failure)) + await gate({"type": "http.response.start", "status": 200, "headers": []}) + await gate({"type": "http.response.body", "body": b"error"}) + assert not gate.pending + assert send.await_args_list[-1].args[0]["body"] == b"error" + + +@pytest.mark.asyncio +async def test_scope_response_preserves_transport_validation_errors(): + from litellm.proxy._experimental.mcp_server.scope_challenge import OAuthScopeResponse + + send = AsyncMock() + gate = OAuthScopeResponse(send, "https://gateway.example", "/mcp/docs") + await gate({"type": "http.response.start", "status": 400, "headers": []}) + await gate({"type": "http.response.body", "body": b"invalid request"}) + assert send.await_args_list[0].args[0]["status"] == 400 + assert send.await_args_list[1].args[0]["body"] == b"invalid request" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("denied", [False, True]) +@pytest.mark.parametrize("stateless", [False, True]) +@pytest.mark.parametrize("padding_size", [0, 131072]) +async def test_streamable_sdk_relays_scope_challenge_before_committing_http_response(denied, stateless, padding_size): + import httpx2 + from mcp.server.lowlevel import Server + from mcp.server.streamable_http_manager import StreamableHTTPSessionManager + from mcp.types import CallToolResult, TextContent + + from litellm.proxy._experimental.mcp_server.legacy_callbacks import get_scope_response + from litellm.proxy._experimental.mcp_server.mcp_context import active_mcp_request_ctx_var + from litellm.proxy._experimental.mcp_server.scope_challenge import SCOPE_RESPONSE_KEY, OAuthScopeResponse, finish_scope_response + + async def call_tool(context, params): + token = active_mcp_request_ctx_var.set(context) + try: + gate = get_scope_response() + assert gate is not None + if denied: + gate.deny(MCPUpstreamAuthError(403, None, "docs", required_scope="tools.write")) + return CallToolResult(content=[TextContent(type="text", text="done")]) + finally: + active_mcp_request_ctx_var.reset(token) + + server = Server("scope-test", on_call_tool=call_tool) + server.middleware.append(finish_scope_response) + manager = StreamableHTTPSessionManager(server, stateless=stateless) + + async def app(scope, receive, send): + response = OAuthScopeResponse(send, "https://gateway.example", "/mcp/docs") + scope[SCOPE_RESPONSE_KEY] = response + await manager.handle_request(scope, receive, response) + + async with manager.run(): + async with httpx2.AsyncClient(transport=httpx2.ASGITransport(app=app)) as client: + headers = {"Accept": "application/json, text/event-stream", "MCP-Protocol-Version": "2025-11-25"} + if not stateless: + initialized = await asyncio.wait_for(client.post("https://gateway.example/mcp/docs", headers=headers, json={ + "jsonrpc": "2.0", "id": 0, "method": "initialize", "params": { + "protocolVersion": "2025-11-25", "capabilities": {}, "clientInfo": {"name": "scope-test" + "x" * padding_size, "version": "1"}, + }, + }), 3) + assert initialized.status_code == 200 + headers["MCP-Session-Id"] = initialized.headers["mcp-session-id"] + notified = await client.post("https://gateway.example/mcp/docs", headers=headers, json={"jsonrpc": "2.0", "method": "notifications/initialized"}) + assert notified.status_code == 202 + response = await asyncio.wait_for(client.post("https://gateway.example/mcp/docs", headers=headers, + json={"jsonrpc": "2.0", "id": 1, "method": "tools/call", "params": {"name": "write", "arguments": {"padding": "x" * padding_size}}}), 3) + assert response.status_code == (403 if denied else 200) + if denied: + assert 'error="insufficient_scope"' in response.headers["www-authenticate"] + assert response.json() == {"detail": "Forbidden"} + else: + assert '"text":"done"' in response.text + assert "www-authenticate" not in response.headers + + +@pytest.mark.asyncio +async def test_scope_challenges_are_isolated_between_concurrent_requests(): + import httpx2 + from mcp.server.context import ServerRequestContext + from starlette.requests import Request + + from litellm.proxy._experimental.mcp_server.legacy_callbacks import get_scope_response, record_upstream_tool_authorization + from litellm.proxy._experimental.mcp_server.mcp_context import active_mcp_request_ctx_var + from litellm.proxy._experimental.mcp_server.scope_challenge import SCOPE_RESPONSE_KEY, OAuthScopeResponse + + async def request_scope(denied): + send = AsyncMock() + gate = OAuthScopeResponse(send, "https://gateway.example", "/mcp/docs") + context = ServerRequestContext(session=MagicMock(), lifespan_context={}, protocol_version="2025-11-25", method="tools/call", request=Request({"type": "http", SCOPE_RESPONSE_KEY: gate})) + token = active_mcp_request_ctx_var.set(context) + try: + await gate({"type": "http.response.start", "status": 200, "headers": []}) + await asyncio.sleep(0) + assert get_scope_response() is gate + if denied: + gate.deny(MCPUpstreamAuthError(403, None, "docs", required_scope="tools.write")) + else: + await record_upstream_tool_authorization(httpx2.Response(200, request=httpx2.Request("POST", "https://upstream.example/mcp", json={"method": "tools/call"})), pending=get_scope_response()) + await gate({"type": "http.response.body", "body": b"done"}) + return send.await_args_list[0].args[0]["status"] + finally: + active_mcp_request_ctx_var.reset(token) + + assert await asyncio.gather(request_scope(True), request_scope(False)) == [403, 200] + assert get_scope_response() is None + + +@pytest.mark.parametrize("body", [b"", b"{", b"[]", b'{"method":"initialize"}']) +def test_non_tool_payloads_do_not_hold_response_headers(body): + from litellm.proxy._experimental.mcp_server.scope_challenge import is_tool_call_request, scope_response_from_request + + assert not is_tool_call_request(body) + assert scope_response_from_request(None) is None + + +@pytest.mark.asyncio +async def test_scope_middleware_preserves_non_http_requests(): + from mcp.server.context import ServerRequestContext + from mcp.types import CallToolResult, TextContent + + from litellm.proxy._experimental.mcp_server.scope_challenge import finish_scope_response + + context = ServerRequestContext(session=MagicMock(), lifespan_context={}, protocol_version="2025-11-25", method="tools/call") + async def call_next(request): + return CallToolResult(content=[TextContent(type="text", text=request.method)]) + + returned = await finish_scope_response(context, call_next) + assert returned.content[0].text == "tools/call" + + +@pytest.mark.parametrize("path", ["/mcp", "/toolset/demo/mcp"]) +def test_managed_step_up_metadata_matches_the_sdk_configured_resource(path): + from mcp.shared.auth_utils import check_resource_allowed + from urllib.parse import urlsplit + + error = MCPUpstreamAuthError(403, None, "docs", required_scope="tools.write") + response = error.to_http_exception("https://gateway.test", path) + challenge = response.headers["www-authenticate"] + metadata = challenge.split('resource_metadata="', 1)[1].split('"', 1)[0] + metadata_path = urlsplit(metadata).path + resource_path = metadata_path.removeprefix("/.well-known/oauth-protected-resource") + assert check_resource_allowed( + requested_resource=f"https://gateway.test{path}", + configured_resource=f"https://gateway.test{resource_path}", + ), "the SDK must accept the advertised resource before it can reauthorize" diff --git a/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py index e670c9980a4..711b357a0ea 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -2211,6 +2211,31 @@ class TestMCPServerManager: assert exc_info.value.www_authenticate == challenge assert mock_client.call_tool.call_args.kwargs.get("raise_on_error") is True + @pytest.mark.asyncio + async def test_managed_scope_challenge_requires_reauthorization_without_retry(self): + server = MCPServer( + server_id="managed-scope", name="managed-scope", alias="managed-scope", + url="https://up.example.com/mcp", transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, oauth2_flow="authorization_code", + ) + upstream_error = self._upstream_status_error( + 403, 'Bearer error="insufficient_scope", scope="tools.write", resource_metadata="https://untrusted.test/metadata"' + ) + manager = MCPServerManager() + client = AsyncMock() + client.call_tool = AsyncMock(side_effect=upstream_error) + manager.create_mcp_client = AsyncMock(return_value=client) + + with pytest.raises(MCPUpstreamAuthError) as raised: + await self._run_call_regular(manager, server) + + response = raised.value.to_http_exception("https://gateway.test", "/managed-scope/mcp") + assert response.status_code == 403 + assert response.headers == { + "www-authenticate": 'Bearer resource_metadata="https://gateway.test/.well-known/oauth-protected-resource/managed-scope/mcp", error="insufficient_scope", scope="tools.write"' + } + assert client.call_tool.await_count == 1 + @pytest.mark.asyncio @pytest.mark.parametrize("is_error", [False, True]) async def test_call_passthrough_returns_tool_result_unchanged(self, is_error): @@ -4533,6 +4558,7 @@ class TestMCPServerManager: f"{issuer}/.well-known/oauth-authorization-server", { "issuer": issuer, + "response_types_supported": ["code"], "authorization_endpoint": "https://idp.example.com/authorize", "token_endpoint": "https://idp.example.com/token", "scopes_supported": ["read", "write"], @@ -4554,6 +4580,11 @@ class TestMCPServerManager: assert result.client_id_metadata_document_supported is True + assert result.authorization_server_metadata is not None + assert result.authorization_server_metadata["scopes_supported"] == ["read", "write"] + merged = manager._merge_discovered_oauth_metadata(self._passthrough_call_server(MCPAuth.oauth2), result) + assert merged.authorization_server_metadata == result.authorization_server_metadata + @pytest.mark.asyncio async def test_fetch_single_authorization_server_metadata_rejects_issuer_mismatch(self): """RFC 8414 §3.3 fail-closed: a document self-attesting a DIFFERENT issuer than the one it was @@ -7255,6 +7286,47 @@ class TestMCPServerManager: def _unrestricted_auth() -> UserAPIKeyAuth: return UserAPIKeyAuth(api_key="sk-test") + @pytest.mark.asyncio + @pytest.mark.parametrize("managed", [False, True]) + async def test_tool_response_waits_only_for_managed_user_oauth_authorization(self, managed): + import asyncio + + from mcp.server.context import ServerRequestContext + from starlette.requests import Request + + from litellm.proxy._experimental.mcp_server.mcp_context import active_mcp_request_ctx_var + from litellm.proxy._experimental.mcp_server.scope_challenge import SCOPE_RESPONSE_KEY, OAuthScopeResponse + + auth = self._unrestricted_auth() + manager, logging = self._manager_ready_for_call_tool( + [MCPTool(name="test_tool", inputSchema={"type": "object"})], + caller=ListedToolsCaller(user_api_key_auth=auth), + ) + if managed: + manager.registry["test-server"] = manager.registry["test-server"].model_copy(update={ + "auth_type": MCPAuth.oauth2, "oauth2_flow": "authorization_code", + }) + send = AsyncMock() + gate = OAuthScopeResponse(send, "https://gateway.example", "/mcp/test-server") + context = ServerRequestContext(session=MagicMock(), lifespan_context={}, protocol_version="2025-11-25", method="tools/call", request=Request({"type": "http", SCOPE_RESPONSE_KEY: gate})) + token = active_mcp_request_ctx_var.set(context) + await gate({"type": "http.response.start", "status": 200, "headers": []}) + progress = asyncio.create_task(gate({"type": "http.response.body", "body": b"progress", "more_body": True})) + + async def upstream_call(*args, **kwargs): + await asyncio.sleep(0) + assert send.await_count == (0 if managed else 2) + return CallToolResult(content=[], is_error=False) + + manager.create_mcp_client.return_value.call_tool.side_effect = upstream_call + try: + result = await manager.call_tool(server_name="test-server", name="test_tool", arguments={}, user_api_key_auth=auth, proxy_logging_obj=logging) + assert not result.is_error + finally: + gate.allow() + await asyncio.wait_for(progress, 1) + active_mcp_request_ctx_var.reset(token) + @pytest.mark.asyncio async def test_call_tool_hands_listed_tool_description_and_schema_to_pre_call_hooks(self): schema = {"type": "object", "properties": {"param": {"type": "string"}}, "required": ["param"]} @@ -18051,7 +18123,10 @@ async def test_unreusable_upstream_session_is_replaced_before_cleanup_finishes( @pytest.fixture def _active_mcp_request_context() -> Iterator[None]: - token: Final = active_mcp_request_ctx_var.set(cast(ServerRequestContext, object())) + context: Final = ServerRequestContext( + session=MagicMock(), lifespan_context={}, protocol_version="2025-11-25", method="tools/call", request=None + ) + token: Final = active_mcp_request_ctx_var.set(context) try: yield finally: @@ -19973,17 +20048,23 @@ def test_oauth_rebuild_retains_only_corroborated_cimd_capability(capability: boo transport=MCPTransport.http, auth_type=MCPAuth.oauth2, authorization_url="https://idp.example.com/authorize", token_url="https://idp.example.com/token", client_id_metadata_document_supported=capability, + authorization_server_metadata={"issuer": "https://idp.example.com", "response_types_supported": ["code"], "scopes_supported": ["offline_access"]}, ) rebuilt: Final = previous.model_copy(update={ "client_id_metadata_document_supported": not capability if rebuild == "fresh_discovery" else None, "authorization_url": "https://changed.example.com/authorize" if rebuild == "repointed" else previous.authorization_url, "token_url": None, + "authorization_server_metadata": {"scopes_supported": []} if rebuild == "fresh_discovery" else None, "issuer": "https://idp.example.com" if rebuild == "anchored" else None, "issuer_is_anchored": rebuild == "anchored", }) carry_forward_resolved_oauth_endpoints(rebuilt, previous) expected: Final = not capability if rebuild == "fresh_discovery" else None if rebuild in ("repointed", "anchored") else capability assert rebuilt.client_id_metadata_document_supported is expected + assert rebuilt.authorization_server_metadata == ( + {"scopes_supported": []} if rebuild == "fresh_discovery" + else previous.authorization_server_metadata if rebuild == "same" else None + ) @pytest.mark.asyncio diff --git a/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_tool_calls_and_headers.py b/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_tool_calls_and_headers.py index 4b46406ac69..ab6ab882095 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_tool_calls_and_headers.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_tool_calls_and_headers.py @@ -264,11 +264,19 @@ async def test_mcp_server_tool_call_strips_custom_litellm_key_header(_mcp_reques @pytest.mark.asyncio -async def test_mcp_server_tool_call_relays_upstream_auth_error_as_iserror(_mcp_request_ctx): - """The MCP session manager serializes handler exceptions as JSON-RPC errors, so a mid-session - tool call cannot emit a raw 401 the way the REST path does. mcp_server_tool_call must turn an - upstream MCPUpstreamAuthError into an explicit isError result naming the status, not a masked - 500 or a raw traceback, so the client still learns it must re-authenticate upstream.""" +@pytest.mark.parametrize("required_scope, http_transport", [(None, False), (None, True), ("tools.write", True)]) +async def test_mcp_server_tool_call_relays_upstream_auth_error_as_iserror( + _mcp_request_ctx, required_scope, http_transport +): + from starlette.requests import Request + + from litellm.proxy._experimental.mcp_server.scope_challenge import SCOPE_RESPONSE_KEY, OAuthScopeResponse + + status_code = 403 if required_scope is not None else 401 + send = AsyncMock() + response = OAuthScopeResponse(send, "https://gateway.example", "/mcp/pt") + request = Request({"type": "http", "headers": [], SCOPE_RESPONSE_KEY: response}) + await response({"type": "http.response.start", "status": 200, "headers": []}) try: from litellm.proxy._experimental.mcp_server.exceptions import MCPUpstreamAuthError from litellm.proxy._experimental.mcp_server.server import ( @@ -284,7 +292,9 @@ async def test_mcp_server_tool_call_relays_upstream_auth_error_as_iserror(_mcp_r return data async def mock_call_mcp_tool(*args, **kwargs): - raise MCPUpstreamAuthError(status_code=401, www_authenticate="Bearer", server_name="pt") + raise MCPUpstreamAuthError( + status_code=status_code, www_authenticate=None, server_name="pt", required_scope=required_scope + ) mock_logger = MagicMock() with patch( @@ -298,20 +308,27 @@ async def test_mcp_server_tool_call_relays_upstream_auth_error_as_iserror(_mcp_r with patch("litellm.proxy.proxy_server.proxy_config", MagicMock()): with patch("litellm.proxy._experimental.mcp_server.operations.verbose_logger", mock_logger): result = await mcp_server_tool_call( - _mcp_request_ctx(), _call_tool_params("test_tool", {"param": "value"}) + _mcp_request_ctx(request=request if http_transport else None), + _call_tool_params("test_tool", {"param": "value"}), ) assert result.is_error is True # The dedicated MCPUpstreamAuthError branch (not the generic Exception fallthrough) produces this # specific message and logs at info, never a traceback via verbose_logger.exception. assert "upstream authentication required" in result.content[0].text - assert "401" in result.content[0].text + assert str(status_code) in result.content[0].text exception_calls = [str(c.args[0]) for c in mock_logger.exception.call_args_list if c.args] assert not any("mcp_server_tool_call" in m for m in exception_calls), ( "must not log a traceback for the expected re-auth" ) info_calls = [str(c.args[0]) for c in mock_logger.info.call_args_list if c.args] assert any("Upstream auth failure" in m for m in info_calls) + response.allow() + await response({"type": "http.response.body", "body": b"tool error"}) + start = send.await_args_list[0].args[0] + assert start["status"] == (403 if required_scope is not None else 200) + if required_scope is not None: + assert b'error="insufficient_scope"' in dict(start["headers"])[b"www-authenticate"] def test_prepare_mcp_server_headers_case_insensitive_extra_headers(): @@ -2464,9 +2481,12 @@ async def test_mcp_routing_caps_body_peek_for_oversized_chunked_body(): (("é", 0), ("é", 1), ("中", 1), ("中", 2), ("😀", 1), ("😀", 2), ("😀", 3)), ) async def test_mcp_routing_peek_survives_multibyte_char_split_at_cap( - method: str, chunked: bool, character: str, bytes_before_cap: int + method: str, chunked: bool, character: str, bytes_before_cap: int, _mcp_request_ctx ) -> None: + from starlette.requests import Request + from litellm.proxy._experimental.mcp_server import server as mcp_module + from litellm.proxy._experimental.mcp_server.scope_challenge import finish_scope_response params: Final = ( { @@ -2490,9 +2510,14 @@ async def test_mcp_routing_peek_survives_multibyte_char_split_at_cap( send: Final = AsyncMock() received: Final[asyncio.Future[bytes]] = asyncio.get_running_loop().create_future() - async def handle_request(_: Scope, downstream_receive: Receive, outgoing: Send) -> None: + async def handle_request(incoming: Scope, downstream_receive: Receive, outgoing: Send) -> None: assert receive.await_count == (2 if chunked else 1) received.set_result(await _drain_body(downstream_receive)) + + async def dispatch(context): + return CallToolResult(content=[]) + + await finish_scope_response(_mcp_request_ctx(request=Request(incoming)), dispatch) await outgoing({"type": "http.response.start", "status": 200, "headers": []}) await outgoing({"type": "http.response.body", "body": b"{}"}) diff --git a/tests/unit/proxy/auth/test_personal_key_team_model_access.py b/tests/unit/proxy/auth/test_personal_key_team_model_access.py new file mode 100644 index 00000000000..38f34adb48e --- /dev/null +++ b/tests/unit/proxy/auth/test_personal_key_team_model_access.py @@ -0,0 +1,377 @@ +from collections.abc import Iterator, Mapping +from typing import Final +from unittest.mock import AsyncMock, MagicMock + +import pytest +from fastapi import Request + +import litellm +from litellm.proxy import proxy_server +from litellm.proxy._types import ( + LiteLLM_TeamTableCachedObj, + LiteLLM_UserTable, + LitellmUserRoles, + ModelAccessDeniedProxyException, + UserAPIKeyAuth, +) +from litellm.proxy.auth import auth_checks +from litellm.proxy.auth.auth_checks import can_key_call_resolved_model, common_checks +from litellm.proxy.utils import get_available_models_for_user + +_TEAMS: Final[Mapping[str, LiteLLM_TeamTableCachedObj]] = { + "team-a": LiteLLM_TeamTableCachedObj(team_id="team-a", models=["gpt-a", "shared"]), + "team-b": LiteLLM_TeamTableCachedObj(team_id="team-b", models=["gpt-b", "shared"]), + "team-open": LiteLLM_TeamTableCachedObj(team_id="team-open", models=[]), + "team-blocked": LiteLLM_TeamTableCachedObj(team_id="team-blocked", models=["gpt-a"], blocked=True), +} + +_ROUTER: Final = litellm.Router( + model_list=[ + {"model_name": name, "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "sk-fake"}} + for name in ("gpt-a", "gpt-b", "shared", "gpt-c") + ] +) + + +@pytest.fixture +def teams(monkeypatch: pytest.MonkeyPatch) -> Iterator[None]: + async def _get_team_object(team_id: str, **_: object) -> LiteLLM_TeamTableCachedObj: + if team_id not in _TEAMS: + raise Exception(f"team {team_id} lookup failed") + return _TEAMS[team_id] + + async def _no_membership(**_: object) -> None: + return None + + monkeypatch.setattr(auth_checks, "get_team_object", _get_team_object) + monkeypatch.setattr(auth_checks, "get_team_membership", _no_membership) + yield + + +def _enable(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "personal_key_model_access_from_teams", True) + + +async def _common_checks( + model: str | list[str], + user: LiteLLM_UserTable | None, + token: UserAPIKeyAuth, + fallbacks: tuple[str, ...] = (), +) -> bool: + request_body: Final = {"model": model} | ({"fallbacks": list(fallbacks)} if fallbacks else {}) + return await common_checks( + request_body=request_body, + team_object=None, + user_object=user, + end_user_object=None, + global_proxy_spend=None, + general_settings={}, + route="/chat/completions", + llm_router=_ROUTER, + proxy_logging_obj=MagicMock(), + valid_token=token, + request=MagicMock(spec=Request), + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "model, allowed", + [ + ("gpt-a", True), + ("gpt-b", True), + ("shared", True), + ("gpt-c", False), + ], +) +async def test_personal_key_inherits_the_union_of_its_owner_team_models( + monkeypatch: pytest.MonkeyPatch, teams: None, model: str, allowed: bool +) -> None: + _enable(monkeypatch) + user: Final = LiteLLM_UserTable(user_id="u1", teams=["team-a", "team-b"]) + token: Final = UserAPIKeyAuth(token="k1", user_id="u1") + + if allowed: + assert await _common_checks(model, user, token) is True + else: + with pytest.raises(ModelAccessDeniedProxyException): + await _common_checks(model, user, token) + + +@pytest.mark.asyncio +async def test_personal_key_without_teams_gets_no_models(monkeypatch: pytest.MonkeyPatch, teams: None) -> None: + _enable(monkeypatch) + with pytest.raises(ModelAccessDeniedProxyException): + await _common_checks( + "gpt-a", LiteLLM_UserTable(user_id="u1", teams=[]), UserAPIKeyAuth(token="k1", user_id="u1") + ) + + +@pytest.mark.asyncio +async def test_personal_key_denied_when_owner_not_loaded(monkeypatch: pytest.MonkeyPatch, teams: None) -> None: + _enable(monkeypatch) + with pytest.raises(ModelAccessDeniedProxyException): + await _common_checks("gpt-a", None, UserAPIKeyAuth(token="k1", user_id="u1")) + + +@pytest.mark.asyncio +async def test_proxy_admin_personal_key_is_not_capped(monkeypatch: pytest.MonkeyPatch, teams: None) -> None: + _enable(monkeypatch) + user: Final = LiteLLM_UserTable(user_id="admin", teams=[], user_role=LitellmUserRoles.PROXY_ADMIN.value) + token: Final = UserAPIKeyAuth(token="k1", user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN) + assert await _common_checks("gpt-c", user, token) is True + + +@pytest.mark.asyncio +async def test_personal_key_unrestricted_when_flag_off(monkeypatch: pytest.MonkeyPatch, teams: None) -> None: + monkeypatch.setattr(litellm, "personal_key_model_access_from_teams", False) + user: Final = LiteLLM_UserTable(user_id="u1", teams=[]) + assert await _common_checks("gpt-c", user, UserAPIKeyAuth(token="k1", user_id="u1")) is True + + +@pytest.mark.asyncio +async def test_personal_key_fallbacks_are_capped_by_owner_teams(monkeypatch: pytest.MonkeyPatch, teams: None) -> None: + _enable(monkeypatch) + owner: Final = LiteLLM_UserTable(user_id="u1", teams=["team-a"]) + token: Final = UserAPIKeyAuth(token="k1", user_id="u1") + + with pytest.raises(ModelAccessDeniedProxyException) as exc_info: + await _common_checks("gpt-a", owner, token, fallbacks=("gpt-c",)) + assert str(exc_info.value.code) == "403" + assert await _common_checks("gpt-a", owner, token, fallbacks=("shared",)) is True + + +@pytest.mark.asyncio +async def test_personal_key_list_models_can_be_allowed_by_different_teams( + monkeypatch: pytest.MonkeyPatch, teams: None +) -> None: + _enable(monkeypatch) + owner: Final = LiteLLM_UserTable(user_id="u1", teams=["team-a", "team-b"]) + token: Final = UserAPIKeyAuth(token="k1", user_id="u1") + + assert await _common_checks(["gpt-a", "gpt-b"], owner, token) is True + + +@pytest.mark.asyncio +async def test_personal_key_list_models_denies_the_unavailable_model( + monkeypatch: pytest.MonkeyPatch, teams: None +) -> None: + _enable(monkeypatch) + owner: Final = LiteLLM_UserTable(user_id="u1", teams=["team-a", "team-b"]) + token: Final = UserAPIKeyAuth(token="k1", user_id="u1") + + with pytest.raises(ModelAccessDeniedProxyException) as exc_info: + await _common_checks(["gpt-a", "gpt-c"], owner, token) + + assert exc_info.value.code == "403" + assert "gpt-c" in exc_info.value.message + + +@pytest.mark.asyncio +async def test_personal_key_fallbacks_are_unchanged_when_flag_off(monkeypatch: pytest.MonkeyPatch, teams: None) -> None: + monkeypatch.setattr(litellm, "personal_key_model_access_from_teams", False) + owner: Final = LiteLLM_UserTable(user_id="u1", teams=["team-a"]) + token: Final = UserAPIKeyAuth(token="k1", user_id="u1") + assert await _common_checks("gpt-a", owner, token, fallbacks=("gpt-c",)) is True + + +@pytest.mark.asyncio +async def test_personal_key_fallbacks_load_each_owner_team_once(monkeypatch: pytest.MonkeyPatch, teams: None) -> None: + _enable(monkeypatch) + get_team: Final = AsyncMock(wraps=auth_checks.get_team_object) + monkeypatch.setattr(auth_checks, "get_team_object", get_team) + owner: Final = LiteLLM_UserTable(user_id="u1", teams=["team-a", "team-b"]) + token: Final = UserAPIKeyAuth(token="k1", user_id="u1") + + assert await _common_checks("gpt-a", owner, token, fallbacks=("shared", "gpt-a")) is True + assert sorted(call.kwargs["team_id"] for call in get_team.call_args_list) == ["team-a", "team-b"] + + +@pytest.mark.asyncio +async def test_team_key_is_not_capped_by_owner_teams(monkeypatch: pytest.MonkeyPatch, teams: None) -> None: + _enable(monkeypatch) + user: Final = LiteLLM_UserTable(user_id="u1", teams=[]) + token: Final = UserAPIKeyAuth(token="k1", user_id="u1", team_id="team-open") + assert await _common_checks("gpt-c", user, token, fallbacks=("gpt-c",)) is True + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "teams_of_user, allowed", + [ + (["team-missing", "team-a"], True), + (["team-missing"], False), + (["team-blocked"], False), + (["team-blocked", "team-a"], True), + ], +) +async def test_unloadable_or_blocked_team_grants_nothing( + monkeypatch: pytest.MonkeyPatch, teams: None, teams_of_user: list[str], allowed: bool +) -> None: + _enable(monkeypatch) + user: Final = LiteLLM_UserTable(user_id="u1", teams=teams_of_user) + token: Final = UserAPIKeyAuth(token="k1", user_id="u1") + if allowed: + assert await _common_checks("gpt-a", user, token) is True + else: + with pytest.raises(ModelAccessDeniedProxyException): + await _common_checks("gpt-a", user, token) + + +@pytest.mark.asyncio +async def test_membership_change_applies_to_existing_key(monkeypatch: pytest.MonkeyPatch, teams: None) -> None: + _enable(monkeypatch) + token: Final = UserAPIKeyAuth(token="k1", user_id="u1") + assert await _common_checks("gpt-a", LiteLLM_UserTable(user_id="u1", teams=["team-a"]), token) is True + with pytest.raises(ModelAccessDeniedProxyException): + await _common_checks("gpt-a", LiteLLM_UserTable(user_id="u1", teams=["team-b"]), token) + + +@pytest.mark.asyncio +async def test_resolved_model_check_caps_personal_key(monkeypatch: pytest.MonkeyPatch, teams: None) -> None: + _enable(monkeypatch) + owner: Final = LiteLLM_UserTable(user_id="u1", teams=["team-a"]) + + async def _get_user_object(**_: object) -> LiteLLM_UserTable: + return owner + + monkeypatch.setattr(auth_checks, "get_user_object", _get_user_object) + token: Final = UserAPIKeyAuth(token="k1", user_id="u1") + + await can_key_call_resolved_model(model="gpt-a", llm_model_list=None, valid_token=token, llm_router=_ROUTER) + with pytest.raises(ModelAccessDeniedProxyException): + await can_key_call_resolved_model(model="gpt-b", llm_model_list=None, valid_token=token, llm_router=_ROUTER) + + +@pytest.mark.asyncio +async def test_resolved_model_check_denies_when_personal_key_owner_lookup_fails( + monkeypatch: pytest.MonkeyPatch, teams: None +) -> None: + _enable(monkeypatch) + monkeypatch.setattr(proxy_server, "prisma_client", None) + token: Final = UserAPIKeyAuth(token="k1", user_id="u-no-db") + await proxy_server.user_api_key_cache.async_delete_cache("u-no-db") + assert await proxy_server.user_api_key_cache.async_get_cache(key="u-no-db", model_type=LiteLLM_UserTable) is None + + with pytest.raises(ModelAccessDeniedProxyException) as exc_info: + await can_key_call_resolved_model(model="gpt-a", llm_model_list=None, valid_token=token, llm_router=_ROUTER) + assert str(exc_info.value.code) == "403" + + +async def _list_models(monkeypatch: pytest.MonkeyPatch, owner: LiteLLM_UserTable, token: UserAPIKeyAuth) -> list[str]: + async def _get_user_object(**_: object) -> LiteLLM_UserTable: + return owner + + monkeypatch.setattr(auth_checks, "get_user_object", _get_user_object) + return await get_available_models_for_user( + user_api_key_dict=token, + llm_router=_ROUTER, + general_settings={}, + user_model=None, + prisma_client=MagicMock(), + proxy_logging_obj=MagicMock(), + user_api_key_cache=MagicMock(), + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "key_models, expected", + [ + ([], {"gpt-a", "gpt-b", "shared"}), + (["gpt-a", "gpt-c"], {"gpt-a"}), + ], +) +async def test_model_listing_matches_team_derived_access( + monkeypatch: pytest.MonkeyPatch, teams: None, key_models: list[str], expected: set[str] +) -> None: + _enable(monkeypatch) + owner: Final = LiteLLM_UserTable(user_id="u1", teams=["team-a", "team-b"]) + listed: Final = await _list_models(monkeypatch, owner, UserAPIKeyAuth(token="k1", user_id="u1", models=key_models)) + assert set(listed) == expected + + +@pytest.mark.asyncio +async def test_model_listing_empty_without_teams(monkeypatch: pytest.MonkeyPatch, teams: None) -> None: + _enable(monkeypatch) + owner: Final = LiteLLM_UserTable(user_id="u1", teams=[]) + assert await _list_models(monkeypatch, owner, UserAPIKeyAuth(token="k1", user_id="u1")) == [] + + +@pytest.mark.asyncio +async def test_model_listing_unchanged_when_flag_off(monkeypatch: pytest.MonkeyPatch, teams: None) -> None: + monkeypatch.setattr(litellm, "personal_key_model_access_from_teams", False) + owner: Final = LiteLLM_UserTable(user_id="u1", teams=[]) + listed: Final = await _list_models(monkeypatch, owner, UserAPIKeyAuth(token="k1", user_id="u1")) + assert set(listed) == {"gpt-a", "gpt-b", "shared", "gpt-c"} + + +@pytest.mark.asyncio +async def test_model_listing_loads_each_team_once(monkeypatch: pytest.MonkeyPatch, teams: None) -> None: + _enable(monkeypatch) + get_team: Final = AsyncMock(wraps=auth_checks.get_team_object) + monkeypatch.setattr(auth_checks, "get_team_object", get_team) + owner: Final = LiteLLM_UserTable(user_id="u1", teams=["team-a", "team-b", "team-a"]) + listed: Final = await _list_models(monkeypatch, owner, UserAPIKeyAuth(token="k1", user_id="u1")) + assert set(listed) == {"gpt-a", "gpt-b", "shared"} + assert sorted(call.kwargs["team_id"] for call in get_team.call_args_list) == ["team-a", "team-b"] + + +@pytest.mark.asyncio +async def test_team_check_error_grants_nothing_for_that_team(monkeypatch: pytest.MonkeyPatch, teams: None) -> None: + _enable(monkeypatch) + check_team: Final = auth_checks.can_team_access_model + + async def _failing_for_team_b(model: str, team_object: LiteLLM_TeamTableCachedObj, **_: object) -> bool: + if team_object.team_id == "team-b": + raise RuntimeError("access group lookup failed") + return await check_team(model=model, team_object=team_object, llm_router=_ROUTER) + + monkeypatch.setattr(auth_checks, "can_team_access_model", _failing_for_team_b) + owner: Final = LiteLLM_UserTable(user_id="u1", teams=["team-a", "team-b"]) + token: Final = UserAPIKeyAuth(token="k1", user_id="u1") + assert await _common_checks("gpt-a", owner, token) is True + with pytest.raises(ModelAccessDeniedProxyException): + await _common_checks("gpt-b", owner, token) + + +async def _owner_lookup_fails(**_: object) -> LiteLLM_UserTable: + raise RuntimeError("database unavailable") + + +async def _owner_missing(**_: object) -> None: + return None + + +@pytest.mark.asyncio +@pytest.mark.parametrize("get_owner", [_owner_lookup_fails, _owner_missing]) +async def test_model_listing_empty_when_owner_not_loaded( + monkeypatch: pytest.MonkeyPatch, teams: None, get_owner: object +) -> None: + _enable(monkeypatch) + monkeypatch.setattr(auth_checks, "get_user_object", get_owner) + listed: Final = await get_available_models_for_user( + user_api_key_dict=UserAPIKeyAuth(token="k1", user_id="u1"), + llm_router=_ROUTER, + general_settings={}, + user_model=None, + prisma_client=MagicMock(), + proxy_logging_obj=MagicMock(), + user_api_key_cache=MagicMock(), + ) + assert listed == [] + + +@pytest.mark.asyncio +async def test_model_listing_empty_without_database(monkeypatch: pytest.MonkeyPatch, teams: None) -> None: + _enable(monkeypatch) + listed: Final = await get_available_models_for_user( + user_api_key_dict=UserAPIKeyAuth(token="k1", user_id="u1"), + llm_router=_ROUTER, + general_settings={}, + user_model=None, + prisma_client=None, + proxy_logging_obj=MagicMock(), + user_api_key_cache=MagicMock(), + ) + assert listed == [] diff --git a/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py b/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py index 492ab408fc8..a0441981447 100644 --- a/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py +++ b/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py @@ -25,6 +25,7 @@ from starlette.datastructures import UploadFile import litellm import litellm.proxy.proxy_server from litellm.caching.dual_cache import DualCache +from litellm.constants import LITELLM_PROXY_MASTER_KEY_ALIAS from litellm.proxy.common_utils.user_api_key_cache import ( NO_TEAM_MEMBERSHIP_SENTINEL, UserApiKeyCache, @@ -4736,6 +4737,98 @@ async def test_centralized_common_checks_runs_for_standard_auth(): setattr(_proxy_server_mod, k, v) +async def _run_checks_for_cached_user(db_user: LiteLLM_UserTable, token: UserAPIKeyAuth) -> None: + """Run the centralized gate for `token` on /chat/completions with `db_user` as + the cached row its user_id resolves to.""" + import litellm.proxy.proxy_server as _proxy_server_mod + from fastapi import Request + from starlette.datastructures import URL + + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + from litellm.proxy.utils import ProxyLogging + + request = Request(scope={"type": "http"}) + request._url = URL(url="/chat/completions") + key_cache = UserApiKeyCache() + key_cache.set_cache(key=db_user.user_id, value=db_user) + attrs = { + **_proxy_attrs_for_centralized_checks(), + "prisma_client": MagicMock(), + "user_api_key_cache": key_cache, + "spend_counter_cache": DualCache(), + "proxy_logging_obj": ProxyLogging(user_api_key_cache=key_cache), + } + originals = {a: getattr(_proxy_server_mod, a, None) for a in attrs} + try: + for k, v in attrs.items(): + setattr(_proxy_server_mod, k, v) + await run_centralized_common_checks( + user_api_key_auth_obj=token, + request=request, + request_data={"model": "gpt-4o", "messages": [{"role": "user", "content": "hi"}]}, + route="/chat/completions", + ) + finally: + for k, v in originals.items(): + setattr(_proxy_server_mod, k, v) + + +@pytest.mark.asyncio +async def test_centralized_common_checks_enforce_personal_budget_on_proxy_admin_token(): + """A proxy_admin token (JWT admin scope, SSO) stays subject to the DB row's + max_budget when the row is forced to the admin role, otherwise the personal + budget check in common_checks silently no-ops for admins.""" + from litellm.proxy.utils import ProxyLogging + + db_user = LiteLLM_UserTable( + user_id="admin-user", user_role="internal_user", spend=25.0, max_budget=10.0, user_email="admin@example.com" + ) + token = UserAPIKeyAuth(api_key="sk-test", user_id="admin-user", user_role=LitellmUserRoles.PROXY_ADMIN) + with ( + patch.object(ProxyLogging, "budget_alerts", new_callable=AsyncMock) as budget_alerts, + pytest.raises(litellm.BudgetExceededError) as exc_info, + ): + await _run_checks_for_cached_user(db_user, token) + assert (exc_info.value.max_budget, exc_info.value.current_cost) == (10.0, 25.0) + user_alerts = [c.kwargs["user_info"] for c in budget_alerts.call_args_list if c.kwargs["type"] == "user_budget"] + assert [a.user_email for a in user_alerts] == ["admin@example.com"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("row_models", [["no-default-models"], ["claude-haiku-4-5"]]) +async def test_centralized_common_checks_proxy_admin_token_ignores_row_model_list(row_models): + """A proxy_admin token keeps every model when its DB row carries a personal + model list, e.g. no-default-models from default_internal_user_params on a + user later promoted to admin; only the row's budget carries over. The same + row still refuses gpt-4o to a non-admin token.""" + db_user = LiteLLM_UserTable(user_id="admin-user", user_role="internal_user", models=row_models, spend=1.0) + await _run_checks_for_cached_user( + db_user, UserAPIKeyAuth(api_key="sk-test", user_id="admin-user", user_role=LitellmUserRoles.PROXY_ADMIN) + ) + with pytest.raises(ProxyException): + await _run_checks_for_cached_user( + db_user, UserAPIKeyAuth(api_key="sk-test", user_id="admin-user", user_role=LitellmUserRoles.INTERNAL_USER) + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("api_key", [LITELLM_PROXY_MASTER_KEY_ALIAS, "sk-ui-session"]) +async def test_centralized_common_checks_admin_name_identity_ignores_its_row_budget(api_key): + """The master key and the UI username/password login both resolve to the + litellm_proxy_admin_name row; that row's max_budget does not apply to them, + even once its spend is past the budget. The same row under any other admin + user_id is refused.""" + db_user = LiteLLM_UserTable(user_id="admin", user_role="proxy_admin", spend=25.0, max_budget=10.0) + await _run_checks_for_cached_user( + db_user, UserAPIKeyAuth(api_key=api_key, user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN) + ) + with pytest.raises(litellm.BudgetExceededError): + await _run_checks_for_cached_user( + db_user.model_copy(update={"user_id": "admin-user"}), + UserAPIKeyAuth(api_key=api_key, user_id="admin-user", user_role=LitellmUserRoles.PROXY_ADMIN), + ) + + @pytest.mark.asyncio @pytest.mark.parametrize( "route", diff --git a/tests/unit/proxy/batches_endpoints/test_endpoints.py b/tests/unit/proxy/batches_endpoints/test_endpoints.py index 3469dbee748..aab2074e49f 100644 --- a/tests/unit/proxy/batches_endpoints/test_endpoints.py +++ b/tests/unit/proxy/batches_endpoints/test_endpoints.py @@ -57,7 +57,7 @@ from litellm.types.llms.openai import BatchJobStatus from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo from litellm.types.utils import CredentialItem, LiteLLMBatch, SpecialEnums -from fastapi import Request, Response +from fastapi import HTTPException, Request, Response # --------------------------------------------------------------------------- # # Fixtures: distinguishable credentials per model so a wrong/hardcoded model_id @@ -1164,6 +1164,34 @@ async def test_create__uses_acreate_batch_route_type(harness, openai_env_creds): assert harness.pre_call.call_args.kwargs["route_type"] == "acreate_batch" +@pytest.mark.asyncio +@pytest.mark.parametrize("loadbalancing", [True, False], ids=["load_balanced", "unified_branch"]) +async def test_create__pre_call_ownership_denial_stops_every_dispatch_branch(harness, loadbalancing: bool): + set_body( + harness, + { + "input_file_id": "litellm_proxy_unified_id", + "model": "vertex-model", + "endpoint": "/v1/chat/completions", + "completion_window": "24h", + }, + ) + harness.is_known_model.return_value = True + harness.pre_call.side_effect = HTTPException(status_code=403, detail="does not have access to the file") + + with ( + patch.object(litellm, "enable_loadbalancing_on_batch_endpoints", loadbalancing), + patch.object(endpoints, "is_base64_encoded_unified_file_id", return_value="unified-xyz"), + patch.object(endpoints, "get_models_from_unified_file_id", return_value=["vertex-model"]), + pytest.raises(ProxyException) as exc, + ): + await call_create(harness) + + assert exc.value.code == "403", exc.value + harness.router_acreate.assert_not_called() + harness.litellm_acreate.assert_not_called() + + def install_managed_files_hook(harness: Harness) -> AsyncMock: prisma_client = AsyncMock() managed_files = PROXY_LiteLLMManagedFiles(MagicMock(async_set_cache=AsyncMock()), prisma_client=prisma_client) diff --git a/tests/unit/proxy/common_utils/test_callback_config_validation.py b/tests/unit/proxy/common_utils/test_callback_config_validation.py index a6544e56c57..f42387d3e76 100644 --- a/tests/unit/proxy/common_utils/test_callback_config_validation.py +++ b/tests/unit/proxy/common_utils/test_callback_config_validation.py @@ -5,7 +5,7 @@ import pytest from litellm.proxy.common_utils.callback_config_validation import ( callback_config_error, - conflicting_capture_error, + conflicting_callback_var_error, conflicting_span_scope_error, cross_entry_family_error, logging_metadata_config_error, @@ -282,6 +282,13 @@ def test_capture_message_content_is_accepted_on_every_otel_v2_destination(callba assert callback_config_error(callback_name, {"capture_message_content": value}) is None +@pytest.mark.usefixtures("otel_v2_on") +@pytest.mark.parametrize("callback_name", ["langfuse_otel", "arize", "weave_otel", "newrelic"]) +@pytest.mark.parametrize("value", ["redis,postgres", "none", " REDIS , postgresql "]) +def test_excluded_services_is_accepted_on_every_otel_v2_destination(callback_name: str, value: str) -> None: + assert callback_config_error(callback_name, {"excluded_services": value}) is None + + def test_capture_message_content_is_rejected_while_otel_v2_is_off(monkeypatch: pytest.MonkeyPatch) -> None: from litellm.integrations.otel.model.config import is_otel_v2_enabled @@ -294,12 +301,31 @@ def test_capture_message_content_is_rejected_while_otel_v2_is_off(monkeypatch: p assert error is not None and "capture_message_content" in error and "LITELLM_OTEL_V2" in error +def test_excluded_services_is_rejected_while_otel_v2_is_off(monkeypatch: pytest.MonkeyPatch) -> None: + from litellm.integrations.otel.model.config import is_otel_v2_enabled + + monkeypatch.delenv("LITELLM_OTEL_V2", raising=False) + is_otel_v2_enabled.cache_clear() + try: + error: Final = callback_config_error("langfuse_otel", {"excluded_services": "redis"}) + finally: + is_otel_v2_enabled.cache_clear() + assert error is not None and "excluded_services" in error and "LITELLM_OTEL_V2" in error + + @pytest.mark.parametrize("callback_name", ["langfuse", "datadog", "otel", "arize_phoenix", None]) def test_capture_message_content_is_rejected_where_it_would_never_take_effect(callback_name: str | None) -> None: error: Final = callback_config_error(callback_name, {"capture_message_content": "no_content"}) assert error is not None and "capture_message_content" in error and "langfuse_otel" in error +@pytest.mark.usefixtures("otel_v2_on") +@pytest.mark.parametrize("callback_name", ["langfuse", "datadog", "otel", "arize_phoenix", None]) +def test_excluded_services_is_rejected_where_it_would_never_take_effect(callback_name: str | None) -> None: + error: Final = callback_config_error(callback_name, {"excluded_services": "redis"}) + assert error is not None and "excluded_services" in error and "langfuse_otel" in error + + @pytest.mark.usefixtures("otel_v2_on") @pytest.mark.parametrize("value", ["full", "event_only", "span_and_event"]) def test_an_unsupported_capture_message_content_is_rejected_on_key_logging_metadata(value: str) -> None: @@ -317,6 +343,38 @@ def test_an_unsupported_capture_message_content_is_rejected_on_key_logging_metad assert error is not None and "Invalid capture_message_content" in error +@pytest.mark.usefixtures("otel_v2_on") +@pytest.mark.parametrize("callback_var", ["capture_message_content", "excluded_services"]) +def test_key_logging_metadata_rejects_non_string_destination_callback_vars(callback_var: str) -> None: + error: Final = logging_metadata_config_error( + { + "logging": [ + { + "callback_name": "newrelic", + "callback_type": "success_and_failure", + "callback_vars": {callback_var: None}, + } + ] + } + ) + assert error == f"{callback_var} must be a string, got NoneType" + + +@pytest.mark.usefixtures("otel_v2_on") +def test_key_logging_metadata_with_non_string_callback_var_key_does_not_raise() -> None: + error: Final = logging_metadata_config_error( + { + "logging": [ + { + "callback_name": "langfuse_otel", + "callback_vars": {1: "x", "excluded_services": "redis"}, + } + ] + } + ) + assert error is None + + @pytest.mark.usefixtures("otel_v2_on") def test_each_entry_keeps_its_own_capture_message_content() -> None: error: Final = logging_metadata_config_error( @@ -357,7 +415,27 @@ def test_entries_for_one_backend_share_one_capture_message_content( stored: list[tuple[str, dict[str, str]]], rejected: bool ) -> None: """The backend's entries merge into one destination, so a second value would silently win or lose.""" - error: Final = conflicting_capture_error("langfuse_otel", {"capture_message_content": "span_only"}, stored) + error: Final = conflicting_callback_var_error( + "langfuse_otel", "capture_message_content", {"capture_message_content": "span_only"}, stored + ) + assert (error is not None) is rejected + + +@pytest.mark.parametrize( + ("stored", "rejected"), + [ + ([("langfuse_otel", {"excluded_services": "none"})], True), + ([("langfuse_otel", {"excluded_services": "redis"})], False), + ([("langfuse_otel", {"langfuse_public_key": "pk"})], False), + ([("newrelic", {"excluded_services": "none"})], False), + ], +) +def test_entries_for_one_backend_share_one_excluded_services_value( + stored: list[tuple[str, dict[str, str]]], rejected: bool +) -> None: + error: Final = conflicting_callback_var_error( + "langfuse_otel", "excluded_services", {"excluded_services": "redis"}, stored + ) assert (error is not None) is rejected @@ -374,3 +452,50 @@ def test_key_logging_entries_for_one_backend_may_not_disagree_on_capture_message {"logging": [entry("success", "no_content"), entry("success_and_failure", "span_only")]} ) assert error is not None and "already set to 'no_content' by another langfuse_otel entry" in error + + +@pytest.mark.usefixtures("otel_v2_on") +def test_key_logging_entries_for_one_backend_may_not_disagree_on_excluded_services() -> None: + def entry(callback_vars: dict[str, str]) -> dict[str, str | dict[str, str]]: + return {"callback_name": "langfuse_otel", "callback_type": "success", "callback_vars": callback_vars} + + error: Final = logging_metadata_config_error( + {"logging": [entry({"excluded_services": "redis"}), entry({"excluded_services": "none"})]} + ) + assert error is not None and "already set to 'redis' by another langfuse_otel entry" in error + + +@pytest.mark.usefixtures("otel_v2_on") +@pytest.mark.parametrize( + ("stored", "incoming", "conflicts"), + [ + ("redis,postgres", "postgres,redis", False), + ("redis", "REDIS", False), + ("postgres", "postgresql", False), + ("none", "NONE", False), + ("redis", "postgres", True), + ("none", "redis", True), + ], +) +def test_key_logging_entries_compare_excluded_services_by_parsed_value( + stored: str, incoming: str, conflicts: bool +) -> None: + def entry(callback_type: str, excluded_services: str) -> dict[str, str | dict[str, str]]: + return { + "callback_name": "newrelic", + "callback_type": callback_type, + "callback_vars": {"newrelic_api_key": "license-key", "excluded_services": excluded_services}, + } + + error: Final = logging_metadata_config_error( + {"logging": [entry("success", stored), entry("success_and_failure", incoming)]} + ) + assert (error is not None) is conflicts + if conflicts: + assert error is not None and f"already set to {stored!r}" in error + + +@pytest.mark.usefixtures("otel_v2_on") +def test_a_failure_only_entry_rejects_excluded_services() -> None: + error: Final = callback_config_error("langfuse_otel", {"excluded_services": "redis"}, "failure") + assert error is not None and "excluded_services" in error and "success_and_failure" in error diff --git a/tests/unit/proxy/decisions_endpoints/test_endpoints.py b/tests/unit/proxy/decisions_endpoints/test_endpoints.py index 81d7f8eb7ac..30f0508b236 100644 --- a/tests/unit/proxy/decisions_endpoints/test_endpoints.py +++ b/tests/unit/proxy/decisions_endpoints/test_endpoints.py @@ -16,6 +16,7 @@ from pydantic import TypeAdapter from starlette.routing import Match import litellm +from litellm.litellm_core_utils.decisions_pricing import apply_decisions_pricing from litellm.proxy._lazy_features import LAZY_FEATURES, LazyFeature, attach_lazy_features from litellm.proxy.decisions_endpoints.endpoints import decisions, router, systemone from litellm.proxy.pass_through_endpoints.pass_through_endpoints import SafeRouteAdder @@ -563,6 +564,16 @@ def test_openai_format_decisions_reach_an_openai_deployment_unchanged_including_ monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter, ) -> None: + monkeypatch.setitem( + litellm.model_cost, + "gpt-6-luna", + { + **litellm.model_cost["gpt-6-luna"], + "input_cost_per_token_decisions": 1e-07, + "cache_read_input_token_cost_decisions": 1e-07, + "output_cost_per_token_decisions": 0.0, + }, + ) monkeypatch.setattr(litellm, "api_base", None) monkeypatch.delenv("OPENAI_BASE_URL", raising=False) monkeypatch.delenv("OPENAI_API_BASE", raising=False) @@ -614,16 +625,123 @@ def test_openai_format_decisions_reach_an_openai_deployment_unchanged_including_ body: Final = response.json() assert body["answers"] == _OPENAI_FORMAT_ANSWERS assert body["usage"] == upstream_usage - luna_cost: Final = litellm.model_cost["gpt-6-luna"] - expected_cost: Final = ( - (_INPUT_TOKENS - cached_tokens) * float(luna_cost["input_cost_per_token"]) - + cached_tokens * float(luna_cost["cache_read_input_token_cost"]) - + _OUTPUT_TOKENS * float(luna_cost["output_cost_per_token"]) + luna_decisions_cost: Final = apply_decisions_pricing(litellm.model_cost["gpt-6-luna"]) + expected_cost: Final = _INPUT_TOKENS * float(luna_decisions_cost["input_cost_per_token"]) + _OUTPUT_TOKENS * float( + luna_decisions_cost["output_cost_per_token"] ) assert expected_cost > 0 assert float(response.headers["x-litellm-response-cost"]) == pytest.approx(expected_cost) +_CLEF_OMNI_MEDIA: Final[Mapping[str, object]] = { + "audio": ["data:audio/wav;base64,YQ==", {"content_type": "audio/webm;codecs=opus", "base64": "Yg=="}], + "videos": ["data:video/mp4;base64,Yw=="], +} +_CLEF_OMNI_RESPONSE: Final[Mapping[str, object]] = { + "success": True, + "result": { + "model": "clef-omni", + "answers": {"heard": {"type": "noul", "noul": 0.9}}, + "usage": {"input_tokens": 100, "output_tokens": 0}, + }, +} +_CLEF_OMNI_UPSTREAM: Final = "https://cf.example/ai/run/@cf/cloudflare/clef-omni" + + +def _route_to_clef_omni(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr( + litellm.proxy.proxy_server, + "llm_router", + litellm.Router( + model_list=[ + { + "model_name": "clef-omni", + "litellm_params": { + "model": "cloudflare/clef-omni", + "api_base": "https://cf.example", + "api_key": "cf-key", + }, + }, + { + "model_name": "decider", + "litellm_params": {"model": "perplexity/pplx-decider-v1-27b", "api_key": "k"}, + }, + ] + ), + ) + + +def _clef_omni_request(endpoint: str) -> Mapping[str, object]: + if endpoint.endswith("systemone"): + return { + "model": "clef-omni", + "state": "Evaluate the clips", + "questions": {"heard": {"type": "noul", "instructions": "Is there speech?"}}, + **_CLEF_OMNI_MEDIA, + } + return { + "model": "clef-omni", + "input": "Evaluate the clips", + "questions": [{"type": "predicate", "name": "heard", "instructions": "Is there speech?"}], + **_CLEF_OMNI_MEDIA, + } + + +@pytest.mark.parametrize("endpoint", ("/v1/systemone", "/systemone", "/v1/decisions", "/decisions")) +def test_proxy_forwards_clef_omni_media_on_every_decisions_route( + client: TestClient, + monkeypatch: pytest.MonkeyPatch, + respx_mock: respx.MockRouter, + endpoint: str, +) -> None: + _route_to_clef_omni(monkeypatch) + upstream: Final = respx_mock.post(_CLEF_OMNI_UPSTREAM).respond(json=_CLEF_OMNI_RESPONSE) + + response: Final = client.post(endpoint, json=_clef_omni_request(endpoint)) + + assert response.status_code == 200, response.text + assert response.json()["answers"] == ( + {"heard": {"type": "noul", "noul": 0.9}} + if endpoint.endswith("systemone") + else [{"type": "predicate", "name": "heard", "probability": 0.9}] + ) + assert json.loads(upstream.calls[0].request.content) == { + "model": "clef-omni", + "state": "Evaluate the clips", + "questions": {"heard": {"type": "noul", "instructions": "Is there speech?"}}, + **_CLEF_OMNI_MEDIA, + } + + +@pytest.mark.parametrize( + ("endpoint", "request_body"), + ( + ("/v1/systemone", {**_REQUEST, **_CLEF_OMNI_MEDIA}), + ("/v1/decisions", {**_clef_omni_request("/v1/decisions"), "audio": [{"content_type": "audio/wav"}]}), + ("/systemone", {**_clef_omni_request("/systemone"), "audio": ["data:image/png;base64,AA=="]}), + ("/decisions", {**_clef_omni_request("/decisions"), "videos": ["data:video/mp4;base64,Yw=="] * 3}), + ), + ids=("model_without_media_support", "clip_missing_base64", "audio_slot_holding_an_image", "too_many_videos"), +) +def test_proxy_refuses_decisions_media_it_cannot_send( + client: TestClient, + monkeypatch: pytest.MonkeyPatch, + respx_mock: respx.MockRouter, + endpoint: str, + request_body: Mapping[str, object], +) -> None: + monkeypatch.setattr(litellm, "drop_params", False) + _route_to_clef_omni(monkeypatch) + clef_upstream: Final = respx_mock.post(_CLEF_OMNI_UPSTREAM).respond(json=_CLEF_OMNI_RESPONSE) + perplexity_upstream: Final = respx_mock.post("https://api.perplexity.ai/v1/decisions").respond(json=_RESPONSE) + + response: Final = client.post(endpoint, json=request_body) + + assert response.status_code == 400, response.text + assert not clef_upstream.called + assert not perplexity_upstream.called + + @pytest.mark.parametrize("safety_identifier", (7, ["end-user-1"]), ids=("numeric", "list")) @pytest.mark.parametrize("deployment_drops_params", (False, True), ids=("strict", "drop_params")) def test_a_non_string_safety_identifier_is_refused_unless_the_deployment_drops_params( @@ -686,6 +804,119 @@ def test_a_non_string_safety_identifier_is_refused_unless_the_deployment_drops_p } +@pytest.mark.parametrize("endpoint", ("/v1/systemone", "/v1/decisions")) +@pytest.mark.parametrize("deployment_drops_params", (False, True), ids=("strict", "drop_params")) +def test_clips_for_a_model_without_media_support_are_refused_unless_the_deployment_drops_params( + client: TestClient, + monkeypatch: pytest.MonkeyPatch, + respx_mock: respx.MockRouter, + endpoint: str, + deployment_drops_params: bool, +) -> None: + monkeypatch.setattr(litellm, "drop_params", False) + monkeypatch.setattr( + litellm.proxy.proxy_server, + "llm_router", + litellm.Router( + model_list=[ + { + "model_name": "clef-flash", + "litellm_params": { + "model": "cloudflare/clef-flash", + "api_base": "https://cf.example", + "api_key": "cf-key", + "drop_params": deployment_drops_params, + }, + } + ] + ), + ) + upstream: Final = respx_mock.post("https://cf.example/ai/run/@cf/cloudflare/clef-flash").respond( + json=_CLEF_OMNI_RESPONSE + ) + + response: Final = client.post(endpoint, json={**_clef_omni_request(endpoint), "model": "clef-flash"}) + + if not deployment_drops_params: + assert response.status_code == 400, response.text + assert "['audio', 'videos']" in response.json()["error"]["message"] + assert not upstream.called + return + assert response.status_code == 200, response.text + assert json.loads(upstream.calls[0].request.content) == { + "model": "clef-flash", + "state": "Evaluate the clips", + "questions": {"heard": {"type": "noul", "instructions": "Is there speech?"}}, + } + + +@pytest.mark.parametrize("endpoint", ("/v1/systemone", "/systemone", "/v1/decisions", "/decisions")) +def test_null_media_on_every_decisions_route_means_no_clips( + client: TestClient, + monkeypatch: pytest.MonkeyPatch, + respx_mock: respx.MockRouter, + endpoint: str, +) -> None: + _route_to_clef_omni(monkeypatch) + upstream: Final = respx_mock.post(_CLEF_OMNI_UPSTREAM).respond(json=_CLEF_OMNI_RESPONSE) + + response: Final = client.post(endpoint, json={**_clef_omni_request(endpoint), "audio": None, "videos": None}) + + assert response.status_code == 200, response.text + assert json.loads(upstream.calls[0].request.content) == { + "model": "clef-omni", + "state": "Evaluate the clips", + "questions": {"heard": {"type": "noul", "instructions": "Is there speech?"}}, + } + + +@pytest.mark.parametrize("endpoint", ("/v1/systemone", "/v1/decisions")) +@pytest.mark.parametrize("deployment_drops_params", (False, True), ids=("strict", "drop_params")) +def test_a_malformed_clip_for_a_model_without_media_support_is_dropped_only_when_the_deployment_drops_params( + client: TestClient, + monkeypatch: pytest.MonkeyPatch, + respx_mock: respx.MockRouter, + endpoint: str, + deployment_drops_params: bool, +) -> None: + monkeypatch.setattr(litellm, "drop_params", False) + monkeypatch.setattr( + litellm.proxy.proxy_server, + "llm_router", + litellm.Router( + model_list=[ + { + "model_name": "clef-flash", + "litellm_params": { + "model": "cloudflare/clef-flash", + "api_base": "https://cf.example", + "api_key": "cf-key", + "drop_params": deployment_drops_params, + }, + } + ] + ), + ) + upstream: Final = respx_mock.post("https://cf.example/ai/run/@cf/cloudflare/clef-flash").respond( + json=_CLEF_OMNI_RESPONSE + ) + malformed: Final = {"audio": ["https://example.com/clip.wav"], "videos": [{"content_type": "video/mp4"}]} + + response: Final = client.post(endpoint, json={**_clef_omni_request(endpoint), "model": "clef-flash", **malformed}) + + if not deployment_drops_params: + assert response.status_code == 400, response.text + assert "Invalid Decisions request" in response.json()["error"]["message"] + assert not upstream.called + return + assert response.status_code == 200, response.text + assert json.loads(upstream.calls[0].request.content) == { + "model": "clef-flash", + "state": "Evaluate the clips", + "questions": {"heard": {"type": "noul", "instructions": "Is there speech?"}}, + } + + def _decisions_feature() -> LazyFeature: return next(feature for feature in LAZY_FEATURES if feature.name == "decisions") diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_decision_model.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_decision_model.py index 54a73690840..a792e93344a 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_decision_model.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_decision_model.py @@ -381,14 +381,22 @@ def test_initialize_applies_streaming_scan_settings(mock_mgr): default_params: Final = _litellm_params() default_instance: Final = initialize_guardrail(default_params, _guardrail(default_params)) - configured_params: Final = _litellm_params(streaming_end_of_stream_only=True, streaming_sampling_rate=3) + configured_params: Final = _litellm_params( + streaming_end_of_stream_only=True, streaming_sampling_rate=3, streaming_buffer_until_moderated=True + ) configured_instance: Final = initialize_guardrail(configured_params, _guardrail(configured_params)) + disabled_params: Final = _litellm_params(streaming_buffer_until_moderated=False) + disabled_instance: Final = initialize_guardrail(disabled_params, _guardrail(disabled_params)) unified: Final = UnifiedLLMGuardrails() assert unified.resolve_streaming_flag(default_instance, "streaming_end_of_stream_only", False) is False assert unified.resolve_streaming_flag(default_instance, "streaming_sampling_rate", 5) == 5 assert unified.resolve_streaming_flag(configured_instance, "streaming_end_of_stream_only", False) is True assert unified.resolve_streaming_flag(configured_instance, "streaming_sampling_rate", 5) == 3 + assert unified.resolve_streaming_flag(default_instance, "streaming_buffer_until_moderated", False) is False + assert unified.resolve_streaming_flag(default_instance, "streaming_buffer_until_moderated", True) is True + assert unified.resolve_streaming_flag(configured_instance, "streaming_buffer_until_moderated", False) is True + assert unified.resolve_streaming_flag(disabled_instance, "streaming_buffer_until_moderated", True) is False @pytest.mark.parametrize( @@ -397,6 +405,7 @@ def test_initialize_applies_streaming_scan_settings(mock_mgr): ({"streaming_end_of_stream_only": "yes"}, TypeError, "streaming_end_of_stream_only"), ({"streaming_sampling_rate": 0}, ValueError, "streaming_sampling_rate"), ({"streaming_sampling_rate": True}, TypeError, "streaming_sampling_rate"), + ({"streaming_buffer_until_moderated": "yes"}, TypeError, "streaming_buffer_until_moderated"), ], ) @patch("litellm.logging_callback_manager") diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/zscaler_ai_guard/test_zscaler_ai_guard.py b/tests/unit/proxy/guardrails/guardrail_hooks/zscaler_ai_guard/test_zscaler_ai_guard.py index 7445db47185..bccc531793f 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/zscaler_ai_guard/test_zscaler_ai_guard.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/zscaler_ai_guard/test_zscaler_ai_guard.py @@ -10,11 +10,15 @@ Tests covering: - resolve-and-execute-policy endpoint (policyId omission) """ +import asyncio import importlib import os +from typing import Final, Literal from unittest.mock import AsyncMock, Mock, patch import pytest +from aiohttp import web +from aiohttp.test_utils import TestServer from fastapi import HTTPException import litellm @@ -561,3 +565,47 @@ def setup_and_teardown(): if hasattr(litellm, "in_memory_llm_clients_cache"): litellm.in_memory_llm_clients_cache.flush_cache() yield + + +@pytest.mark.asyncio +@pytest.mark.parametrize("input_type", ["request", "response"]) +@pytest.mark.parametrize( + ("request_data", "expected_body_user"), + [ + ({"user": "alice@example.com"}, "alice@example.com"), + ({}, None), + ({"user": None}, None), + ({"user": ""}, None), + ({"user": 12345}, None), + ({"user": ["alice@example.com"]}, None), + ], +) +async def test_apply_guardrail_sends_request_user_in_execute_policy_body( + request_data: dict[str, object], + expected_body_user: str | None, + input_type: Literal["request", "response"], +) -> None: + received: Final = asyncio.Queue[object]() + + async def execute_policy(request: web.Request) -> web.Response: + await received.put(await request.json()) + return web.json_response({"statusCode": 200, "action": "ALLOW"}) + + app: Final = web.Application() + app.router.add_post("/v1/detection/execute-policy", execute_policy) + async with TestServer(app) as server: + guardrail: Final = ZscalerAIGuard( + api_key="test_key", + api_base=str(server.make_url("/v1/detection/execute-policy")), + policy_id=1, + ) + await guardrail.apply_guardrail({"texts": ["hello"]}, request_data, input_type) + + body: Final = received.get_nowait() + assert received.empty() + assert body == { + "direction": "OUT" if input_type == "response" else "IN", + "content": "hello", + "policyId": 1, + **({"user": expected_body_user} if expected_body_user is not None else {}), + } diff --git a/tests/unit/proxy/health_endpoints/test_health_endpoints.py b/tests/unit/proxy/health_endpoints/test_health_endpoints.py index 1ca07eafb7d..2c6e013d075 100644 --- a/tests/unit/proxy/health_endpoints/test_health_endpoints.py +++ b/tests/unit/proxy/health_endpoints/test_health_endpoints.py @@ -7,7 +7,7 @@ from contextlib import contextmanager from datetime import datetime, timedelta from importlib import import_module from types import MappingProxyType, SimpleNamespace -from typing import Final +from typing import Final, Literal from unittest.mock import AsyncMock, MagicMock, patch import httpx @@ -5013,7 +5013,9 @@ def test_test_model_connection_accepts_image_edit_mode(monkeypatch): @pytest.mark.parametrize("mode", ("decisions", "evaluation"), ids=("decisions", "legacy-evaluation")) -def test_test_model_connection_accepts_decision_modes(monkeypatch, mode: str): +def test_test_model_connection_accepts_decision_modes( + monkeypatch, mode: Literal["evaluation", "decisions"] +): monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) litellm.in_memory_llm_clients_cache.flush_cache() diff --git a/tests/unit/proxy/hooks/test_batch_rate_limiter.py b/tests/unit/proxy/hooks/test_batch_rate_limiter.py index 2f71b177705..b4e972ea34f 100644 --- a/tests/unit/proxy/hooks/test_batch_rate_limiter.py +++ b/tests/unit/proxy/hooks/test_batch_rate_limiter.py @@ -11,6 +11,7 @@ import time from collections.abc import Iterator, Sequence from datetime import datetime, timezone from typing import Final, Literal +from unittest.mock import AsyncMock, MagicMock import httpx import pytest @@ -481,10 +482,9 @@ async def test_batch_rate_limiter_without_user_context(openai_files: respx.MockR assert usage_without_context.total_tokens == usage_with_context.total_tokens == _token_counter_total(rows) -@pytest.mark.asyncio -async def test_managed_file_is_read_through_the_managed_files_hook_with_user_context( - openai_files: respx.MockRouter, monkeypatch: pytest.MonkeyPatch -): +async def _install_managed_file_hook( + openai_files: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, rows: str +) -> tuple[respx.Route, respx.Route]: from litellm_enterprise.proxy.hooks.managed_files import _PROXY_LiteLLMManagedFiles from litellm import Router @@ -494,24 +494,22 @@ async def test_managed_file_is_read_through_the_managed_files_hook_with_user_con from litellm.proxy.utils import ProxyLogging assert is_base64_encoded_unified_file_id(_MANAGED_FILE_ID) - rows: Final = _batch_rows(("Test message for regression",)) provider_route: Final = _serve_file(openai_files, "file-provider-1", rows) - standard_route: Final = _serve_file(openai_files, "file-abc123", rows) unrouted_managed_read: Final = _serve_file(openai_files, _MANAGED_FILE_ID, rows) - file_cache: Final = InternalUsageCache(dual_cache=DualCache()) - await file_cache.async_set_cache( - key=_MANAGED_FILE_ID, - value=LiteLLM_ManagedFileTable( - unified_file_id=_MANAGED_FILE_ID, - model_mappings={"deployment-1": "file-provider-1"}, - flat_model_file_ids=["file-provider-1"], - created_by="test-user-regression", - ).model_dump(), - litellm_parent_otel_span=None, + row: Final = LiteLLM_ManagedFileTable( + unified_file_id=_MANAGED_FILE_ID, + model_mappings={"deployment-1": "file-provider-1"}, + flat_model_file_ids=["file-provider-1"], + created_by="test-user-regression", + team_id="team-owner", ) + file_cache: Final = InternalUsageCache(dual_cache=DualCache()) + await file_cache.async_set_cache(key=_MANAGED_FILE_ID, value=row.model_dump(), litellm_parent_otel_span=None) + prisma_client: Final = MagicMock() + prisma_client.db.litellm_managedfiletable.find_first = AsyncMock(return_value=row) proxy_logging: Final = ProxyLogging(user_api_key_cache=DualCache()) proxy_logging.proxy_hook_mapping["managed_files"] = _PROXY_LiteLLMManagedFiles( - internal_usage_cache=file_cache, prisma_client=None + internal_usage_cache=file_cache, prisma_client=prisma_client ) router: Final = Router( model_list=[ @@ -524,6 +522,16 @@ async def test_managed_file_is_read_through_the_managed_files_hook_with_user_con ) monkeypatch.setattr(proxy_server, "proxy_logging_obj", proxy_logging) monkeypatch.setattr(proxy_server, "llm_router", router) + return provider_route, unrouted_managed_read + + +@pytest.mark.asyncio +async def test_managed_file_is_read_through_the_managed_files_hook_with_user_context( + openai_files: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + rows: Final = _batch_rows(("Test message for regression",)) + provider_route, unrouted_managed_read = await _install_managed_file_hook(openai_files, monkeypatch, rows) + standard_route: Final = _serve_file(openai_files, "file-abc123", rows) _, _, batch_limiter = _make_limiters() user_api_key_dict: Final = UserAPIKeyAuth( api_key="test-key-regression", user_id="test-user-regression", tpm_limit=1000, rpm_limit=10 @@ -541,3 +549,25 @@ async def test_managed_file_is_read_through_the_managed_files_hook_with_user_con assert not unrouted_managed_read.called assert managed_usage.request_count == standard_usage.request_count == 1 assert managed_usage.total_tokens == standard_usage.total_tokens == _token_counter_total(rows) + + +@pytest.mark.asyncio +async def test_managed_file_of_another_tenant_is_denied_before_it_is_read( + openai_files: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + rows: Final = _batch_rows(("Test message for regression",)) + provider_route, unrouted_managed_read = await _install_managed_file_hook(openai_files, monkeypatch, rows) + _, _, batch_limiter = _make_limiters() + outsider: Final = UserAPIKeyAuth(api_key="test-key-outsider", user_id="outsider", team_id="team-other", tpm_limit=1) + + with pytest.raises(HTTPException) as denied: + await batch_limiter.async_pre_call_hook( + user_api_key_dict=outsider, + cache=DualCache(), + data=dict(_create_batch_data(_MANAGED_FILE_ID)), + call_type="acreate_batch", + ) + + assert denied.value.status_code == 403 + assert not provider_route.called + assert not unrouted_managed_read.called diff --git a/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py b/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py index 100e79e371d..833af4c03a3 100644 --- a/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py +++ b/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py @@ -457,6 +457,35 @@ async def test_async_post_call_failure_hook_non_llm_route(): mock_update_database.assert_not_called() +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("team_id", "route", "written"), + [ + ("litellm-dashboard", "/key/list", False), + ("litellm-dashboard", "/v2/user/info", False), + ("litellm-dashboard", "/chat/completions", True), + ("team-1", "/key/list", True), + (None, "/key/list", True), + ], +) +async def test_async_post_call_failure_hook_skips_dashboard_session_info_route_failures( # test-quality-ok: update_database is the DB write boundary, so whether it is awaited is the observable outcome + team_id: str | None, route: str, written: bool +) -> None: + """An expired Admin UI session key refetching dashboard info routes must not + write a failure SpendLog row, while its LLM calls and other keys' info calls still do.""" + with patch( + "litellm.proxy.db.db_spend_update_writer.DBSpendUpdateWriter.update_database", + new_callable=AsyncMock, + ) as mock_update_database: + await ProxyDBLogger().async_post_call_failure_hook( + request_data={}, + original_exception=Exception("Authentication Error - Expired Key"), + user_api_key_dict=UserAPIKeyAuth(api_key="sk-session", team_id=team_id, request_route=route), + ) + + assert mock_update_database.await_count == (1 if written else 0) + + @pytest.mark.asyncio async def test_async_post_call_failure_hook_releases_budget_reservation_before_route_skip(): logger = ProxyDBLogger() diff --git a/tests/unit/proxy/management_endpoints/test_cost_tracking_settings.py b/tests/unit/proxy/management_endpoints/test_cost_tracking_settings.py index e155ef6fbaa..06a468bf944 100644 --- a/tests/unit/proxy/management_endpoints/test_cost_tracking_settings.py +++ b/tests/unit/proxy/management_endpoints/test_cost_tracking_settings.py @@ -166,6 +166,80 @@ class TestCostTrackingSettings: # Verify litellm.cost_discount_config was updated assert litellm.cost_discount_config == test_discount_config + @pytest.mark.asyncio + async def test_update_cost_discount_config_model_pattern_key(self): + """ + Test PATCH /config/cost_discount_config accepts / keys. + """ + mock_proxy_config = AsyncMock() + mock_proxy_config.get_config = AsyncMock(return_value={"litellm_settings": {}}) + mock_proxy_config.save_config = AsyncMock() + + mock_prisma_client = MagicMock() + + test_discount_config = {"vertex_ai/claude-*": 0.2} + + with ( + patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client), + patch("litellm.proxy.proxy_server.proxy_config", mock_proxy_config), + patch("litellm.proxy.proxy_server.store_model_in_db", True), + patch.object(litellm, "cost_discount_config", {}), + ): + response = client.patch( + "/config/cost_discount_config", + json=test_discount_config, + headers={"Authorization": "Bearer sk-9876"}, + ) + + assert response.status_code == 200 + assert response.json()["values"]["vertex_ai/claude-*"] == 0.2 + mock_proxy_config.save_config.assert_called_once() + assert litellm.cost_discount_config == test_discount_config + + @pytest.mark.asyncio + async def test_update_cost_discount_config_invalid_provider_with_pattern(self): + """ + Test PATCH /config/cost_discount_config rejects invalid providers in pattern keys. + """ + mock_proxy_config = AsyncMock() + mock_prisma_client = MagicMock() + + with ( + patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client), + patch("litellm.proxy.proxy_server.proxy_config", mock_proxy_config), + patch("litellm.proxy.proxy_server.store_model_in_db", True), + ): + response = client.patch( + "/config/cost_discount_config", + json={"notaprovider/claude-*": 0.2}, + headers={"Authorization": "Bearer sk-9876"}, + ) + + assert response.status_code == 400 + assert "notaprovider" in response.json()["detail"]["error"] + + @pytest.mark.asyncio + async def test_update_cost_discount_config_empty_model_pattern(self): + """ + Test PATCH /config/cost_discount_config rejects a key ending in a bare slash. + """ + mock_proxy_config = AsyncMock() + mock_prisma_client = MagicMock() + + with ( + patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client), + patch("litellm.proxy.proxy_server.proxy_config", mock_proxy_config), + patch("litellm.proxy.proxy_server.store_model_in_db", True), + ): + response = client.patch( + "/config/cost_discount_config", + json={"vertex_ai/": 0.2}, + headers={"Authorization": "Bearer sk-9876"}, + ) + + assert response.status_code == 400 + assert "cannot be empty" in response.json()["detail"] + @pytest.mark.asyncio async def test_update_cost_discount_config_invalid_provider(self): """ diff --git a/tests/unit/proxy/management_endpoints/test_team_callback_endpoints.py b/tests/unit/proxy/management_endpoints/test_team_callback_endpoints.py index ad5b6903ece..ac78085ca47 100644 --- a/tests/unit/proxy/management_endpoints/test_team_callback_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_team_callback_endpoints.py @@ -8,11 +8,12 @@ redacted audit rows for callback mutations. import json from collections.abc import Iterator -from typing import Final +from typing import Final, cast from unittest.mock import AsyncMock, MagicMock, Mock, patch import pytest from fastapi import HTTPException, Request +from pydantic import ValidationError import litellm from litellm.proxy._types import ( @@ -21,6 +22,7 @@ from litellm.proxy._types import ( LitellmUserRoles, UserAPIKeyAuth, ) +from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_utils.callback_config_validation import cross_entry_family_error from litellm.proxy.management.teams.authz import TeamAccess from litellm.proxy.management_endpoints.team_callback_endpoints import ( @@ -28,6 +30,7 @@ from litellm.proxy.management_endpoints.team_callback_endpoints import ( delete_team_callback, disable_team_logging, get_team_callbacks, + router, ) @@ -1434,7 +1437,7 @@ async def test_delete_team_callback_route_accepts_team_ids_containing_slashes(): } mock_prisma = _patch_prisma(_team_row(team_id=team_id, metadata=metadata)) - app = FastAPI() + app: Final = FastAPI() app.include_router(router) app.dependency_overrides[user_api_key_auth] = _admin_auth @@ -1683,6 +1686,72 @@ async def test_a_second_entry_may_not_flip_the_span_scope(patched_prisma, caller patched_prisma.db.litellm_teamtable.update.assert_awaited_once() +@pytest.mark.usefixtures("otel_v2_on") +@pytest.mark.asyncio +async def test_a_second_newrelic_entry_may_not_change_excluded_services(patched_prisma) -> None: + patched_prisma.get_data = AsyncMock( + return_value=_team_row( + metadata={ + "logging": [ + { + "callback_name": "newrelic", + "callback_type": "success", + "callback_vars": {"newrelic_api_key": "license-key", "excluded_services": "redis"}, + } + ] + } + ) + ) + data: Final = AddTeamCallback( + callback_name="newrelic", + callback_type="success_and_failure", + callback_vars={"newrelic_api_key": "license-key", "excluded_services": "postgres"}, + ) + with pytest.raises(HTTPException) as exc: + await add_team_callbacks( + data=data, + http_request=Mock(spec=Request), + team_id="team-victim", + user_api_key_dict=_admin_auth(), + ) + assert exc.value.status_code == 400 + assert "excluded_services" in str(exc.value.detail) + patched_prisma.db.litellm_teamtable.update.assert_not_called() + + +@pytest.mark.usefixtures("otel_v2_on") +@pytest.mark.asyncio +async def test_a_second_newrelic_entry_accepts_equivalent_excluded_services(patched_prisma) -> None: + patched_prisma.get_data = AsyncMock( + return_value=_team_row( + metadata={ + "logging": [ + { + "callback_name": "newrelic", + "callback_type": "success", + "callback_vars": { + "newrelic_api_key": "license-key", + "excluded_services": "redis,postgres", + }, + } + ] + } + ) + ) + data: Final = AddTeamCallback( + callback_name="newrelic", + callback_type="success_and_failure", + callback_vars={"newrelic_api_key": "license-key", "excluded_services": "postgres,redis"}, + ) + await add_team_callbacks( + data=data, + http_request=Mock(spec=Request), + team_id="team-victim", + user_api_key_dict=_admin_auth(), + ) + patched_prisma.db.litellm_teamtable.update.assert_awaited_once() + + @pytest.mark.asyncio async def test_add_team_callbacks_rejects_out_of_range_arize_sampling_rate(patched_prisma): data = AddTeamCallback( @@ -1772,3 +1841,59 @@ async def test_add_team_callbacks_stores_capture_message_content_in_the_existing patched_prisma.db.litellm_teamtable.update.assert_awaited_once() stored: Final = json.loads(patched_prisma.db.litellm_teamtable.update.await_args.kwargs["data"]["metadata"]) assert [entry["callback_vars"].get("capture_message_content") for entry in stored["logging"]] == [capture] + + +@pytest.mark.usefixtures("otel_v2_on") +@pytest.mark.asyncio +async def test_add_team_callbacks_stores_excluded_services_in_the_existing_callback_metadata(patched_prisma) -> None: + data: Final = AddTeamCallback( + callback_name="langfuse_otel", + callback_type="success", + callback_vars={ + "langfuse_public_key": "pk", + "langfuse_secret_key": "sk", + "excluded_services": "redis,postgres", + }, + ) + await add_team_callbacks( + data=data, + http_request=Mock(spec=Request), + team_id="team-victim", + user_api_key_dict=_admin_auth(), + ) + stored: Final = json.loads(patched_prisma.db.litellm_teamtable.update.await_args.kwargs["data"]["metadata"]) + assert [entry["callback_vars"].get("excluded_services") for entry in stored["logging"]] == ["redis,postgres"] + + +def test_add_team_callback_returns_http_422_for_unknown_excluded_services() -> None: + from fastapi import FastAPI + from fastapi.testclient import TestClient + + app = FastAPI() + app.include_router(router) + app.dependency_overrides[user_api_key_auth] = _admin_auth + + response: Final = TestClient(app).post( + "/team/team-victim/callback", + json={ + "callback_name": "langfuse_otel", + "callback_type": "success", + "callback_vars": { + "langfuse_public_key": "pk", + "langfuse_secret_key": "sk", + "excluded_services": "auth", + }, + }, + ) + + assert response.status_code == 422 + assert "excluded_services" in response.text + + +@pytest.mark.parametrize("callback_var", ["capture_message_content", "excluded_services"]) +@pytest.mark.parametrize("value", [None, 1]) +def test_add_team_callback_rejects_non_string_destination_callback_vars(callback_var: str, value: object) -> None: + callback_vars: Final = cast("dict[str, str]", {callback_var: value}) + with pytest.raises(ValidationError) as exc: + AddTeamCallback(callback_name="newrelic", callback_vars=callback_vars) + assert f"{callback_var} must be a string, got {type(value).__name__}" in str(exc.value) diff --git a/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_typesafe_passthrough_logging_handler.py b/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_typesafe_passthrough_logging_handler.py index 7d9f4ae33c8..79209a32f28 100644 --- a/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_typesafe_passthrough_logging_handler.py +++ b/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_typesafe_passthrough_logging_handler.py @@ -1,5 +1,5 @@ from datetime import datetime -from typing import Final +from typing import Final, Literal from unittest.mock import MagicMock import httpx @@ -72,6 +72,33 @@ def test_uses_registry_pricing_and_standard_usage(): assert response["kwargs"]["combined_usage_object"].total_tokens == 360 +def test_uses_decisions_pricing_fields_when_the_row_has_them(monkeypatch: pytest.MonkeyPatch) -> None: + model: Final = "nimble-latest" + pricing: Final = { + "input_cost_per_token": 1.0, + "input_cost_per_token_decisions": 2.0, + "output_cost_per_token": 1.0, + "output_cost_per_token_decisions": 3.0, + "litellm_provider": "bespoke", + "mode": "chat", + } + monkeypatch.setitem(litellm.model_cost, f"bespoke/{model}", pricing) + result: Final = TypeSafePassthroughLoggingHandler.typesafe_passthrough_handler( + httpx_response=_response(), + response_body={"model": model, "usage": {"input_tokens": 2, "output_tokens": 1}}, + logging_obj=_logging_obj(), + url_route="https://api.bespoke.example/v1/systemone", + result='{"answers": {}}', + start_time=datetime(2000, 1, 1), + end_time=datetime(2000, 1, 1), + cache_hit=False, + request_body={"model": model}, + custom_llm_provider="bespoke", + ) + + assert result["kwargs"]["response_cost"] == pytest.approx(7.0) + + def test_falls_back_to_request_model_when_response_model_is_missing(): result = _handler_result( {"usage": {"input_tokens": 10, "output_tokens": 2}}, @@ -143,68 +170,111 @@ def test_success_handler_dispatches_to_typesafe_handler(): @pytest.mark.parametrize("mode", ("decisions", "evaluation"), ids=("decisions", "legacy-evaluation")) @pytest.mark.parametrize("guardrail_cost", [0.0, 0.25]) @pytest.mark.parametrize("metadata_slot", ["metadata", "litellm_metadata"]) -@pytest.mark.parametrize("provider,requested,routing_model", [ - ("laya", "english", "multilingual"), ("laya", "english", None), - ("bespoke", "nimble-latest", None), - ("bespoke", "bespokelabs/Bespoke-Nimble-9B", None), -]) +@pytest.mark.parametrize( + "provider,requested,routing_model", + [ + ("laya", "english", "multilingual"), + ("laya", "english", None), + ("bespoke", "nimble-latest", None), + ("bespoke", "bespokelabs/Bespoke-Nimble-9B", None), + ], +) async def test_oss_gateway_accounts_for_checkpoint_usage_and_registered_cost( - monkeypatch: pytest.MonkeyPatch, routing_model: str | None, metadata_slot: str, guardrail_cost: float, - provider: str, requested: str, mode: str + monkeypatch: pytest.MonkeyPatch, + routing_model: str | None, + metadata_slot: str, + guardrail_cost: float, + provider: str, + requested: str, + mode: Literal["decisions", "evaluation"], ) -> None: checkpoint: Final = routing_model or requested model: Final = f"{provider}/{checkpoint}" input_rate: Final = 0.002 output_rate: Final = 0.005 - monkeypatch.setitem(litellm.model_cost, model, { - "input_cost_per_token": input_rate, "output_cost_per_token": output_rate, - "litellm_provider": provider, "mode": mode, - }) + pricing_row: Final = { + "input_cost_per_token": input_rate, + "output_cost_per_token": output_rate, + "litellm_provider": provider, + "mode": mode, + } + monkeypatch.setitem(litellm.model_cost, model, pricing_row) start: Final = datetime.now() logging_obj: Final = Logging( - model=requested, messages=[], stream=False, call_type="pass_through_endpoint", - start_time=start, litellm_call_id="oss-accounting", function_id="oss-accounting", kwargs={}, + model=requested, + messages=[], + stream=False, + call_type="pass_through_endpoint", + start_time=start, + litellm_call_id="oss-accounting", + function_id="oss-accounting", + kwargs={}, ) from fastapi import Request + from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.pass_through_endpoints.pass_through_endpoints import HttpPassThroughEndpointHelpers - request: Final = Request({ - "type": "http", "method": "POST", "path": f"/{provider}/v1/systemone", - "headers": [], "query_string": b"", - }) + request: Final = Request( + { + "type": "http", + "method": "POST", + "path": f"/{provider}/v1/systemone", + "headers": [], + "query_string": b"", + } + ) auth: Final = UserAPIKeyAuth( - api_key="oss-budget-key", token="oss-budget-key", + api_key="oss-budget-key", + token="oss-budget-key", model_max_budget={f"{provider}/{requested}": {"budget_limit": 0.01, "time_period": "1d"}}, ) request_body: Final = {"model": requested, metadata_slot: {"model_group": "unbounded-client-choice"}} logging_kwargs: Final = HttpPassThroughEndpointHelpers._init_kwargs_for_pass_through_endpoint( - request=request, user_api_key_dict=auth, logging_obj=logging_obj, - passthrough_logging_payload={"url": f"https://{provider}.test/v1/systemone"}, _parsed_body=request_body, + request=request, + user_api_key_dict=auth, + logging_obj=logging_obj, + passthrough_logging_payload={"url": f"https://{provider}.test/v1/systemone"}, + _parsed_body=request_body, ) logging_kwargs["litellm_params"]["metadata"]["standard_logging_guardrail_information"] = [ {"guardrail_name": "trusted-hook", "guardrail_cost": guardrail_cost}, ] logging_obj.update_environment_variables( - model=requested, user="unknown", optional_params={}, - litellm_params=logging_kwargs["litellm_params"], call_type="pass_through_endpoint", + model=requested, + user="unknown", + optional_params={}, + litellm_params=logging_kwargs["litellm_params"], + call_type="pass_through_endpoint", ) body: Final = { - "model": "laya-rl-agent" if provider == "laya" else requested, "usage": {"input_tokens": 10, "output_tokens": 3}, + "model": "laya-rl-agent" if provider == "laya" else requested, + "usage": {"input_tokens": 10, "output_tokens": 3}, **({"routing": {"model": routing_model}} if routing_model else {}), } normalized: Final = PassThroughEndpointLogging().normalize_llm_passthrough_logging_payload( - httpx_response=httpx.Response(200, request=httpx.Request("POST", f"https://{provider}.test/v1/systemone"), json=body), - response_body=body, request_body={"model": requested}, logging_obj=logging_obj, - url_route=f"https://{provider}.test/v1/systemone", result="{}", start_time=start, - end_time=datetime.now(), cache_hit=False, custom_llm_provider=provider, **logging_kwargs, + httpx_response=httpx.Response( + 200, request=httpx.Request("POST", f"https://{provider}.test/v1/systemone"), json=body + ), + response_body=body, + request_body={"model": requested}, + logging_obj=logging_obj, + url_route=f"https://{provider}.test/v1/systemone", + result="{}", + start_time=start, + end_time=datetime.now(), + cache_hit=False, + custom_llm_provider=provider, + **logging_kwargs, ) logged: Final = normalized["kwargs"] expected_cost: Final = 10 * input_rate + 3 * output_rate assert (logged["model"], logged["custom_llm_provider"]) == (model, provider) assert logged["response_cost"] == pytest.approx(expected_cost) assert logged["combined_usage_object"].model_dump(exclude_none=True) == { - "prompt_tokens": 10, "completion_tokens": 3, "total_tokens": 13, + "prompt_tokens": 10, + "completion_tokens": 3, + "total_tokens": 13, } assert logging_obj.model_call_details["model"] == model assert logging_obj.model_call_details["response_cost"] == pytest.approx(expected_cost) diff --git a/tests/unit/proxy/test_fallback_management_endpoints.py b/tests/unit/proxy/test_fallback_management_endpoints.py index 3ee1b64580d..530d22f02be 100644 --- a/tests/unit/proxy/test_fallback_management_endpoints.py +++ b/tests/unit/proxy/test_fallback_management_endpoints.py @@ -17,9 +17,9 @@ from unittest.mock import AsyncMock, MagicMock, patch import pytest from fastapi import HTTPException +from pydantic import JsonValue, TypeAdapter from litellm import Router -from litellm.proxy.utils import evict_config_param, get_config_param from litellm.proxy.management_endpoints.fallback_management_endpoints import ( FallbackCreateRequest, FallbackDeleteResponse, @@ -28,6 +28,8 @@ from litellm.proxy.management_endpoints.fallback_management_endpoints import ( delete_fallback, get_fallback, ) +from litellm.proxy.proxy_server import ProxyConfig +from litellm.proxy.utils import evict_config_param, get_config_param from litellm.types.router import DeploymentTypedDict, LiteLLMParamsTypedDict @@ -122,6 +124,7 @@ FALLBACK_DEPLOYMENT_ID: Final = "team-fallback-id" FallbackRules = list[dict[str, list[str]] | dict[str, object] | str] RouterSettings = dict[str, FallbackRules | None] +_ROUTER_SETTINGS_JSON: Final = TypeAdapter(dict[str, JsonValue]) def _litellm_params(mock_response: str | None) -> LiteLLMParamsTypedDict: @@ -254,12 +257,14 @@ def _config_row(router_settings: RouterSettings) -> SimpleNamespace: class _StoredRouterSettings: """The LiteLLM_Config router_settings row, with the proxy objects that read it through the config cache""" - def __init__(self, router_settings: RouterSettings | None) -> None: + def __init__(self, router_settings: RouterSettings | None, config_file: RouterSettings | None = None) -> None: self.row: SimpleNamespace | None = None if router_settings is None else _config_row(router_settings) self.prisma_client: Final = MagicMock() self.prisma_client.get_generic_data = AsyncMock(side_effect=lambda **_: self.row) self.prisma_client.db.litellm_config.upsert = AsyncMock(side_effect=self._upsert) - self.proxy_config: Final = MagicMock() + self.proxy_config: Final = MagicMock() if config_file is None else ProxyConfig() + if config_file is not None: + self.proxy_config.router_settings.load_yaml(_ROUTER_SETTINGS_JSON.validate_python(config_file)) self.proxy_config.get_config = AsyncMock(side_effect=self._get_config) async def _upsert(self, where: dict[str, str], data: dict[str, dict[str, str]]) -> None: @@ -376,6 +381,102 @@ class TestFallbackWritesSeeTheLatestStoredRules: assert stored.stored_fallbacks() == [GATEWAY_RULE] +CONFIG_RULE: Final = {"gpt-5.4-mini": ["team-fallback"]} +CONFIG_WINDOW_RULE: Final = {"team-primary": ["gpt-5.4-mini"]} + + +class TestFallbackWritesHonourTheConfigFile: + """A fallbacks list set in config.yaml owns the key: /fallback writes that would change it are refused and + the router keeps serving the file's rules""" + + @pytest.fixture(autouse=True) + async def clean_config_cache(self) -> AsyncIterator[None]: + await evict_config_param("router_settings") + yield + await evict_config_param("router_settings") + + @pytest.fixture + def router(self) -> Router: + return _team_router() + + @pytest.fixture + def stored(self) -> _StoredRouterSettings: + return _StoredRouterSettings({"fallbacks": [CONFIG_RULE]}, config_file={"fallbacks": [CONFIG_RULE]}) + + async def _create( + self, request: FallbackCreateRequest, router: Router, stored: _StoredRouterSettings + ) -> FallbackResponse: + with ( + patch("litellm.proxy.proxy_server.llm_router", router), + patch("litellm.proxy.proxy_server.prisma_client", stored.prisma_client), + patch("litellm.proxy.proxy_server.proxy_config", stored.proxy_config), + patch("litellm.proxy.proxy_server.store_model_in_db", True), + ): + return await create_fallback(request, MagicMock()) + + async def _delete(self, model: str, router: Router, stored: _StoredRouterSettings) -> FallbackDeleteResponse: + with ( + patch("litellm.proxy.proxy_server.llm_router", router), + patch("litellm.proxy.proxy_server.prisma_client", stored.prisma_client), + patch("litellm.proxy.proxy_server.proxy_config", stored.proxy_config), + patch("litellm.proxy.proxy_server.store_model_in_db", True), + ): + return await delete_fallback(model, "general", MagicMock()) + + async def test_a_create_that_adds_a_rule_to_a_config_owned_list_is_refused( + self, router: Router, stored: _StoredRouterSettings + ) -> None: + router.fallbacks = [CONFIG_RULE] + request: Final = FallbackCreateRequest(model="team-primary", fallback_models=["team-fallback"]) + + with pytest.raises(HTTPException) as refused: + await self._create(request, router, stored) + + assert refused.value.status_code == 400 + assert "key 'fallbacks' is set in the config file" in str(refused.value.detail) + assert stored.stored_fallbacks() == [CONFIG_RULE] + stored.prisma_client.db.litellm_config.upsert.assert_not_awaited() + assert router.fallbacks == [CONFIG_RULE] + + async def test_a_delete_of_a_config_owned_rule_is_refused( + self, router: Router, stored: _StoredRouterSettings + ) -> None: + router.fallbacks = [CONFIG_RULE] + + with pytest.raises(HTTPException) as refused: + await self._delete("gpt-5.4-mini", router, stored) + + assert refused.value.status_code == 400 + assert "key 'fallbacks' is set in the config file" in str(refused.value.detail) + assert stored.stored_fallbacks() == [CONFIG_RULE] + stored.prisma_client.db.litellm_config.upsert.assert_not_awaited() + assert router.fallbacks == [CONFIG_RULE] + + async def test_a_create_that_restates_the_config_rule_is_accepted( + self, router: Router, stored: _StoredRouterSettings + ) -> None: + request: Final = FallbackCreateRequest(model="gpt-5.4-mini", fallback_models=["team-fallback"]) + + response: Final = await self._create(request, router, stored) + + assert response.fallback_models == ["team-fallback"] + assert stored.stored_fallbacks() == [CONFIG_RULE] + + async def test_a_sibling_list_the_config_file_does_not_set_stays_writable( + self, router: Router, stored: _StoredRouterSettings + ) -> None: + request: Final = FallbackCreateRequest( + model="team-primary", fallback_models=["gpt-5.4-mini"], fallback_type="context_window" + ) + + response: Final = await self._create(request, router, stored) + + assert response.fallback_type == "context_window" + assert stored.row is not None + assert stored.row.param_value["context_window_fallbacks"] == [CONFIG_WINDOW_RULE] + assert router.context_window_fallbacks == [CONFIG_WINDOW_RULE] + + @pytest.mark.asyncio class TestCreateFallback: """Test the create_fallback endpoint""" diff --git a/tests/unit/proxy/utils/proxy_logging/test_post_call_failure_hook.py b/tests/unit/proxy/utils/proxy_logging/test_post_call_failure_hook.py index 0e007429789..84aa1798da4 100644 --- a/tests/unit/proxy/utils/proxy_logging/test_post_call_failure_hook.py +++ b/tests/unit/proxy/utils/proxy_logging/test_post_call_failure_hook.py @@ -66,6 +66,48 @@ def test_is_proxy_only_llm_api_truth_table(proxy_logging): } +def test_is_proxy_only_llm_api_skips_dashboard_session_info_route_failures(proxy_logging): + """An expired Admin UI session key refetching dashboard info routes must not + write Request Log failure rows, while its LLM calls and other keys still do.""" + + def decide(route: str, team_id: str | None) -> bool: + return proxy_logging._is_proxy_only_llm_api_error( + original_exception=Exception("Authentication Error - Expired Key"), + error_type=ProxyErrorTypes.auth_error, + route=route, + team_id=team_id, + ) + + assert { + "dashboard_key_list": decide("/key/list", "litellm-dashboard"), + "dashboard_user_info": decide("/v2/user/info", "litellm-dashboard"), + "dashboard_model_info": decide("/v2/model/info", "litellm-dashboard"), + "dashboard_chat": decide("/chat/completions", "litellm-dashboard"), + "team_key_list": decide("/key/list", "team-1"), + "teamless_key_list": decide("/key/list", None), + } == { + "dashboard_key_list": False, + "dashboard_user_info": False, + "dashboard_model_info": False, + "dashboard_chat": True, + "team_key_list": True, + "teamless_key_list": True, + } + + +@pytest.mark.asyncio +async def test_post_call_failure_hook_skips_logging_for_dashboard_session_info_route(proxy_logging): + proxy_logging._handle_logging_proxy_only_error = AsyncMock() + await proxy_logging.post_call_failure_hook( + request_data={}, + original_exception=Exception("Authentication Error - Expired Key"), + user_api_key_dict=UserAPIKeyAuth(team_id="litellm-dashboard", request_route="/key/list"), + error_type=ProxyErrorTypes.auth_error, + route="/key/list", + ) + proxy_logging._handle_logging_proxy_only_error.assert_not_awaited() + + def test_is_proxy_only_llm_api_missing_exception_raises(proxy_logging): """Passing nothing should TypeError on the missing positional kwarg.""" with pytest.raises(TypeError): diff --git a/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py b/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py index ac02354e4d9..7006feb22b2 100644 --- a/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py +++ b/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py @@ -3,6 +3,7 @@ from copy import deepcopy from typing import Final, Literal import pytest +from openai.types.responses import ResponseFunctionToolCall from openai.types.responses.response_function_web_search import ( ActionFind, ActionOpenPage, @@ -39,11 +40,7 @@ class TestLiteLLMCompletionResponsesConfig: input_item = {"type": "input_file", "file_id": "file-abc123xyz"} # Execute - result = ( - LiteLLMCompletionResponsesConfig._transform_input_file_item_to_file_item( - input_item - ) - ) + result = LiteLLMCompletionResponsesConfig._transform_input_file_item_to_file_item(input_item) # Assert expected = {"type": "file", "file": {"file_id": "file-abc123xyz"}} @@ -58,11 +55,7 @@ class TestLiteLLMCompletionResponsesConfig: input_item = {"type": "input_file", "file_data": file_data} # Execute - result = ( - LiteLLMCompletionResponsesConfig._transform_input_file_item_to_file_item( - input_item - ) - ) + result = LiteLLMCompletionResponsesConfig._transform_input_file_item_to_file_item(input_item) # Assert expected = {"type": "file", "file": {"file_data": file_data}} @@ -80,11 +73,7 @@ class TestLiteLLMCompletionResponsesConfig: } # Execute - result = ( - LiteLLMCompletionResponsesConfig._transform_input_file_item_to_file_item( - input_item - ) - ) + result = LiteLLMCompletionResponsesConfig._transform_input_file_item_to_file_item(input_item) # Assert expected = { @@ -102,11 +91,7 @@ class TestLiteLLMCompletionResponsesConfig: input_item = {"type": "input_file"} # Execute - result = ( - LiteLLMCompletionResponsesConfig._transform_input_file_item_to_file_item( - input_item - ) - ) + result = LiteLLMCompletionResponsesConfig._transform_input_file_item_to_file_item(input_item) # Assert expected = {"type": "file", "file": {}} @@ -125,11 +110,7 @@ class TestLiteLLMCompletionResponsesConfig: } # Execute - result = ( - LiteLLMCompletionResponsesConfig._transform_input_file_item_to_file_item( - input_item - ) - ) + result = LiteLLMCompletionResponsesConfig._transform_input_file_item_to_file_item(input_item) # Assert expected = {"type": "file", "file": {"file_id": "file-abc123xyz"}} @@ -139,14 +120,12 @@ class TestLiteLLMCompletionResponsesConfig: def test_transform_input_file_item_to_file_item_keeps_filename(self): """OpenAI rejects file_data with no filename beside it, so dropping it 400s the request""" - result = ( - LiteLLMCompletionResponsesConfig._transform_input_file_item_to_file_item( - { - "type": "input_file", - "filename": "report.pdf", - "file_data": "data:application/pdf;base64,JVBERi0=", - } - ) + result = LiteLLMCompletionResponsesConfig._transform_input_file_item_to_file_item( + { + "type": "input_file", + "filename": "report.pdf", + "file_data": "data:application/pdf;base64,JVBERi0=", + } ) assert result == { "type": "file", @@ -158,10 +137,8 @@ class TestLiteLLMCompletionResponsesConfig: def test_transform_input_file_item_to_file_item_with_file_url(self): """file_url should be mapped to file_id for downstream URL handling""" - result = ( - LiteLLMCompletionResponsesConfig._transform_input_file_item_to_file_item( - {"type": "input_file", "file_url": "https://example.com/doc.pdf"} - ) + result = LiteLLMCompletionResponsesConfig._transform_input_file_item_to_file_item( + {"type": "input_file", "file_url": "https://example.com/doc.pdf"} ) assert result == { "type": "file", @@ -170,14 +147,12 @@ class TestLiteLLMCompletionResponsesConfig: def test_transform_input_file_item_file_id_takes_precedence_over_file_url(self): """explicit file_id should not be overwritten by file_url""" - result = ( - LiteLLMCompletionResponsesConfig._transform_input_file_item_to_file_item( - { - "type": "input_file", - "file_id": "file-abc123", - "file_url": "https://example.com/doc.pdf", - } - ) + result = LiteLLMCompletionResponsesConfig._transform_input_file_item_to_file_item( + { + "type": "input_file", + "file_id": "file-abc123", + "file_url": "https://example.com/doc.pdf", + } ) assert result == {"type": "file", "file": {"file_id": "file-abc123"}} @@ -188,11 +163,7 @@ class TestLiteLLMCompletionResponsesConfig: input_item = {"type": "input_image", "image_url": image_url, "detail": "high"} # Execute - result = ( - LiteLLMCompletionResponsesConfig._transform_input_image_item_to_image_item( - input_item - ) - ) + result = LiteLLMCompletionResponsesConfig._transform_input_image_item_to_image_item(input_item) # Assert expected = { @@ -211,11 +182,7 @@ class TestLiteLLMCompletionResponsesConfig: input_item = {"type": "input_image", "image_url": image_url, "detail": "high"} # Execute - result = ( - LiteLLMCompletionResponsesConfig._transform_input_image_item_to_image_item( - input_item - ) - ) + result = LiteLLMCompletionResponsesConfig._transform_input_image_item_to_image_item(input_item) # Assert expected = { @@ -234,11 +201,7 @@ class TestLiteLLMCompletionResponsesConfig: input_item = {"type": "input_image", "image_url": image_url} # Execute - result = ( - LiteLLMCompletionResponsesConfig._transform_input_image_item_to_image_item( - input_item - ) - ) + result = LiteLLMCompletionResponsesConfig._transform_input_image_item_to_image_item(input_item) # Assert expected = { @@ -256,11 +219,7 @@ class TestLiteLLMCompletionResponsesConfig: input_item = {"type": "input_image"} # Execute - result = ( - LiteLLMCompletionResponsesConfig._transform_input_image_item_to_image_item( - input_item - ) - ) + result = LiteLLMCompletionResponsesConfig._transform_input_image_item_to_image_item(input_item) # Assert expected = {"type": "image_url", "image_url": {"url": "", "detail": "auto"}} @@ -280,11 +239,7 @@ class TestLiteLLMCompletionResponsesConfig: } # Execute - result = ( - LiteLLMCompletionResponsesConfig._transform_input_image_item_to_image_item( - input_item - ) - ) + result = LiteLLMCompletionResponsesConfig._transform_input_image_item_to_image_item(input_item) # Assert expected = { @@ -320,28 +275,26 @@ class TestLiteLLMCompletionResponsesConfig: ) # Execute - responses_api_response = LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( - request_input="What is the meaning of life?", - responses_api_request={}, - chat_completion_response=chat_completion_response, + responses_api_response = ( + LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( + request_input="What is the meaning of life?", + responses_api_request={}, + chat_completion_response=chat_completion_response, + ) ) # Assert assert hasattr(responses_api_response, "output") assert len(responses_api_response.output) >= 2 - reasoning_items = [ - item for item in responses_api_response.output if item.type == "reasoning" - ] + reasoning_items = [item for item in responses_api_response.output if item.type == "reasoning"] assert len(reasoning_items) == 1, "Should have exactly one reasoning item" reasoning_item = reasoning_items[0] # Note: ID auto-generation was disabled, so reasoning items may not have IDs # Only assert ID format if an ID is present if hasattr(reasoning_item, "id") and reasoning_item.id: - assert reasoning_item.id.startswith( - "rs_" - ), f"Expected ID to start with 'rs_', got: {reasoning_item.id}" + assert reasoning_item.id.startswith("rs_"), f"Expected ID to start with 'rs_', got: {reasoning_item.id}" assert reasoning_item.status == "completed" assert reasoning_item.role == "assistant" assert len(reasoning_item.content) == 1 @@ -349,9 +302,7 @@ class TestLiteLLMCompletionResponsesConfig: assert "step by step" in reasoning_item.content[0].text assert "42" in reasoning_item.content[0].text - message_items = [ - item for item in responses_api_response.output if item.type == "message" - ] + message_items = [item for item in responses_api_response.output if item.type == "message"] assert len(message_items) == 1, "Should have exactly one message item" message_item = message_items[0] @@ -378,21 +329,19 @@ class TestLiteLLMCompletionResponsesConfig: ) # Execute - responses_api_response = LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( - request_input="A simple question?", - responses_api_request={}, - chat_completion_response=chat_completion_response, + responses_api_response = ( + LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( + request_input="A simple question?", + responses_api_request={}, + chat_completion_response=chat_completion_response, + ) ) # Assert - reasoning_items = [ - item for item in responses_api_response.output if item.type == "reasoning" - ] + reasoning_items = [item for item in responses_api_response.output if item.type == "reasoning"] assert len(reasoning_items) == 0, "Should have no reasoning items" - message_items = [ - item for item in responses_api_response.output if item.type == "message" - ] + message_items = [item for item in responses_api_response.output if item.type == "message"] assert len(message_items) == 1, "Should have exactly one message item" assert message_items[0].content[0].text == "Just a regular answer." assert responses_api_response.object == "response" @@ -428,22 +377,20 @@ class TestLiteLLMCompletionResponsesConfig: ) # Execute - responses_api_response = LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( - request_input="A question with multiple answers?", - responses_api_request={}, - chat_completion_response=chat_completion_response, + responses_api_response = ( + LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( + request_input="A question with multiple answers?", + responses_api_request={}, + chat_completion_response=chat_completion_response, + ) ) # Assert - reasoning_items = [ - item for item in responses_api_response.output if item.type == "reasoning" - ] + reasoning_items = [item for item in responses_api_response.output if item.type == "reasoning"] assert len(reasoning_items) == 1, "Should have exactly one reasoning item" assert reasoning_items[0].content[0].text == "First reasoning process." - message_items = [ - item for item in responses_api_response.output if item.type == "message" - ] + message_items = [item for item in responses_api_response.output if item.type == "message"] assert len(message_items) == 2, "Should have two message items" def test_signature_only_thinking_block_still_emits_reasoning_item(self): @@ -460,23 +407,21 @@ class TestLiteLLMCompletionResponsesConfig: content="10", role="assistant", reasoning_content="", - thinking_blocks=[ - {"type": "thinking", "thinking": "", "signature": "signature-payload"} - ], + thinking_blocks=[{"type": "thinking", "thinking": "", "signature": "signature-payload"}], ), ) ], ) - responses_api_response = LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( - request_input="Test input", - responses_api_request={}, - chat_completion_response=response, + responses_api_response = ( + LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( + request_input="Test input", + responses_api_request={}, + chat_completion_response=response, + ) ) - reasoning_items = [ - item for item in responses_api_response.output if item.type == "reasoning" - ] + reasoning_items = [item for item in responses_api_response.output if item.type == "reasoning"] assert len(reasoning_items) == 1, "Signature-only thinking should still surface a reasoning item" assert reasoning_items[0].content == [] assert "signature-payload" in reasoning_items[0].encrypted_content @@ -500,15 +445,15 @@ class TestLiteLLMCompletionResponsesConfig: ], ) - responses_api_response = LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( - request_input="Test input", - responses_api_request={}, - chat_completion_response=response, + responses_api_response = ( + LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( + request_input="Test input", + responses_api_request={}, + chat_completion_response=response, + ) ) - reasoning_items = [ - item for item in responses_api_response.output if item.type == "reasoning" - ] + reasoning_items = [item for item in responses_api_response.output if item.type == "reasoning"] assert len(reasoning_items) == 1 assert "redacted-payload" in reasoning_items[0].encrypted_content @@ -526,23 +471,21 @@ class TestLiteLLMCompletionResponsesConfig: content="10", role="assistant", reasoning_content="counting the primes", - thinking_blocks=[ - {"type": "thinking", "thinking": "counting the primes", "signature": "sig"} - ], + thinking_blocks=[{"type": "thinking", "thinking": "counting the primes", "signature": "sig"}], ), ) ], ) - responses_api_response = LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( - request_input="Test input", - responses_api_request={}, - chat_completion_response=response, + responses_api_response = ( + LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( + request_input="Test input", + responses_api_request={}, + chat_completion_response=response, + ) ) - reasoning_items = [ - item for item in responses_api_response.output if item.type == "reasoning" - ] + reasoning_items = [item for item in responses_api_response.output if item.type == "reasoning"] assert len(reasoning_items) == 1 assert reasoning_items[0].content[0].text == "counting the primes" assert "sig" in reasoning_items[0].encrypted_content @@ -571,10 +514,12 @@ class TestLiteLLMCompletionResponsesConfig: ], ) - responses_api_response = LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( - request_input="this is a test", - responses_api_request={}, - chat_completion_response=chat_completion_response, + responses_api_response = ( + LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( + request_input="this is a test", + responses_api_request={}, + chat_completion_response=chat_completion_response, + ) ) assert responses_api_response.status == "completed" @@ -610,15 +555,15 @@ class TestLiteLLMCompletionResponsesConfig: ], ) - responses_api_response = LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( - request_input="this is a test", - responses_api_request={}, - chat_completion_response=chat_completion_response, + responses_api_response = ( + LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( + request_input="this is a test", + responses_api_request={}, + chat_completion_response=chat_completion_response, + ) ) - message_items = [ - item for item in responses_api_response.output if item.type == "message" - ] + message_items = [item for item in responses_api_response.output if item.type == "message"] assert len(message_items) > 0 for item in message_items: @@ -653,10 +598,12 @@ class TestLiteLLMCompletionResponsesConfig: ], ) - responses_api_response = LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( - request_input="this is a test", - responses_api_request={}, - chat_completion_response=chat_completion_response, + responses_api_response = ( + LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( + request_input="this is a test", + responses_api_request={}, + chat_completion_response=chat_completion_response, + ) ) assert responses_api_response.status == "incomplete" @@ -691,10 +638,12 @@ class TestLiteLLMCompletionResponsesConfig: ], ) - responses_api_response = LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( - request_input="what's the weather in SF?", - responses_api_request={}, - chat_completion_response=chat_completion_response, + responses_api_response = ( + LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( + request_input="what's the weather in SF?", + responses_api_request={}, + chat_completion_response=chat_completion_response, + ) ) output_types = [item.type for item in responses_api_response.output] @@ -717,10 +666,12 @@ class TestLiteLLMCompletionResponsesConfig: ], ) - responses_api_response = LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( - request_input="what's the weather in SF?", - responses_api_request={}, - chat_completion_response=chat_completion_response, + responses_api_response = ( + LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( + request_input="what's the weather in SF?", + responses_api_request={}, + chat_completion_response=chat_completion_response, + ) ) message_items = [item for item in responses_api_response.output if item.type == "message"] @@ -754,10 +705,12 @@ class TestLiteLLMCompletionResponsesConfig: } # Execute - responses_api_response = LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( - request_input="Test", - responses_api_request={}, - chat_completion_response=chat_completion_response, + responses_api_response = ( + LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( + request_input="Test", + responses_api_request={}, + chat_completion_response=chat_completion_response, + ) ) # Assert @@ -826,10 +779,12 @@ class TestLiteLLMCompletionResponsesConfig: ) # Execute - responses_api_response = LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( - request_input="Test", - responses_api_request={}, - chat_completion_response=chat_completion_response, + responses_api_response = ( + LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( + request_input="Test", + responses_api_request={}, + chat_completion_response=chat_completion_response, + ) ) # Assert - should default to empty dict @@ -866,36 +821,34 @@ class TestLiteLLMCompletionResponsesConfig: ) try: - responses_api_response = LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( - request_input="Spawn an agent", - responses_api_request={ - "tools": [ - { - "type": "namespace", - "name": "collaboration", - "tools": [ - { - "type": "function", - "name": "spawn_agent", - "parameters": { - "type": "object", - "properties": {}, - }, - } - ], - } - ] - }, - chat_completion_response=chat_completion_response, + responses_api_response = ( + LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( + request_input="Spawn an agent", + responses_api_request={ + "tools": [ + { + "type": "namespace", + "name": "collaboration", + "tools": [ + { + "type": "function", + "name": "spawn_agent", + "parameters": { + "type": "object", + "properties": {}, + }, + } + ], + } + ] + }, + chat_completion_response=chat_completion_response, + ) ) finally: TOOL_CALLS_CACHE.delete_cache(key=tool_call_id) - tool_calls = [ - item - for item in responses_api_response.output - if item.type == "function_call" - ] + tool_calls = [item for item in responses_api_response.output if item.type == "function_call"] assert len(tool_calls) == 1 assert tool_calls[0].name == "spawn_agent" assert tool_calls[0].namespace == "collaboration" @@ -933,33 +886,30 @@ class TestLiteLLMCompletionResponsesConfig: ) try: - responses_api_response = LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( - request_input="What is the weather in Paris?", - responses_api_request={ - "tools": [ - { - "type": "function", - "name": "get_weather", - "parameters": {"type": "object", "properties": {}}, - } - ] - }, - chat_completion_response=chat_completion_response, + responses_api_response = ( + LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( + request_input="What is the weather in Paris?", + responses_api_request={ + "tools": [ + { + "type": "function", + "name": "get_weather", + "parameters": {"type": "object", "properties": {}}, + } + ] + }, + chat_completion_response=chat_completion_response, + ) ) finally: TOOL_CALLS_CACHE.delete_cache(key=tool_call_id) - tool_calls = [ - item - for item in responses_api_response.output - if item.type == "function_call" - ] + tool_calls = [item for item in responses_api_response.output if item.type == "function_call"] assert len(tool_calls) == 1 assert tool_calls[0].name == "get_weather" assert tool_calls[0].namespace is None assert "namespace" not in tool_calls[0].model_fields_set - def test_transform_top_level_function_collision_stays_unnamespaced(self): tool_call_id = "call_top_level_collision" chat_completion_response = ModelResponse( @@ -1033,26 +983,14 @@ class TestFunctionCallTransformation: regular_message = {"type": "message", "role": "user", "content": "Hello"} # Test function_call detection - assert LiteLLMCompletionResponsesConfig._is_input_item_function_call( - function_call_item - ) - assert not LiteLLMCompletionResponsesConfig._is_input_item_function_call( - function_call_output_item - ) - assert not LiteLLMCompletionResponsesConfig._is_input_item_function_call( - regular_message - ) + assert LiteLLMCompletionResponsesConfig._is_input_item_function_call(function_call_item) + assert not LiteLLMCompletionResponsesConfig._is_input_item_function_call(function_call_output_item) + assert not LiteLLMCompletionResponsesConfig._is_input_item_function_call(regular_message) # Test function_call_output detection (should still work) - assert LiteLLMCompletionResponsesConfig._is_input_item_tool_call_output( - function_call_output_item - ) - assert not LiteLLMCompletionResponsesConfig._is_input_item_tool_call_output( - function_call_item - ) - assert not LiteLLMCompletionResponsesConfig._is_input_item_tool_call_output( - regular_message - ) + assert LiteLLMCompletionResponsesConfig._is_input_item_tool_call_output(function_call_output_item) + assert not LiteLLMCompletionResponsesConfig._is_input_item_tool_call_output(function_call_item) + assert not LiteLLMCompletionResponsesConfig._is_input_item_tool_call_output(regular_message) def test_function_call_transformation(self): """Test that function_call items are correctly transformed to assistant messages with tool calls""" @@ -1198,9 +1136,7 @@ class TestFunctionCallTransformation: tool_msg = messages[2] assert tool_msg.get("role") == "tool" assert tool_msg.get("content") == "Rainy" - assert ( - tool_msg.get("tool_call_id") == "call_1fe70e2a-a596-45ef-b72c-9b8567c460e5" - ) + assert tool_msg.get("tool_call_id") == "call_1fe70e2a-a596-45ef-b72c-9b8567c460e5" def test_complete_request_transformation_with_function_calls(self): """Test the complete request transformation that would be used by the responses API""" @@ -1466,9 +1402,7 @@ class TestToolChoiceTransformation: Test that {"type": "tool"} is transformed to "required". This fixes the Anthropic error: "tool_choice.tool.name: Field required" """ - result = LiteLLMCompletionResponsesConfig._transform_tool_choice( - {"type": "tool"} - ) + result = LiteLLMCompletionResponsesConfig._transform_tool_choice({"type": "tool"}) assert result == "required" def test_transform_tool_choice_preserves_function_with_name(self): @@ -1480,9 +1414,7 @@ class TestToolChoiceTransformation: def test_transform_tool_choice_responses_flat_function_name(self): """Responses-API forced-function with a top-level name maps to the nested Chat Completions shape instead of degrading to required and dropping the name""" - result = LiteLLMCompletionResponsesConfig._transform_tool_choice( - {"type": "function", "name": "get_weather"} - ) + result = LiteLLMCompletionResponsesConfig._transform_tool_choice({"type": "function", "name": "get_weather"}) assert result == {"type": "function", "function": {"name": "get_weather"}} def test_transform_tool_choice_custom_follows_function_downgrade(self): @@ -1492,9 +1424,7 @@ class TestToolChoiceTransformation: function tool_choice naming the same tool or it references a tool type absent from the converted request. """ - flat = LiteLLMCompletionResponsesConfig._transform_tool_choice( - {"type": "custom", "name": "ApplyPatch"} - ) + flat = LiteLLMCompletionResponsesConfig._transform_tool_choice({"type": "custom", "name": "ApplyPatch"}) assert flat == {"type": "function", "function": {"name": "ApplyPatch"}} nested = LiteLLMCompletionResponsesConfig._transform_tool_choice( @@ -1508,16 +1438,12 @@ class TestToolChoiceTransformation: def test_transform_tool_choice_function_without_name_falls_back_to_required(self): """A function-type dict with no name still falls back to required""" - result = LiteLLMCompletionResponsesConfig._transform_tool_choice( - {"type": "function"} - ) + result = LiteLLMCompletionResponsesConfig._transform_tool_choice({"type": "function"}) assert result == "required" def test_transform_tool_choice_function_empty_name_falls_back_to_required(self): """An empty top-level name is falsy and must not produce an empty function name""" - result = LiteLLMCompletionResponsesConfig._transform_tool_choice( - {"type": "function", "name": ""} - ) + result = LiteLLMCompletionResponsesConfig._transform_tool_choice({"type": "function", "name": ""}) assert result == "required" @pytest.mark.parametrize( @@ -1571,10 +1497,12 @@ class TestToolChoiceTransformation: ], ) - responses_api_response: Final = LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( - request_input="Run the command pwd.", - responses_api_request={"tool_choice": {"type": "function", "name": "run_command"}}, - chat_completion_response=chat_completion_response, + responses_api_response: Final = ( + LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( + request_input="Run the command pwd.", + responses_api_request={"tool_choice": {"type": "function", "name": "run_command"}}, + chat_completion_response=chat_completion_response, + ) ) assert responses_api_response.tool_choice == {"type": "function", "name": "run_command"} @@ -1594,10 +1522,12 @@ class TestToolChoiceTransformation: ], ) - responses_api_response: Final = LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( - request_input="Run the command pwd.", - responses_api_request={"tool_choice": "any"}, - chat_completion_response=chat_completion_response, + responses_api_response: Final = ( + LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( + request_input="Run the command pwd.", + responses_api_request={"tool_choice": "any"}, + chat_completion_response=chat_completion_response, + ) ) assert responses_api_response.tool_choice == "auto" @@ -1611,20 +1541,12 @@ class TestContentTypeTransformation: Test that 'tool_result' content type is transformed to 'text'. This fixes: Invalid user message - content type 'tool_result' not valid. """ - result = ( - LiteLLMCompletionResponsesConfig._get_chat_completion_request_content_type( - "tool_result" - ) - ) + result = LiteLLMCompletionResponsesConfig._get_chat_completion_request_content_type("tool_result") assert result == "text" def test_input_text_content_type_transformed_to_text(self): """Test that 'input_text' content type is transformed to 'text'""" - result = ( - LiteLLMCompletionResponsesConfig._get_chat_completion_request_content_type( - "input_text" - ) - ) + result = LiteLLMCompletionResponsesConfig._get_chat_completion_request_content_type("input_text") assert result == "text" def test_none_text_blocks_filtered_out(self): @@ -1638,9 +1560,7 @@ class TestContentTypeTransformation: {"type": "text", "text": None}, # Should be filtered out {"type": "text", "text": "another valid"}, ] - result = LiteLLMCompletionResponsesConfig._transform_responses_api_content_to_chat_completion_content( - content - ) + result = LiteLLMCompletionResponsesConfig._transform_responses_api_content_to_chat_completion_content(content) assert len(result) == 2 assert result[0]["text"] == "valid text" assert result[1]["text"] == "another valid" @@ -1659,9 +1579,7 @@ class TestContentTypeTransformation: }, ] - result = LiteLLMCompletionResponsesConfig._transform_responses_api_content_to_chat_completion_content( - content - ) + result = LiteLLMCompletionResponsesConfig._transform_responses_api_content_to_chat_completion_content(content) assert result == [ {"type": "text", "text": "Payload:\n"}, @@ -1685,9 +1603,7 @@ class TestToolTransformation: ( result_tools, web_search_options, - ) = LiteLLMCompletionResponsesConfig.transform_responses_api_tools_to_chat_completion_tools( - tools=tools - ) + ) = LiteLLMCompletionResponsesConfig.transform_responses_api_tools_to_chat_completion_tools(tools=tools) # Assert assert len(result_tools) == 1 @@ -1709,9 +1625,7 @@ class TestToolTransformation: ( result_tools, web_search_options, - ) = LiteLLMCompletionResponsesConfig.transform_responses_api_tools_to_chat_completion_tools( - tools=tools - ) + ) = LiteLLMCompletionResponsesConfig.transform_responses_api_tools_to_chat_completion_tools(tools=tools) # Assert assert len(result_tools) == 1 @@ -1739,9 +1653,7 @@ class TestToolTransformation: ( result_tools, web_search_options, - ) = LiteLLMCompletionResponsesConfig.transform_responses_api_tools_to_chat_completion_tools( - tools=tools - ) + ) = LiteLLMCompletionResponsesConfig.transform_responses_api_tools_to_chat_completion_tools(tools=tools) # Assert - computer_use has no Chat Completions equivalent, so it is dropped assert len(result_tools) == 0 @@ -1842,9 +1754,7 @@ class TestToolTransformation: ( result_tools, web_search_options, - ) = LiteLLMCompletionResponsesConfig.transform_responses_api_tools_to_chat_completion_tools( - tools=tools - ) + ) = LiteLLMCompletionResponsesConfig.transform_responses_api_tools_to_chat_completion_tools(tools=tools) # Assert - custom tool is converted to a function tool assert len(result_tools) == 1 @@ -1869,9 +1779,7 @@ class TestToolTransformation: ( result_tools, web_search_options, - ) = LiteLLMCompletionResponsesConfig.transform_responses_api_tools_to_chat_completion_tools( - tools=tools - ) + ) = LiteLLMCompletionResponsesConfig.transform_responses_api_tools_to_chat_completion_tools(tools=tools) # Assert assert len(result_tools) == 1 @@ -1895,9 +1803,7 @@ class TestToolTransformation: ( result_tools, _, - ) = LiteLLMCompletionResponsesConfig.transform_responses_api_tools_to_chat_completion_tools( - tools=tools - ) + ) = LiteLLMCompletionResponsesConfig.transform_responses_api_tools_to_chat_completion_tools(tools=tools) # Assert assert len(result_tools) == 1 @@ -1918,9 +1824,7 @@ class TestToolTransformation: ( result_tools, _, - ) = LiteLLMCompletionResponsesConfig.transform_responses_api_tools_to_chat_completion_tools( - tools=tools - ) + ) = LiteLLMCompletionResponsesConfig.transform_responses_api_tools_to_chat_completion_tools(tools=tools) # Assert assert len(result_tools) == 1 @@ -1939,9 +1843,7 @@ class TestToolTransformation: tools = [custom_tool] with pytest.raises(ValueError, match="allowed_callers must be a list of strings"): - LiteLLMCompletionResponsesConfig.transform_responses_api_tools_to_chat_completion_tools( - tools=tools - ) + LiteLLMCompletionResponsesConfig.transform_responses_api_tools_to_chat_completion_tools(tools=tools) def test_transform_web_search_tools_to_web_search_options(self): """Test that web_search tools are converted to web_search_options""" @@ -1957,9 +1859,7 @@ class TestToolTransformation: ( result_tools, web_search_options, - ) = LiteLLMCompletionResponsesConfig.transform_responses_api_tools_to_chat_completion_tools( - tools=tools - ) + ) = LiteLLMCompletionResponsesConfig.transform_responses_api_tools_to_chat_completion_tools(tools=tools) # Assert assert len(result_tools) == 0 # Web search is not added to tools @@ -1990,9 +1890,7 @@ class TestToolTransformation: ( result_tools, web_search_options, - ) = LiteLLMCompletionResponsesConfig.transform_responses_api_tools_to_chat_completion_tools( - tools=tools - ) + ) = LiteLLMCompletionResponsesConfig.transform_responses_api_tools_to_chat_completion_tools(tools=tools) # Assert assert len(result_tools) == 1 @@ -2022,9 +1920,7 @@ class TestToolTransformation: ( result_tools, _, - ) = LiteLLMCompletionResponsesConfig.transform_responses_api_tools_to_chat_completion_tools( - tools=tools - ) + ) = LiteLLMCompletionResponsesConfig.transform_responses_api_tools_to_chat_completion_tools(tools=tools) # Assert assert len(result_tools) == 1 @@ -2050,9 +1946,7 @@ class TestToolTransformation: ( result_tools, _, - ) = LiteLLMCompletionResponsesConfig.transform_responses_api_tools_to_chat_completion_tools( - tools=tools - ) + ) = LiteLLMCompletionResponsesConfig.transform_responses_api_tools_to_chat_completion_tools(tools=tools) # Assert assert len(result_tools) == 1 @@ -2091,8 +1985,8 @@ class TestToolTransformation: "eager_input_streaming": eager_input_streaming, } - result_tools: Final = LiteLLMCompletionResponsesConfig.transform_chat_completion_tool_params_to_responses_api_tools( - [chat_tool] + result_tools: Final = ( + LiteLLMCompletionResponsesConfig.transform_chat_completion_tool_params_to_responses_api_tools([chat_tool]) ) assert result_tools[0]["eager_input_streaming"] is eager_input_streaming @@ -2110,9 +2004,7 @@ class TestToolTransformation: ( result_tools, _, - ) = LiteLLMCompletionResponsesConfig.transform_responses_api_tools_to_chat_completion_tools( - tools=tools - ) + ) = LiteLLMCompletionResponsesConfig.transform_responses_api_tools_to_chat_completion_tools(tools=tools) # Assert assert len(result_tools) == 1 @@ -2136,9 +2028,7 @@ class TestToolTransformation: ( result_tools, _, - ) = LiteLLMCompletionResponsesConfig.transform_responses_api_tools_to_chat_completion_tools( - tools=tools - ) + ) = LiteLLMCompletionResponsesConfig.transform_responses_api_tools_to_chat_completion_tools(tools=tools) # Assert assert len(result_tools) == 2 @@ -2170,14 +2060,10 @@ class TestToolTransformation: ( result_tools, web_search_options, - ) = LiteLLMCompletionResponsesConfig.transform_responses_api_tools_to_chat_completion_tools( - tools=tools - ) + ) = LiteLLMCompletionResponsesConfig.transform_responses_api_tools_to_chat_completion_tools(tools=tools) # Assert - assert ( - len(result_tools) == 3 - ) # function, mcp, vertex (web_search becomes options) + assert len(result_tools) == 3 # function, mcp, vertex (web_search becomes options) assert web_search_options is not None # Check function tool @@ -2207,9 +2093,7 @@ class TestToolTransformation: ( result_tools, _, - ) = LiteLLMCompletionResponsesConfig.transform_responses_api_tools_to_chat_completion_tools( - tools=tools - ) + ) = LiteLLMCompletionResponsesConfig.transform_responses_api_tools_to_chat_completion_tools(tools=tools) # Assert assert len(result_tools) == 1 @@ -2219,7 +2103,11 @@ class TestToolTransformation: def test_transform_function_tools_parameters_keep_client_key_order(self): tools = [ - {"type": "function", "name": "a", "parameters": {"properties": {"arg": {"type": "string"}}, "required": ["arg"]}}, + { + "type": "function", + "name": "a", + "parameters": {"properties": {"arg": {"type": "string"}}, "required": ["arg"]}, + }, {"type": "function", "name": "b", "parameters": {"type": "object", "properties": {}}}, ] @@ -2245,9 +2133,7 @@ class TestToolTransformation: ( result_tools, _, - ) = LiteLLMCompletionResponsesConfig.transform_responses_api_tools_to_chat_completion_tools( - tools=tools - ) + ) = LiteLLMCompletionResponsesConfig.transform_responses_api_tools_to_chat_completion_tools(tools=tools) # Assert assert len(result_tools) == 1 @@ -2268,9 +2154,7 @@ class TestToolTransformation: ( result_tools, _, - ) = LiteLLMCompletionResponsesConfig.transform_responses_api_tools_to_chat_completion_tools( - tools=tools - ) + ) = LiteLLMCompletionResponsesConfig.transform_responses_api_tools_to_chat_completion_tools(tools=tools) # Assert assert len(result_tools) == 1 @@ -2292,19 +2176,14 @@ class TestToolTransformation: ( result_tools, _, - ) = LiteLLMCompletionResponsesConfig.transform_responses_api_tools_to_chat_completion_tools( - tools=tools - ) + ) = LiteLLMCompletionResponsesConfig.transform_responses_api_tools_to_chat_completion_tools(tools=tools) # Assert assert len(result_tools) == 1 result_tool = result_tools[0] assert result_tool["function"]["parameters"]["type"] == "object" assert "properties" in result_tool["function"]["parameters"] - assert ( - result_tool["function"]["parameters"]["properties"]["arg"]["type"] - == "string" - ) + assert result_tool["function"]["parameters"]["properties"]["arg"]["type"] == "string" def test_bedrock_anthropic_drops_derived_web_search_options(self): """ @@ -2512,7 +2391,6 @@ class TestToolTransformation: assert result_tools[0]["allowed_callers"] == ["code_execution_20250825"] - @pytest.mark.parametrize("nested", [True, False]) def test_transform_namespace_tools_rejects_invalid_allowed_callers(self, nested): function_tool = { @@ -2536,7 +2414,6 @@ class TestToolTransformation: tools=[namespace_tool] ) - def test_transform_flat_namespace_tools_to_function_tools(self): namespace_tool = { "type": "namespace", @@ -2582,9 +2459,7 @@ class TestToolTransformation: ], } - result = LiteLLMCompletionResponsesConfig.namespace_tool_name_map( - [namespace_tool] - ) + result = LiteLLMCompletionResponsesConfig.namespace_tool_name_map([namespace_tool]) assert result["collaboration__wait_agent"] == ("collaboration", "wait_agent") assert result["wait_agent"] == ("collaboration", "wait_agent") @@ -2628,9 +2503,7 @@ class TestToolTransformation: }, ] - result = LiteLLMCompletionResponsesConfig.namespace_tool_name_map( - tools - ) + result = LiteLLMCompletionResponsesConfig.namespace_tool_name_map(tools) assert result["alpha__run"] == ("alpha", "run") assert result["beta__run"] == ("beta", "run") @@ -2658,7 +2531,6 @@ class TestToolTransformation: assert result["admin__run"] == ("admin", "run") assert "run" not in result - def test_transform_tools_rejects_flattened_name_collision(self): tools = [ { @@ -2685,7 +2557,6 @@ class TestToolTransformation: ): LiteLLMCompletionResponsesConfig.transform_responses_api_tools_to_chat_completion_tools(tools=tools) - def test_restore_namespace_tool_name_leaves_unknown_tool_unchanged(self): tool_name, namespace = LiteLLMCompletionResponsesConfig._restore_namespace_tool_name( "mcp__node_repl", @@ -2711,10 +2582,8 @@ class TestToolTransformation: ], } - result_tools, _ = ( - LiteLLMCompletionResponsesConfig.transform_responses_api_tools_to_chat_completion_tools( - tools=[namespace_tool] - ) + result_tools, _ = LiteLLMCompletionResponsesConfig.transform_responses_api_tools_to_chat_completion_tools( + tools=[namespace_tool] ) assert len(result_tools) == 1 @@ -2739,10 +2608,8 @@ class TestToolTransformation: ], } - result_tools, _ = ( - LiteLLMCompletionResponsesConfig.transform_responses_api_tools_to_chat_completion_tools( - tools=[namespace_tool] - ) + result_tools, _ = LiteLLMCompletionResponsesConfig.transform_responses_api_tools_to_chat_completion_tools( + tools=[namespace_tool] ) assert len(result_tools) == 1 @@ -2964,11 +2831,7 @@ class TestToolTransformation: assert "web_search_options" not in result bedrock_tool_blocks = _bedrock_tools_pt(tools=result["tools"], model=model) - names = [ - block["toolSpec"]["name"] - for block in bedrock_tool_blocks - if "toolSpec" in block - ] + names = [block["toolSpec"]["name"] for block in bedrock_tool_blocks if "toolSpec" in block] assert names == ["noop"] assert not any(name.startswith("litellm_unnamed_tool_") for name in names) @@ -3303,9 +3166,7 @@ class TestUsageTransformation: Choices( finish_reason="stop", index=0, - message=Message( - content="Here is the generated image.", role="assistant" - ), + message=Message(content="Here is the generated image.", role="assistant"), ) ], ) @@ -3530,9 +3391,9 @@ class TestStreamingIDConsistency: # Verify the cached ID is set and matches assert iterator._cached_item_id is not None, "Iterator should cache the item_id" assert iterator._cached_item_id == item_id_1, "Cached ID should match event IDs" - assert iterator._cached_item_id.startswith( - "msg_" - ), "Message item IDs must use the Responses API msg_ prefix (issue #27333)" + assert iterator._cached_item_id.startswith("msg_"), ( + "Message item IDs must use the Responses API msg_ prefix (issue #27333)" + ) def test_streaming_iterator_initial_events_use_cached_id(self): """ @@ -3622,9 +3483,7 @@ class TestStreamingIDConsistency: # Create done events text_done_event = iterator.create_output_text_done_event(complete_response) - content_done_event = iterator.create_output_content_part_done_event( - complete_response - ) + content_done_event = iterator.create_output_content_part_done_event(complete_response) item_done_event = iterator.create_output_item_done_event(complete_response) # Extract IDs @@ -3684,27 +3543,19 @@ class TestStreamingIDConsistency: input=input_items ) - roles = [ - m.get("role") if isinstance(m, dict) else getattr(m, "role", None) - for m in messages - ] + roles = [m.get("role") if isinstance(m, dict) else getattr(m, "role", None) for m in messages] # Must not have two consecutive assistant messages for i in range(len(roles) - 1): - assert not ( - roles[i] == "assistant" and roles[i + 1] == "assistant" - ), f"Consecutive assistant messages at indices {i} and {i+1}: {roles}" + assert not (roles[i] == "assistant" and roles[i + 1] == "assistant"), ( + f"Consecutive assistant messages at indices {i} and {i + 1}: {roles}" + ) # The single assistant message must contain BOTH tool_calls assistant_messages = [ - m - for m in messages - if (m.get("role") if isinstance(m, dict) else getattr(m, "role", None)) - == "assistant" + m for m in messages if (m.get("role") if isinstance(m, dict) else getattr(m, "role", None)) == "assistant" ] - assert ( - len(assistant_messages) == 1 - ), f"Expected 1 assistant message, got {len(assistant_messages)}" + assert len(assistant_messages) == 1, f"Expected 1 assistant message, got {len(assistant_messages)}" assistant_msg = assistant_messages[0] tool_calls = ( @@ -3712,27 +3563,19 @@ class TestStreamingIDConsistency: if isinstance(assistant_msg, dict) else getattr(assistant_msg, "tool_calls", None) ) - assert ( - tool_calls is not None and len(tool_calls) == 2 - ), f"Expected 2 tool_calls in the merged assistant message, got: {tool_calls}" + assert tool_calls is not None and len(tool_calls) == 2, ( + f"Expected 2 tool_calls in the merged assistant message, got: {tool_calls}" + ) - call_ids = [ - (tc.get("id") if isinstance(tc, dict) else getattr(tc, "id", None)) - for tc in tool_calls - ] + call_ids = [(tc.get("id") if isinstance(tc, dict) else getattr(tc, "id", None)) for tc in tool_calls] assert "toolu_01" in call_ids, f"toolu_01 missing from tool_calls: {call_ids}" assert "toolu_02" in call_ids, f"toolu_02 missing from tool_calls: {call_ids}" # Both tool messages must be present tool_messages = [ - m - for m in messages - if (m.get("role") if isinstance(m, dict) else getattr(m, "role", None)) - == "tool" + m for m in messages if (m.get("role") if isinstance(m, dict) else getattr(m, "role", None)) == "tool" ] - assert ( - len(tool_messages) == 2 - ), f"Expected 2 tool messages, got {len(tool_messages)}" + assert len(tool_messages) == 2, f"Expected 2 tool messages, got {len(tool_messages)}" def test_single_tool_call_still_works_after_merge_fix(self): """ @@ -3754,20 +3597,14 @@ class TestStreamingIDConsistency: input=input_items ) - roles = [ - m.get("role") if isinstance(m, dict) else getattr(m, "role", None) - for m in messages - ] + roles = [m.get("role") if isinstance(m, dict) else getattr(m, "role", None) for m in messages] assert "user" in roles assert "assistant" in roles assert "tool" in roles assistant_messages = [ - m - for m in messages - if (m.get("role") if isinstance(m, dict) else getattr(m, "role", None)) - == "assistant" + m for m in messages if (m.get("role") if isinstance(m, dict) else getattr(m, "role", None)) == "assistant" ] assert len(assistant_messages) == 1 @@ -3938,9 +3775,7 @@ class TestEnsureOutputItemContentPartAdded: LiteLLMCompletionStreamingIterator, ) - iterator = LiteLLMCompletionStreamingIterator.__new__( - LiteLLMCompletionStreamingIterator - ) + iterator = LiteLLMCompletionStreamingIterator.__new__(LiteLLMCompletionStreamingIterator) iterator.sent_output_item_added_event = False iterator.sent_content_part_added_event = False iterator._sequence_number = 0 @@ -4181,7 +4016,9 @@ class TestEnsureOutputItemContentPartAdded: function_items = [item for item in items if item.call_id == "call_fn"] custom_items = [item for item in items if item.call_id == "call_custom"] assert len(function_items) == 2 and len(custom_items) == 2 - assert all((item.type, item.name, item.namespace) == ("function_call", "run", "beta") for item in function_items) + assert all( + (item.type, item.name, item.namespace) == ("function_call", "run", "beta") for item in function_items + ) assert function_items[-1].arguments == '{"job_id":"42"}' assert all(item.type == "custom_tool_call" and item.name == "run" for item in custom_items) assert all(getattr(item, "namespace", None) is None for item in custom_items) @@ -4380,7 +4217,6 @@ class TestEnsureOutputItemContentPartAdded: assert added.item.name == "run" assert getattr(added.item, "namespace", None) is None - def test_streaming_namespace_map_is_built_once(self): from unittest.mock import MagicMock, patch @@ -4422,7 +4258,6 @@ class TestEnsureOutputItemContentPartAdded: namespace_map.assert_called_once_with(request["tools"]) - def test_emit_terminal_event_uses_stream_finish_reason(self): """ The terminal stream event type must track the assembled response @@ -4512,9 +4347,7 @@ class TestCacheControlPreservation: "cache_control": {"type": "ephemeral"}, } ] - result = LiteLLMCompletionResponsesConfig._transform_responses_api_content_to_chat_completion_content( - content - ) + result = LiteLLMCompletionResponsesConfig._transform_responses_api_content_to_chat_completion_content(content) assert isinstance(result, list) assert len(result) == 1 assert result[0]["cache_control"] == {"type": "ephemeral"} @@ -4522,9 +4355,7 @@ class TestCacheControlPreservation: def test_content_without_cache_control_unaffected(self): """Content blocks that don't have cache_control should be unaffected.""" content = [{"type": "text", "text": "hello"}] - result = LiteLLMCompletionResponsesConfig._transform_responses_api_content_to_chat_completion_content( - content - ) + result = LiteLLMCompletionResponsesConfig._transform_responses_api_content_to_chat_completion_content(content) assert isinstance(result, list) assert len(result) == 1 assert "cache_control" not in result[0] @@ -4546,9 +4377,7 @@ class TestCacheControlPreservation: ) assert len(messages) == 1 msg_content = ( - messages[0].get("content") - if isinstance(messages[0], dict) - else getattr(messages[0], "content", None) + messages[0].get("content") if isinstance(messages[0], dict) else getattr(messages[0], "content", None) ) assert isinstance(msg_content, list) assert msg_content[0]["cache_control"] == {"type": "ephemeral"} @@ -4561,9 +4390,7 @@ class TestCacheControlPreservation: "cache_control": {"type": "ephemeral"}, } ] - result = LiteLLMCompletionResponsesConfig._transform_responses_api_content_to_chat_completion_content( - content - ) + result = LiteLLMCompletionResponsesConfig._transform_responses_api_content_to_chat_completion_content(content) assert isinstance(result, list) assert len(result) == 1 assert result[0]["cache_control"] == {"type": "ephemeral"} @@ -4576,9 +4403,7 @@ class TestCacheControlPreservation: "cache_control": {"type": "ephemeral"}, } ] - result = LiteLLMCompletionResponsesConfig._transform_responses_api_content_to_chat_completion_content( - content - ) + result = LiteLLMCompletionResponsesConfig._transform_responses_api_content_to_chat_completion_content(content) assert isinstance(result, list) assert len(result) == 1 assert result[0]["cache_control"] == {"type": "ephemeral"} @@ -4594,18 +4419,12 @@ def test_function_call_tool_id_falls_back_to_unique_id_for_degenerate_call_id(): for the bedrock-mantle gpt-5.5 non-streaming path.""" from types import SimpleNamespace - convert = ( - LiteLLMCompletionResponsesConfig.convert_response_function_tool_call_to_chat_completion_tool_call - ) + convert = LiteLLMCompletionResponsesConfig.convert_response_function_tool_call_to_chat_completion_tool_call - mantle = SimpleNamespace( - id="fc_unique_abc123", call_id="call_0", name="get_weather", arguments="{}" - ) + mantle = SimpleNamespace(id="fc_unique_abc123", call_id="call_0", name="get_weather", arguments="{}") assert convert(mantle)["id"] == "fc_unique_abc123" - openai = SimpleNamespace( - id="fc_2", call_id="call_tokyo", name="get_weather", arguments="{}" - ) + openai = SimpleNamespace(id="fc_2", call_id="call_tokyo", name="get_weather", arguments="{}") assert convert(openai)["id"] == "call_tokyo" @@ -4670,7 +4489,7 @@ class TestHostedWebSearchReplay: text_parts: Final = tuple(block["text"] for block in content if block.get("type") == "text") assert text_parts[0] == "I found a forecast source." replayed_searches: Final = tuple( - ResponseFunctionWebSearch.model_validate_json(text[text.index("{"):]) + ResponseFunctionWebSearch.model_validate_json(text[text.index("{") :]) for text in text_parts if "web_search_call" in text ) @@ -4705,8 +4524,10 @@ class TestHostedWebSearchReplay: input_item: Final = search.model_dump(exclude_none=True) original: Final = deepcopy(input_item) - messages: Final = LiteLLMCompletionResponsesConfig._transform_responses_api_input_item_to_chat_completion_message( - input_item=input_item + messages: Final = ( + LiteLLMCompletionResponsesConfig._transform_responses_api_input_item_to_chat_completion_message( + input_item=input_item + ) ) assert len(messages) == 1 @@ -4714,7 +4535,7 @@ class TestHostedWebSearchReplay: assert not messages[0].get("tool_calls") content: Final = messages[0].get("content") assert isinstance(content, str) - replayed: Final = ResponseFunctionWebSearch.model_validate_json(content[content.index("{"):]) + replayed: Final = ResponseFunctionWebSearch.model_validate_json(content[content.index("{") :]) assert replayed == search assert input_item == original @@ -4787,7 +4608,7 @@ class TestHostedWebSearchReplay: assert text_parts[0] == "I will check the forecast." assert text_parts[-1] == "I found two sources." replayed_searches: Final = tuple( - ResponseFunctionWebSearch.model_validate_json(text[text.index("{"):]) + ResponseFunctionWebSearch.model_validate_json(text[text.index("{") :]) for text in text_parts if "web_search_call" in text ) @@ -4802,13 +4623,16 @@ class TestHostedWebSearchReplay: assistant_blocks: Final = provider_messages[1]["content"] result_blocks: Final = provider_messages[2]["content"] assert [block["id"] for block in assistant_blocks if block.get("type") == "tool_use"] == [ - "call_weather", "call_time" + "call_weather", + "call_time", ] assert [block["tool_use_id"] for block in result_blocks if block.get("type") == "tool_result"] == [ - "call_weather", "call_time" + "call_weather", + "call_time", ] assert [block["content"] for block in result_blocks if block.get("type") == "tool_result"] == [ - "Paris is sunny.", "12:00" + "Paris is sunny.", + "12:00", ] assert [block["text"] for block in assistant_blocks if block.get("type") == "text"] == list(text_parts) assert history == original @@ -5046,9 +4870,7 @@ class TestStreamingSnapshotItemIds: def test_incremental_item_id_uses_msg_prefix(self): iterator = self._make_iterator() - event = iterator._transform_chat_completion_chunk_to_response_api_chunk( - self._make_chunk("apple") - ) + event = iterator._transform_chat_completion_chunk_to_response_api_chunk(self._make_chunk("apple")) assert event is not None assert event.item_id.startswith("msg_") @@ -5057,14 +4879,10 @@ class TestStreamingSnapshotItemIds: def test_completed_snapshot_reuses_streamed_item_id(self): iterator = self._make_iterator() - streamed_event = iterator._transform_chat_completion_chunk_to_response_api_chunk( - self._make_chunk("apple") - ) + streamed_event = iterator._transform_chat_completion_chunk_to_response_api_chunk(self._make_chunk("apple")) assert streamed_event is not None - completed_event = iterator._emit_terminal_response_event( - _bridged_chat_completion_response() - ) + completed_event = iterator._emit_terminal_response_event(_bridged_chat_completion_response()) assert completed_event is not None message_items = _bridged_output_items(completed_event.response, "message") @@ -5073,13 +4891,9 @@ class TestStreamingSnapshotItemIds: def test_completed_snapshot_item_id_is_replayable(self): iterator = self._make_iterator() - iterator._transform_chat_completion_chunk_to_response_api_chunk( - self._make_chunk("apple") - ) + iterator._transform_chat_completion_chunk_to_response_api_chunk(self._make_chunk("apple")) - completed_event = iterator._emit_terminal_response_event( - _bridged_chat_completion_response() - ) + completed_event = iterator._emit_terminal_response_event(_bridged_chat_completion_response()) assert completed_event is not None for item in _bridged_output_items(completed_event.response, "message"): @@ -5106,18 +4920,14 @@ class TestStreamingSnapshotItemIds: def _reasoning_chat_completion_response(self): message = Message(role="assistant", content="apple") message.reasoning_content = "thinking about fruit" - return _bridged_chat_completion_response( - choices=[Choices(index=0, finish_reason="stop", message=message)] - ) + return _bridged_chat_completion_response(choices=[Choices(index=0, finish_reason="stop", message=message)]) def test_reasoning_delta_events_share_one_item_id(self): """The old rs_{hash(text)} ID changed with every delta, so a client accumulating reasoning by item ID saw a new item per chunk.""" iterator = self._make_iterator() - first = iterator._transform_chat_completion_chunk_to_response_api_chunk( - self._make_reasoning_chunk("thinking ") - ) + first = iterator._transform_chat_completion_chunk_to_response_api_chunk(self._make_reasoning_chunk("thinking ")) second = iterator._transform_chat_completion_chunk_to_response_api_chunk( self._make_reasoning_chunk("about fruit") ) @@ -5134,9 +4944,7 @@ class TestStreamingSnapshotItemIds: ) assert streamed_event is not None - completed_event = iterator._emit_terminal_response_event( - self._reasoning_chat_completion_response() - ) + completed_event = iterator._emit_terminal_response_event(self._reasoning_chat_completion_response()) assert completed_event is not None reasoning_items = _bridged_output_items(completed_event.response, "reasoning") @@ -5352,3 +5160,34 @@ async def test_bridge_rejects_untranslatable_tool_choice_with_a_400(stream: bool ) assert exc_info.value.status_code == 400 assert "tool_choice={'type': 'file_search'}" in str(exc_info.value) + + +@pytest.mark.parametrize("index", [0, 1, 2]) +def test_convert_response_function_tool_call_keeps_the_index_it_is_given(index: int) -> None: + item: Final = ResponseFunctionToolCall( + id="fc_x", + type="function_call", + status="completed", + arguments='{"city":"Paris"}', + call_id="call_x", + name="get_weather", + ) + + tool_call: Final = ( + LiteLLMCompletionResponsesConfig.convert_response_function_tool_call_to_chat_completion_tool_call( + tool_call_item=item, index=index + ) + ) + + assert tool_call["index"] == index + + +def test_transform_request_reads_service_tier_from_responses_api_request() -> None: + result: Final = LiteLLMCompletionResponsesConfig.transform_responses_api_request_to_chat_completion_request( + model="databricks/dbrx", + input="hello", + responses_api_request={"service_tier": "priority"}, + custom_llm_provider="databricks", + ) + + assert result["service_tier"] == "priority" diff --git a/tests/unit/test_check_test_quality.py b/tests/unit/test_check_test_quality.py index bf05775d09d..a0b46cdcdaa 100644 --- a/tests/unit/test_check_test_quality.py +++ b/tests/unit/test_check_test_quality.py @@ -505,6 +505,7 @@ def test_a_named_iterable_that_is_not_a_module_constant_is_skipped_quietly(tmp_p def _written(tmp_path, source, name="conftest.py"): path = tmp_path / name + path.parent.mkdir(parents=True, exist_ok=True) path.write_text(source, encoding="utf-8") return path @@ -637,6 +638,8 @@ _VIOLATING_SNIPPETS: Final = MappingProxyType( "test_snippet.py", 'import subprocess, sys\nsubprocess.run([sys.executable, "-c", "pass"])\n', ), + "TQ010": ("tests/integration/chaos/test_snippet.py", "import os, signal\n\nos.kill(1, signal.SIGKILL)\n"), + "TQ011": ("tests/integration/chaos/test_snippet.py", "def test_fast():\n elapsed = run()\n assert elapsed < 2\n"), } ) @@ -675,3 +678,134 @@ def test_non_interpreter_subprocess_call_is_untouched(tmp_path): def test_popen_sys_executable_tuple_is_flagged(tmp_path): source = 'import subprocess, sys\nsubprocess.Popen((sys.executable, "script.py"))\n' assert _codes(tmp_path, source) == ["TQ009"] + + +_CHAOS_TEST: Final = "tests/integration/observability/test_snippet.py" + + +def _chaos_codes(tmp_path, source, name=_CHAOS_TEST): + return [v.code for v in checker.check_file(_written(tmp_path, source, name))] + + +@pytest.mark.parametrize( + "call", + ( + "os.kill(pid, signal.SIGKILL)", + "os.killpg(pid, signal.SIGTERM)", + "os.kill(pid, 0)", + "worker.send_signal(signal.SIGTERM)", + "signal_group(pid, signal.SIGKILL)", + "redis.signal(SIGSTOP)", + "redis.signal(signal.SIGCONT)", + "worker.kill()", + "worker.terminate()", + "worker.suspend()", + "worker.resume()", + ), +) +def test_an_integration_test_that_signals_a_process_itself_is_flagged_once_per_call(tmp_path, call): + assert _chaos_codes(tmp_path, f"import os, signal\nfrom signal import SIGSTOP\n\n{call}\n") == ["TQ010"] + + +@pytest.mark.parametrize( + "source", + ( + "assert record['exit_before_stop'] == -signal.SIGKILL\n", + "signal.signal(signal.SIGTERM, handler)\n", + "task.cancel()\n", + "slot.kill(reason)\n", + "owned.kill_worker(pid)\n", + "stop_child(process)\n", + ), +) +def test_reading_a_signal_number_or_calling_a_harness_helper_is_not_signalling(tmp_path, source): + assert _chaos_codes(tmp_path, f"import signal\n\n{source}") == [] + + +@pytest.mark.parametrize( + "name", + ( + "tests/integration/_support/process.py", + "tests/unit/test_snippet.py", + "tests/e2e/integration_like/test_snippet.py", + ), +) +def test_the_signal_rule_only_applies_to_integration_tests_outside_the_harness(tmp_path, name): + assert _chaos_codes(tmp_path, "import os, signal\n\nos.kill(1, signal.SIGKILL)\nworker.suspend()\n", name) == [] + + +def test_the_signal_rule_is_suppressible_with_a_reason(tmp_path): + source = "import os, signal\n\nos.kill(1, signal.SIGKILL) # test-quality-ok: reaps a leaked fake helper\n" + assert _chaos_codes(tmp_path, source) == [] + + +@pytest.mark.parametrize( + "assertion", + ( + "assert elapsed < 2", + "assert 2 > elapsed", + "assert elapsed <= 2.5", + "assert time.monotonic() - started < 5", + "assert all(duration < 2 for _, duration in results)", + "assert max(latencies) < 1", + "assert response.elapsed.total_seconds() < 5", + "assert outcome.elapsed < BUDGET_SECONDS", + "assert elapsed < BUDGET_SECONDS - 1", + "assert 59 <= elapsed <= 75", + "assert taken < 3", + ), +) +def test_an_elapsed_time_capped_by_a_number_is_flagged(tmp_path, assertion): + source = ( + "import time\n\nBUDGET_SECONDS = 10\n\n\n" + "def test_fast():\n" + " started = time.monotonic()\n" + " taken = time.perf_counter() - started\n" + f" {assertion}\n" + ) + assert _chaos_codes(tmp_path, source) == ["TQ011"] + + +@pytest.mark.parametrize( + "statement", + ( + "assert elapsed >= 2", + "assert 2 <= elapsed", + "assert elapsed < 25 * control_elapsed", + "assert elapsed < control + injected_delay", + "assert_enforced_within(elapsed, configured=CAP_SECONDS, slack=15, not_before=CAP_SECONDS - 1)", + "assert_cut_short(time.monotonic() - started, injected_delay=30)", + "assert after - baseline < BUDGET_SECONDS", + "assert len(durations) == 3", + "assert count < 2", + "while time.monotonic() - started < 30:\n poll()", + ), +) +def test_lower_bounds_relative_bounds_counts_and_polling_deadlines_are_left_alone(tmp_path, statement): + source = ( + "import time\n\nBUDGET_SECONDS = 10\n\n\n" + "def _relative():\n" + " started = time.monotonic()\n" + " baseline = settled_reads(read(), time.monotonic())\n" + f" {statement}\n" + ) + assert _chaos_codes(tmp_path, source) == [] + + +def test_a_clock_bound_name_is_only_an_elapsed_time_in_the_function_that_binds_it(tmp_path): + source = ( + "import time\n\n\n" + "def _measure():\n" + " taken = time.monotonic() - 1\n" + " return taken\n\n\n" + "def test_counts():\n" + " taken = len(rows())\n" + " assert taken < 3\n" + ) + assert _chaos_codes(tmp_path, source) == [] + + +def test_the_wall_clock_rule_only_applies_to_integration_tests_outside_the_harness(tmp_path): + source = "def test_fast():\n elapsed = run()\n assert elapsed < 2\n" + assert _chaos_codes(tmp_path, source, "tests/unit/test_snippet.py") == [] + assert _chaos_codes(tmp_path, source, "tests/integration/_support/timing.py") == [] diff --git a/tests/unit/test_circleci_path_filter.py b/tests/unit/test_circleci_path_filter.py index b432f358f85..0856115e39b 100644 --- a/tests/unit/test_circleci_path_filter.py +++ b/tests/unit/test_circleci_path_filter.py @@ -124,8 +124,10 @@ def test_classify_decisions(category: str, changed: list[str], expected: str) -> ("litellm/caching/redis_cache.py", "run"), ("tests/unit/caching/test_redis_cluster_cache.py", "run"), ("tests/integration/sdk/test_redis_cluster_iam_auth.py", "run"), - (".circleci/config.yml", "run"), + (".github/workflows/test-unit.yml", "run"), + ("tests/integration/run.py", "run"), ("uv.lock", "run"), + (".circleci/config.yml", "skip"), ("litellm/router.py", "skip"), ("docs/my-website/redis.md", "skip"), ), diff --git a/tests/unit/test_cost_calculator.py b/tests/unit/test_cost_calculator.py index 8dbef772472..5d5d53d793e 100644 --- a/tests/unit/test_cost_calculator.py +++ b/tests/unit/test_cost_calculator.py @@ -186,6 +186,47 @@ def test_cost_calculator_with_response_cost_in_additional_headers(): assert result == 1000 +@pytest.mark.parametrize( + ("input_cost_per_token", "output_cost_per_token", "expected_cost"), + ((0.0, 0.0, 0.0), (0.001, 0.002, 10 * 0.001 + 253 * 0.002)), +) +def test_response_cost_calculator_custom_pricing_ignores_upstream_response_cost_header( + monkeypatch: pytest.MonkeyPatch, + input_cost_per_token: float, + output_cost_per_token: float, + expected_cost: float, +) -> None: + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url="")) + litellm.register_model( + model_cost={ + "cc-custom-priced-chat": { + "litellm_provider": "openai", + "mode": "chat", + "input_cost_per_token": input_cost_per_token, + "output_cost_per_token": output_cost_per_token, + } + } + ) + response: Final = ModelResponse( + model="cc-custom-priced-chat", + choices=[Choices(message=Message(content="ok"))], + usage=Usage(prompt_tokens=10, completion_tokens=253, total_tokens=263), + ) + response._hidden_params = {"additional_headers": {"llm_provider-x-litellm-response-cost": 0.000413}} + + cost: Final = response_cost_calculator( + response_object=response, + model="cc-custom-priced-chat", + custom_llm_provider="openai", + call_type="completion", + optional_params={}, + custom_pricing=True, + ) + + assert cost == pytest.approx(expected_cost) + + def test_response_cost_calculator_keeps_optional_params_out_of_hidden_params(): class MockResponse(BaseModel): pass @@ -1059,9 +1100,9 @@ def test_per_query_priced_rerank_deployment_completion_cost_is_nonzero(): router: Final = Router( model_list=[ { - "model_name": "semantic-ranker-default-004", + "model_name": "semantic-ranker-custom-test", "litellm_params": { - "model": "vertex_ai/semantic-ranker-default-004", + "model": "vertex_ai/semantic-ranker-custom-test", "vertex_project": "test-project", "vertex_location": "us-east5", }, @@ -1070,7 +1111,7 @@ def test_per_query_priced_rerank_deployment_completion_cost_is_nonzero(): ] ) router_model_id: Final = router.model_list[0]["model_info"]["id"] - assert litellm.model_cost["vertex_ai/semantic-ranker-default-004"].get("input_cost_per_query") is None + assert litellm.model_cost.get("vertex_ai/semantic-ranker-custom-test", {}).get("input_cost_per_query") is None response: Final = RerankResponse( id="vertex_ai_rerank_test", @@ -1080,7 +1121,7 @@ def test_per_query_priced_rerank_deployment_completion_cost_is_nonzero(): cost: Final = completion_cost( completion_response=response, - model="vertex_ai/semantic-ranker-default-004", + model="vertex_ai/semantic-ranker-custom-test", custom_llm_provider="vertex_ai", call_type="arerank", custom_pricing=True, @@ -1090,6 +1131,39 @@ def test_per_query_priced_rerank_deployment_completion_cost_is_nonzero(): assert cost == pytest.approx(3 * 0.001) +@pytest.mark.parametrize( + "model", + ( + "vertex_ai/semantic-ranker-default", + "vertex_ai/semantic-ranker-default@latest", + "vertex_ai/semantic-ranker-default-004", + "vertex_ai/semantic-ranker-default-003", + "vertex_ai/semantic-ranker-default-002", + "vertex_ai/semantic-ranker-fast", + "vertex_ai/semantic-ranker-fast@latest", + "vertex_ai/semantic-ranker-fast-004", + ), +) +def test_vertex_semantic_ranker_rerank_completion_cost_is_nonzero( + model: str, _local_model_cost_map: None +) -> None: + response: Final = RerankResponse( + id="vertex_ai_rerank_pricing_test", + results=[], + meta={"billed_units": {"search_units": 3}}, + ) + + cost: Final = completion_cost( + completion_response=response, + model=model, + custom_llm_provider="vertex_ai", + call_type="rerank", + ) + + assert cost > 0 + assert cost == pytest.approx(3 * litellm.model_cost[model]["input_cost_per_query"]) + + def test_azure_realtime_cost_calculator(_local_model_cost_map): cost = handle_realtime_stream_cost_calculation( @@ -1696,6 +1770,132 @@ def test_cost_discount_vertex_ai(monkeypatch): print(f" - Savings: ${cost_without_discount - cost_with_discount:.6f}") +def test_cost_discount_model_pattern_beats_bare_provider(monkeypatch): + """ + Test that a / discount key applies only to matching models + """ + from litellm import completion_cost + from litellm.types.utils import Usage + + claude_response = ModelResponse( + id="test-id", + choices=[], + created=1234567890, + model="claude-sonnet-4-5", + object="chat.completion", + usage=Usage(prompt_tokens=100, completion_tokens=50, total_tokens=150), + ) + gemini_response = ModelResponse( + id="test-id", + choices=[], + created=1234567890, + model="gemini-3-pro-preview", + object="chat.completion", + usage=Usage(prompt_tokens=100, completion_tokens=50, total_tokens=150), + ) + + monkeypatch.setattr(litellm, "cost_discount_config", {}) + claude_undiscounted = completion_cost( + completion_response=claude_response, + model="vertex_ai/claude-sonnet-4-5", + custom_llm_provider="vertex_ai", + ) + gemini_undiscounted = completion_cost( + completion_response=gemini_response, + model="vertex_ai/gemini-3-pro-preview", + custom_llm_provider="vertex_ai", + ) + + monkeypatch.setattr(litellm, "cost_discount_config", {"vertex_ai/claude-*": 0.2, "vertex_ai": 0.05}) + claude_discounted = completion_cost( + completion_response=claude_response, + model="vertex_ai/claude-sonnet-4-5", + custom_llm_provider="vertex_ai", + ) + gemini_discounted = completion_cost( + completion_response=gemini_response, + model="vertex_ai/gemini-3-pro-preview", + custom_llm_provider="vertex_ai", + ) + + assert claude_discounted == pytest.approx(claude_undiscounted * 0.8, rel=1e-9) + assert gemini_discounted == pytest.approx(gemini_undiscounted * 0.95, rel=1e-9) + + +def test_cost_discount_model_pattern_matches_custom_priced_deployment(monkeypatch): + """ + A custom-priced deployment prices against its router model id, but the discount still matches the request model + """ + from litellm import Router, completion_cost + from litellm.types.utils import Choices, Message, Usage + + router: Final = Router( + model_list=[ + { + "model_name": "qwen-3.7-plus", + "litellm_params": {"model": "dashscope/qwen3.7-plus", "api_key": "sk-fake"}, + "model_info": {"input_cost_per_token": 4e-07, "output_cost_per_token": 1.6e-06}, + } + ] + ) + router_model_id: Final = router.model_list[0]["model_info"]["id"] + response: Final = ModelResponse( + model="dashscope/qwen3.7-plus", + choices=[Choices(index=0, message=Message(role="assistant", content="hi"))], + usage=Usage(prompt_tokens=12, completion_tokens=377, total_tokens=389), + ) + + monkeypatch.setattr(litellm, "cost_discount_config", {"dashscope/qwen*": 0.2}) + discounted: Final = completion_cost( + completion_response=response, + model="dashscope/qwen3.7-plus", + custom_llm_provider="dashscope", + custom_pricing=True, + router_model_id=router_model_id, + ) + + assert discounted == pytest.approx((12 * 4e-07 + 377 * 1.6e-06) * 0.8, rel=1e-9) + + +def test_cost_discount_together_ai_pattern_matches_unrewritten_model(monkeypatch): + """ + Test that a / discount matches the request model, not the + pricing category the Together AI cost lookup rewrites it to + """ + from litellm import completion_cost + from litellm.types.utils import Usage + + response = ModelResponse( + id="test-id", + choices=[], + created=1234567890, + model="together_ai/my-org/Custom-70B-Instruct", + object="chat.completion", + usage=Usage(prompt_tokens=100, completion_tokens=50, total_tokens=150), + ) + + monkeypatch.setattr(litellm, "cost_discount_config", {}) + undiscounted = completion_cost( + completion_response=response, + model="together_ai/my-org/Custom-70B-Instruct", + custom_llm_provider="together_ai", + ) + assert undiscounted > 0 + + monkeypatch.setattr( + litellm, + "cost_discount_config", + {"together_ai": 0.05, "together_ai/my-org/Custom-*": 0.20}, + ) + discounted = completion_cost( + completion_response=response, + model="together_ai/my-org/Custom-70B-Instruct", + custom_llm_provider="together_ai", + ) + + assert discounted == pytest.approx(undiscounted * 0.8, rel=1e-9) + + def test_cost_discount_not_applied_to_other_providers(monkeypatch): """ Test that cost discount only applies to configured providers @@ -3957,6 +4157,41 @@ def test_completion_cost_region_name_prices_mantle_on_the_regional_row(_local_mo ) == pytest.approx(expected_flat) +def test_completion_cost_discount_matches_region_stripped_model(monkeypatch, _local_model_cost_map): + """A discount pattern written against the bare model name must still match when the + cost-selected name carries a // prefix from a priced regional row.""" + + response = litellm.ModelResponse( + id="x", + choices=[{"index": 0, "message": {"role": "assistant", "content": "hi"}, "finish_reason": "stop"}], + model="xai.grok-4.3", + usage={"prompt_tokens": 38, "completion_tokens": 20, "total_tokens": 58}, + ) + + monkeypatch.setattr(litellm, "cost_discount_config", {}) + undiscounted = litellm.completion_cost( + completion_response=response, + model="xai.grok-4.3", + custom_llm_provider="bedrock_mantle", + region_name="us-gov-west-1", + ) + assert undiscounted > 0 + + monkeypatch.setattr( + litellm, + "cost_discount_config", + {"bedrock_mantle": 0.05, "bedrock_mantle/xai.grok-*": 0.20}, + ) + discounted = litellm.completion_cost( + completion_response=response, + model="xai.grok-4.3", + custom_llm_provider="bedrock_mantle", + region_name="us-gov-west-1", + ) + + assert discounted == pytest.approx(undiscounted * 0.8, rel=1e-9) + + def test_cost_per_token_region_name_applies_to_provider_prefixed_model(_local_model_cost_map): """A provider-prefixed model must still find its bedrock_mantle// row instead of composing the region key with the provider segment twice.""" diff --git a/tests/unit/test_cost_map_guard.py b/tests/unit/test_cost_map_guard.py index e127aa08d29..1b9c6950ba9 100644 --- a/tests/unit/test_cost_map_guard.py +++ b/tests/unit/test_cost_map_guard.py @@ -1,8 +1,8 @@ -from collections.abc import Mapping import importlib.util import json import subprocess import sys +from collections.abc import Mapping from pathlib import Path from types import ModuleType from typing import Final @@ -84,30 +84,29 @@ def test_in_sync_files_pass_for_humans_and_bots() -> None: EVALUATION_MODEL: Final = "azure_ai/Microsoft-Decision-1" -EVALUATION_OUTPUT_PRICE_FAILURE: Final = ( - f"{EVALUATION_MODEL}: mode evaluation rows bill input only; output_cost_per_token must be 0 unless the vendor " - "prices output tokens for this model (add the source to EVALUATION_OUTPUT_PRICE_SOURCES)" +DECISIONS_MODEL: Final = "azure_ai/decisions/Microsoft-Decision-1" +INPUT_ONLY_OUTPUT_PRICE_FAILURE: Final = ( + "mode evaluation and decisions rows bill input only; output_cost_per_token must be 0 unless the vendor prices " + "output tokens for this model (add the source to EVALUATION_OUTPUT_PRICE_SOURCES)" ) +@pytest.mark.parametrize(("model", "mode"), ((EVALUATION_MODEL, "evaluation"), (DECISIONS_MODEL, "decisions"))) @pytest.mark.parametrize("bot", (True, False)) -def test_evaluation_rows_with_output_price_fail_for_humans_and_bots(bot: bool) -> None: - head: Final = _snapshot( - {**BASE_MAP, EVALUATION_MODEL: _entry(mode="evaluation", output_cost_per_token=4.2e-08)} - ) - assert _failures(head, bot=bot) == (EVALUATION_OUTPUT_PRICE_FAILURE,) +def test_evaluation_and_decisions_rows_with_output_price_fail_for_humans_and_bots( + model: str, mode: str, bot: bool +) -> None: + head: Final = _snapshot({**BASE_MAP, model: _entry(mode=mode, output_cost_per_token=4.2e-08)}) + assert _failures(head, bot=bot) == (f"{model}: {INPUT_ONLY_OUTPUT_PRICE_FAILURE}",) -def test_evaluation_rows_with_zero_or_missing_output_price_pass() -> None: - zero_output: Final = _snapshot( - {**BASE_MAP, EVALUATION_MODEL: _entry(mode="evaluation", output_cost_per_token=0.0)} - ) +@pytest.mark.parametrize(("model", "mode"), ((EVALUATION_MODEL, "evaluation"), (DECISIONS_MODEL, "decisions"))) +def test_evaluation_and_decisions_rows_with_zero_or_missing_output_price_pass(model: str, mode: str) -> None: + zero_output: Final = _snapshot({**BASE_MAP, model: _entry(mode=mode, output_cost_per_token=0.0)}) missing_output_entry: Final = { - key: value - for key, value in _entry(mode="evaluation").items() - if key != "output_cost_per_token" + key: value for key, value in _entry(mode=mode).items() if key != "output_cost_per_token" } - missing_output: Final = _snapshot({**BASE_MAP, EVALUATION_MODEL: missing_output_entry}) + missing_output: Final = _snapshot({**BASE_MAP, model: missing_output_entry}) assert _failures(zero_output) == () assert _failures(missing_output) == () @@ -120,13 +119,12 @@ def test_chat_rows_with_output_price_pass() -> None: assert _failures(head) == () -def test_evaluation_output_price_source_allowlist_passes() -> None: - head: Final = _snapshot( - {**BASE_MAP, EVALUATION_MODEL: _entry(mode="evaluation", output_cost_per_token=4.2e-08)} - ) +@pytest.mark.parametrize(("model", "mode"), ((EVALUATION_MODEL, "evaluation"), (DECISIONS_MODEL, "decisions"))) +def test_evaluation_and_decisions_output_price_source_allowlist_passes(model: str, mode: str) -> None: + head: Final = _snapshot({**BASE_MAP, model: _entry(mode=mode, output_cost_per_token=4.2e-08)}) source: Final = 'https://vendor.example/pricing "Output tokens are charged per token."' - assert _failures(head, evaluation_output_price_sources={EVALUATION_MODEL: source}) == () + assert _failures(head, evaluation_output_price_sources={model: source}) == () def test_bot_may_add_and_reprice_models() -> None: diff --git a/tests/unit/test_integration_run.py b/tests/unit/test_integration_run.py index 36612525572..4a8551e07a4 100644 --- a/tests/unit/test_integration_run.py +++ b/tests/unit/test_integration_run.py @@ -1,6 +1,8 @@ from typing import Final -from tests.integration.run import select, uncollected +import pytest + +from tests.integration.run import GROUPS, select, uncollected, worker_settings _GROUP: Final = ( "tests/integration/cost_calculation/test_cost_tracking.py", @@ -34,3 +36,13 @@ def test_a_node_id_whose_file_collected_tests_is_not_empty() -> None: def test_a_selected_file_that_collected_nothing_is_reported() -> None: assert uncollected(_GROUP, frozenset({_CELL})) == ("tests/integration/cost_calculation/test_rollups.py",) + + +@pytest.mark.parametrize("group", tuple(GROUPS)) +def test_every_group_run_by_several_workers_gives_each_worker_its_own_stack(group: str) -> None: + assert worker_settings(group, 2)["INTEGRATION_WORKER_ISOLATION"] == "1" + + +@pytest.mark.parametrize("group", tuple(GROUPS)) +def test_a_single_worker_run_keeps_the_supplied_deployment(group: str) -> None: + assert worker_settings(group, 1) == {} diff --git a/tests/unit/test_muse_spark_1_3_model_metadata.py b/tests/unit/test_muse_spark_1_3_model_metadata.py index 4392553fcc3..049fb988fd5 100644 --- a/tests/unit/test_muse_spark_1_3_model_metadata.py +++ b/tests/unit/test_muse_spark_1_3_model_metadata.py @@ -6,9 +6,11 @@ import pytest import litellm from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider from litellm.litellm_core_utils.llm_cost_calc.tool_call_cost_tracking import StandardBuiltInToolCostTracking +from litellm.router_utils.reasoning_effort_capability import resolve_supported_reasoning_efforts MUSE_SPARK_STANDARD = "meta/muse-spark-1.3" MUSE_SPARK_CONTRIBUTOR = "meta/muse-spark-1.3-contributor" +MUSE_WONDER = "meta/muse-wonder-1.0" WEB_SEARCH_COST_PER_QUERY = 0.0025 PRICING = ( @@ -22,7 +24,7 @@ def _load_cost_map(filename: str = "model_prices_and_context_window.json") -> di return json.load(f) -@pytest.mark.parametrize("model", (MUSE_SPARK_STANDARD, MUSE_SPARK_CONTRIBUTOR)) +@pytest.mark.parametrize("model", (MUSE_SPARK_STANDARD, MUSE_SPARK_CONTRIBUTOR, MUSE_WONDER)) def test_muse_spark_1_3_routes_to_meta_model_api(model: str): routed_model, provider, _, api_base = get_llm_provider(model=model, api_key="sk-test") @@ -31,7 +33,7 @@ def test_muse_spark_1_3_routes_to_meta_model_api(model: str): assert api_base == "https://api.meta.ai/v1" -@pytest.mark.parametrize("model", (MUSE_SPARK_STANDARD, MUSE_SPARK_CONTRIBUTOR)) +@pytest.mark.parametrize("model", (MUSE_SPARK_STANDARD, MUSE_SPARK_CONTRIBUTOR, MUSE_WONDER)) def test_muse_spark_1_3_backup_matches_main(model: str): """Ensure the bundled model cost map stays in sync with the canonical file.""" main_cost = _load_cost_map() @@ -40,6 +42,18 @@ def test_muse_spark_1_3_backup_matches_main(model: str): assert backup_cost.get(model) == main_cost.get(model), f"{model} differs between main and backup model cost maps" +def test_muse_wonder_reasoning_efforts_include_max_and_exclude_none(): + entry = dict(_load_cost_map()[MUSE_WONDER]) + entry["key"] = MUSE_WONDER + + resolved = resolve_supported_reasoning_efforts(entry, deployment_is_mapped=True) + + assert resolved is not None + assert "max" in resolved + assert "xhigh" in resolved + assert "none" not in resolved + + def test_muse_spark_contributor_tier_is_cheaper_than_standard(): cost_map = _load_cost_map() standard = cost_map[MUSE_SPARK_STANDARD] diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index 466f3c988e4..6ad4f75f97f 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -438,6 +438,26 @@ def test_supports_audio_output_reads_its_own_cost_map_flag( assert litellm.supports_audio_output(model) is audio_output +@pytest.mark.parametrize(("audio_input", "video_input"), [(True, False), (False, True)]) +def test_supports_video_input_reads_its_own_cost_map_flag( + local_model_cost_map: None, monkeypatch: pytest.MonkeyPatch, audio_input: bool, video_input: bool +) -> None: + model: Final = "openai/video-flags-disagree-model" + monkeypatch.setitem( + litellm.model_cost, + model, + { + "litellm_provider": "openai", + "mode": "chat", + "supports_audio_input": audio_input, + "supports_video_input": video_input, + }, + ) + assert litellm.supports_video_input(model) is video_input + assert litellm.supports_audio_input(model) is audio_input + assert litellm.get_model_info(model)["supports_video_input"] is video_input + + def test_get_optional_params_image_gen(): from litellm.llms.azure.image_generation import AzureGPTImageGenerationConfig @@ -918,6 +938,14 @@ def test_aaamodel_prices_and_context_window_json_is_valid(): "cache_creation_input_token_cost_above_200k_tokens_batches": {"type": "number"}, "cache_creation_input_token_cost_above_272k_tokens_batches": {"type": "number"}, "cache_creation_input_token_cost_batches": {"type": "number"}, + "cache_creation_input_token_cost_decisions": {"type": "number"}, + "cache_creation_input_token_cost_above_272k_tokens_decisions": {"type": "number"}, + "cache_read_input_token_cost_decisions": {"type": "number"}, + "cache_read_input_token_cost_above_272k_tokens_decisions": {"type": "number"}, + "input_cost_per_token_decisions": {"type": "number"}, + "input_cost_per_token_above_272k_tokens_decisions": {"type": "number"}, + "output_cost_per_token_decisions": {"type": "number"}, + "output_cost_per_token_above_272k_tokens_decisions": {"type": "number"}, "cache_creation_input_token_cost_flex": {"type": "number"}, "cache_creation_input_token_cost_priority": {"type": "number"}, "cache_creation_input_token_cost_ultrafast": {"type": "number"}, @@ -1397,6 +1425,25 @@ def test_get_model_info_bedrock_regional_inference_profile_pricing(local_model_c assert control["key"] == "eu.amazon.nova-pro-v1:0" +@pytest.mark.parametrize("route", ["bedrock/", "bedrock/converse/", "bedrock/invoke/", "bedrock/eu-west-1/"]) +def test_get_model_info_bedrock_inference_profile_arn_keeps_regional_row(local_model_cost_map, route): + regional = litellm.model_cost["eu.amazon.nova-pro-v1:0"] + base = litellm.model_cost["amazon.nova-pro-v1:0"] + assert regional["input_cost_per_token"] > base["input_cost_per_token"] + arn_prefix = "arn:aws:bedrock:eu-west-1:123456789012:inference-profile/" + + info = litellm.get_model_info(model=f"{route}{arn_prefix}eu.amazon.nova-pro-v1:0") + assert info["key"] == "eu.amazon.nova-pro-v1:0" + assert info["input_cost_per_token"] == regional["input_cost_per_token"] + assert info["output_cost_per_token"] == regional["output_cost_per_token"] + + bare = litellm.get_model_info(model=f"{route}{arn_prefix}amazon.nova-pro-v1:0") + assert bare["key"] == "amazon.nova-pro-v1:0" + assert bare["input_cost_per_token"] == base["input_cost_per_token"] + + + + @pytest.mark.parametrize( "bare_key", [ diff --git a/tests/unit/types/test_utils.py b/tests/unit/types/test_utils.py index 09063d894e7..32479c62372 100644 --- a/tests/unit/types/test_utils.py +++ b/tests/unit/types/test_utils.py @@ -1,8 +1,15 @@ import pytest +from litellm.litellm_core_utils.initialize_dynamic_callback_params import validate_excluded_services_value from litellm.types.utils import CAPTURE_MESSAGE_CONTENT_VALUES, CaptureMessageContent, captures_span_content +@pytest.mark.parametrize("value", ["", " ", "auth", "none,redis", "redis,", ",postgres"]) +def test_invalid_excluded_services_values_are_rejected(value: str) -> None: + with pytest.raises(ValueError, match=r"Invalid excluded_services.*none"): + validate_excluded_services_value(value) + + @pytest.mark.parametrize( "mode", [ diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/add_provider_form.integration.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/add_provider_form.integration.test.tsx index 4680a4d504a..ec19625bfc9 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/add_provider_form.integration.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/add_provider_form.integration.test.tsx @@ -4,7 +4,6 @@ import { beforeEach, describe, expect, it, vi } from "vitest"; import { renderWithProviders } from "../../../../../tests/test-utils"; import AddProviderForm from "./add_provider_form"; -import { DiscountConfig } from "./types"; const onAddProvider = vi.fn(); const onParentFinish = vi.fn(); @@ -18,11 +17,12 @@ const ParentOwnedForm = () => ( className="space-y-6" > diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/add_provider_form.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/add_provider_form.test.tsx index 94213f923c9..e82c61adcd8 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/add_provider_form.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/add_provider_form.test.tsx @@ -4,15 +4,15 @@ import { screen } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; import { renderWithProviders } from "../../../../../tests/test-utils"; import AddProviderForm from "./add_provider_form"; -import { DiscountConfig } from "./types"; import { Providers, providerLogoMap } from "@/components/provider_info_helpers"; const DEFAULT_PROPS = { - discountConfig: {} as DiscountConfig, selectedProvider: undefined, newDiscount: "", + modelPattern: "", onProviderChange: vi.fn(), onDiscountChange: vi.fn(), + onModelPatternChange: vi.fn(), onAddProvider: vi.fn(), }; @@ -71,6 +71,20 @@ describe("AddProviderForm", () => { expect(onDiscountChange).toHaveBeenCalledWith("15"); }); + it("should render the model pattern input field", () => { + renderWithProviders(); + expect(screen.getByPlaceholderText("claude-*")).toBeInTheDocument(); + }); + + it("should report the edited model pattern as the user types", async () => { + const onModelPatternChange = vi.fn(); + const user = userEvent.setup(); + renderWithProviders(); + + await user.type(screen.getByPlaceholderText("claude-*"), "g"); + expect(onModelPatternChange).toHaveBeenCalledWith("g"); + }); + it("should show the percent sign next to the discount input", () => { renderWithProviders(); expect(screen.getByText("%")).toBeInTheDocument(); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/add_provider_form.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/add_provider_form.tsx index 05748652180..8f3813e1705 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/add_provider_form.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/add_provider_form.tsx @@ -2,7 +2,7 @@ import React from "react"; import { CircleHelp } from "lucide-react"; import { Logo } from "@/components/molecules/logo/Logo"; -import { Providers, provider_map } from "@/components/provider_info_helpers"; +import { Providers } from "@/components/provider_info_helpers"; import { Field, FieldGroup, FieldLabel } from "@/components/ui/field"; import { Button } from "@/components/ui/button"; import { @@ -16,7 +16,6 @@ import { import { Input } from "@/components/ui/input"; import { InputGroupAddon } from "@/components/ui/input-group"; import { Tooltip, TooltipContent, TooltipProvider, TooltipTrigger } from "@/components/ui/tooltip"; -import { DiscountConfig } from "./types"; interface ProviderOption { value: string; @@ -24,24 +23,20 @@ interface ProviderOption { } interface AddProviderFormProps { - discountConfig: DiscountConfig; selectedProvider: string | undefined; newDiscount: string; + modelPattern: string; onProviderChange: (provider: string | undefined) => void; onDiscountChange: (discount: string) => void; + onModelPatternChange: (value: string) => void; onAddProvider: () => void; } const PROVIDER_FIELD_ID = "add-provider-discount-provider"; +const MODEL_PATTERN_FIELD_ID = "add-provider-discount-model-pattern"; const DISCOUNT_FIELD_ID = "add-provider-discount-percentage"; -const providerOptionsWithoutDiscount = (discountConfig: DiscountConfig): ProviderOption[] => - Object.entries(Providers) - .filter(([providerEnum]) => { - const providerValue = provider_map[providerEnum as keyof typeof provider_map]; - return !(providerValue && discountConfig[providerValue]); - }) - .map(([value, label]) => ({ value, label })); +const providerOptions = (): ProviderOption[] => Object.entries(Providers).map(([value, label]) => ({ value, label })); const selectedProviderOption = (selectedProvider: string | undefined): ProviderOption | null => { if (!selectedProvider) { @@ -62,14 +57,15 @@ const labelWithHint = (label: string, hint: string): React.ReactNode => ( ); const AddProviderForm: React.FC = ({ - discountConfig, selectedProvider, newDiscount, + modelPattern, onProviderChange, onDiscountChange, + onModelPatternChange, onAddProvider, }) => { - const options = providerOptionsWithoutDiscount(discountConfig); + const options = providerOptions(); const selectedOption = selectedProviderOption(selectedProvider); return ( @@ -110,6 +106,22 @@ const AddProviderForm: React.FC = ({ + + + {labelWithHint( + "Model pattern", + "Leave empty to discount every model from this provider. Use a glob such as claude-* or gemini-2* to discount only matching models", + )} + + onModelPatternChange(event.target.value)} + className="rounded-lg" + /> + + {labelWithHint("Discount Percentage", "Enter a percentage value (e.g., 5 for 5% discount)")} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/cost_tracking_settings.integration.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/cost_tracking_settings.integration.test.tsx index 3de88fb6f57..75785d9ab87 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/cost_tracking_settings.integration.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/cost_tracking_settings.integration.test.tsx @@ -92,7 +92,21 @@ describe("CostTrackingSettings submit paths", () => { await waitFor(() => expect(stableDiscountCallbacks.handleAddProvider).toHaveBeenCalled()); expect(stableDiscountCallbacks.handleAddProvider).toHaveBeenCalledTimes(1); - expect(stableDiscountCallbacks.handleAddProvider).toHaveBeenCalledWith("OpenAI", "5"); + expect(stableDiscountCallbacks.handleAddProvider).toHaveBeenCalledWith("OpenAI", "5", ""); + }); + + it("passes the model pattern through to the add handler", async () => { + const user = userEvent.setup(); + await openDiscountModal(user); + + await user.click(screen.getAllByRole("combobox")[0]); + await user.click((await screen.findAllByRole("option"))[0]); + fireEvent.change(screen.getByLabelText(/Model pattern/i), { target: { value: "claude-*" } }); + fireEvent.change(screen.getByLabelText(/Discount Percentage/i), { target: { value: "20" } }); + await user.click(submitDiscount()); + + await waitFor(() => expect(stableDiscountCallbacks.handleAddProvider).toHaveBeenCalled()); + expect(stableDiscountCallbacks.handleAddProvider).toHaveBeenCalledWith("OpenAI", "20", "claude-*"); }); it("requests the margin exactly once per click", async () => { @@ -127,7 +141,7 @@ describe("CostTrackingSettings submit paths", () => { await waitFor(() => expect(stableDiscountCallbacks.handleAddProvider).toHaveBeenCalled()); expect(stableDiscountCallbacks.handleAddProvider).toHaveBeenCalledTimes(1); - expect(stableDiscountCallbacks.handleAddProvider).toHaveBeenCalledWith("OpenAI", "5"); + expect(stableDiscountCallbacks.handleAddProvider).toHaveBeenCalledWith("OpenAI", "5", ""); }); it("requests the margin exactly once when Enter is pressed in the percentage field", async () => { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/cost_tracking_settings.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/cost_tracking_settings.tsx index b8e51939dd0..717e1b53d13 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/cost_tracking_settings.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/cost_tracking_settings.tsx @@ -58,6 +58,7 @@ const SectionHeader: React.FC<{ title: string; description: string }> = ({ title const CostTrackingSettings: React.FC = ({ userID, userRole, accessToken }) => { const [selectedProvider, setSelectedProvider] = useState(undefined); const [newDiscount, setNewDiscount] = useState(""); + const [modelPattern, setModelPattern] = useState(""); const [isFetching, setIsFetching] = useState(true); const [isModalVisible, setIsModalVisible] = useState(false); const [isMarginModalVisible, setIsMarginModalVisible] = useState(false); @@ -115,10 +116,11 @@ const CostTrackingSettings: React.FC = ({ userID, use }, [accessToken, fetchDiscountConfig, fetchMarginConfig, fetchBlockUnpriced]); const handleAddProvider = async () => { - const success = await addProvider(selectedProvider, newDiscount); + const success = await addProvider(selectedProvider, newDiscount, modelPattern); if (success) { setSelectedProvider(undefined); setNewDiscount(""); + setModelPattern(""); setIsModalVisible(false); } }; @@ -127,6 +129,7 @@ const CostTrackingSettings: React.FC = ({ userID, use setIsModalVisible(false); setSelectedProvider(undefined); setNewDiscount(""); + setModelPattern(""); }; const handleRemoveProvider = (provider: string, providerDisplayName: string) => { @@ -386,11 +389,12 @@ const CostTrackingSettings: React.FC = ({ userID, use

event.preventDefault()} className="space-y-6"> diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/provider_discount_table.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/provider_discount_table.test.tsx index 8d1ee100a2a..6713351face 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/provider_discount_table.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/provider_discount_table.test.tsx @@ -295,6 +295,78 @@ describe("ProviderDiscountTable", () => { expect(screen.getByRole("button", { name: "Cancel editing discount for OpenAI" })).toBeInTheDocument(); }); + it("should render a Models column", () => { + renderWithProviders( + , + ); + expect(screen.getByRole("columnheader", { name: "Models" })).toBeInTheDocument(); + }); + + it("should show the model pattern for a provider/pattern key and All models for a bare key", () => { + renderWithProviders( + , + ); + expect(screen.getByText("claude-*")).toBeInTheDocument(); + expect(screen.getByText("All models")).toBeInTheDocument(); + }); + + it("should split a nested pattern key on the first slash only", () => { + renderWithProviders( + , + ); + expect(screen.getByText("anthropic/claude-*")).toBeInTheDocument(); + expect(screen.queryByText("All models")).not.toBeInTheDocument(); + }); + + it("should include the pattern in the row action labels for a pattern key", async () => { + const user = userEvent.setup(); + renderWithProviders( + , + ); + + await user.click(screen.getByRole("button", { name: /Remove discount for Vertex AI .* \(claude-\*\)/ })); + + expect(onRemoveProvider).toHaveBeenCalledWith( + "vertex_ai/claude-*", + "Vertex AI (Anthropic, Gemini, etc.) (claude-*)", + ); + }); + + it("should call onDiscountChange with the full config key when editing a pattern row", async () => { + const user = userEvent.setup(); + renderWithProviders( + , + ); + + await user.click(screen.getByRole("button", { name: /Edit discount for Vertex AI .* \(claude-\*\)/ })); + const input = screen.getByPlaceholderText("5"); + await user.clear(input); + fireEvent.change(input, { target: { value: "30" } }); + await user.click(screen.getByRole("button", { name: /Save discount for Vertex AI .* \(claude-\*\)/ })); + + expect(onDiscountChange).toHaveBeenCalledWith("vertex_ai/claude-*", "0.3"); + }); + it("should render the empty message when no discounts are configured", () => { renderWithProviders( { + const slashIndex = key.indexOf("/"); + if (slashIndex < 0) { + return { provider: key, modelPattern: null }; + } + return { provider: key.slice(0, slashIndex), modelPattern: key.slice(slashIndex + 1) }; +}; + const ProviderDiscountTable: React.FC = ({ discountConfig, onDiscountChange, onRemoveProvider, }) => { - const [editingProvider, setEditingProvider] = useState(null); + const [editingKey, setEditingKey] = useState(null); const [editValue, setEditValue] = useState(""); - const handleStartEdit = (provider: string, currentDiscount: number) => { - setEditingProvider(provider); + const handleStartEdit = (key: string, currentDiscount: number) => { + setEditingKey(key); setEditValue((currentDiscount * 100).toString()); }; - const handleSaveEdit = (provider: string) => { + const handleSaveEdit = (key: string) => { const percentValue = parseFloat(editValue); if (!isNaN(percentValue) && percentValue >= 0 && percentValue <= 100) { - onDiscountChange(provider, (percentValue / 100).toString()); + onDiscountChange(key, (percentValue / 100).toString()); } - setEditingProvider(null); + setEditingKey(null); setEditValue(""); }; const handleCancelEdit = () => { - setEditingProvider(null); + setEditingKey(null); setEditValue(""); }; - const handleKeyDown = (e: React.KeyboardEvent, provider: string) => { + const handleKeyDown = (e: React.KeyboardEvent, key: string) => { if (e.key === "Enter") { - handleSaveEdit(provider); + handleSaveEdit(key); } else if (e.key === "Escape") { handleCancelEdit(); } }; + const rowLabel = (row: ProviderDiscountRow): string => { + const { displayName } = getProviderLogoAndName(row.provider); + return row.modelPattern ? `${displayName} (${row.modelPattern})` : displayName; + }; + // Convert discount config to array and sort const data: ProviderDiscountRow[] = Object.entries(discountConfig) - .map(([provider, discount]) => ({ provider, discount })) + .map(([key, discount]) => ({ key, ...splitDiscountKey(key), discount })) .sort((a, b) => { const displayA = getProviderLogoAndName(a.provider).displayName; const displayB = getProviderLogoAndName(b.provider).displayName; - return displayA.localeCompare(displayB); + return displayA.localeCompare(displayB) || (a.modelPattern ?? "").localeCompare(b.modelPattern ?? ""); }); return ( @@ -78,19 +93,28 @@ const ProviderDiscountTable: React.FC = ({ ); }, }, + { + header: "Models", + cell: (row) => + row.modelPattern ? ( + {row.modelPattern} + ) : ( + All models + ), + }, { header: "Discount Percentage", numeric: true, cell: (row) => { - const { displayName } = getProviderLogoAndName(row.provider); + const label = rowLabel(row); return (
- {editingProvider === row.provider ? ( + {editingKey === row.key ? ( <> setEditValue(e.target.value)} - onKeyDown={(e) => handleKeyDown(e, row.provider)} + onKeyDown={(e) => handleKeyDown(e, row.key)} placeholder="5" className="w-20" autoFocus @@ -99,8 +123,8 @@ const ProviderDiscountTable: React.FC = ({ ); } +export function QuestionCardHeader({ + index, + removeLabel, + onRemove, +}: { + index: number; + removeLabel: string; + onRemove: () => void; +}) { + return ( +
+ Question {index + 1} + +
+ ); +} + export function FormFrame({ validation, children }: { validation: PayloadValidation; children: ReactNode }) { const issuesId = useId(); return ( -
-
- Request form - -
-
{children}
+
+
{children}
); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/systemOneUI/OpenAIDecisionsForm.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/systemOneUI/OpenAIDecisionsForm.tsx index e7a1e14ca02..3fa6a75959f 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/systemOneUI/OpenAIDecisionsForm.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/systemOneUI/OpenAIDecisionsForm.tsx @@ -6,7 +6,7 @@ import { Label } from "@/components/ui/label"; import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select"; import { Textarea } from "@/components/ui/textarea"; import { useId, useMemo } from "react"; -import { AddButton, FormFrame, RemoveButton } from "./DecisionFormControls"; +import { AddButton, FormFrame, QuestionCardHeader, RemoveButton } from "./DecisionFormControls"; import { freeName, removeAt, replaceAt } from "./lib/formPayload"; import { readDecisionsForm, @@ -124,38 +124,52 @@ interface QuestionCardProps extends QuestionEditorProps { onRemove: () => void; } +const blankDecisionsQuestion = (existing: DecisionsFormQuestion[]): DecisionsFormQuestion => ({ + type: "predicate", + name: freeName( + existing.map((q) => q.name ?? ""), + "question", + ), + instructions: "", +}); + function QuestionCard({ question, index, onChange, onRemove }: QuestionCardProps) { const instructionsId = useId(); return ( -
-
- onChange({ ...question, name: event.target.value })} - /> - - +
+ +
+
+ Name + onChange({ ...question, name: event.target.value })} + /> +
+
+ Answer type + +
@@ -229,21 +243,7 @@ export default function OpenAIDecisionsForm({ value, onChange, validation, onRes onRemove={() => writeQuestions(removeAt(questions, index))} /> ))} - - writeQuestions([ - ...questions, - { - type: "predicate", - name: freeName( - questions.map((q) => q.name ?? ""), - "question", - ), - instructions: "", - }, - ]) - } - > + writeQuestions([...questions, blankDecisionsQuestion(questions)])}> Add question
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/systemOneUI/ResponseView.integration.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/systemOneUI/ResponseView.integration.test.tsx index 141fdc43374..f6c5794fa5c 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/systemOneUI/ResponseView.integration.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/systemOneUI/ResponseView.integration.test.tsx @@ -106,7 +106,7 @@ describe("ResponseView integration", () => { expect(screen.getByText("99% yes")).toBeInTheDocument(); expect(screen.getByRole("meter", { name: "Yes probability" })).toHaveAttribute("aria-valuenow", "99"); expect(screen.getByRole("meter", { name: "2: angry probability" })).toBeInTheDocument(); - expect(screen.getByText("378 input / 63 output tokens")).toBeInTheDocument(); + expect(screen.getByText("378 in / 63 out")).toBeInTheDocument(); }); it("selects the score level nearest a fractional calibrated score", () => { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/systemOneUI/ResponseView.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/systemOneUI/ResponseView.tsx index 538d2f659a6..cc73b44265c 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/systemOneUI/ResponseView.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/systemOneUI/ResponseView.tsx @@ -1,9 +1,9 @@ import { Alert, AlertDescription, AlertTitle } from "@/components/ui/alert"; import { Badge } from "@/components/ui/badge"; import { Button } from "@/components/ui/button"; -import { Card, CardAction, CardContent, CardDescription, CardHeader, CardTitle } from "@/components/ui/card"; +import { Card, CardContent, CardHeader, CardTitle } from "@/components/ui/card"; import { Collapsible, CollapsibleContent, CollapsibleTrigger } from "@/components/ui/collapsible"; -import { ChevronDown, LoaderCircle } from "lucide-react"; +import { ChevronRight, Inbox, LoaderCircle } from "lucide-react"; import { useState } from "react"; import type { SystemOneAnswer, SystemOneResponse } from "./lib/schemas"; import type { OpenAIDecisionsResponse } from "./lib/openAIDecisions"; @@ -175,17 +175,19 @@ function AnswerCards({ answers }: { answers: SystemOneResponse["answers"] | Open type: answer.type, details: , })); - return cards.map(({ id, name, type, details }) => ( - - - {name} - - {type} - - - {details} - - )); + return ( +
    + {cards.map(({ id, name, type, details }) => ( +
  • +
    +

    {name}

    + {type} +
    + {details} +
    • + ))} +
+ ); } export default function ResponseView({ @@ -220,11 +222,16 @@ export default function ResponseView({ if (!response) { return ( - - - Calibrated probabilities - Send a request to see decision answers and probabilities. + + + Response + +
+ +
+

Send a request to see answers here

+ ); } @@ -232,36 +239,43 @@ export default function ResponseView({ const model = response.model ?? fallbackModel; return ( - - -
-
- {view === "curl" ? "Response JSON" : "Calibrated probabilities"} - {model && {model}} -
-
- {latencyMs !== undefined && {Math.round(latencyMs)} ms} - {response.usage && ( - - {response.usage.input_tokens} input / {response.usage.output_tokens} output tokens - - )} -
+ + + {view === "curl" ? "Response JSON" : "Response"} +
+ {latencyMs !== undefined && ( + + {Math.round(latencyMs)} ms + + )} + {model && ( + + {model} + + )} + {response.usage && ( + + {response.usage.input_tokens} in / {response.usage.output_tokens} out + + )}
- + {view === "curl" ? : } {view === "form" && ( - + - + } /> - + diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/systemOneUI/SystemOneForm.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/systemOneUI/SystemOneForm.tsx index a241e68cb9e..1e3ae9e07fe 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/systemOneUI/SystemOneForm.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/systemOneUI/SystemOneForm.tsx @@ -6,7 +6,7 @@ import { Label } from "@/components/ui/label"; import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select"; import { Textarea } from "@/components/ui/textarea"; import { useId, useMemo, useState } from "react"; -import { AddButton, FormFrame, RemoveButton } from "./DecisionFormControls"; +import { AddButton, FormFrame, QuestionCardHeader, RemoveButton } from "./DecisionFormControls"; import { blankQuestion, freeName, @@ -70,7 +70,7 @@ function KeyInput({ value, taken, onCommit, label, placeholder }: KeyInputProps) const errorId = useId(); const [draft, setDraft] = useState(value); const next = draft.trim(); - const error = keyError(next, value, taken); + const error = next === value ? undefined : keyError(next, value, taken); function commit() { if (error === undefined && next !== value) { @@ -203,34 +203,49 @@ function CriteriaEditor({ question, onChange }: QuestionEditorProps { name: string; + index: number; taken: readonly string[]; onRename: (name: string) => void; onRemove: () => void; } -function QuestionCard({ name, question, taken, onChange, onRename, onRemove }: QuestionCardProps) { +function QuestionCard({ name, index, question, taken, onChange, onRename, onRemove }: QuestionCardProps) { const instructionsId = useId(); + const title = name === "" ? String(index + 1) : name; return ( -
-
- - - +
+ +
+
+ Name + +
+
+ Answer type + +
@@ -286,6 +301,7 @@ export default function SystemOneForm({ value, onChange, validation, onReset }: writeQuestions({ ...questions, [name]: next })} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/systemOneUI/SystemOneUI.integration.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/systemOneUI/SystemOneUI.integration.test.tsx index 8158e5782bb..680fea187c9 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/systemOneUI/SystemOneUI.integration.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/systemOneUI/SystemOneUI.integration.test.tsx @@ -75,9 +75,13 @@ const selectEndpoint = async (endpoint: "/v1/decisions" | "/v1/systemone") => { }; const selectPreset = async (label: string) => { - screen.getByRole("combobox", { name: "Example" }).focus(); - await userEvent.keyboard("{ArrowDown}"); - await userEvent.click(await screen.findByRole("option", { name: label })); + await userEvent.click(screen.getByRole("button", { name: "Load example" })); + await userEvent.click(await screen.findByRole("menuitem", { name: label })); +}; + +const renderPopulated = async (ui: ReactElement) => { + render(ui); + await selectPreset("Triage a bug report"); }; const createResponse = (body: SystemOneResponse | OpenAIDecisionsResponse, status = 200, errorText = "") => @@ -108,36 +112,38 @@ describe("SystemOneUI integration", () => { document.cookie = "token=; path=/; max-age=0"; }); - it("switches between example presets and reads Custom once the request is edited", async () => { + it("starts empty with one blank question and loads examples from the menu", async () => { render(); - const input = screen.getByRole("textbox", { name: "Input" }) as HTMLTextAreaElement; - const example = input.value; - const picker = screen.getByRole("combobox", { name: "Example" }); - expect(picker).toHaveTextContent("Triage a bug report"); + expect(screen.getByRole("textbox", { name: "Input" })).toHaveValue(""); + expect(screen.getAllByRole("group", { name: /Question / })).toHaveLength(1); + expect(screen.getByRole("button", { name: "Clear form" })).toBeDisabled(); + expect(screen.getByRole("button", { name: "Send" })).toBeDisabled(); - fireEvent.change(input, { target: { value: "Something else" } }); - expect(picker).toHaveTextContent("Custom"); await selectPreset("Triage a bug report"); - expect(screen.getByRole("textbox", { name: "Input" })).toHaveValue(example); - expect(picker).toHaveTextContent("Triage a bug report"); + const input = screen.getByRole("textbox", { name: "Input" }) as HTMLTextAreaElement; + expect(input).toHaveDisplayValue(/streaming responses/); + expect(screen.getByRole("button", { name: "Clear form" })).toBeEnabled(); await selectPreset("Score a sales lead"); - expect(picker).toHaveTextContent("Score a sales lead"); expect(await readCurlPayload()).toMatchObject({ questions: { fit: { type: "score" } } }); }); it("clears the whole request and lets a preset be picked from nothing", async () => { const user = userEvent.setup(); - render(); - const clear = screen.getByRole("button", { name: "Clear" }); + await renderPopulated(); + const clear = screen.getByRole("button", { name: "Clear form" }); expect(screen.getByRole("textbox", { name: "Input" })).toHaveDisplayValue(/streaming responses/); await user.click(clear); expect(screen.getByRole("textbox", { name: "Input" })).toHaveValue(""); - expect(screen.queryAllByRole("group", { name: /Question / })).toHaveLength(0); + expect(screen.getAllByRole("group", { name: /Question / })).toHaveLength(1); + expect(screen.getByRole("textbox", { name: "Question name" })).toHaveValue(""); + expect(screen.queryByText("Enter a name")).not.toBeInTheDocument(); + expect(screen.getByRole("textbox", { name: "Instructions" })).toHaveValue(""); + expect(screen.queryByRole("textbox", { name: "Option label" })).not.toBeInTheDocument(); + expect(screen.getByText("1 issue to fix")).toBeInTheDocument(); expect(screen.getByRole("button", { name: "Send" })).toBeDisabled(); expect(clear).toBeDisabled(); - expect(screen.getByRole("combobox", { name: "Example" })).toHaveTextContent("Custom"); await selectPreset("Route a support ticket"); expect(screen.getByRole("textbox", { name: "Input" })).toHaveDisplayValue(/charged twice/); @@ -148,26 +154,28 @@ describe("SystemOneUI integration", () => { it("keeps the preset label after picking a model and keeps that model when clearing", async () => { const user = userEvent.setup(); mockModelLookup.mockResolvedValue(modelGroupInfoResponse(DECISION_AND_CHAT_MODELS)); - render(); + await renderPopulated(); const modelPicker = screen.getByRole("combobox", { name: "Decision model" }); await waitFor(() => expect(modelPicker).toHaveValue("jev-latest")); - const picker = screen.getByRole("combobox", { name: "Example" }); await user.clear(modelPicker); await user.type(modelPicker, "p"); await user.click(await screen.findByRole("option", { name: "pplx-decider" })); - expect(picker).toHaveTextContent("Triage a bug report"); expect(await readCurlPayload()).toMatchObject({ model: "pplx-decider" }); - await user.click(screen.getByRole("button", { name: "Clear" })); - expect(await readCurlPayload()).toEqual({ model: "pplx-decider", state: "", questions: {} }); + await user.click(screen.getByRole("button", { name: "Clear form" })); + expect(await readCurlPayload()).toEqual({ + model: "pplx-decider", + state: "", + questions: { "": { type: "choice", instructions: "", criteria: {} } }, + }); await selectPreset("Score a sales lead"); expect(await readCurlPayload()).toMatchObject({ model: "pplx-decider", questions: { fit: { type: "score" } } }); }); it("offers the same presets on /v1/decisions in the OpenAI shape and clears to an empty input", async () => { const user = userEvent.setup(); - render(); + await renderPopulated(); await selectEndpoint("/v1/decisions"); await selectPreset("Moderate a forum comment"); const payload = await readCurlPayload(); @@ -176,13 +184,17 @@ describe("SystemOneUI integration", () => { { type: "predicate", name: "breaks_rules", instructions: "Does this comment break the community guidelines?" }, ]); - await user.click(screen.getByRole("button", { name: "Clear" })); - expect(await readCurlPayload()).toEqual({ input: "", questions: [] }); + await user.click(screen.getByRole("button", { name: "Clear form" })); + expect(await readCurlPayload()).toEqual({ + input: "", + questions: [{ type: "predicate", name: "", instructions: "" }], + }); + expect(screen.getByRole("textbox", { name: "Question name" })).toHaveValue(""); expect(screen.getByRole("button", { name: "Send" })).toBeDisabled(); }); it("offers both literal proxy endpoints without TypeSafe or product prefixes", async () => { - render(); + await renderPopulated(); expect(screen.getByRole("combobox", { name: "Decision endpoint" })).toHaveTextContent("/v1/systemone"); screen.getByRole("combobox", { name: "Decision endpoint" }).focus(); @@ -204,7 +216,7 @@ describe("SystemOneUI integration", () => { it("opens on /v1/systemone with the first decision model the key can call", async () => { const user = userEvent.setup(); mockModelLookup.mockResolvedValue(modelGroupInfoResponse(DECISION_AND_CHAT_MODELS)); - render(); + await renderPopulated(); await waitFor(() => expect(screen.getByRole("combobox", { name: "Decision model" })).toHaveValue("jev-latest")); expect(lookupAuthHeaders(mockModelLookup.mock.calls)).toEqual([expect.arrayContaining(["Bearer session-key"])]); @@ -222,12 +234,13 @@ describe("SystemOneUI integration", () => { it("sends each endpoint's real shape and switches existing responses between cards and JSON without resending", async () => { const user = userEvent.setup(); mockModelLookup.mockResolvedValue(modelGroupInfoResponse(DECISION_AND_CHAT_MODELS)); - render(); + await renderPopulated(); await waitFor(() => expect(screen.getByRole("combobox", { name: "Decision model" })).toHaveValue("jev-latest")); fireEvent.change(screen.getByRole("textbox", { name: "Input" }), { target: { value: "An outage" } }); await selectEndpoint("/v1/decisions"); expect(screen.getByRole("textbox", { name: "Input" })).not.toHaveValue("An outage"); + await selectPreset("Triage a bug report"); expect(await readCurlPayload()).toMatchObject({ input: expect.any(String), questions: expect.arrayContaining([expect.objectContaining({ type: "predicate" })]), @@ -261,7 +274,9 @@ describe("SystemOneUI integration", () => { await selectEndpoint("/v1/systemone"); expect(screen.getByRole("textbox", { name: "Input" })).toHaveValue("An outage"); - expect(screen.queryByRole("region", { name: "Decisions results" })).not.toBeInTheDocument(); + expect( + within(screen.getByRole("region", { name: "Decisions results" })).getByText("Send a request to see answers here"), + ).toBeInTheDocument(); await user.click(screen.getByRole("button", { name: "Send" })); expect(await screen.findByText("Selected choice")).toBeInTheDocument(); expect(screen.getByRole("meter", { name: "technical probability" })).toHaveAttribute("aria-valuenow", "90"); @@ -284,9 +299,9 @@ describe("SystemOneUI integration", () => { }); it("validates /v1/decisions independently and preserves an invalid draft across switches", async () => { - render(); + await renderPopulated(); await selectEndpoint("/v1/decisions"); - await userEvent.click(screen.getByRole("button", { name: "Clear" })); + await userEvent.click(screen.getByRole("button", { name: "Clear form" })); expect(screen.getByRole("button", { name: "Send" })).toBeDisabled(); await selectEndpoint("/v1/systemone"); expect(screen.getByRole("button", { name: "Send" })).toBeEnabled(); @@ -305,8 +320,9 @@ describe("SystemOneUI integration", () => { usage: { input_tokens: 1, output_tokens: 1, total_tokens: 2 }, }), ); - render(); + await renderPopulated(); await selectEndpoint("/v1/decisions"); + await selectPreset("Triage a bug report"); expect(screen.getByRole("tab", { name: "Form" })).toHaveAttribute("aria-selected", "true"); fireEvent.change(screen.getByRole("textbox", { name: "Input" }), { target: { value: "Login is broken" } }); const firstQuestion = within(screen.getByRole("group", { name: "Question 1" })); @@ -341,8 +357,9 @@ describe("SystemOneUI integration", () => { it("adds and removes Decisions questions, boolean choices and score levels without stale fields", async () => { const user = userEvent.setup(); - render(); + await renderPopulated(); await selectEndpoint("/v1/decisions"); + await selectPreset("Triage a bug report"); const question = (index: number) => within(screen.getByRole("group", { name: `Question ${index}` })); const pick = async (control: HTMLElement, label: string) => { control.focus(); @@ -410,6 +427,7 @@ describe("SystemOneUI integration", () => { , ); + await selectPreset("Triage a bug report"); await userEvent.click(screen.getByRole("button", { name: "Send" })); await screen.findByRole("button", { name: "Cancel request" }); await selectEndpoint("/v1/decisions"); @@ -426,14 +444,15 @@ describe("SystemOneUI integration", () => { expect(screen.queryByRole("alert")).not.toBeInTheDocument(); expect(screen.queryByText("stale-model")).not.toBeInTheDocument(); expect(screen.getByRole("tab", { name: "Form" })).toHaveAttribute("aria-selected", "true"); - expect(screen.getByRole("button", { name: "Send" })).toBeEnabled(); + expect(screen.queryByRole("button", { name: "Cancel request" })).not.toBeInTheDocument(); }); it("offers only decision models and writes the picked one into the request", async () => { const user = userEvent.setup(); mockModelLookup.mockResolvedValue(modelGroupInfoResponse(DECISION_AND_CHAT_MODELS)); - render(); - const picker = await screen.findByRole("combobox", { name: "Decision model" }); + await renderPopulated(); + + const picker = screen.getByRole("combobox", { name: "Decision model" }); await user.clear(picker); await user.type(picker, "p"); @@ -456,7 +475,7 @@ describe("SystemOneUI integration", () => { mockModelLookup.mockResolvedValue( modelGroupInfoResponse([{ model_group: "typesafe/jev-latest", mode: "evaluation" }]), ); - render(); + await renderPopulated(); await waitFor(() => expect(screen.getByRole("combobox", { name: "Decision model" })).toHaveValue("typesafe/jev-latest"), ); @@ -470,7 +489,7 @@ describe("SystemOneUI integration", () => { }); it("links to Add Model when the key cannot call any decision model", async () => { - render(); + await renderPopulated(); expect(await screen.findByRole("link", { name: "Add a decision model" })).toHaveAttribute( "href", @@ -496,7 +515,7 @@ describe("SystemOneUI integration", () => { it("keeps the admin signed in when a typed virtual key has expired", async () => { document.cookie = "token=ui-session; path=/"; mockModelLookup.mockResolvedValue(EXPIRED_KEY_RESPONSE); - render(); + await renderPopulated(); fireEvent.change(screen.getByLabelText("Virtual Key", { exact: true }), { target: { value: "sk-expired" } }); @@ -511,7 +530,7 @@ describe("SystemOneUI integration", () => { it("keeps an edited request in place when decision models load afterwards", async () => { const models = Promise.withResolvers(); mockModelLookup.mockReturnValue(models.promise); - render(); + await renderPopulated(); fireEvent.change(screen.getByRole("textbox", { name: "Input" }), { target: { value: "edited" } }); await act(async () => { @@ -535,9 +554,10 @@ describe("SystemOneUI integration", () => { usage: { input_tokens: 1, output_tokens: 1, total_tokens: 2 }, }), ); - render(); + await renderPopulated(); if (endpoint === "/v1/decisions") { await selectEndpoint(endpoint); + await selectPreset("Triage a bug report"); } await userEvent.click(screen.getByRole("button", { name: "Send" })); expect(await screen.findByText("jev-latest")).toBeInTheDocument(); @@ -554,7 +574,7 @@ describe("SystemOneUI integration", () => { it("shows a read-only cURL command for the proxy with a key placeholder instead of the real key", async () => { const user = userEvent.setup(); sessionStorage.setItem("customProxyBaseUrl", "https://proxy.example/"); - render(); + await renderPopulated(); await user.click(screen.getByRole("tab", { name: "cURL" })); const command = screen.getByLabelText("cURL request"); @@ -568,35 +588,68 @@ describe("SystemOneUI integration", () => { expect(screen.getByRole("button", { name: "Copy cURL" })).toBeInTheDocument(); await selectEndpoint("/v1/decisions"); + await selectPreset("Triage a bug report"); expect(screen.getByLabelText("cURL request")).toHaveTextContent("https://proxy.example/v1/decisions"); expect(screen.getByRole("button", { name: "Send" })).toBeEnabled(); }); it("keeps Send below the form and shows results only after a request", async () => { const user = userEvent.setup(); - render(); + await renderPopulated(); const send = screen.getByRole("button", { name: "Send" }); const editor = screen.getByRole("region", { name: "Decisions request editor" }); const form = within(editor).getByRole("textbox", { name: "Input" }); expect(form.compareDocumentPosition(send) & Node.DOCUMENT_POSITION_FOLLOWING).toBeTruthy(); - expect(screen.queryByRole("region", { name: "Decisions results" })).not.toBeInTheDocument(); + expect( + within(screen.getByRole("region", { name: "Decisions results" })).getByText("Send a request to see answers here"), + ).toBeInTheDocument(); expect(screen.queryByRole("note")).not.toBeInTheDocument(); await user.click(send); const results = await screen.findByRole("region", { name: "Decisions results" }); - expect(within(results).getByText("Calibrated probabilities")).toBeInTheDocument(); + expect(within(results).getByText("Response")).toBeInTheDocument(); expect(editor.compareDocumentPosition(results) & Node.DOCUMENT_POSITION_FOLLOWING).toBeTruthy(); fireEvent.change(within(editor).getByRole("textbox", { name: "Input" }), { target: { value: "Edited state" } }); - expect(screen.queryByRole("region", { name: "Decisions results" })).not.toBeInTheDocument(); + expect( + within(screen.getByRole("region", { name: "Decisions results" })).getByText("Send a request to see answers here"), + ).toBeInTheDocument(); + }); + + it("ties each setup label to its control, numbers question cards and keeps the actions in the editor footer", async () => { + await renderPopulated(); + + expect(screen.getByText("Endpoint")).toHaveAttribute("for", screen.getByLabelText("Decision endpoint").id); + expect(screen.getByText("Model")).toHaveAttribute("for", screen.getByLabelText("Decision model").id); + expect(screen.getByText("API key")).toHaveAttribute("for", screen.getByLabelText("Virtual Key Source").id); + + const question = screen.getByRole("group", { name: "Question area" }); + expect(within(question).getByText("Question 1")).toBeInTheDocument(); + expect(within(question).getByText("Name")).toBeInTheDocument(); + expect(within(question).getByText("Answer type")).toBeInTheDocument(); + + const editor = screen.getByRole("region", { name: "Decisions request editor" }); + const card = within(editor).getByRole("group", { name: "Request editor" }); + const actions = within(card).getByRole("toolbar", { name: "Request actions" }); + const status = within(actions).getByRole("status"); + const send = within(actions).getByRole("button", { name: "Send" }); + expect(status).toHaveTextContent("Ready to send"); + expect(status.compareDocumentPosition(send) & Node.DOCUMENT_POSITION_FOLLOWING).toBeTruthy(); + expect(within(actions).queryByRole("button", { name: "Clear form" })).not.toBeInTheDocument(); + const clear = within(card).getByRole("button", { name: "Clear form" }); + expect(clear.compareDocumentPosition(actions) & Node.DOCUMENT_POSITION_FOLLOWING).toBeTruthy(); + expect( + within(card).getByRole("button", { name: "Load example" }).compareDocumentPosition(clear) & + Node.DOCUMENT_POSITION_FOLLOWING, + ).toBeTruthy(); }); it("sends to /v1/systemone even without discovered models and renders calibrated answers", async () => { const user = userEvent.setup(); sessionStorage.setItem("customProxyBaseUrl", "https://stale.example.com/"); - render(); + await renderPopulated(); sessionStorage.setItem("customProxyBaseUrl", "https://tenant.example.com/"); await user.click(screen.getByRole("button", { name: "Send" })); @@ -622,7 +675,7 @@ describe("SystemOneUI integration", () => { answers: responseBody.answers, }), ); - render(); + await renderPopulated(); await waitFor(() => expect(screen.getByRole("combobox", { name: "Decision model" })).toHaveValue("jev-latest")); @@ -634,7 +687,7 @@ describe("SystemOneUI integration", () => { it("sends a native request without model so the proxy can select its default", async () => { const user = userEvent.setup(); - render(); + await renderPopulated(); fireEvent.change(screen.getByRole("textbox", { name: "Input" }), { target: { value: "An outage" } }); expect(screen.getByRole("button", { name: "Send" })).toBeEnabled(); await user.click(screen.getByRole("button", { name: "Send" })); @@ -655,6 +708,7 @@ describe("SystemOneUI integration", () => { , ); + await selectPreset("Triage a bug report"); await user.click(screen.getByRole("button", { name: "Send" })); await screen.findByRole("button", { name: "Cancel request" }); fireEvent.change(screen.getByRole("textbox", { name: "Input" }), { target: { value: "A new request" } }); @@ -685,7 +739,7 @@ describe("SystemOneUI integration", () => { it("uses a custom virtual key when personal key creation is disabled", async () => { const user = userEvent.setup(); - render(); + await renderPopulated(); expect(screen.getByRole("button", { name: "Send" })).toBeDisabled(); fireEvent.change(screen.getByLabelText("Virtual Key", { exact: true }), { target: { value: "test-virtual-key" }, @@ -698,7 +752,7 @@ describe("SystemOneUI integration", () => { it("renders upstream errors inline", async () => { const user = userEvent.setup(); mockFetch.mockResolvedValueOnce(createResponse(responseBody, 401, "Virtual key rejected")); - render(); + await renderPopulated(); await user.click(screen.getByRole("button", { name: "Send" })); @@ -721,7 +775,7 @@ describe("SystemOneUI integration", () => { }, }), } as Response); - render(); + await renderPopulated(); await user.click(screen.getByRole("button", { name: "Send" })); @@ -731,7 +785,7 @@ describe("SystemOneUI integration", () => { it("clears a previous answer when the request editor changes", async () => { const user = userEvent.setup(); - render(); + await renderPopulated(); await user.click(screen.getByRole("button", { name: "Send" })); expect(await screen.findByText("Selected choice")).toBeInTheDocument(); @@ -752,6 +806,8 @@ describe("SystemOneUI integration", () => { , ); + await selectPreset("Triage a bug report"); + await selectPreset("Triage a bug report"); await user.click(screen.getByRole("button", { name: "Send" })); expect(await screen.findByText("Selected choice")).toBeInTheDocument(); @@ -775,7 +831,7 @@ describe("SystemOneUI integration", () => { }); }), ); - render(); + await renderPopulated(); await user.click(screen.getByRole("button", { name: "Send" })); await user.click(await screen.findByRole("button", { name: "Cancel request" })); @@ -803,7 +859,7 @@ describe("SystemOneUI integration", () => { it("explains the yes / no answer type option", async () => { const user = userEvent.setup(); - render(); + await renderPopulated(); question("has_repro_steps").getByRole("combobox", { name: "Answer type" }).focus(); await user.keyboard("{ArrowDown}"); @@ -816,7 +872,7 @@ describe("SystemOneUI integration", () => { it("sends a choice and a yes/no question built in the form without opening JSON", async () => { const user = userEvent.setup(); - render(); + await renderPopulated(); const input = screen.getByRole("textbox", { name: "Input" }); fill(input, "The login page is blank after the upgrade"); for (const name of ["area", "has_repro_steps", "severity"]) { @@ -826,6 +882,8 @@ describe("SystemOneUI integration", () => { await user.click(screen.getByRole("button", { name: "Add question" })); rename(question("question_1").getByRole("textbox", { name: "Question name" }), "team"); fill(question("team").getByRole("textbox", { name: "Instructions" }), "Which team owns this?"); + await user.click(question("team").getByRole("button", { name: "Add option" })); + await user.click(question("team").getByRole("button", { name: "Add option" })); const [first, second] = question("team").getAllByRole("textbox", { name: "Option label" }); rename(first, "frontend"); rename(second, "backend"); @@ -862,7 +920,7 @@ describe("SystemOneUI integration", () => { it("edits score levels and keeps instructions when a question changes type", async () => { const user = userEvent.setup(); mockModelLookup.mockResolvedValue(modelGroupInfoResponse(DECISION_AND_CHAT_MODELS)); - render(); + await renderPopulated(); await waitFor(() => expect(screen.getByRole("combobox", { name: "Decision model" })).toHaveValue("jev-latest")); const level = question("severity").getByRole("textbox", { name: "Level 0" }); @@ -889,14 +947,14 @@ describe("SystemOneUI integration", () => { expect(sentBody().questions.has_repro_steps).toEqual({ type: "choice", instructions: "Does the issue include steps someone could follow to reproduce it?", - criteria: { option_1: "", option_2: "", option_3: "" }, + criteria: { option_1: "" }, }); expect(Object.keys(sentBody().questions.area.criteria)).toEqual(["backend", "sdk", "ui"]); }); it("refuses an empty or taken question name and trims a new one", async () => { const user = userEvent.setup(); - render(); + await renderPopulated(); const name = question("severity").getByRole("textbox", { name: "Question name" }); expect(name).toHaveAttribute("aria-invalid", "false"); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/systemOneUI/SystemOneUI.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/systemOneUI/SystemOneUI.tsx index 5ece80abaf3..43d9d4f3024 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/systemOneUI/SystemOneUI.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/systemOneUI/SystemOneUI.tsx @@ -3,6 +3,12 @@ import { Button, buttonVariants } from "@/components/ui/button"; import { Input } from "@/components/ui/input"; import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select"; +import { + DropdownMenu, + DropdownMenuContent, + DropdownMenuItem, + DropdownMenuTrigger, +} from "@/components/ui/dropdown-menu"; import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs"; import { SimpleTooltip } from "@/components/ui/tooltip"; import { SearchSelect } from "@/components/shared/SearchSelect"; @@ -11,13 +17,14 @@ import { DECISIONS_DOCS_URL } from "@/lib/decisionModels"; import { getGlobalLitellmHeaderName, getProxyBaseUrl } from "@/components/networking"; import { uiHref } from "@/utils/uiHref"; import { useMutation } from "@tanstack/react-query"; -import { BookOpen, Eraser, LoaderCircle, Send } from "lucide-react"; -import { useEffect, useMemo, useRef, useState } from "react"; +import { BookOpen, ChevronDown, Eraser, FileText, LoaderCircle, Send } from "lucide-react"; +import { useEffect, useMemo, useRef, useState, type ReactNode } from "react"; import { makePlaygroundDecisionRequest, type PlaygroundDecisionRequest } from "../../llm_calls/system_one"; -import { DECISION_PRESETS, emptyPayload, matchPreset, presetPayload, type DecisionPreset } from "./lib/example"; +import { DECISION_PRESETS, emptyPayload, presetPayload } from "./lib/example"; import { payloadModel, withPayloadModel } from "./lib/payloadModel"; import CurlPreview from "./CurlPreview"; import ResponseView from "./ResponseView"; +import { ValidationSummary } from "./ValidationFeedback"; import SystemOneForm from "./SystemOneForm"; import OpenAIDecisionsForm from "./OpenAIDecisionsForm"; import { validateSystemOnePayload, validateOpenAIDecisionsPayload, errorCount } from "./lib/validatePayload"; @@ -36,32 +43,33 @@ interface SystemOneSendVariables { signal: AbortSignal; } -const CUSTOM_PRESET = "custom"; - -interface PresetPickerProps { - activePreset: DecisionPreset | undefined; - onPick: (id: string | null) => void; +function SetupField({ label, htmlFor, children }: { label: string; htmlFor: string; children: ReactNode }) { + return ( +
+ + {children} +
+ ); } -function PresetPicker({ activePreset, onPick }: PresetPickerProps) { +function ExamplesMenu({ onPick }: { onPick: (id: string) => void }) { return ( - + + ); } @@ -81,9 +89,8 @@ export default function SystemOneUI({ accessToken, disabledPersonalKeyCreation = const [endpoint, setEndpoint] = useState("/v1/systemone"); const [drafts, setDrafts] = useState>>({}); const [view, setView] = useState("form"); - const rawPayload = drafts[endpoint] ?? presetPayload(DECISION_PRESETS[0], endpoint, decisionModels[0]); + const rawPayload = drafts[endpoint] ?? emptyPayload(endpoint, decisionModels[0]); const draftModel = drafts[endpoint] === undefined ? decisionModels[0] : payloadModel(rawPayload); - const activePreset = matchPreset(endpoint, rawPayload); const clearedPayload = emptyPayload(endpoint, draftModel); const activeController = useRef(null); const validated = useMemo( @@ -100,7 +107,6 @@ export default function SystemOneUI({ accessToken, disabledPersonalKeyCreation = makePlaygroundDecisionRequest(request, apiKey, getCustomProxyBaseUrl(), { signal }), }); const isLoading = systemOne.isPending; - const hasResult = isLoading || systemOne.data !== undefined || systemOne.error !== null; const { reset: resetSystemOne } = systemOne; useEffect(() => () => activeController.current?.abort(), []); @@ -128,7 +134,7 @@ export default function SystemOneUI({ accessToken, disabledPersonalKeyCreation = } } - function handlePresetPick(id: string | null) { + function handlePresetPick(id: string) { const preset = DECISION_PRESETS.find((candidate) => candidate.id === id); if (preset !== undefined) { handlePayloadChange(presetPayload(preset, endpoint, draftModel)); @@ -160,33 +166,31 @@ export default function SystemOneUI({ accessToken, disabledPersonalKeyCreation = return (
-
-
-
- -
+
+
+
+ + + + ({ label: model, value: model }))} value={payloadModel(rawPayload) ?? null} @@ -195,55 +199,73 @@ export default function SystemOneUI({ accessToken, disabledPersonalKeyCreation = emptyText="No decision models found" allowClear={false} /> -
- {decisionModelsLoaded && decisionModels.length === 0 && ( - - Add a decision model - - )} - - {apiKeySource === "custom" && ( - setCustomApiKey(event.target.value)} - placeholder="Enter Virtual Key" - className="w-56" - /> - )} + {decisionModelsLoaded && decisionModels.length === 0 && ( + + Add a decision model + + )} + + + + {apiKeySource === "custom" && ( + setCustomApiKey(event.target.value)} + placeholder="Enter Virtual Key" + /> + )} +
(value === "form" || value === "curl") && setView(value)} - className="min-h-0 flex-1" + className={cn( + "flex min-h-0 flex-1 flex-col gap-0 overflow-hidden rounded-lg border bg-card", + !validation.isValid && "border-destructive/60", + )} > -
- {hasResult && ( -
- -
- )} +
+ +
); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/systemOneUI/ValidationFeedback.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/systemOneUI/ValidationFeedback.test.tsx index 24f93c58af9..a2f6c24781a 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/systemOneUI/ValidationFeedback.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/systemOneUI/ValidationFeedback.test.tsx @@ -1,6 +1,6 @@ import { render, screen } from "@testing-library/react"; import { describe, expect, it } from "vitest"; -import { IssueList, ValidationStatus } from "./ValidationFeedback"; +import { IssueList, ValidationStatus, ValidationSummary } from "./ValidationFeedback"; import type { PayloadValidation } from "./lib/validatePayload"; const validation = (issues: PayloadValidation["issues"]): PayloadValidation => ({ @@ -29,12 +29,28 @@ describe("ValidationStatus", () => { }); describe("IssueList", () => { - it("lists every issue with its path and says ready when there are none", () => { + it("lists every issue under its field label and renders nothing when there are none", () => { const { rerender } = render( , ); - expect(screen.getByRole("list", { name: "Payload validation issues" })).toHaveTextContent("stateRequired"); + expect(screen.getByRole("list", { name: "Payload validation issues" })).toHaveTextContent("InputRequired"); rerender(); - expect(screen.getByRole("status")).toHaveTextContent("Ready to send"); + expect(screen.queryByRole("list", { name: "Payload validation issues" })).not.toBeInTheDocument(); + }); +}); + +describe("ValidationSummary", () => { + it("says ready when nothing blocks the request and counts the errors otherwise", () => { + const { rerender } = render(); + expect(screen.getByRole("status")).toHaveTextContent("Ready to send"); + rerender( + , + ); + expect(screen.getByRole("status")).toHaveTextContent("1 issue to fix"); }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/systemOneUI/ValidationFeedback.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/systemOneUI/ValidationFeedback.tsx index 369c4243857..e8152456d15 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/systemOneUI/ValidationFeedback.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/systemOneUI/ValidationFeedback.tsx @@ -1,5 +1,6 @@ import { Badge } from "@/components/ui/badge"; import { CircleAlert, CircleCheck, TriangleAlert } from "lucide-react"; +import { issueLabel } from "./lib/issueLabel"; import { errorCount as countErrors, type PayloadValidation } from "./lib/validatePayload"; export function ValidationStatus({ validation }: { validation: PayloadValidation }) { @@ -14,14 +15,25 @@ export function ValidationStatus({ validation }: { validation: PayloadValidation return Valid payload; } +export function ValidationSummary({ validation }: { validation: PayloadValidation }) { + const errorCount = countErrors(validation); + return ( +

+ {errorCount === 0 ? ( + + ) : ( + + )} + + {errorCount === 0 ? "Ready to send" : `${errorCount} ${errorCount === 1 ? "issue" : "issues"} to fix`} + +

+ ); +} + export function IssueList({ id, validation }: { id: string; validation: PayloadValidation }) { if (validation.issues.length === 0) { - return ( -

- - Ready to send -

- ); + return null; } return (
    @@ -32,7 +44,7 @@ export function IssueList({ id, validation }: { id: string; validation: PayloadV ) : ( )} - {issue.path} + {issueLabel(issue.path)} {issue.message} ))} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/systemOneUI/lib/example.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/systemOneUI/lib/example.test.ts index e6b10fda545..e24cf8895b3 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/systemOneUI/lib/example.test.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/systemOneUI/lib/example.test.ts @@ -54,7 +54,10 @@ describe("decision presets", () => { it("omits the model from payloads when none is known", () => { expect(JSON.parse(presetPayload(DECISION_PRESETS[1], "/v1/systemone"))).not.toHaveProperty("model"); - expect(JSON.parse(emptyPayload("/v1/decisions"))).toEqual({ input: "", questions: [] }); + expect(JSON.parse(emptyPayload("/v1/decisions"))).toEqual({ + input: "", + questions: [{ type: "predicate", name: "", instructions: "" }], + }); expect(validateSystemOnePayload(emptyPayload("/v1/systemone"), "/v1/systemone").isValid).toBe(false); }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/systemOneUI/lib/example.ts b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/systemOneUI/lib/example.ts index c8cbeaf2e5c..226d47f4a3b 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/systemOneUI/lib/example.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/systemOneUI/lib/example.ts @@ -1,6 +1,7 @@ import type { DecisionRequest, SystemOneRequest } from "./schemas"; import type { OpenAIDecisionsRequest } from "./openAIDecisions"; import { parseJson } from "./validatePayload"; +import { blankQuestion } from "./formPayload"; type PresetChoice = { type: "choice"; @@ -192,7 +193,9 @@ export const matchPreset = (endpoint: PresetEndpoint, raw: string): DecisionPres export const emptyPayload = (endpoint: PresetEndpoint, model?: string): string => JSON.stringify( - endpoint === "/v1/decisions" ? { model, input: "", questions: [] } : { model, state: "", questions: {} }, + endpoint === "/v1/decisions" + ? { model, input: "", questions: [{ type: "predicate", name: "", instructions: "" }] } + : { model, state: "", questions: { "": blankQuestion() } }, null, 2, ); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/systemOneUI/lib/formPayload.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/systemOneUI/lib/formPayload.test.ts index 8339abeadab..ef52876ef52 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/systemOneUI/lib/formPayload.test.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/systemOneUI/lib/formPayload.test.ts @@ -82,7 +82,7 @@ describe("retype", () => { type: "choice", instructions: "Is it urgent?", note: "kept", - criteria: { option_1: "", option_2: "" }, + criteria: {}, }; const asScore = { type: "score", instructions: "Is it urgent?", note: "kept", criteria: ["", ""] }; expect(retype(noul, "choice")).toEqual(asChoice); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/systemOneUI/lib/formPayload.ts b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/systemOneUI/lib/formPayload.ts index 754476b8ff1..1906ad80cd9 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/systemOneUI/lib/formPayload.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/systemOneUI/lib/formPayload.ts @@ -64,7 +64,7 @@ export const removeAt = (items: readonly T[], index: number): T[] => export const blankQuestion = (): ChoiceQuestion => ({ type: "choice", instructions: "", - criteria: { option_1: "", option_2: "" }, + criteria: {}, }); export const retype = (question: FormQuestion, type: QuestionType): FormQuestion => { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/systemOneUI/lib/issueLabel.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/systemOneUI/lib/issueLabel.test.ts new file mode 100644 index 00000000000..29ae0f60ae4 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/systemOneUI/lib/issueLabel.test.ts @@ -0,0 +1,31 @@ +import { describe, expect, it } from "vitest"; +import { issueLabel } from "./issueLabel"; + +describe("issueLabel", () => { + it("names the Input field for both endpoints instead of the JSON key", () => { + expect(issueLabel("state")).toBe("Input"); + expect(issueLabel("input")).toBe("Input"); + }); + + it("names top-level fields and the request itself", () => { + expect(issueLabel("model")).toBe("Model"); + expect(issueLabel("questions")).toBe("Questions"); + expect(issueLabel("safety_identifier")).toBe("Safety identifier"); + expect(issueLabel("root")).toBe("Request"); + expect(issueLabel("syntax")).toBe("Request"); + }); + + it("names a question by its key or its 1-based position and the field inside it", () => { + expect(issueLabel("questions.urgency")).toBe('Question "urgency"'); + expect(issueLabel("questions.urgency.criteria")).toBe('Question "urgency" options'); + expect(issueLabel("questions.urgency.criteria.1")).toBe('Question "urgency" options'); + expect(issueLabel("questions.category.instructions")).toBe('Question "category" instructions'); + expect(issueLabel("questions.0.name")).toBe("Question 1 name"); + expect(issueLabel("questions.2.type")).toBe("Question 3 answer type"); + }); + + it("falls back to the humanized key for anything unknown", () => { + expect(issueLabel("max_tokens")).toBe("max tokens"); + expect(issueLabel("questions.a.max_items")).toBe('Question "a" max items'); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/systemOneUI/lib/issueLabel.ts b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/systemOneUI/lib/issueLabel.ts new file mode 100644 index 00000000000..76acc266117 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/systemOneUI/lib/issueLabel.ts @@ -0,0 +1,30 @@ +const TOP_LEVEL: Record = { + state: "Input", + input: "Input", + model: "Model", + questions: "Questions", + safety_identifier: "Safety identifier", + root: "Request", + syntax: "Request", +}; + +const QUESTION_FIELD: Record = { + name: "name", + type: "answer type", + instructions: "instructions", + criteria: "options", +}; + +const humanize = (segment: string): string => segment.replace(/_/g, " "); + +export function issueLabel(path: string): string { + const [head, question, field] = path.split("."); + if (head !== "questions" || question === undefined) { + return TOP_LEVEL[head] ?? humanize(head); + } + const title = /^\d+$/.test(question) ? `Question ${Number(question) + 1}` : `Question "${question}"`; + if (field === undefined) { + return title; + } + return `${title} ${QUESTION_FIELD[field] ?? humanize(field)}`; +} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/systemOneUI/lib/validatePayload.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/systemOneUI/lib/validatePayload.test.ts index 603c22c0830..6e91f09f8bd 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/systemOneUI/lib/validatePayload.test.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/systemOneUI/lib/validatePayload.test.ts @@ -1,6 +1,6 @@ import { describe, expect, it } from "vitest"; import { SYSTEM_ONE_EXAMPLE } from "./example"; -import { validateSystemOnePayload } from "./validatePayload"; +import { validateOpenAIDecisionsPayload, validateSystemOnePayload } from "./validatePayload"; const requestWith = (fields: Record = {}) => JSON.stringify({ @@ -331,3 +331,45 @@ describe("validateSystemOnePayload", () => { expect(validateSystemOnePayload(JSON.stringify(SYSTEM_ONE_EXAMPLE))).toMatchObject({ isValid: true, issues: [] }); }); }); + +describe("blank context", () => { + const question = { type: "choice", instructions: "Pick one", criteria: { a: "", b: "" } }; + + it.each(["", " "])("rejects /v1/systemone with state %j even when the schema passes", (state) => { + const result = validateSystemOnePayload(JSON.stringify({ state, questions: { q: question } }), "/v1/systemone"); + expect(result).toMatchObject({ isValid: false, payload: expect.anything() }); + expect(result.issues.map((issue) => issue.path)).toEqual(["state"]); + }); + + it("reports only the missing state while the blank question is still invalid", () => { + const raw = JSON.stringify({ state: "", questions: { "": { type: "choice", instructions: "", criteria: {} } } }); + expect(validateSystemOnePayload(raw, "/v1/systemone").issues.map((issue) => issue.path)).toEqual(["state"]); + const filled = validateSystemOnePayload(raw.replace('"state":""', '"state":"A ticket"'), "/v1/systemone"); + expect(filled.isValid).toBe(false); + expect(filled.issues).toContainEqual({ + path: "questions", + message: "Give every question a name.", + severity: "error", + }); + }); + + it("accepts /v1/systemone once the state carries text", () => { + const raw = JSON.stringify({ state: "A customer wrote in", questions: { q: question } }); + expect(validateSystemOnePayload(raw, "/v1/systemone")).toMatchObject({ isValid: true, issues: [] }); + }); + + it.each(["", " ", []])("rejects /v1/decisions with %j input even when the schema passes", (input) => { + const raw = JSON.stringify({ input, questions: [{ type: "predicate", name: "q", instructions: "Is it?" }] }); + const result = validateOpenAIDecisionsPayload(raw); + expect(result.isValid).toBe(false); + expect(result.issues.map((issue) => issue.path)).toEqual(["input"]); + }); + + it("accepts /v1/decisions once the input carries text", () => { + const raw = JSON.stringify({ + input: "Ship it?", + questions: [{ type: "predicate", name: "q", instructions: "Is it?" }], + }); + expect(validateOpenAIDecisionsPayload(raw)).toMatchObject({ isValid: true, issues: [] }); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/systemOneUI/lib/validatePayload.ts b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/systemOneUI/lib/validatePayload.ts index d74923bb3e6..1fb32efd158 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/systemOneUI/lib/validatePayload.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/systemOneUI/lib/validatePayload.ts @@ -26,6 +26,11 @@ export type SystemOnePayloadValidation = PayloadValidation; export const errorCount = (validation: PayloadValidation): number => validation.issues.filter((issue) => issue.severity === "error").length; +const toIssue = (path: PropertyKey[], message: string): SystemOnePayloadIssue => + path[0] === "questions" && path[1] === "" + ? { path: "questions", message: "Give every question a name.", severity: "error" } + : { path: path.length > 0 ? path.join(".") : "root", message, severity: "error" }; + const invalid = (path: string, message: string) => ({ isValid: false, issues: [{ path, message, severity: "error" as const }], @@ -62,26 +67,47 @@ function validatePayload(raw: string, schema: z.ZodType): PayloadValidatio if (!result.success) { return { isValid: false, - issues: result.error.issues.map((issue) => ({ - path: issue.path.length > 0 ? issue.path.join(".") : "root", - message: issue.message, - severity: "error", - })), + issues: result.error.issues.map((issue) => toIssue(issue.path, issue.message)), }; } return { isValid: true, payload: result.data, issues: [] }; } +const isBlankContext = (context: unknown): boolean => + (typeof context === "string" && context.trim() === "") || (Array.isArray(context) && context.length === 0); + +const isMissingContext = (raw: string, path: "state" | "input"): boolean => { + const json = parseJson(raw); + if (!json.ok || typeof json.value !== "object" || json.value === null) { + return false; + } + return isBlankContext((json.value as Record)[path]); +}; + +const requireContext = (raw: string, path: "state" | "input", result: PayloadValidation): PayloadValidation => + isMissingContext(raw, path) + ? { + ...result, + isValid: false, + issues: [{ path, message: "Add the text the model should decide about.", severity: "error" }], + } + : result; + export function validateSystemOnePayload( raw: string, endpoint: DecisionEndpoint = "/typesafe/v1/systemone", ): SystemOnePayloadValidation { const schema = endpoint === "/v1/systemone" ? decisionsRequestSchema : systemOneRequestSchema; const result = validatePayload(raw, schema); - return result.payload ? { ...result, issues: scoreLevelWarnings(result.payload) } : result; + return requireContext( + raw, + "state", + result.payload ? { ...result, issues: scoreLevelWarnings(result.payload) } : result, + ); } export function validateOpenAIDecisionsPayload(raw: string): PayloadValidation { - return validatePayload(raw, openAIDecisionsRequestSchema); + const result = validatePayload(raw, openAIDecisionsRequestSchema); + return requireContext(raw, "input", result); } diff --git a/ui/litellm-dashboard/src/components/AIHub/ModelHubTable.tsx b/ui/litellm-dashboard/src/components/AIHub/ModelHubTable.tsx index 41b57a78eef..2007020e117 100644 --- a/ui/litellm-dashboard/src/components/AIHub/ModelHubTable.tsx +++ b/ui/litellm-dashboard/src/components/AIHub/ModelHubTable.tsx @@ -44,7 +44,7 @@ import { checkTokenValidity } from "@/utils/jwtUtils"; import { getCookie } from "@/utils/cookieUtils"; import { getLoginUrl } from "@/utils/returnUrlUtils"; import { uiHref } from "@/utils/uiHref"; -import { DECISIONS_DOCS_URL, SYSTEM_ONE_PLAYGROUND_ROUTE, isDecisionMode } from "@/lib/decisionModels"; +import { DECISIONS_DOCS_URL, SYSTEM_ONE_PLAYGROUND_ROUTE, displayMode, isDecisionMode } from "@/lib/decisionModels"; interface ModelHubTableProps { accessToken: string | null; @@ -672,7 +672,7 @@ const ModelHubTable: React.FC = ({

Mode:

-

{selectedModel.mode || "Not specified"}

+

{displayMode(selectedModel.mode) ?? "Not specified"}

Providers:

diff --git a/ui/litellm-dashboard/src/components/AIHub/ModelHubTableColumns.test.tsx b/ui/litellm-dashboard/src/components/AIHub/ModelHubTableColumns.test.tsx index 678b811921c..dbe0371241a 100644 --- a/ui/litellm-dashboard/src/components/AIHub/ModelHubTableColumns.test.tsx +++ b/ui/litellm-dashboard/src/components/AIHub/ModelHubTableColumns.test.tsx @@ -37,6 +37,12 @@ describe("getModelHubTableColumns", () => { expect(screen.getByText("gpt-4o")).toBeInTheDocument(); }); + it("shows the legacy evaluation mode as decisions", () => { + renderTable([{ ...mockModel, mode: "evaluation" }]); + expect(screen.getByText("decisions")).toBeInTheDocument(); + expect(screen.queryByText("evaluation")).not.toBeInTheDocument(); + }); + it("shows the first two providers and '+1' for overflow", () => { renderTable([mockModel]); expect(screen.getByText("openai")).toBeInTheDocument(); diff --git a/ui/litellm-dashboard/src/components/AIHub/ModelHubTableColumns.tsx b/ui/litellm-dashboard/src/components/AIHub/ModelHubTableColumns.tsx index 3d6423f2d44..559c18930c7 100644 --- a/ui/litellm-dashboard/src/components/AIHub/ModelHubTableColumns.tsx +++ b/ui/litellm-dashboard/src/components/AIHub/ModelHubTableColumns.tsx @@ -14,6 +14,7 @@ import { DropdownMenuTrigger, } from "@/components/ui/dropdown-menu"; import { cn } from "@/lib/cva.config"; +import { displayMode } from "@/lib/decisionModels"; import { copyToClipboard } from "@/utils/dataUtils"; export interface ModelHubData { @@ -136,7 +137,7 @@ export const getModelHubTableColumns = ({ onModelClick }: ModelHubTableColumnsDe sortingFn: "alphanumeric", cell: ({ row }) => row.original.mode ? ( - {row.original.mode} + {displayMode(row.original.mode)} ) : ( - ), diff --git a/ui/litellm-dashboard/src/components/Settings/RouterSettings/Fallbacks/AddFallbacks.tsx b/ui/litellm-dashboard/src/components/Settings/RouterSettings/Fallbacks/AddFallbacks.tsx index aca47e6b5c3..6ad7624ec60 100644 --- a/ui/litellm-dashboard/src/components/Settings/RouterSettings/Fallbacks/AddFallbacks.tsx +++ b/ui/litellm-dashboard/src/components/Settings/RouterSettings/Fallbacks/AddFallbacks.tsx @@ -6,6 +6,7 @@ import React, { useEffect, useState } from "react"; import { Button } from "@/components/ui/button"; +import { Tooltip, TooltipContent, TooltipTrigger } from "@/components/ui/tooltip"; import { UiLoadingSpinner } from "@/components/ui/ui-loading-spinner"; import { toast } from "@/lib/toast"; import { fetchAvailableModels, ModelGroup } from "@/components/llm_calls/fetch_models"; @@ -20,9 +21,10 @@ interface AddFallbacksProps { accessToken: string; value?: Fallbacks; // Current fallbacks value from form onChange?: (fallbacks: Fallbacks) => Promise; // Callback to update form value + disabledReason?: string; // When set, the button is disabled and this explains why } -export default function AddFallbacks({ accessToken, value = [], onChange }: AddFallbacksProps) { +export default function AddFallbacks({ accessToken, value = [], onChange, disabledReason }: AddFallbacksProps) { const [isModalVisible, setIsModalVisible] = useState(false); const [modelInfo, setModelInfo] = useState([]); const [modalKey, setModalKey] = useState(0); // Key to force remount of form when modal opens @@ -116,10 +118,22 @@ export default function AddFallbacks({ accessToken, value = [], onChange }: AddF return (
- + {disabledReason ? ( + + }> + + + {disabledReason} + + ) : ( + + )} ({ getCallbacksCall: vi.fn(), + getRouterSettingsCall: vi.fn(), setCallbacksCall: vi.fn(), })); @@ -57,7 +58,7 @@ vi.mock("../../../common_components/DeleteResourceModal", () => ({ vi.mock("./AddFallbacks", () => ({ __esModule: true, - default: ({ value, onChange }: any) => { + default: ({ value, onChange, disabledReason }: any) => { const handleClick = async () => { if (onChange) { try { @@ -69,7 +70,7 @@ vi.mock("./AddFallbacks", () => ({ } }; return ( - ); @@ -107,6 +108,7 @@ describe("Fallbacks", () => { vi.mocked(networkingModule.getCallbacksCall).mockResolvedValue({ router_settings: mockRouterSettings, }); + vi.mocked(networkingModule.getRouterSettingsCall).mockResolvedValue({ source: { fallbacks: "db" } }); vi.mocked(networkingModule.setCallbacksCall).mockResolvedValue(undefined); vi.mocked(fetchModelsModule.fetchAvailableModels).mockResolvedValue([ { model_group: "gpt-4", mode: "chat" }, @@ -175,6 +177,67 @@ describe("Fallbacks", () => { }); }); + describe("when the router settings source cannot be loaded", () => { + beforeEach(() => { + vi.mocked(networkingModule.getRouterSettingsCall).mockRejectedValue(new Error("500 Internal Server Error")); + }); + + it("keeps add, edit and delete grayed out instead of guessing the list is editable", async () => { + renderWithQueryClient(); + + expect((await screen.findAllByText("gpt-4")).length).toBeGreaterThan(0); + expect(screen.getByTestId("add-fallbacks-button")).toBeDisabled(); + for (const button of screen.getAllByTestId("edit-fallback-button")) { + expect(button).toHaveAttribute("aria-disabled", "true"); + } + for (const button of screen.getAllByTestId("delete-fallback-button")) { + expect(button).toHaveAttribute("aria-disabled", "true"); + } + expect(screen.queryByTestId("config-owned-badge")).not.toBeInTheDocument(); + }); + }); + + describe("when config.yaml owns fallbacks", () => { + beforeEach(() => { + vi.mocked(networkingModule.getRouterSettingsCall).mockResolvedValue({ source: { fallbacks: "config" } }); + }); + + it("shows the config rules but grays out add, edit and delete", async () => { + renderWithQueryClient(); + + await waitFor(() => { + expect(screen.getByTestId("config-owned-badge")).toBeInTheDocument(); + }); + expect(screen.getAllByText("gpt-4").length).toBeGreaterThan(0); + expect(screen.getByTestId("add-fallbacks-button")).toBeDisabled(); + for (const button of screen.getAllByTestId("edit-fallback-button")) { + expect(button).toHaveAttribute("aria-disabled", "true"); + } + for (const button of screen.getAllByTestId("delete-fallback-button")) { + expect(button).toHaveAttribute("aria-disabled", "true"); + } + }); + + it("does not open the edit or delete modals from click or keyboard", async () => { + const user = userEvent.setup(); + renderWithQueryClient(); + await waitFor(() => { + expect(screen.getByTestId("config-owned-badge")).toBeInTheDocument(); + }); + + await user.click(screen.getAllByTestId("delete-fallback-button")[0]); + screen.getAllByTestId("delete-fallback-button")[0].focus(); + await user.keyboard("{Enter}"); + expect(screen.queryByTestId("delete-modal")).not.toBeInTheDocument(); + + await user.click(screen.getAllByTestId("edit-fallback-button")[0]); + screen.getAllByTestId("edit-fallback-button")[0].focus(); + await user.keyboard("{Enter}"); + expect(screen.queryByText("Configure Model Fallbacks")).not.toBeInTheDocument(); + expect(networkingModule.setCallbacksCall).not.toHaveBeenCalled(); + }); + }); + it("should open delete modal when delete icon is clicked", async () => { const user = userEvent.setup(); renderWithQueryClient(); diff --git a/ui/litellm-dashboard/src/components/Settings/RouterSettings/Fallbacks/Fallbacks.tsx b/ui/litellm-dashboard/src/components/Settings/RouterSettings/Fallbacks/Fallbacks.tsx index 2ccd1a70dc5..03fd1d43510 100644 --- a/ui/litellm-dashboard/src/components/Settings/RouterSettings/Fallbacks/Fallbacks.tsx +++ b/ui/litellm-dashboard/src/components/Settings/RouterSettings/Fallbacks/Fallbacks.tsx @@ -7,7 +7,14 @@ import React, { useEffect, useState } from "react"; import DeleteResourceModal from "../../../common_components/DeleteResourceModal"; import { ProviderLogo } from "../../../molecules/models/ProviderLogo"; import { toast } from "@/lib/toast"; -import { getCallbacksCall, setCallbacksCall } from "../../../networking"; +import { getCallbacksCall, getRouterSettingsCall, setCallbacksCall } from "../../../networking"; +import { + CONFIG_OWNED_HINT, + SOURCE_UNKNOWN_HINT, + ConfigOwnedBadge, + RouterSettingsSource, + isOwnedByConfig, +} from "../../../router_settings/ConfigOwned"; import { isProxyAdminRole } from "@/utils/roles"; import AddFallbacks from "./AddFallbacks"; import EditFallbacks from "./EditFallbacks"; @@ -119,12 +126,19 @@ async function testFallbackModelResponse(selectedModel: string, accessToken: str } } +function fallbacksWriteHint(source: RouterSettingsSource | null, ownedByConfig: boolean): string | undefined { + if (ownedByConfig) return CONFIG_OWNED_HINT; + if (source === null) return SOURCE_UNKNOWN_HINT; + return undefined; +} + const Fallbacks: React.FC = ({ accessToken, userRole, userID }) => { const [routerSettings, setRouterSettings] = useState<{ [key: string]: any }>({}); const [isDeleting, setIsDeleting] = useState(false); const [fallbackToDelete, setFallbackToDelete] = useState(null); const [isDeleteModalOpen, setIsDeleteModalOpen] = useState(false); const [fallbackToEdit, setFallbackToEdit] = useState(null); + const [routerSettingsSource, setRouterSettingsSource] = useState(null); const { data: modelCostMapData } = useModelCostMap(); const getProviderFromModel = (model: string): string => { @@ -145,6 +159,9 @@ const Fallbacks: React.FC = ({ accessToken, userRole, userID }) } setRouterSettings(router_settings); }); + getRouterSettingsCall(accessToken) + .then((data) => setRouterSettingsSource(data.source ?? {})) + .catch(() => setRouterSettingsSource(null)); }, [accessToken, userRole, userID]); const handleDeleteClick = (fallbackEntry: FallbackEntry) => { @@ -250,15 +267,24 @@ const Fallbacks: React.FC = ({ accessToken, userRole, userID }) const hasFallbacks = Array.isArray(routerSettings.fallbacks) && routerSettings.fallbacks.length > 0; // Admin Viewer follows the read-parity rule: see fallbacks, no writes. const canModify = isProxyAdminRole(userRole ?? ""); + const ownedByConfig = routerSettingsSource !== null && isOwnedByConfig(routerSettingsSource, "fallbacks"); + const writeHint = fallbacksWriteHint(routerSettingsSource, ownedByConfig); + const canWrite = canModify && writeHint === undefined; + const writeActionClass = (hoverClass: string) => + canWrite ? `${iconWrapperClass} cursor-pointer ${hoverClass}` : `${iconWrapperClass} cursor-not-allowed opacity-50`; return ( {canModify && ( - +
+ + {ownedByConfig && } +
)} {!hasFallbacks ? (
@@ -309,15 +335,16 @@ const Fallbacks: React.FC = ({ accessToken, userRole, userID }) data-testid="edit-fallback-button" role="button" tabIndex={0} - onClick={() => handleEditClick(item)} - onKeyDown={(e) => e.key === "Enter" && handleEditClick(item)} - className={`${iconWrapperClass} cursor-pointer hover:text-info`} + aria-disabled={!canWrite} + onClick={() => canWrite && handleEditClick(item)} + onKeyDown={(e) => canWrite && e.key === "Enter" && handleEditClick(item)} + className={writeActionClass("hover:text-info")} /> } > - Edit fallback + {writeHint ?? "Edit fallback"} = ({ accessToken, userRole, userID }) data-testid="delete-fallback-button" role="button" tabIndex={0} - onClick={() => handleDeleteClick(item)} - onKeyDown={(e) => e.key === "Enter" && handleDeleteClick(item)} - className={`${iconWrapperClass} cursor-pointer hover:text-destructive`} + aria-disabled={!canWrite} + onClick={() => canWrite && handleDeleteClick(item)} + onKeyDown={(e) => canWrite && e.key === "Enter" && handleDeleteClick(item)} + className={writeActionClass("hover:text-destructive")} /> } > - Delete fallback + {writeHint ?? "Delete fallback"} )} diff --git a/ui/litellm-dashboard/src/components/add_model/AddModelForm.integration.test.tsx b/ui/litellm-dashboard/src/components/add_model/AddModelForm.integration.test.tsx index 067adcaeb56..681217cc918 100644 --- a/ui/litellm-dashboard/src/components/add_model/AddModelForm.integration.test.tsx +++ b/ui/litellm-dashboard/src/components/add_model/AddModelForm.integration.test.tsx @@ -429,7 +429,7 @@ describe("AddModelForm", () => { expect(await screen.findByRole("button", { name: "Add Model" })).toBeInTheDocument(); }); - it("offers the Evaluation decisions mode", async () => { + it("offers the Decisions mode", async () => { const mockUseAuthorized = vi.mocked(await import("@/app/(dashboard)/hooks/useAuthorized")); mockUseAuthorized.default.mockReturnValue(mockAuthorizedUser("proxy_admin", "user-1", true)); @@ -438,7 +438,7 @@ describe("AddModelForm", () => { await screen.findByText("Provider"); await userEvent.click(screen.getByRole("combobox", { name: "Mode" })); - expect(await screen.findByRole("option", { name: "Evaluation - /v1/decisions", exact: true })).toBeInTheDocument(); + expect(await screen.findByRole("option", { name: "Decisions - /v1/decisions", exact: true })).toBeInTheDocument(); }); it("shows only the Close button in the connection test dialog footer", async () => { diff --git a/ui/litellm-dashboard/src/components/add_model/add_model_modes.tsx b/ui/litellm-dashboard/src/components/add_model/add_model_modes.tsx index 289fd390c62..45ffe969a21 100644 --- a/ui/litellm-dashboard/src/components/add_model/add_model_modes.tsx +++ b/ui/litellm-dashboard/src/components/add_model/add_model_modes.tsx @@ -13,7 +13,7 @@ export const TEST_MODES = [ { value: "batch", label: "Batch - /batch" }, { value: "anthropic_messages", label: "Anthropic Messages - /v1/messages" }, { value: "ocr", label: "OCR - /ocr" }, - { value: "evaluation", label: "Evaluation - /v1/decisions" }, + { value: "decisions", label: "Decisions - /v1/decisions" }, ]; // Define the available auto router routing strategies diff --git a/ui/litellm-dashboard/src/components/chat/ConnectFlowBanner.integration.test.tsx b/ui/litellm-dashboard/src/components/chat/ConnectFlowBanner.integration.test.tsx new file mode 100644 index 00000000000..51647247f75 --- /dev/null +++ b/ui/litellm-dashboard/src/components/chat/ConnectFlowBanner.integration.test.tsx @@ -0,0 +1,91 @@ +import { webcrypto } from "node:crypto"; +import { afterEach, expect, it, vi } from "vitest"; +import { render, screen, waitFor } from "@testing-library/react"; +import ConnectFlowBanner from "./ConnectFlowBanner"; + +const locationDescriptor = Object.getOwnPropertyDescriptor(window, "location")!; + +afterEach(() => { + vi.unstubAllGlobals(); + window.sessionStorage.clear(); + Object.defineProperty(window, "location", locationDescriptor); +}); + +it("carries a challenged scope through the connect button into vendor authorization", async () => { + vi.stubGlobal("crypto", webcrypto); + vi.stubGlobal( + "fetch", + vi.fn(async () => Response.json({ client_id: "registered-client" })), + ); + Object.defineProperty(window, "location", { + configurable: true, + value: { + href: "https://gateway.example/ui/connect?connect_flow=flow-handle", + origin: "https://gateway.example", + pathname: "/ui/connect", + }, + }); + render( + , + ); + await waitFor(() => expect(window.location.href).toContain("scope=tools.write")); + expect(new URL(window.location.href, window.location.origin).pathname).toBe( + "/v1/mcp/server/oauth/server-1/authorize", + ); + expect(screen.queryByRole("button", { name: /finish connecting/i })).not.toBeInTheDocument(); + expect(screen.getByRole("button", { name: "Cancel" })).toBeInTheDocument(); +}); + +it("starts consent for the next upstream in the same connect flow", async () => { + vi.stubGlobal("crypto", webcrypto); + vi.stubGlobal( + "fetch", + vi.fn(async () => Response.json({ client_id: "registered-client" })), + ); + Object.defineProperty(window, "location", { + configurable: true, + value: { + href: "https://gateway.example/ui/connect?connect_flow=flow-handle", + origin: "https://gateway.example", + pathname: "/ui/connect", + }, + }); + const props = { + flowHandle: "flow-handle", + accessToken: "caller", + onConnected: vi.fn(), + failed: false, + }; + const first = { + state: "interactive" as const, + client_origin: "https://client.example", + server_id: "first", + server_name: "First", + connected: false, + requested_scopes: ["first.write"], + }; + const { rerender } = render(); + await waitFor(() => expect(window.location.href).toContain("scope=first.write")); + rerender( + , + ); + await waitFor(() => expect(window.location.href).toContain("scope=second.write")); + expect(new URL(window.location.href, window.location.origin).pathname).toBe("/v1/mcp/server/oauth/second/authorize"); + expect(screen.queryByRole("button", { name: /finish connecting/i })).not.toBeInTheDocument(); +}); diff --git a/ui/litellm-dashboard/src/components/chat/ConnectFlowBanner.tsx b/ui/litellm-dashboard/src/components/chat/ConnectFlowBanner.tsx index 0d6e708f734..e32ae028097 100644 --- a/ui/litellm-dashboard/src/components/chat/ConnectFlowBanner.tsx +++ b/ui/litellm-dashboard/src/components/chat/ConnectFlowBanner.tsx @@ -76,11 +76,13 @@ const ConnectFlowBanner: React.FC = ({ flowHandle, flow, accessToken, onC
{vendorServer !== null && ( )}
diff --git a/ui/litellm-dashboard/src/components/chat/MCPAppsPanel.tsx b/ui/litellm-dashboard/src/components/chat/MCPAppsPanel.tsx index 1fee8923e94..d8386afeb93 100644 --- a/ui/litellm-dashboard/src/components/chat/MCPAppsPanel.tsx +++ b/ui/litellm-dashboard/src/components/chat/MCPAppsPanel.tsx @@ -31,6 +31,7 @@ interface OAuth2ConnectButtonProps { onConnect: (serverId: string) => void; variant?: "badge" | "button"; autoStartKey?: string | null; + scopes?: string[]; } export const OAuth2ConnectButton: React.FC = ({ @@ -39,14 +40,17 @@ export const OAuth2ConnectButton: React.FC = ({ onConnect, variant = "badge", autoStartKey = null, + scopes, }) => { const name = server.server_name ?? server.alias ?? server.server_id; - const { startOAuthFlow, status } = useUserMcpOAuthFlow({ + const flowOptions = { accessToken, serverId: server.server_id, serverAlias: name, + scopes, onSuccess: useCallback(() => onConnect(server.server_id), [onConnect, server.server_id]), - }); + }; + const { startOAuthFlow, status } = useUserMcpOAuthFlow(flowOptions); useEffect(() => { if (autoStartKey === null || status !== "idle" || getSecureItem(autoStartKey) !== null) return; diff --git a/ui/litellm-dashboard/src/components/liteadmin/LiteAdmin.integration.test.tsx b/ui/litellm-dashboard/src/components/liteadmin/LiteAdmin.integration.test.tsx index c2b8805b892..2bc586cd1bb 100644 --- a/ui/litellm-dashboard/src/components/liteadmin/LiteAdmin.integration.test.tsx +++ b/ui/litellm-dashboard/src/components/liteadmin/LiteAdmin.integration.test.tsx @@ -9,7 +9,7 @@ import userEvent from "@testing-library/user-event"; import { useEffect, type ComponentType } from "react"; import SidebarAccountMenu from "@/components/SidebarAccountMenu/SidebarAccountMenu"; import UserDropdown from "@/components/Navbar/UserDropdown/UserDropdown"; -import LiteAdmin, { LiteAdminFrame } from "./LiteAdmin"; +import LiteAdmin, { LiteAdminFrame, LiteAdminHome } from "./LiteAdmin"; import { MAX_INPUT_LENGTH } from "./agent"; import type { ModelGroup } from "@/components/llm_calls/fetch_models"; @@ -40,7 +40,7 @@ const keyArguments = { duration: null, }; -type RecordedRequest = { url: string; headers: Headers; body: Record }; +type RecordedRequest = { url: string; headers: Headers; body: Record; signal: AbortSignal }; type ModelReply = { role: "assistant"; content: string | null; @@ -98,9 +98,9 @@ function Page() { return null; } -function renderWidget(Menu?: ComponentType<{ onLogout: () => void }>) { +function renderWidget(Menu?: ComponentType<{ onLogout: () => void }>, showHome = false) { const client = new QueryClient({ defaultOptions: { queries: { retry: false, gcTime: 0 } } }); - const tree = () => ( + const tree = (home = showHome) => ( @@ -108,13 +108,14 @@ function renderWidget(Menu?: ComponentType<{ onLogout: () => void }>) { {Menu && undefined} />} + {home && } ); const view = render(tree()); - return { ...view, refresh: () => view.rerender(tree()), client }; + return { ...view, refresh: (home = showHome) => view.rerender(tree(home)), client }; } interface GatewayOptions { @@ -130,7 +131,8 @@ function gateway(replies: (ModelReply | Promise)[], options: Gateway transport.mockImplementation(async (input, init) => { const request = input instanceof Request ? input : new Request(new URL(String(input), "http://localhost"), init); const body: Record = request.method === "GET" ? {} : await request.clone().json(); - requests.push({ url: request.url, headers: request.headers, body }); + const recorded = { url: request.url, headers: request.headers, body, signal: request.signal }; + requests.push(recorded); const path = new URL(request.url).pathname; if (path.endsWith("/litellm-ui-config")) return json({ proxy_base_url: MANAGEMENT, server_root_path: "", admin_ui_disabled: false }); @@ -170,10 +172,10 @@ function gateway(replies: (ModelReply | Promise)[], options: Gateway return requests; } -async function selectModel() { +async function selectModel(model = "chat-model") { const user = userEvent.setup(); await user.click(await screen.findByRole("combobox", { name: "LiteAdmin model" })); - await user.click(await screen.findByRole("option", { name: "chat-model" })); + await user.click(await screen.findByRole("option", { name: model })); } async function openWidget() { @@ -202,13 +204,388 @@ afterEach(() => { }); describe("LiteAdmin in the gateway", () => { + it.each(["keyboard", "click"])("continues one conversation from Home via %s and the dock", async (method) => { + const requests = gateway([answer("First reply."), answer("Second reply."), answer("Third reply.")]); + renderWidget(undefined, true); + const home = await screen.findByRole("region", { name: "Ask LiteAdmin" }); + await selectModel(); + const input = within(home).getByPlaceholderText("Ask LiteAdmin…"); + fireEvent.change(input, { target: { value: "First question" } }); + if (method === "keyboard") fireEvent.keyDown(input, { key: "Enter" }); + else fireEvent.click(within(home).getByRole("button", { name: "Send message" })); + expect(await screen.findByText("First reply.")).toBeInTheDocument(); + expect(home).not.toBeVisible(); + expect(screen.getAllByRole("complementary", { name: "LiteAdmin" })).toHaveLength(1); + expect(requests.filter((request) => request.url.endsWith("/chat/completions"))).toHaveLength(1); + send("Second question"); + expect(await screen.findByText("Second reply.")).toBeInTheDocument(); + fireEvent.click(screen.getByRole("button", { name: "Close LiteAdmin" })); + expect(home).toBeVisible(); + send("Third question"); + expect(await screen.findByText("Third reply.")).toBeInTheDocument(); + const completions = requests.filter((request) => request.url.endsWith("/chat/completions")); + expect(completions).toHaveLength(3); + expect(completions[2].body.messages).toEqual([ + { role: "system", content: expect.any(String) }, + { role: "user", content: "First question" }, + { role: "assistant", content: "First reply." }, + { role: "user", content: "Second question" }, + { role: "assistant", content: "Second reply." }, + { role: "user", content: "Third question" }, + ]); + }); + + it("preserves the draft and chosen model when opening, closing and leaving Home", async () => { + const requests = gateway([answer("Draft received.")], { + models: [ + { model_group: "claude-sonnet-5-5", mode: "chat" }, + { model_group: "chat-model", mode: "chat" }, + ], + }); + const view = renderWidget(undefined, true); + await selectModel(); + fireEvent.change(screen.getByPlaceholderText("Ask LiteAdmin…"), { target: { value: "Keep this draft" } }); + fireEvent.click(screen.getByRole("button", { name: "LiteAdmin" })); + expect(screen.getByPlaceholderText("Ask LiteAdmin…")).toHaveValue("Keep this draft"); + fireEvent.click(screen.getByRole("button", { name: "Close LiteAdmin" })); + const home = screen.getByRole("region", { name: "Ask LiteAdmin" }); + expect(within(home).getByPlaceholderText("Ask LiteAdmin…")).toHaveValue("Keep this draft"); + view.refresh(false); + expect(home).not.toBeInTheDocument(); + fireEvent.click(screen.getByRole("button", { name: "LiteAdmin" })); + expect(screen.getByPlaceholderText("Ask LiteAdmin…")).toHaveValue("Keep this draft"); + fireEvent.click(screen.getByRole("button", { name: "Send message" })); + expect(await screen.findByText("Draft received.")).toBeInTheDocument(); + const completion = requests.find((request) => request.url.endsWith("/chat/completions")); + expect(completion?.body.model).toBe("chat-model"); + expect(completion?.body.messages).toContainEqual({ role: "user", content: "Keep this draft" }); + }); + + it("keeps each tab's history, model and draft independent through the Home composer", async () => { + const requests = gateway( + [answer("First answer."), answer("Second answer."), answer("Second follow-up."), answer("First follow-up.")], + { + models: [ + { model_group: "chat-model", mode: "chat" }, + { model_group: "other-chat-model", mode: "chat" }, + ], + }, + ); + renderWidget(undefined, true); + await selectModel(); + send("First question"); + await screen.findByText("First answer."); + expect(within(screen.getByRole("tablist", { name: "LiteAdmin chats" })).getAllByRole("tab")).toHaveLength(1); + fireEvent.click(screen.getByRole("button", { name: "New chat" })); + expect(screen.getByRole("tab", { name: "Chat 2" })).toHaveAttribute("aria-selected", "true"); + expect(screen.queryByText("First answer.")).not.toBeInTheDocument(); + expect(screen.getByPlaceholderText("Ask LiteAdmin…")).toHaveValue(""); + await selectModel("other-chat-model"); + send("Second question"); + await screen.findByText("Second answer."); + fireEvent.change(screen.getByPlaceholderText("Ask LiteAdmin…"), { target: { value: "Second draft" } }); + + fireEvent.click(screen.getByRole("tab", { name: "Chat 1" })); + expect(screen.getByText("First answer.")).toBeInTheDocument(); + expect(screen.queryByText("Second answer.")).not.toBeInTheDocument(); + expect(screen.getByPlaceholderText("Ask LiteAdmin…")).toHaveValue(""); + expect(screen.getByRole("combobox", { name: "LiteAdmin model" })).toHaveValue("chat-model"); + fireEvent.change(screen.getByPlaceholderText("Ask LiteAdmin…"), { target: { value: "First draft" } }); + fireEvent.click(screen.getByRole("button", { name: "Close LiteAdmin" })); + const home = screen.getByRole("region", { name: "Ask LiteAdmin" }); + expect(within(home).getByPlaceholderText("Ask LiteAdmin…")).toHaveValue("First draft"); + expect(screen.getAllByPlaceholderText("Ask LiteAdmin…")).toHaveLength(1); + + fireEvent.click(screen.getByRole("button", { name: "LiteAdmin" })); + fireEvent.click(screen.getByRole("tab", { name: "Chat 2" })); + expect(screen.getByPlaceholderText("Ask LiteAdmin…")).toHaveValue("Second draft"); + expect(screen.getByRole("combobox", { name: "LiteAdmin model" })).toHaveValue("other-chat-model"); + expect(screen.getAllByLabelText("LiteAdmin conversation")).toHaveLength(1); + fireEvent.click(screen.getByRole("button", { name: "Close LiteAdmin" })); + expect(within(home).getByPlaceholderText("Ask LiteAdmin…")).toHaveValue("Second draft"); + fireEvent.click(within(home).getByRole("button", { name: "Send message" })); + await screen.findByText("Second follow-up."); + expect(screen.getByRole("tab", { name: "Chat 2" })).toHaveAttribute("aria-selected", "true"); + fireEvent.click(screen.getByRole("tab", { name: "Chat 1" })); + expect(screen.getByPlaceholderText("Ask LiteAdmin…")).toHaveValue("First draft"); + fireEvent.click(screen.getByRole("button", { name: "Send message" })); + await screen.findByText("First follow-up."); + + const completions = requests.filter((request) => request.url.endsWith("/chat/completions")); + expect(completions.map((request) => request.body.model)).toEqual([ + "chat-model", + "other-chat-model", + "other-chat-model", + "chat-model", + ]); + expect(completions[1].body.messages).toEqual([ + { role: "system", content: expect.any(String) }, + { role: "user", content: "Second question" }, + ]); + expect(completions[2].body.messages).toEqual([ + { role: "system", content: expect.any(String) }, + { role: "user", content: "Second question" }, + { role: "assistant", content: "Second answer." }, + { role: "user", content: "Second draft" }, + ]); + expect(completions[3].body.messages).toEqual([ + { role: "system", content: expect.any(String) }, + { role: "user", content: "First question" }, + { role: "assistant", content: "First answer." }, + { role: "user", content: "First draft" }, + ]); + }); + + it("supports keyboard tab switching and keeps focus on a surviving tab after closing", async () => { + gateway([]); + renderWidget(); + await openWidget(); + const user = userEvent.setup(); + await user.click(screen.getByRole("button", { name: "New chat" })); + await user.click(screen.getByRole("tab", { name: "Chat 1" })); + await user.keyboard("{ArrowRight}{Enter}"); + expect(screen.getByRole("tab", { name: "Chat 2" })).toHaveAttribute("aria-selected", "true"); + expect(screen.getByRole("tab", { name: "Chat 2" })).toHaveFocus(); + await user.click(screen.getByRole("button", { name: "Close Chat 2" })); + expect(screen.getByRole("tab", { name: "Chat 1" })).toHaveFocus(); + await user.click(screen.getByRole("button", { name: "Close Chat 1" })); + expect(screen.getByRole("tab", { name: "Chat 3" })).toHaveFocus(); + expect(screen.getByRole("tab", { name: "Chat 3" })).toHaveAttribute("aria-selected", "true"); + }); + + it("keeps tab focus when returning to a saved approval and focuses it when reopening the panel", async () => { + const requests = gateway([toolReply("key_create", keyArguments)]); + renderWidget(); + await openWidget(); + send("Prepare a key for approval"); + const review = await screen.findByRole("region", { name: "Create a virtual key" }); + expect(review).toHaveFocus(); + const user = userEvent.setup(); + await user.click(screen.getByRole("button", { name: "New chat" })); + await user.click(screen.getByRole("tab", { name: "Chat 2" })); + await user.keyboard("{ArrowLeft}{Enter}"); + expect(screen.getByRole("tab", { name: "Chat 1" })).toHaveFocus(); + expect(screen.getByRole("button", { name: "Confirm change" })).toBeInTheDocument(); + await user.click(screen.getByRole("tab", { name: "Chat 2" })); + await user.click(screen.getByRole("button", { name: "Close Chat 2" })); + expect(screen.getByRole("tab", { name: "Chat 1" })).toHaveFocus(); + await user.click(screen.getByRole("button", { name: "Close LiteAdmin" })); + await user.click(screen.getByRole("button", { name: "LiteAdmin" })); + expect(screen.getByRole("region", { name: "Create a virtual key" })).toHaveFocus(); + expect(requests.filter((request) => request.url.endsWith("/key/generate"))).toHaveLength(0); + }); + + it("lets simultaneous turns finish independently without a background review stealing focus", async () => { + const first = deferred(); + const second = deferred(); + const requests = gateway([first.promise, second.promise]); + renderWidget(); + await openWidget(); + send("First pending question"); + await waitFor(() => + expect(requests.filter((request) => request.url.endsWith("/chat/completions"))).toHaveLength(1), + ); + fireEvent.click(screen.getByRole("button", { name: "New chat" })); + await selectModel(); + send("Second pending question"); + await waitFor(() => + expect(requests.filter((request) => request.url.endsWith("/chat/completions"))).toHaveLength(2), + ); + const completions = requests.filter((request) => request.url.endsWith("/chat/completions")); + expect(completions.every((request) => !request.signal.aborted)).toBe(true); + await act(async () => second.resolve(answer("Second completed first."))); + await screen.findByText("Second completed first."); + const composer = screen.getByPlaceholderText("Ask LiteAdmin…"); + fireEvent.change(composer, { target: { value: "Keep typing in the second chat" } }); + act(() => composer.focus()); + + await act(async () => first.resolve(toolReply("key_create", keyArguments))); + expect(composer).toHaveFocus(); + expect(composer).toHaveValue("Keep typing in the second chat"); + expect(composer).toBeEnabled(); + expect(screen.getByRole("tab", { name: "Chat 2" })).toHaveAttribute("aria-selected", "true"); + expect(screen.queryByRole("button", { name: "Confirm change" })).not.toBeInTheDocument(); + expect(screen.queryByText("First pending question")).not.toBeInTheDocument(); + fireEvent.click(screen.getByRole("tab", { name: "Chat 1" })); + expect(await screen.findByRole("region", { name: "Create a virtual key" })).toHaveTextContent("Widget key"); + expect(screen.getByText("First pending question")).toBeInTheDocument(); + expect(screen.queryByText("Second completed first.")).not.toBeInTheDocument(); + fireEvent.click(screen.getByRole("tab", { name: "Chat 2" })); + expect(screen.getByPlaceholderText("Ask LiteAdmin…")).toHaveValue("Keep typing in the second chat"); + expect(screen.getByText("Second completed first.")).toBeInTheDocument(); + expect(requests.filter((request) => request.url.endsWith("/key/generate"))).toHaveLength(0); + }); + + it("cancels only a closed tab's approval and prevents closing an applying write even from another tab", async () => { + const response = deferred(); + const requests = gateway( + [ + toolReply("key_create", keyArguments), + answer("Second chat is usable."), + toolReply("key_create", { ...keyArguments, key_alias: "Second chat key" }), + answer("Third chat is usable."), + answer("Second chat write completed."), + ], + { write: () => response.promise }, + ); + renderWidget(); + await openWidget(); + send("Propose a first key"); + await screen.findByRole("button", { name: "Confirm change" }); + fireEvent.click(screen.getByRole("button", { name: "New chat" })); + await selectModel(); + expect(screen.queryByRole("button", { name: "Confirm change" })).not.toBeInTheDocument(); + send("Read in the second chat"); + await screen.findByText("Second chat is usable."); + fireEvent.click(screen.getByRole("tab", { name: "Chat 1" })); + expect(screen.getByRole("region", { name: "Create a virtual key" })).toHaveTextContent("Widget key"); + fireEvent.click(screen.getByRole("tab", { name: "Chat 2" })); + fireEvent.click(screen.getByRole("button", { name: "Close Chat 1" })); + expect(screen.queryByRole("tab", { name: "Chat 1" })).not.toBeInTheDocument(); + expect(screen.getByRole("tab", { name: "Chat 2" })).toHaveAttribute("aria-selected", "true"); + expect(requests.filter((request) => request.url.endsWith("/key/generate"))).toHaveLength(0); + + send("Create the second chat key"); + fireEvent.click(await screen.findByRole("button", { name: "Confirm change" })); + await waitFor(() => expect(requests.filter((request) => request.url.endsWith("/key/generate"))).toHaveLength(1)); + expect(screen.getByRole("button", { name: "Close Chat 2" })).toBeDisabled(); + fireEvent.click(screen.getByRole("button", { name: "Close Chat 2" })); + expect(screen.getByRole("tab", { name: "Chat 2" })).toHaveAttribute("aria-selected", "true"); + expect(screen.getByRole("button", { name: "New chat" })).toBeEnabled(); + fireEvent.click(screen.getByRole("button", { name: "New chat" })); + expect(screen.getByRole("button", { name: "Close Chat 2" })).toBeDisabled(); + fireEvent.click(screen.getByRole("button", { name: "Close Chat 2" })); + expect(screen.getByRole("tab", { name: "Chat 2" })).toBeInTheDocument(); + await selectModel(); + send("Read while the other tab writes"); + await screen.findByText("Third chat is usable."); + await act(async () => response.resolve(json({ key: NEW_KEY }))); + await waitFor(() => expect(screen.getByRole("button", { name: "Close Chat 2" })).toBeEnabled()); + expect(screen.getByRole("tab", { name: "Chat 3" })).toHaveAttribute("aria-selected", "true"); + expect(screen.queryByText(NEW_KEY)).not.toBeInTheDocument(); + fireEvent.click(screen.getByRole("tab", { name: "Chat 2" })); + expect(await screen.findByText("Second chat write completed.")).toBeInTheDocument(); + expect(screen.getByText(NEW_KEY)).toBeInTheDocument(); + expect(screen.getByRole("region", { name: "Create a virtual key" })).toHaveTextContent("Completed"); + expect(requests.filter((request) => request.url.endsWith("/key/generate")).map((request) => request.body)).toEqual([ + { key_alias: "Second chat key", team_id: "team-1", max_budget: 40 }, + ]); + expect(screen.queryByText(/A submitted change may have completed/)).not.toBeInTheDocument(); + }); + + it("closes inactive, active and last tabs without cancelling survivors or reusing closed histories", async () => { + const first = deferred(); + const second = deferred(); + const requests = gateway([first.promise, second.promise, answer("Third answer."), answer("Fresh answer.")]); + renderWidget(); + await openWidget(); + send("First pending question"); + await waitFor(() => + expect(requests.filter((request) => request.url.endsWith("/chat/completions"))).toHaveLength(1), + ); + fireEvent.click(screen.getByRole("button", { name: "New chat" })); + await selectModel(); + send("Second pending question"); + await waitFor(() => + expect(requests.filter((request) => request.url.endsWith("/chat/completions"))).toHaveLength(2), + ); + const pending = requests.filter((request) => request.url.endsWith("/chat/completions")); + fireEvent.click(screen.getByRole("button", { name: "Close Chat 1" })); + await waitFor(() => expect(pending[0].signal.aborted).toBe(true)); + expect(pending[1].signal.aborted).toBe(false); + expect(screen.getByRole("tab", { name: "Chat 2" })).toHaveAttribute("aria-selected", "true"); + await act(async () => first.resolve(toolReply("key_create", keyArguments))); + expect(screen.queryByRole("button", { name: "Confirm change" })).not.toBeInTheDocument(); + expect(screen.getByRole("status", { name: "LiteAdmin is working" })).toBeInTheDocument(); + + fireEvent.click(screen.getByRole("button", { name: "New chat" })); + await selectModel(); + send("Third question"); + await screen.findByText("Third answer."); + fireEvent.click(screen.getByRole("button", { name: "Close Chat 3" })); + expect(screen.getByRole("tab", { name: "Chat 2" })).toHaveAttribute("aria-selected", "true"); + expect(screen.getByText("Second pending question")).toBeInTheDocument(); + expect(pending[1].signal.aborted).toBe(false); + fireEvent.click(screen.getByRole("button", { name: "Close Chat 2" })); + await waitFor(() => expect(pending[1].signal.aborted).toBe(true)); + expect(within(screen.getByRole("tablist", { name: "LiteAdmin chats" })).getAllByRole("tab")).toHaveLength(1); + expect(screen.getByRole("tab", { name: "Chat 4" })).toHaveAttribute("aria-selected", "true"); + expect(screen.getByPlaceholderText("Ask LiteAdmin…")).toHaveValue(""); + expect(screen.queryByRole("button", { name: "Stop request" })).not.toBeInTheDocument(); + await act(async () => second.resolve(answer("Late closed answer."))); + expect(screen.queryByText("Late closed answer.")).not.toBeInTheDocument(); + expect(screen.queryByText("Third answer.")).not.toBeInTheDocument(); + await selectModel(); + send("Fresh question"); + await screen.findByText("Fresh answer."); + expect(requests.filter((request) => request.url.endsWith("/chat/completions")).at(-1)?.body.messages).toEqual([ + { role: "system", content: expect.any(String) }, + { role: "user", content: "Fresh question" }, + ]); + expect(requests.filter((request) => request.url.endsWith("/key/generate"))).toHaveLength(0); + }); + + it("keeps Home drafts unsent without a model or with empty or oversized input", async () => { + const options: GatewayOptions = { models: [] }; + const requests = gateway([], options); + const { client } = renderWidget(undefined, true); + const home = await screen.findByRole("region", { name: "Ask LiteAdmin" }); + expect(await within(home).findByText("Add a chat model to your gateway to use LiteAdmin.")).toBeInTheDocument(); + expect(within(home).getByPlaceholderText("Ask LiteAdmin…")).toBeDisabled(); + expect(within(home).getByRole("button", { name: "Send message" })).toBeDisabled(); + options.models = [{ model_group: "chat-model", mode: "chat" }]; + await act(async () => client.invalidateQueries({ queryKey: ["liteadmin-models"] })); + await selectModel(); + for (const text of [" ", "x".repeat(MAX_INPUT_LENGTH + 1)]) { + const input = within(home).getByPlaceholderText("Ask LiteAdmin…"); + fireEvent.change(input, { target: { value: text } }); + fireEvent.keyDown(input, { key: "Enter" }); + const submit = within(home).getByRole("button", { name: "Send message" }); + expect(submit).toBeDisabled(); + fireEvent.click(submit); + expect(input).toHaveValue(text); + } + expect(home).toBeVisible(); + expect(within(home).getByRole("alert")).toHaveTextContent("8,000 characters"); + expect(screen.queryByRole("complementary", { name: "LiteAdmin" })).not.toBeInTheDocument(); + expect(requests.filter((request) => request.url.endsWith("/chat/completions"))).toHaveLength(0); + }); + + it("shows destination consent in Home and discards its draft when the destination changes", async () => { + const settings = { target: EXTERNAL_INFERENCE, status: 200 }; + const settingsReady = deferred(); + const requests = gateway([], { + settings: async () => { + await settingsReady.promise; + return json({ PROXY_BASE_URL: MANAGEMENT, LITELLM_UI_API_DOC_BASE_URL: settings.target }); + }, + }); + const { client } = renderWidget(undefined, true); + const home = await screen.findByRole("region", { name: "Ask LiteAdmin" }); + expect(await within(home).findByLabelText("Loading gateway settings")).toBeInTheDocument(); + await act(async () => settingsReady.resolve()); + const consent = await within(home).findByRole("button", { name: "Use configured gateway" }); + expect(within(home).getByText(EXTERNAL_INFERENCE)).toBeInTheDocument(); + expect(within(home).queryByPlaceholderText("Ask LiteAdmin…")).not.toBeInTheDocument(); + fireEvent.click(consent); + await selectModel(); + fireEvent.change(within(home).getByPlaceholderText("Ask LiteAdmin…"), { target: { value: "Old gateway draft" } }); + settings.target = `${EXTERNAL_INFERENCE}/changed`; + await act(async () => client.invalidateQueries({ queryKey: ["proxySettings"] })); + fireEvent.click(await within(home).findByRole("button", { name: "Use configured gateway" })); + await selectModel(); + expect(within(home).getByPlaceholderText("Ask LiteAdmin…")).toHaveValue(""); + expect(home).toBeVisible(); + expect(requests.some((request) => request.url.endsWith("/chat/completions"))).toBe(false); + }); + it.each([ ["sidebar", SidebarAccountMenu], ["navbar", UserDropdown], ] as const)("persists Hide LiteAdmin only for enterprise users in the %s account menu", async (_name, Menu) => { gateway([]); const user = userEvent.setup(); - const view = renderWidget(Menu); + const view = renderWidget(Menu, true); await screen.findByRole("button", { name: "LiteAdmin" }); await user.click(screen.getByRole("button", { name: /account menu/i })); const toggle = await screen.findByRole("switch", { name: "Toggle hide LiteAdmin" }); @@ -216,9 +593,10 @@ describe("LiteAdmin in the gateway", () => { await user.click(toggle); expect(toggle).toBeChecked(); expect(screen.queryByRole("button", { name: "LiteAdmin" })).not.toBeInTheDocument(); + expect(screen.queryByRole("region", { name: "Ask LiteAdmin", hidden: true })).not.toBeInTheDocument(); view.unmount(); - const restored = renderWidget(Menu); + const restored = renderWidget(Menu, true); await screen.findByText("Session ready"); await waitFor(() => expect(restored.client.isFetching()).toBe(0)); expect(screen.queryByRole("button", { name: "LiteAdmin" })).not.toBeInTheDocument(); @@ -227,6 +605,7 @@ describe("LiteAdmin in the gateway", () => { expect(savedToggle).toBeChecked(); await user.click(savedToggle); expect(await screen.findByRole("button", { name: "LiteAdmin" })).toBeInTheDocument(); + expect(await screen.findByRole("region", { name: "Ask LiteAdmin" })).toBeVisible(); for (const premiumUser of [false, null, undefined, "true"]) { session("proxy_admin", "first-admin", { premium_user: premiumUser }); @@ -241,10 +620,11 @@ describe("LiteAdmin in the gateway", () => { async (premiumUser) => { session("proxy_admin", "first-admin", { premium_user: premiumUser }); const requests = gateway([]); - const { client } = renderWidget(); + const { client } = renderWidget(undefined, true); await screen.findByText("Session ready"); await waitFor(() => expect(client.isFetching()).toBe(0)); expect(screen.queryByRole("button", { name: "LiteAdmin" })).not.toBeInTheDocument(); + expect(screen.queryByRole("region", { name: "Ask LiteAdmin", hidden: true })).not.toBeInTheDocument(); for (const modifiers of [{ metaKey: true }, { ctrlKey: true }]) { expect(fireEvent.keyDown(document, { key: "j", ...modifiers })).toBe(true); expect(screen.queryByRole("complementary", { name: "LiteAdmin" })).not.toBeInTheDocument(); @@ -387,10 +767,11 @@ describe("LiteAdmin in the gateway", () => { async (role) => { session(role); const requests = gateway([]); - const { client } = renderWidget(); + const { client } = renderWidget(undefined, true); await screen.findByText("Session ready"); await waitFor(() => expect(client.isFetching()).toBe(0)); expect(screen.queryByRole("button", { name: "LiteAdmin" })).not.toBeInTheDocument(); + expect(screen.queryByRole("region", { name: "Ask LiteAdmin", hidden: true })).not.toBeInTheDocument(); expect(requests.every((request) => request.url.endsWith("/litellm-ui-config"))).toBe(true); }, ); @@ -490,7 +871,8 @@ describe("LiteAdmin in the gateway", () => { const confirm = await screen.findByRole("button", { name: "Confirm change" }); fireEvent.click(confirm); fireEvent.click(confirm); - expect(screen.getByRole("button", { name: "New chat" })).toBeDisabled(); + expect(screen.getByRole("button", { name: "Close Chat 1" })).toBeDisabled(); + expect(screen.getByRole("button", { name: "New chat" })).toBeEnabled(); expect(screen.queryByRole("button", { name: "Stop request" })).not.toBeInTheDocument(); fireEvent.click(screen.getByRole("button", { name: "Close LiteAdmin" })); await act(async () => response.resolve(json({ key: NEW_KEY }))); @@ -573,25 +955,73 @@ describe("LiteAdmin in the gateway", () => { expect(screen.queryByText(/A submitted change may have completed/)).not.toBeInTheDocument(); }); - it.each(["account", "target"])("discards an unconfirmed review when the %s changes", async (transition) => { - const settings = { target: INFERENCE, status: 200 }; - const requests = gateway([toolReply("key_create", keyArguments)], { settings }); - const view = renderWidget(); - await openWidget(); - send("Old request"); - await screen.findByRole("button", { name: "Confirm change" }); - if (transition === "account") { - session("proxy_admin", "second-admin"); - view.refresh(); - } else { - settings.target = `${INFERENCE}/updated`; - await act(async () => view.client.invalidateQueries({ queryKey: ["proxySettings"] })); - } - await waitFor(() => expect(screen.queryByRole("button", { name: "Confirm change" })).not.toBeInTheDocument()); - expect(screen.queryByText("Old request")).not.toBeInTheDocument(); - expect(screen.queryByText(/A submitted change may have completed/)).not.toBeInTheDocument(); - expect(requests.filter((request) => request.url.endsWith("/key/generate"))).toHaveLength(0); - }); + it.each(["account", "gateway", "destination", "permission"])( + "resets every chat when the %s changes", + async (transition) => { + const pending = deferred(); + const settings = { target: INFERENCE, status: 200 }; + const requests = gateway([pending.promise, toolReply("key_create", keyArguments), answer("Fresh session.")], { + settings, + }); + const view = renderWidget(undefined, true); + await openWidget(); + send("Old request"); + await waitFor(() => + expect(requests.filter((request) => request.url.endsWith("/chat/completions"))).toHaveLength(1), + ); + const original = requests.find((request) => request.url.endsWith("/chat/completions")); + fireEvent.click(screen.getByRole("button", { name: "New chat" })); + await selectModel(); + send("Old approval"); + await screen.findByRole("button", { name: "Confirm change" }); + fireEvent.click(screen.getByRole("button", { name: "New chat" })); + await selectModel(); + fireEvent.change(screen.getByPlaceholderText("Ask LiteAdmin…"), { target: { value: "Old draft" } }); + if (transition === "account") { + session("proxy_admin", "second-admin"); + view.refresh(); + } else if (transition === "gateway") { + settings.target = "https://worker.test/inference"; + switchToWorkerUrl("https://worker.test/proxy"); + view.refresh(); + } else if (transition === "destination") { + settings.target = `${INFERENCE}/updated`; + await act(async () => view.client.invalidateQueries({ queryKey: ["proxySettings"] })); + } else { + session("proxy_admin_viewer"); + view.refresh(); + expect(screen.queryByRole("button", { name: "LiteAdmin" })).not.toBeInTheDocument(); + expect(screen.queryByRole("region", { name: "Ask LiteAdmin", hidden: true })).not.toBeInTheDocument(); + expect(screen.queryByRole("tablist", { name: "LiteAdmin chats", hidden: true })).not.toBeInTheDocument(); + fireEvent.keyDown(document, { key: "j", ctrlKey: true }); + expect(screen.queryByRole("complementary", { name: "LiteAdmin" })).not.toBeInTheDocument(); + session(); + view.refresh(); + } + await waitFor(() => expect(original?.signal.aborted).toBe(true)); + const trigger = await screen.findByRole("button", { name: "LiteAdmin" }); + if (trigger.getAttribute("aria-expanded") === "false") fireEvent.click(trigger); + const tabs = await screen.findByRole("tablist", { name: "LiteAdmin chats" }); + expect(within(tabs).getAllByRole("tab")).toHaveLength(1); + expect(screen.getByRole("tab", { name: "Chat 1" })).toHaveAttribute("aria-selected", "true"); + await waitFor(() => expect(screen.queryByRole("button", { name: "Confirm change" })).not.toBeInTheDocument()); + expect(screen.queryByText("Old request")).not.toBeInTheDocument(); + expect(screen.queryByText("Old approval")).not.toBeInTheDocument(); + expect(screen.getByPlaceholderText("Ask LiteAdmin…")).toHaveValue(""); + expect(screen.getByRole("combobox", { name: "LiteAdmin model" })).toHaveValue(""); + await act(async () => pending.resolve(toolReply("key_create", keyArguments))); + expect(screen.queryByRole("button", { name: "Confirm change" })).not.toBeInTheDocument(); + expect(screen.queryByText(/A submitted change may have completed/)).not.toBeInTheDocument(); + expect(requests.filter((request) => request.url.endsWith("/key/generate"))).toHaveLength(0); + await selectModel(); + send("Fresh question"); + await screen.findByText("Fresh session."); + expect(requests.filter((request) => request.url.endsWith("/chat/completions")).at(-1)?.body.messages).toEqual([ + { role: "system", content: expect.any(String) }, + { role: "user", content: "Fresh question" }, + ]); + }, + ); it.each(["account", "target"])( "warns about an interrupted write when the %s changes and rejects its late key", @@ -641,9 +1071,9 @@ describe("LiteAdmin in the gateway", () => { it("blocks inference when settings cannot be loaded", async () => { const requests = gateway([], { settings: { target: INFERENCE, status: 503 } }); - renderWidget(); - fireEvent.click(await screen.findByRole("button", { name: "LiteAdmin" })); - expect(await screen.findByText("Could not load gateway settings.")).toBeInTheDocument(); + renderWidget(undefined, true); + const home = await screen.findByRole("region", { name: "Ask LiteAdmin" }); + expect(await within(home).findByText("Could not load gateway settings.")).toBeInTheDocument(); expect(requests.some((request) => request.url.endsWith("/chat/completions"))).toBe(false); }); diff --git a/ui/litellm-dashboard/src/components/liteadmin/LiteAdmin.tsx b/ui/litellm-dashboard/src/components/liteadmin/LiteAdmin.tsx index 38e4334763f..bf296d47746 100644 --- a/ui/litellm-dashboard/src/components/liteadmin/LiteAdmin.tsx +++ b/ui/litellm-dashboard/src/components/liteadmin/LiteAdmin.tsx @@ -1,8 +1,19 @@ "use client"; -import { createContext, useContext, useEffect, useRef, useState, type ReactNode, type RefObject } from "react"; -import { useQuery } from "@tanstack/react-query"; -import { RotateCcw, Sparkles, X } from "lucide-react"; +import { + createContext, + useCallback, + useContext, + useEffect, + useLayoutEffect, + useRef, + useState, + type ReactNode, + type RefObject, +} from "react"; +import { createPortal } from "react-dom"; +import { useQuery, type UseQueryResult } from "@tanstack/react-query"; +import { Plus, Sparkles, X } from "lucide-react"; import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; import { useDisableLiteAdmin } from "@/app/(dashboard)/hooks/useDisableLiteAdmin"; import { useProxySettingsQuery } from "@/app/(dashboard)/hooks/proxySettings/useProxySettings"; @@ -16,14 +27,23 @@ import { Button } from "@/components/ui/button"; import { Card, CardContent, CardFooter, CardHeader, CardTitle } from "@/components/ui/card"; import { FieldError } from "@/components/ui/field"; import { Skeleton } from "@/components/ui/skeleton"; +import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs"; import { cn } from "@/lib/cva.config"; import { isProxyAdminRole } from "@/utils/roles"; import { getPreferredLiteAdminModel, MAX_INPUT_LENGTH, resolveInferenceTarget } from "./agent"; import { LiteAdminConversation } from "./LiteAdminConversation"; -import { useLiteAdmin, type LiteAdminSession } from "./useLiteAdmin"; +import { useLiteAdmin, type ConversationEntry, type LiteAdminSession } from "./useLiteAdmin"; type ManagementSession = Omit; -type LiteAdminState = { open: boolean; toggle: () => void }; +type ChatPhase = ReturnType["phase"]; +type ChatTab = { id: number; phase: ChatPhase }; +type LiteAdminState = { + open: boolean; + toggle: () => void; + show: () => void; + homeContainer: HTMLDivElement | null; + setHomeContainer: (container: HTMLDivElement | null) => void; +}; const LiteAdminContext = createContext(null); @@ -43,12 +63,16 @@ export function LiteAdminFrame({ children }: { children: ReactNode }) { const configured = useLiteAdminSession(); const sessionKey = configured?.key ?? null; const [open, setOpen] = useState(false); + const [homeContainer, setHomeContainer] = useState(null); const [openedSession, setOpenedSession] = useState(sessionKey); if (openedSession !== sessionKey) { setOpenedSession(sessionKey); setOpen(false); } const toggle = () => setOpen((current) => !current); + const show = () => { + if (!open) setOpen(true); + }; useEffect(() => { if (sessionKey === null) return; const onKeyDown = (event: KeyboardEvent) => { @@ -62,7 +86,7 @@ export function LiteAdminFrame({ children }: { children: ReactNode }) { return () => document.removeEventListener("keydown", onKeyDown); }, [sessionKey]); return ( - + {children} {configured && ( setOpen(false)} /> @@ -71,6 +95,29 @@ export function LiteAdminFrame({ children }: { children: ReactNode }) { ); } +export function LiteAdminHome() { + const state = useContext(LiteAdminContext); + if (!state) return null; + const { open, show, setHomeContainer } = state; + return ( + + ); +} + export default function LiteAdminTrigger() { const state = useContext(LiteAdminContext); if (!state) return null; @@ -129,39 +176,43 @@ function Destination({ }) { const [approved, setApproved] = useState(false); if (target?.baseUrl && (!target.requiresConsent || approved)) { - return ; + return ; } return ( - + + {loading && } + {!loading && !target?.baseUrl && ( + + {target?.error || "Could not load gateway settings."} + + + )} + {!loading && target?.baseUrl && ( + + + Connect to the configured gateway + + +

{target.baseUrl}

+

+ This gateway uses a different address. LiteAdmin will send your existing session credential to the + address shown above. +

+ + + + + + )} +
+ } + > -
- {loading && } - {!loading && !target?.baseUrl && ( - - {target?.error || "Could not load gateway settings."} - - - )} - {!loading && target?.baseUrl && ( - - - Connect to the configured gateway - - -

{target.baseUrl}

-

- This gateway uses a different address. LiteAdmin will send your existing session credential to the - address shown above. -

- - - - - - )} -
); } @@ -170,12 +221,16 @@ function Destination({ function Panel({ open, initialFocus, + composer, children, }: { open: boolean; initialFocus?: RefObject; + composer: ReactNode; children: ReactNode; }) { + const state = useContext(LiteAdminContext); + const homeContainer = !open ? state?.homeContainer : null; const ref = useRef(null); useEffect(() => { const panel = ref.current; @@ -196,19 +251,154 @@ function Panel({ )} > {children} + {homeContainer ? createPortal(composer, homeContainer) : composer} ); } -function LiteAdminChat({ session, open, close }: { session: LiteAdminSession; open: boolean; close: () => void }) { +function LiteAdminChats({ session, open, close }: { session: LiteAdminSession; open: boolean; close: () => void }) { + const [tabs, setTabs] = useState<{ chats: ChatTab[]; activeId: number; nextId: number }>({ + chats: [{ id: 1, phase: "idle" }], + activeId: 1, + nextId: 2, + }); + const reviewRef = useRef(null); + const activeTabRef = useRef(null); + const focusAfterClose = useRef(false); + useLayoutEffect(() => { + if (!focusAfterClose.current) return; + focusAfterClose.current = false; + activeTabRef.current?.focus(); + }); + const addChat = () => { + setTabs((current) => ({ + chats: [...current.chats, { id: current.nextId, phase: "idle" }], + activeId: current.nextId, + nextId: current.nextId + 1, + })); + }; + const closeChat = (id: number) => { + if (tabs.chats.some((chat) => chat.id === id && chat.phase === "applying")) return; + focusAfterClose.current = true; + setTabs((current) => { + const closing = current.chats.find((chat) => chat.id === id); + if (!closing || closing.phase === "applying") return current; + const remaining = current.chats.filter((chat) => chat.id !== id); + const next = remaining.find((chat) => chat.id > id) ?? remaining.at(-1); + if (!next) { + return { + chats: [{ id: current.nextId, phase: "idle" }], + activeId: current.nextId, + nextId: current.nextId + 1, + }; + } + return { ...current, chats: remaining, activeId: current.activeId === id ? next.id : current.activeId }; + }); + }; + const updatePhase = useCallback((id: number, phase: ChatPhase) => { + setTabs((current) => { + if (!current.chats.some((chat) => chat.id === id && chat.phase !== phase)) return current; + return { ...current, chats: current.chats.map((chat) => (chat.id === id ? { ...chat, phase } : chat)) }; + }); + }, []); + return ( + { + if (typeof id === "number") setTabs((current) => ({ ...current, activeId: id })); + }} + className="contents" + > + + + + + + {tabs.chats.map((tab) => ( +
+ + Chat {tab.id} + {tab.phase !== "idle" && } + + +
+ ))} + + {tabs.chats.map((tab) => ( + + ))} + + + ); +} + +function useForegroundReviewFocus( + entries: ConversationEntry[], + foreground: boolean, + reviewRef: RefObject, +) { + const lastReviewId = useRef(null); + const pendingReviewId = entries.find((entry) => entry.kind === "action" && entry.status === "review")?.id ?? null; + useEffect(() => { + if (lastReviewId.current === pendingReviewId) return; + lastReviewId.current = pendingReviewId; + if (foreground && pendingReviewId) reviewRef.current?.focus({ preventScroll: true }); + }, [foreground, pendingReviewId, reviewRef]); +} + +function LiteAdminChat({ + id, + session, + active, + open, + reviewRef, + onPhaseChange, +}: { + id: number; + session: LiteAdminSession; + active: boolean; + open: boolean; + reviewRef: RefObject; + onPhaseChange: (id: number, phase: ChatPhase) => void; +}) { + const state = useContext(LiteAdminContext); const chat = useLiteAdmin(session); const [input, setInput] = useState(""); const [model, setModel] = useState(null); - const reviewRef = useRef(null); + useForegroundReviewFocus(chat.entries, active && open, reviewRef); + useEffect(() => onPhaseChange(id, chat.phase), [id, chat.phase, onPhaseChange]); const modelQuery = { queryKey: ["liteadmin-models", session.accessToken, session.managementBaseUrl], queryFn: () => fetchAvailableModels(session.accessToken), - enabled: open, + enabled: active && (open || Boolean(state?.homeContainer)), select: (available: ModelGroup[]) => available.filter((item) => isModeCompatibleWithEndpoint(item.mode, EndpointType.CHAT)), }; @@ -221,26 +411,31 @@ function LiteAdminChat({ session, open, close }: { session: LiteAdminSession; op const submitDisabled = busy || models.isError || !hasValidInput; const send = () => { if (submitDisabled || !selectedModel) return; + state?.show(); void chat.send(input, selectedModel); setInput(""); }; + if (!active) return null; + const homeContainer = !open ? state?.homeContainer : null; + const composer = ( +
+ + {tooLong && Keep your message within {MAX_INPUT_LENGTH.toLocaleString()} characters.} + +

Use a model you trust with your gateway data.

+
+ ); return ( - - - - + -
- {models.isPending ? ( - - ) : ( - ({ value: item.model_group, label: item.model_group }))} - value={selectedModel} - onValueChange={setModel} - placeholder="Choose a chat model" - disabled={busy} - allowClear={false} - /> - )} - {models.isError && ( - - - Could not load models.{" "} - - - - )} - {models.isSuccess && models.data.length === 0 && ( - - Add a chat model to your gateway to use LiteAdmin. - - )} - {tooLong && Keep your message within {MAX_INPUT_LENGTH.toLocaleString()} characters.} - + ); +} + +function ModelPicker({ + models, + value, + onChange, + disabled, +}: { + models: UseQueryResult; + value: string | null; + onChange: (value: string | null) => void; + disabled: boolean; +}) { + return ( + <> + {models.isPending ? ( + + ) : ( + ({ value: item.model_group, label: item.model_group }))} + value={value} + onValueChange={onChange} + placeholder="Choose a chat model" + disabled={disabled} + allowClear={false} /> -

Use a model you trust with your gateway data.

-
- + )} + {models.isError && ( + + + Could not load models.{" "} + + + + )} + {models.isSuccess && models.data.length === 0 && ( + + Add a chat model to your gateway to use LiteAdmin. + + )} + ); } diff --git a/ui/litellm-dashboard/src/components/liteadmin/LiteAdminConversation.tsx b/ui/litellm-dashboard/src/components/liteadmin/LiteAdminConversation.tsx index 826db57a736..014597dba39 100644 --- a/ui/litellm-dashboard/src/components/liteadmin/LiteAdminConversation.tsx +++ b/ui/litellm-dashboard/src/components/liteadmin/LiteAdminConversation.tsx @@ -99,8 +99,7 @@ function ActionCard({ useEffect(() => { if (!open || !review) return; scrollToEnd(); - reviewRef.current?.focus({ preventScroll: true }); - }, [entry.id, open, review, reviewRef, scrollToEnd]); + }, [entry.id, open, review, scrollToEnd]); const answer = (approved: boolean) => { scrollToEnd(); onAnswer(entry.id, approved); diff --git a/ui/litellm-dashboard/src/components/model_filters.tsx b/ui/litellm-dashboard/src/components/model_filters.tsx index 0ce82d6d050..1eef5083f16 100644 --- a/ui/litellm-dashboard/src/components/model_filters.tsx +++ b/ui/litellm-dashboard/src/components/model_filters.tsx @@ -1,5 +1,6 @@ import React, { useState, useEffect, useMemo, useRef } from "react"; import { Card } from "@/components/ui/card"; +import { displayMode } from "@/lib/decisionModels"; import { matchesSearchTerm } from "@/utils/searchUtils"; interface ModelGroupInfo { @@ -51,7 +52,8 @@ const ModelFilters: React.FC = ({ const getUniqueModes = (data: ModelGroupInfo[]) => { const modes = new Set(); data.forEach((model) => { - if (model.mode) modes.add(model.mode); + const mode = displayMode(model.mode); + if (mode) modes.add(mode); }); return Array.from(modes); }; @@ -79,7 +81,7 @@ const ModelFilters: React.FC = ({ modelHubData?.filter((model) => { const matchesSearch = matchesSearchTerm(searchTerm, [model.model_group]); const matchesProvider = selectedProvider === "" || model.providers.includes(selectedProvider); - const matchesMode = selectedMode === "" || model.mode === selectedMode; + const matchesMode = selectedMode === "" || displayMode(model.mode) === selectedMode; // Check if model has the selected feature const matchesFeature = diff --git a/ui/litellm-dashboard/src/components/model_group_alias_settings.test.tsx b/ui/litellm-dashboard/src/components/model_group_alias_settings.test.tsx new file mode 100644 index 00000000000..fe59e6eb139 --- /dev/null +++ b/ui/litellm-dashboard/src/components/model_group_alias_settings.test.tsx @@ -0,0 +1,66 @@ +/* @vitest-environment jsdom */ +import { render, screen } from "@testing-library/react"; +import userEvent from "@testing-library/user-event"; +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; +import { setCallbacksCall } from "@/components/networking"; +import ModelGroupAliasSettings from "./model_group_alias_settings"; + +vi.mock("@/components/networking", () => ({ + setCallbacksCall: vi.fn(), +})); + +const aliases = { + "fast-model": "claude-haiku", + "smart-model": { model: "claude-sonnet", hidden: true }, +}; + +describe("ModelGroupAliasSettings", () => { + beforeEach(() => { + vi.clearAllMocks(); + vi.mocked(setCallbacksCall).mockResolvedValue({}); + }); + + afterEach(() => { + vi.restoreAllMocks(); + }); + + it("locks config-owned aliases while keeping their names and targets visible", async () => { + render(); + + expect( + screen.getByText( + "These aliases are defined in config.yaml and are read-only here. Edit the config file to change them.", + ), + ).toBeInTheDocument(); + expect(screen.getByPlaceholderText("e.g., gpt-4o")).toBeDisabled(); + expect(screen.getByPlaceholderText("e.g., gpt-4o-mini-openai")).toBeDisabled(); + expect(screen.getByRole("button", { name: "Add Alias" })).toBeDisabled(); + expect(screen.getByRole("button", { name: "Edit alias fast-model" })).toBeDisabled(); + expect(screen.getByRole("button", { name: "Delete alias fast-model" })).toBeDisabled(); + expect(screen.getByRole("button", { name: "Edit alias smart-model" })).toBeDisabled(); + expect(screen.getByRole("button", { name: "Delete alias smart-model" })).toBeDisabled(); + + expect(await screen.findByText("fast-model")).toBeInTheDocument(); + expect(screen.getByText("claude-haiku")).toBeInTheDocument(); + expect(screen.getByText("smart-model")).toBeInTheDocument(); + expect(screen.getByText("claude-sonnet")).toBeInTheDocument(); + + await userEvent.setup().click(screen.getByRole("button", { name: "Delete alias fast-model" })); + expect(setCallbacksCall).not.toHaveBeenCalled(); + }); + + it("keeps aliases editable by default and saves the current string representation after delete", async () => { + const user = userEvent.setup(); + render(); + + expect(await screen.findByText("fast-model")).toBeInTheDocument(); + expect(screen.getByRole("button", { name: "Edit alias fast-model" })).toBeEnabled(); + expect(screen.getByRole("button", { name: "Delete alias fast-model" })).toBeEnabled(); + + await user.click(screen.getByRole("button", { name: "Delete alias fast-model" })); + + expect(setCallbacksCall).toHaveBeenCalledWith("test-token", { + router_settings: { model_group_alias: { "smart-model": "claude-sonnet" } }, + }); + }); +}); diff --git a/ui/litellm-dashboard/src/components/model_group_alias_settings.tsx b/ui/litellm-dashboard/src/components/model_group_alias_settings.tsx index 1e943da25c5..d3cebc76e98 100644 --- a/ui/litellm-dashboard/src/components/model_group_alias_settings.tsx +++ b/ui/litellm-dashboard/src/components/model_group_alias_settings.tsx @@ -11,6 +11,8 @@ interface ModelGroupAliasSettingsProps { accessToken: string; initialModelGroupAlias?: Record; onAliasUpdate?: (updatedAlias: { [key: string]: string }) => void; + managedByConfig?: boolean; + ownershipLoading?: boolean; } interface AliasItem { @@ -23,10 +25,13 @@ const ModelGroupAliasSettings: React.FC = ({ accessToken, initialModelGroupAlias = {}, onAliasUpdate, + managedByConfig = false, + ownershipLoading = false, }) => { const [aliases, setAliases] = useState([]); const [newAlias, setNewAlias] = useState({ aliasName: "", targetModelGroup: "" }); const [editingAlias, setEditingAlias] = useState(null); + const writesDisabled = managedByConfig || ownershipLoading; const [isExpanded, setIsExpanded] = useState(true); useEffect(() => { @@ -73,6 +78,8 @@ const ModelGroupAliasSettings: React.FC = ({ }; const handleAddAlias = async () => { + if (writesDisabled) return; + if (!newAlias.aliasName || !newAlias.targetModelGroup) { toast.fromError("Please provide both alias name and target model group"); return; @@ -100,10 +107,13 @@ const ModelGroupAliasSettings: React.FC = ({ }; const handleEditAlias = (alias: AliasItem) => { + if (writesDisabled) return; + setEditingAlias({ ...alias }); }; const handleUpdateAlias = async () => { + if (writesDisabled) return; if (!editingAlias) return; if (!editingAlias.aliasName || !editingAlias.targetModelGroup) { @@ -131,6 +141,8 @@ const ModelGroupAliasSettings: React.FC = ({ }; const deleteAlias = async (aliasId: string) => { + if (writesDisabled) return; + const updatedAliases = aliases.filter((alias) => alias.id !== aliasId); if (await saveAliasesToBackend(updatedAliases)) { @@ -170,6 +182,11 @@ const ModelGroupAliasSettings: React.FC = ({ {isExpanded && (
+ {managedByConfig && ( +

+ These aliases are defined in config.yaml and are read-only here. Edit the config file to change them. +

+ )}

Add New Alias

@@ -177,6 +194,7 @@ const ModelGroupAliasSettings: React.FC = ({ setNewAlias({ ...newAlias, @@ -184,7 +202,7 @@ const ModelGroupAliasSettings: React.FC = ({ }) } placeholder="e.g., gpt-4o" - className="w-full px-3 py-2 border border-border rounded-md text-sm" + className="w-full px-3 py-2 border border-border rounded-md text-sm disabled:bg-muted disabled:text-muted-foreground disabled:cursor-not-allowed" />
@@ -192,6 +210,7 @@ const ModelGroupAliasSettings: React.FC = ({ setNewAlias({ ...newAlias, @@ -199,14 +218,14 @@ const ModelGroupAliasSettings: React.FC = ({ }) } placeholder="e.g., gpt-4o-mini-openai" - className="w-full px-3 py-2 border border-border rounded-md text-sm" + className="w-full px-3 py-2 border border-border rounded-md text-sm disabled:bg-muted disabled:text-muted-foreground disabled:cursor-not-allowed" />
diff --git a/ui/litellm-dashboard/src/components/networking.tsx b/ui/litellm-dashboard/src/components/networking.tsx index 8b8a84aba7b..d1e0fa81b06 100644 --- a/ui/litellm-dashboard/src/components/networking.tsx +++ b/ui/litellm-dashboard/src/components/networking.tsx @@ -4782,6 +4782,7 @@ export interface ConnectFlowStatus { server_id: string | null; server_name: string | null; connected: boolean | null; + requested_scopes?: string[]; } /** diff --git a/ui/litellm-dashboard/src/components/provider_info_helpers.test.tsx b/ui/litellm-dashboard/src/components/provider_info_helpers.test.tsx index 52e13642020..b27be2ba966 100644 --- a/ui/litellm-dashboard/src/components/provider_info_helpers.test.tsx +++ b/ui/litellm-dashboard/src/components/provider_info_helpers.test.tsx @@ -201,6 +201,7 @@ describe("provider_info_helpers", () => { Providers.CLARIFAI, Providers.Cognition, Providers.COMPACTIFAI, + Providers.CORALBRICKS, Providers.DATAROBOT, Providers.DOCKER_MODEL_RUNNER, Providers.DOTPROMPT, @@ -361,6 +362,10 @@ describe("provider_info_helpers", () => { expect(getPlaceholder(Providers.Tencent)).toBe("tencent/deepseek-v4-pro"); }); + it("should return a coralbricks/ placeholder for the CoralBricks provider", () => { + expect(getPlaceholder(Providers.CORALBRICKS)).toBe("coralbricks/glm-5.3-fast"); + }); + it("should return decision model placeholders for the TypeSafe and StrandsDecider dropdown keys", () => { expect(getPlaceholder("TypeSafe")).toBe("typesafe/jev-latest"); expect(getPlaceholder("StrandsDecider")).toBe("strands_decider/strands-decider-2B-hobson-v19"); diff --git a/ui/litellm-dashboard/src/components/provider_info_helpers.tsx b/ui/litellm-dashboard/src/components/provider_info_helpers.tsx index fa774d46f0f..7b6c7d24706 100644 --- a/ui/litellm-dashboard/src/components/provider_info_helpers.tsx +++ b/ui/litellm-dashboard/src/components/provider_info_helpers.tsx @@ -99,6 +99,7 @@ export enum Providers { COHERE_CHAT = "Cohere Chat", COMETAPI = "Cometapi", COMPACTIFAI = "Compactifai", + CORALBRICKS = "CoralBricks", Cursor = "Cursor", Dashscope = "Dashscope", Databricks = "Databricks (Qwen API)", @@ -221,6 +222,7 @@ export const provider_map: Record = { COHERE_CHAT: "cohere_chat", COMETAPI: "cometapi", COMPACTIFAI: "compactifai", + CORALBRICKS: "coralbricks", Cursor: "cursor", Dashscope: "dashscope", Databricks: "databricks", @@ -459,6 +461,7 @@ const providerPlaceholderMap: Partial> = { [Providers.Bedrock]: "claude-3-opus", [Providers.CHATGPT]: "chatgpt/gpt-5.4", [Providers.Cognition]: "cognition/swe-1.7", + [Providers.CORALBRICKS]: "coralbricks/glm-5.3-fast", [Providers.Cursor]: "cursor/claude-4-sonnet", [Providers.DeepInfra]: "deepinfra/", [Providers.EDENAI]: "edenai/openai/gpt-mini-latest", diff --git a/ui/litellm-dashboard/src/components/router_settings/ConfigOwned.tsx b/ui/litellm-dashboard/src/components/router_settings/ConfigOwned.tsx new file mode 100644 index 00000000000..4587b62c2ce --- /dev/null +++ b/ui/litellm-dashboard/src/components/router_settings/ConfigOwned.tsx @@ -0,0 +1,28 @@ +import { Badge } from "@/components/ui/badge"; +import { Tooltip, TooltipContent, TooltipTrigger } from "@/components/ui/tooltip"; +import { Lock } from "lucide-react"; +import React from "react"; + +export type RouterSettingSource = "config" | "db" | "env" | "default" | "unset"; +export type RouterSettingsSource = { [key: string]: RouterSettingSource }; + +export const CONFIG_OWNED_HINT = "Set in config.yaml. Edit the file to change it."; +export const SOURCE_UNKNOWN_HINT = "Could not load which settings config.yaml owns. Reload the page to make changes."; + +export function isOwnedByConfig(source: RouterSettingsSource, key: string): boolean { + return source[key] === "config"; +} + +export function ConfigOwnedBadge(): React.ReactElement { + return ( + + } + > + + config.yaml + + {CONFIG_OWNED_HINT} + + ); +} diff --git a/ui/litellm-dashboard/src/components/router_settings/LatencyBasedConfiguration.tsx b/ui/litellm-dashboard/src/components/router_settings/LatencyBasedConfiguration.tsx index 61f3abbbb0f..1212a732697 100644 --- a/ui/litellm-dashboard/src/components/router_settings/LatencyBasedConfiguration.tsx +++ b/ui/litellm-dashboard/src/components/router_settings/LatencyBasedConfiguration.tsx @@ -1,4 +1,5 @@ import React from "react"; +import { CONFIG_OWNED_HINT, ConfigOwnedBadge } from "./ConfigOwned"; import { Input } from "@/components/ui/input"; interface routingStrategyArgs { @@ -13,9 +14,13 @@ const defaultLowestLatencyArgs: routingStrategyArgs = { interface LatencyBasedConfigurationProps { routingStrategyArgs: { [key: string]: any }; + ownedByConfig?: boolean; } -const LatencyBasedConfiguration: React.FC = ({ routingStrategyArgs }) => { +const LatencyBasedConfiguration: React.FC = ({ + routingStrategyArgs, + ownedByConfig = false, +}) => { const paramExplanation: { [key: string]: string } = { ttl: "Sliding window to look back over when calculating the average latency of a deployment. Default - 1 hour (in seconds).", lowest_latency_buffer: @@ -26,7 +31,10 @@ const LatencyBasedConfiguration: React.FC = ({ r <>
-

Latency-Based Configuration

+

+ Latency-Based Configuration + {ownedByConfig && } +

Fine-tune latency-based routing behavior

@@ -41,6 +49,8 @@ const LatencyBasedConfiguration: React.FC = ({ r diff --git a/ui/litellm-dashboard/src/components/router_settings/ReliabilityRetriesSection.tsx b/ui/litellm-dashboard/src/components/router_settings/ReliabilityRetriesSection.tsx index 09bc7e65c57..1ca5d89f662 100644 --- a/ui/litellm-dashboard/src/components/router_settings/ReliabilityRetriesSection.tsx +++ b/ui/litellm-dashboard/src/components/router_settings/ReliabilityRetriesSection.tsx @@ -1,14 +1,17 @@ import React from "react"; +import { CONFIG_OWNED_HINT, ConfigOwnedBadge, RouterSettingsSource, isOwnedByConfig } from "./ConfigOwned"; import { Input } from "@/components/ui/input"; interface ReliabilityRetriesSectionProps { routerSettings: { [key: string]: any }; routerFieldsMetadata: { [key: string]: any }; + routerSettingsSource?: RouterSettingsSource; } const ReliabilityRetriesSection: React.FC = ({ routerSettings, routerFieldsMetadata, + routerSettingsSource = {}, }) => { return (
@@ -36,6 +39,7 @@ const ReliabilityRetriesSection: React.FC = ({ {routerFieldsMetadata[param]?.ui_field_name || param} + {isOwnedByConfig(routerSettingsSource, param) && }

{routerFieldsMetadata[param]?.field_description || ""}

@@ -49,6 +53,8 @@ const ReliabilityRetriesSection: React.FC = ({ : value?.toString() || "" } placeholder="—" + disabled={isOwnedByConfig(routerSettingsSource, param)} + title={isOwnedByConfig(routerSettingsSource, param) ? CONFIG_OWNED_HINT : undefined} className="font-mono text-sm w-full" /> diff --git a/ui/litellm-dashboard/src/components/router_settings/RouterSettingsForm.test.tsx b/ui/litellm-dashboard/src/components/router_settings/RouterSettingsForm.integration.test.tsx similarity index 73% rename from ui/litellm-dashboard/src/components/router_settings/RouterSettingsForm.test.tsx rename to ui/litellm-dashboard/src/components/router_settings/RouterSettingsForm.integration.test.tsx index 64f9f113fcc..06649f67a6e 100644 --- a/ui/litellm-dashboard/src/components/router_settings/RouterSettingsForm.test.tsx +++ b/ui/litellm-dashboard/src/components/router_settings/RouterSettingsForm.integration.test.tsx @@ -88,6 +88,38 @@ describe("RouterSettingsForm", () => { expect(onChange).toHaveBeenCalledWith(expect.objectContaining({ enableTagFiltering: true })); }); + it("disables every control whose key config.yaml owns and leaves the rest editable", () => { + const props = { + ...baseProps, + value: { + ...defaultValue, + selectedStrategy: "latency-based-routing", + routerSettings: { + routing_strategy: "latency-based-routing", + routing_strategy_args: { ttl: 3600 }, + num_retries: 3, + timeout: 60, + }, + }, + availableRoutingStrategies: ["simple-shuffle", "latency-based-routing"], + routerSettingsSource: { + routing_strategy: "config" as const, + routing_strategy_args: "config" as const, + enable_tag_filtering: "config" as const, + num_retries: "config" as const, + timeout: "db" as const, + }, + }; + render(); + + expect(screen.getByRole("combobox")).toBeDisabled(); + expect(screen.getByRole("switch")).toHaveAttribute("aria-disabled", "true"); + expect(document.querySelector('input[name="ttl"]')).toBeDisabled(); + expect(document.querySelector('input[name="num_retries"]')).toBeDisabled(); + expect(document.querySelector('input[name="timeout"]')).toBeEnabled(); + expect(screen.getAllByTestId("config-owned-badge")).toHaveLength(4); + }); + it("should show the Reliability & Retries section", () => { render(); expect(screen.getByText("Reliability & Retries")).toBeInTheDocument(); diff --git a/ui/litellm-dashboard/src/components/router_settings/RouterSettingsForm.tsx b/ui/litellm-dashboard/src/components/router_settings/RouterSettingsForm.tsx index 05e38993575..8e20296302f 100644 --- a/ui/litellm-dashboard/src/components/router_settings/RouterSettingsForm.tsx +++ b/ui/litellm-dashboard/src/components/router_settings/RouterSettingsForm.tsx @@ -1,5 +1,6 @@ import React from "react"; import LatencyBasedConfiguration from "./LatencyBasedConfiguration"; +import { RouterSettingsSource, isOwnedByConfig } from "./ConfigOwned"; import ReliabilityRetriesSection from "./ReliabilityRetriesSection"; import RoutingStrategySelector from "./RoutingStrategySelector"; import TagFilteringToggle from "./TagFilteringToggle"; @@ -16,6 +17,7 @@ interface RouterSettingsFormProps { routerFieldsMetadata: { [key: string]: any }; availableRoutingStrategies: string[]; routingStrategyDescriptions: { [key: string]: string }; + routerSettingsSource?: RouterSettingsSource; } const RouterSettingsForm: React.FC = ({ @@ -24,6 +26,7 @@ const RouterSettingsForm: React.FC = ({ routerFieldsMetadata, availableRoutingStrategies, routingStrategyDescriptions, + routerSettingsSource = {}, }) => { const handleStrategyChange = (strategy: string) => { onChange({ @@ -56,6 +59,7 @@ const RouterSettingsForm: React.FC = ({ routingStrategyDescriptions={routingStrategyDescriptions} routerFieldsMetadata={routerFieldsMetadata} onStrategyChange={handleStrategyChange} + ownedByConfig={isOwnedByConfig(routerSettingsSource, "routing_strategy")} /> )} @@ -64,6 +68,7 @@ const RouterSettingsForm: React.FC = ({ enabled={value.enableTagFiltering} routerFieldsMetadata={routerFieldsMetadata} onToggle={handleTagFilteringToggle} + ownedByConfig={isOwnedByConfig(routerSettingsSource, "enable_tag_filtering")} />
@@ -72,11 +77,18 @@ const RouterSettingsForm: React.FC = ({ {/* Strategy-Specific Args - Show immediately after strategy if latency-based */} {value.selectedStrategy === "latency-based-routing" && ( - + )} {/* Other Settings */} - +
); }; diff --git a/ui/litellm-dashboard/src/components/router_settings/RoutingStrategySelector.tsx b/ui/litellm-dashboard/src/components/router_settings/RoutingStrategySelector.tsx index 313732b8b3d..b24cadd63a3 100644 --- a/ui/litellm-dashboard/src/components/router_settings/RoutingStrategySelector.tsx +++ b/ui/litellm-dashboard/src/components/router_settings/RoutingStrategySelector.tsx @@ -1,4 +1,5 @@ import React from "react"; +import { CONFIG_OWNED_HINT, ConfigOwnedBadge } from "./ConfigOwned"; import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select"; interface RoutingStrategySelectorProps { @@ -7,6 +8,7 @@ interface RoutingStrategySelectorProps { routingStrategyDescriptions: { [key: string]: string }; routerFieldsMetadata: { [key: string]: any }; onStrategyChange: (strategy: string) => void; + ownedByConfig?: boolean; } const RoutingStrategySelector: React.FC = ({ @@ -15,6 +17,7 @@ const RoutingStrategySelector: React.FC = ({ routingStrategyDescriptions, routerFieldsMetadata, onStrategyChange, + ownedByConfig = false, }) => { return (
@@ -22,6 +25,7 @@ const RoutingStrategySelector: React.FC = ({ + {ownedByConfig && }

{routerFieldsMetadata["routing_strategy"]?.field_description || ""}

@@ -30,8 +34,9 @@ const RoutingStrategySelector: React.FC = ({