scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-12 03:48:56 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 57

Merge remote-tracking branch 'origin/main' into litellm_otel_v2_team_excluded_services
Some checks are pending
LiteLLM Rust / rust-lint (push) Waiting to run
Details
LiteLLM Rust / rust-test (push) Waiting to run
Details
LiteLLM Rust / rust-wheel (push) Waiting to run
Details

Browse source
This commit is contained in:
mrinal 2026-10-11 17:59:50 +00:00
commit 9c74e16086
256 changed files with 6183 additions and 17378 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

338
.circleci/config.yml
View file

@ -38,7 +38,7 @@ commands:
parameters:
category:
type: enum
enum: ["backend", "client", "provider-harness", "redis-compat"]
enum: ["backend", "client", "redis-compat"]
default: "backend"
steps:
- run:
@ -619,7 +619,7 @@ jobs:
name: Run Basic Proxy Startup Tests (Health Readiness and Chat Completion)
command: |
mkdir -p test-results
TEST_FILES=$(circleci tests glob "tests/basic_proxy_startup_tests/**/test_*.py")
TEST_FILES=$(circleci tests glob "tests/integration/configuration/test_basic_proxy_startup.py")
echo "$TEST_FILES" | circleci tests run \
--verbose \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest --tb=short \
@ -696,7 +696,7 @@ jobs:
name: Run second round of tests
command: |
mkdir -p test-results
TEST_FILES=$(circleci tests glob "tests/basic_proxy_startup_tests/**/test_*.py")
TEST_FILES=$(circleci tests glob "tests/integration/configuration/test_basic_proxy_startup.py")
echo "$TEST_FILES" | circleci tests run \
--verbose \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest --tb=short \
@ -708,84 +708,46 @@ jobs:
# Store test results
- store_test_results:
path: test-results
proxy_pass_through_endpoint_tests:
machine:
image: ubuntu-2204:2024.04.1
e2e_ui_build:
docker:
- image: cimg/python:3.12-browsers@sha256:b432899af01c9a311bf74f4f22e9ada2e5306d4b1b4383f8d29e1228a5844ef2
auth:
username: ${DOCKERHUB_USERNAME}
password: ${DOCKERHUB_PASSWORD}
resource_class: large
working_directory: ~/project
steps:
- checkout
- run:
name: Generate LiteLLM master key
command: |
key="$(openssl rand -hex 16)"
printf 'export LITELLM_MASTER_KEY=sk-%s\n' "$key" >> "$BASH_ENV"
- skip_if_unrelated_changes
- skip_if_unrelated_changes:
category: client
- setup_google_dns
- install_uv
- install_rust
- run:
name: Install Dependencies
command: |
uv sync --frozen --all-groups --all-extras --python 3.12
- start_postgres
- attach_workspace:
at: ~/project
- run:
name: Load Docker Database Image
command: |
zstd -d litellm-docker-database.tar.zst --stdout | docker load
docker images | grep litellm-docker-database
- run:
name: Run Docker container
command: |
docker run -d \
-p 4000:4000 \
-e DATABASE_URL=postgresql://postgres:postgres@host.docker.internal:5432/circle_test \
-e LITELLM_MASTER_KEY="$LITELLM_MASTER_KEY" \
-e OPENAI_API_KEY=$OPENAI_API_KEY \
-e GEMINI_API_KEY=$GEMINI_API_KEY \
-e ANTHROPIC_API_KEY=$ANTHROPIC_API_KEY \
-e ASSEMBLYAI_API_KEY=$ASSEMBLYAI_API_KEY \
-e AZURE_API_KEY=$AZURE_API_KEY \
-e AZURE_API_BASE=$AZURE_API_BASE \
-e USE_DDTRACE=True \
-e DD_API_KEY=$DD_API_KEY \
-e DD_SITE=$DD_SITE \
-e LITELLM_LICENSE=$LITELLM_LICENSE \
-e LITELLM_USE_CHAT_COMPLETIONS_URL_FOR_ANTHROPIC_MESSAGES=true \
-e LITELLM_LOG=ERROR \
--add-host host.docker.internal:host-gateway \
--name my-app \
-v $(pwd)/litellm/proxy/example_config_yaml/pass_through_config.yaml:/app/config.yaml \
-v $(pwd)/litellm/proxy/example_config_yaml/custom_auth_basic.py:/app/custom_auth_basic.py \
litellm-docker-database:ci \
--config /app/config.yaml \
--port 4000
- run:
name: Start outputting logs
command: docker logs -f my-app
background: true
- wait_for_service:
url: http://localhost:4000
timeout: "300"
# Install Node.js directly from nodejs.org with SHA256 verification,
# instead of piping NodeSource's setup_24.x apt-repo installer into
# sudo bash (which runs a mutable upstream script unattended).
- install_node
- restore_cache:
keys:
- ui-dashboard-node-deps-v1-{{ checksum "ui/litellm-dashboard/package-lock.json" }}
- run:
name: Install Node.js test dependencies
name: Install dashboard dependencies
command: |
cd tests/pass_through_tests
cd ui/litellm-dashboard
npm ci
- save_cache:
key: ui-dashboard-node-deps-v1-{{ checksum "ui/litellm-dashboard/package-lock.json" }}
paths:
- ui/litellm-dashboard/node_modules
- run:
name: Run Vertex AI, Google AI Studio Node.js tests
name: Build UI bundle
command: |
cd tests/pass_through_tests
NODE_OPTIONS=--experimental-vm-modules npx jest . --verbose
no_output_timeout: 30m
cd ui/litellm-dashboard
npm run build
mkdir -p ../../ui-bundle
mv out ../../ui-bundle/out
find ../../ui-bundle/out -name '*.html' ! -name 'index.html' | while read -r f; do
d="${f%.html}"; mkdir -p "$d"; mv "$f" "$d/index.html"
done
- persist_to_workspace:
root: ~/project
paths:
- ui-bundle/out
e2e_ui_testing:
docker:
@ -798,7 +760,7 @@ jobs:
POSTGRES_USER: e2euser
POSTGRES_PASSWORD: e2epassword
POSTGRES_DB: litellm_e2e
resource_class: large
resource_class: xlarge
working_directory: ~/project
environment:
DATABASE_URL: "postgresql://e2euser:e2epassword@localhost:5432/litellm_e2e"
@ -817,6 +779,8 @@ jobs:
printf 'export LITELLM_MASTER_KEY=sk-%s\n' "$key" >> "$BASH_ENV"
- skip_if_unrelated_changes:
category: client
- attach_workspace:
at: ~/project
- setup_google_dns
- install_node
- install_uv
@ -835,7 +799,7 @@ jobs:
- ~/.cache/uv
- restore_cache:
keys:
- ui-e2e-node-deps-v4-{{ checksum "ui/litellm-dashboard/package-lock.json" }}-{{ checksum "tests/e2e/ui/package-lock.json" }}
- ui-playwright-node-deps-v1-{{ checksum "tests/e2e/ui/package-lock.json" }}
- run:
name: Install Node dependencies and Playwright
# The cimg/python:3.12-browsers image already ships the Chromium system
@ -844,36 +808,19 @@ jobs:
# 5-10 minutes to the job and frequently stalls on flaky Ubuntu mirrors,
# so we install just the browser binary.
command: |
cd ui/litellm-dashboard
npm ci
cd ../../tests/e2e/ui
cd tests/e2e/ui
npm ci
npx playwright install chromium
- save_cache:
key: ui-e2e-node-deps-v4-{{ checksum "ui/litellm-dashboard/package-lock.json" }}-{{ checksum "tests/e2e/ui/package-lock.json" }}
key: ui-playwright-node-deps-v1-{{ checksum "tests/e2e/ui/package-lock.json" }}
paths:
- ui/litellm-dashboard/node_modules
- tests/e2e/ui/node_modules
- ~/.cache/ms-playwright
- run:
name: Build UI from source
# Prior version used `cp -r out/ ../../litellm/proxy/_experimental/out/`.
# GNU cp (used on CircleCI's Ubuntu image) interprets that as "copy the
# source directory as a child of the destination" when the destination
# already exists — silently creating `_experimental/out/out/` instead of
# replacing the served bundle. The proxy continued serving whatever was
# checked into `_experimental/out/*`, so this job was effectively testing
# the pre-build bundle on every run. Replace-and-move guarantees the
# freshly built bundle is what the proxy actually serves.
name: Use built UI bundle
command: |
cd ui/litellm-dashboard
npm run build
rm -rf ../../litellm/proxy/_experimental/out
mv out ../../litellm/proxy/_experimental/out
# Restructure HTML so extensionless routes work (login.html -> login/index.html)
find ../../litellm/proxy/_experimental/out -name '*.html' ! -name 'index.html' | while read -r f; do
d="${f%.html}"; mkdir -p "$d"; mv "$f" "$d/index.html"
done
rm -rf litellm/proxy/_experimental/out
cp -r ui-bundle/out litellm/proxy/_experimental/out
- wait_for_service:
url: tcp://localhost:5432
timeout: "30"
@ -940,7 +887,7 @@ jobs:
command: |
cd tests/e2e/ui
LITELLM_LICENSE="$LITELLM_LICENSE" \
npx playwright test --config playwright.config.ts
npx playwright test --config playwright.config.ts --workers=3 --retries=0
no_output_timeout: 10m
- store_artifacts:
path: tests/e2e/ui/test-results
@ -978,6 +925,8 @@ jobs:
printf 'export LITELLM_MASTER_KEY=sk-%s\n' "$key" >> "$BASH_ENV"
- skip_if_unrelated_changes:
category: client
- attach_workspace:
at: ~/project
- setup_google_dns
- install_node
- install_uv
@ -996,31 +945,23 @@ jobs:
- ~/.cache/uv
- restore_cache:
keys:
- ui-e2e-node-deps-v4-{{ checksum "ui/litellm-dashboard/package-lock.json" }}-{{ checksum "tests/e2e/ui/package-lock.json" }}
- ui-playwright-node-deps-v1-{{ checksum "tests/e2e/ui/package-lock.json" }}
- run:
name: Install Node dependencies and Playwright
command: |
cd ui/litellm-dashboard
npm ci
cd ../../tests/e2e/ui
cd tests/e2e/ui
npm ci
npx playwright install chromium
- save_cache:
key: ui-e2e-node-deps-v4-{{ checksum "ui/litellm-dashboard/package-lock.json" }}-{{ checksum "tests/e2e/ui/package-lock.json" }}
key: ui-playwright-node-deps-v1-{{ checksum "tests/e2e/ui/package-lock.json" }}
paths:
- ui/litellm-dashboard/node_modules
- tests/e2e/ui/node_modules
- ~/.cache/ms-playwright
- run:
name: Build UI from source
name: Use built UI bundle
command: |
cd ui/litellm-dashboard
npm run build
rm -rf ../../litellm/proxy/_experimental/out
mv out ../../litellm/proxy/_experimental/out
find ../../litellm/proxy/_experimental/out -name '*.html' ! -name 'index.html' | while read -r f; do
d="${f%.html}"; mkdir -p "$d"; mv "$f" "$d/index.html"
done
rm -rf litellm/proxy/_experimental/out
cp -r ui-bundle/out litellm/proxy/_experimental/out
- wait_for_service:
url: tcp://localhost:5432
timeout: "30"
@ -1271,45 +1212,6 @@ jobs:
! docker exec my-app true 2>/dev/null || fail "container is still running"
echo "Proxy refused to serve (exit $exit_code) and never reached startup. Test passed."
provider_replay_harness:
docker:
- &python312_image
image: cimg/python:3.12@sha256:9c796c23c84e84a66a964acb508d39dc5433c81a47e07efd56dccbbc2427e07c
auth:
username: ${DOCKERHUB_USERNAME}
password: ${DOCKERHUB_PASSWORD}
- image: redis@sha256:e2debfb7956fa12c7ddc79d7e645c8cf26b30c99a6e9161ea9bf4171e1668a5f
working_directory: ~/project
resource_class: medium
environment:
E2E_CACHE_TEST_REDIS_URL: redis://127.0.0.1:6379/0
E2E_PROVIDER_CACHE: "0"
E2E_FIXTURE_MODE: live
steps:
- checkout
- run:
name: Generate LiteLLM master key
command: |
key="$(openssl rand -hex 16)"
printf 'export LITELLM_MASTER_KEY=sk-%s\n' "$key" >> "$BASH_ENV"
- skip_if_unrelated_changes:
category: provider-harness
- setup_litellm_test_deps
- wait_for_service:
url: tcp://localhost:6379
- run:
name: Test provider capture and replay harness
command: |
mkdir -p test-results/provider-replay-harness
uv run --no-sync pytest --tb=short -q --noconftest -o addopts= -o "pythonpath=tests/e2e tests/e2e_harness" -p no:rerunfailures \
--junitxml=test-results/provider-replay-harness/junit.xml \
tests/e2e_harness/test_provider_edge.py tests/e2e_harness/test_fixture_bundle.py \
tests/e2e_harness/test_fixture_canonical.py tests/e2e_harness/test_fixture_mode.py \
tests/code_coverage_tests/test_provider_replay_harness.py \
tests/code_coverage_tests/test_provider_cache.py
- store_test_results:
path: test-results/provider-replay-harness
integration_contracts:
parameters:
suite:
@ -1682,9 +1584,14 @@ jobs:
path: test-results
workflows:
migration_startup:
when: << pipeline.parameters.run_migration_tests >>
jobs: &migration_jobs
migration_startup_scheduled:
triggers:
- schedule:
cron: "17 0,6,12,18 * * *"
filters:
branches:
only: main
jobs:
- build_docker_database_image:
migration_qualification: true
- migration_startup_tests:
@ -1709,33 +1616,11 @@ workflows:
suite: shaped
baseline: true
requires: [build_docker_database_image]
migration_startup_scheduled:
triggers:
- schedule:
cron: "17 0,6,12,18 * * *"
filters:
branches:
only: main
jobs: *migration_jobs
routing_parity:
when:
not:
equal: ["", << pipeline.parameters.routing_parity_base >>]
jobs:
- routing_parity:
name: routing-parity-<< matrix.suite >>
matrix:
parameters:
suite: [management, accounting, database, providers, extensions, cost, mcp]
integration:
unless:
or:
- << pipeline.parameters.run_migration_tests >>
- not:
equal: ["", << pipeline.parameters.routing_parity_base >>]
jobs:
- integration_contracts:
name: integration-<< matrix.suite >>
filters: &integration_default_run not pipeline.parameters.run_migration_tests and pipeline.parameters.routing_parity_base == ""
matrix:
parameters:
suite: [sdk, cost, browser]
@ -1745,116 +1630,131 @@ workflows:
suite: security
parallelism: 1
resource_class: xlarge
filters: *integration_default_run
- integration_contracts:
name: integration-management
suite: management
parallelism: 4
resource_class: xlarge
filters: *integration_default_run
- integration_shard_coverage:
name: integration-management-coverage
suite: management
shards: 4
requires: [integration-management]
filters: *integration_default_run
- integration_contracts:
name: integration-accounting
suite: accounting
parallelism: 4
resource_class: xlarge
filters: *integration_default_run
- integration_shard_coverage:
name: integration-accounting-coverage
suite: accounting
shards: 4
requires: [integration-accounting]
filters: *integration_default_run
- integration_contracts:
name: integration-providers
suite: providers
parallelism: 4
resource_class: xlarge
filters: *integration_default_run
- integration_shard_coverage:
name: integration-providers-coverage
suite: providers
shards: 4
requires:
- integration-providers
filters: *integration_default_run
- integration_contracts:
name: integration-database
suite: database
parallelism: 2
filters: *integration_default_run
- integration_shard_coverage:
name: integration-database-coverage
suite: database
shards: 2
requires: [integration-database]
filters: *integration_default_run
- integration_contracts:
name: integration-mcp
suite: mcp
resource_class: xlarge
parallelism: 2
filters: *integration_default_run
- integration_shard_coverage:
name: integration-mcp-coverage
suite: mcp
shards: 2
requires:
- integration-mcp
filters: *integration_default_run
- integration_contracts:
name: integration-extensions
suite: extensions
parallelism: 8
resource_class: xlarge
filters: *integration_default_run
- integration_shard_coverage:
name: integration-extensions-coverage
suite: extensions
shards: 8
requires:
- integration-extensions
filters: *integration_default_run
- integration_contracts:
name: integration-management-replica
suite: management
mode: replica
parallelism: 4
resource_class: xlarge
filters: *integration_default_run
- integration_shard_coverage:
name: integration-management-replica-coverage
suite: management
mode: replica
shards: 4
requires: [integration-management-replica]
filters: *integration_default_run
- integration_contracts:
name: integration-database-replica
suite: database
mode: replica
parallelism: 2
filters: *integration_default_run
- integration_shard_coverage:
name: integration-database-replica-coverage
suite: database
mode: replica
shards: 2
requires: [integration-database-replica]
- postgres_suite:
name: proxy-behavior
test_path: tests/proxy_behavior
seed: true
coverage_flag: lens-postgres
timeout_minutes: 25
filters: *integration_default_run
- postgres_suite:
name: proxy-security
test_path: tests/proxy_security_tests
seed: true
timeout_minutes: 15
filters: *integration_default_run
- postgres_suite:
name: schema-migration
test_path: tests/proxy_migration_tests
seed: false
timeout_minutes: 20
workers: 2
filters: *integration_default_run
- postgres_suite:
name: roi-database
test_path: tests/integration/database/test_roi_observed.py
seed: false
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:
@ -1864,31 +1764,71 @@ workflows:
- "6.4.0"
- "7.4.1"
- "8.0.1"
- windows_release_wheel
- base_sdk_install
build_and_test:
unless:
or:
- << pipeline.parameters.run_migration_tests >>
- not:
equal: ["", << pipeline.parameters.routing_parity_base >>]
jobs:
- provider_replay_harness
- build_docker_database_image
- e2e_ui_testing
- e2e_ui_testing_server_root_path
filters: *integration_default_run
- windows_release_wheel:
filters: *integration_default_run
- base_sdk_install:
filters: *integration_default_run
- build_docker_database_image:
filters: *integration_default_run
- proxy_logging_guardrails_model_info_tests:
requires:
- build_docker_database_image
- proxy_pass_through_endpoint_tests:
requires:
- build_docker_database_image
filters: *integration_default_run
- db_migration_disable_update_check:
requires:
- build_docker_database_image
- helm_chart_testing:
requires:
- build_docker_database_image
filters: *integration_default_run
- test_bad_database_url:
requires:
- build_docker_database_image
filters: *integration_default_run
- helm_chart_testing:
requires:
- build_docker_database_image
filters: *integration_default_run
- e2e_ui_build:
filters: *integration_default_run
- e2e_ui_testing:
requires: [e2e_ui_build]
filters: *integration_default_run
- e2e_ui_testing_server_root_path:
requires: [e2e_ui_build]
filters: *integration_default_run
- build_docker_database_image:
name: build_docker_database_image-migration
migration_qualification: true
filters: pipeline.parameters.run_migration_tests
- migration_startup_tests:
name: migration-startup
suite: startup
requires: [build_docker_database_image-migration]
filters: pipeline.parameters.run_migration_tests
- migration_startup_tests:
name: migration-recovery
suite: recovery
requires: [build_docker_database_image-migration]
filters: pipeline.parameters.run_migration_tests
- migration_startup_tests:
name: migration-legacy-and-pooling
suite: legacy
requires: [build_docker_database_image-migration]
filters: pipeline.parameters.run_migration_tests
- migration_startup_tests:
name: migration-upgrade
suite: upgrade
baseline: true
requires: [build_docker_database_image-migration]
filters: pipeline.parameters.run_migration_tests
- migration_startup_tests:
name: migration-upgrade-shaped
suite: shaped
baseline: true
requires: [build_docker_database_image-migration]
filters: pipeline.parameters.run_migration_tests
- routing_parity:
name: routing-parity-<< matrix.suite >>
matrix:
parameters:
suite: [management, accounting, database, providers, extensions, cost, mcp]
filters: pipeline.parameters.routing_parity_base != ""

13
.circleci/scripts/classify_changes.sh
View file

@ -1,12 +1,11 @@
#!/usr/bin/env bash
set -uo pipefail
category="${1:?usage: classify_changes.sh <backend|client|ui|provider-harness|cost-map-only|mcp-dependencies|windows-release|redis-compat>}"
category="${1:?usage: classify_changes.sh <backend|client|ui|cost-map-only|mcp-dependencies|windows-release|redis-compat>}"
has_client=false
has_backend=false
has_ci=false
has_provider_harness=false
has_cost_map=false
has_mcp_dependencies=false
has_windows_release=false
@ -23,11 +22,6 @@ while IFS= read -r file || [ -n "$file" ]; do
pyproject.toml | */pyproject.toml | uv.lock | uv.toml | .python-version | rust-toolchain.toml | litellm-rust/* | litellm/__init__.py | litellm/proxy/proxy_server.py | litellm/*mcp* | tests/*mcp* | litellm/integrations/arize/* | tests/base_sdk_tests/* | scripts/check_mcp_sdk_install.py | .github/workflows/test-mcp-dependency-resolution.yml | .github/actions/detect-changes/* | .github/actions/setup-uv-with-retries/* | .github/actions/cache-cargo-build/* | .github/scripts/detect_changes.sh | .github/scripts/uv_sync_with_retries.sh | .circleci/scripts/classify_changes.sh | tests/unit/test_circleci_path_filter.py | tests/unit/test_detect_changes.py)
has_mcp_dependencies=true ;;
esac
case "$file" in
tests/e2e/*/*.py | tests/e2e_harness/*/*.py) : ;;
tests/e2e/*.py | tests/e2e_harness/*.py | tests/code_coverage_tests/test_provider_cache.py | tests/code_coverage_tests/test_provider_replay_harness.py | tests/unit/test_circleci_path_filter.py | .circleci/* | pyproject.toml | uv.lock)
has_provider_harness=true ;;
esac
case "$file" in
litellm-rust/* | litellm/rust_bridge/* | rust-toolchain.toml | pyproject.toml | uv.lock | tests/windows_tests/* | .circleci/*)
has_windows_release=true ;;
@ -41,7 +35,7 @@ while IFS= read -r file || [ -n "$file" ]; do
case "$file" in
model_prices_and_context_window.json | litellm/model_prices_and_context_window_backup.json | model_prices_and_context_window.schema.json)
has_cost_map=true ;;
tests/test_litellm/* | tests/proxy_unit_tests/* | tests/unit/proxy/*) : ;;
tests/unit/proxy/*) : ;;
*) outside_cost_map_set=true ;;
esac
done
@ -53,9 +47,6 @@ case "$category" in
cost-map-only)
{ [ "$has_cost_map" = true ] && [ "$outside_cost_map_set" = false ]; } && echo run || echo skip
;;
provider-harness)
[ "$has_provider_harness" = true ] && echo run || echo skip
;;
windows-release)
[ "$has_windows_release" = true ] && echo run || echo skip
;;

2
.circleci/scripts/path_filter.sh
View file

@ -1,7 +1,7 @@
#!/usr/bin/env bash
set -uo pipefail
category="${1:?usage: path_filter.sh <backend|client|provider-harness|redis-compat>}"
category="${1:?usage: path_filter.sh <backend|client|redis-compat>}"
here="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
run_full() {

2
.circleci/scripts/run_migration_tests.py
View file

@ -10,7 +10,7 @@ from typing import Final
from xml.etree import ElementTree
SUITES: Final = {
"startup": (("test_startup.py",), 12),
"startup": (("test_startup.py",), 13),
"recovery": (("test_recovery.py",), 15),
"legacy": (("test_legacy.py", "test_pooling.py"), 11),
"upgrade": (("test_upgrade.py", "test_rolling_upgrade.py"), 5),

54
.github/ci-coverage-allowlist.yml vendored
View file

@ -23,18 +23,6 @@ test_paths:
and that pull request gives it a new home
paths:
- tests/test_rust_python_harness.py
- reason: >-
Live-provider caching cases in tests/local_testing that remain outside CI. No CircleCI job
globs that directory anymore, and the legacy jobs that did always deselected these files.
The two Redis Cluster IAM authentication tests moved to tests/integration/sdk and remain
selected by node ID in redis_compat, where redis-server is installed.
The five keyless files now run in the caching-local shard, leaving live cases in these two.
Measured 2026-08-21 with no provider credentials: test_disk_cache_unit_tests.py needs
OPENAI_API_KEY for 2 of its 4 tests, and test_gcs_cache_unit_tests.py needs GCS credentials
for all 4. Keep these live cases outside PR CI rather than making a job red by construction
paths:
- tests/local_testing/test_disk_cache_unit_tests.py
- tests/local_testing/test_gcs_cache_unit_tests.py
- reason: >-
The end-to-end suite runs against a deployed proxy from its own in-cluster rig rather than
from a pull request; it needs a live gateway and provider credentials no PR job holds
@ -60,48 +48,6 @@ test_paths:
choice it informed is settled
paths:
- tests/code_coverage_tests/test_aio_http_image_conversion.py
- reason: >-
No job invokes this suite and its files mix pure transformation tests with ones driving live
vendor vector stores, so assigning them needs a per-file decision
paths:
- tests/vector_store_tests/rag/test_rag_bedrock.py
- tests/vector_store_tests/rag/test_rag_openai.py
- tests/vector_store_tests/rag/test_rag_s3_vectors.py
- tests/vector_store_tests/rag/test_rag_vertex_ai.py
- tests/vector_store_tests/test_azure_ai_vector_store.py
- tests/vector_store_tests/test_azure_vector_store.py
- tests/vector_store_tests/test_bedrock_vector_store.py
- tests/vector_store_tests/test_gemini_vector_store.py
- tests/vector_store_tests/test_milvus_vector_store.py
- tests/vector_store_tests/test_openai_vector_store.py
- tests/vector_store_tests/test_ragflow_vector_store.py
- tests/vector_store_tests/test_s3_vectors_vector_store.py
- tests/vector_store_tests/test_vertex_ai_search_api_vector_store.py
- tests/vector_store_tests/test_vertex_ai_vector_store.py
- reason: >-
Throughput and memory-growth measurements whose runtime and variance make them unsuitable for
a per-pull-request job
paths:
- tests/load_tests/test_datadog_load_test.py
- tests/load_tests/test_langsmith_load_test.py
- tests/load_tests/test_linear_memory_growth.py
- tests/load_tests/test_memory_usage.py
- tests/load_tests/test_otel_load_test.py
- tests/load_tests/test_vertex_embeddings_load_test.py
- tests/load_tests/test_vertex_load_tests.py
- reason: >-
Env-gated saturation benchmark requires a live proxy and provider credentials, so it is run
locally rather than in pull-request jobs
paths:
- tests/load_tests/test_granian_admission_saturation.py
- reason: >-
A local-only agent rig: test_a2a_completion_bridge.py needs a LangGraph server on
localhost:2024 and test_a2a.py drives a live A2A endpoint, so neither can run in a
pull request job. Until 2026-08-20 the CircleCI agent job hid them behind a grep -v
that this census could not see; the glob now excludes them structurally and this entry
is the decision on the record. Revisit when the A2A bridge gets a recorded-wire fixture
paths:
- tests/agent_tests/local_only_agent_tests
- reason: >-
Third-party integration tests that skip themselves without OCI configuration or sandbox
credentials, neither of which a pull request job holds

1
.github/scripts/assert_ci_coverage.py vendored
View file

@ -37,7 +37,6 @@ GLOB_CHARS = frozenset("*?")
# tests has to be named by some shard or it runs nowhere. A child listed here is
# itself decomposed one level deeper and is checked through its own entry.
SHARDED_ROOTS: tuple[str, ...] = (
"tests/test_litellm",
"tests/unit",
"tests/unit/enterprise",
"tests/unit/enterprise/enterprise_callbacks",

30
.github/workflows/deploy-dev-on-main.yml vendored Normal file
View file

@ -0,0 +1,30 @@
name: Deploy main to gateway-dev
on:
push:
branches: [main]
permissions: {}
concurrency:
group: deploy-dev-dispatch
cancel-in-progress: true
jobs:
dispatch:
if: github.repository == 'BerriAI/litellm'
runs-on: ubuntu-latest
timeout-minutes: 5
steps:
- uses: actions/create-github-app-token@fee1f7d63c2ff003460e3d139729b119787bc349 # v2.2.2
id: token
with:
app-id: ${{ secrets.LITELLM_OPS_DISPATCH_APP_ID }}
private-key: ${{ secrets.LITELLM_OPS_DISPATCH_APP_PRIVATE_KEY }}
owner: BerriAI
repositories: litellm-ops
permission-actions: write
- name: Dispatch the litellm-ops dev deploy
env:
GH_TOKEN: ${{ steps.token.outputs.token }}
run: gh api --method POST repos/BerriAI/litellm-ops/actions/workflows/deploy-dev.yml/dispatches -f ref=main

16
.github/workflows/test-linting.yml vendored
View file

@ -22,6 +22,16 @@ jobs:
actions: read
runs-on: ubuntu-latest
timeout-minutes: 15
services:
redis:
image: redis@sha256:e2debfb7956fa12c7ddc79d7e645c8cf26b30c99a6e9161ea9bf4171e1668a5f
ports:
- 6379:6379
options: >-
--health-cmd "redis-cli ping"
--health-interval 10s
--health-timeout 5s
--health-retries 5
steps:
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
@ -189,8 +199,9 @@ jobs:
if: steps.changes.outputs.decision != 'skip'
env:
LITELLM_MASTER_KEY: sk-e2e-harness-tests-reach-no-proxy
E2E_CACHE_TEST_REDIS_URL: redis://127.0.0.1:6379/0
run: |
if ! git diff --name-only --diff-filter=ACMRD "$GATE_BASE_SHA" HEAD -- tests/e2e tests/e2e_harness ':(exclude)tests/e2e/ui' pyproject.toml uv.lock .github/workflows/test-linting.yml | grep -q .; then
if ! git diff --name-only --diff-filter=ACMRD "$GATE_BASE_SHA" HEAD -- tests/e2e tests/e2e_harness ':(exclude)tests/e2e/ui' tests/code_coverage_tests/test_provider_replay_harness.py tests/code_coverage_tests/test_provider_cache.py pyproject.toml uv.lock .github/workflows/test-linting.yml | grep -q .; then
echo "No changed e2e harness files; skipping."
exit 0
fi
@ -198,6 +209,9 @@ jobs:
CLAUDE_VERSION="$(retry uv run --no-sync python tests/e2e/claude_code/pr_gate_version_resolver.py)"
tests/e2e/claude_code/cron_vm/install_claude_code.sh "$CLAUDE_VERSION" "$RUNNER_TEMP/claude-cli"
PATH="$RUNNER_TEMP/claude-cli:$PATH" uv run --no-sync pytest -q tests/e2e_harness
uv run --no-sync pytest --tb=short -q --noconftest -o addopts= -o "pythonpath=tests/e2e tests/e2e_harness" -p no:rerunfailures \
tests/code_coverage_tests/test_provider_replay_harness.py \
tests/code_coverage_tests/test_provider_cache.py
- name: Check for circular imports
if: steps.changes.outputs.decision != 'skip'

25
.github/workflows/test-merge-smoke.yml vendored
View file

@ -234,7 +234,30 @@ jobs:
exit 1
- name: Run the startup tests
run: .venv/bin/python -m pytest --tb=short -vv tests/basic_proxy_startup_tests
shell: bash
run: |
set -euo pipefail
base_url=http://localhost:4000
readiness_response="$(curl -sS --fail-with-body --write-out $'\n%{http_code}' "${base_url}/health/readiness")"
readiness_status="${readiness_response##*$'\n'}"
readiness_body="${readiness_response%$'\n'*}"
[[ "${readiness_status}" == "200" ]]
jq -e '.status == "healthy" or .status == "connected"' <<< "${readiness_body}" > /dev/null
liveness_status="$(curl -sS --fail-with-body --output /dev/null --write-out '%{http_code}' "${base_url}/health/liveness")"
[[ "${liveness_status}" == "200" ]]
completion_response="$(curl -sS --fail-with-body --write-out $'\n%{http_code}' \
--header "Authorization: Bearer ${LITELLM_MASTER_KEY}" \
--header 'Content-Type: application/json' \
--data '{"model":"gpt-4","messages":[{"role":"system","content":"You are a helpful assistant."},{"role":"user","content":"Hello!"}]}' \
"${base_url}/chat/completions")"
completion_status="${completion_response##*$'\n'}"
completion_body="${completion_response%$'\n'*}"
[[ "${completion_status}" == "200" ]]
jq -e '
(.choices | length) == 1 and
.choices[0].message.role == "assistant" and
.choices[0].message.content == "Hello! This is a mock response from the fake OpenAI endpoint."
' <<< "${completion_body}" > /dev/null
- name: Show the proxy logs
if: always()

2
.github/workflows/test-rust.yml vendored
View file

@ -18,7 +18,6 @@ on:
- "litellm/ocr/**"
- "litellm/llms/base_llm/ocr/**"
- "litellm/llms/custom_httpx/llm_http_handler.py"
- "tests/test_litellm/conftest.py"
- "Makefile"
- ".cargo/**"
- "pyproject.toml"
@ -49,7 +48,6 @@ on:
- "litellm/ocr/**"
- "litellm/llms/base_llm/ocr/**"
- "litellm/llms/custom_httpx/llm_http_handler.py"
- "tests/test_litellm/conftest.py"
- "Makefile"
- ".cargo/**"
- "pyproject.toml"

20
.github/workflows/test-unit.yml vendored
View file

@ -174,8 +174,6 @@ jobs:
- shard: integrations
test-path: |-
tests/test_litellm/integrations
tests/test_litellm/tracing
tests/unit/integrations
python-version: "3.12"
workers: 4
@ -292,6 +290,7 @@ jobs:
tests/unit/llms/manus
tests/unit/llms/meta_llama
tests/unit/llms/microsoft_365_copilot
tests/unit/llms/milvus
tests/unit/llms/minimax
tests/unit/llms/mistral
tests/unit/llms/modelscope
@ -376,7 +375,6 @@ jobs:
- shard: root
test-path: >-
tests/test_litellm/test_*.py
tests/unit/test_*.py
python-version: "3.12"
workers: 4
@ -395,16 +393,6 @@ jobs:
job-timeout-minutes: 60
dist: loadscope
- shard: router-unit-tests
test-path: >-
tests/router_unit_tests
python-version: "3.12"
workers: 0
reruns: 0
timeout-minutes: 10
job-timeout-minutes: 60
dist: loadscope
- shard: endpoints
test-path: >-
tests/unit/a2a_protocol
@ -775,7 +763,7 @@ jobs:
- shard: proxy-server-core
test-path: |-
tests/proxy_unit_tests/test_proxy_server_gemini_pass_through.py
tests/unit/proxy/pass_through_endpoints/test_llm_passthrough_endpoints.py
tests/unit/proxy/test__lazy_features.py
tests/unit/proxy/test_aproxy_startup.py
tests/unit/proxy/test_proxy_server.py
@ -816,7 +804,7 @@ jobs:
- shard: custom-logging
test-path: |-
tests/proxy_unit_tests/test_proxy_custom_logger.py
tests/unit/proxy/test_proxy_server_custom_logger.py
tests/unit/proxy/test_custom_callback_input.py
tests/unit/proxy/test_custom_logger_s3_gcs.py
python-version: "3.12"
@ -880,7 +868,7 @@ jobs:
- shard: endpoints-and-responses
test-path: |-
tests/proxy_unit_tests/test_proxy_exception_mapping.py
tests/unit/proxy/test_proxy_server_exception_mapping.py
tests/unit/proxy/lens
tests/unit/proxy/auth/test_models_fallback_endpoint.py
tests/unit/proxy/common_utils/test_check_batch_cost.py

8
Makefile
View file

@ -39,7 +39,7 @@ help:
@echo " make check-circular-imports - Check for circular imports"
@echo " make check-import-safety - Check import safety"
@echo " make test - Run all tests"
@echo " make test-unit - Run unit tests (tests/unit and tests/test_litellm)"
@echo " make test-unit - Run unit tests (tests/unit)"
@echo " make test-unit-llms - Run LLM provider tests (~225 files)"
@echo " make test-unit-proxy-guardrails - Run proxy guardrails+mgmt tests (~51 files)"
@echo " make test-unit-proxy-core - Run proxy auth+client+db+hooks tests (~52 files)"
@ -300,7 +300,7 @@ test: install-test-deps
$(UV_RUN) pytest tests/
test-unit: install-test-deps
$(UV_RUN) pytest tests/unit tests/test_litellm -x -vv -n 4
$(UV_RUN) pytest tests/unit -x -vv -n 4
# Matrix test targets (matching CI workflow groups)
test-unit-llms: install-test-deps
@ -328,7 +328,7 @@ test-unit-other: install-test-deps
$(UV_RUN) pytest tests/unit/caching tests/unit/responses tests/unit/secret_managers tests/unit/vector_stores tests/unit/a2a_protocol tests/unit/completion_extras tests/unit/containers tests/unit/enterprise tests/unit/experimental_mcp_client tests/unit/google_genai tests/unit/images tests/unit/interactions tests/unit/router_strategy tests/unit/router_utils tests/unit/types --tb=short -vv -n 4 --durations=20
test-unit-root: install-test-deps
$(UV_RUN) pytest tests/unit/test_*.py tests/test_litellm/test_*.py --tb=short -vv -n 4 --durations=20
$(UV_RUN) pytest tests/unit/test_*.py --tb=short -vv -n 4 --durations=20
# Proxy unit tests (tests/unit/proxy split alphabetically)
test-proxy-unit-a: install-test-deps
@ -338,7 +338,7 @@ test-proxy-unit-b: install-test-deps
$(UV_RUN) pytest tests/unit/proxy/test_[p-z]*.py tests/unit/skills --tb=short -vv -n 2 --durations=20
test-integration: install-test-deps
$(UV_RUN) pytest tests/ -k "not test_litellm"
$(UV_RUN) pytest tests/
test-unit-helm: install-helm-unittest
helm unittest -f 'tests/*.yaml' helm/litellm-helm

2
codecov.yaml
View file

@ -71,6 +71,8 @@ flag_management:
carryforward: false
- name: lens-python-310
carryforward: false
- name: lens-postgres
carryforward: false
- name: proxy-db-key-generation
carryforward: false
- name: proxy-db-auth-checks

5
litellm/__init__.py
View file

@ -440,10 +440,7 @@ autorouter_presets_url: str = os.getenv(
"LITELLM_AUTOROUTER_PRESETS_URL",
"https://raw.githubusercontent.com/BerriAI/litellm/main/litellm/proxy/public_endpoints/autorouter_presets.json",
)
whats_new_url: str = os.getenv(
"LITELLM_WHATS_NEW_URL",
"https://raw.githubusercontent.com/BerriAI/litellm/main/litellm/proxy/public_endpoints/whats_new.json",
)
whats_new_url: str | None = os.getenv("LITELLM_WHATS_NEW_URL")
suppress_debug_info: bool = False
dynamodb_table_name: Optional[str] = None
s3_callback_params: Optional[Dict] = None

12
litellm/litellm_core_utils/litellm_logging.py
View file

@ -2507,7 +2507,7 @@ class Logging(LiteLLMLoggingBaseClass):
self.callback_duration_ms += (time.time() - _start) * 1000
return payload
def _transform_usage_objects(self, result):
def _transform_usage_objects(self, result: object) -> object:
if isinstance(result, ResponsesAPIResponse):
result = result.model_copy()
transformed_usage = ResponseAPILoggingUtils.transform_response_api_usage_to_chat_usage(result.usage)
@ -2545,13 +2545,13 @@ class Logging(LiteLLMLoggingBaseClass):
def _success_handler_helper_fn(
self,
result=None,
start_time=None,
end_time=None,
cache_hit=None,
result: object = None,
start_time: datetime.datetime | None = None,
end_time: datetime.datetime | None = None,
cache_hit: bool | None = None,
standard_logging_object: StandardLoggingPayload | None = None,
build_logging_payload: bool = True,
):
) -> tuple[datetime.datetime, datetime.datetime, object]:
try:
if start_time is None:
start_time = self.start_time

84
litellm/litellm_core_utils/streaming_chunk_builder_utils.py
View file

@ -1,9 +1,10 @@
import base64
import time
from collections.abc import Callable, Iterator, Mapping, Sequence
from dataclasses import dataclass, replace
from itertools import groupby
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, Optional, TypeAlias, TypedDict, Union, cast
from typing import TYPE_CHECKING, Final, Optional, TypeAlias, TypedDict, Union, cast
from pydantic import BaseModel
from typing_extensions import ReadOnly, Required
@ -154,6 +155,15 @@ class _ToolCallChoice(TypedDict, total=False):
_ToolCallKey: TypeAlias = tuple[int, int]
@dataclass(frozen=True, slots=True)
class _ToolCallAccumulator:
id: str | None = None
name: str | None = None
type: str | None = None
custom_name: str | None = None
provider_specific_fields: Mapping[str, object] | None = None
class _ToolCallChunk(TypedDict):
choices: ReadOnly[Sequence[_ToolCallChoice]]
@ -503,7 +513,7 @@ class ChunkProcessor:
self, tool_call_chunks: Sequence["_ToolCallChunk"]
) -> list[ChatCompletionMessageToolCall | ChatCompletionMessageCustomToolCall]:
tool_calls_list: list[ChatCompletionMessageToolCall | ChatCompletionMessageCustomToolCall] = []
tool_call_map: Final[dict[_ToolCallKey, dict[str, Any]]] = {}
tool_call_map: Final[dict[_ToolCallKey, _ToolCallAccumulator]] = {}
for chunk in tool_call_chunks:
choices = chunk["choices"]
@ -537,49 +547,43 @@ class ChunkProcessor:
index = (choice_index, getattr(tool_call, "index", 0))
if index not in tool_call_map:
tool_call_map[index] = {
"id": None,
"name": None,
"type": None,
"custom_name": None,
"provider_specific_fields": None,
}
tool_call_map[index] = _ToolCallAccumulator()
# Extract id, type, and function data (handle both dict and object)
if isinstance(tool_call, dict):
if fragment_id := tool_call.get("id"):
tool_call_map[index]["id"] = fragment_id
tool_call_map[index] = replace(tool_call_map[index], id=fragment_id)
if fragment_type := tool_call.get("type"):
tool_call_map[index]["type"] = fragment_type
tool_call_map[index] = replace(tool_call_map[index], type=fragment_type)
function = tool_call.get("function", {})
if isinstance(function, dict):
if fragment_name := function.get("name"):
tool_call_map[index]["name"] = fragment_name
tool_call_map[index] = replace(tool_call_map[index], name=fragment_name)
else:
# function is an object
if function_name := getattr(function, "name", None):
tool_call_map[index]["name"] = function_name
tool_call_map[index] = replace(tool_call_map[index], name=function_name)
custom = tool_call.get("custom")
if isinstance(custom, dict):
if custom_name := custom.get("name"):
tool_call_map[index]["custom_name"] = custom_name
tool_call_map[index] = replace(tool_call_map[index], custom_name=custom_name)
else:
# tool_call is an object
if hasattr(tool_call, "id") and tool_call.id:
tool_call_map[index]["id"] = tool_call.id
tool_call_map[index] = replace(tool_call_map[index], id=tool_call.id)
if hasattr(tool_call, "type") and tool_call.type:
tool_call_map[index]["type"] = tool_call.type
tool_call_map[index] = replace(tool_call_map[index], type=tool_call.type)
if object_function_name := getattr(getattr(tool_call, "function", None), "name", None):
tool_call_map[index]["name"] = object_function_name
tool_call_map[index] = replace(tool_call_map[index], name=object_function_name)
object_custom: ChatCompletionDeltaCustomToolCallPayload | None = getattr(
tool_call, "custom", None
)
if object_custom is not None:
if getattr(object_custom, "name", None):
tool_call_map[index]["custom_name"] = object_custom.name
tool_call_map[index] = replace(tool_call_map[index], custom_name=object_custom.name)
# Preserve provider_specific_fields from streaming chunks
provider_fields: object = None
@ -602,12 +606,14 @@ class ChunkProcessor:
if provider_fields:
# Merge provider_specific_fields if multiple chunks have them
merged_provider_fields = tool_call_map[index]["provider_specific_fields"]
if merged_provider_fields is None:
merged_provider_fields = {}
tool_call_map[index]["provider_specific_fields"] = merged_provider_fields
if isinstance(provider_fields, dict):
merged_provider_fields.update(provider_fields)
tool_call_map[index] = replace(
tool_call_map[index],
provider_specific_fields=(
{**(tool_call_map[index].provider_specific_fields or {}), **provider_fields}
if isinstance(provider_fields, dict)
else dict(tool_call_map[index].provider_specific_fields or {})
),
)
joined_fragments: Final = self._join_fragments_by_key_and_field(
self._iter_tool_call_fragments(tool_call_chunks)
@ -616,37 +622,37 @@ class ChunkProcessor:
# Convert the map to a list of tool calls
for index in sorted(tool_call_map.keys()):
tool_call_data = tool_call_map[index]
if tool_call_data["id"] and tool_call_data["custom_name"]:
if tool_call_data.id and tool_call_data.custom_name:
tool_calls_list.append(
ChatCompletionMessageCustomToolCall(
id=tool_call_data["id"],
id=tool_call_data.id,
custom=ChatCompletionCustomToolCallPayload(
name=tool_call_data["custom_name"],
name=tool_call_data.custom_name,
input=joined_fragments.get((index, "custom_input"), ""),
),
)
)
elif tool_call_data["id"] and tool_call_data["name"]:
elif tool_call_data.id and tool_call_data.name:
combined_arguments = joined_fragments.get((index, "arguments"), "") or "{}"
# Build function - provider_specific_fields should be on tool_call level, not function level
function = Function(
arguments=combined_arguments,
name=tool_call_data["name"],
name=tool_call_data.name,
)
# Prepare params for ChatCompletionMessageToolCall
tool_call_params = {
"id": tool_call_data["id"],
"function": function,
"type": tool_call_data["type"] or "function",
}
# Add provider_specific_fields if present (for thought signatures in Gemini 3)
if tool_call_data.get("provider_specific_fields"):
tool_call_params["provider_specific_fields"] = tool_call_data["provider_specific_fields"]
tool_call = ChatCompletionMessageToolCall(**tool_call_params)
tool_call = ChatCompletionMessageToolCall(
id=tool_call_data.id,
function=function,
type=tool_call_data.type or "function",
**(
{"provider_specific_fields": dict(tool_call_data.provider_specific_fields)}
if tool_call_data.provider_specific_fields
else {}
),
)
tool_calls_list.append(tool_call)
return tool_calls_list

6
litellm/llms/bedrock/chat/agentcore/transformation.py
View file

@ -474,7 +474,7 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM):
Returns:
AgentCoreParsedResponse: Parsed response data
"""
content_type: Final = raw_response.headers.get("content-type", "").lower()
content_type: Final = ", ".join(httpx.Headers(raw_response.headers).get_list("content-type")).lower()
verbose_logger.debug("AgentCore response Content-Type: %s", content_type)
# Parse response based on content type
@ -698,7 +698,7 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM):
)
# Check if response is JSON (agent used sync return) instead of SSE
content_type: Final = response.headers.get("content-type", "").lower()
content_type: Final = ", ".join(httpx.Headers(response.headers).get_list("content-type")).lower()
if "application/json" in content_type:
verbose_logger.debug(
"AgentCore streaming: received JSON response instead of SSE, converting to single-chunk stream"
@ -920,7 +920,7 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM):
)
# Check if response is JSON (agent used sync return) instead of SSE
content_type: Final = response.headers.get("content-type", "").lower()
content_type: Final = ", ".join(httpx.Headers(response.headers).get_list("content-type")).lower()
if "application/json" in content_type:
verbose_logger.debug(
"AgentCore streaming: received JSON response instead of SSE, converting to single-chunk stream"

19
litellm/llms/gemini/count_tokens/handler.py
View file

@ -1,6 +1,9 @@
from collections.abc import Mapping, Sequence
from typing import TYPE_CHECKING, Any, Final
import httpx
from pydantic import ConfigDict, RootModel, with_config
from typing_extensions import NotRequired, ReadOnly, TypedDict
import litellm
from litellm.llms.custom_httpx.http_handler import get_async_httpx_client
@ -12,6 +15,16 @@ else:
GenerateContentContentListUnionDict = Any
@with_config(ConfigDict(extra="allow", strict=True, hide_input_in_errors=True))
class _GeminiCountTokensResponse(TypedDict):
totalTokens: NotRequired[ReadOnly[int | float]]
tokenizer_used: NotRequired[ReadOnly[str]]
class _GeminiCountTokensResponseRoot(RootModel[_GeminiCountTokensResponse]):
model_config = ConfigDict(frozen=True, hide_input_in_errors=True)
class GoogleAIStudioTokenCounter:
def _clean_contents_for_gemini_api(self, contents: Any) -> object:
"""
@ -79,7 +92,7 @@ class GoogleAIStudioTokenCounter:
async def acount_tokens(
self,
contents: Any,
contents: Sequence[Mapping[str, object]] | None,
model: str,
api_key: str | None = None,
api_base: str | None = None,
@ -143,8 +156,8 @@ class GoogleAIStudioTokenCounter:
response.raise_for_status()
# Parse response
result: Final = response.json()
return result
result: Final = _GeminiCountTokensResponseRoot.model_validate(response.json()).root
return dict(result)
except httpx.HTTPStatusError as e:
error_msg = f"Google Gen AI Studio API error: {e.response.status_code} - {e.response.text}"

62
litellm/llms/openrouter/image_edit/transformation.py
View file

@ -41,11 +41,14 @@ Response format:
"""
import base64
from collections.abc import Mapping, Sequence
from io import BufferedReader, BytesIO
from typing import TYPE_CHECKING, Any, Final, cast
import httpx
from httpx._types import RequestFiles
from pydantic import ConfigDict, RootModel, with_config
from typing_extensions import NotRequired, ReadOnly, TypedDict
import litellm
from litellm.images.utils import ImageEditRequestUtils
@ -71,6 +74,52 @@ else:
LiteLLMLoggingObj = Any
@with_config(ConfigDict(extra="allow", strict=True, hide_input_in_errors=True))
class _ImageUrl(TypedDict):
url: ReadOnly[NotRequired[str | None]]
@with_config(ConfigDict(extra="allow", strict=True, hide_input_in_errors=True))
class _Image(TypedDict):
image_url: ReadOnly[NotRequired[_ImageUrl]]
@with_config(ConfigDict(extra="allow", strict=True, hide_input_in_errors=True))
class _Message(TypedDict):
images: ReadOnly[NotRequired[Sequence[_Image]]]
@with_config(ConfigDict(extra="allow", strict=True, hide_input_in_errors=True))
class _Choice(TypedDict):
message: ReadOnly[NotRequired[_Message]]
@with_config(ConfigDict(extra="allow", strict=True, hide_input_in_errors=True))
class _TokensDetails(TypedDict):
image_tokens: ReadOnly[NotRequired[int | float]]
@with_config(ConfigDict(extra="allow", strict=True, hide_input_in_errors=True))
class _Usage(TypedDict):
prompt_tokens: ReadOnly[NotRequired[int | float]]
total_tokens: ReadOnly[NotRequired[int | float]]
completion_tokens_details: ReadOnly[NotRequired[_TokensDetails]]
prompt_tokens_details: ReadOnly[NotRequired[_TokensDetails | None]]
cost: ReadOnly[NotRequired[int | float | str | None]]
cost_details: ReadOnly[NotRequired[Mapping[str, object] | None]]
@with_config(ConfigDict(extra="allow", strict=True, hide_input_in_errors=True))
class _OpenRouterImageEditResponse(TypedDict):
choices: ReadOnly[NotRequired[Sequence[_Choice]]]
usage: ReadOnly[NotRequired[_Usage | None]]
model: ReadOnly[NotRequired[object]]
class _OpenRouterImageEditResponseRoot(RootModel[_OpenRouterImageEditResponse]):
model_config = ConfigDict(frozen=True, hide_input_in_errors=True)
class OpenRouterImageEditConfig(BaseImageEditConfig):
"""
Configuration for OpenRouter image editing via chat completions.
@ -201,7 +250,7 @@ class OpenRouterImageEditConfig(BaseImageEditConfig):
logging_obj: LiteLLMLoggingObj,
) -> ImageResponse:
try:
response_json: Final = raw_response.json()
response_json: Final[object] = raw_response.json()
except Exception as e:
raise OpenRouterException(
message=f"Error parsing OpenRouter response: {e}",
@ -213,7 +262,8 @@ class OpenRouterImageEditConfig(BaseImageEditConfig):
model_response.data = []
try:
choices: Final = response_json.get("choices", [])
payload: Final = _OpenRouterImageEditResponseRoot.model_validate(response_json).root
choices: Final = payload.get("choices", [])
for choice in choices:
message = choice.get("message", {})
@ -252,7 +302,7 @@ class OpenRouterImageEditConfig(BaseImageEditConfig):
headers={},
)
self._set_usage_and_cost(model_response, response_json, model)
self._set_usage_and_cost(model_response, payload, model)
return model_response
def get_error_class(self, error_message: str, status_code: int, headers: dict | httpx.Headers) -> BaseLLMException:
@ -303,11 +353,11 @@ class OpenRouterImageEditConfig(BaseImageEditConfig):
def _set_usage_and_cost(
self,
model_response: ImageResponse,
response_json: dict,
payload: _OpenRouterImageEditResponse,
model: str,
) -> None:
"""Extract and set usage and cost information from OpenRouter response."""
usage_data: Final = response_json.get("usage", {})
usage_data: Final = payload.get("usage", {})
if usage_data:
prompt_tokens: Final = usage_data.get("prompt_tokens", 0)
total_tokens: Final = usage_data.get("total_tokens", 0)
@ -347,7 +397,7 @@ class OpenRouterImageEditConfig(BaseImageEditConfig):
)
response_cost_details.update(cost_details)
model_response.hidden_params["model"] = response_json.get("model", model)
model_response.hidden_params["model"] = payload.get("model", model)
def _read_image_bytes(self, image: FileTypes) -> bytes:
"""Read raw bytes from various image input types."""

24
litellm/llms/ragflow/vector_stores/transformation.py
View file

@ -2,6 +2,8 @@ from collections.abc import Mapping
from typing import TYPE_CHECKING, Any, Final
import httpx
from pydantic import ConfigDict, RootModel, with_config
from typing_extensions import NotRequired, ReadOnly, TypedDict
from litellm.llms.base_llm.vector_store.transformation import BaseVectorStoreConfig
from litellm.secret_managers.main import get_secret_str
@ -22,6 +24,25 @@ else:
LiteLLMLoggingObj = Any
@with_config(ConfigDict(extra="allow", strict=True, hide_input_in_errors=True))
class _RAGFlowDataset(TypedDict):
id: ReadOnly[NotRequired[str | None]]
name: ReadOnly[NotRequired[str | None]]
create_time: ReadOnly[NotRequired[int | float | None]]
@with_config(ConfigDict(extra="allow", strict=True, hide_input_in_errors=True))
class _RAGFlowCreateResponse(TypedDict):
data: ReadOnly[NotRequired[_RAGFlowDataset]]
class _RAGFlowCreateResponseRoot(RootModel[_RAGFlowCreateResponse]):
model_config = ConfigDict(frozen=True, hide_input_in_errors=True)
_RAGFLOW_EMPTY_DATA: Final[_RAGFlowDataset] = {}
class RAGFlowVectorStoreConfig(BaseVectorStoreConfig):
"""Vector store configuration for RAGFlow datasets."""
@ -185,7 +206,8 @@ class RAGFlowVectorStoreConfig(BaseVectorStoreConfig):
headers=response.headers,
)
data: Final = response_json.get("data", {})
create_response: Final = _RAGFlowCreateResponseRoot.model_validate(response_json).root
data: Final = create_response.get("data", _RAGFLOW_EMPTY_DATA)
# Extract dataset ID
dataset_id: Final = data.get("id")

67
litellm/llms/sagemaker/common_utils.py
View file

@ -4,7 +4,8 @@ from collections.abc import AsyncIterator, Iterator, Mapping
from typing import Final
import httpx
from pydantic import ConfigDict, TypeAdapter
from pydantic import ConfigDict, RootModel, TypeAdapter, with_config
from typing_extensions import ReadOnly, Required, TypedDict
import litellm
from litellm import verbose_logger
@ -15,6 +16,38 @@ from litellm.types.utils import StreamingChatCompletionChunk
_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True))
@with_config(ConfigDict(extra="allow", strict=True, hide_input_in_errors=True))
class _SageMakerToken(TypedDict, total=False):
text: ReadOnly[str | None]
@with_config(ConfigDict(extra="allow", strict=True, hide_input_in_errors=True))
class _SageMakerChunkData(TypedDict, total=False):
token: ReadOnly[_SageMakerToken | None]
index: ReadOnly[int | None]
@with_config(ConfigDict(extra="allow", strict=True, hide_input_in_errors=True))
class _SageMakerMessagesChunkData(TypedDict, total=False):
choices: ReadOnly[Required[list[Mapping[str, object]]]]
class _SageMakerChunkDataRoot(RootModel[_SageMakerChunkData]):
model_config = ConfigDict(frozen=True, hide_input_in_errors=True)
class _SageMakerMessagesChunkDataRoot(RootModel[_SageMakerMessagesChunkData]):
model_config = ConfigDict(frozen=True, hide_input_in_errors=True)
def _parse_sagemaker_chunk_data(data: str) -> _SageMakerChunkData:
return _SageMakerChunkDataRoot.model_validate(json.loads(data)).root
def _parse_sagemaker_messages_chunk_data(data: str) -> _SageMakerMessagesChunkData:
return _SageMakerMessagesChunkDataRoot.model_validate(json.loads(data)).root
def _load_sagemaker_response_stream_shape():
try:
from botocore.loaders import Loader
@ -62,12 +95,12 @@ class AWSEventStreamDecoder:
self.content_blocks: list = []
self.is_messages_api = is_messages_api
def _chunk_parser_messages_api(self, chunk_data: dict) -> StreamingChatCompletionChunk:
def _chunk_parser_messages_api(self, chunk_data: _SageMakerMessagesChunkData) -> StreamingChatCompletionChunk:
openai_chunk: Final = StreamingChatCompletionChunk(**{"model": self.model, **chunk_data})
return openai_chunk
def _chunk_parser(self, chunk_data: dict) -> GChunk:
def _chunk_parser(self, chunk_data: _SageMakerChunkData) -> GChunk:
verbose_logger.debug("in sagemaker chunk parser, chunk_data %s", chunk_data)
_token: Final = chunk_data.get("token", {}) or {}
_index: Final = chunk_data.get("index", None) or 0
@ -113,11 +146,12 @@ class AWSEventStreamDecoder:
# Try to parse the accumulated JSON
try:
_data = json.loads(accumulated_json)
if self.is_messages_api:
yield self._chunk_parser_messages_api(chunk_data=_data)
yield self._chunk_parser_messages_api(
chunk_data=_parse_sagemaker_messages_chunk_data(accumulated_json)
)
else:
yield self._chunk_parser(chunk_data=_data)
yield self._chunk_parser(chunk_data=_parse_sagemaker_chunk_data(accumulated_json))
# Reset accumulated_json after successful parsing
accumulated_json = ""
except json.JSONDecodeError:
@ -127,11 +161,12 @@ class AWSEventStreamDecoder:
# Handle any remaining data after the iterator is exhausted
if accumulated_json:
try:
_data = json.loads(accumulated_json)
if self.is_messages_api:
yield self._chunk_parser_messages_api(chunk_data=_data)
yield self._chunk_parser_messages_api(
chunk_data=_parse_sagemaker_messages_chunk_data(accumulated_json)
)
else:
yield self._chunk_parser(chunk_data=_data)
yield self._chunk_parser(chunk_data=_parse_sagemaker_chunk_data(accumulated_json))
except json.JSONDecodeError:
# Handle or log any unparseable data at the end
verbose_logger.error("Warning: Unparseable JSON data remained: %s", accumulated_json)
@ -161,11 +196,12 @@ class AWSEventStreamDecoder:
accumulated_json += message
# Try to parse the accumulated JSON
_data = json.loads(accumulated_json)
if self.is_messages_api:
yield self._chunk_parser_messages_api(chunk_data=_data)
yield self._chunk_parser_messages_api(
chunk_data=_parse_sagemaker_messages_chunk_data(accumulated_json)
)
else:
yield self._chunk_parser(chunk_data=_data)
yield self._chunk_parser(chunk_data=_parse_sagemaker_chunk_data(accumulated_json))
# Reset accumulated_json after successful parsing
accumulated_json = ""
except json.JSONDecodeError:
@ -181,11 +217,12 @@ class AWSEventStreamDecoder:
# Handle any remaining data after the iterator is exhausted
if accumulated_json:
try:
_data = json.loads(accumulated_json)
if self.is_messages_api:
yield self._chunk_parser_messages_api(chunk_data=_data)
yield self._chunk_parser_messages_api(
chunk_data=_parse_sagemaker_messages_chunk_data(accumulated_json)
)
else:
yield self._chunk_parser(chunk_data=_data)
yield self._chunk_parser(chunk_data=_parse_sagemaker_chunk_data(accumulated_json))
except json.JSONDecodeError:
# Handle or log any unparseable data at the end
verbose_logger.error("Warning: Unparseable JSON data remained: %s", accumulated_json)

68
litellm/llms/vertex_ai/context_caching/vertex_ai_context_caching.py
View file

@ -1,6 +1,9 @@
from collections.abc import Sequence
from typing import Final, Literal
import httpx
from pydantic import ConfigDict, RootModel, with_config
from typing_extensions import NotRequired, ReadOnly, TypedDict
import litellm
from litellm._logging import verbose_logger
@ -13,10 +16,6 @@ from litellm.llms.custom_httpx.http_handler import (
get_async_httpx_client,
)
from litellm.llms.openai.openai import AllMessageValues
from litellm.types.llms.vertex_ai import (
CachedContentListAllResponseBody,
VertexAICachedContentResponseObject,
)
from litellm.utils import is_prompt_caching_valid_prompt
from ..common_utils import VertexAIError, get_vertex_base_url
@ -32,6 +31,31 @@ local_cache_obj: Final = Cache(type=LiteLLMCacheType.LOCAL) # only used for cal
MAX_PAGINATION_PAGES: Final = 100 # Reasonable upper bound for pagination
@with_config(ConfigDict(extra="allow", strict=True, hide_input_in_errors=True))
class _CachedContentItem(TypedDict):
displayName: ReadOnly[NotRequired[object]]
name: ReadOnly[NotRequired[str | None]]
@with_config(ConfigDict(extra="allow", strict=True, hide_input_in_errors=True))
class _CachedContentListBody(TypedDict):
cachedContents: ReadOnly[NotRequired[Sequence[_CachedContentItem]]]
nextPageToken: ReadOnly[NotRequired[str | None]]
@with_config(ConfigDict(extra="allow", strict=True, hide_input_in_errors=True))
class _CreatedCachedContent(TypedDict):
name: ReadOnly[NotRequired[str | None]]
class _CachedContentListBodyRoot(RootModel[_CachedContentListBody]):
model_config = ConfigDict(frozen=True, hide_input_in_errors=True)
class _CreatedCachedContentRoot(RootModel[_CreatedCachedContent]):
model_config = ConfigDict(frozen=True, hide_input_in_errors=True)
class ContextCachingEndpoints(VertexBase):
"""
Covers context caching endpoints for Vertex AI + Google AI Studio
@ -157,25 +181,20 @@ class ContextCachingEndpoints(VertexBase):
except Exception as e:
raise VertexAIError(status_code=500, message=str(e))
raw_response = resp.json()
raw_response = _CachedContentListBodyRoot.model_validate(resp.json()).root
logging_obj.post_call(original_response=raw_response)
if "cachedContents" not in raw_response:
return None
all_cached_items = CachedContentListAllResponseBody(**raw_response)
if "cachedContents" not in all_cached_items:
return None
# Check current page for matching cache_key
for cached_item in all_cached_items["cachedContents"]:
for cached_item in raw_response["cachedContents"]:
display_name = cached_item.get("displayName")
if display_name is not None and display_name == cache_key:
return cached_item.get("name")
# Check if there are more pages
page_token = all_cached_items.get("nextPageToken")
page_token = raw_response.get("nextPageToken")
if not page_token:
# No more pages, cache not found
break
@ -249,25 +268,20 @@ class ContextCachingEndpoints(VertexBase):
except Exception as e:
raise VertexAIError(status_code=500, message=str(e))
raw_response = resp.json()
raw_response = _CachedContentListBodyRoot.model_validate(resp.json()).root
logging_obj.post_call(original_response=raw_response)
if "cachedContents" not in raw_response:
return None
all_cached_items = CachedContentListAllResponseBody(**raw_response)
if "cachedContents" not in all_cached_items:
return None
# Check current page for matching cache_key
for cached_item in all_cached_items["cachedContents"]:
for cached_item in raw_response["cachedContents"]:
display_name = cached_item.get("displayName")
if display_name is not None and display_name == cache_key:
return cached_item.get("name")
# Check if there are more pages
page_token = all_cached_items.get("nextPageToken")
page_token = raw_response.get("nextPageToken")
if not page_token:
# No more pages, cache not found
break
@ -424,14 +438,11 @@ class ContextCachingEndpoints(VertexBase):
except httpx.TimeoutException:
raise VertexAIError(status_code=408, message="Timeout error occurred.")
raw_response_cached: Final = response.json()
cached_content_response_obj: Final = VertexAICachedContentResponseObject(
name=raw_response_cached.get("name"), model=raw_response_cached.get("model")
)
raw_response_cached: Final = _CreatedCachedContentRoot.model_validate(response.json()).root
return (
non_cached_messages,
optional_params,
cached_content_response_obj["name"],
raw_response_cached.get("name"),
)
async def async_check_and_create_cache(
@ -580,14 +591,11 @@ class ContextCachingEndpoints(VertexBase):
except httpx.TimeoutException:
raise VertexAIError(status_code=408, message="Timeout error occurred.")
raw_response_cached: Final = response.json()
cached_content_response_obj: Final = VertexAICachedContentResponseObject(
name=raw_response_cached.get("name"), model=raw_response_cached.get("model")
)
raw_response_cached: Final = _CreatedCachedContentRoot.model_validate(response.json()).root
return (
non_cached_messages,
optional_params,
cached_content_response_obj["name"],
raw_response_cached.get("name"),
)
def get_cache(self):

8
litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py
View file

@ -2209,10 +2209,10 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
@staticmethod
def _set_stream_metadata_on_response(
model_response: Any,
grounding_metadata: list[dict],
url_context_metadata: list[dict],
safety_ratings: list[dict],
citation_metadata: list[dict],
grounding_metadata: Sequence[Mapping[str, object]],
url_context_metadata: Sequence[Mapping[str, object]],
safety_ratings: Sequence[Mapping[str, object]],
citation_metadata: Sequence[Mapping[str, object]],
) -> None:
setattr(model_response, "vertex_ai_grounding_metadata", grounding_metadata)
if grounding_metadata:

2
litellm/model_prices_and_context_window_backup.json
View file

@ -63565,7 +63565,7 @@
},
"azure_ai/kimi-k2.7-code": {
"cache_read_input_token_cost": 1.9e-07,
"deprecation_date": "2026-10-03",
"deprecation_date": "2027-04-16",
"input_cost_per_token": 9.5e-07,
"litellm_provider": "azure_ai",
"max_input_tokens": 262144,

76
litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py
View file

@ -5,6 +5,8 @@ from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal
import httpx
from fastapi import HTTPException
from pydantic import ConfigDict, RootModel, with_config
from typing_extensions import NotRequired, ReadOnly, TypedDict
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
@ -65,6 +67,55 @@ from litellm.types.utils import (
GUARDRAIL_NAME: Final = "model_armor"
@with_config(ConfigDict(extra="allow", strict=True, hide_input_in_errors=True))
class _MatchResult(TypedDict):
matchState: NotRequired[ReadOnly[str]]
@with_config(ConfigDict(extra="allow", strict=True, hide_input_in_errors=True))
class _DeidentifyData(TypedDict):
text: NotRequired[ReadOnly[str]]
@with_config(ConfigDict(extra="allow", strict=True, hide_input_in_errors=True))
class _DeidentifyResult(TypedDict):
matchState: NotRequired[ReadOnly[str]]
data: NotRequired[ReadOnly[_DeidentifyData]]
@with_config(ConfigDict(extra="allow", strict=True, hide_input_in_errors=True))
class _SdpResult(TypedDict):
inspectResult: NotRequired[ReadOnly[_MatchResult]]
deidentifyResult: NotRequired[ReadOnly[_DeidentifyResult]]
@with_config(ConfigDict(extra="allow", strict=True, hide_input_in_errors=True))
class _Filter(TypedDict):
raiFilterResult: NotRequired[ReadOnly[_MatchResult]]
piAndJailbreakFilterResult: NotRequired[ReadOnly[_MatchResult]]
maliciousUriFilterResult: NotRequired[ReadOnly[_MatchResult]]
csamFilterFilterResult: NotRequired[ReadOnly[_MatchResult]]
virusScanFilterResult: NotRequired[ReadOnly[_MatchResult]]
sdpFilterResult: NotRequired[ReadOnly[_SdpResult | None]]
@with_config(ConfigDict(extra="allow", strict=True, hide_input_in_errors=True))
class _SanitizationResult(TypedDict):
filterResults: NotRequired[ReadOnly[Mapping[str, _Filter] | Sequence[_Filter] | str | int | float | bool | None]]
@with_config(ConfigDict(extra="allow", strict=True, hide_input_in_errors=True))
class _ModelArmorResponse(TypedDict):
sanitizationResult: NotRequired[ReadOnly[_SanitizationResult]]
sanitizedText: NotRequired[ReadOnly[str | None]]
text: NotRequired[ReadOnly[str | None]]
class _ModelArmorResponseRoot(RootModel[_ModelArmorResponse]):
model_config = ConfigDict(frozen=True, hide_input_in_errors=True)
# Only these carry the finished output; response.created carries an empty body
_RESPONSES_TERMINAL_EVENT_TYPES: Final = frozenset({"response.completed", "response.incomplete", "response.failed"})
@ -99,10 +150,11 @@ class ModelArmorAPIError(Exception):
_SCANNED_CONTENT_KEYS: Final = frozenset({"text", "sanitizedText", "findings", "maliciousUriMatchedItems"})
RedactablePayload = dict | list | str | int | float | bool | None
def _redact_scanned_content(payload: RedactablePayload, depth: int = 0) -> RedactablePayload:
def _redact_scanned_content(payload: object, depth: int = 0) -> object:
if depth >= DEFAULT_MAX_RECURSE_DEPTH:
return "[REDACTED]"
if isinstance(payload, dict):
@ -223,12 +275,12 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase):
return f"Model Armor API error (upstream {status_code})"
return f"Model Armor API error (upstream {status_code}): {response_text}"
def _build_block_error_detail(self, message: str, armor_response: RedactablePayload) -> dict:
def _build_block_error_detail(self, message: str, armor_response: Mapping[str, object]) -> dict:
if self.sanitize_error_detail:
return {"error": message}
return {"error": message, "model_armor_response": armor_response}
def _build_logging_response(self, armor_response: RedactablePayload) -> RedactablePayload:
def _build_logging_response(self, armor_response: Mapping[str, object]) -> object:
if self.sanitize_error_detail:
return _redact_scanned_content(armor_response)
return armor_response
@ -286,7 +338,7 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase):
request_data: dict | None = None,
file_bytes: bytes | None = None,
file_type: str | None = None,
) -> dict:
) -> _ModelArmorResponse:
"""
Make request to Model Armor API. Supports both text and file prompt sanitization.
If file_bytes and file_type are provided, file prompt sanitization is performed.
@ -361,8 +413,8 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase):
json_response: Final = response.json()
if hasattr(json_response, "__await__"):
return await json_response
return json_response
return _ModelArmorResponseRoot.model_validate(await json_response).root
return _ModelArmorResponseRoot.model_validate(json_response).root
def sanitize_file_prompt(self, file_bytes: bytes, file_type: str, source: str = "user_prompt") -> dict:
"""
@ -378,7 +430,7 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase):
else:
return {"modelResponseData": {"byteItem": {"byteDataType": file_type, "byteData": base64_data}}}
def _should_block_content(self, armor_response: Mapping[str, object], allow_sanitization: bool = False) -> bool:
def _should_block_content(self, armor_response: _ModelArmorResponse, allow_sanitization: bool = False) -> bool:
"""Check if Model Armor response indicates content should be blocked, including both inspectResult and deidentifyResult."""
for filt in self._filter_result_items(armor_response):
# Check RAI, PI/Jailbreak, Malicious URI, CSAM, Virus scan as before
@ -404,7 +456,7 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase):
# Fallback dict code removed; all cases handled above
return False
def _get_sanitized_content(self, armor_response: Mapping[str, Any]) -> str | None:
def _get_sanitized_content(self, armor_response: _ModelArmorResponse) -> str | None:
"""
Get the sanitized content from a Model Armor response, if available.
Looks for sanitized text in deidentifyResult, and falls back to root-level fields if not found.
@ -434,7 +486,7 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase):
return armor_response.get("sanitizedText") or armor_response.get("text")
@staticmethod
def _filter_result_items(armor_response: Mapping[str, Any]) -> Sequence[Any]:
def _filter_result_items(armor_response: _ModelArmorResponse) -> Sequence[_Filter]:
"""Every filter result in a scan response.
filterResults is a dict of named filters on most templates and a list on some, so both
@ -447,7 +499,7 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase):
return filter_results
return []
def _has_deidentify_match(self, armor_response: Mapping[str, object]) -> bool:
def _has_deidentify_match(self, armor_response: _ModelArmorResponse) -> bool:
"""Whether an SDP de-identify filter matched, i.e. Model Armor owes this response a redaction."""
for filter_entry in self._filter_result_items(armor_response):
sdp = filter_entry.get("sdpFilterResult")
@ -457,7 +509,7 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase):
def _resolve_streaming_outcome(
self,
armor_response: Mapping[str, object],
armor_response: _ModelArmorResponse,
assembled_response: object,
content: str,
) -> tuple[bool, str | None]:
@ -489,7 +541,7 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase):
return False, sanitized_content
@staticmethod
def _append_armor_response(existing: object, armor_response: Mapping[str, object]) -> object:
def _append_armor_response(existing: object, armor_response: object) -> object:
"""Accumulate scan responses so a later text scan does not drop an earlier file scan.
Returns the single response on its own (backward compatible) and a list once a request

6
litellm/proxy/management_endpoints/model_management_endpoints.py
View file

@ -245,7 +245,7 @@ class _TransactionFactory(Protocol):
class _ModelTransactionClient(LiteLLMBaseModel):
model_config = ConfigDict(arbitrary_types_allowed=True, from_attributes=True)
model_config = ConfigDict(arbitrary_types_allowed=True, from_attributes=True, hide_input_in_errors=True)
tx: _TransactionFactory
@ -1856,8 +1856,8 @@ async def delete_team_models(
Returns the model_ids that were deleted.
"""
deleted_model_ids: Final[list[str]] = []
async with prisma_client.db.tx() as tx_ctx:
tx: Final[_TxModelTables] = tx_ctx
transaction_client: Final = _ModelTransactionClient.model_validate(prisma_client.db)
async with transaction_client.tx() as tx:
for team_id in team_ids:
rows = await _get_team_deployments(team_id, prisma_client, table=tx.litellm_proxymodeltable)
model_ids = [row.model_id for row in rows]

12
litellm/proxy/management_endpoints/ui_sso.py
View file

@ -41,7 +41,7 @@ if TYPE_CHECKING:
import jwt
from fastapi import APIRouter, Depends, Header, HTTPException, Request, Response, status
from fastapi.responses import RedirectResponse
from pydantic import TypeAdapter, ValidationError
from pydantic import ConfigDict, TypeAdapter, ValidationError
import litellm
from litellm._internal_context import with_service_target
@ -237,6 +237,12 @@ class CliSsoTeamDetail(LiteLLMBaseModel):
_CLI_SSO_TEAM_DETAILS_ADAPTER: Final = TypeAdapter(tuple[CliSsoTeamDetail, ...])
_CLI_SSO_FLOW_JSON_ADAPTER: Final[TypeAdapter[dict[str, object]]] = TypeAdapter(
dict[str, object], config=ConfigDict(strict=True, hide_input_in_errors=True)
)
_SSO_USERINFO_RESPONSE_ADAPTER: Final[TypeAdapter[dict[str, object] | None]] = TypeAdapter(
dict[str, object] | None, config=ConfigDict(strict=True, hide_input_in_errors=True)
)
_TEAMLESS_CLI_SSO_TEAM_DETAIL: Final = CliSsoTeamDetail(team_models=())
@ -371,7 +377,7 @@ def _get_cli_sso_flow_or_raise(login_id: str | None, cache: DualCache) -> dict:
flow = _read_cli_sso_flow(cache, _get_cli_sso_flow_cache_key(cast(str, login_id)))
if isinstance(flow, str):
try:
flow = _as_object(json.loads(flow))
flow = _CLI_SSO_FLOW_JSON_ADAPTER.validate_python(json.loads(flow))
except ValueError:
flow = None
if not isinstance(flow, dict) or "poll_secret_hash" not in flow:
@ -4199,7 +4205,7 @@ class SSOAuthenticationHandler:
)
if resp.status_code == 200:
try:
userinfo_raw: Final[dict[str, object] | None] = resp.json()
userinfo_raw: Final = _SSO_USERINFO_RESPONSE_ADAPTER.validate_python(resp.json())
if not userinfo_raw:
# JSON null (None) or empty dict ({}) — no identity claims.
# Treat as failure so id_token fallback can be attempted.

14
litellm/proxy/public_endpoints/public_endpoints.py
View file

@ -591,8 +591,8 @@ async def _fetch_remote_whats_new(url: str) -> WhatsNewResponse:
return WhatsNewResponse.model_validate_json(await _fetch_remote_bytes(url))
async def _resolve_whats_new(url: str, fetch: Callable[[str], Awaitable[WhatsNewResponse]]) -> WhatsNewResponse:
if os.getenv("LITELLM_LOCAL_WHATS_NEW", "").lower() == "true":
async def _resolve_whats_new(url: str | None, fetch: Callable[[str], Awaitable[WhatsNewResponse]]) -> WhatsNewResponse:
if url is None or os.getenv("LITELLM_LOCAL_WHATS_NEW", "").lower() == "true":
return _load_bundled_whats_new()
try:
return await fetch(url)
@ -611,7 +611,7 @@ class _WhatsNewCache:
async def get_whats_new(
url: str,
url: str | None,
fetch: Callable[[str], Awaitable[WhatsNewResponse]] = _fetch_remote_whats_new,
) -> WhatsNewResponse:
cached: Final = _WhatsNewCache.launches
@ -637,10 +637,10 @@ async def get_public_whats_new() -> WhatsNewResponse:
"""
Return the launches the dashboard Home page shows under What's new.
Resolved once per process: fetched from ``litellm.whats_new_url`` (override with ``LITELLM_WHATS_NEW_URL``)
on the first request, falling back to the list bundled with the package on any failure, so an airgapped
proxy serves the bundled list without retrying. Set ``LITELLM_LOCAL_WHATS_NEW=True`` to skip the fetch.
A restart picks up a newly published list.
Serves the list bundled with this release. Set ``LITELLM_WHATS_NEW_URL`` to fetch a remote list instead,
resolved once per process on the first request and falling back to the bundled list on any failure, so an
airgapped proxy serves the bundled list without retrying. ``LITELLM_LOCAL_WHATS_NEW=True`` forces the
bundled list even when a URL is set. A restart picks up a newly published list.
"""
return await get_whats_new(url=litellm.whats_new_url)

2
model_prices_and_context_window.json
View file

@ -63565,7 +63565,7 @@
},
"azure_ai/kimi-k2.7-code": {
"cache_read_input_token_cost": 1.9e-07,
"deprecation_date": "2026-10-03",
"deprecation_date": "2027-04-16",
"input_cost_per_token": 9.5e-07,
"litellm_provider": "azure_ai",
"max_input_tokens": 262144,

157
tests/agent_tests/local_only_agent_tests/local_vertex_agent.py
View file

@ -1,157 +0,0 @@
"""
Test script for Vertex AI Reasoning Engine.
This script demonstrates how to:
1. Authenticate with Google Cloud
2. Send queries to a Vertex AI Reasoning Engine using the :query endpoint
Usage:
python local_vertex_agent.py
Requirements:
pip install httpx google-auth
"""
import asyncio
import json
from uuid import uuid4
from google.auth import default
from google.auth.transport.requests import Request
import httpx
# Configuration - update these for your agent
PROJECT_ID = "test-gcp-project-id-123" # Your GCP project ID (test value)
LOCATION = "us-central1" # Your agent's location
# For Reasoning Engines, use just the numeric ID at the end
REASONING_ENGINE_ID = "8263861224643493888"
# The project number from the resource name
PROJECT_NUMBER = "1060139831167"
async def main():
"""Main function to test Vertex AI Reasoning Engine."""
# Step 1: Authenticate with Google Cloud
print("Step 1: Authenticating with Google Cloud...")
credentials, project = default(
scopes=["https://www.googleapis.com/auth/cloud-platform"]
)
credentials.refresh(Request())
print(f"Authenticated! Project: {project}")
print(f"Token (first 20 chars): {credentials.token[:20]}...")
# Step 2: Build the endpoint URL
base_url = f"https://{LOCATION}-aiplatform.googleapis.com"
resource_path = f"projects/{PROJECT_NUMBER}/locations/{LOCATION}/reasoningEngines/{REASONING_ENGINE_ID}"
# The Reasoning Engine uses :query endpoint with specific format
query_url = f"{base_url}/v1beta1/{resource_path}:query"
stream_url = f"{base_url}/v1beta1/{resource_path}:streamQuery"
print(f"\nQuery URL: {query_url}")
print(f"Stream URL: {stream_url}")
# Step 3: Create authenticated httpx client
print("\nStep 2: Creating authenticated HTTP client...")
client = httpx.AsyncClient(
headers={
"Authorization": f"Bearer {credentials.token}",
"Content-Type": "application/json",
},
timeout=120.0,
)
# Step 4: Build the query request (non-streaming)
# Note: For non-streaming, we need to:
# 1. Create a session
# 2. Use the streaming endpoint with stream_query method
# The :query endpoint only supports session management methods
user_id = f"test-user-{uuid4().hex[:8]}"
# First create a session
create_session_request = {
"class_method": "async_create_session",
"input": {
"user_id": user_id,
},
}
print(f"\nStep 3: Creating session...")
print(f"User ID: {user_id}")
async with client:
# Create session
print(f"\nSending to: {query_url}")
response = await client.post(query_url, json=create_session_request)
print(f"Create session status: {response.status_code}")
if response.status_code == 200:
session_data = response.json()
print(f"Session created:\n{json.dumps(session_data, indent=2)}")
# Extract session_id from response
session_id = session_data.get("output", {}).get("id") or session_data.get(
"output", {}
).get("session_id")
print(f"\nSession ID: {session_id}")
# Now send the actual query via streamQuery
query_request = {
"class_method": "stream_query",
"input": {
"message": "Hello! What can you do?",
"user_id": user_id,
"session_id": session_id,
},
}
print(f"\nStep 4: Sending query via streamQuery...")
print(f"Request:\n{json.dumps(query_request, indent=2)}")
# Use streaming endpoint but collect full response
async with client.stream(
"POST", stream_url, json=query_request
) as stream_response:
print(f"Query status: {stream_response.status_code}")
if stream_response.status_code == 200:
print("\nResponse:")
full_response = ""
async for line in stream_response.aiter_lines():
if line:
full_response = line # Keep last line (full response)
# Parse and display
try:
data = json.loads(full_response)
# Extract the text from the response
content = data.get("content", {})
parts = content.get("parts", [])
for part in parts:
if "text" in part:
print(f"\nAgent response:\n{part['text']}")
except:
print(full_response)
else:
content = await stream_response.aread()
print(f"Error: {content.decode()}")
else:
print(f"Error creating session: {response.text}")
if __name__ == "__main__":
print("=" * 60)
print("Vertex AI Reasoning Engine Test Script")
print("=" * 60)
print(f"\nConfiguration:")
print(f" PROJECT_ID: {PROJECT_ID}")
print(f" PROJECT_NUMBER: {PROJECT_NUMBER}")
print(f" LOCATION: {LOCATION}")
print(f" REASONING_ENGINE_ID: {REASONING_ENGINE_ID}")
print()
asyncio.run(main())

353
tests/agent_tests/local_only_agent_tests/test_a2a.py
View file

@ -1,353 +0,0 @@
"""
Test for LiteLLM A2A module.
Run with:
pytest tests/agent_tests/test_a2a.py -v -s
"""
import asyncio
import json
from typing import Optional
from uuid import uuid4
import pytest
import litellm
from litellm.integrations.custom_logger import CustomLogger
from litellm.types.utils import StandardLoggingPayload
from a2a.types import MessageSendParams, SendMessageRequest
@pytest.mark.asyncio
async def test_asend_message_with_client_decorator():
"""
Test asend_message standalone function with @client decorator.
This tests the LiteLLM logging integration.
"""
litellm.turn_on_debug()
from litellm.a2a_protocol import asend_message, create_a2a_client
# Create the A2A client first
a2a_client = await create_a2a_client(base_url="http://localhost:10001")
# Build the request matching A2A SDK spec
send_message_payload = {
"message": {
"role": "user",
"parts": [
{
"kind": "text",
"text": "Hello from @client decorated function!",
}
],
"messageId": uuid4().hex,
},
}
request = SendMessageRequest(
id=str(uuid4()),
params=MessageSendParams(**send_message_payload),
)
# Send message using standalone function with @client decorator
response = await asend_message(a2a_client=a2a_client, request=request)
# Print response for debugging
print("\n=== A2A Response (standalone with @client) ===")
print(response.model_dump(mode="json", exclude_none=True))
# Basic assertions
assert response is not None
class TestA2ALogger(CustomLogger):
"""Custom logger to capture A2A logging payloads for testing."""
def __init__(self):
self.standard_logging_payload: Optional[StandardLoggingPayload] = None
self.logged_kwargs: Optional[dict] = None
self.log_success_called = False
super().__init__()
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
print("TestA2ALogger: async_log_success_event called")
self.log_success_called = True
self.logged_kwargs = kwargs
self.standard_logging_payload = kwargs.get("standard_logging_object", None)
print(f"Captured standard_logging_payload: {self.standard_logging_payload}")
@pytest.mark.asyncio
async def test_a2a_logging_payload():
"""
Test that A2A calls create a standard logging payload.
Validates the @client decorator integration with LiteLLM logging.
"""
# Reset callbacks and set up custom logger
litellm.logging_callback_manager._reset_all_callbacks()
test_logger = TestA2ALogger()
litellm.callbacks = [test_logger]
from litellm.a2a_protocol import asend_message, create_a2a_client
# Create the A2A client first
a2a_client = await create_a2a_client(base_url="http://localhost:10001")
# Build the request
send_message_payload = {
"message": {
"role": "user",
"parts": [
{
"kind": "text",
"text": "Hello! Testing logging payload.",
}
],
"messageId": uuid4().hex,
},
}
request = SendMessageRequest(
id=str(uuid4()),
params=MessageSendParams(**send_message_payload),
)
# Send message
response = await asend_message(a2a_client=a2a_client, request=request)
# Give async logging time to complete
await asyncio.sleep(1)
# Print debug info
print("\n=== Logging Validation ===")
print(f"log_success_called: {test_logger.log_success_called}")
print(f"standard_logging_payload: {test_logger.standard_logging_payload}")
print(
f"logged kwargs: {json.dumps(test_logger.logged_kwargs, indent=4, default=str)}"
)
# Verify logging was called
assert test_logger.log_success_called is True
assert test_logger.standard_logging_payload is not None
# Verify standard_logging_payload exists
slp = test_logger.standard_logging_payload
assert slp is not None
# Get values from standard logging payload
logged_model = (
slp.get("model") if isinstance(slp, dict) else getattr(slp, "model", None)
)
logged_provider = (
slp.get("custom_llm_provider")
if isinstance(slp, dict)
else getattr(slp, "custom_llm_provider", None)
)
call_type = (
slp.get("call_type")
if isinstance(slp, dict)
else getattr(slp, "call_type", None)
)
response_cost = (
slp.get("response_cost")
if isinstance(slp, dict)
else getattr(slp, "response_cost", None)
)
print(f"\n=== Standard Logging Payload Validation ===")
print(f"model: {logged_model}")
print(f"custom_llm_provider: {logged_provider}")
print(f"call_type: {call_type}")
print(f"response_cost: {response_cost}")
# Verify model and custom_llm_provider are set correctly
assert logged_model is not None, "model should be set"
assert (
"a2a_agent/" in logged_model
), f"model should contain 'a2a_agent/', got: {logged_model}"
assert (
logged_provider == "a2a_agent"
), f"custom_llm_provider should be 'a2a_agent', got: {logged_provider}"
# Verify call_type is correct for A2A
assert (
call_type == "asend_message"
), f"call_type should be 'asend_message', got: {call_type}"
# Verify response_cost is set to 0.0 (not None, not an error)
# This confirms the A2A cost calculator is working
assert response_cost is not None, "response_cost should not be None"
assert (
response_cost == 0.0
), f"response_cost should be 0.0 for A2A, got: {response_cost}"
@pytest.mark.asyncio
async def test_pydantic_ai_non_streaming():
"""
Test non-streaming requests to Pydantic AI agents.
Pydantic AI agents follow A2A protocol but don't support streaming.
This test validates non-streaming requests work correctly.
"""
litellm.turn_on_debug()
from litellm.a2a_protocol import asend_message
# Build the request
send_message_payload = {
"message": {
"role": "user",
"parts": [
{
"kind": "text",
"text": "Hello from Pydantic AI test!",
}
],
"messageId": uuid4().hex,
},
}
request = SendMessageRequest(
id=str(uuid4()),
params=MessageSendParams(**send_message_payload),
)
# Send message using Pydantic AI provider
response = await asend_message(
request=request,
api_base="http://localhost:9999",
litellm_params={"custom_llm_provider": "pydantic_ai_agents"},
)
# Print response for debugging
print("\n=== Pydantic AI Non-Streaming Response ===")
print(response.model_dump(mode="json", exclude_none=True))
# Basic assertions
assert response is not None
assert hasattr(response, "result")
# Verify result structure
result = response.result
assert result is not None
# Pydantic AI returns a task with history/artifacts, not a direct message
# Check for either format
result_dict = (
result
if isinstance(result, dict)
else result.model_dump(mode="python", exclude_none=True)
)
has_message = "message" in result_dict
has_history = "history" in result_dict
has_artifacts = "artifacts" in result_dict
assert (
has_message or has_history or has_artifacts
), f"Result should contain 'message', 'history', or 'artifacts'. Got: {list(result_dict.keys())}"
# If it's a task response (Pydantic AI style), verify we got agent response
if has_history:
history = result_dict.get("history", [])
agent_messages = [m for m in history if m.get("role") == "agent"]
assert (
len(agent_messages) > 0
), "Should have at least one agent message in history"
# Verify agent message has text content
agent_msg = agent_messages[-1]
parts = agent_msg.get("parts", [])
text_parts = [p for p in parts if p.get("kind") == "text"]
assert len(text_parts) > 0, "Agent message should have text content"
print(f"\nAgent response: {text_parts[0].get('text')}")
@pytest.mark.asyncio
async def test_pydantic_ai_fake_streaming():
"""
Test fake streaming for Pydantic AI agents.
Pydantic AI agents don't support streaming natively.
This test validates that fake streaming works by converting
non-streaming responses into streaming chunks.
"""
litellm.turn_on_debug()
from litellm.a2a_protocol import asend_message_streaming
# Build the request
from a2a.types import SendStreamingMessageRequest
send_message_payload = {
"message": {
"role": "user",
"parts": [
{
"kind": "text",
"text": "Hello from Pydantic AI streaming test!",
}
],
"messageId": uuid4().hex,
},
}
request = SendStreamingMessageRequest(
id=str(uuid4()),
params=MessageSendParams(**send_message_payload),
)
# Send streaming message using Pydantic AI provider
print("\n=== Pydantic AI Fake Streaming Response ===")
chunks_received = 0
task_event_received = False
working_event_received = False
artifact_event_received = False
completed_event_received = False
async for chunk in asend_message_streaming(
request=request,
api_base="http://localhost:9999",
litellm_params={"custom_llm_provider": "pydantic_ai_agents"},
):
chunks_received += 1
print(f"\nChunk {chunks_received}:")
# Convert chunk to dict for inspection
chunk_dict = (
chunk.model_dump(mode="json", exclude_none=True)
if hasattr(chunk, "model_dump")
else chunk
)
print(json.dumps(chunk_dict, indent=2))
# Check event types
result = chunk_dict.get("result", {})
kind = result.get("kind")
if kind == "task":
task_event_received = True
elif kind == "status-update":
status = result.get("status", {})
state = status.get("state")
if state == "working":
working_event_received = True
elif state == "completed":
completed_event_received = True
elif kind == "artifact-update":
artifact_event_received = True
print(f"\n=== Streaming Summary ===")
print(f"Total chunks received: {chunks_received}")
print(f"Task event received: {task_event_received}")
print(f"Working event received: {working_event_received}")
print(f"Artifact event received: {artifact_event_received}")
print(f"Completed event received: {completed_event_received}")
# Verify we received chunks
assert chunks_received > 0, "Should receive at least one chunk"
# Verify all required event types were received
assert task_event_received, "Should receive task event"
assert working_event_received, "Should receive working status event"
assert artifact_event_received, "Should receive artifact update event"
assert completed_event_received, "Should receive completed status event"

287
tests/agent_tests/local_only_agent_tests/test_a2a_completion_bridge.py
View file

@ -1,287 +0,0 @@
"""
Test for A2A to LiteLLM Completion Bridge.
Tests the SDK-level functions that route A2A requests through litellm.acompletion.
Run with:
pytest tests/agent_tests/test_a2a_completion_bridge.py -v -s
Prerequisites:
- LangGraph server running on localhost:2024
"""
from uuid import uuid4
import pytest
import litellm
from a2a.types import MessageSendParams, SendMessageRequest, SendStreamingMessageRequest
@pytest.mark.asyncio
async def test_a2a_completion_bridge_non_streaming():
"""
Test non-streaming A2A request via the completion bridge with LangGraph provider.
"""
from litellm.a2a_protocol import asend_message
litellm.turn_on_debug()
send_message_payload = {
"message": {
"role": "user",
"parts": [{"kind": "text", "text": "What is 2 + 2?"}],
"messageId": uuid4().hex,
}
}
request = SendMessageRequest(
id=str(uuid4()),
params=MessageSendParams(**send_message_payload), # type: ignore
)
response = await asend_message(
request=request,
api_base="http://localhost:2024",
litellm_params={"custom_llm_provider": "langgraph", "model": "agent"},
)
# Validate response is LiteLLMSendMessageResponse
assert response.jsonrpc == "2.0"
assert response.id is not None
assert response.result is not None
assert response.result.get("kind") == "message"
message = response.result
assert "role" in message
assert message["role"] == "agent"
assert "parts" in message
assert len(message["parts"]) > 0
assert message["parts"][0]["kind"] == "text"
assert len(message["parts"][0]["text"]) > 0
print(f"Response: {response.model_dump(mode='json', exclude_none=True)}")
@pytest.mark.asyncio
async def test_a2a_completion_bridge_streaming():
"""
Test streaming A2A request via the completion bridge with LangGraph provider.
Validates proper A2A streaming format with events:
1. Task event (kind: "task") - Initial task with status "submitted"
2. Status update (kind: "status-update") - Status "working"
3. Artifact update (kind: "artifact-update") - Content delivery
4. Status update (kind: "status-update") - Final "completed" status
"""
from litellm.a2a_protocol import asend_message_streaming
litellm.turn_on_debug()
send_message_payload = {
"message": {
"role": "user",
"parts": [{"kind": "text", "text": "Count from 1 to 5."}],
"messageId": uuid4().hex,
}
}
request = SendStreamingMessageRequest(
id=str(uuid4()),
params=MessageSendParams(**send_message_payload), # type: ignore
)
chunks = []
async for chunk in asend_message_streaming(
request=request,
api_base="http://localhost:2024",
litellm_params={"custom_llm_provider": "langgraph", "model": "agent"},
):
chunks.append(chunk)
print(f"Chunk: {chunk}")
# Validate we received proper A2A streaming events
assert (
len(chunks) >= 4
), f"Expected at least 4 chunks (task, working, artifact, completed), got {len(chunks)}"
# Validate chunk structure follows A2A spec
for chunk in chunks:
assert "jsonrpc" in chunk
assert chunk["jsonrpc"] == "2.0"
assert "id" in chunk
assert "result" in chunk
# Validate first chunk is task event
task_chunk = chunks[0]
assert task_chunk["result"]["kind"] == "task", "First chunk should be task event"
assert task_chunk["result"]["status"]["state"] == "submitted"
assert "contextId" in task_chunk["result"]
assert "id" in task_chunk["result"] # task id
assert "history" in task_chunk["result"]
# Validate second chunk is working status update
working_chunk = chunks[1]
assert (
working_chunk["result"]["kind"] == "status-update"
), "Second chunk should be status-update"
assert working_chunk["result"]["status"]["state"] == "working"
assert "taskId" in working_chunk["result"]
assert "contextId" in working_chunk["result"]
assert working_chunk["result"]["final"] is False
# Validate artifact update chunk
artifact_chunk = chunks[2]
assert (
artifact_chunk["result"]["kind"] == "artifact-update"
), "Third chunk should be artifact-update"
assert "artifact" in artifact_chunk["result"]
assert "artifactId" in artifact_chunk["result"]["artifact"]
assert "parts" in artifact_chunk["result"]["artifact"]
assert len(artifact_chunk["result"]["artifact"]["parts"]) > 0
assert artifact_chunk["result"]["artifact"]["parts"][0]["kind"] == "text"
# Validate final chunk is completed status update
final_chunk = chunks[-1]
assert (
final_chunk["result"]["kind"] == "status-update"
), "Last chunk should be status-update"
assert final_chunk["result"]["status"]["state"] == "completed"
assert final_chunk["result"]["final"] is True
print(f"Received {len(chunks)} chunks with proper A2A streaming format")
@pytest.mark.asyncio
async def test_a2a_completion_bridge_bedrock_agentcore():
"""
Test A2A request via the completion bridge with Bedrock AgentCore provider.
Uses the AgentCore runtime ARN to call a hosted agent.
"""
from litellm.a2a_protocol import asend_message_streaming
litellm.turn_on_debug()
# Bedrock AgentCore ARN (streaming-capable runtime)
agentcore_arn = "arn:aws:bedrock-agentcore:us-west-2:888602223428:runtime/hosted_agent_r9jvp-3ySZuRHjLC"
send_message_payload = {
"message": {
"role": "user",
"parts": [
{"kind": "text", "text": "Explain machine learning in simple terms"}
],
"messageId": uuid4().hex,
}
}
request = SendStreamingMessageRequest(
id=str(uuid4()),
params=MessageSendParams(**send_message_payload), # type: ignore
)
chunks = []
async for chunk in asend_message_streaming(
request=request,
api_base=None, # Not needed for Bedrock AgentCore
litellm_params={
"custom_llm_provider": "bedrock",
"model": f"bedrock/agentcore/{agentcore_arn}",
},
):
chunks.append(chunk)
print(f"Chunk: {chunk}")
# Validate we received proper A2A streaming events
assert len(chunks) >= 4, f"Expected at least 4 chunks, got {len(chunks)}"
# Validate first chunk is task event
assert chunks[0]["result"]["kind"] == "task"
assert chunks[0]["result"]["status"]["state"] == "submitted"
# Validate final chunk is completed status
assert chunks[-1]["result"]["kind"] == "status-update"
assert chunks[-1]["result"]["status"]["state"] == "completed"
assert chunks[-1]["result"]["final"] is True
print(f"Received {len(chunks)} chunks from Bedrock AgentCore")
# ============================================================
# Vertex AI Agent Engine Tests
# ============================================================
# Configuration - update these for your Vertex AI Reasoning Engine
VERTEX_AGENT_RESOURCE_NAME = (
"projects/1060139831167/locations/us-central1/reasoningEngines/8263861224643493888"
)
@pytest.mark.asyncio
async def test_vertex_agent_engine_non_streaming():
"""
Test non-streaming request to Vertex AI Agent Engine via litellm.acompletion.
Uses the Reasoning Engine resource ID to call a hosted agent.
"""
litellm.turn_on_debug()
# Call via litellm.acompletion with vertex_ai/agent_engine/ prefix
response = await litellm.acompletion(
model=f"vertex_ai/agent_engine/{VERTEX_AGENT_RESOURCE_NAME}",
messages=[{"role": "user", "content": "Hello! What can you do?"}],
stream=False,
)
print(f"\n=== Vertex Agent Engine Non-Streaming Response ===")
print(f"Response: {response}")
# Basic assertions
assert response is not None
assert hasattr(response, "choices")
assert len(response.choices) > 0
assert response.choices[0].message is not None
assert response.choices[0].message.content is not None
assert len(response.choices[0].message.content) > 0
print(f"Agent response: {response.choices[0].message.content[:200]}...")
@pytest.mark.asyncio
async def test_vertex_agent_engine_streaming():
"""
Test streaming request to Vertex AI Agent Engine via litellm.acompletion.
Uses the Reasoning Engine resource ID to call a hosted agent with streaming.
"""
# litellm.turn_on_debug()
# Call via litellm.acompletion with streaming
response = await litellm.acompletion(
model=f"vertex_ai/agent_engine/{VERTEX_AGENT_RESOURCE_NAME}",
messages=[{"role": "user", "content": "Hello! What can you do?"}],
stream=True,
)
print(f"\n=== Vertex Agent Engine Streaming Response ===")
chunks = []
full_content = ""
async for chunk in response:
print(f"Chunk: {chunk}")
# chunks.append(chunk)
# if hasattr(chunk, "choices") and len(chunk.choices) > 0:
# delta = chunk.choices[0].delta
# if hasattr(delta, "content") and delta.content:
# full_content += delta.content
# print(f"Chunk: {delta.content}", end="", flush=True)
# # print(f"\n\nReceived {len(chunks)} chunks")
# print(f"Full content: {full_content[:200]}...")
# # Basic assertions
# assert len(chunks) > 0
# assert len(full_content) > 0

54
tests/basic_proxy_startup_tests/test_basic_proxy_startup.py
View file

@ -1,54 +0,0 @@
"""
This test ensures that the proxy starts and serves requests even with a bad license.
in ci/cd config.yml, we set the license to "bad-license"
"""
import os
import pytest
import aiohttp
from typing import Optional
@pytest.mark.asyncio
async def test_health_and_chat_completion():
"""
Test health endpoints and chat completion:
1. Check /health/readiness
2. Check /health/liveness
3. Make a chat completion call
"""
async with aiohttp.ClientSession() as session:
# Test readiness endpoint
async with session.get("http://0.0.0.0:4000/health/readiness") as response:
assert response.status == 200
readiness_response = await response.json()
# Accept both "healthy" (new format) and "connected" (legacy format)
# since this test runs against both source builds and packaged installs
assert readiness_response["status"] in ("healthy", "connected")
# Test liveness endpoint
async with session.get("http://0.0.0.0:4000/health/liveness") as response:
assert response.status == 200
liveness_response = await response.json()
print("liveness_response", liveness_response)
# Make a chat completion call
url = "http://0.0.0.0:4000/chat/completions"
headers = {
"Authorization": f"Bearer {os.environ['LITELLM_MASTER_KEY']}",
"Content-Type": "application/json",
}
data = {
"model": "gpt-4",
"messages": [
{"role": "system", "content": "You are a helpful assistant."},
{"role": "user", "content": "Hello!"},
],
}
async with session.post(url, headers=headers, json=data) as response:
assert response.status == 200
completion_response = await response.json()
assert "choices" in completion_response

2
tests/code_coverage_tests/code_qa_check_tests.py
View file

@ -174,7 +174,7 @@ def main():
f"This can cause import issues and test failures. Files: {problematic_files}"
)
else:
print("✓ No litellm module deletion patterns found in tests/test_litellm or tests/unit.")
print("✓ No litellm module deletion patterns found in tests/unit.")
if __name__ == "__main__":

4
tests/code_coverage_tests/router_code_coverage.py
View file

@ -27,11 +27,11 @@ def get_function_names_from_file(file_path):
def get_all_functions_called_in_tests(base_dir):
"""
Returns a set of function names that are called in test functions
inside 'local_testing' and 'router_unit_test' directories,
inside the unit test directories,
specifically in files containing the word 'router'.
"""
called_functions = set()
test_dirs = ["local_testing", "router_unit_tests", "test_litellm", "unit"]
test_dirs = ["unit"]
for test_dir in test_dirs:
dir_path = os.path.join(base_dir, test_dir)

2
tests/e2e/CONTRIBUTING.md
View file

@ -260,7 +260,7 @@ The semantic header set is `content-type`, `accept`, `anthropic-version`, `anthr
Excluded transport and telemetry headers are `host`, `content-length`, `connection`, `accept-encoding`, `user-agent`, `traceparent`, `tracestate`, `x-request-id`, `x-client-request-id` and `x-stainless-*`. Inbound transfer-encoding is unsupported; send JSON with content-length framing. The destination represents host identity and the relay carries original body bytes. Replay does not verify credentials, SDK timeout/retry behavior, transport performance, model availability or stateful remote IDs. Live relay uses original request bytes and header values, never the stored identity
Strict replay harness regression tests live in `tests/code_coverage_tests/test_provider_replay_harness.py`. The CircleCI `provider_replay_harness` job runs them alongside the provider-edge and fixture tests at the root of `tests/e2e_harness/` with `--noconftest -o "pythonpath=tests/e2e tests/e2e_harness"`; they need only synthetic HTTP providers and temporary fixture storage
Strict replay harness regression tests live in `tests/code_coverage_tests/test_provider_replay_harness.py`. The `python` job in `.github/workflows/test-linting.yml` runs them alongside the provider-edge and fixture tests at the root of `tests/e2e_harness/` with `--noconftest -o addopts= -o "pythonpath=tests/e2e tests/e2e_harness" -p no:rerunfailures`; they need only synthetic HTTP providers and temporary fixture storage
## MCP OAuth happy path

19
tests/e2e/migrations/checks.py
View file

@ -39,6 +39,25 @@ def assert_completed(database: Database, migration: Migration = COMPLETE) -> Non
assert database.query("SELECT id FROM migration_effect") == ((1,),)
@step("Check that migration {migration.name} rolled back lock timeouts and completed exactly once")
def assert_completed_after_lock_timeouts(database: Database, migration: Migration) -> None:
rows: Final = database.query(
"SELECT finished_at IS NOT NULL, rolled_back_at IS NOT NULL, applied_steps_count FROM "
"_prisma_migrations WHERE migration_name = %s ORDER BY started_at",
(migration.name,),
)
finished: Final = tuple(row for row in rows if row[0])
rolled_back: Final = tuple(row for row in rows if row[1])
assert finished == ((True, False, 1),), "Expected exactly one completed migration attempt"
assert rolled_back and all(row == (False, True, 0) for row in rolled_back), (
"Expected every failed migration attempt to be rolled back"
)
assert all(is_finished != is_rolled_back for is_finished, is_rolled_back, _ in rows), (
"Expected every migration attempt to be either completed or rolled back"
)
assert database.query("SELECT id FROM migration_effect") == ((1,),)
@step("Apply the test migration by hand and record it in _prisma_migrations")
def confirmed_history(database: Database) -> str:
database.execute(COMPLETE_SQL)

10
tests/e2e/migrations/containers.py
View file

@ -41,6 +41,7 @@ def until(description: str, condition: Callable[[], bool], seconds: float = 150)
@dataclass(frozen=True, slots=True)
class Replica:
name: str
identity: str
transport: HttpTransport
output: Path
@ -171,7 +172,7 @@ class Containers:
"model_list: []\ngeneral_settings:\n master_key: os.environ/LITELLM_MASTER_KEY\n"
)
env: Final = {
"DATABASE_URL": prisma_url(database.container_url, database.schema),
"DATABASE_URL": prisma_url(database.container_url, database.schema, application_name=name),
"LITELLM_MASTER_KEY": MASTER_KEY,
"LITELLM_SALT_KEY": MASTER_KEY,
"LITELLM_LOCAL_MODEL_COST_MAP": "True",
@ -209,7 +210,12 @@ class Containers:
"4000",
)
port: Final = int(docker("port", name, "4000/tcp").rsplit(":", 1)[1])
replica: Final = Replica(name, HttpTransport(f"http://127.0.0.1:{port}", MASTER_KEY, 15), directory)
replica: Final = Replica(
name=name,
identity=name,
transport=HttpTransport(f"http://127.0.0.1:{port}", MASTER_KEY, 15),
output=directory,
)
yield replica
finally:
try:

37
tests/e2e/migrations/database.py
View file

@ -3,6 +3,7 @@ from __future__ import annotations
from collections.abc import Generator
from contextlib import contextmanager
from dataclasses import dataclass
from itertools import chain
from typing import Final, LiteralString
from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit
from uuid import uuid4
@ -24,10 +25,19 @@ def connect_url(url: str, name: str) -> str:
return urlunsplit(urlsplit(url)._replace(path=f"/{name}", query=""))
def prisma_url(url: str, schema: str) -> str:
def prisma_url(url: str, schema: str, application_name: str | None = None) -> str:
parsed: Final = urlsplit(url)
query: Final = tuple((key, value) for key, value in parse_qsl(parsed.query) if key != "schema")
return urlunsplit(parsed._replace(query=urlencode((*query, ("schema", schema)))))
query: Final = tuple(
(key, value)
for key, value in parse_qsl(parsed.query)
if key != "schema" and (application_name is None or key != "application_name")
)
parameters: Final = query + (
(("schema", schema),)
if application_name is None
else (("schema", schema), ("application_name", application_name))
)
return urlunsplit(parsed._replace(query=urlencode(parameters)))
@dataclass(frozen=True, slots=True)
@ -79,6 +89,27 @@ class Database:
(key >> 32, key & 0xFFFFFFFF),
)
@step("List replica identities waiting on any of the advisory locks")
def waiting_identities(self, keys: tuple[int, ...]) -> frozenset[str]:
if not keys:
return frozenset()
key_pairs: Final = tuple((key >> 32, key & 0xFFFFFFFF) for key in keys)
key_predicates: Final = sql.SQL(" OR ").join(
sql.SQL("(waiting_lock.classid = %s AND waiting_lock.objid = %s)") for _ in key_pairs
)
parameters: Final = tuple(chain.from_iterable(key_pairs))
statement: Final = sql.SQL(
"SELECT DISTINCT activity.application_name "
"FROM pg_stat_activity AS activity "
"JOIN pg_locks AS waiting_lock ON waiting_lock.pid = activity.pid "
"WHERE activity.datname = current_database() "
"AND waiting_lock.locktype = 'advisory' AND NOT waiting_lock.granted "
"AND waiting_lock.database = (SELECT oid FROM pg_database WHERE datname = current_database()) "
"AND ({})"
).format(key_predicates)
identities: Final = self.query(statement, parameters)
return frozenset(row[0] for row in identities if isinstance(row[0], str))
@step("Hold an advisory lock on the test database")
@contextmanager
def lock(self, key: int = GATE_KEY) -> Generator[None]:

65
tests/e2e/migrations/test_startup.py
View file

@ -1,17 +1,31 @@
from contextlib import ExitStack
from itertools import accumulate, repeat
from typing import Final
import pytest
from .checks import COMPLETE, FATAL, GATED, assert_completed, start_replicas
from .checks import (
COMPLETE,
FATAL,
GATED,
assert_completed,
assert_completed_after_lock_timeouts,
start_replicas,
)
from .containers import Containers, failed, ready, until, waiting
from .database import PRISMA_LOCK, Database, Databases, restricted_user
from .database import GATE_KEY, PRISMA_LOCK, Database, Databases, restricted_user
from .startup_models import Migration
from e2e_metadata import Domain, Subject, meta
pytestmark: Final = [pytest.mark.e2e, pytest.mark.migration_startup]
def union_waiting_identities(
observed: frozenset[str], snapshot: frozenset[str]
) -> frozenset[str]:
return observed | snapshot
class TestMigrationStartup:
@pytest.mark.parametrize("replicas,v2", ((1, True), (3, True), (1, False)))
@meta(
@ -42,16 +56,34 @@ class TestMigrationStartup:
domain=Domain.DB,
)
)
def test_waiters_survive_prolonged_contention(self, containers: Containers, database: Database) -> None:
def test_waiters_retry_lock_timeouts_until_contention_clears(
self, containers: Containers, database: Database
) -> None:
with ExitStack() as stack:
with database.lock():
owner: Final = stack.enter_context(containers.start(database, (GATED,)))
until("owner blocked in migration SQL", lambda: bool(database.blocked()))
followers: Final = start_replicas(stack, containers, database, (GATED,), count=2)
until("both followers attempted Prisma locking", lambda: len(database.blocked(PRISMA_LOCK)) == 2)
waiting((owner, *followers), 120)
waiting_identities: Final = accumulate(
(database.waiting_identities((PRISMA_LOCK, GATE_KEY)) for _ in repeat(None)),
union_waiting_identities,
initial=frozenset(),
)
follower_identities: Final = frozenset(follower.identity for follower in followers)
until(
"both followers entered migration locking",
lambda: follower_identities <= next(waiting_identities),
)
until(
"a replica logged a migration lock-timeout retry",
lambda: any(
"timed out waiting for a table lock" in replica.logs()
for replica in (owner, *followers)
),
)
waiting((owner, *followers), 20)
ready((owner, *followers), database)
assert_completed(database, GATED)
assert_completed_after_lock_timeouts(database, GATED)
@meta(
Subject(
@ -72,6 +104,27 @@ class TestMigrationStartup:
ready((restarted,), database)
assert_completed(database)
@meta(
Subject(
domain=Domain.DB,
)
)
def test_lock_timeout_exhaustion_then_restart(self, containers: Containers, database: Database) -> None:
with database.lock():
with containers.start(
database, (GATED,), environment={"LITELLM_MIGRATION_DDL_LOCK_TIMEOUT": "2"}
) as replica:
until("owner blocked in migration SQL", lambda: bool(database.blocked()))
failed((replica,), "attempts that made no progress")
assert not database.exists("migration_effect")
assert database.query(
"SELECT count(*) FROM _prisma_migrations WHERE migration_name = %s AND finished_at IS NOT NULL",
(GATED.name,),
) == ((0,),)
with containers.start(database, (GATED,)) as restarted:
ready((restarted,), database)
assert_completed_after_lock_timeouts(database, GATED)
@meta(
Subject(
domain=Domain.DB,

2
tests/e2e/quota_management/spend_tracking/SPEND_TRACKING_COVERAGE_MATRIX.md
View file

@ -32,7 +32,7 @@ proxy + SpendLogs rows. Status: `covered` / `partial` / `gap`.
| Chat (non-stream) | `test_cost_calculator.py`, `local_testing/test_completion_cost.py` | covered | yes (`test_chat_completion_writes_nonzero_spend_row`) |
| Chat (streaming) | `test_streaming_interrupt_spend_tracking.py` | partial | yes (`test_streaming_chat_completion_tracks_spend`) |
| Embedding | `test_cost_calculator.py` (#29956) | partial | yes (`test_embedding_writes_nonzero_spend_row`) |
| Pass-through (gemini/anthropic) | `pass_through_tests/*.test.js` + `llm_translation/` suite | covered | yes (llm_translation suite) |
| Pass-through (Gemini/Vertex) | `tests/integration/observability/test_passthrough_upstream_error_visibility.py`, `tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py`, and `tests/e2e/llm_translation/` | covered | yes (integration and e2e suites) |
| Image / audio / rerank / responses / realtime | per-provider unit cost tests | partial/gap | gap |
## Entity spend aggregation

3
tests/e2e/ui/migration.serverRootPath.config.ts
View file

@ -12,7 +12,8 @@ export default defineConfig({
testMatch: ["migratedPages.spec.ts"],
fullyParallel: true,
forbidOnly: !!process.env.CI,
retries: process.env.CI ? 2 : 0,
retries: 0,
grep: /navigates directly between two migrated pages/,
workers: process.env.CI ? 1 : undefined,
reporter: "list",
use: {

6
tests/e2e/ui/tests/logs/logsFilters.spec.ts
View file

@ -143,10 +143,12 @@ test.describe("Logs page filters", () => {
await applyComboboxFilter(page, drawer, "Search a key alias", mine.alias!);
await expect(requestLogsRows(page).filter({ hasText: theirRequestId })).toHaveCount(0, { timeout: 30_000 });
// A filter you cannot clear is a page that looks empty forever, which is how it reads to a user.
await page.getByRole("button", { name: "Reset Filters" }).filter({ visible: true }).click();
const search = visibleTestId(page, "datatable-search");
await search.fill(theirRequestId);
await expect(requestLogsRows(page).filter({ hasText: theirRequestId })).toHaveCount(1, { timeout: 30_000 });
await expect(requestLogsRows(page).filter({ hasText: myRequestId })).toHaveCount(1, { timeout: 10_000 });
await search.fill(myRequestId);
await expect(requestLogsRows(page).filter({ hasText: myRequestId })).toHaveCount(1, { timeout: 30_000 });
});
});

2
tests/e2e/ui/tests/proxy-admin/keyBudgetWindow.spec.ts
View file

@ -24,6 +24,8 @@ const test = base.extend<{ scopedKey: ScopedKey }>({
});
test.describe("Proxy Admin - Key budget window", () => {
// TODO: drop once https://github.com/BerriAI/litellm/pull/41350 lands, a background refresh can wipe the edit draft
test.describe.configure({ retries: 2 });
test.use({ storageState: ADMIN_STORAGE_PATH });
test("a monthly spend cap survives a reload, and clearing the window keeps the cap", async ({ page, scopedKey }) => {

2
tests/e2e_harness/AGENTS.md
View file

@ -14,4 +14,4 @@ LITELLM_MASTER_KEY=sk-harness uv run pytest tests/e2e_harness
Rules: no `e2e` marker and no `@meta`, since nothing here drives the proxy. `@pytest.mark.covers` only where the test proves the collector or the JUnit properties read it. Inputs via arguments or env vars (setting an env var through pytest's `monkeypatch` fixture is fine, patching a function, class or module is not). And the same typing bar as the suite, `make lint-e2e-basedpyright` covers this folder and allows zero errors. The raw HTTP client ban (`tests/code_coverage_tests/check_e2e_no_raw_requests.py`) applies here too
CI: the `python` job in `.github/workflows/test-linting.yml` runs this folder whenever anything under `tests/e2e/` (except `ui/`) or `tests/e2e_harness/` changes, with the `claude` CLI installed. The CircleCI `provider_replay_harness` job also runs the provider-edge and fixture tests at the root of this folder next to `tests/code_coverage_tests/test_provider_replay_harness.py`, which imports helpers from `test_provider_edge.py`
CI: the `python` job in `.github/workflows/test-linting.yml` runs this folder whenever anything under `tests/e2e/` (except `ui/`) or `tests/e2e_harness/` changes, with the `claude` CLI installed. It also runs `tests/code_coverage_tests/test_provider_replay_harness.py` and `tests/code_coverage_tests/test_provider_cache.py` with `--noconftest -o addopts= -o "pythonpath=tests/e2e tests/e2e_harness" -p no:rerunfailures`; the replay test imports helpers from `test_provider_edge.py`

443
tests/integration/_support/actors.py Normal file
View file

@ -0,0 +1,443 @@
import uuid
from collections.abc import AsyncIterator, Iterator, Mapping, Sequence
from contextlib import asynccontextmanager, contextmanager
from dataclasses import dataclass
from enum import Enum
from pathlib import Path
from types import MappingProxyType
from typing import Final
import httpx
import yaml
from prisma import Prisma
from pydantic import JsonValue
from litellm.proxy.utils import PrismaClient, hash_token
from tests.integration._support.client import Gateway, Scenario, string_value
from tests.integration._support.prisma_clients import connected_proxy_prisma
from tests.integration._support.process import owned_proxy
TEAM_ALPHA: Final = "behavior-pin-team-alpha"
TEAM_BETA: Final = "behavior-pin-team-beta"
TEAM_GAMMA: Final = "behavior-pin-team-gamma"
ORG_A: Final = "behavior-pin-org-a"
ORG_B: Final = "behavior-pin-org-b"
class Actor(str, Enum):
PROXY_ADMIN = "proxy_admin"
ORG_ADMIN = "org_admin"
TEAM_ADMIN = "team_admin"
INTERNAL_USER = "internal_user"
OWNER = "owner"
UNRELATED_SAME_ORG = "unrelated_same_org"
CROSS_ORG_USER = "cross_org_user"
SERVICE_ACCOUNT = "service_account"
ORG_B_ADMIN = "org_b_admin"
@dataclass(frozen=True, slots=True)
class SeededKey:
user_id: str
cleartext: str
hashed: str
@dataclass(frozen=True, slots=True)
class ActorSet:
org_a_id: str
org_b_id: str
team_alpha_id: str
team_beta_id: str
team_gamma_id: str
keys: Mapping[Actor, SeededKey]
@dataclass(frozen=True, slots=True)
class ScratchResources:
prefix: str
def tag(self, suffix: str = "") -> str:
return f"{self.prefix}-{suffix}" if suffix else self.prefix
@dataclass(frozen=True, slots=True)
class GatewayClient:
gateway: Gateway
replacements: Mapping[str, str]
async def request(
self,
method: str,
path: str,
*,
headers: Mapping[str, str] | None = None,
json: Mapping[str, JsonValue] | None = None,
params: Mapping[str, str] | None = None,
) -> httpx.Response:
resolved_path: Final = _replace_text(path, self.replacements)
resolved_body: Final = _replace_body(json, self.replacements) if json is not None else None
return self.gateway.request(method, resolved_path, resolved_body, headers=headers, params=params)
async def post(
self,
path: str,
*,
headers: Mapping[str, str] | None = None,
json: Mapping[str, JsonValue] | None = None,
) -> httpx.Response:
return await self.request("POST", path, headers=headers, json=json)
async def get(
self,
path: str,
*,
headers: Mapping[str, str] | None = None,
params: Mapping[str, str] | None = None,
) -> httpx.Response:
return await self.request("GET", path, headers=headers, params=params)
async def put(
self,
path: str,
*,
headers: Mapping[str, str] | None = None,
json: Mapping[str, JsonValue] | None = None,
) -> httpx.Response:
return await self.request("PUT", path, headers=headers, json=json)
async def delete(
self,
path: str,
*,
headers: Mapping[str, str] | None = None,
json: Mapping[str, JsonValue] | None = None,
) -> httpx.Response:
return await self.request("DELETE", path, headers=headers, json=json)
@contextmanager
def owned_settings_client(
client: GatewayClient, directory: Path, litellm_settings: Mapping[str, JsonValue]
) -> Iterator[GatewayClient]:
base: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
config: Final = {**base, "litellm_settings": {**base["litellm_settings"], **litellm_settings}}
path: Final = directory / "owned-settings.yaml"
path.write_text(yaml.safe_dump(config))
with owned_proxy(client.gateway, directory, {}, config=path) as candidate:
yield GatewayClient(candidate, client.replacements)
@dataclass(frozen=True, slots=True)
class ManagementDatabase:
prisma: PrismaClient
gateway: Gateway
@property
def db(self) -> Prisma:
return self.prisma.db
def _replace_text(value: str, replacements: Mapping[str, str]) -> str:
return value.replace(TEAM_ALPHA, replacements[TEAM_ALPHA]).replace(
TEAM_BETA, replacements[TEAM_BETA]
).replace(TEAM_GAMMA, replacements[TEAM_GAMMA]).replace(ORG_A, replacements[ORG_A]).replace(
ORG_B, replacements[ORG_B]
)
def _replace_body(
value: JsonValue, replacements: Mapping[str, str]
) -> JsonValue:
if isinstance(value, str):
return _replace_text(value, replacements)
if isinstance(value, list):
return [_replace_body(item, replacements) for item in value]
if isinstance(value, dict):
return {key: _replace_body(item, replacements) for key, item in value.items()}
return value
def _seeded_key(
scenario: Scenario,
actor: Actor,
user_id: str,
*,
team_id: str | None = None,
organization_id: str | None = None,
) -> SeededKey:
token: Final = scenario.key(
key_alias=f"integration-{actor.value}-{uuid.uuid4().hex}",
user_id=user_id,
team_id=team_id,
organization_id=organization_id,
)
return SeededKey(user_id=user_id, cleartext=token, hashed=hash_token(token))
async def create_actor_set(
gateway: Gateway, scenario: Scenario, db: ManagementDatabase
) -> ActorSet:
run_id: Final = uuid.uuid4().hex
org_a_id: Final = scenario.organization(organization_alias=f"alpha-{run_id}")
org_b_id: Final = scenario.organization(organization_alias=f"beta-{run_id}")
proxy_admin_id: Final = scenario.user(user_role="proxy_admin")
org_admin_a_id: Final = scenario.org_member(org_a_id, "org_admin")
org_admin_b_id: Final = scenario.org_member(org_b_id, "org_admin")
await db.db.litellm_usertable.update(
where={"user_id": org_admin_a_id}, data={"user_role": "org_admin"}
)
await db.db.litellm_usertable.update(
where={"user_id": org_admin_b_id}, data={"user_role": "org_admin"}
)
team_alpha_id: Final = scenario.team(
team_id=f"integration-{run_id}-alpha",
team_alias=f"alpha-{run_id}",
organization_id=org_a_id,
team_member_permissions=["/key/health"],
)
team_beta_id: Final = scenario.team(
team_id=f"integration-{run_id}-beta",
team_alias=f"beta-{run_id}",
organization_id=org_b_id,
team_member_permissions=["/key/health"],
)
team_gamma_id: Final = scenario.team(
team_id=f"integration-{run_id}-gamma",
team_alias=f"gamma-{run_id}",
organization_id=org_a_id,
team_member_permissions=["/key/health"],
)
team_admin_id: Final = scenario.member(team_alpha_id, role="admin")
internal_user_id: Final = scenario.member(team_alpha_id)
owner_id: Final = scenario.member(team_alpha_id)
unrelated_id: Final = scenario.member(team_alpha_id)
service_account_id: Final = scenario.member(team_alpha_id)
cross_org_id: Final = scenario.member(team_beta_id)
keys: Final[Mapping[Actor, SeededKey]] = MappingProxyType(
{
Actor.PROXY_ADMIN: _seeded_key(scenario, Actor.PROXY_ADMIN, proxy_admin_id),
Actor.ORG_ADMIN: _seeded_key(
scenario, Actor.ORG_ADMIN, org_admin_a_id, organization_id=org_a_id
),
Actor.TEAM_ADMIN: _seeded_key(
scenario, Actor.TEAM_ADMIN, team_admin_id, team_id=team_alpha_id, organization_id=org_a_id
),
Actor.INTERNAL_USER: _seeded_key(
scenario, Actor.INTERNAL_USER, internal_user_id, team_id=team_alpha_id, organization_id=org_a_id
),
Actor.OWNER: _seeded_key(
scenario, Actor.OWNER, owner_id, team_id=team_alpha_id, organization_id=org_a_id
),
Actor.UNRELATED_SAME_ORG: _seeded_key(
scenario, Actor.UNRELATED_SAME_ORG, unrelated_id, team_id=team_alpha_id, organization_id=org_a_id
),
Actor.CROSS_ORG_USER: _seeded_key(
scenario, Actor.CROSS_ORG_USER, cross_org_id, team_id=team_beta_id, organization_id=org_b_id
),
Actor.SERVICE_ACCOUNT: _seeded_key(
scenario, Actor.SERVICE_ACCOUNT, service_account_id, team_id=team_alpha_id, organization_id=org_a_id
),
Actor.ORG_B_ADMIN: _seeded_key(
scenario, Actor.ORG_B_ADMIN, org_admin_b_id, organization_id=org_b_id
),
}
)
return ActorSet(org_a_id, org_b_id, team_alpha_id, team_beta_id, team_gamma_id, keys)
def create_gateway_client(gateway: Gateway, actors: ActorSet) -> GatewayClient:
replacements: Final[Mapping[str, str]] = MappingProxyType(
{
TEAM_ALPHA: actors.team_alpha_id,
TEAM_BETA: actors.team_beta_id,
TEAM_GAMMA: actors.team_gamma_id,
ORG_A: actors.org_a_id,
ORG_B: actors.org_b_id,
}
)
return GatewayClient(gateway, replacements)
async def create_scratch_key(
gateway_client: GatewayClient,
seeder_cleartext: str,
scratch_prefix: str,
*,
user_id: str,
team_id: str | None = None,
organization_id: str | None = None,
key_alias: str | None = None,
) -> str:
body: Final[dict[str, JsonValue]] = {
"key_alias": key_alias or scratch_prefix,
"user_id": user_id,
**({"team_id": team_id} if team_id is not None else {}),
**({"organization_id": organization_id} if organization_id is not None else {}),
}
created: Final = gateway_client.gateway.post(
"/key/generate", body, key=seeder_cleartext
)
return string_value(created["key"])
async def create_scratch_team(
db: ManagementDatabase,
team_id: str,
*,
organization_id: str | None = None,
admin_user_ids: Sequence[str] = (),
member_user_ids: Sequence[str] = (),
team_member_permissions: Sequence[str] | None = None,
models: Sequence[str] | None = None,
max_budget: float | None = None,
tpm_limit: int | None = None,
rpm_limit: int | None = None,
metadata: Mapping[str, JsonValue] | None = None,
) -> str:
admins: Final = list(admin_user_ids)
members: Final = list(member_user_ids)
members_with_roles: Final = [
{"user_id": identity, "role": "admin"} for identity in admins
] + [{"user_id": identity, "role": "user"} for identity in members]
fields: Final[dict[str, JsonValue]] = {
"team_id": team_id,
"team_alias": team_id,
"admins": admins,
"members": admins + members,
"members_with_roles": members_with_roles,
**({"organization_id": organization_id} if organization_id is not None else {}),
**(
{"team_member_permissions": list(team_member_permissions)}
if team_member_permissions is not None
else {}
),
**({"models": list(models)} if models is not None else {}),
**({"max_budget": max_budget} if max_budget is not None else {}),
**({"tpm_limit": tpm_limit} if tpm_limit is not None else {}),
**({"rpm_limit": rpm_limit} if rpm_limit is not None else {}),
**({"metadata": dict(metadata)} if metadata is not None else {}),
}
created: Final = db.gateway.post("/team/new", fields)
return string_value(created["team_id"])
async def create_scratch_user(
db: ManagementDatabase,
scratch_prefix: str,
*,
suffix: str,
user_email: str | None = None,
) -> str:
user_id: Final = f"{scratch_prefix}-{suffix}"
fields: Final[dict[str, JsonValue]] = {
"user_id": user_id,
"user_role": "internal_user",
"auto_create_key": False,
**({"user_email": user_email} if user_email is not None else {}),
}
created: Final = db.gateway.post("/user/new", fields)
return string_value(created["user_id"])
async def create_scratch_org(
db: ManagementDatabase,
scratch_prefix: str,
*,
max_budget: float | None = None,
tpm_limit: int | None = None,
rpm_limit: int | None = None,
models: Sequence[str] | None = None,
metadata: Mapping[str, JsonValue] | None = None,
suffix: str = "org",
) -> str:
organization_id: Final = f"{scratch_prefix}-{suffix}"
budget_fields: Final[dict[str, JsonValue]] = {
"budget_id": f"{organization_id}-budget",
**({"max_budget": max_budget} if max_budget is not None else {}),
**({"tpm_limit": tpm_limit} if tpm_limit is not None else {}),
**({"rpm_limit": rpm_limit} if rpm_limit is not None else {}),
}
budget: Final = db.gateway.post("/budget/new", budget_fields)
budget_id: Final = string_value(budget["budget_id"])
fields: Final[dict[str, JsonValue]] = {
"organization_id": organization_id,
"organization_alias": organization_id,
"budget_id": budget_id,
**({"models": list(models)} if models is not None else {}),
**({"metadata": dict(metadata)} if metadata is not None else {}),
}
created: Final = db.gateway.post("/organization/new", fields)
return string_value(created["organization_id"])
async def create_scratch_actor(
db: ManagementDatabase,
scratch_prefix: str,
*,
user_role: str,
org_admin_of: Sequence[str] = (),
organization_id: str | None = None,
suffix: str = "actor",
) -> SeededKey:
api_user_role: Final = (
user_role
if user_role in {"proxy_admin", "proxy_admin_viewer", "internal_user", "internal_user_viewer"}
else "internal_user"
)
created_user: Final = db.gateway.post(
"/user/new",
{
"user_id": f"{scratch_prefix}-{suffix}",
"user_role": api_user_role,
"auto_create_key": False,
**({"organization_id": organization_id} if organization_id is not None else {}),
},
)
user_id: Final = string_value(created_user["user_id"])
token: Final = string_value(
db.gateway.post(
"/key/generate",
{"user_id": user_id, "key_alias": f"{scratch_prefix}-{suffix}-key"},
)["key"]
)
tuple(
db.gateway.post(
"/organization/member_add",
{"organization_id": org_id, "member": {"user_id": user_id, "role": "org_admin"}},
)
for org_id in org_admin_of
)
return SeededKey(user_id, token, hash_token(token))
@asynccontextmanager
async def management_database(gateway: Gateway) -> AsyncIterator[ManagementDatabase]:
async with connected_proxy_prisma() as client:
yield ManagementDatabase(client, gateway)
@asynccontextmanager
async def scratch_resources(db: ManagementDatabase) -> AsyncIterator[ScratchResources]:
resource: Final = ScratchResources(f"integration-{uuid.uuid4().hex}")
yield resource
prefix: Final = resource.prefix
await db.db.litellm_verificationtoken.delete_many(
where={"OR": [{"key_alias": {"startswith": prefix}}, {"key_name": {"startswith": prefix}}]}
)
await db.db.litellm_teammembership.delete_many(where={"team_id": {"startswith": prefix}})
await db.db.litellm_organizationmembership.delete_many(where={"user_id": {"startswith": prefix}})
await db.db.litellm_teamtable.delete_many(where={"team_id": {"startswith": prefix}})
await db.db.litellm_usertable.delete_many(
where={"OR": [{"user_id": {"startswith": prefix}}, {"user_email": {"startswith": prefix}}]}
)
await db.db.litellm_organizationtable.delete_many(where={"organization_id": {"startswith": prefix}})
await db.db.litellm_budgettable.delete_many(where={"budget_id": {"startswith": prefix}})
users: Final = await db.db.litellm_usertable.find_many(where={"teams": {"isEmpty": False}})
for user in users:
teams: Final = [team_id for team_id in user.teams if not team_id.startswith(prefix)]
if teams != list(user.teams):
await db.db.litellm_usertable.update(
where={"user_id": user.user_id},
data={"teams": {"set": teams}},
)

35
tests/integration/_support/prisma_clients.py Normal file
View file

@ -0,0 +1,35 @@
import os
from collections.abc import AsyncIterator
from contextlib import asynccontextmanager
from typing import Final
from prisma import Prisma
from prisma.types import DatasourceOverride
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
from litellm.proxy.utils import PrismaClient, ProxyLogging
@asynccontextmanager
async def connected_prisma(database_url: str | None = None) -> AsyncIterator[Prisma]:
client: Final = Prisma(
datasource=DatasourceOverride(url=database_url)
) if database_url is not None else Prisma()
await client.connect()
try:
yield client
finally:
await client.disconnect()
@asynccontextmanager
async def connected_proxy_prisma(database_url: str | None = None) -> AsyncIterator[PrismaClient]:
client: Final = PrismaClient(
database_url=database_url or os.environ["DATABASE_URL"],
proxy_logging_obj=ProxyLogging(UserApiKeyCache()),
)
await client.connect()
try:
yield client
finally:
await client.disconnect()

0
tests/proxy_behavior/__init__.py → tests/integration/authorization/__init__.py
View file

10
tests/proxy_behavior/auth/test_auth_object_prefetch.py → tests/integration/authorization/test_auth_object_prefetch.py
View file

@ -3,10 +3,12 @@ team and hand the getters rows they validate. The per-regime round-trip counts a
tests/unit/proxy/auth/test_auth_object_prefetch.py."""
import json
from collections.abc import AsyncIterator
from unittest.mock import AsyncMock, MagicMock
from uuid import uuid4
import pytest
import pytest_asyncio
from litellm.caching.in_memory_cache import InMemoryCache
from litellm.proxy.auth.auth_checks import (
@ -18,10 +20,18 @@ from litellm.proxy.auth.auth_checks import (
from litellm.proxy.auth.auth_object_prefetch import AuthObjectRefs, prefetch_auth_objects
from litellm.proxy.auth.team_grants import team_model_aliases
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
from litellm.proxy.utils import PrismaClient
from tests.integration._support.prisma_clients import connected_proxy_prisma
pytestmark = pytest.mark.asyncio(loop_scope="session")
@pytest_asyncio.fixture(scope="module", loop_scope="session")
async def prisma() -> AsyncIterator[PrismaClient]:
async with connected_proxy_prisma() as client:
yield client
def _frozen_cache() -> UserApiKeyCache:
"""The org entries carry a 5s TTL; a frozen clock keeps a slow first call from expiring them mid-test."""
return UserApiKeyCache(in_memory_cache=InMemoryCache(clock=lambda: 1_000_000.0), redis_cache=None)

67
tests/load_tests/test_granian_admission_saturation.py → tests/integration/compatibility/test_granian_admission_saturation.py
View file

@ -12,10 +12,9 @@ import pytest
from tests._master_key import MASTER_KEY
pytestmark = pytest.mark.skipif(
os.environ.get("LITELLM_RUN_SATURATION_BENCHMARK") != "1",
reason="set LITELLM_RUN_SATURATION_BENCHMARK=1 to run the saturation benchmark",
)
MAX_IN_FLIGHT: Final = 8
MAX_QUEUED: Final = 8
REQUESTS: Final = 200
def _free_port() -> int:
@ -24,27 +23,24 @@ def _free_port() -> int:
return int(listener.getsockname()[1])
def _percentile(values: list[float], percentile: float) -> float:
return sorted(values)[min(int(len(values) * percentile), len(values) - 1)]
@pytest.mark.asyncio
async def test_granian_admission_control_saturation(tmp_path: Path) -> None:
fake_port: Final = _free_port()
proxy_port: Final = _free_port()
fake_script: Final = Path(__file__).parents[1] / "_fake_openai_endpoint_server.py"
fake_script: Final = Path(__file__).parents[2] / "_fake_openai_endpoint_server.py"
config_path: Final = tmp_path / "saturation_config.yaml"
proxy_env: Final = {**os.environ, "LITELLM_MASTER_KEY": MASTER_KEY}
proxy_env: Final = {**{k: v for k, v in os.environ.items() if k != "DATABASE_URL"}, "LITELLM_MASTER_KEY": MASTER_KEY}
config_path.write_text(
f"""model_list:
- model_name: slow-endpoint
litellm_params:
model: openai/slow-endpoint
api_base: http://127.0.0.1:{fake_port}/v1
api_key: fake-upstream-key
general_settings:
master_key: os.environ/LITELLM_MASTER_KEY
max_in_flight_requests_per_worker: 8
max_queued_requests_per_worker: 8
max_in_flight_requests_per_worker: {MAX_IN_FLIGHT}
max_queued_requests_per_worker: {MAX_QUEUED}
admission_queue_timeout_seconds: 0.5
"""
)
@ -85,22 +81,7 @@ general_settings:
else:
raise AssertionError("Granian proxy did not become healthy")
liveness_latencies: Final[list[float]] = []
stop_sampling: Final = asyncio.Event()
async def sample_liveness() -> None:
while not stop_sampling.is_set():
start: Final = time.perf_counter()
try:
response = await client.get("/health/liveliness", timeout=2)
response.raise_for_status()
liveness_latencies.append(time.perf_counter() - start)
except httpx.HTTPError:
pass
await asyncio.sleep(0.05)
async def send_completion() -> tuple[int, float, bool]:
start: Final = time.perf_counter()
async def send_completion() -> tuple[int, bool]:
response = await client.post(
"/chat/completions",
headers={"Authorization": f"Bearer {MASTER_KEY}"},
@ -110,34 +91,16 @@ general_settings:
},
timeout=10,
)
return response.status_code, time.perf_counter() - start, "retry-after" in response.headers
return response.status_code, "retry-after" in response.headers
sampler: Final = asyncio.create_task(sample_liveness())
results: Final = await asyncio.gather(*(send_completion() for _ in range(200)))
stop_sampling.set()
await sampler
results: Final = await asyncio.gather(*(send_completion() for _ in range(REQUESTS)))
statuses: Final = [result[0] for result in results]
latencies: Final = [result[1] for result in results]
rejected: Final = [result for result in results if result[0] == 503]
assert set(statuses) <= {200, 503}
assert rejected
assert all(result[2] for result in rejected)
assert _percentile(latencies, 0.99) < 5
assert liveness_latencies
assert _percentile(liveness_latencies, 0.95) < 0.5
duration: Final = max(latencies)
print(
"\nmetric value\n"
f"rps {len(results) / duration:.2f}\n"
f"200 count {statuses.count(200)}\n"
f"503 count {statuses.count(503)}\n"
f"p50 {_percentile(latencies, 0.50):.3f}s\n"
f"p95 {_percentile(latencies, 0.95):.3f}s\n"
f"p99 {_percentile(latencies, 0.99):.3f}s\n"
f"liveness p95 {_percentile(liveness_latencies, 0.95):.3f}s"
)
assert set(statuses) <= {200, 503}, statuses
assert 1 <= statuses.count(200) <= MAX_IN_FLIGHT + MAX_QUEUED, statuses
assert len(rejected) >= REQUESTS - MAX_IN_FLIGHT - MAX_QUEUED, statuses
assert all(result[1] for result in rejected), rejected
finally:
proxy_process.terminate()
try:

58
tests/integration/configuration/test_basic_proxy_startup.py Normal file
View file

@ -0,0 +1,58 @@
"""
This test ensures that the proxy starts and serves requests even with a bad license.
in ci/cd config.yml, we set the license to "bad-license"
"""
import os
from typing import Final
import aiohttp
import pytest
from tests.integration._support.provider import SharedProvider
from tests.integration._support.wire import Reply
PROXY_URL: Final = os.environ.get("INTEGRATION_PROXY_URL", "http://0.0.0.0:4000")
INTEGRATION_RIG: Final = "INTEGRATION_PROXY_URL" in os.environ
MODEL: Final = "openai/gpt-5.4" if INTEGRATION_RIG else "gpt-4"
@pytest.mark.asyncio
async def test_health_and_chat_completion(provider: SharedProvider) -> None:
base_url: Final = PROXY_URL.rstrip("/")
async with aiohttp.ClientSession() as session:
async with session.get(f"{base_url}/health/readiness") as response:
assert response.status == 200, await response.text()
readiness_response: Final = await response.json()
assert readiness_response["status"] in ("healthy", "connected")
async with session.get(f"{base_url}/health/liveness") as response:
assert response.status == 200, await response.text()
headers: Final = {
"Authorization": f"Bearer {os.environ['LITELLM_MASTER_KEY']}",
"Content-Type": "application/json",
}
if INTEGRATION_RIG:
provider.expect(
Reply(
status=200,
body=b'{"id":"chatcmpl-scripted","object":"chat.completion","created":1700000000,"model":"gpt-5.4","choices":[{"index":0,"message":{"role":"assistant","content":"scripted integration reply"},"finish_reason":"stop"}],"usage":{"prompt_tokens":4,"completion_tokens":5,"total_tokens":9}}',
)
)
data: Final = {
"model": MODEL,
"messages": [
{"role": "system", "content": "You are a helpful assistant."},
{"role": "user", "content": "Hello!"},
],
}
async with session.post(f"{base_url}/chat/completions", headers=headers, json=data) as response:
assert response.status == 200, await response.text()
completion_response: Final = await response.json()
assert "choices" in completion_response
if INTEGRATION_RIG:
requests: Final = provider.received()
assert len(requests) == 1, requests
assert requests[0].method == "POST", requests[0]
assert requests[0].target.endswith("/chat/completions"), requests[0]

2
tests/test_litellm/test_conftest.py → tests/integration/configuration/test_unit_conftest_isolation.py
View file

@ -4,7 +4,7 @@ import sys
from pathlib import Path
from typing import Final
REPO_ROOT: Final = Path(__file__).resolve().parents[2]
REPO_ROOT: Final = Path(__file__).resolve().parents[3]
PROXY_BASE_URL_SENSITIVE_NODE: Final = (
"tests/unit/proxy/management_endpoints/test_mcp_management_endpoints.py"

59
tests/integration/conftest.py
View file

@ -4,20 +4,21 @@ import hashlib
import json
import os
import re
from collections.abc import Generator, Iterator, Mapping, Sequence
from collections.abc import AsyncIterator, Generator, Iterator, Mapping, Sequence
from contextlib import ExitStack
from copy import copy
from importlib.metadata import version
from pathlib import Path
from typing import Final, Protocol, runtime_checkable
from typing import TYPE_CHECKING, Final, Protocol, runtime_checkable
import httpx
import pytest
import pytest_asyncio
from execnet import XSpec
from pydantic import TypeAdapter
from redis import Redis
from tests.integration._support.client import Gateway, eventually, gateway_from_environment
from tests.integration._support.client import Gateway, Scenario, eventually, gateway_from_environment
from tests.integration._support.generation import LIFECYCLE_SETTINGS
from tests.integration._support.manifest import OWNED_DIRECTORIES
from tests.integration._support.ordering import collected_case, fixture_groups, scheduling_groups
@ -26,6 +27,9 @@ from tests.integration._support.routing import RoutingPlugin
from tests.integration._support.workers import worker_endpoints, worker_gateway
from tests.integration.run import GITHUB_FILES
if TYPE_CHECKING:
from tests.integration._support.actors import ActorSet, GatewayClient, ManagementDatabase, ScratchResources
COLLECTED: Final = pytest.StashKey[tuple[str, ...]]()
INVENTORY: Final = pytest.StashKey[tuple[str, ...]]()
REPORTS: Final = pytest.StashKey[list[pytest.TestReport]]()
@ -293,6 +297,55 @@ def gateway() -> Iterator[Gateway]:
yield value
@pytest.fixture(scope="session")
def session_gateway() -> Iterator[Gateway]:
with gateway_from_environment() as value:
yield value
@pytest.fixture
def scenario(gateway: Gateway) -> Iterator[Scenario]:
value: Final[Scenario]
with gateway.scenario() as value:
yield value
@pytest_asyncio.fixture(scope="session", loop_scope="session")
async def actors(
session_gateway: Gateway, db: ManagementDatabase
) -> AsyncIterator[ActorSet]:
with session_gateway.scenario() as scenario:
from tests.integration._support.actors import create_actor_set
value: Final[ActorSet] = await create_actor_set(session_gateway, scenario, db)
yield value
@pytest.fixture(scope="session")
def gateway_client(session_gateway: Gateway, actors: ActorSet) -> GatewayClient:
from tests.integration._support.actors import create_gateway_client
return create_gateway_client(session_gateway, actors)
@pytest_asyncio.fixture(scope="session", loop_scope="session")
async def db(session_gateway: Gateway) -> AsyncIterator[ManagementDatabase]:
from tests.integration._support.actors import management_database
value: Final[ManagementDatabase]
async with management_database(session_gateway) as value:
yield value
@pytest_asyncio.fixture(scope="function", loop_scope="session")
async def resources(db: ManagementDatabase) -> AsyncIterator[ScratchResources]:
from tests.integration._support.actors import scratch_resources
value: Final[ScratchResources]
async with scratch_resources(db) as value:
yield value
@pytest.fixture(scope="session")
def shared_provider_server() -> Iterator[SharedProvider]:
if os.environ.get("PYTEST_XDIST_WORKER") and os.environ.get("INTEGRATION_WORKER_ISOLATION") != "1":

0
tests/proxy_behavior/auth/__init__.py → tests/integration/management/__init__.py
View file

29
tests/proxy_behavior/management/test_credential_migration_endpoint.py → tests/integration/management/test_credential_migration_endpoint.py
View file

@ -13,16 +13,15 @@ the ``user_api_key_auth`` layer (401) before the endpoint's own admin guard runs
import pytest
from .conftest import MASTER_KEY
pytestmark = pytest.mark.asyncio(loop_scope="session")
async def test_migrate_encryption_check_as_admin_is_read_only(proxy_client):
async def test_migrate_encryption_check_as_admin_is_read_only(gateway_client):
"""GET /credentials/migrate-encryption/check returns a residual report (no writes)."""
resp = await proxy_client.get(
resp = await gateway_client.get(
"/credentials/migrate-encryption/check",
headers={"Authorization": f"Bearer {MASTER_KEY}"},
headers={"Authorization": f"Bearer {gateway_client.gateway.key}"},
)
assert resp.status_code == 200, resp.text
body = resp.json()
@ -30,41 +29,41 @@ async def test_migrate_encryption_check_as_admin_is_read_only(proxy_client):
assert "residual_legacy" in body["report"]
async def test_migrate_encryption_check_requires_admin(proxy_client, scratch):
async def test_migrate_encryption_check_requires_admin(gateway_client, resources):
"""A non-admin key cannot reach the residual scan (auth layer rejects, 401)."""
gen = await proxy_client.post(
gen = await gateway_client.post(
"/key/generate",
headers={"Authorization": f"Bearer {MASTER_KEY}"},
json={"key_alias": scratch.tag("check"), "user_id": scratch.tag("check-user")},
headers={"Authorization": f"Bearer {gateway_client.gateway.key}"},
json={"key_alias": resources.tag("check"), "user_id": resources.tag("check-user")},
)
assert gen.status_code == 200, gen.text
nonadmin_key = gen.json()["key"]
resp = await proxy_client.get(
resp = await gateway_client.get(
"/credentials/migrate-encryption/check",
headers={"Authorization": f"Bearer {nonadmin_key}"},
)
assert resp.status_code == 401, resp.text
async def test_migrate_encryption_requires_admin(proxy_client, scratch):
async def test_migrate_encryption_requires_admin(gateway_client, resources):
"""A non-admin key cannot trigger the migration (auth layer rejects, 401).
Rejection happens before any write: the admin-only route check fires in
``user_api_key_auth``, ahead of the endpoint body.
"""
gen = await proxy_client.post(
gen = await gateway_client.post(
"/key/generate",
headers={"Authorization": f"Bearer {MASTER_KEY}"},
headers={"Authorization": f"Bearer {gateway_client.gateway.key}"},
json={
"key_alias": scratch.tag("migrate"),
"user_id": scratch.tag("migrate-user"),
"key_alias": resources.tag("migrate"),
"user_id": resources.tag("migrate-user"),
},
)
assert gen.status_code == 200, gen.text
nonadmin_key = gen.json()["key"]
resp = await proxy_client.post(
resp = await gateway_client.post(
"/credentials/migrate-encryption",
headers={"Authorization": f"Bearer {nonadmin_key}"},
)

199
tests/proxy_behavior/management/test_f7_coverage_closeout.py → tests/integration/management/test_f7_coverage_closeout.py
View file

@ -1,3 +1,4 @@
from typing import Final
"""Phase 4 F7 — coverage gap-closer scenarios picked from the un-covered
ranges left after F1–F6 landed.
@ -27,8 +28,8 @@ from prisma import Json
from litellm.proxy.utils import hash_token
from .actors import TEAM_ALPHA, Actor
from .conftest import create_scratch_org, create_scratch_team
from tests.integration._support.actors import TEAM_ALPHA, Actor
from tests.integration._support.actors import create_scratch_org, create_scratch_team
pytestmark = pytest.mark.asyncio(loop_scope="session")
@ -43,11 +44,11 @@ pytestmark = pytest.mark.asyncio(loop_scope="session")
async def test_org_team_guaranteed_throughput_model_over_bound_rejected(
proxy_client, prisma, scratch, world
gateway_client, db, resources, actors
):
org_id = await create_scratch_org(
prisma,
scratch.prefix,
db,
resources.prefix,
models=["gpt-4"],
metadata={"model_rpm_limit": {"gpt-4": 30}},
)
@ -55,13 +56,13 @@ async def test_org_team_guaranteed_throughput_model_over_bound_rejected(
# forces _check_team_model_specific_limits's `model_specific_rpm_limit`
# accumulator to actually accumulate (covers lines 468–478).
await create_scratch_team(
prisma,
team_id=scratch.tag("sibling"),
db,
team_id=resources.tag("sibling"),
organization_id=org_id,
metadata={"model_rpm_limit": {"gpt-4": 20}},
)
seeder = world.keys[Actor.PROXY_ADMIN].cleartext
team_id = scratch.tag("new")
seeder = actors.keys[Actor.PROXY_ADMIN].cleartext
team_id = resources.tag("new")
body = {
"team_id": team_id,
"team_alias": team_id,
@ -70,7 +71,7 @@ async def test_org_team_guaranteed_throughput_model_over_bound_rejected(
"model_rpm_limit": {"gpt-4": 100},
"rpm_limit_type": "guaranteed_throughput",
}
resp = await proxy_client.post(
resp = await gateway_client.post(
"/team/new",
headers={"Authorization": f"Bearer {seeder}"},
json=body,
@ -78,30 +79,30 @@ async def test_org_team_guaranteed_throughput_model_over_bound_rejected(
# 20 (sibling) + 100 (new) > 30 (org cap) → guard fires.
assert resp.status_code == 400, resp.text
assert "RPM" in resp.text, resp.text
rows = await prisma.db.litellm_teamtable.find_many(where={"team_id": team_id})
rows = await db.db.litellm_teamtable.find_many(where={"team_id": team_id})
assert rows == []
async def test_org_team_guaranteed_throughput_model_tpm_over_bound_rejected(
proxy_client, prisma, scratch, world
gateway_client, db, resources, actors
):
"""Mirror of the rpm scenario but for the model_tpm side — pins
_check_team_model_specific_limits's tpm branch (lines 503–521) which
the rpm scenario doesn't exercise."""
org_id = await create_scratch_org(
prisma,
scratch.prefix,
db,
resources.prefix,
models=["gpt-4"],
metadata={"model_tpm_limit": {"gpt-4": 500}},
)
await create_scratch_team(
prisma,
team_id=scratch.tag("sibling"),
db,
team_id=resources.tag("sibling"),
organization_id=org_id,
metadata={"model_tpm_limit": {"gpt-4": 200}},
)
seeder = world.keys[Actor.PROXY_ADMIN].cleartext
team_id = scratch.tag("new")
seeder = actors.keys[Actor.PROXY_ADMIN].cleartext
team_id = resources.tag("new")
body = {
"team_id": team_id,
"team_alias": team_id,
@ -110,51 +111,48 @@ async def test_org_team_guaranteed_throughput_model_tpm_over_bound_rejected(
"model_tpm_limit": {"gpt-4": 5000},
"tpm_limit_type": "guaranteed_throughput",
}
resp = await proxy_client.post(
resp = await gateway_client.post(
"/team/new",
headers={"Authorization": f"Bearer {seeder}"},
json=body,
)
assert resp.status_code == 400, resp.text
assert "TPM" in resp.text, resp.text
rows = await prisma.db.litellm_teamtable.find_many(where={"team_id": team_id})
rows = await db.db.litellm_teamtable.find_many(where={"team_id": team_id})
assert rows == []
async def test_org_team_guaranteed_throughput_aggregate_runs(
proxy_client, prisma, scratch, world
):
"""Aggregate guard's no-op path (line 549-566 — entity_rpm_limit is None
because include_budget_table=False at the call site). The helper still
executes its `allocated_tpm = sum(...)` and `allocated_rpm = sum(...)`
lines, which is the coverage target."""
org_id = await create_scratch_org(prisma, scratch.prefix, models=["m"])
await create_scratch_team(
prisma,
team_id=scratch.tag("sibling"),
organization_id=org_id,
tpm_limit=100,
rpm_limit=10,
)
seeder = world.keys[Actor.PROXY_ADMIN].cleartext
team_id = scratch.tag("new")
body = {
"team_id": team_id,
"team_alias": team_id,
"organization_id": org_id,
"models": ["m"],
"tpm_limit": 50,
"rpm_limit": 5,
"tpm_limit_type": "guaranteed_throughput",
}
resp = await proxy_client.post(
"/team/new",
headers={"Authorization": f"Bearer {seeder}"},
json=body,
)
# No org budget table loaded → check is no-op → 200.
assert resp.status_code == 200, resp.text
@pytest.mark.parametrize("requested_tpm,expected_status", [(50, 400), (20, 200)], ids=["over", "within"])
async def test_org_team_guaranteed_throughput_counts_sibling_allocations(
requested_tpm: int, expected_status: int, gateway_client, db, resources, actors
):
org_id = await create_scratch_org(db, resources.prefix, models=["m"], tpm_limit=120)
await create_scratch_team(
db,
team_id=resources.tag("sibling"),
organization_id=org_id,
models=["m"],
tpm_limit=100,
)
team_id = resources.tag("new")
resp = await gateway_client.post(
"/team/new",
headers={"Authorization": f"Bearer {actors.keys[Actor.PROXY_ADMIN].cleartext}"},
json={
"team_id": team_id,
"team_alias": team_id,
"organization_id": org_id,
"models": ["m"],
"tpm_limit": requested_tpm,
"tpm_limit_type": "guaranteed_throughput",
},
)
assert resp.status_code == expected_status, resp.text
rows = await db.db.litellm_teamtable.find_many(where={"team_id": team_id})
assert len(rows) == (1 if expected_status == 200 else 0)
if expected_status == 400:
assert "Allocated TPM limit=100" in resp.text, resp.text
# ---------------------------------------------------------------------------
# /key/generate — _check_project_key_limits project-not-found branch
@ -164,14 +162,14 @@ async def test_org_team_guaranteed_throughput_aggregate_runs(
async def test_key_generate_with_unknown_project_id_rejected(
proxy_client, prisma, scratch, world
gateway_client, db, resources, actors
):
seeder = world.keys[Actor.PROXY_ADMIN].cleartext
seeder = actors.keys[Actor.PROXY_ADMIN].cleartext
body = {
"key_alias": scratch.prefix,
"project_id": f"{scratch.prefix}-ghost-project",
"key_alias": resources.prefix,
"project_id": f"{resources.prefix}-ghost-project",
}
resp = await proxy_client.post(
resp = await gateway_client.post(
"/key/generate",
headers={"Authorization": f"Bearer {seeder}"},
json=body,
@ -180,8 +178,8 @@ async def test_key_generate_with_unknown_project_id_rejected(
# The route's exception handler may wrap it; pin both shapes.
assert resp.status_code in (400, 404), resp.text
assert "project" in resp.text.lower() or "not found" in resp.text.lower(), resp.text
rows = await prisma.db.litellm_verificationtoken.find_many(
where={"key_alias": scratch.prefix}
rows = await db.db.litellm_verificationtoken.find_many(
where={"key_alias": resources.prefix}
)
assert rows == [], "rejected key leaked a row"
@ -195,10 +193,10 @@ async def test_key_generate_with_unknown_project_id_rejected(
async def _seed_key_for_relocation(
prisma, scratch_prefix: str, *, user_id: str, team_id: str
db, scratch_prefix: str, *, user_id: str, team_id: str
) -> str:
cleartext = "sk-" + uuid.uuid4().hex
await prisma.db.litellm_verificationtoken.create(
await db.db.litellm_verificationtoken.create(
data={
"token": hash_token(cleartext),
"key_alias": f"{scratch_prefix}-key",
@ -212,39 +210,39 @@ async def _seed_key_for_relocation(
async def test_team_new_with_team_member_budget_creates_budget_row(
proxy_client, prisma, scratch, world
gateway_client, db, resources, actors
):
"""/team/new with `team_member_budget` routes through
TeamMemberBudgetHandler.create_team_member_budget_table (lines 196–248).
Observable end-state: a litellm_budgettable row is created and the
team's metadata.team_member_budget_id points at it."""
seeder = world.keys[Actor.PROXY_ADMIN].cleartext
team_id = scratch.tag("team-with-mbudget")
resp = await proxy_client.post(
seeder = actors.keys[Actor.PROXY_ADMIN].cleartext
team_id = resources.tag("team-with-mbudget")
resp = await gateway_client.post(
"/team/new",
headers={"Authorization": f"Bearer {seeder}"},
json={
"team_id": team_id,
"team_alias": scratch.tag("alias"),
"team_alias": resources.tag("alias"),
"team_member_budget": 10.0,
"team_member_rpm_limit": 100,
"team_member_tpm_limit": 1000,
},
)
assert resp.status_code == 200, resp.text
team_row = await prisma.db.litellm_teamtable.find_unique(where={"team_id": team_id})
team_row = await db.db.litellm_teamtable.find_unique(where={"team_id": team_id})
assert team_row is not None
budget_id = (team_row.metadata or {}).get("team_member_budget_id")
assert (
budget_id is not None
), f"team_member_budget_id not stamped on team metadata: {team_row.metadata!r}"
# The handler writes the per-member budget row under a non-scratch id
# The handler writes the per-member budget row under a non-resources id
# pattern (`team-<alias>-budget-<uuid>`), so the prefix sweep can't
# reclaim it. Cleanup must run even when a downstream assertion fires;
# otherwise orphan rows accumulate across CI re-runs.
try:
budget_row = await prisma.db.litellm_budgettable.find_unique(
budget_row = await db.db.litellm_budgettable.find_unique(
where={"budget_id": budget_id}
)
assert budget_row is not None, "team_member_budget row was not created"
@ -252,23 +250,19 @@ async def test_team_new_with_team_member_budget_creates_budget_row(
assert budget_row.rpm_limit == 100
assert budget_row.tpm_limit == 1000
finally:
await prisma.db.litellm_teamtable.update(
await db.db.litellm_teamtable.update(
where={"team_id": team_id},
data={"metadata": Json({})},
)
await prisma.db.litellm_budgettable.delete(where={"budget_id": budget_id})
await db.db.litellm_budgettable.delete(where={"budget_id": budget_id})
async def test_team_update_team_member_budget_upserts(
proxy_client, prisma, scratch, world
gateway_client, db, resources, actors
):
"""/team/update with `team_member_budget` against a team that has no
pre-existing team_member_budget_id routes through
TeamMemberBudgetHandler.upsert_team_member_budget_table's else-branch
(lines 294–303), which in turn calls create_team_member_budget_table."""
team_id = await create_scratch_team(prisma, team_id=scratch.tag("team"))
seeder = world.keys[Actor.PROXY_ADMIN].cleartext
resp = await proxy_client.post(
team_id = await create_scratch_team(db, team_id=resources.tag("team"))
seeder = actors.keys[Actor.PROXY_ADMIN].cleartext
resp = await gateway_client.post(
"/team/update",
headers={"Authorization": f"Bearer {seeder}"},
json={
@ -277,69 +271,70 @@ async def test_team_update_team_member_budget_upserts(
},
)
assert resp.status_code == 200, resp.text
team_row = await prisma.db.litellm_teamtable.find_unique(where={"team_id": team_id})
team_row = await db.db.litellm_teamtable.find_unique(where={"team_id": team_id})
assert team_row is not None
budget_id = (team_row.metadata or {}).get("team_member_budget_id")
assert budget_id is not None, "team_member_budget_id not upserted"
budget_row = await db.db.litellm_budgettable.find_unique(where={"budget_id": budget_id})
assert budget_row is not None and budget_row.max_budget == 5.0, budget_row
# See comment on the sibling test — non-prefixed budget row, cleanup
# must survive an assertion failure to avoid orphan accumulation.
try:
# (No further assertions today, but the try/finally keeps the
# cleanup contract uniform with the sibling test and is robust to
# future asserts being added here.)
pass
budget_row: Final = await db.db.litellm_budgettable.find_unique(
where={"budget_id": budget_id}
)
assert budget_row is not None, "team_member_budget row was not created"
assert budget_row.max_budget == 5.0
finally:
await prisma.db.litellm_teamtable.update(
await db.db.litellm_teamtable.update(
where={"team_id": team_id},
data={"metadata": Json({})},
)
await prisma.db.litellm_budgettable.delete(where={"budget_id": budget_id})
await db.db.litellm_budgettable.delete(where={"budget_id": budget_id})
async def test_key_team_change_accepted_by_target_team_admin(
proxy_client, prisma, scratch, world
gateway_client, db, resources, actors
):
"""Caller is admin of the destination team AND the key's owner — the
common_key_access_checks gate requires user_id-match for non-proxy-admin
callers, so the team-admin-accepts branch of validate_key_team_change
can only be reached when the team admin is also the key holder.
Hits validate_key_team_change line 3007–3011."""
actor_user_id = f"{scratch.prefix}-self-admin"
actor_user_id = f"{resources.prefix}-self-admin"
actor_cleartext = "sk-" + uuid.uuid4().hex
await prisma.db.litellm_usertable.create(
await db.db.litellm_usertable.create(
data={"user_id": actor_user_id, "user_role": "internal_user"}
)
await prisma.db.litellm_verificationtoken.create(
await db.db.litellm_verificationtoken.create(
data={
"token": hash_token(actor_cleartext),
"key_alias": f"{scratch.prefix}-actor-key",
"key_name": f"{scratch.prefix}-actor-key",
"key_alias": f"{resources.prefix}-actor-key",
"key_name": f"{resources.prefix}-actor-key",
"user_id": actor_user_id,
"models": [],
"allowed_routes": ["/key/update"],
}
)
source_team = await create_scratch_team(
prisma,
team_id=scratch.tag("source"),
db,
team_id=resources.tag("source"),
admin_user_ids=[actor_user_id],
)
target_team = await create_scratch_team(
prisma,
team_id=scratch.tag("target"),
db,
team_id=resources.tag("target"),
admin_user_ids=[actor_user_id],
)
key_cleartext = await _seed_key_for_relocation(
prisma, scratch.prefix, user_id=actor_user_id, team_id=source_team
db, resources.prefix, user_id=actor_user_id, team_id=source_team
)
resp = await proxy_client.post(
resp = await gateway_client.post(
"/key/update",
headers={"Authorization": f"Bearer {actor_cleartext}"},
json={"key": key_cleartext, "team_id": target_team},
)
assert resp.status_code == 200, resp.text
row = await prisma.db.litellm_verificationtoken.find_unique(
row = await db.db.litellm_verificationtoken.find_unique(
where={"token": hash_token(key_cleartext)}
)
assert row is not None

550
tests/proxy_behavior/management/test_f7_key_coverage_push.py → tests/integration/management/test_f7_key_coverage_push.py
View file

@ -1,38 +1,20 @@
"""Phase 4 F7-extension — additional payload pins pushing
`key_management_endpoints.py` past the 70 % stretch.
Same pattern as `test_f7_coverage_closeout.py`: each scenario cites the
file:line range it pins and asserts observable end-state, not response-body
snapshots. Targets the largest non-deferred un-covered ranges left after
the first F7 pass:
* 5942–6063 /key/health logging-metadata path (test_key_logging body)
* 4630–4692 /key/reset_spend happy path + 404 + admin gate
* 4565–4596 _validate_reset_spend_value branches
* 4421–4476 /key/regenerate ghost-key 404 + premium gate
* 6118–6133 _enforce_unique_key_alias duplicate-alias rejection
* 6148–6169 validate_model_max_budget malformed payload rejection
* 4708–4789 validate_key_list_check user/team/org/key_hash branches
Excluded: `_rotate_master_key` (lines 3997–4123) — deferred per plan §6.
"""
import uuid
from typing import Any, Dict, Optional
from datetime import datetime, timedelta, timezone
from typing import Any, Dict, Final, Optional
import pytest
from prisma import Json
from litellm.proxy.utils import hash_token
from .actors import TEAM_ALPHA, TEAM_BETA, Actor
from .conftest import create_scratch_team
from tests.integration._support.actors import TEAM_ALPHA, TEAM_BETA, Actor
from tests.integration._support.actors import create_scratch_key, create_scratch_team
pytestmark = pytest.mark.asyncio(loop_scope="session")
async def _seed_token(
prisma,
db,
scratch_prefix: str,
*,
suffix: str = "tok",
@ -57,7 +39,7 @@ async def _seed_token(
data["metadata"] = Json(metadata)
if team_id is not None:
data["team_id"] = team_id
await prisma.db.litellm_verificationtoken.create(data=data)
await db.db.litellm_verificationtoken.create(data=data)
return cleartext
@ -71,19 +53,19 @@ async def _seed_token(
async def test_key_health_with_logging_metadata_runs_test_logging(
proxy_client, prisma, scratch, world
gateway_client, db, resources, actors
):
cleartext = await _seed_token(
prisma,
scratch.prefix,
user_id=world.keys[Actor.PROXY_ADMIN].user_id,
db,
resources.prefix,
user_id=actors.keys[Actor.INTERNAL_USER].user_id,
metadata={
"logging": [
{"callback_name": "noop-scratch-callback"},
{"callback_name": "noop-resources-callback"},
]
},
)
resp = await proxy_client.post(
resp = await gateway_client.post(
"/key/health",
headers={"Authorization": f"Bearer {cleartext}"},
)
@ -92,21 +74,21 @@ async def test_key_health_with_logging_metadata_runs_test_logging(
# Either healthy or unhealthy — both pin the logging branch.
assert body["key"] in ("healthy", "unhealthy")
assert "logging_callbacks" in body
assert body["logging_callbacks"]["callbacks"] == ["noop-scratch-callback"]
assert body["logging_callbacks"]["callbacks"] == ["noop-resources-callback"]
async def test_key_health_with_missing_callback_name_rejected(
proxy_client, prisma, scratch, world
gateway_client, db, resources, actors
):
"""test_key_logging raises ValueError if a callback dict lacks
callback_name — wrapped by the outer try/except into a 500."""
cleartext = await _seed_token(
prisma,
scratch.prefix,
user_id=world.keys[Actor.PROXY_ADMIN].user_id,
db,
resources.prefix,
user_id=actors.keys[Actor.PROXY_ADMIN].user_id,
metadata={"logging": [{"not_callback_name": "x"}]},
)
resp = await proxy_client.post(
resp = await gateway_client.post(
"/key/health",
headers={"Authorization": f"Bearer {cleartext}"},
)
@ -124,25 +106,25 @@ async def test_key_health_with_missing_callback_name_rejected(
# ---------------------------------------------------------------------------
async def test_reset_spend_ghost_key_404(proxy_client, scratch, world):
seeder = world.keys[Actor.PROXY_ADMIN].cleartext
resp = await proxy_client.post(
f"/key/sk-{scratch.prefix}-ghost/reset_spend",
async def test_reset_spend_ghost_key_404(gateway_client, resources, actors):
seeder = actors.keys[Actor.PROXY_ADMIN].cleartext
resp = await gateway_client.post(
f"/key/sk-{resources.prefix}-ghost/reset_spend",
headers={"Authorization": f"Bearer {seeder}"},
json={"reset_to": 0.0},
)
assert resp.status_code == 404, resp.text
async def test_reset_spend_happy_path(proxy_client, prisma, scratch, world):
async def test_reset_spend_happy_path(gateway_client, db, resources, actors):
cleartext = await _seed_token(
prisma,
scratch.prefix,
user_id=world.keys[Actor.PROXY_ADMIN].user_id,
db,
resources.prefix,
user_id=actors.keys[Actor.PROXY_ADMIN].user_id,
spend=5.0,
)
seeder = world.keys[Actor.PROXY_ADMIN].cleartext
resp = await proxy_client.post(
seeder = actors.keys[Actor.PROXY_ADMIN].cleartext
resp = await gateway_client.post(
f"/key/{cleartext}/reset_spend",
headers={"Authorization": f"Bearer {seeder}"},
json={"reset_to": 1.0},
@ -151,7 +133,7 @@ async def test_reset_spend_happy_path(proxy_client, prisma, scratch, world):
assert resp.json()["spend"] == 1.0
assert resp.json()["previous_spend"] == 5.0
row = await prisma.db.litellm_verificationtoken.find_unique(
row = await db.db.litellm_verificationtoken.find_unique(
where={"token": hash_token(cleartext)}
)
assert row is not None
@ -171,16 +153,16 @@ async def test_reset_spend_happy_path(proxy_client, prisma, scratch, world):
ids=["negative", "above_current_spend"],
)
async def test_reset_spend_validate_value_branches(
reset_to, expected_detail: str, proxy_client, prisma, scratch, world
reset_to, expected_detail: str, gateway_client, db, resources, actors
):
cleartext = await _seed_token(
prisma,
scratch.prefix,
user_id=world.keys[Actor.PROXY_ADMIN].user_id,
db,
resources.prefix,
user_id=actors.keys[Actor.PROXY_ADMIN].user_id,
spend=5.0,
)
seeder = world.keys[Actor.PROXY_ADMIN].cleartext
resp = await proxy_client.post(
seeder = actors.keys[Actor.PROXY_ADMIN].cleartext
resp = await gateway_client.post(
f"/key/{cleartext}/reset_spend",
headers={"Authorization": f"Bearer {seeder}"},
json={"reset_to": reset_to},
@ -188,50 +170,28 @@ async def test_reset_spend_validate_value_branches(
assert resp.status_code == 400, resp.text
assert expected_detail in resp.text, resp.text
# Row spend must be unchanged.
row = await prisma.db.litellm_verificationtoken.find_unique(
row = await db.db.litellm_verificationtoken.find_unique(
where={"token": hash_token(cleartext)}
)
assert row.spend == 5.0, "spend mutated despite validation rejection"
async def test_reset_spend_non_numeric_caught_by_pydantic(
proxy_client, prisma, scratch, world
):
"""Pin: non-float reset_to is rejected at the Pydantic layer with 422
before reaching _validate_reset_spend_value. This documents that the
helper's `isinstance(reset_to, (int, float))` guard at line 4568 is
structurally unreachable via HTTP."""
cleartext = await _seed_token(
prisma,
scratch.prefix,
user_id=world.keys[Actor.PROXY_ADMIN].user_id,
spend=5.0,
)
seeder = world.keys[Actor.PROXY_ADMIN].cleartext
resp = await proxy_client.post(
f"/key/{cleartext}/reset_spend",
headers={"Authorization": f"Bearer {seeder}"},
json={"reset_to": "not-a-number"},
)
assert resp.status_code == 422, resp.text
async def test_reset_spend_non_admin_caller_rejected(
proxy_client, prisma, scratch, world
gateway_client, db, resources, actors
):
"""_check_proxy_or_team_admin_for_key raises 403 when caller is neither
proxy admin nor admin of the key's team."""
# Seed a key against TEAM_BETA (no actor in our world is admin of beta).
# Seed a key against actors.team_beta_id (no actor in our actors is admin of beta).
cleartext = await _seed_token(
prisma,
scratch.prefix,
user_id=world.keys[Actor.CROSS_ORG_USER].user_id,
db,
resources.prefix,
user_id=actors.keys[Actor.CROSS_ORG_USER].user_id,
spend=5.0,
team_id=TEAM_BETA,
team_id=actors.team_beta_id,
)
# Caller is a member of TEAM_ALPHA but not admin of TEAM_BETA.
initiator = world.keys[Actor.INTERNAL_USER].cleartext
resp = await proxy_client.post(
# Caller is a member of actors.team_alpha_id but not admin of actors.team_beta_id.
initiator = actors.keys[Actor.INTERNAL_USER].cleartext
resp = await gateway_client.post(
f"/key/{cleartext}/reset_spend",
headers={"Authorization": f"Bearer {initiator}"},
json={"reset_to": 1.0},
@ -239,7 +199,7 @@ async def test_reset_spend_non_admin_caller_rejected(
# Route-level admin gate may fire first (401) or the helper's own 403 —
# both prove the path is guarded; pin either as rejection.
assert resp.status_code in (401, 403), resp.text
row = await prisma.db.litellm_verificationtoken.find_unique(
row = await db.db.litellm_verificationtoken.find_unique(
where={"token": hash_token(cleartext)}
)
assert row.spend == 5.0, "spend mutated despite rejection"
@ -252,37 +212,38 @@ async def test_reset_spend_non_admin_caller_rejected(
# ---------------------------------------------------------------------------
async def test_regenerate_ghost_key_404(proxy_client, scratch, world):
seeder = world.keys[Actor.PROXY_ADMIN].cleartext
resp = await proxy_client.post(
f"/key/sk-{scratch.prefix}-ghost/regenerate",
async def test_regenerate_ghost_key_404(gateway_client, resources, actors):
seeder = actors.keys[Actor.PROXY_ADMIN].cleartext
resp = await gateway_client.post(
f"/key/sk-{resources.prefix}-ghost/regenerate",
headers={"Authorization": f"Bearer {seeder}"},
json={},
)
assert resp.status_code == 404, resp.text
async def test_regenerate_no_key_supplied_400(proxy_client, world):
"""POST /key/regenerate with no key in path AND no key in body."""
seeder = world.keys[Actor.PROXY_ADMIN].cleartext
resp = await proxy_client.post(
async def test_regenerate_no_key_supplied_400(gateway_client, actors):
seeder = actors.keys[Actor.PROXY_ADMIN].cleartext
resp = await gateway_client.post(
"/key/regenerate",
headers={"Authorization": f"Bearer {seeder}"},
json={},
)
assert resp.status_code == 400, resp.text
assert "No key passed in" in resp.text or "key" in resp.text.lower()
error = resp.json()["error"]
assert error["code"] == "400", resp.text
assert "No key passed in." in error["message"], resp.text
async def test_regenerate_happy_path(proxy_client, prisma, scratch, world):
async def test_regenerate_happy_path(gateway_client, db, resources, actors):
cleartext = await _seed_token(
prisma,
scratch.prefix,
user_id=world.keys[Actor.PROXY_ADMIN].user_id,
db,
resources.prefix,
user_id=actors.keys[Actor.PROXY_ADMIN].user_id,
)
old_hash = hash_token(cleartext)
seeder = world.keys[Actor.PROXY_ADMIN].cleartext
resp = await proxy_client.post(
seeder = actors.keys[Actor.PROXY_ADMIN].cleartext
resp = await gateway_client.post(
f"/key/{cleartext}/regenerate",
headers={"Authorization": f"Bearer {seeder}"},
json={},
@ -294,35 +255,35 @@ async def test_regenerate_happy_path(proxy_client, prisma, scratch, world):
assert new_hash != old_hash
# Old token should no longer be in active tokens (deleted by regenerate).
old_row = await prisma.db.litellm_verificationtoken.find_unique(
old_row = await db.db.litellm_verificationtoken.find_unique(
where={"token": old_hash}
)
assert old_row is None, "old token still present after regenerate"
# New token should be active.
new_row = await prisma.db.litellm_verificationtoken.find_unique(
new_row = await db.db.litellm_verificationtoken.find_unique(
where={"token": new_hash}
)
assert new_row is not None, "new token not written after regenerate"
async def test_regenerate_with_explicit_new_key(proxy_client, prisma, scratch, world):
async def test_regenerate_with_explicit_new_key(gateway_client, db, resources, actors):
cleartext = await _seed_token(
prisma,
scratch.prefix,
user_id=world.keys[Actor.PROXY_ADMIN].user_id,
db,
resources.prefix,
user_id=actors.keys[Actor.PROXY_ADMIN].user_id,
suffix="explicit",
)
seeder = world.keys[Actor.PROXY_ADMIN].cleartext
seeder = actors.keys[Actor.PROXY_ADMIN].cleartext
new_key = "sk-" + uuid.uuid4().hex
resp = await proxy_client.post(
resp = await gateway_client.post(
f"/key/{cleartext}/regenerate",
headers={"Authorization": f"Bearer {seeder}"},
json={"new_key": new_key},
)
assert resp.status_code == 200, resp.text
assert resp.json()["key"] == new_key
new_row = await prisma.db.litellm_verificationtoken.find_unique(
new_row = await db.db.litellm_verificationtoken.find_unique(
where={"token": hash_token(new_key)}
)
assert new_row is not None
@ -334,16 +295,16 @@ async def test_regenerate_with_explicit_new_key(proxy_client, prisma, scratch, w
# ---------------------------------------------------------------------------
async def test_generate_duplicate_alias_rejected(proxy_client, prisma, scratch, world):
seeder = world.keys[Actor.PROXY_ADMIN].cleartext
alias = scratch.prefix + "-shared-alias"
first = await proxy_client.post(
async def test_generate_duplicate_alias_rejected(gateway_client, db, resources, actors):
seeder = actors.keys[Actor.PROXY_ADMIN].cleartext
alias = resources.prefix + "-shared-alias"
first = await gateway_client.post(
"/key/generate",
headers={"Authorization": f"Bearer {seeder}"},
json={"key_alias": alias},
)
assert first.status_code == 200, first.text
second = await proxy_client.post(
second = await gateway_client.post(
"/key/generate",
headers={"Authorization": f"Bearer {seeder}"},
json={"key_alias": alias},
@ -359,14 +320,14 @@ async def test_generate_duplicate_alias_rejected(proxy_client, prisma, scratch,
async def test_generate_with_invalid_model_max_budget_rejected(
proxy_client, scratch, world
gateway_client, resources, actors
):
seeder = world.keys[Actor.PROXY_ADMIN].cleartext
resp = await proxy_client.post(
seeder = actors.keys[Actor.PROXY_ADMIN].cleartext
resp = await gateway_client.post(
"/key/generate",
headers={"Authorization": f"Bearer {seeder}"},
json={
"key_alias": scratch.prefix,
"key_alias": resources.prefix,
# Wrong shape — budget_limit should be numeric; passing a dict
# for the model value bypasses the BudgetConfig parse and trips
# the validator's exception wrap.
@ -401,14 +362,14 @@ async def test_generate_with_invalid_model_max_budget_rejected(
ids=["user_mismatch", "team_mismatch", "org_mismatch"],
)
async def test_key_list_non_admin_authz_branches(
filter_kwarg, expected_substring: str, proxy_client, world
filter_kwarg, expected_substring: str, gateway_client, actors
):
"""Non-admin caller hits validate_key_list_check's three rejection
branches. INTERNAL_USER (Org A, TEAM_ALPHA member) is the caller; each
branches. INTERNAL_USER (Org A, actors.team_alpha_id member) is the caller; each
filter targets a foreign user/team/org and trips the matching guard."""
caller = world.keys[Actor.INTERNAL_USER].cleartext
caller = actors.keys[Actor.INTERNAL_USER].cleartext
qs = "&".join(f"{k}={v}" for k, v in filter_kwarg.items())
resp = await proxy_client.get(
resp = await gateway_client.get(
f"/key/list?{qs}",
headers={"Authorization": f"Bearer {caller}"},
)
@ -424,22 +385,22 @@ async def test_key_list_non_admin_authz_branches(
async def test_key_bulk_update_mixed_success_and_failure(
proxy_client, prisma, scratch, world
gateway_client, db, resources, actors
):
seeder = world.keys[Actor.PROXY_ADMIN].cleartext
seeder = actors.keys[Actor.PROXY_ADMIN].cleartext
existing = await _seed_token(
prisma,
scratch.prefix,
user_id=world.keys[Actor.PROXY_ADMIN].user_id,
db,
resources.prefix,
user_id=actors.keys[Actor.PROXY_ADMIN].user_id,
suffix="bulk1",
)
resp = await proxy_client.post(
resp = await gateway_client.post(
"/key/bulk_update",
headers={"Authorization": f"Bearer {seeder}"},
json={
"keys": [
{"key": existing, "max_budget": 5.0},
{"key": "sk-ghost-" + scratch.prefix, "max_budget": 5.0},
{"key": "sk-ghost-" + resources.prefix, "max_budget": 5.0},
]
},
)
@ -449,16 +410,16 @@ async def test_key_bulk_update_mixed_success_and_failure(
assert len(body["successful_updates"]) == 1
assert len(body["failed_updates"]) == 1
# Re-read confirms the successful key was actually updated.
row = await prisma.db.litellm_verificationtoken.find_unique(
row = await db.db.litellm_verificationtoken.find_unique(
where={"token": hash_token(existing)}
)
assert row.max_budget == 5.0
async def test_key_bulk_update_non_admin_rejected(proxy_client, world):
async def test_key_bulk_update_non_admin_rejected(gateway_client, actors):
"""Lines 2631–2635 — admin-only role gate."""
caller = world.keys[Actor.INTERNAL_USER].cleartext
resp = await proxy_client.post(
caller = actors.keys[Actor.INTERNAL_USER].cleartext
resp = await gateway_client.post(
"/key/bulk_update",
headers={"Authorization": f"Bearer {caller}"},
json={"keys": [{"key": "sk-x", "max_budget": 1.0}]},
@ -466,10 +427,10 @@ async def test_key_bulk_update_non_admin_rejected(proxy_client, world):
assert resp.status_code in (401, 403), resp.text
async def test_key_bulk_update_empty_keys_rejected(proxy_client, world):
async def test_key_bulk_update_empty_keys_rejected(gateway_client, actors):
"""Line 2643–2647 — empty keys list rejected."""
seeder = world.keys[Actor.PROXY_ADMIN].cleartext
resp = await proxy_client.post(
seeder = actors.keys[Actor.PROXY_ADMIN].cleartext
resp = await gateway_client.post(
"/key/bulk_update",
headers={"Authorization": f"Bearer {seeder}"},
json={"keys": []},
@ -478,12 +439,12 @@ async def test_key_bulk_update_empty_keys_rejected(proxy_client, world):
assert "No keys" in resp.text
async def test_key_bulk_update_exceeds_max_batch_rejected(proxy_client, world):
async def test_key_bulk_update_exceeds_max_batch_rejected(gateway_client, actors):
"""Lines 2649–2656 — over-batch-size rejection. 501 ghost keys are fine
here because validation fires before any update runs."""
seeder = world.keys[Actor.PROXY_ADMIN].cleartext
seeder = actors.keys[Actor.PROXY_ADMIN].cleartext
over_batch = [{"key": f"sk-{i}", "max_budget": 1.0} for i in range(501)]
resp = await proxy_client.post(
resp = await gateway_client.post(
"/key/bulk_update",
headers={"Authorization": f"Bearer {seeder}"},
json={"keys": over_batch},
@ -501,15 +462,15 @@ async def test_key_bulk_update_exceeds_max_batch_rejected(proxy_client, world):
async def test_key_update_with_duration_and_budget_duration(
proxy_client, prisma, scratch, world
gateway_client, db, resources, actors
):
cleartext = await _seed_token(
prisma,
scratch.prefix,
user_id=world.keys[Actor.PROXY_ADMIN].user_id,
db,
resources.prefix,
user_id=actors.keys[Actor.PROXY_ADMIN].user_id,
)
seeder = world.keys[Actor.PROXY_ADMIN].cleartext
resp = await proxy_client.post(
seeder = actors.keys[Actor.PROXY_ADMIN].cleartext
resp = await gateway_client.post(
"/key/update",
headers={"Authorization": f"Bearer {seeder}"},
json={
@ -519,35 +480,48 @@ async def test_key_update_with_duration_and_budget_duration(
},
)
assert resp.status_code == 200, resp.text
row = await prisma.db.litellm_verificationtoken.find_unique(
row = await db.db.litellm_verificationtoken.find_unique(
where={"token": hash_token(cleartext)}
)
assert row.expires is not None, "expires not stamped from duration"
assert row.budget_reset_at is not None, "budget_reset_at not stamped"
async def test_key_update_with_clear_duration(proxy_client, prisma, scratch, world):
"""`duration: -1` clears expires (line 1796–1798)."""
cleartext = await _seed_token(
prisma,
scratch.prefix,
user_id=world.keys[Actor.PROXY_ADMIN].user_id,
assert row is not None
assert row.budget_duration == "30d"
now = datetime.now(timezone.utc)
assert timedelta(minutes=55) < row.expires - now <= timedelta(hours=1), row.expires
assert now < row.budget_reset_at <= now + timedelta(days=31), row.budget_reset_at
info = await gateway_client.get(
f"/key/info?key={hash_token(cleartext)}",
headers={"Authorization": f"Bearer {seeder}"},
)
seeder = world.keys[Actor.PROXY_ADMIN].cleartext
# First set an expiry
await proxy_client.post(
assert info.status_code == 200, info.text
assert info.json()["info"]["expires"] == row.expires.isoformat()
assert info.json()["info"]["budget_reset_at"] == row.budget_reset_at.isoformat()
async def test_key_update_with_clear_duration(gateway_client, db, resources, actors):
cleartext = await _seed_token(
db,
resources.prefix,
user_id=actors.keys[Actor.PROXY_ADMIN].user_id,
)
seeder = actors.keys[Actor.PROXY_ADMIN].cleartext
expiring = await gateway_client.post(
"/key/update",
headers={"Authorization": f"Bearer {seeder}"},
json={"key": cleartext, "duration": "1h"},
)
# Then clear it
resp = await proxy_client.post(
assert expiring.status_code == 200, expiring.text
expiring_row = await db.db.litellm_verificationtoken.find_unique(
where={"token": hash_token(cleartext)}
)
assert expiring_row.expires is not None, "duration=1h did not set expires"
resp = await gateway_client.post(
"/key/update",
headers={"Authorization": f"Bearer {seeder}"},
json={"key": cleartext, "duration": "-1"},
)
assert resp.status_code == 200, resp.text
row = await prisma.db.litellm_verificationtoken.find_unique(
row = await db.db.litellm_verificationtoken.find_unique(
where={"token": hash_token(cleartext)}
)
assert row.expires is None, "expires not cleared by duration=-1"
@ -564,27 +538,15 @@ async def test_key_update_with_clear_duration(proxy_client, prisma, scratch, wor
# ---------------------------------------------------------------------------
async def test_team_key_bulk_update_missing_team_id_rejected(proxy_client, world):
"""Pydantic catches missing team_id at the request layer → 422 before
the handler's own `if not data.team_id` guard at line 2803."""
seeder = world.keys[Actor.PROXY_ADMIN].cleartext
resp = await proxy_client.post(
"/team/key/bulk_update",
headers={"Authorization": f"Bearer {seeder}"},
json={"update_fields": {"max_budget": 5.0}},
)
assert resp.status_code == 422, resp.text
async def test_team_key_bulk_update_no_selector_rejected(proxy_client, world):
async def test_team_key_bulk_update_no_selector_rejected(gateway_client, actors):
"""Pydantic root-validator catches missing key_ids/all_keys_in_team
at the request layer → 422."""
seeder = world.keys[Actor.PROXY_ADMIN].cleartext
resp = await proxy_client.post(
seeder = actors.keys[Actor.PROXY_ADMIN].cleartext
resp = await gateway_client.post(
"/team/key/bulk_update",
headers={"Authorization": f"Bearer {seeder}"},
json={
"team_id": TEAM_ALPHA,
"team_id": actors.team_alpha_id,
"update_fields": {"max_budget": 5.0},
},
)
@ -593,31 +555,31 @@ async def test_team_key_bulk_update_no_selector_rejected(proxy_client, world):
async def test_team_key_bulk_update_all_keys_in_team(
proxy_client, prisma, scratch, world
gateway_client, db, resources, actors
):
"""Pin the all_keys_in_team branch (lines 2818–2841) plus the
per-key success path. Seed two scratch keys against the scratch team."""
seeder = world.keys[Actor.PROXY_ADMIN].cleartext
per-key success path. Seed two resources keys against the resources team."""
seeder = actors.keys[Actor.PROXY_ADMIN].cleartext
team_id = await create_scratch_team(
prisma,
team_id=scratch.tag("team"),
admin_user_ids=[world.keys[Actor.PROXY_ADMIN].user_id],
db,
team_id=resources.tag("team"),
admin_user_ids=[actors.keys[Actor.PROXY_ADMIN].user_id],
)
k1 = await _seed_token(
prisma,
scratch.prefix,
user_id=world.keys[Actor.PROXY_ADMIN].user_id,
db,
resources.prefix,
user_id=actors.keys[Actor.PROXY_ADMIN].user_id,
suffix="t1",
team_id=team_id,
)
k2 = await _seed_token(
prisma,
scratch.prefix,
user_id=world.keys[Actor.PROXY_ADMIN].user_id,
db,
resources.prefix,
user_id=actors.keys[Actor.PROXY_ADMIN].user_id,
suffix="t2",
team_id=team_id,
)
resp = await proxy_client.post(
resp = await gateway_client.post(
"/team/key/bulk_update",
headers={"Authorization": f"Bearer {seeder}"},
json={
@ -631,36 +593,36 @@ async def test_team_key_bulk_update_all_keys_in_team(
assert body["total_requested"] == 2
# Re-read both keys; max_budget should be set.
for k in (k1, k2):
row = await prisma.db.litellm_verificationtoken.find_unique(
row = await db.db.litellm_verificationtoken.find_unique(
where={"token": hash_token(k)}
)
assert row.max_budget == 7.0
async def test_team_key_bulk_update_explicit_key_ids_mixed(
proxy_client, prisma, scratch, world
gateway_client, db, resources, actors
):
"""Pin the explicit-key_ids dedupe + per-key 404 path (lines 2842–2944).
Send a real key + a ghost key under the same team_id."""
seeder = world.keys[Actor.PROXY_ADMIN].cleartext
seeder = actors.keys[Actor.PROXY_ADMIN].cleartext
team_id = await create_scratch_team(
prisma,
team_id=scratch.tag("team"),
admin_user_ids=[world.keys[Actor.PROXY_ADMIN].user_id],
db,
team_id=resources.tag("team"),
admin_user_ids=[actors.keys[Actor.PROXY_ADMIN].user_id],
)
real_key = await _seed_token(
prisma,
scratch.prefix,
user_id=world.keys[Actor.PROXY_ADMIN].user_id,
db,
resources.prefix,
user_id=actors.keys[Actor.PROXY_ADMIN].user_id,
suffix="real",
team_id=team_id,
)
resp = await proxy_client.post(
resp = await gateway_client.post(
"/team/key/bulk_update",
headers={"Authorization": f"Bearer {seeder}"},
json={
"team_id": team_id,
"key_ids": [real_key, real_key, "sk-ghost-" + scratch.prefix],
"key_ids": [real_key, real_key, "sk-ghost-" + resources.prefix],
"update_fields": {"max_budget": 9.0},
},
)
@ -686,24 +648,24 @@ async def test_team_key_bulk_update_explicit_key_ids_mixed(
# ---------------------------------------------------------------------------
async def test_key_info_ghost_key_404(proxy_client, scratch, world):
seeder = world.keys[Actor.PROXY_ADMIN].cleartext
resp = await proxy_client.get(
f"/key/info?key=sk-{scratch.prefix}-ghost",
async def test_key_info_ghost_key_404(gateway_client, resources, actors):
seeder = actors.keys[Actor.PROXY_ADMIN].cleartext
resp = await gateway_client.get(
f"/key/info?key=sk-{resources.prefix}-ghost",
headers={"Authorization": f"Bearer {seeder}"},
)
assert resp.status_code == 404, resp.text
async def test_key_info_explicit_existing_key(proxy_client, prisma, scratch, world):
async def test_key_info_explicit_existing_key(gateway_client, db, resources, actors):
cleartext = await _seed_token(
prisma,
scratch.prefix,
user_id=world.keys[Actor.PROXY_ADMIN].user_id,
db,
resources.prefix,
user_id=actors.keys[Actor.PROXY_ADMIN].user_id,
suffix="info-test",
)
seeder = world.keys[Actor.PROXY_ADMIN].cleartext
resp = await proxy_client.get(
seeder = actors.keys[Actor.PROXY_ADMIN].cleartext
resp = await gateway_client.get(
f"/key/info?key={cleartext}",
headers={"Authorization": f"Bearer {seeder}"},
)
@ -715,15 +677,18 @@ async def test_key_info_explicit_existing_key(proxy_client, prisma, scratch, wor
assert "token" not in body["info"]
async def test_key_info_no_key_uses_auth_header(proxy_client, world):
"""Pin line 3260: `key = key or user_api_key_dict.api_key` — caller's
own key info is returned when no `?key=` is supplied."""
caller = world.keys[Actor.INTERNAL_USER].cleartext
resp = await proxy_client.get(
async def test_key_info_no_key_uses_auth_header(gateway_client, actors):
caller = actors.keys[Actor.INTERNAL_USER].cleartext
resp = await gateway_client.get(
"/key/info",
headers={"Authorization": f"Bearer {caller}"},
)
assert resp.status_code == 200, resp.text
body: Final = resp.json()
assert body["key"] == hash_token(caller)
assert body["info"]["user_id"] == actors.keys[Actor.INTERNAL_USER].user_id
assert body["info"]["status"] == "active"
assert "token" not in body["info"]
# ---------------------------------------------------------------------------
@ -732,13 +697,13 @@ async def test_key_info_no_key_uses_auth_header(proxy_client, world):
# ---------------------------------------------------------------------------
async def test_key_generate_with_budget_limits(proxy_client, prisma, scratch, world):
seeder = world.keys[Actor.PROXY_ADMIN].cleartext
resp = await proxy_client.post(
async def test_key_generate_with_budget_limits(gateway_client, db, resources, actors):
seeder = actors.keys[Actor.PROXY_ADMIN].cleartext
resp = await gateway_client.post(
"/key/generate",
headers={"Authorization": f"Bearer {seeder}"},
json={
"key_alias": scratch.prefix,
"key_alias": resources.prefix,
"budget_limits": [
{"max_budget": 5.0, "budget_duration": "1d"},
{"max_budget": 20.0, "budget_duration": "30d"},
@ -746,17 +711,27 @@ async def test_key_generate_with_budget_limits(proxy_client, prisma, scratch, wo
},
)
assert resp.status_code == 200, resp.text
cleartext: Final = resp.json()["key"]
info: Final = await gateway_client.get(
f"/key/info?key={hash_token(cleartext)}",
headers={"Authorization": f"Bearer {seeder}"},
)
assert info.status_code == 200, info.text
windows: Final = info.json()["info"]["budget_limits"]
assert [
(window["max_budget"], window["budget_duration"]) for window in windows
] == [(5.0, "1d"), (20.0, "30d")]
async def test_key_aliases_proxy_admin_unscoped(proxy_client, prisma, scratch, world):
seeder = world.keys[Actor.PROXY_ADMIN].cleartext
alias = scratch.prefix + "-aliases-test"
await proxy_client.post(
async def test_key_aliases_proxy_admin_unscoped(gateway_client, db, resources, actors):
seeder = actors.keys[Actor.PROXY_ADMIN].cleartext
alias = resources.prefix + "-aliases-test"
await gateway_client.post(
"/key/generate",
headers={"Authorization": f"Bearer {seeder}"},
json={"key_alias": alias},
)
resp = await proxy_client.get(
resp = await gateway_client.get(
"/key/aliases?size=10",
headers={"Authorization": f"Bearer {seeder}"},
)
@ -771,36 +746,23 @@ async def test_key_aliases_proxy_admin_unscoped(proxy_client, prisma, scratch, w
), f"newly created alias not in list: {body['aliases']}"
async def test_key_aliases_with_search_filter(proxy_client, prisma, scratch, world):
async def test_key_aliases_with_search_filter(gateway_client, db, resources, actors):
"""Pin the `search` ILIKE branch (lines 5166–5168)."""
seeder = world.keys[Actor.PROXY_ADMIN].cleartext
unique_alias = scratch.prefix + "-search-unique-tag"
await proxy_client.post(
seeder = actors.keys[Actor.PROXY_ADMIN].cleartext
unique_alias = resources.prefix + "-search-unique-tag"
await gateway_client.post(
"/key/generate",
headers={"Authorization": f"Bearer {seeder}"},
json={"key_alias": unique_alias},
)
resp = await proxy_client.get(
f"/key/aliases?search={scratch.prefix}-search",
resp = await gateway_client.get(
f"/key/aliases?search={resources.prefix}-search",
headers={"Authorization": f"Bearer {seeder}"},
)
assert resp.status_code == 200, resp.text
assert unique_alias in resp.json()["aliases"]
async def test_key_aliases_non_admin_scoped(proxy_client, world):
"""Non-admin caller routes through _apply_non_admin_alias_scope (line
5161–5164). The exact alias visibility depends on team membership; the
pin is that the call succeeds with a scoped result."""
caller = world.keys[Actor.INTERNAL_USER].cleartext
resp = await proxy_client.get(
"/key/aliases?size=10",
headers={"Authorization": f"Bearer {caller}"},
)
assert resp.status_code == 200, resp.text
assert "aliases" in resp.json()
# ---------------------------------------------------------------------------
# /key/list with extended filters — pins _build_filter_conditions branches
# at lines 5280–5388. Each scenario varies one filter so a different
@ -809,23 +771,51 @@ async def test_key_aliases_non_admin_scoped(proxy_client, world):
@pytest.mark.parametrize(
"query",
"filter_name",
[
"include_created_by_keys=true",
"include_team_keys=true",
"return_full_object=true",
"sort_by=created_at&sort_order=asc",
"key_alias=behavior",
"created_by",
"team_keys",
"full_object",
"sort",
"alias_substring",
],
ids=["created_by", "team_keys", "full_object", "sort", "alias_substring"],
)
async def test_key_list_filter_branches(query: str, proxy_client, world):
seeder = world.keys[Actor.PROXY_ADMIN].cleartext
resp = await proxy_client.get(
f"/key/list?{query}",
headers={"Authorization": f"Bearer {seeder}"},
async def test_key_list_filter_branches(filter_name: str, gateway_client, resources, actors):
admin: Final = actors.keys[Actor.PROXY_ADMIN]
key_alias: Final = resources.tag("filter-branch")
cleartext: Final = await create_scratch_key(
gateway_client,
admin.cleartext,
resources.prefix,
user_id=admin.user_id,
key_alias=key_alias,
team_id=actors.team_alpha_id,
)
query: Final = {
"created_by": "include_created_by_keys=true",
"team_keys": "include_team_keys=true",
"full_object": "return_full_object=true",
"sort": "sort_by=created_at&sort_order=asc",
"alias_substring": f"key_alias={key_alias[:-5]}&substring_matching=true",
}[filter_name]
list_query: Final = (
query
if filter_name == "alias_substring"
else f"{query}&key_alias={key_alias}"
)
resp: Final = await gateway_client.get(
f"/key/list?{list_query}",
headers={"Authorization": f"Bearer {admin.cleartext}"},
)
assert resp.status_code == 200, resp.text
response_body: Final = resp.json()
entries: Final = response_body["keys"]
assert response_body["total_count"] == 1
returned_hashes: Final = frozenset(
entry["token"] if isinstance(entry, dict) else entry for entry in entries
)
assert returned_hashes == {hash_token(cleartext)}
# ---------------------------------------------------------------------------
@ -836,17 +826,17 @@ async def test_key_list_filter_branches(query: str, proxy_client, world):
async def test_regenerate_with_grace_period_inserts_deprecated_row(
proxy_client, prisma, scratch, world
gateway_client, db, resources, actors
):
cleartext = await _seed_token(
prisma,
scratch.prefix,
user_id=world.keys[Actor.PROXY_ADMIN].user_id,
db,
resources.prefix,
user_id=actors.keys[Actor.PROXY_ADMIN].user_id,
suffix="grace",
)
old_hash = hash_token(cleartext)
seeder = world.keys[Actor.PROXY_ADMIN].cleartext
resp = await proxy_client.post(
seeder = actors.keys[Actor.PROXY_ADMIN].cleartext
resp = await gateway_client.post(
f"/key/{cleartext}/regenerate",
headers={"Authorization": f"Bearer {seeder}"},
json={"grace_period": "1h"},
@ -854,45 +844,45 @@ async def test_regenerate_with_grace_period_inserts_deprecated_row(
assert resp.status_code == 200, resp.text
new_key = resp.json()["key"]
# Old token should now be in the deprecated table.
deprecated_row = await prisma.db.litellm_deprecatedverificationtoken.find_unique(
deprecated_row = await db.db.litellm_deprecatedverificationtoken.find_unique(
where={"token": old_hash}
)
assert deprecated_row is not None, "old token not retained in deprecated table"
assert deprecated_row.active_token_id == hash_token(new_key)
assert deprecated_row.revoke_at is not None
# Manual cleanup — scratch prefix sweep doesn't cover this table.
await prisma.db.litellm_deprecatedverificationtoken.delete(
# Manual cleanup — resources prefix sweep doesn't cover this table.
await db.db.litellm_deprecatedverificationtoken.delete(
where={"token": old_hash}
)
async def test_regenerate_with_invalid_grace_period_format(
proxy_client, prisma, scratch, world
gateway_client, db, resources, actors
):
"""Invalid grace_period format falls through silently (line 4170–4175);
regenerate still succeeds but no deprecated row is inserted."""
cleartext = await _seed_token(
prisma,
scratch.prefix,
user_id=world.keys[Actor.PROXY_ADMIN].user_id,
db,
resources.prefix,
user_id=actors.keys[Actor.PROXY_ADMIN].user_id,
suffix="grace-bad",
)
old_hash = hash_token(cleartext)
seeder = world.keys[Actor.PROXY_ADMIN].cleartext
resp = await proxy_client.post(
seeder = actors.keys[Actor.PROXY_ADMIN].cleartext
resp = await gateway_client.post(
f"/key/{cleartext}/regenerate",
headers={"Authorization": f"Bearer {seeder}"},
json={"grace_period": "totally-not-a-duration"},
)
assert resp.status_code == 200, resp.text
deprecated_row = await prisma.db.litellm_deprecatedverificationtoken.find_unique(
deprecated_row = await db.db.litellm_deprecatedverificationtoken.find_unique(
where={"token": old_hash}
)
assert deprecated_row is None, "deprecated row created despite invalid grace_period"
async def test_key_list_key_hash_filter_unauthorized(
proxy_client, prisma, scratch, world
gateway_client, db, resources, actors
):
"""validate_key_list_check's key_hash branch (lines 4766–4789): a
cross-tenant non-admin caller asks for a key_hash they don't own → 403.
@ -903,13 +893,13 @@ async def test_key_list_key_hash_filter_unauthorized(
the key owner, team member, nor admin — i.e. CROSS_ORG_USER.
"""
cleartext = await _seed_token(
prisma,
scratch.prefix,
user_id=world.keys[Actor.OWNER].user_id,
team_id=TEAM_ALPHA,
db,
resources.prefix,
user_id=actors.keys[Actor.OWNER].user_id,
team_id=actors.team_alpha_id,
)
caller = world.keys[Actor.CROSS_ORG_USER].cleartext
resp = await proxy_client.get(
caller = actors.keys[Actor.CROSS_ORG_USER].cleartext
resp = await gateway_client.get(
f"/key/list?key_hash={hash_token(cleartext)}",
headers={"Authorization": f"Bearer {caller}"},
)

50
tests/proxy_behavior/management/test_key_aliases.py → tests/integration/management/test_key_aliases.py
View file

@ -5,7 +5,7 @@ import pytest
from litellm.proxy.utils import hash_token
from .actors import TEAM_ALPHA, TEAM_BETA, Actor
from tests.integration._support.actors import TEAM_ALPHA, TEAM_BETA, Actor
pytestmark = pytest.mark.asyncio(loop_scope="session")
@ -16,11 +16,11 @@ pytestmark = pytest.mark.asyncio(loop_scope="session")
# own — owned by INTERNAL_USER, no team -> user_id scope only
# alpha — owned by OWNER, team TEAM_ALPHA -> team scope for alpha members
# beta — owned by CROSS_ORG_USER, TEAM_BETA
async def _seed_alias_keys(prisma, prefix: str, world) -> dict:
async def _seed_alias_keys(db, prefix: str, actors) -> dict:
spec = {
"own": (Actor.INTERNAL_USER, None),
"alpha": (Actor.OWNER, TEAM_ALPHA),
"beta": (Actor.CROSS_ORG_USER, TEAM_BETA),
"alpha": (Actor.OWNER, actors.team_alpha_id),
"beta": (Actor.CROSS_ORG_USER, actors.team_beta_id),
}
out = {}
for tag, (owner, team_id) in spec.items():
@ -29,18 +29,18 @@ async def _seed_alias_keys(prisma, prefix: str, world) -> dict:
"token": hash_token("sk-" + uuid.uuid4().hex),
"key_name": f"{prefix}-{tag}-key",
"key_alias": alias,
"user_id": world.keys[owner].user_id,
"user_id": actors.keys[owner].user_id,
"models": [],
}
if team_id is not None:
data["team_id"] = team_id
await prisma.db.litellm_verificationtoken.create(data=data)
await db.db.litellm_verificationtoken.create(data=data)
out[tag] = alias
return out
async def _fetch_aliases(proxy_client, caller_cleartext: str, query: str) -> set:
resp = await proxy_client.get(
async def _fetch_aliases(gateway_client, caller_cleartext: str, query: str) -> set:
resp = await gateway_client.get(
f"/key/aliases?{query}&size=100",
headers={"Authorization": f"Bearer {caller_cleartext}"},
)
@ -73,17 +73,17 @@ async def test_key_aliases_visibility(
actor: Actor,
expected_status: int,
expected_tags: FrozenSet[str],
proxy_client,
prisma,
scratch,
world,
gateway_client,
db,
resources,
actors,
):
aliases = await _seed_alias_keys(prisma, scratch.prefix, world)
aliases = await _seed_alias_keys(db, resources.prefix, actors)
known = {v: k for k, v in aliases.items()}
resp = await proxy_client.get(
f"/key/aliases?search={scratch.prefix}&size=100",
headers={"Authorization": f"Bearer {world.keys[actor].cleartext}"},
resp = await gateway_client.get(
f"/key/aliases?search={resources.prefix}&size=100",
headers={"Authorization": f"Bearer {actors.keys[actor].cleartext}"},
)
assert (
resp.status_code == expected_status
@ -97,23 +97,23 @@ async def test_key_aliases_visibility(
), f"{actor.value}: expected {sorted(expected_tags)}, got {sorted(visible)}"
async def test_key_aliases_team_id_filter(proxy_client, prisma, scratch, world):
async def test_key_aliases_team_id_filter(gateway_client, db, resources, actors):
"""team_id filter narrows the result to keys of that team."""
aliases = await _seed_alias_keys(prisma, scratch.prefix, world)
aliases = await _seed_alias_keys(db, resources.prefix, actors)
returned = await _fetch_aliases(
proxy_client,
world.keys[Actor.PROXY_ADMIN].cleartext,
f"search={scratch.prefix}&team_id={TEAM_ALPHA}",
gateway_client,
actors.keys[Actor.PROXY_ADMIN].cleartext,
f"search={resources.prefix}&team_id={actors.team_alpha_id}",
)
assert returned & set(aliases.values()) == {aliases["alpha"]}
async def test_key_aliases_search_filter(proxy_client, prisma, scratch, world):
async def test_key_aliases_search_filter(gateway_client, db, resources, actors):
"""search is a case-insensitive substring match on key_alias."""
aliases = await _seed_alias_keys(prisma, scratch.prefix, world)
aliases = await _seed_alias_keys(db, resources.prefix, actors)
returned = await _fetch_aliases(
proxy_client,
world.keys[Actor.PROXY_ADMIN].cleartext,
gateway_client,
actors.keys[Actor.PROXY_ADMIN].cleartext,
f"search={aliases['beta']}",
)
assert returned & set(aliases.values()) == {aliases["beta"]}

60
tests/proxy_behavior/management/test_key_block_unblock.py → tests/integration/management/test_key_block_unblock.py
View file

@ -4,8 +4,8 @@ import pytest
from litellm.proxy.utils import hash_token
from .actors import TEAM_ALPHA, TEAM_BETA, Actor
from .conftest import create_scratch_key
from tests.integration._support.actors import TEAM_ALPHA, TEAM_BETA, Actor
from tests.integration._support.actors import create_scratch_key
pytestmark = pytest.mark.asyncio(loop_scope="session")
@ -40,26 +40,26 @@ _SCENARIOS = [
]
async def _seed_target(proxy_client, seeder, scratch_prefix, world, shape, caller):
async def _seed_target(gateway_client, seeder, scratch_prefix, actors, shape, caller):
if shape == "self":
return await create_scratch_key(
proxy_client, seeder, scratch_prefix, user_id=caller.user_id
gateway_client, seeder, scratch_prefix, user_id=caller.user_id
)
if shape == "owner":
return await create_scratch_key(
proxy_client,
gateway_client,
seeder,
scratch_prefix,
user_id=world.keys[Actor.OWNER].user_id,
team_id=TEAM_ALPHA,
user_id=actors.keys[Actor.OWNER].user_id,
team_id=actors.team_alpha_id,
)
if shape == "cross_org":
return await create_scratch_key(
proxy_client,
gateway_client,
seeder,
scratch_prefix,
user_id=world.keys[Actor.CROSS_ORG_USER].user_id,
team_id=TEAM_BETA,
user_id=actors.keys[Actor.CROSS_ORG_USER].user_id,
team_id=actors.team_beta_id,
)
pytest.fail(f"unknown shape={shape}") # pragma: no cover
@ -75,25 +75,25 @@ async def test_key_block_unblock_authz_matrix(
actor: Actor,
shape: str,
expected_status: int,
proxy_client,
prisma,
scratch,
world,
gateway_client,
db,
resources,
actors,
):
caller = world.keys[actor]
seeder = world.keys[Actor.PROXY_ADMIN].cleartext
caller = actors.keys[actor]
seeder = actors.keys[Actor.PROXY_ADMIN].cleartext
target_cleartext = await _seed_target(
proxy_client, seeder, scratch.prefix, world, shape, caller
gateway_client, seeder, resources.prefix, actors, shape, caller
)
target_hashed = hash_token(target_cleartext)
# /unblock starts from a blocked row so a 200 is observable as True->False.
if route == "unblock":
await prisma.db.litellm_verificationtoken.update(
await db.db.litellm_verificationtoken.update(
where={"token": target_hashed}, data={"blocked": True}
)
resp = await proxy_client.post(
resp = await gateway_client.post(
f"/key/{route}",
headers={"Authorization": f"Bearer {caller.cleartext}"},
json={"key": target_cleartext},
@ -102,7 +102,7 @@ async def test_key_block_unblock_authz_matrix(
resp.status_code == expected_status
), f"{route} {actor.value} {shape}: {resp.status_code} {resp.text}"
row = await prisma.db.litellm_verificationtoken.find_unique(
row = await db.db.litellm_verificationtoken.find_unique(
where={"token": target_hashed}
)
assert row is not None
@ -114,27 +114,27 @@ async def test_key_block_unblock_authz_matrix(
assert bool(row.blocked) is (route == "unblock"), "denied but blocked mutated"
async def test_key_block_unblock_round_trip(proxy_client, prisma, scratch, world):
async def test_key_block_unblock_round_trip(gateway_client, db, resources, actors):
"""PROXY_ADMIN block then unblock flips the blocked column True then False."""
admin = world.keys[Actor.PROXY_ADMIN]
admin = actors.keys[Actor.PROXY_ADMIN]
target = await create_scratch_key(
proxy_client, admin.cleartext, scratch.prefix, user_id=admin.user_id
gateway_client, admin.cleartext, resources.prefix, user_id=admin.user_id
)
hashed = hash_token(target)
headers = {"Authorization": f"Bearer {admin.cleartext}"}
blocked = await proxy_client.post(
blocked = await gateway_client.post(
"/key/block", headers=headers, json={"key": target}
)
assert blocked.status_code == 200, blocked.text
row = await prisma.db.litellm_verificationtoken.find_unique(where={"token": hashed})
row = await db.db.litellm_verificationtoken.find_unique(where={"token": hashed})
assert row is not None and row.blocked is True
unblocked = await proxy_client.post(
unblocked = await gateway_client.post(
"/key/unblock", headers=headers, json={"key": target}
)
assert unblocked.status_code == 200, unblocked.text
row = await prisma.db.litellm_verificationtoken.find_unique(where={"token": hashed})
row = await db.db.litellm_verificationtoken.find_unique(where={"token": hashed})
assert row is not None and row.blocked is False
@ -143,13 +143,13 @@ async def test_key_block_unblock_round_trip(proxy_client, prisma, scratch, world
"actor", [Actor.PROXY_ADMIN, Actor.TEAM_ADMIN], ids=["proxy_admin", "team_admin"]
)
async def test_key_block_unblock_missing_key_returns_404(
route: str, actor: Actor, proxy_client, world
route: str, actor: Actor, gateway_client, actors
):
"""A well-formed but unseeded key is 404 — not 401/403 — for both the
PROXY_ADMIN existence check and the non-admin _check_key_admin_access path."""
caller = world.keys[actor]
caller = actors.keys[actor]
missing = "sk-" + uuid.uuid4().hex
resp = await proxy_client.post(
resp = await gateway_client.post(
f"/key/{route}",
headers={"Authorization": f"Bearer {caller.cleartext}"},
json={"key": missing},

88
tests/proxy_behavior/management/test_key_budget_limits.py → tests/integration/management/test_key_budget_limits.py
View file

@ -26,8 +26,8 @@ from typing import Any, Dict
import pytest
from .actors import Actor
from .conftest import create_scratch_org, create_scratch_team
from tests.integration._support.actors import Actor
from tests.integration._support.actors import create_scratch_org, create_scratch_team
pytestmark = pytest.mark.asyncio(loop_scope="session")
@ -93,24 +93,24 @@ async def test_check_team_key_limits_aggregate(
body_extras: Dict[str, Any],
expected_status: int,
detail_substring,
proxy_client,
prisma,
scratch,
world,
gateway_client,
db,
resources,
actors,
):
team_id = await create_scratch_team(
prisma,
team_id=scratch.tag("team"),
db,
team_id=resources.tag("team"),
tpm_limit=team_tpm,
rpm_limit=team_rpm,
)
seeder = world.keys[Actor.PROXY_ADMIN].cleartext
seeder = actors.keys[Actor.PROXY_ADMIN].cleartext
body: Dict[str, Any] = {
"key_alias": scratch.prefix,
"key_alias": resources.prefix,
"team_id": team_id,
**body_extras,
}
resp = await proxy_client.post(
resp = await gateway_client.post(
"/key/generate",
headers={"Authorization": f"Bearer {seeder}"},
json=body,
@ -121,8 +121,8 @@ async def test_check_team_key_limits_aggregate(
if detail_substring is not None:
assert detail_substring in resp.text, resp.text
rows = await prisma.db.litellm_verificationtoken.find_many(
where={"key_alias": scratch.prefix}
rows = await db.db.litellm_verificationtoken.find_many(
where={"key_alias": resources.prefix}
)
if expected_status == 200:
assert len(rows) == 1
@ -178,23 +178,23 @@ async def test_check_team_key_limits_model_specific(
body_extras: Dict[str, Any],
expected_status: int,
detail_substring,
proxy_client,
prisma,
scratch,
world,
gateway_client,
db,
resources,
actors,
):
team_id = await create_scratch_team(
prisma,
team_id=scratch.tag("team"),
db,
team_id=resources.tag("team"),
metadata=team_metadata,
)
seeder = world.keys[Actor.PROXY_ADMIN].cleartext
seeder = actors.keys[Actor.PROXY_ADMIN].cleartext
body: Dict[str, Any] = {
"key_alias": scratch.prefix,
"key_alias": resources.prefix,
"team_id": team_id,
**body_extras,
}
resp = await proxy_client.post(
resp = await gateway_client.post(
"/key/generate",
headers={"Authorization": f"Bearer {seeder}"},
json=body,
@ -204,8 +204,8 @@ async def test_check_team_key_limits_model_specific(
), f"{body!r} → {resp.status_code}: {resp.text}"
if detail_substring is not None:
assert detail_substring in resp.text, resp.text
rows = await prisma.db.litellm_verificationtoken.find_many(
where={"key_alias": scratch.prefix}
rows = await db.db.litellm_verificationtoken.find_many(
where={"key_alias": resources.prefix}
)
assert len(rows) == (1 if expected_status == 200 else 0)
@ -273,19 +273,19 @@ async def test_check_org_key_limits_aggregate(
body_extras: Dict[str, Any],
expected_status: int,
detail_substring,
proxy_client,
prisma,
scratch,
world,
gateway_client,
db,
resources,
actors,
):
org_id = await create_scratch_org(prisma, scratch.prefix, **org_budget)
seeder = world.keys[Actor.PROXY_ADMIN].cleartext
org_id = await create_scratch_org(db, resources.prefix, **org_budget)
seeder = actors.keys[Actor.PROXY_ADMIN].cleartext
body: Dict[str, Any] = {
"key_alias": scratch.prefix,
"key_alias": resources.prefix,
"organization_id": org_id,
**body_extras,
}
resp = await proxy_client.post(
resp = await gateway_client.post(
"/key/generate",
headers={"Authorization": f"Bearer {seeder}"},
json=body,
@ -295,8 +295,8 @@ async def test_check_org_key_limits_aggregate(
), f"{body!r} → {resp.status_code}: {resp.text}"
if detail_substring is not None:
assert detail_substring in resp.text, resp.text
rows = await prisma.db.litellm_verificationtoken.find_many(
where={"key_alias": scratch.prefix}
rows = await db.db.litellm_verificationtoken.find_many(
where={"key_alias": resources.prefix}
)
assert len(rows) == (1 if expected_status == 200 else 0)
@ -347,19 +347,19 @@ async def test_check_org_key_limits_model_specific(
body_extras: Dict[str, Any],
expected_status: int,
detail_substring,
proxy_client,
prisma,
scratch,
world,
gateway_client,
db,
resources,
actors,
):
org_id = await create_scratch_org(prisma, scratch.prefix, metadata=org_metadata)
seeder = world.keys[Actor.PROXY_ADMIN].cleartext
org_id = await create_scratch_org(db, resources.prefix, metadata=org_metadata)
seeder = actors.keys[Actor.PROXY_ADMIN].cleartext
body: Dict[str, Any] = {
"key_alias": scratch.prefix,
"key_alias": resources.prefix,
"organization_id": org_id,
**body_extras,
}
resp = await proxy_client.post(
resp = await gateway_client.post(
"/key/generate",
headers={"Authorization": f"Bearer {seeder}"},
json=body,
@ -369,8 +369,8 @@ async def test_check_org_key_limits_model_specific(
), f"{body!r} → {resp.status_code}: {resp.text}"
if detail_substring is not None:
assert detail_substring in resp.text, resp.text
rows = await prisma.db.litellm_verificationtoken.find_many(
where={"key_alias": scratch.prefix}
rows = await db.db.litellm_verificationtoken.find_many(
where={"key_alias": resources.prefix}
)
assert len(rows) == (1 if expected_status == 200 else 0)

44
tests/proxy_behavior/management/test_key_bulk_update.py → tests/integration/management/test_key_bulk_update.py
View file

@ -4,8 +4,8 @@ import pytest
from litellm.proxy.utils import hash_token
from .actors import Actor
from .conftest import create_scratch_key
from tests.integration._support.actors import Actor
from tests.integration._support.actors import create_scratch_key
pytestmark = pytest.mark.asyncio(loop_scope="session")
@ -37,19 +37,19 @@ _MATRIX = [
async def test_key_bulk_update_authz_matrix(
actor: Actor,
expected_status: int,
proxy_client,
prisma,
scratch,
world,
gateway_client,
db,
resources,
actors,
):
caller = world.keys[actor]
seeder = world.keys[Actor.PROXY_ADMIN].cleartext
caller = actors.keys[actor]
seeder = actors.keys[Actor.PROXY_ADMIN].cleartext
target = await create_scratch_key(
proxy_client, seeder, scratch.prefix, user_id=caller.user_id
gateway_client, seeder, resources.prefix, user_id=caller.user_id
)
hashed = hash_token(target)
resp = await proxy_client.post(
resp = await gateway_client.post(
"/key/bulk_update",
headers={"Authorization": f"Bearer {caller.cleartext}"},
json={"keys": [{"key": target, "max_budget": _MARKER_BUDGET}]},
@ -58,7 +58,7 @@ async def test_key_bulk_update_authz_matrix(
resp.status_code == expected_status
), f"{actor.value}: {resp.status_code} {resp.text}"
row = await prisma.db.litellm_verificationtoken.find_unique(where={"token": hashed})
row = await db.db.litellm_verificationtoken.find_unique(where={"token": hashed})
assert row is not None
if expected_status == 200:
body = resp.json()
@ -69,39 +69,39 @@ async def test_key_bulk_update_authz_matrix(
assert row.max_budget != _MARKER_BUDGET, "denied but key mutated"
async def test_key_bulk_update_empty_keys_is_400(proxy_client, world):
async def test_key_bulk_update_empty_keys_is_400(gateway_client, actors):
"""An empty batch is rejected 400 before any per-key processing."""
resp = await proxy_client.post(
resp = await gateway_client.post(
"/key/bulk_update",
headers={"Authorization": f"Bearer {world.keys[Actor.PROXY_ADMIN].cleartext}"},
headers={"Authorization": f"Bearer {actors.keys[Actor.PROXY_ADMIN].cleartext}"},
json={"keys": []},
)
assert resp.status_code == 400, resp.text
async def test_key_bulk_update_over_max_batch_is_400(proxy_client, world):
async def test_key_bulk_update_over_max_batch_is_400(gateway_client, actors):
"""A batch larger than the 500-key cap is rejected 400."""
items = [{"key": "sk-" + uuid.uuid4().hex} for _ in range(501)]
resp = await proxy_client.post(
resp = await gateway_client.post(
"/key/bulk_update",
headers={"Authorization": f"Bearer {world.keys[Actor.PROXY_ADMIN].cleartext}"},
headers={"Authorization": f"Bearer {actors.keys[Actor.PROXY_ADMIN].cleartext}"},
json={"keys": items},
)
assert resp.status_code == 400, resp.text
async def test_key_bulk_update_per_key_failure_is_isolated(
proxy_client, prisma, scratch, world
gateway_client, db, resources, actors
):
"""One bad key in the batch does not abort the others — it lands in
failed_updates while the valid key is still updated."""
admin = world.keys[Actor.PROXY_ADMIN]
admin = actors.keys[Actor.PROXY_ADMIN]
valid = await create_scratch_key(
proxy_client, admin.cleartext, scratch.prefix, user_id=admin.user_id
gateway_client, admin.cleartext, resources.prefix, user_id=admin.user_id
)
missing = "sk-" + uuid.uuid4().hex
resp = await proxy_client.post(
resp = await gateway_client.post(
"/key/bulk_update",
headers={"Authorization": f"Bearer {admin.cleartext}"},
json={
@ -117,7 +117,7 @@ async def test_key_bulk_update_per_key_failure_is_isolated(
assert len(body["successful_updates"]) == 1
assert len(body["failed_updates"]) == 1
row = await prisma.db.litellm_verificationtoken.find_unique(
row = await db.db.litellm_verificationtoken.find_unique(
where={"token": hash_token(valid)}
)
assert row is not None and row.max_budget == _MARKER_BUDGET

46
tests/proxy_behavior/management/test_key_delete.py → tests/integration/management/test_key_delete.py
View file

@ -4,8 +4,8 @@ import pytest
from litellm.proxy.utils import hash_token
from .actors import TEAM_ALPHA, TEAM_BETA, Actor
from .conftest import create_scratch_key
from tests.integration._support.actors import TEAM_ALPHA, TEAM_BETA, Actor
from tests.integration._support.actors import create_scratch_key
pytestmark = pytest.mark.asyncio(loop_scope="session")
@ -46,40 +46,40 @@ async def test_key_delete_authz_matrix(
actor: Actor,
target_shape: str,
expected_status: int,
proxy_client,
prisma,
scratch,
world,
gateway_client,
db,
resources,
actors,
):
caller = world.keys[actor]
seeder = world.keys[Actor.PROXY_ADMIN].cleartext
caller = actors.keys[actor]
seeder = actors.keys[Actor.PROXY_ADMIN].cleartext
if target_shape == "self":
target_cleartext = await create_scratch_key(
proxy_client, seeder, scratch.prefix, user_id=caller.user_id
gateway_client, seeder, resources.prefix, user_id=caller.user_id
)
elif target_shape == "owner":
target_cleartext = await create_scratch_key(
proxy_client,
gateway_client,
seeder,
scratch.prefix,
user_id=world.keys[Actor.OWNER].user_id,
team_id=TEAM_ALPHA,
resources.prefix,
user_id=actors.keys[Actor.OWNER].user_id,
team_id=actors.team_alpha_id,
)
elif target_shape == "cross_org":
target_cleartext = await create_scratch_key(
proxy_client,
gateway_client,
seeder,
scratch.prefix,
user_id=world.keys[Actor.CROSS_ORG_USER].user_id,
team_id=TEAM_BETA,
resources.prefix,
user_id=actors.keys[Actor.CROSS_ORG_USER].user_id,
team_id=actors.team_beta_id,
)
else:
pytest.fail(f"unknown target_shape={target_shape}")
target_hashed = hash_token(target_cleartext)
resp = await proxy_client.post(
resp = await gateway_client.post(
"/key/delete",
headers={"Authorization": f"Bearer {caller.cleartext}"},
json={"keys": [target_cleartext]},
@ -88,10 +88,10 @@ async def test_key_delete_authz_matrix(
resp.status_code == expected_status
), f"{actor.value} {target_shape}: {resp.status_code} {resp.text}"
row = await prisma.db.litellm_verificationtoken.find_unique(
row = await db.db.litellm_verificationtoken.find_unique(
where={"token": target_hashed}
)
auth_check = await proxy_client.get(
auth_check = await gateway_client.get(
"/key/info", headers={"Authorization": f"Bearer {target_cleartext}"}
)
@ -103,11 +103,11 @@ async def test_key_delete_authz_matrix(
assert auth_check.status_code == 200
async def test_key_delete_missing_key_is_404(proxy_client, world):
async def test_key_delete_missing_key_is_404(gateway_client, actors):
"""Deleting a key absent from the DB is a 404 — not 401/403."""
resp = await proxy_client.post(
resp = await gateway_client.post(
"/key/delete",
headers={"Authorization": f"Bearer {world.keys[Actor.PROXY_ADMIN].cleartext}"},
headers={"Authorization": f"Bearer {actors.keys[Actor.PROXY_ADMIN].cleartext}"},
json={"keys": ["sk-" + uuid.uuid4().hex]},
)
assert resp.status_code == 404, resp.text

20
tests/proxy_behavior/management/test_key_generate.py → tests/integration/management/test_key_generate.py
View file

@ -2,7 +2,7 @@ from typing import Any, Dict
import pytest
from .actors import TEAM_ALPHA, TEAM_BETA, Actor
from tests.integration._support.actors import TEAM_ALPHA, TEAM_BETA, Actor
pytestmark = pytest.mark.asyncio(loop_scope="session")
@ -42,15 +42,15 @@ async def test_key_generate_authz_matrix(
actor: Actor,
body_extras: Dict[str, Any],
expected_status: int,
proxy_client,
prisma,
scratch,
world,
gateway_client,
db,
resources,
actors,
):
seeded = world.keys[actor]
body: Dict[str, Any] = {"key_alias": scratch.prefix, **body_extras}
seeded = actors.keys[actor]
body: Dict[str, Any] = {"key_alias": resources.prefix, **body_extras}
resp = await proxy_client.post(
resp = await gateway_client.post(
"/key/generate",
headers={"Authorization": f"Bearer {seeded.cleartext}"},
json=body,
@ -59,8 +59,8 @@ async def test_key_generate_authz_matrix(
resp.status_code == expected_status
), f"{actor.value} {body!r} → {resp.status_code}: {resp.text}"
rows = await prisma.db.litellm_verificationtoken.find_many(
where={"key_alias": scratch.prefix}
rows = await db.db.litellm_verificationtoken.find_many(
where={"key_alias": resources.prefix}
)
if expected_status == 200:
cleartext = resp.json()["key"]

14
tests/proxy_behavior/management/test_key_health.py → tests/integration/management/test_key_health.py
View file

@ -1,17 +1,17 @@
import pytest
from .actors import Actor
from tests.integration._support.actors import Actor
pytestmark = pytest.mark.asyncio(loop_scope="session")
# POST /key/health has no role gate — it reflects the caller's OWN key logging
# metadata. The world keys carry no "logging" metadata, so every authenticated
# metadata. The actors keys carry no "logging" metadata, so every authenticated
# actor gets 200 with key="healthy". This pins auth-required + route coverage.
@pytest.mark.parametrize("actor", list(Actor), ids=[a.value for a in Actor])
async def test_key_health_each_actor_is_healthy(actor: Actor, proxy_client, world):
caller = world.keys[actor]
resp = await proxy_client.post(
async def test_key_health_each_actor_is_healthy(actor: Actor, gateway_client, actors):
caller = actors.keys[actor]
resp = await gateway_client.post(
"/key/health",
headers={"Authorization": f"Bearer {caller.cleartext}"},
)
@ -19,6 +19,6 @@ async def test_key_health_each_actor_is_healthy(actor: Actor, proxy_client, worl
assert resp.json()["key"] == "healthy"
async def test_key_health_requires_auth(proxy_client):
resp = await proxy_client.post("/key/health")
async def test_key_health_requires_auth(gateway_client):
resp = await gateway_client.post("/key/health", headers={"Authorization": "Bearer invalid"})
assert resp.status_code == 401, resp.text

10
tests/proxy_behavior/management/test_key_info.py → tests/integration/management/test_key_info.py
View file

@ -1,6 +1,6 @@
import pytest
from .actors import Actor
from tests.integration._support.actors import Actor
pytestmark = pytest.mark.asyncio(loop_scope="session")
@ -53,12 +53,12 @@ _SCENARIOS = [
ids=[s[0] for s in _SCENARIOS],
)
async def test_key_info_authz_matrix(
actor: Actor, target_actor: Actor, expected_status: int, proxy_client, world
actor: Actor, target_actor: Actor, expected_status: int, gateway_client, actors
):
caller = world.keys[actor]
target = world.keys[target_actor]
caller = actors.keys[actor]
target = actors.keys[target_actor]
resp = await proxy_client.get(
resp = await gateway_client.get(
f"/key/info?key={target.cleartext}",
headers={"Authorization": f"Bearer {caller.cleartext}"},
)

28
tests/proxy_behavior/management/test_key_info_v2.py → tests/integration/management/test_key_info_v2.py
View file

@ -2,16 +2,16 @@ import uuid
import pytest
from .actors import Actor
from tests.integration._support.actors import Actor
pytestmark = pytest.mark.asyncio(loop_scope="session")
# POST /v2/key/info resolves the posted keys, then drops any key the caller
# cannot see via _can_user_query_key_info — silently, no 403. A non-admin sees
# a key it owns (user_id match) or a key whose team it belongs to. The world's
# a key it owns (user_id match) or a key whose team it belongs to. The actors's
# TEAM_ALPHA members all see each other's keys; CROSS_ORG_USER and the org
# admins see only their own. The request is posted with every world key, and
# admins see only their own. The request is posted with every actors key, and
# the returned info set is asserted to equal the visible subset.
_ALPHA_KEYS = frozenset(
{
@ -40,14 +40,14 @@ _VISIBILITY = {
list(_VISIBILITY.items()),
ids=[a.value for a in _VISIBILITY],
)
async def test_key_info_v2_visibility(actor, expected_visible, proxy_client, world):
caller = world.keys[actor]
user_id_to_actor = {world.keys[a].user_id: a for a in Actor}
async def test_key_info_v2_visibility(actor, expected_visible, gateway_client, actors):
caller = actors.keys[actor]
user_id_to_actor = {actors.keys[a].user_id: a for a in Actor}
resp = await proxy_client.post(
resp = await gateway_client.post(
"/v2/key/info",
headers={"Authorization": f"Bearer {caller.cleartext}"},
json={"keys": [world.keys[a].cleartext for a in Actor]},
json={"keys": [actors.keys[a].cleartext for a in Actor]},
)
assert resp.status_code == 200, f"{actor.value}: {resp.status_code} {resp.text}"
@ -62,20 +62,20 @@ async def test_key_info_v2_visibility(actor, expected_visible, proxy_client, wor
)
async def test_key_info_v2_no_body_is_422(proxy_client, world):
async def test_key_info_v2_no_body_is_422(gateway_client, actors):
"""A request with no body is a 422 — the handler has no keys to resolve."""
resp = await proxy_client.post(
resp = await gateway_client.post(
"/v2/key/info",
headers={"Authorization": f"Bearer {world.keys[Actor.PROXY_ADMIN].cleartext}"},
headers={"Authorization": f"Bearer {actors.keys[Actor.PROXY_ADMIN].cleartext}"},
)
assert resp.status_code == 422, resp.text
async def test_key_info_v2_unknown_key_returns_empty_info(proxy_client, world):
async def test_key_info_v2_unknown_key_returns_empty_info(gateway_client, actors):
"""Keys that resolve to no rows yield an empty info list, not an error."""
resp = await proxy_client.post(
resp = await gateway_client.post(
"/v2/key/info",
headers={"Authorization": f"Bearer {world.keys[Actor.PROXY_ADMIN].cleartext}"},
headers={"Authorization": f"Bearer {actors.keys[Actor.PROXY_ADMIN].cleartext}"},
json={"keys": ["sk-" + uuid.uuid4().hex]},
)
assert resp.status_code == 200, resp.text

104
tests/proxy_behavior/management/test_key_list.py → tests/integration/management/test_key_list.py
View file

@ -4,8 +4,8 @@ import pytest
from litellm.proxy.utils import hash_token
from .actors import TEAM_ALPHA, Actor
from .conftest import create_scratch_key
from tests.integration._support.actors import TEAM_ALPHA, Actor
from tests.integration._support.actors import create_scratch_key
pytestmark = pytest.mark.asyncio(loop_scope="session")
@ -24,13 +24,13 @@ _VISIBILITY = {
}
async def _all_visible_hashes(proxy_client, caller_cleartext) -> set:
async def _all_visible_hashes(gateway_client, caller_cleartext) -> set:
"""Walk every /key/list page — size is capped at 100 by the endpoint, so a
single request can truncate PROXY_ADMIN's view on a non-fresh DB."""
hashes: set = set()
page = 1
while True:
resp = await proxy_client.get(
resp = await gateway_client.get(
f"/key/list?page={page}&size=100",
headers={"Authorization": f"Bearer {caller_cleartext}"},
)
@ -51,12 +51,12 @@ async def _all_visible_hashes(proxy_client, caller_cleartext) -> set:
ids=[a.value for a in _VISIBILITY],
)
async def test_key_list_visibility(
actor: Actor, expected_visible: FrozenSet[Actor], proxy_client, world
actor: Actor, expected_visible: FrozenSet[Actor], gateway_client, actors
):
caller = world.keys[actor]
hashed_to_actor = {world.keys[a].hashed: a for a in Actor}
caller = actors.keys[actor]
hashed_to_actor = {actors.keys[a].hashed: a for a in Actor}
returned_hashes = await _all_visible_hashes(proxy_client, caller.cleartext)
returned_hashes = await _all_visible_hashes(gateway_client, caller.cleartext)
visible_seeded = {
hashed_to_actor[h] for h in returned_hashes if h in hashed_to_actor
}
@ -66,8 +66,8 @@ async def test_key_list_visibility(
)
async def _list_hashes(proxy_client, caller_cleartext: str, query: str) -> set:
resp = await proxy_client.get(
async def _list_hashes(gateway_client, caller_cleartext: str, query: str) -> set:
resp = await gateway_client.get(
f"/key/list?{query}&size=100",
headers={"Authorization": f"Bearer {caller_cleartext}"},
)
@ -80,125 +80,125 @@ async def _list_hashes(proxy_client, caller_cleartext: str, query: str) -> set:
return hashes
async def test_key_list_admin_key_alias_substring_match(proxy_client, scratch, world):
async def test_key_list_admin_key_alias_substring_match(gateway_client, resources, actors):
"""A PROXY_ADMIN's key_alias filter is a case-insensitive substring match
when substring_matching=true is requested (the dashboard search box); a
narrower fragment selects the subset whose alias contains it. Substring
matching is opt-in: without the flag the filter is exact (see
test_key_list_admin_key_alias_exact_without_substring_flag)."""
admin = world.keys[Actor.PROXY_ADMIN]
admin = actors.keys[Actor.PROXY_ADMIN]
a = await create_scratch_key(
proxy_client,
gateway_client,
admin.cleartext,
scratch.prefix,
resources.prefix,
user_id=admin.user_id,
key_alias=f"{scratch.prefix}-sub-a",
key_alias=f"{resources.prefix}-sub-a",
)
b = await create_scratch_key(
proxy_client,
gateway_client,
admin.cleartext,
scratch.prefix,
resources.prefix,
user_id=admin.user_id,
key_alias=f"{scratch.prefix}-sub-b",
key_alias=f"{resources.prefix}-sub-b",
)
seeded = {hash_token(a), hash_token(b)}
broad = await _list_hashes(
proxy_client,
gateway_client,
admin.cleartext,
f"key_alias={scratch.prefix}-sub&substring_matching=true",
f"key_alias={resources.prefix}-sub&substring_matching=true",
)
assert broad & seeded == seeded
narrow = await _list_hashes(
proxy_client,
gateway_client,
admin.cleartext,
f"key_alias={scratch.prefix}-sub-a&substring_matching=true",
f"key_alias={resources.prefix}-sub-a&substring_matching=true",
)
assert narrow & seeded == {hash_token(a)}
async def test_key_list_admin_key_alias_exact_without_substring_flag(
proxy_client, scratch, world
gateway_client, resources, actors
):
"""Regression guard for the prior exact-match contract: without
substring_matching, even a PROXY_ADMIN's key_alias filter is exact, so a
fragment of a seeded alias does not select it."""
admin = world.keys[Actor.PROXY_ADMIN]
full_alias = f"{scratch.prefix}-exactflag"
admin = actors.keys[Actor.PROXY_ADMIN]
full_alias = f"{resources.prefix}-exactflag"
key = await create_scratch_key(
proxy_client,
gateway_client,
admin.cleartext,
scratch.prefix,
resources.prefix,
user_id=admin.user_id,
key_alias=full_alias,
)
key_hash = hash_token(key)
exact = await _list_hashes(proxy_client, admin.cleartext, f"key_alias={full_alias}")
exact = await _list_hashes(gateway_client, admin.cleartext, f"key_alias={full_alias}")
assert key_hash in exact
fragment = await _list_hashes(
proxy_client, admin.cleartext, f"key_alias={scratch.prefix}-exactfla"
gateway_client, admin.cleartext, f"key_alias={resources.prefix}-exactfla"
)
assert key_hash not in fragment
async def test_key_list_non_admin_key_alias_is_exact_match(
proxy_client, scratch, world
gateway_client, resources, actors
):
"""A non-admin's key_alias filter is exact-match only — substring filtering
is restricted to admins. The full alias matches; a fragment does not."""
caller = world.keys[Actor.INTERNAL_USER]
alias = f"{scratch.prefix}-exact"
caller = actors.keys[Actor.INTERNAL_USER]
alias = f"{resources.prefix}-exact"
key = await create_scratch_key(
proxy_client,
world.keys[Actor.PROXY_ADMIN].cleartext,
scratch.prefix,
gateway_client,
actors.keys[Actor.PROXY_ADMIN].cleartext,
resources.prefix,
user_id=caller.user_id,
key_alias=alias,
)
key_hash = hash_token(key)
exact = await _list_hashes(proxy_client, caller.cleartext, f"key_alias={alias}")
exact = await _list_hashes(gateway_client, caller.cleartext, f"key_alias={alias}")
assert key_hash in exact
fragment = await _list_hashes(
proxy_client, caller.cleartext, f"key_alias={scratch.prefix}-exac"
gateway_client, caller.cleartext, f"key_alias={resources.prefix}-exac"
)
assert key_hash not in fragment
async def test_key_list_team_id_filter(proxy_client, scratch, world):
async def test_key_list_team_id_filter(gateway_client, resources, actors):
"""A team_id filter narrows the listing to keys of that team."""
admin = world.keys[Actor.PROXY_ADMIN]
admin = actors.keys[Actor.PROXY_ADMIN]
team_key = await create_scratch_key(
proxy_client,
gateway_client,
admin.cleartext,
scratch.prefix,
user_id=world.keys[Actor.OWNER].user_id,
team_id=TEAM_ALPHA,
key_alias=f"{scratch.prefix}-team",
resources.prefix,
user_id=actors.keys[Actor.OWNER].user_id,
team_id=actors.team_alpha_id,
key_alias=f"{resources.prefix}-team",
)
no_team_key = await create_scratch_key(
proxy_client,
gateway_client,
admin.cleartext,
scratch.prefix,
resources.prefix,
user_id=admin.user_id,
key_alias=f"{scratch.prefix}-noteam",
key_alias=f"{resources.prefix}-noteam",
)
hashes = await _list_hashes(proxy_client, admin.cleartext, f"team_id={TEAM_ALPHA}")
hashes = await _list_hashes(gateway_client, admin.cleartext, f"team_id={actors.team_alpha_id}")
assert hash_token(team_key) in hashes
assert hash_token(no_team_key) not in hashes
async def test_key_list_non_admin_cannot_filter_other_team(proxy_client, world):
async def test_key_list_non_admin_cannot_filter_other_team(gateway_client, actors):
"""A non-admin filtering by a team it does not belong to is rejected 403."""
resp = await proxy_client.get(
f"/key/list?team_id={world.team_beta_id}",
resp = await gateway_client.get(
f"/key/list?team_id={actors.team_beta_id}",
headers={
"Authorization": f"Bearer {world.keys[Actor.INTERNAL_USER].cleartext}"
"Authorization": f"Bearer {actors.keys[Actor.INTERNAL_USER].cleartext}"
},
)
assert resp.status_code == 403, resp.text

107
tests/proxy_behavior/management/test_key_regenerate.py → tests/integration/management/test_key_regenerate.py
View file

@ -1,12 +1,7 @@
import litellm
import pytest
from litellm.types.proxy.management_endpoints.ui_sso import (
LiteLLM_UpperboundKeyGenerateParams,
)
from .actors import TEAM_ALPHA, TEAM_BETA, Actor
from .conftest import create_scratch_key
from tests.integration._support.actors import TEAM_ALPHA, TEAM_BETA, Actor
from tests.integration._support.actors import create_scratch_key, owned_settings_client
pytestmark = pytest.mark.asyncio(loop_scope="session")
@ -39,8 +34,8 @@ _SCENARIOS = [
]
async def _info(proxy_client, cleartext: str):
return await proxy_client.get(
async def _info(gateway_client, cleartext: str):
return await gateway_client.get(
"/key/info", headers={"Authorization": f"Bearer {cleartext}"}
)
@ -54,37 +49,37 @@ async def test_key_regenerate_authz_matrix(
actor: Actor,
target_shape: str,
expected_status: int,
proxy_client,
scratch,
world,
gateway_client,
resources,
actors,
):
caller = world.keys[actor]
seeder = world.keys[Actor.PROXY_ADMIN].cleartext
caller = actors.keys[actor]
seeder = actors.keys[Actor.PROXY_ADMIN].cleartext
if target_shape == "self":
target_cleartext = await create_scratch_key(
proxy_client, seeder, scratch.prefix, user_id=caller.user_id
gateway_client, seeder, resources.prefix, user_id=caller.user_id
)
elif target_shape == "owner":
target_cleartext = await create_scratch_key(
proxy_client,
gateway_client,
seeder,
scratch.prefix,
user_id=world.keys[Actor.OWNER].user_id,
team_id=TEAM_ALPHA,
resources.prefix,
user_id=actors.keys[Actor.OWNER].user_id,
team_id=actors.team_alpha_id,
)
elif target_shape == "cross_org":
target_cleartext = await create_scratch_key(
proxy_client,
gateway_client,
seeder,
scratch.prefix,
user_id=world.keys[Actor.CROSS_ORG_USER].user_id,
team_id=TEAM_BETA,
resources.prefix,
user_id=actors.keys[Actor.CROSS_ORG_USER].user_id,
team_id=actors.team_beta_id,
)
else:
pytest.fail(f"unknown target_shape={target_shape}")
resp = await proxy_client.post(
resp = await gateway_client.post(
"/key/regenerate",
headers={"Authorization": f"Bearer {caller.cleartext}"},
json={"key": target_cleartext},
@ -96,21 +91,21 @@ async def test_key_regenerate_authz_matrix(
if expected_status == 200:
new_cleartext = resp.json()["key"]
assert new_cleartext.startswith("sk-") and new_cleartext != target_cleartext
assert (await _info(proxy_client, target_cleartext)).status_code == 401
assert (await _info(proxy_client, new_cleartext)).status_code == 200
assert (await _info(gateway_client, target_cleartext)).status_code == 401
assert (await _info(gateway_client, new_cleartext)).status_code == 200
else:
# Denied: rotation must not have leaked — old cleartext still works.
assert (await _info(proxy_client, target_cleartext)).status_code == 200
assert (await _info(gateway_client, target_cleartext)).status_code == 200
async def test_key_path_regenerate_smoke(proxy_client, scratch, world):
async def test_key_path_regenerate_smoke(gateway_client, resources, actors):
"""Pins that POST /key/{key:path}/regenerate shares the same handler."""
caller = world.keys[Actor.PROXY_ADMIN]
caller = actors.keys[Actor.PROXY_ADMIN]
target_cleartext = await create_scratch_key(
proxy_client, caller.cleartext, scratch.prefix, user_id=caller.user_id
gateway_client, caller.cleartext, resources.prefix, user_id=caller.user_id
)
resp = await proxy_client.post(
resp = await gateway_client.post(
f"/key/{target_cleartext}/regenerate",
headers={"Authorization": f"Bearer {caller.cleartext}"},
json={},
@ -118,48 +113,40 @@ async def test_key_path_regenerate_smoke(proxy_client, scratch, world):
assert resp.status_code == 200, resp.text
new_cleartext = resp.json()["key"]
assert new_cleartext.startswith("sk-") and new_cleartext != target_cleartext
assert (await _info(proxy_client, target_cleartext)).status_code == 401
assert (await _info(proxy_client, new_cleartext)).status_code == 200
assert (await _info(gateway_client, target_cleartext)).status_code == 401
assert (await _info(gateway_client, new_cleartext)).status_code == 200
async def test_key_regenerate_enforces_upperbound_key_params(
proxy_client, scratch, world, monkeypatch
gateway_client, resources, actors, tmp_path
):
"""Regenerate runs _enforce_upperbound_key_params: a max_budget above
litellm.upperbound_key_generate_params is rejected 400, a value within the
bound is accepted. Pins #26340 (db8ef44323) — regenerate previously
bypassed the upperbound. upperbound_key_generate_params is module-level
litellm.* state, so monkeypatch save/restores it."""
admin = world.keys[Actor.PROXY_ADMIN]
admin = actors.keys[Actor.PROXY_ADMIN]
over_key = await create_scratch_key(
proxy_client,
gateway_client,
admin.cleartext,
scratch.prefix,
resources.prefix,
user_id=admin.user_id,
key_alias=f"{scratch.prefix}-over",
key_alias=f"{resources.prefix}-over",
)
within_key = await create_scratch_key(
proxy_client,
gateway_client,
admin.cleartext,
scratch.prefix,
resources.prefix,
user_id=admin.user_id,
key_alias=f"{scratch.prefix}-within",
)
monkeypatch.setattr(
litellm,
"upperbound_key_generate_params",
LiteLLM_UpperboundKeyGenerateParams(max_budget=100.0),
key_alias=f"{resources.prefix}-within",
)
headers = {"Authorization": f"Bearer {admin.cleartext}"}
over = await proxy_client.post(
"/key/regenerate", headers=headers, json={"key": over_key, "max_budget": 500.0}
)
with owned_settings_client(
gateway_client, tmp_path, {"upperbound_key_generate_params": {"max_budget": 100.0}}
) as client:
over = await client.post(
"/key/regenerate", headers=headers, json={"key": over_key, "max_budget": 500.0}
)
within = await client.post(
"/key/regenerate",
headers=headers,
json={"key": within_key, "max_budget": 50.0},
)
assert over.status_code == 400, over.text
within = await proxy_client.post(
"/key/regenerate",
headers=headers,
json={"key": within_key, "max_budget": 50.0},
)
assert within.status_code == 200, within.text

54
tests/proxy_behavior/management/test_key_reset_spend.py → tests/integration/management/test_key_reset_spend.py
View file

@ -4,8 +4,8 @@ import pytest
from litellm.proxy.utils import hash_token
from .actors import TEAM_ALPHA, TEAM_BETA, Actor
from .conftest import create_scratch_key
from tests.integration._support.actors import TEAM_ALPHA, TEAM_BETA, Actor
from tests.integration._support.actors import create_scratch_key
pytestmark = pytest.mark.asyncio(loop_scope="session")
@ -42,26 +42,26 @@ _SCENARIOS = [
]
async def _seed_target(proxy_client, seeder, prefix, world, shape, caller) -> str:
async def _seed_target(gateway_client, seeder, prefix, actors, shape, caller) -> str:
if shape == "self":
return await create_scratch_key(
proxy_client, seeder, prefix, user_id=caller.user_id
gateway_client, seeder, prefix, user_id=caller.user_id
)
if shape == "team_alpha":
return await create_scratch_key(
proxy_client,
gateway_client,
seeder,
prefix,
user_id=world.keys[Actor.OWNER].user_id,
team_id=TEAM_ALPHA,
user_id=actors.keys[Actor.OWNER].user_id,
team_id=actors.team_alpha_id,
)
if shape == "team_beta":
return await create_scratch_key(
proxy_client,
gateway_client,
seeder,
prefix,
user_id=world.keys[Actor.CROSS_ORG_USER].user_id,
team_id=TEAM_BETA,
user_id=actors.keys[Actor.CROSS_ORG_USER].user_id,
team_id=actors.team_beta_id,
)
pytest.fail(f"unknown shape={shape}") # pragma: no cover
@ -75,22 +75,22 @@ async def test_key_reset_spend_authz_matrix(
actor: Actor,
shape: str,
expected_status: int,
proxy_client,
prisma,
scratch,
world,
gateway_client,
db,
resources,
actors,
):
caller = world.keys[actor]
seeder = world.keys[Actor.PROXY_ADMIN].cleartext
caller = actors.keys[actor]
seeder = actors.keys[Actor.PROXY_ADMIN].cleartext
target = await _seed_target(
proxy_client, seeder, scratch.prefix, world, shape, caller
gateway_client, seeder, resources.prefix, actors, shape, caller
)
hashed = hash_token(target)
await prisma.db.litellm_verificationtoken.update(
await db.db.litellm_verificationtoken.update(
where={"token": hashed}, data={"spend": _SEED_SPEND}
)
resp = await proxy_client.post(
resp = await gateway_client.post(
f"/key/{target}/reset_spend",
headers={"Authorization": f"Bearer {caller.cleartext}"},
json={"reset_to": _RESET_TO},
@ -99,7 +99,7 @@ async def test_key_reset_spend_authz_matrix(
resp.status_code == expected_status
), f"{actor.value} {shape}: {resp.status_code} {resp.text}"
row = await prisma.db.litellm_verificationtoken.find_unique(where={"token": hashed})
row = await db.db.litellm_verificationtoken.find_unique(where={"token": hashed})
assert row is not None
if expected_status == 200:
assert row.spend == _RESET_TO
@ -110,25 +110,25 @@ async def test_key_reset_spend_authz_matrix(
@pytest.mark.parametrize(
"actor", [Actor.PROXY_ADMIN, Actor.TEAM_ADMIN], ids=["proxy_admin", "team_admin"]
)
async def test_key_reset_spend_missing_key_is_404(actor: Actor, proxy_client, world):
async def test_key_reset_spend_missing_key_is_404(actor: Actor, gateway_client, actors):
"""A well-formed but unseeded key is 404 before any spend validation."""
resp = await proxy_client.post(
resp = await gateway_client.post(
f"/key/sk-{uuid.uuid4().hex}/reset_spend",
headers={"Authorization": f"Bearer {world.keys[actor].cleartext}"},
headers={"Authorization": f"Bearer {actors.keys[actor].cleartext}"},
json={"reset_to": 0.0},
)
assert resp.status_code == 404, resp.text
async def test_key_reset_spend_above_current_spend_is_400(
proxy_client, prisma, scratch, world
gateway_client, db, resources, actors
):
"""reset_to above the key's current spend is rejected 400."""
admin = world.keys[Actor.PROXY_ADMIN]
admin = actors.keys[Actor.PROXY_ADMIN]
target = await create_scratch_key(
proxy_client, admin.cleartext, scratch.prefix, user_id=admin.user_id
gateway_client, admin.cleartext, resources.prefix, user_id=admin.user_id
)
resp = await proxy_client.post(
resp = await gateway_client.post(
f"/key/{target}/reset_spend",
headers={"Authorization": f"Bearer {admin.cleartext}"},
json={"reset_to": 1.0},

36
tests/proxy_behavior/management/test_key_service_account_generate.py → tests/integration/management/test_key_service_account_generate.py
View file

@ -1,6 +1,6 @@
import pytest
from .actors import Actor
from tests.integration._support.actors import Actor
pytestmark = pytest.mark.asyncio(loop_scope="session")
@ -45,23 +45,23 @@ async def test_key_service_account_generate_authz_matrix(
actor: Actor,
team_target: str,
expected_status: int,
proxy_client,
prisma,
scratch,
world,
gateway_client,
db,
resources,
actors,
):
caller = world.keys[actor]
caller = actors.keys[actor]
team_id = {
"own": world.team_alpha_id,
"cross_org": world.team_beta_id,
"own": actors.team_alpha_id,
"cross_org": actors.team_beta_id,
"none": None,
}[team_target]
body = {"key_alias": scratch.prefix}
body = {"key_alias": resources.prefix}
if team_id is not None:
body["team_id"] = team_id
resp = await proxy_client.post(
resp = await gateway_client.post(
"/key/service-account/generate",
headers={"Authorization": f"Bearer {caller.cleartext}"},
json=body,
@ -70,8 +70,8 @@ async def test_key_service_account_generate_authz_matrix(
resp.status_code == expected_status
), f"{actor.value} {team_target}: {resp.status_code} {resp.text}"
rows = await prisma.db.litellm_verificationtoken.find_many(
where={"key_alias": scratch.prefix}
rows = await db.db.litellm_verificationtoken.find_many(
where={"key_alias": resources.prefix}
)
if expected_status == 200:
assert len(rows) == 1
@ -83,16 +83,16 @@ async def test_key_service_account_generate_authz_matrix(
async def test_key_service_account_generate_unknown_team_is_400(
proxy_client, prisma, scratch, world
gateway_client, db, resources, actors
):
"""A team_id absent from the database is rejected 400."""
resp = await proxy_client.post(
resp = await gateway_client.post(
"/key/service-account/generate",
headers={"Authorization": f"Bearer {world.keys[Actor.PROXY_ADMIN].cleartext}"},
json={"key_alias": scratch.prefix, "team_id": scratch.tag("no-such-team")},
headers={"Authorization": f"Bearer {actors.keys[Actor.PROXY_ADMIN].cleartext}"},
json={"key_alias": resources.prefix, "team_id": resources.tag("no-such-team")},
)
assert resp.status_code == 400, resp.text
rows = await prisma.db.litellm_verificationtoken.find_many(
where={"key_alias": scratch.prefix}
rows = await db.db.litellm_verificationtoken.find_many(
where={"key_alias": resources.prefix}
)
assert rows == []

84
tests/proxy_behavior/management/test_key_team_change.py → tests/integration/management/test_key_team_change.py
View file

@ -21,8 +21,8 @@ from prisma import Json
from litellm.proxy.utils import hash_token
from .actors import TEAM_ALPHA, Actor
from .conftest import create_scratch_team
from tests.integration._support.actors import TEAM_ALPHA, Actor
from tests.integration._support.actors import create_scratch_team
pytestmark = pytest.mark.asyncio(loop_scope="session")
@ -30,7 +30,7 @@ KEY_MODEL = "phase4-f2-key-model"
async def _seed_key_with_limits(
prisma,
db,
scratch_prefix: str,
*,
user_id: str,
@ -39,7 +39,7 @@ async def _seed_key_with_limits(
tpm_limit: Optional[int] = None,
rpm_limit: Optional[int] = None,
) -> str:
"""Raw-seed a scratch key with explicit models / tpm / rpm — /key/generate
"""Raw-seed a resources key with explicit models / tpm / rpm — /key/generate
can't set these against a non-throughput team without firing F1's guards.
Returns the cleartext key for /key/update calls."""
cleartext = "sk-" + uuid.uuid4().hex
@ -55,45 +55,45 @@ async def _seed_key_with_limits(
data["tpm_limit"] = tpm_limit
if rpm_limit is not None:
data["rpm_limit"] = rpm_limit
await prisma.db.litellm_verificationtoken.create(data=data)
await db.db.litellm_verificationtoken.create(data=data)
return cleartext
# ---------------------------------------------------------------------------
# Accepted — proxy admin moves a key into a scratch team that has the model,
# Accepted — proxy admin moves a key into a resources team that has the model,
# accommodates the limits, and has the key owner as a member.
# ---------------------------------------------------------------------------
async def test_key_team_change_accepted(proxy_client, prisma, scratch, world):
owner_id = world.keys[Actor.OWNER].user_id
async def test_key_team_change_accepted(gateway_client, db, resources, actors):
owner_id = actors.keys[Actor.OWNER].user_id
target_team = await create_scratch_team(
prisma,
team_id=scratch.tag("target"),
db,
team_id=resources.tag("target"),
member_user_ids=[owner_id],
models=[KEY_MODEL],
tpm_limit=10_000,
rpm_limit=1_000,
)
key_cleartext = await _seed_key_with_limits(
prisma,
scratch.prefix,
db,
resources.prefix,
user_id=owner_id,
team_id=TEAM_ALPHA,
team_id=actors.team_alpha_id,
models=[KEY_MODEL],
tpm_limit=500,
rpm_limit=50,
)
seeder = world.keys[Actor.PROXY_ADMIN].cleartext
seeder = actors.keys[Actor.PROXY_ADMIN].cleartext
resp = await proxy_client.post(
resp = await gateway_client.post(
"/key/update",
headers={"Authorization": f"Bearer {seeder}"},
json={"key": key_cleartext, "team_id": target_team},
)
assert resp.status_code == 200, resp.text
row = await prisma.db.litellm_verificationtoken.find_unique(
row = await db.db.litellm_verificationtoken.find_unique(
where={"token": hash_token(key_cleartext)}
)
assert row is not None
@ -115,9 +115,9 @@ async def test_key_team_change_accepted(proxy_client, prisma, scratch, world):
],
)
async def test_key_team_change_rejected_guards(
scenario: str, proxy_client, prisma, scratch, world
scenario: str, gateway_client, db, resources, actors
):
owner_id = world.keys[Actor.OWNER].user_id
owner_id = actors.keys[Actor.OWNER].user_id
team_kwargs: Dict[str, Any] = {
"member_user_ids": [owner_id],
@ -141,20 +141,20 @@ async def test_key_team_change_rejected_guards(
team_kwargs["member_user_ids"] = []
target_team = await create_scratch_team(
prisma,
team_id=scratch.tag("target"),
db,
team_id=resources.tag("target"),
**team_kwargs,
)
key_cleartext = await _seed_key_with_limits(
prisma,
scratch.prefix,
db,
resources.prefix,
user_id=owner_id,
team_id=TEAM_ALPHA,
team_id=actors.team_alpha_id,
**key_kwargs,
)
seeder = world.keys[Actor.PROXY_ADMIN].cleartext
seeder = actors.keys[Actor.PROXY_ADMIN].cleartext
resp = await proxy_client.post(
resp = await gateway_client.post(
"/key/update",
headers={"Authorization": f"Bearer {seeder}"},
json={"key": key_cleartext, "team_id": target_team},
@ -168,11 +168,11 @@ async def test_key_team_change_rejected_guards(
403,
), f"{scenario}: expected 400/403, got {resp.status_code}: {resp.text}"
row = await prisma.db.litellm_verificationtoken.find_unique(
row = await db.db.litellm_verificationtoken.find_unique(
where={"token": hash_token(key_cleartext)}
)
assert row is not None
assert row.team_id == TEAM_ALPHA, (
assert row.team_id == actors.team_alpha_id, (
f"{scenario}: row team_id mutated despite rejection — " f"got {row.team_id!r}"
)
@ -186,34 +186,34 @@ async def test_key_team_change_rejected_guards(
async def test_key_team_change_rejected_initiator_not_admin(
proxy_client, prisma, scratch, world
gateway_client, db, resources, actors
):
owner_id = world.keys[Actor.OWNER].user_id
owner_id = actors.keys[Actor.OWNER].user_id
target_team = await create_scratch_team(
prisma,
team_id=scratch.tag("target"),
db,
team_id=resources.tag("target"),
# Owner is the key's user_id; member of the team to clear membership
# guard. Internal-user role on world.keys[INTERNAL_USER] is also a
# member of TEAM_ALPHA but never an admin → role-gate path fires.
member_user_ids=[owner_id, world.keys[Actor.INTERNAL_USER].user_id],
# guard. Internal-user role on actors.keys[INTERNAL_USER] is also a
# member of actors.team_alpha_id but never an admin → role-gate path fires.
member_user_ids=[owner_id, actors.keys[Actor.INTERNAL_USER].user_id],
models=[KEY_MODEL],
tpm_limit=10_000,
rpm_limit=1_000,
)
key_cleartext = await _seed_key_with_limits(
prisma,
scratch.prefix,
db,
resources.prefix,
user_id=owner_id,
team_id=TEAM_ALPHA,
team_id=actors.team_alpha_id,
models=[KEY_MODEL],
tpm_limit=500,
rpm_limit=50,
)
# The internal_user actor is the initiator: a TEAM_ALPHA member, but not
# The internal_user actor is the initiator: a actors.team_alpha_id member, but not
# an admin anywhere, and has no team_member_permissions for /key/update.
initiator = world.keys[Actor.INTERNAL_USER].cleartext
initiator = actors.keys[Actor.INTERNAL_USER].cleartext
resp = await proxy_client.post(
resp = await gateway_client.post(
"/key/update",
headers={"Authorization": f"Bearer {initiator}"},
json={"key": key_cleartext, "team_id": target_team},
@ -222,8 +222,8 @@ async def test_key_team_change_rejected_initiator_not_admin(
401,
403,
), f"expected 401/403, got {resp.status_code}: {resp.text}"
row = await prisma.db.litellm_verificationtoken.find_unique(
row = await db.db.litellm_verificationtoken.find_unique(
where={"token": hash_token(key_cleartext)}
)
assert row is not None
assert row.team_id == TEAM_ALPHA, "row team_id mutated despite rejection"
assert row.team_id == actors.team_alpha_id, "row team_id mutated despite rejection"

78
tests/proxy_behavior/management/test_key_update.py → tests/integration/management/test_key_update.py
View file

@ -4,8 +4,8 @@ import pytest
from litellm.proxy.utils import hash_token
from .actors import TEAM_ALPHA, TEAM_BETA, Actor
from .conftest import create_scratch_key
from tests.integration._support.actors import TEAM_ALPHA, TEAM_BETA, Actor
from tests.integration._support.actors import create_scratch_key
pytestmark = pytest.mark.asyncio(loop_scope="session")
@ -50,40 +50,40 @@ async def test_key_update_authz_matrix(
actor: Actor,
target_shape: str,
expected_status: int,
proxy_client,
prisma,
scratch,
world,
gateway_client,
db,
resources,
actors,
):
caller = world.keys[actor]
seeder = world.keys[Actor.PROXY_ADMIN].cleartext
caller = actors.keys[actor]
seeder = actors.keys[Actor.PROXY_ADMIN].cleartext
if target_shape == "self":
target_cleartext = await create_scratch_key(
proxy_client, seeder, scratch.prefix, user_id=caller.user_id
gateway_client, seeder, resources.prefix, user_id=caller.user_id
)
elif target_shape == "owner":
target_cleartext = await create_scratch_key(
proxy_client,
gateway_client,
seeder,
scratch.prefix,
user_id=world.keys[Actor.OWNER].user_id,
team_id=TEAM_ALPHA,
resources.prefix,
user_id=actors.keys[Actor.OWNER].user_id,
team_id=actors.team_alpha_id,
)
elif target_shape == "cross_org":
target_cleartext = await create_scratch_key(
proxy_client,
gateway_client,
seeder,
scratch.prefix,
user_id=world.keys[Actor.CROSS_ORG_USER].user_id,
team_id=TEAM_BETA,
resources.prefix,
user_id=actors.keys[Actor.CROSS_ORG_USER].user_id,
team_id=actors.team_beta_id,
)
else:
pytest.fail(f"unknown target_shape={target_shape}")
target_hashed = hash_token(target_cleartext)
resp = await proxy_client.post(
resp = await gateway_client.post(
"/key/update",
headers={"Authorization": f"Bearer {caller.cleartext}"},
json={"key": target_cleartext, "models": [MARKER_MODEL]},
@ -92,7 +92,7 @@ async def test_key_update_authz_matrix(
resp.status_code == expected_status
), f"{actor.value} {target_shape}: {resp.status_code} {resp.text}"
row = await prisma.db.litellm_verificationtoken.find_unique(
row = await db.db.litellm_verificationtoken.find_unique(
where={"token": target_hashed}
)
assert row is not None
@ -102,35 +102,35 @@ async def test_key_update_authz_matrix(
assert row.models != [MARKER_MODEL], "denied but row mutated"
async def _seed_shape(proxy_client, seeder, prefix, world, shape, caller) -> str:
async def _seed_shape(gateway_client, seeder, prefix, actors, shape, caller) -> str:
if shape == "self":
return await create_scratch_key(
proxy_client, seeder, prefix, user_id=caller.user_id
gateway_client, seeder, prefix, user_id=caller.user_id
)
if shape == "owner":
return await create_scratch_key(
proxy_client,
gateway_client,
seeder,
prefix,
user_id=world.keys[Actor.OWNER].user_id,
team_id=TEAM_ALPHA,
user_id=actors.keys[Actor.OWNER].user_id,
team_id=actors.team_alpha_id,
)
if shape == "cross_org":
return await create_scratch_key(
proxy_client,
gateway_client,
seeder,
prefix,
user_id=world.keys[Actor.CROSS_ORG_USER].user_id,
team_id=TEAM_BETA,
user_id=actors.keys[Actor.CROSS_ORG_USER].user_id,
team_id=actors.team_beta_id,
)
pytest.fail(f"unknown shape={shape}") # pragma: no cover
async def test_key_update_missing_key_is_404(proxy_client, world):
async def test_key_update_missing_key_is_404(gateway_client, actors):
"""An update targeting a key absent from the DB is a 404 — not 401/403."""
resp = await proxy_client.post(
resp = await gateway_client.post(
"/key/update",
headers={"Authorization": f"Bearer {world.keys[Actor.PROXY_ADMIN].cleartext}"},
headers={"Authorization": f"Bearer {actors.keys[Actor.PROXY_ADMIN].cleartext}"},
json={"key": "sk-" + uuid.uuid4().hex, "models": [MARKER_MODEL]},
)
assert resp.status_code == 404, resp.text
@ -154,19 +154,19 @@ async def test_key_update_denied_does_not_touch_budget_counters(
actor: Actor,
target_shape: str,
expected_status: int,
proxy_client,
prisma,
scratch,
world,
gateway_client,
db,
resources,
actors,
):
caller = world.keys[actor]
seeder = world.keys[Actor.PROXY_ADMIN].cleartext
caller = actors.keys[actor]
seeder = actors.keys[Actor.PROXY_ADMIN].cleartext
target = await _seed_shape(
proxy_client, seeder, scratch.prefix, world, target_shape, caller
gateway_client, seeder, resources.prefix, actors, target_shape, caller
)
target_hashed = hash_token(target)
resp = await proxy_client.post(
resp = await gateway_client.post(
"/key/update",
headers={"Authorization": f"Bearer {caller.cleartext}"},
json={"key": target, "max_budget": 999.0, "tpm_limit": 888, "rpm_limit": 777},
@ -175,7 +175,7 @@ async def test_key_update_denied_does_not_touch_budget_counters(
resp.status_code == expected_status
), f"{actor.value} {target_shape}: {resp.status_code} {resp.text}"
row = await prisma.db.litellm_verificationtoken.find_unique(
row = await db.db.litellm_verificationtoken.find_unique(
where={"token": target_hashed}
)
assert row is not None

22
tests/integration/management/test_smoke.py Normal file
View file

@ -0,0 +1,22 @@
import pytest
pytestmark = pytest.mark.asyncio(loop_scope="session")
async def test_key_generate_lands_in_db(gateway_client, db, resources):
from litellm.proxy.utils import hash_token
resp = await gateway_client.post(
"/key/generate",
headers={"Authorization": f"Bearer {gateway_client.gateway.key}"},
json={"key_alias": resources.prefix},
)
assert resp.status_code == 200, resp.text
cleartext = resp.json()["key"]
assert cleartext.startswith("sk-")
hashed = hash_token(cleartext)
row = await db.db.litellm_verificationtoken.find_unique(where={"token": hashed})
assert row is not None
assert row.token == hashed != cleartext

10
tests/proxy_behavior/management/test_team_available.py → tests/integration/management/test_team_available.py
View file

@ -1,21 +1,21 @@
import pytest
from .actors import Actor
from tests.integration._support.actors import Actor
pytestmark = pytest.mark.asyncio(loop_scope="session")
# GET /team/available lists teams from
# litellm.default_internal_user_params["available_teams"]. The behavior world
# litellm.default_internal_user_params["available_teams"]. The behavior actors
# configures no available_teams, so the handler returns [] for every actor
# before it even reads the caller — this is the route-coverage + default-path
# pin. /team/available is an info route, so every authenticated actor reaches
# the handler.
@pytest.mark.parametrize("actor", list(Actor), ids=[a.value for a in Actor])
async def test_team_available_default_is_empty(actor: Actor, proxy_client, world):
resp = await proxy_client.get(
async def test_team_available_default_is_empty(actor: Actor, gateway_client, actors):
resp = await gateway_client.get(
"/team/available",
headers={"Authorization": f"Bearer {world.keys[actor].cleartext}"},
headers={"Authorization": f"Bearer {actors.keys[actor].cleartext}"},
)
assert resp.status_code == 200, f"{actor.value}: {resp.status_code} {resp.text}"
assert resp.json() == []

64
tests/proxy_behavior/management/test_team_block_unblock.py → tests/integration/management/test_team_block_unblock.py
View file

@ -1,7 +1,7 @@
import pytest
from .actors import Actor
from .conftest import create_scratch_team
from tests.integration._support.actors import Actor
from tests.integration._support.actors import create_scratch_team
pytestmark = pytest.mark.asyncio(loop_scope="session")
@ -28,10 +28,10 @@ _MATRIX = [
]
async def _seed_target(prisma, world, shape: str, team_id: str) -> str:
"""Raw-seed the scratch target team; returns its organization_id."""
org_id = world.org_a_id if shape == "alpha" else world.org_b_id
await create_scratch_team(prisma, team_id, organization_id=org_id)
async def _seed_target(db, actors, shape: str, team_id: str) -> str:
"""Raw-seed the resources target team; returns its organization_id."""
org_id = actors.org_a_id if shape == "alpha" else actors.org_b_id
await create_scratch_team(db, team_id, organization_id=org_id)
return org_id
@ -46,31 +46,31 @@ async def test_team_block_unblock_authz_matrix(
actor: Actor,
shape: str,
expected_status: int,
proxy_client,
prisma,
scratch,
world,
gateway_client,
db,
resources,
actors,
):
org_id = await _seed_target(prisma, world, shape, scratch.prefix)
caller = world.keys[actor]
org_id = await _seed_target(db, actors, shape, resources.prefix)
caller = actors.keys[actor]
# /unblock starts from a blocked row so a 200 is observable as True->False.
if route == "unblock":
await prisma.db.litellm_teamtable.update(
where={"team_id": scratch.prefix}, data={"blocked": True}
await db.db.litellm_teamtable.update(
where={"team_id": resources.prefix}, data={"blocked": True}
)
resp = await proxy_client.post(
resp = await gateway_client.post(
f"/team/{route}",
headers={"Authorization": f"Bearer {caller.cleartext}"},
json={"team_id": scratch.prefix, "organization_id": org_id},
json={"team_id": resources.prefix, "organization_id": org_id},
)
assert (
resp.status_code == expected_status
), f"{route} {actor.value} {shape}: {resp.status_code} {resp.text}"
row = await prisma.db.litellm_teamtable.find_unique(
where={"team_id": scratch.prefix}
row = await db.db.litellm_teamtable.find_unique(
where={"team_id": resources.prefix}
)
assert row is not None
if expected_status == 200:
@ -79,36 +79,36 @@ async def test_team_block_unblock_authz_matrix(
assert bool(row.blocked) is (route == "unblock"), "denied but blocked mutated"
async def test_team_block_unblock_round_trip(proxy_client, prisma, scratch, world):
async def test_team_block_unblock_round_trip(gateway_client, db, resources, actors):
"""PROXY_ADMIN block then unblock flips the blocked column True then False."""
await create_scratch_team(prisma, scratch.prefix, organization_id=world.org_a_id)
headers = {"Authorization": f"Bearer {world.keys[Actor.PROXY_ADMIN].cleartext}"}
await create_scratch_team(db, resources.prefix, organization_id=actors.org_a_id)
headers = {"Authorization": f"Bearer {actors.keys[Actor.PROXY_ADMIN].cleartext}"}
blocked = await proxy_client.post(
"/team/block", headers=headers, json={"team_id": scratch.prefix}
blocked = await gateway_client.post(
"/team/block", headers=headers, json={"team_id": resources.prefix}
)
assert blocked.status_code == 200, blocked.text
row = await prisma.db.litellm_teamtable.find_unique(
where={"team_id": scratch.prefix}
row = await db.db.litellm_teamtable.find_unique(
where={"team_id": resources.prefix}
)
assert row is not None and row.blocked is True
unblocked = await proxy_client.post(
"/team/unblock", headers=headers, json={"team_id": scratch.prefix}
unblocked = await gateway_client.post(
"/team/unblock", headers=headers, json={"team_id": resources.prefix}
)
assert unblocked.status_code == 200, unblocked.text
row = await prisma.db.litellm_teamtable.find_unique(
where={"team_id": scratch.prefix}
row = await db.db.litellm_teamtable.find_unique(
where={"team_id": resources.prefix}
)
assert row is not None and row.blocked is False
@pytest.mark.parametrize("route", ["block", "unblock"])
async def test_team_block_unblock_missing_team_is_404(route: str, proxy_client, world):
async def test_team_block_unblock_missing_team_is_404(route: str, gateway_client, actors):
"""A team_id absent from the DB is 404 — the existence check precedes authz."""
resp = await proxy_client.post(
resp = await gateway_client.post(
f"/team/{route}",
headers={"Authorization": f"Bearer {world.keys[Actor.PROXY_ADMIN].cleartext}"},
headers={"Authorization": f"Bearer {actors.keys[Actor.PROXY_ADMIN].cleartext}"},
json={"team_id": "behavior-pin-no-such-team"},
)
assert resp.status_code == 404, resp.text

142
tests/proxy_behavior/management/test_team_budget_limits.py → tests/integration/management/test_team_budget_limits.py
View file

@ -26,14 +26,14 @@ import pytest
from litellm.proxy.utils import hash_token
from .actors import Actor
from .conftest import MASTER_KEY, create_scratch_org, create_scratch_team
from tests.integration._support.actors import Actor
from tests.integration._support.actors import create_scratch_org, create_scratch_team
pytestmark = pytest.mark.asyncio(loop_scope="session")
async def _seed_scratch_actor_with_caps(
prisma,
db,
scratch_prefix: str,
*,
user_role: str = "internal_user",
@ -42,7 +42,7 @@ async def _seed_scratch_actor_with_caps(
tpm_limit: Optional[int] = None,
rpm_limit: Optional[int] = None,
) -> str:
"""Raw-seed a scratch actor + verification token, with the token carrying
"""Raw-seed a resources actor + verification token, with the token carrying
explicit caps that flow into `user_api_key_dict` at request time.
Returns cleartext key. Used by F3 user-team-limit scenarios where the
@ -54,7 +54,7 @@ async def _seed_scratch_actor_with_caps(
"""
user_id = f"{scratch_prefix}-team-creator"
cleartext = "sk-" + uuid.uuid4().hex
await prisma.db.litellm_usertable.create(
await db.db.litellm_usertable.create(
data={
"user_id": user_id,
"user_role": user_role,
@ -73,7 +73,7 @@ async def _seed_scratch_actor_with_caps(
token_data["tpm_limit"] = tpm_limit
if rpm_limit is not None:
token_data["rpm_limit"] = rpm_limit
await prisma.db.litellm_verificationtoken.create(data=token_data)
await db.db.litellm_verificationtoken.create(data=token_data)
return cleartext
@ -112,28 +112,28 @@ async def test_check_org_team_limits_models_subset(
org_models,
body_extras: Dict[str, Any],
expected_status: int,
proxy_client,
prisma,
scratch,
world,
gateway_client,
db,
resources,
actors,
):
org_id = await create_scratch_org(prisma, scratch.prefix, models=org_models)
seeder = world.keys[Actor.PROXY_ADMIN].cleartext
team_id = scratch.tag("team")
org_id = await create_scratch_org(db, resources.prefix, models=org_models)
seeder = actors.keys[Actor.PROXY_ADMIN].cleartext
team_id = resources.tag("team")
body: Dict[str, Any] = {
"team_id": team_id,
"team_alias": scratch.prefix,
"team_alias": resources.prefix,
"organization_id": org_id,
**body_extras,
}
resp = await proxy_client.post(
resp = await gateway_client.post(
"/team/new",
headers={"Authorization": f"Bearer {seeder}"},
json=body,
)
assert resp.status_code == expected_status, f"{body!r} → {resp.status_code}: {resp.text}"
rows = await prisma.db.litellm_teamtable.find_many(where={"team_id": team_id})
rows = await db.db.litellm_teamtable.find_many(where={"team_id": team_id})
assert len(rows) == (1 if expected_status == 200 else 0)
@ -191,26 +191,26 @@ async def test_check_org_team_limits_budget_enforced(
org_budget,
body_extras: Dict[str, Any],
expected_status: int,
proxy_client,
prisma,
scratch,
world,
gateway_client,
db,
resources,
actors,
):
org_id = await create_scratch_org(prisma, scratch.prefix, **org_budget)
seeder = world.keys[Actor.PROXY_ADMIN].cleartext
team_id = scratch.tag("team")
resp = await proxy_client.post(
org_id = await create_scratch_org(db, resources.prefix, **org_budget)
seeder = actors.keys[Actor.PROXY_ADMIN].cleartext
team_id = resources.tag("team")
resp = await gateway_client.post(
"/team/new",
headers={"Authorization": f"Bearer {seeder}"},
json={
"team_id": team_id,
"team_alias": scratch.prefix,
"team_alias": resources.prefix,
"organization_id": org_id,
**body_extras,
},
)
assert resp.status_code == expected_status, f"{body_extras!r} → {resp.status_code}: {resp.text}"
rows = await prisma.db.litellm_teamtable.find_many(where={"team_id": team_id})
rows = await db.db.litellm_teamtable.find_many(where={"team_id": team_id})
assert len(rows) == (1 if expected_status == 200 else 0)
@ -223,21 +223,21 @@ async def test_check_org_team_limits_budget_enforced_on_update(
org_budget,
body_extras: Dict[str, Any],
expected_status: int,
proxy_client,
prisma,
scratch,
world,
gateway_client,
db,
resources,
actors,
):
org_id = await create_scratch_org(prisma, scratch.prefix, **org_budget)
team_id = await create_scratch_team(prisma, scratch.tag("team"), organization_id=org_id)
seeder = world.keys[Actor.PROXY_ADMIN].cleartext
resp = await proxy_client.post(
org_id = await create_scratch_org(db, resources.prefix, **org_budget)
team_id = await create_scratch_team(db, resources.tag("team"), organization_id=org_id)
seeder = actors.keys[Actor.PROXY_ADMIN].cleartext
resp = await gateway_client.post(
"/team/update",
headers={"Authorization": f"Bearer {seeder}"},
json={"team_id": team_id, **body_extras},
)
assert resp.status_code == expected_status, f"{body_extras!r} → {resp.status_code}: {resp.text}"
row = await prisma.db.litellm_teamtable.find_unique(where={"team_id": team_id})
row = await db.db.litellm_teamtable.find_unique(where={"team_id": team_id})
assert row is not None
persisted = {field: getattr(row, field) for field in body_extras}
assert (persisted == body_extras) == (expected_status == 200)
@ -309,18 +309,18 @@ async def test_check_user_team_limits(
actor_caps,
body_extras: Dict[str, Any],
expected_status: int,
proxy_client,
prisma,
scratch,
gateway_client,
db,
resources,
):
caller = await _seed_scratch_actor_with_caps(prisma, scratch.prefix, **actor_caps)
team_id = scratch.tag("team")
resp = await proxy_client.post(
caller = await _seed_scratch_actor_with_caps(db, resources.prefix, **actor_caps)
team_id = resources.tag("team")
resp = await gateway_client.post(
"/team/new",
headers={"Authorization": f"Bearer {caller}"},
json={
"team_id": team_id,
"team_alias": scratch.prefix,
"team_alias": resources.prefix,
# Standalone team — no organization_id, so user-limit guard fires.
**body_extras,
},
@ -329,7 +329,7 @@ async def test_check_user_team_limits(
f"caps={actor_caps} body={body_extras} → {resp.status_code}: {resp.text}"
)
rows = await prisma.db.litellm_teamtable.find_many(where={"team_id": team_id})
rows = await db.db.litellm_teamtable.find_many(where={"team_id": team_id})
assert len(rows) == (1 if expected_status == 200 else 0)
@ -351,16 +351,16 @@ async def test_check_user_team_limits(
ids=["raise_with_generous_personal_budget", "lower_with_tiny_personal_budget"],
)
async def test_team_admin_cannot_change_budget_while_max_budget_is_not_editable(
proxy_client, prisma, scratch, personal_budget: float, requested_budget: float
gateway_client, db, resources, personal_budget: float, requested_budget: float
):
caller_cleartext = await _seed_scratch_actor_with_caps(prisma, scratch.prefix, max_budget=personal_budget)
caller_cleartext = await _seed_scratch_actor_with_caps(db, resources.prefix, max_budget=personal_budget)
team_id = await create_scratch_team(
prisma,
team_id=scratch.tag("team"),
admin_user_ids=[f"{scratch.prefix}-team-creator"],
db,
team_id=resources.tag("team"),
admin_user_ids=[f"{resources.prefix}-team-creator"],
max_budget=500.0,
)
resp = await proxy_client.post(
resp = await gateway_client.post(
"/team/update",
headers={"Authorization": f"Bearer {caller_cleartext}"},
json={"team_id": team_id, "max_budget": requested_budget},
@ -368,69 +368,69 @@ async def test_team_admin_cannot_change_budget_while_max_budget_is_not_editable(
assert resp.status_code == 403, resp.text
assert "Team admin editable fields" in resp.text, resp.text
row = await prisma.db.litellm_teamtable.find_unique(where={"team_id": team_id})
row = await db.db.litellm_teamtable.find_unique(where={"team_id": team_id})
assert row is not None
assert row.max_budget == 500.0, "a refused update must leave the team budget unchanged"
async def test_proxy_admin_raise_budget_allowed(proxy_client, prisma, scratch):
async def test_proxy_admin_raise_budget_allowed(gateway_client, db, resources):
"""A proxy admin may raise a team's budget."""
team_id = await create_scratch_team(
prisma,
team_id=scratch.tag("team"),
admin_user_ids=[f"{scratch.prefix}-team-creator"],
db,
team_id=resources.tag("team"),
admin_user_ids=[f"{resources.prefix}-team-creator"],
max_budget=50.0,
)
# MASTER_KEY acts as proxy admin.
resp = await proxy_client.post(
resp = await gateway_client.post(
"/team/update",
headers={"Authorization": f"Bearer {MASTER_KEY}"},
headers={"Authorization": f"Bearer {gateway_client.gateway.key}"},
json={"team_id": team_id, "max_budget": 999.0},
)
assert resp.status_code == 200, resp.text
row = await prisma.db.litellm_teamtable.find_unique(where={"team_id": team_id})
row = await db.db.litellm_teamtable.find_unique(where={"team_id": team_id})
assert row is not None
assert row.max_budget == 999.0, "proxy admin should be able to raise the budget"
async def test_team_admin_remove_budget_cap_blocked(proxy_client, prisma, scratch):
async def test_team_admin_remove_budget_cap_blocked(gateway_client, db, resources):
"""A team admin cannot strip the team's cap (max_budget=null); removing the
ceiling is the strongest possible raise -> proxy-admin only."""
caller_cleartext = await _seed_scratch_actor_with_caps(prisma, scratch.prefix, max_budget=100000.0)
caller_cleartext = await _seed_scratch_actor_with_caps(db, resources.prefix, max_budget=100000.0)
team_id = await create_scratch_team(
prisma,
team_id=scratch.tag("team"),
admin_user_ids=[f"{scratch.prefix}-team-creator"],
db,
team_id=resources.tag("team"),
admin_user_ids=[f"{resources.prefix}-team-creator"],
max_budget=50.0,
)
resp = await proxy_client.post(
resp = await gateway_client.post(
"/team/update",
headers={"Authorization": f"Bearer {caller_cleartext}"},
json={"team_id": team_id, "max_budget": None},
)
assert resp.status_code == 403, resp.text
row = await prisma.db.litellm_teamtable.find_unique(where={"team_id": team_id})
row = await db.db.litellm_teamtable.find_unique(where={"team_id": team_id})
assert row is not None
assert row.max_budget == 50.0, "team budget cap must not be removed by a team admin"
async def test_proxy_admin_remove_budget_cap_allowed(proxy_client, prisma, scratch):
async def test_proxy_admin_remove_budget_cap_allowed(gateway_client, db, resources):
"""A proxy admin may remove a team's cap (max_budget=null)."""
team_id = await create_scratch_team(
prisma,
team_id=scratch.tag("team"),
admin_user_ids=[f"{scratch.prefix}-team-creator"],
db,
team_id=resources.tag("team"),
admin_user_ids=[f"{resources.prefix}-team-creator"],
max_budget=50.0,
)
resp = await proxy_client.post(
resp = await gateway_client.post(
"/team/update",
headers={"Authorization": f"Bearer {MASTER_KEY}"},
headers={"Authorization": f"Bearer {gateway_client.gateway.key}"},
json={"team_id": team_id, "max_budget": None},
)
assert resp.status_code == 200, resp.text
row = await prisma.db.litellm_teamtable.find_unique(where={"team_id": team_id})
row = await db.db.litellm_teamtable.find_unique(where={"team_id": team_id})
assert row is not None
assert row.max_budget is None, "proxy admin should be able to remove the cap"

105
tests/integration/management/test_team_bulk_member_add.py Normal file
View file

@ -0,0 +1,105 @@
import pytest
from tests.integration._support.actors import Actor
from tests.integration._support.actors import create_scratch_team
pytestmark = pytest.mark.asyncio(loop_scope="session")
def _member_ids(row) -> list:
return [m["user_id"] for m in (row.members_with_roles or [])]
async def test_team_bulk_member_add_proxy_admin_adds_explicit_members(
gateway_client, db, resources, actors
):
"""PROXY_ADMIN bulk-adds an explicit member list to a resources team."""
await create_scratch_team(db, resources.prefix, organization_id=actors.org_a_id)
new_member = resources.tag("m1")
resp = await gateway_client.post(
"/team/bulk_member_add",
headers={"Authorization": f"Bearer {actors.keys[Actor.PROXY_ADMIN].cleartext}"},
json={
"team_id": resources.prefix,
"members": [{"user_id": new_member, "role": "user"}],
},
)
assert resp.status_code == 200, resp.text
row = await db.db.litellm_teamtable.find_unique(
where={"team_id": resources.prefix}
)
assert row is not None and new_member in _member_ids(row)
async def test_team_bulk_member_add_empty_members_is_400(
gateway_client, db, resources, actors
):
"""An empty member list (with all_users unset) is rejected 400."""
await create_scratch_team(db, resources.prefix, organization_id=actors.org_a_id)
resp = await gateway_client.post(
"/team/bulk_member_add",
headers={"Authorization": f"Bearer {actors.keys[Actor.PROXY_ADMIN].cleartext}"},
json={"team_id": resources.prefix, "members": []},
)
assert resp.status_code == 400, resp.text
async def test_team_bulk_member_add_over_max_batch_is_400(
gateway_client, db, resources, actors
):
"""A member list larger than the 500-member cap is rejected 400."""
await create_scratch_team(db, resources.prefix, organization_id=actors.org_a_id)
members = [
{"user_id": f"{resources.prefix}-u{i}", "role": "user"} for i in range(501)
]
resp = await gateway_client.post(
"/team/bulk_member_add",
headers={"Authorization": f"Bearer {actors.keys[Actor.PROXY_ADMIN].cleartext}"},
json={"team_id": resources.prefix, "members": members},
)
assert resp.status_code == 400, resp.text
@pytest.mark.parametrize(
"actor",
[Actor.TEAM_ADMIN, Actor.INTERNAL_USER],
ids=["team_admin", "internal_user"],
)
async def test_team_bulk_member_add_non_admin_is_401(
actor: Actor, gateway_client, db, resources, actors
):
"""/team/bulk_member_add is neither an internal_user nor a self-managed
route — a non-proxy-admin with no org context is 401 at the route gate."""
await create_scratch_team(db, resources.prefix, organization_id=actors.org_a_id)
resp = await gateway_client.post(
"/team/bulk_member_add",
headers={"Authorization": f"Bearer {actors.keys[actor].cleartext}"},
json={
"team_id": resources.prefix,
"members": [{"user_id": resources.tag("m"), "role": "user"}],
},
)
assert resp.status_code == 401, f"{actor.value}: {resp.status_code} {resp.text}"
async def test_team_bulk_member_add_all_users_proxy_admin(
gateway_client, db, resources, actors
):
"""all_users=True pulls every user in the DB into the team. The route is
reachable only by PROXY_ADMIN (the route gate 401s every other actor — even
an org admin with organization_id in the body), so the handler's own
all_users PROXY_ADMIN gate is never the deciding check at the boundary."""
await create_scratch_team(db, resources.prefix, organization_id=actors.org_a_id)
resp = await gateway_client.post(
"/team/bulk_member_add",
headers={"Authorization": f"Bearer {actors.keys[Actor.PROXY_ADMIN].cleartext}"},
json={"team_id": resources.prefix, "all_users": True},
)
assert resp.status_code == 200, resp.text
row = await db.db.litellm_teamtable.find_unique(
where={"team_id": resources.prefix}
)
assert row is not None
member_ids = _member_ids(row)
# every actors actor is a user in the DB, so all are now team members
assert actors.keys[Actor.INTERNAL_USER].user_id in member_ids

140
tests/proxy_behavior/management/test_team_bulk_member_delete.py → tests/integration/management/test_team_bulk_member_delete.py
View file

@ -1,8 +1,8 @@
import pytest
from prisma import Json
from .actors import Actor
from .conftest import create_scratch_team, create_scratch_user
from tests.integration._support.actors import Actor
from tests.integration._support.actors import create_scratch_team, create_scratch_user
pytestmark = pytest.mark.asyncio(loop_scope="session")
@ -28,20 +28,20 @@ _MATRIX = [
]
async def _seed_target(prisma, world, shape: str, team_id: str, victim_ids: list) -> None:
async def _seed_target(db, actors, shape: str, team_id: str, victim_ids: list) -> None:
if shape == "alpha":
await create_scratch_team(
prisma,
db,
team_id,
organization_id=world.org_a_id,
admin_user_ids=[world.keys[Actor.TEAM_ADMIN].user_id],
organization_id=actors.org_a_id,
admin_user_ids=[actors.keys[Actor.TEAM_ADMIN].user_id],
member_user_ids=victim_ids,
)
elif shape == "beta":
await create_scratch_team(
prisma,
db,
team_id,
organization_id=world.org_b_id,
organization_id=actors.org_b_id,
member_user_ids=victim_ids,
)
else: # pragma: no cover - guard
@ -61,18 +61,18 @@ async def test_team_bulk_member_delete_authz_matrix(
actor: Actor,
shape: str,
expected_status: int,
proxy_client,
prisma,
scratch,
world,
gateway_client,
db,
resources,
actors,
):
victims = [scratch.tag("v1"), scratch.tag("v2")]
keep = scratch.tag("keep")
await _seed_target(prisma, world, shape, scratch.prefix, victims + [keep])
caller = world.keys[actor]
victims = [resources.tag("v1"), resources.tag("v2")]
keep = resources.tag("keep")
await _seed_target(db, actors, shape, resources.prefix, victims + [keep])
caller = actors.keys[actor]
resp = await proxy_client.post(
f"/management/v1/teams/{scratch.prefix}/members/bulk_delete",
resp = await gateway_client.post(
f"/management/v1/teams/{resources.prefix}/members/bulk_delete",
headers={"Authorization": f"Bearer {caller.cleartext}"},
json={"members": [{"user_id": v} for v in victims]},
)
@ -81,7 +81,7 @@ async def test_team_bulk_member_delete_authz_matrix(
assert resp.headers["content-type"] == "application/problem+json"
assert resp.json()["type"] == "urn:litellm:error:forbidden"
row = await prisma.db.litellm_teamtable.find_unique(where={"team_id": scratch.prefix})
row = await db.db.litellm_teamtable.find_unique(where={"team_id": resources.prefix})
assert row is not None
assert keep in _member_ids(row), "unrelated member removed"
if expected_status == 200:
@ -91,15 +91,15 @@ async def test_team_bulk_member_delete_authz_matrix(
assert set(victims) <= set(_member_ids(row)), "denied but members removed"
async def test_team_bulk_member_delete_reports_each_row_in_order(proxy_client, prisma, scratch, world):
victim = scratch.tag("victim")
keep = scratch.tag("keep")
stranger = scratch.tag("stranger")
await create_scratch_team(prisma, scratch.prefix, organization_id=world.org_a_id, member_user_ids=[victim, keep])
async def test_team_bulk_member_delete_reports_each_row_in_order(gateway_client, db, resources, actors):
victim = resources.tag("victim")
keep = resources.tag("keep")
stranger = resources.tag("stranger")
await create_scratch_team(db, resources.prefix, organization_id=actors.org_a_id, member_user_ids=[victim, keep])
resp = await proxy_client.post(
f"/management/v1/teams/{scratch.prefix}/members/bulk_delete",
headers={"Authorization": f"Bearer {world.keys[Actor.PROXY_ADMIN].cleartext}"},
resp = await gateway_client.post(
f"/management/v1/teams/{resources.prefix}/members/bulk_delete",
headers={"Authorization": f"Bearer {actors.keys[Actor.PROXY_ADMIN].cleartext}"},
json={"members": [{"user_id": stranger}, {"user_id": victim}]},
)
assert resp.status_code == 200, resp.text
@ -111,47 +111,47 @@ async def test_team_bulk_member_delete_reports_each_row_in_order(proxy_client, p
]
assert body["data"][0]["error"] == "User not found in team"
row = await prisma.db.litellm_teamtable.find_unique(where={"team_id": scratch.prefix})
assert row is not None and _member_ids(row) == [keep]
row = await db.db.litellm_teamtable.find_unique(where={"team_id": resources.prefix})
assert row is not None and set(_member_ids(row)) == {keep, "default_user_id"}
async def test_team_bulk_member_delete_by_id_removes_a_legacy_email_only_roster_entry(
proxy_client, prisma, scratch, world
gateway_client, db, resources, actors
):
email = f"{scratch.prefix}@example.com"
victim = await create_scratch_user(prisma, scratch.prefix, suffix="victim", user_email=email)
keep = scratch.tag("keep")
await prisma.db.litellm_teamtable.create(
email = f"{resources.prefix}@example.com"
victim = await create_scratch_user(db, resources.prefix, suffix="victim", user_email=email)
keep = resources.tag("keep")
await db.db.litellm_teamtable.create(
data={
"team_id": scratch.prefix,
"team_alias": scratch.prefix,
"organization_id": world.org_a_id,
"team_id": resources.prefix,
"team_alias": resources.prefix,
"organization_id": actors.org_a_id,
"members_with_roles": Json([{"user_email": email, "role": "user"}, {"user_id": keep, "role": "user"}]),
}
)
await prisma.db.litellm_usertable.update(where={"user_id": victim}, data={"teams": [scratch.prefix]})
await db.db.litellm_usertable.update(where={"user_id": victim}, data={"teams": [resources.prefix]})
resp = await proxy_client.post(
f"/management/v1/teams/{scratch.prefix}/members/bulk_delete",
headers={"Authorization": f"Bearer {world.keys[Actor.PROXY_ADMIN].cleartext}"},
resp = await gateway_client.post(
f"/management/v1/teams/{resources.prefix}/members/bulk_delete",
headers={"Authorization": f"Bearer {actors.keys[Actor.PROXY_ADMIN].cleartext}"},
json={"members": [{"user_id": victim}]},
)
assert resp.status_code == 200, resp.text
assert [(r["user_id"], r["success"]) for r in resp.json()["data"]] == [(victim, True)]
row = await prisma.db.litellm_teamtable.find_unique(where={"team_id": scratch.prefix})
row = await db.db.litellm_teamtable.find_unique(where={"team_id": resources.prefix})
assert row is not None and [(m["user_id"], m.get("user_email")) for m in row.members_with_roles] == [(keep, None)]
user = await prisma.db.litellm_usertable.find_unique(where={"user_id": victim})
user = await db.db.litellm_usertable.find_unique(where={"user_id": victim})
assert user is not None and user.teams == []
async def test_team_bulk_member_delete_row_naming_both_identifiers_is_422(proxy_client, prisma, scratch, world):
victim = scratch.tag("victim")
await create_scratch_team(prisma, scratch.prefix, organization_id=world.org_a_id, member_user_ids=[victim])
async def test_team_bulk_member_delete_row_naming_both_identifiers_is_422(gateway_client, db, resources, actors):
victim = resources.tag("victim")
await create_scratch_team(db, resources.prefix, organization_id=actors.org_a_id, member_user_ids=[victim])
resp = await proxy_client.post(
f"/management/v1/teams/{scratch.prefix}/members/bulk_delete",
headers={"Authorization": f"Bearer {world.keys[Actor.PROXY_ADMIN].cleartext}"},
resp = await gateway_client.post(
f"/management/v1/teams/{resources.prefix}/members/bulk_delete",
headers={"Authorization": f"Bearer {actors.keys[Actor.PROXY_ADMIN].cleartext}"},
json={"members": [{"user_id": victim, "user_email": f"{victim}@example.com"}]},
)
assert resp.status_code == 422, resp.text
@ -162,48 +162,48 @@ async def test_team_bulk_member_delete_row_naming_both_identifiers_is_422(proxy_
== "members.0: Value error, Each member must be identified by exactly one of user_id or user_email"
)
row = await prisma.db.litellm_teamtable.find_unique(where={"team_id": scratch.prefix})
row = await db.db.litellm_teamtable.find_unique(where={"team_id": resources.prefix})
assert row is not None and victim in _member_ids(row)
async def test_team_bulk_member_delete_unknown_query_param_is_400(proxy_client, prisma, scratch, world):
victim = scratch.tag("victim")
await create_scratch_team(prisma, scratch.prefix, organization_id=world.org_a_id, member_user_ids=[victim])
async def test_team_bulk_member_delete_unknown_query_param_is_400(gateway_client, db, resources, actors):
victim = resources.tag("victim")
await create_scratch_team(db, resources.prefix, organization_id=actors.org_a_id, member_user_ids=[victim])
resp = await proxy_client.post(
f"/management/v1/teams/{scratch.prefix}/members/bulk_delete?dry_run=1",
headers={"Authorization": f"Bearer {world.keys[Actor.PROXY_ADMIN].cleartext}"},
resp = await gateway_client.post(
f"/management/v1/teams/{resources.prefix}/members/bulk_delete?dry_run=1",
headers={"Authorization": f"Bearer {actors.keys[Actor.PROXY_ADMIN].cleartext}"},
json={"members": [{"user_id": victim}]},
)
assert resp.status_code == 400, resp.text
assert resp.headers["content-type"] == "application/problem+json"
assert "dry_run" in resp.json()["detail"]
row = await prisma.db.litellm_teamtable.find_unique(where={"team_id": scratch.prefix})
row = await db.db.litellm_teamtable.find_unique(where={"team_id": resources.prefix})
assert row is not None and victim in _member_ids(row)
async def test_team_bulk_member_delete_unknown_body_field_is_422(proxy_client, prisma, scratch, world):
victim = scratch.tag("victim")
await create_scratch_team(prisma, scratch.prefix, organization_id=world.org_a_id, member_user_ids=[victim])
async def test_team_bulk_member_delete_unknown_body_field_is_422(gateway_client, db, resources, actors):
victim = resources.tag("victim")
await create_scratch_team(db, resources.prefix, organization_id=actors.org_a_id, member_user_ids=[victim])
resp = await proxy_client.post(
f"/management/v1/teams/{scratch.prefix}/members/bulk_delete",
headers={"Authorization": f"Bearer {world.keys[Actor.PROXY_ADMIN].cleartext}"},
json={"team_id": scratch.prefix, "members": [{"user_id": victim}]},
resp = await gateway_client.post(
f"/management/v1/teams/{resources.prefix}/members/bulk_delete",
headers={"Authorization": f"Bearer {actors.keys[Actor.PROXY_ADMIN].cleartext}"},
json={"team_id": resources.prefix, "members": [{"user_id": victim}]},
)
assert resp.status_code == 422, resp.text
assert "team_id" in resp.json()["detail"]
row = await prisma.db.litellm_teamtable.find_unique(where={"team_id": scratch.prefix})
row = await db.db.litellm_teamtable.find_unique(where={"team_id": resources.prefix})
assert row is not None and victim in _member_ids(row)
async def test_team_bulk_member_delete_unknown_team_is_404_problem(proxy_client, scratch, world):
resp = await proxy_client.post(
f"/management/v1/teams/{scratch.tag('missing')}/members/bulk_delete",
headers={"Authorization": f"Bearer {world.keys[Actor.PROXY_ADMIN].cleartext}"},
json={"members": [{"user_id": scratch.tag("victim")}]},
async def test_team_bulk_member_delete_unknown_team_is_404_problem(gateway_client, resources, actors):
resp = await gateway_client.post(
f"/management/v1/teams/{resources.tag('missing')}/members/bulk_delete",
headers={"Authorization": f"Bearer {actors.keys[Actor.PROXY_ADMIN].cleartext}"},
json={"members": [{"user_id": resources.tag("victim")}]},
)
assert resp.status_code == 404, resp.text
assert resp.headers["content-type"] == "application/problem+json"

32
tests/proxy_behavior/management/test_team_daily_activity.py → tests/integration/management/test_team_daily_activity.py
View file

@ -1,6 +1,9 @@
from typing import Final
import pytest
from .actors import Actor
from tests.integration._support.actors import Actor
from tests.integration._support.database import write_rows
pytestmark = pytest.mark.asyncio(loop_scope="session")
@ -41,6 +44,21 @@ _CASES = [
_DATES = "start_date=2024-01-01&end_date=2024-12-31"
_SEEDED_SPEND: Final = {"alpha": 1.25, "beta": 2.5}
@pytest.fixture(scope="module")
def seeded_team_spend(actors) -> None:
team_ids: Final = {"alpha": actors.team_alpha_id, "beta": actors.team_beta_id}
for team, spend in _SEEDED_SPEND.items():
write_rows(
'INSERT INTO "LiteLLM_DailyTeamSpend" (id, team_id, date, api_key, model, spend, updated_at) '
"VALUES (%s, %s, '2024-06-01', 'integration-daily-activity', 'gpt-4o', %s, now()) "
"ON CONFLICT (id) DO NOTHING",
(f"integration-daily-activity-{team}", team_ids[team], str(spend)),
)
@pytest.mark.parametrize(
"endpoint",
("/team/daily/activity", "/team/daily/activity/aggregated"),
@ -52,18 +70,20 @@ _DATES = "start_date=2024-01-01&end_date=2024-12-31"
ids=[c[0] for c in _CASES],
)
async def test_team_daily_activity_matrix(
actor: Actor, team: str, expected_status: int, endpoint: str, proxy_client, world
actor: Actor, team: str, expected_status: int, endpoint: str, gateway_client, actors, seeded_team_spend
):
query = _DATES
if team == "alpha":
query += f"&team_ids={world.team_alpha_id}"
query += f"&team_ids={actors.team_alpha_id}"
elif team == "beta":
query += f"&team_ids={world.team_beta_id}"
query += f"&team_ids={actors.team_beta_id}"
resp = await proxy_client.get(
resp = await gateway_client.get(
f"{endpoint}?{query}",
headers={"Authorization": f"Bearer {world.keys[actor].cleartext}"},
headers={"Authorization": f"Bearer {actors.keys[actor].cleartext}"},
)
assert (
resp.status_code == expected_status
), f"{actor.value} -> {team}: {resp.status_code} {resp.text}"
if actor == Actor.PROXY_ADMIN and team in _SEEDED_SPEND:
assert resp.json()["metadata"]["total_spend"] == _SEEDED_SPEND[team], resp.text

40
tests/proxy_behavior/management/test_team_delete.py → tests/integration/management/test_team_delete.py
View file

@ -1,7 +1,7 @@
import pytest
from .actors import Actor
from .conftest import create_scratch_team
from tests.integration._support.actors import Actor
from tests.integration._support.actors import create_scratch_team
pytestmark = pytest.mark.asyncio(loop_scope="session")
@ -33,26 +33,26 @@ async def test_team_delete_authz_matrix(
actor: Actor,
shape: str,
expected_status: int,
proxy_client,
prisma,
scratch,
world,
gateway_client,
db,
resources,
actors,
):
org_id = world.org_a_id if shape == "alpha" else world.org_b_id
await create_scratch_team(prisma, scratch.prefix, organization_id=org_id)
caller = world.keys[actor]
org_id = actors.org_a_id if shape == "alpha" else actors.org_b_id
await create_scratch_team(db, resources.prefix, organization_id=org_id)
caller = actors.keys[actor]
resp = await proxy_client.post(
resp = await gateway_client.post(
"/team/delete",
headers={"Authorization": f"Bearer {caller.cleartext}"},
json={"team_ids": [scratch.prefix], "organization_id": org_id},
json={"team_ids": [resources.prefix], "organization_id": org_id},
)
assert (
resp.status_code == expected_status
), f"{actor.value} {shape}: {resp.status_code} {resp.text}"
row = await prisma.db.litellm_teamtable.find_unique(
where={"team_id": scratch.prefix}
row = await db.db.litellm_teamtable.find_unique(
where={"team_id": resources.prefix}
)
if expected_status == 200:
assert row is None, "deleted but team row survives"
@ -61,18 +61,18 @@ async def test_team_delete_authz_matrix(
async def test_team_delete_batch_with_missing_id_deletes_nothing(
proxy_client, prisma, scratch, world
gateway_client, db, resources, actors
):
"""A batch is validated whole before any deletion: one missing team_id
fails the request 404 and the accessible team in the batch survives."""
await create_scratch_team(prisma, scratch.prefix, organization_id=world.org_a_id)
resp = await proxy_client.post(
await create_scratch_team(db, resources.prefix, organization_id=actors.org_a_id)
resp = await gateway_client.post(
"/team/delete",
headers={"Authorization": f"Bearer {world.keys[Actor.PROXY_ADMIN].cleartext}"},
json={"team_ids": [scratch.prefix, "behavior-pin-no-such-team"]},
headers={"Authorization": f"Bearer {actors.keys[Actor.PROXY_ADMIN].cleartext}"},
json={"team_ids": [resources.prefix, "behavior-pin-no-such-team"]},
)
assert resp.status_code == 404, resp.text
row = await prisma.db.litellm_teamtable.find_unique(
where={"team_id": scratch.prefix}
row = await db.db.litellm_teamtable.find_unique(
where={"team_id": resources.prefix}
)
assert row is not None, "batch aborted but the accessible team was deleted"

22
tests/proxy_behavior/management/test_team_filter_ui.py → tests/integration/management/test_team_filter_ui.py
View file

@ -1,6 +1,6 @@
import pytest
from .actors import Actor
from tests.integration._support.actors import Actor
pytestmark = pytest.mark.asyncio(loop_scope="session")
@ -12,10 +12,10 @@ pytestmark = pytest.mark.asyncio(loop_scope="session")
# route) and 401s every non-proxy-admin before the handler runs. PROXY_ADMIN
# reaches the unscoped find_many and sees teams across every org.
@pytest.mark.parametrize("actor", list(Actor), ids=[a.value for a in Actor])
async def test_team_filter_ui_is_proxy_admin_only(actor: Actor, proxy_client, world):
resp = await proxy_client.get(
async def test_team_filter_ui_is_proxy_admin_only(actor: Actor, gateway_client, actors):
resp = await gateway_client.get(
"/team/filter/ui",
headers={"Authorization": f"Bearer {world.keys[actor].cleartext}"},
headers={"Authorization": f"Bearer {actors.keys[actor].cleartext}"},
)
expected = 200 if actor == Actor.PROXY_ADMIN else 401
assert (
@ -23,17 +23,17 @@ async def test_team_filter_ui_is_proxy_admin_only(actor: Actor, proxy_client, wo
), f"{actor.value}: {resp.status_code} {resp.text}"
async def test_team_filter_ui_proxy_admin_sees_cross_org_teams(proxy_client, world):
async def test_team_filter_ui_proxy_admin_sees_cross_org_teams(gateway_client, actors):
"""The handler runs an unscoped query — PROXY_ADMIN sees teams from every
org, including the three seeded world teams."""
resp = await proxy_client.get(
org, including the three seeded actors teams."""
resp = await gateway_client.get(
"/team/filter/ui",
headers={"Authorization": f"Bearer {world.keys[Actor.PROXY_ADMIN].cleartext}"},
headers={"Authorization": f"Bearer {actors.keys[Actor.PROXY_ADMIN].cleartext}"},
)
assert resp.status_code == 200, resp.text
team_ids = {t.get("team_id") for t in resp.json() if isinstance(t, dict)}
assert {
world.team_alpha_id,
world.team_beta_id,
world.team_gamma_id,
actors.team_alpha_id,
actors.team_beta_id,
actors.team_gamma_id,
} <= team_ids

22
tests/proxy_behavior/management/test_team_info.py → tests/integration/management/test_team_info.py
View file

@ -1,6 +1,6 @@
import pytest
from .actors import Actor
from tests.integration._support.actors import Actor
pytestmark = pytest.mark.asyncio(loop_scope="session")
@ -47,16 +47,16 @@ _SCENARIOS = [
ids=[s[0] for s in _SCENARIOS],
)
async def test_team_info_authz_matrix(
actor: Actor, target: str, expected_status: int, proxy_client, world
actor: Actor, target: str, expected_status: int, gateway_client, actors
):
caller = world.keys[actor]
caller = actors.keys[actor]
target_team_id = {
"alpha": world.team_alpha_id,
"gamma": world.team_gamma_id,
"beta": world.team_beta_id,
"alpha": actors.team_alpha_id,
"gamma": actors.team_gamma_id,
"beta": actors.team_beta_id,
}[target]
resp = await proxy_client.get(
resp = await gateway_client.get(
f"/team/info?team_id={target_team_id}",
headers={"Authorization": f"Bearer {caller.cleartext}"},
)
@ -74,10 +74,10 @@ async def test_team_info_authz_matrix(
# alpha/org_b_admin already covers the branch in the matrix; this guard
# turns a silent rename of the exception detail into a CI red, which is the
# behavior tripwire that the matrix's status-only assertion cannot catch.
async def test_team_info_org_admin_cross_org_rejection_detail(proxy_client, world):
resp = await proxy_client.get(
f"/team/info?team_id={world.team_alpha_id}",
headers={"Authorization": f"Bearer {world.keys[Actor.ORG_B_ADMIN].cleartext}"},
async def test_team_info_org_admin_cross_org_rejection_detail(gateway_client, actors):
resp = await gateway_client.get(
f"/team/info?team_id={actors.team_alpha_id}",
headers={"Authorization": f"Bearer {actors.keys[Actor.ORG_B_ADMIN].cleartext}"},
)
assert resp.status_code == 403, resp.text
# validate_membership() at GET /team/info raises with this exact phrase.

96
tests/proxy_behavior/management/test_team_key_bulk_update.py → tests/integration/management/test_team_key_bulk_update.py
View file

@ -5,8 +5,8 @@ import pytest
from litellm.proxy._types import KeyManagementRoutes
from litellm.proxy.utils import hash_token
from .actors import Actor
from .conftest import create_scratch_key, create_scratch_team
from tests.integration._support.actors import Actor
from tests.integration._support.actors import create_scratch_key, create_scratch_team
pytestmark = pytest.mark.asyncio(loop_scope="session")
@ -31,20 +31,20 @@ _MATRIX = [
]
async def _seed_team_key(prisma, proxy_client, prefix: str, world, shape: str) -> str:
"""Raw-seed the scratch team for `shape`, return a team key's cleartext."""
internal = world.keys[Actor.INTERNAL_USER].user_id
owner = world.keys[Actor.OWNER].user_id
async def _seed_team_key(db, gateway_client, prefix: str, actors, shape: str) -> str:
"""Raw-seed the resources team for `shape`, return a team key's cleartext."""
internal = actors.keys[Actor.INTERNAL_USER].user_id
owner = actors.keys[Actor.OWNER].user_id
if shape == "admin":
await create_scratch_team(
prisma, prefix, organization_id=world.org_a_id, admin_user_ids=[internal]
db, prefix, organization_id=actors.org_a_id, admin_user_ids=[internal]
)
key_owner = internal
elif shape == "member_allowed":
await create_scratch_team(
prisma,
db,
prefix,
organization_id=world.org_a_id,
organization_id=actors.org_a_id,
admin_user_ids=[owner],
member_user_ids=[internal],
team_member_permissions=[_KEY_UPDATE],
@ -52,9 +52,9 @@ async def _seed_team_key(prisma, proxy_client, prefix: str, world, shape: str) -
key_owner = owner
elif shape == "member_denied":
await create_scratch_team(
prisma,
db,
prefix,
organization_id=world.org_a_id,
organization_id=actors.org_a_id,
admin_user_ids=[owner],
member_user_ids=[internal],
team_member_permissions=[],
@ -62,14 +62,14 @@ async def _seed_team_key(prisma, proxy_client, prefix: str, world, shape: str) -
key_owner = owner
elif shape == "nonmember":
await create_scratch_team(
prisma, prefix, organization_id=world.org_a_id, admin_user_ids=[owner]
db, prefix, organization_id=actors.org_a_id, admin_user_ids=[owner]
)
key_owner = owner
else:
pytest.fail(f"unknown shape={shape}") # pragma: no cover
return await create_scratch_key(
proxy_client,
world.keys[Actor.PROXY_ADMIN].cleartext,
gateway_client,
actors.keys[Actor.PROXY_ADMIN].cleartext,
prefix,
user_id=key_owner,
team_id=prefix,
@ -85,20 +85,20 @@ async def test_team_key_bulk_update_authz_matrix(
actor: Actor,
shape: str,
expected_status: int,
proxy_client,
prisma,
scratch,
world,
gateway_client,
db,
resources,
actors,
):
key = await _seed_team_key(prisma, proxy_client, scratch.prefix, world, shape)
key = await _seed_team_key(db, gateway_client, resources.prefix, actors, shape)
hashed = hash_token(key)
caller = world.keys[actor]
caller = actors.keys[actor]
resp = await proxy_client.post(
resp = await gateway_client.post(
"/team/key/bulk_update",
headers={"Authorization": f"Bearer {caller.cleartext}"},
json={
"team_id": scratch.prefix,
"team_id": resources.prefix,
"key_ids": [key],
"update_fields": {"max_budget": _MARKER_BUDGET},
},
@ -107,7 +107,7 @@ async def test_team_key_bulk_update_authz_matrix(
resp.status_code == expected_status
), f"{actor.value} {shape}: {resp.status_code} {resp.text}"
row = await prisma.db.litellm_verificationtoken.find_unique(where={"token": hashed})
row = await db.db.litellm_verificationtoken.find_unique(where={"token": hashed})
assert row is not None
if expected_status == 200:
assert len(resp.json()["successful_updates"]) == 1
@ -117,12 +117,12 @@ async def test_team_key_bulk_update_authz_matrix(
async def test_team_key_bulk_update_requires_team_id(
proxy_client, prisma, scratch, world
gateway_client, db, resources, actors
):
"""An empty team_id is rejected 400."""
resp = await proxy_client.post(
resp = await gateway_client.post(
"/team/key/bulk_update",
headers={"Authorization": f"Bearer {world.keys[Actor.PROXY_ADMIN].cleartext}"},
headers={"Authorization": f"Bearer {actors.keys[Actor.PROXY_ADMIN].cleartext}"},
json={
"team_id": "",
"key_ids": ["sk-" + uuid.uuid4().hex],
@ -133,28 +133,28 @@ async def test_team_key_bulk_update_requires_team_id(
async def test_team_key_bulk_update_all_keys_in_team(
proxy_client, prisma, scratch, world
gateway_client, db, resources, actors
):
"""all_keys_in_team=True broadcasts the update to every key in the team."""
admin = world.keys[Actor.PROXY_ADMIN].cleartext
await create_scratch_team(prisma, scratch.prefix, organization_id=world.org_a_id)
admin = actors.keys[Actor.PROXY_ADMIN].cleartext
await create_scratch_team(db, resources.prefix, organization_id=actors.org_a_id)
keys = [
await create_scratch_key(
proxy_client,
gateway_client,
admin,
scratch.prefix,
user_id=world.keys[Actor.OWNER].user_id,
team_id=scratch.prefix,
key_alias=f"{scratch.prefix}-k{i}",
resources.prefix,
user_id=actors.keys[Actor.OWNER].user_id,
team_id=resources.prefix,
key_alias=f"{resources.prefix}-k{i}",
)
for i in range(2)
]
resp = await proxy_client.post(
resp = await gateway_client.post(
"/team/key/bulk_update",
headers={"Authorization": f"Bearer {admin}"},
json={
"team_id": scratch.prefix,
"team_id": resources.prefix,
"all_keys_in_team": True,
"update_fields": {"max_budget": _MARKER_BUDGET},
},
@ -162,22 +162,22 @@ async def test_team_key_bulk_update_all_keys_in_team(
assert resp.status_code == 200, resp.text
assert len(resp.json()["successful_updates"]) == 2
for key in keys:
row = await prisma.db.litellm_verificationtoken.find_unique(
row = await db.db.litellm_verificationtoken.find_unique(
where={"token": hash_token(key)}
)
assert row is not None and row.max_budget == _MARKER_BUDGET
async def test_team_key_bulk_update_no_keys_found_is_404(
proxy_client, prisma, scratch, world
gateway_client, db, resources, actors
):
"""all_keys_in_team=True on a team with no keys is a top-level 404."""
await create_scratch_team(prisma, scratch.prefix, organization_id=world.org_a_id)
resp = await proxy_client.post(
await create_scratch_team(db, resources.prefix, organization_id=actors.org_a_id)
resp = await gateway_client.post(
"/team/key/bulk_update",
headers={"Authorization": f"Bearer {world.keys[Actor.PROXY_ADMIN].cleartext}"},
headers={"Authorization": f"Bearer {actors.keys[Actor.PROXY_ADMIN].cleartext}"},
json={
"team_id": scratch.prefix,
"team_id": resources.prefix,
"all_keys_in_team": True,
"update_fields": {"max_budget": _MARKER_BUDGET},
},
@ -186,21 +186,21 @@ async def test_team_key_bulk_update_no_keys_found_is_404(
async def test_team_key_bulk_update_missing_key_is_isolated(
proxy_client, prisma, scratch, world
gateway_client, db, resources, actors
):
"""A key_id absent from the team lands in failed_updates; the batch still
returns 200 and the real key is updated."""
admin = world.keys[Actor.PROXY_ADMIN].cleartext
admin = actors.keys[Actor.PROXY_ADMIN].cleartext
real = await _seed_team_key(
prisma, proxy_client, scratch.prefix, world, "nonmember"
db, gateway_client, resources.prefix, actors, "nonmember"
)
missing = "sk-" + uuid.uuid4().hex
resp = await proxy_client.post(
resp = await gateway_client.post(
"/team/key/bulk_update",
headers={"Authorization": f"Bearer {admin}"},
json={
"team_id": scratch.prefix,
"team_id": resources.prefix,
"key_ids": [real, missing],
"update_fields": {"max_budget": _MARKER_BUDGET},
},
@ -211,7 +211,7 @@ async def test_team_key_bulk_update_missing_key_is_isolated(
assert len(body["successful_updates"]) == 1
assert len(body["failed_updates"]) == 1
row = await prisma.db.litellm_verificationtoken.find_unique(
row = await db.db.litellm_verificationtoken.find_unique(
where={"token": hash_token(real)}
)
assert row is not None and row.max_budget == _MARKER_BUDGET

28
tests/proxy_behavior/management/test_team_list.py → tests/integration/management/test_team_list.py
View file

@ -2,18 +2,18 @@ from typing import FrozenSet, Optional
import pytest
from .actors import Actor
from tests.integration._support.actors import Actor
pytestmark = pytest.mark.asyncio(loop_scope="session")
# The behavior DB may hold teams beyond the three seeded ones, so every
# assertion intersects the returned team_ids with the known seeded set.
def _seeded_visible(resp_json, world) -> set:
def _seeded_visible(resp_json, actors) -> set:
known = {
world.team_alpha_id: "alpha",
world.team_beta_id: "beta",
world.team_gamma_id: "gamma",
actors.team_alpha_id: "alpha",
actors.team_beta_id: "beta",
actors.team_gamma_id: "gamma",
}
return {
known[entry["team_id"]]
@ -47,11 +47,11 @@ async def test_team_list_bare_authz(
actor: Actor,
expected_status: int,
expected_visible: Optional[set],
proxy_client,
world,
gateway_client,
actors,
):
caller = world.keys[actor]
resp = await proxy_client.get(
caller = actors.keys[actor]
resp = await gateway_client.get(
"/team/list",
headers={"Authorization": f"Bearer {caller.cleartext}"},
)
@ -60,7 +60,7 @@ async def test_team_list_bare_authz(
), f"{actor.value}: {resp.status_code} {resp.text}"
if expected_status == 200:
visible = _seeded_visible(resp.json(), world)
visible = _seeded_visible(resp.json(), actors)
assert visible == expected_visible, (
f"{actor.value}: expected {sorted(expected_visible)}, "
f"got {sorted(visible)}"
@ -90,16 +90,16 @@ _OWN = {
ids=[a.value for a in _OWN],
)
async def test_team_list_own_query(
actor: Actor, expected_visible: FrozenSet[str], proxy_client, world
actor: Actor, expected_visible: FrozenSet[str], gateway_client, actors
):
caller = world.keys[actor]
resp = await proxy_client.get(
caller = actors.keys[actor]
resp = await gateway_client.get(
f"/team/list?user_id={caller.user_id}",
headers={"Authorization": f"Bearer {caller.cleartext}"},
)
assert resp.status_code == 200, f"{actor.value}: {resp.status_code} {resp.text}"
visible = _seeded_visible(resp.json(), world)
visible = _seeded_visible(resp.json(), actors)
assert visible == set(expected_visible), (
f"{actor.value}: expected {sorted(expected_visible)}, " f"got {sorted(visible)}"
)

54
tests/proxy_behavior/management/test_team_list_v2.py → tests/integration/management/test_team_list_v2.py
View file

@ -2,26 +2,26 @@ from typing import FrozenSet, Optional
import pytest
from .actors import Actor
from tests.integration._support.actors import Actor
pytestmark = pytest.mark.asyncio(loop_scope="session")
def _seeded(team_ids: set, world) -> set:
def _seeded(team_ids: set, actors) -> set:
known = {
world.team_alpha_id: "alpha",
world.team_beta_id: "beta",
world.team_gamma_id: "gamma",
actors.team_alpha_id: "alpha",
actors.team_beta_id: "beta",
actors.team_gamma_id: "gamma",
}
return {known[t] for t in team_ids if t in known}
async def _v2_team_ids(proxy_client, caller_cleartext: str, extra: str = "") -> set:
async def _v2_team_ids(gateway_client, caller_cleartext: str, extra: str = "") -> set:
"""Walk every /v2/team/list page and collect the returned team_ids."""
ids: set = set()
page = 1
while True:
resp = await proxy_client.get(
resp = await gateway_client.get(
f"/v2/team/list?page={page}&page_size=100{extra}",
headers={"Authorization": f"Bearer {caller_cleartext}"},
)
@ -64,19 +64,19 @@ async def test_team_list_v2_bare(
actor: Actor,
expected_status: int,
expected_visible: Optional[FrozenSet[str]],
proxy_client,
world,
gateway_client,
actors,
):
caller = world.keys[actor]
caller = actors.keys[actor]
if expected_status != 200:
resp = await proxy_client.get(
resp = await gateway_client.get(
"/v2/team/list",
headers={"Authorization": f"Bearer {caller.cleartext}"},
)
assert resp.status_code == expected_status, resp.text
return
visible = _seeded(await _v2_team_ids(proxy_client, caller.cleartext), world)
visible = _seeded(await _v2_team_ids(gateway_client, caller.cleartext), actors)
assert visible == set(
expected_visible
), f"{actor.value}: expected {sorted(expected_visible)}, got {sorted(visible)}"
@ -98,44 +98,44 @@ _OWN = {
"actor,expected_visible", list(_OWN.items()), ids=[a.value for a in _OWN]
)
async def test_team_list_v2_own_user_id_query(
actor: Actor, expected_visible: FrozenSet[str], proxy_client, world
actor: Actor, expected_visible: FrozenSet[str], gateway_client, actors
):
caller = world.keys[actor]
caller = actors.keys[actor]
visible = _seeded(
await _v2_team_ids(
proxy_client, caller.cleartext, f"&user_id={caller.user_id}"
gateway_client, caller.cleartext, f"&user_id={caller.user_id}"
),
world,
actors,
)
assert visible == set(
expected_visible
), f"{actor.value}: expected {sorted(expected_visible)}, got {sorted(visible)}"
async def test_team_list_v2_user_id_filter_other_user_is_401(proxy_client, world):
async def test_team_list_v2_user_id_filter_other_user_is_401(gateway_client, actors):
"""A regular user filtering by another user's user_id is rejected 401."""
resp = await proxy_client.get(
f"/v2/team/list?user_id={world.keys[Actor.OWNER].user_id}",
resp = await gateway_client.get(
f"/v2/team/list?user_id={actors.keys[Actor.OWNER].user_id}",
headers={
"Authorization": f"Bearer {world.keys[Actor.INTERNAL_USER].cleartext}"
"Authorization": f"Bearer {actors.keys[Actor.INTERNAL_USER].cleartext}"
},
)
assert resp.status_code == 401, resp.text
async def test_team_list_v2_org_filter_foreign_org_is_403(proxy_client, world):
async def test_team_list_v2_org_filter_foreign_org_is_403(gateway_client, actors):
"""An org admin filtering by an organization it does not administer is 403."""
resp = await proxy_client.get(
f"/v2/team/list?organization_id={world.org_b_id}",
headers={"Authorization": f"Bearer {world.keys[Actor.ORG_ADMIN].cleartext}"},
resp = await gateway_client.get(
f"/v2/team/list?organization_id={actors.org_b_id}",
headers={"Authorization": f"Bearer {actors.keys[Actor.ORG_ADMIN].cleartext}"},
)
assert resp.status_code == 403, resp.text
async def test_team_list_v2_invalid_status_is_400(proxy_client, world):
async def test_team_list_v2_invalid_status_is_400(gateway_client, actors):
"""status accepts only 'deleted' — any other value is 400."""
resp = await proxy_client.get(
resp = await gateway_client.get(
"/v2/team/list?status=bogus",
headers={"Authorization": f"Bearer {world.keys[Actor.PROXY_ADMIN].cleartext}"},
headers={"Authorization": f"Bearer {actors.keys[Actor.PROXY_ADMIN].cleartext}"},
)
assert resp.status_code == 400, resp.text

135
tests/proxy_behavior/management/test_team_member_add.py → tests/integration/management/test_team_member_add.py
View file

@ -1,8 +1,7 @@
import litellm
import pytest
from .actors import Actor
from .conftest import create_scratch_team, create_scratch_user
from tests.integration._support.actors import Actor
from tests.integration._support.actors import create_scratch_team, create_scratch_user, owned_settings_client
pytestmark = pytest.mark.asyncio(loop_scope="session")
@ -40,16 +39,16 @@ _MATRIX = [
]
async def _seed_target(prisma, world, shape: str, team_id: str) -> None:
async def _seed_target(db, actors, shape: str, team_id: str) -> None:
if shape == "alpha":
await create_scratch_team(
prisma,
db,
team_id,
organization_id=world.org_a_id,
admin_user_ids=[world.keys[Actor.TEAM_ADMIN].user_id],
organization_id=actors.org_a_id,
admin_user_ids=[actors.keys[Actor.TEAM_ADMIN].user_id],
)
elif shape == "beta":
await create_scratch_team(prisma, team_id, organization_id=world.org_b_id)
await create_scratch_team(db, team_id, organization_id=actors.org_b_id)
else: # pragma: no cover - guard
pytest.fail(f"unknown shape={shape}")
@ -67,22 +66,22 @@ async def test_team_member_add_authz_matrix(
actor: Actor,
shape: str,
expected_status: int,
proxy_client,
prisma,
scratch,
world,
gateway_client,
db,
resources,
actors,
):
await _seed_target(prisma, world, shape, scratch.prefix)
caller = world.keys[actor]
await _seed_target(db, actors, shape, resources.prefix)
caller = actors.keys[actor]
new_member_id = await create_scratch_user(
prisma, scratch.prefix, suffix="newmember"
db, resources.prefix, suffix="newmember"
)
resp = await proxy_client.post(
resp = await gateway_client.post(
"/team/member_add",
headers={"Authorization": f"Bearer {caller.cleartext}"},
json={
"team_id": scratch.prefix,
"team_id": resources.prefix,
"member": {"user_id": new_member_id, "role": "user"},
},
)
@ -90,8 +89,8 @@ async def test_team_member_add_authz_matrix(
resp.status_code == expected_status
), f"{actor.value} {shape}: {resp.status_code} {resp.text}"
row = await prisma.db.litellm_teamtable.find_unique(
where={"team_id": scratch.prefix}
row = await db.db.litellm_teamtable.find_unique(
where={"team_id": resources.prefix}
)
assert row is not None
if expected_status == 200:
@ -121,28 +120,28 @@ _UNPROVISIONED = [
async def test_team_member_add_unprovisioned_user_id_is_proxy_admin_only(
actor: Actor,
expected_status: int,
proxy_client,
prisma,
scratch,
world,
gateway_client,
db,
resources,
actors,
):
await _seed_target(prisma, world, "alpha", scratch.prefix)
caller = world.keys[actor]
await _seed_target(db, actors, "alpha", resources.prefix)
caller = actors.keys[actor]
# Deliberately NOT seeded — create_scratch_user is what the matrix above
# calls, and its absence here is the whole scenario.
unprovisioned_id = scratch.tag("unprovisioned")
unprovisioned_id = resources.tag("unprovisioned")
assert (
await prisma.db.litellm_usertable.find_unique(
await db.db.litellm_usertable.find_unique(
where={"user_id": unprovisioned_id}
)
is None
), "setup: user must not exist yet"
resp = await proxy_client.post(
resp = await gateway_client.post(
"/team/member_add",
headers={"Authorization": f"Bearer {caller.cleartext}"},
json={
"team_id": scratch.prefix,
"team_id": resources.prefix,
"member": {"user_id": unprovisioned_id, "role": "user"},
},
)
@ -150,11 +149,11 @@ async def test_team_member_add_unprovisioned_user_id_is_proxy_admin_only(
resp.status_code == expected_status
), f"{actor.value}: {resp.status_code} {resp.text}"
row = await prisma.db.litellm_teamtable.find_unique(
where={"team_id": scratch.prefix}
row = await db.db.litellm_teamtable.find_unique(
where={"team_id": resources.prefix}
)
assert row is not None
created = await prisma.db.litellm_usertable.find_unique(
created = await db.db.litellm_usertable.find_unique(
where={"user_id": unprovisioned_id}
)
if expected_status == 200:
@ -167,10 +166,10 @@ async def test_team_member_add_unprovisioned_user_id_is_proxy_admin_only(
async def test_team_member_add_email_invite_open_to_team_admin(
proxy_client,
prisma,
scratch,
world,
gateway_client,
db,
resources,
actors,
):
"""The escape hatch the provisioning 403 names must actually work.
@ -179,34 +178,31 @@ async def test_team_member_add_email_invite_open_to_team_admin(
it stays open to team admins; if it ever closed, the 403 would be sending
callers down a dead end.
"""
await _seed_target(prisma, world, "alpha", scratch.prefix)
caller = world.keys[Actor.TEAM_ADMIN]
email = f"{scratch.prefix}-invitee@example.com"
await _seed_target(db, actors, "alpha", resources.prefix)
caller = actors.keys[Actor.TEAM_ADMIN]
email = f"{resources.prefix}-invitee@example.com"
resp = await proxy_client.post(
resp = await gateway_client.post(
"/team/member_add",
headers={"Authorization": f"Bearer {caller.cleartext}"},
json={
"team_id": scratch.prefix,
"team_id": resources.prefix,
"member": {"user_email": email, "role": "user"},
},
)
assert resp.status_code == 200, resp.text
invited = await prisma.db.litellm_usertable.find_first(
invited = await db.db.litellm_usertable.find_first(
where={"user_email": email}
)
assert invited is not None, "invite did not create the user"
row = await prisma.db.litellm_teamtable.find_unique(
where={"team_id": scratch.prefix}
row = await db.db.litellm_teamtable.find_unique(
where={"team_id": resources.prefix}
)
assert row is not None
assert invited.user_id in _member_ids(row)
# Available-team self-join: a non-admin caller may add ITSELF to a team listed
# in litellm.default_internal_user_params["available_teams"], but the bypass
# must not escalate to role=admin or inject another user.
_SELF_JOIN = [
("self_as_user", "self", "user", 200),
("self_as_admin", "self", "admin", 403),
@ -223,36 +219,33 @@ async def test_team_member_add_available_team_self_join(
who: str,
role: str,
expected_status: int,
proxy_client,
prisma,
scratch,
world,
monkeypatch,
gateway_client,
db,
resources,
actors,
tmp_path,
):
# Org-less team with no admins: the INTERNAL_USER caller is neither team
# nor org admin, so it lands on the available-team branch.
await create_scratch_team(prisma, scratch.prefix)
monkeypatch.setattr(
litellm, "default_internal_user_params", {"available_teams": [scratch.prefix]}
)
await create_scratch_team(db, resources.prefix)
caller = actors.keys[Actor.INTERNAL_USER]
member_id = caller.user_id if who == "self" else actors.keys[Actor.OWNER].user_id
caller = world.keys[Actor.INTERNAL_USER]
member_id = caller.user_id if who == "self" else world.keys[Actor.OWNER].user_id
resp = await proxy_client.post(
"/team/member_add",
headers={"Authorization": f"Bearer {caller.cleartext}"},
json={
"team_id": scratch.prefix,
"member": {"user_id": member_id, "role": role},
},
)
with owned_settings_client(
gateway_client, tmp_path, {"default_internal_user_params": {"available_teams": [resources.prefix]}}
) as client:
resp = await client.post(
"/team/member_add",
headers={"Authorization": f"Bearer {caller.cleartext}"},
json={
"team_id": resources.prefix,
"member": {"user_id": member_id, "role": role},
},
)
assert (
resp.status_code == expected_status
), f"{who}/{role}: {resp.status_code} {resp.text}"
row = await prisma.db.litellm_teamtable.find_unique(
where={"team_id": scratch.prefix}
row = await db.db.litellm_teamtable.find_unique(
where={"team_id": resources.prefix}
)
assert row is not None
if expected_status == 200:

40
tests/proxy_behavior/management/test_team_member_delete.py → tests/integration/management/test_team_member_delete.py
View file

@ -1,12 +1,12 @@
import pytest
from .actors import Actor
from .conftest import create_scratch_team
from tests.integration._support.actors import Actor
from tests.integration._support.actors import create_scratch_team
pytestmark = pytest.mark.asyncio(loop_scope="session")
# POST /team/member_delete — actor x team-shape matrix. The scratch team is
# POST /team/member_delete — actor x team-shape matrix. The resources team is
# raw-seeded with a victim member already in it; PROXY_ADMIN, the team's team
# admin, or an org admin of the team's org may remove members; else 403.
_MATRIX = [
@ -31,20 +31,20 @@ _MATRIX = [
]
async def _seed_target(prisma, world, shape: str, team_id: str, victim_id: str) -> None:
async def _seed_target(db, actors, shape: str, team_id: str, victim_id: str) -> None:
if shape == "alpha":
await create_scratch_team(
prisma,
db,
team_id,
organization_id=world.org_a_id,
admin_user_ids=[world.keys[Actor.TEAM_ADMIN].user_id],
organization_id=actors.org_a_id,
admin_user_ids=[actors.keys[Actor.TEAM_ADMIN].user_id],
member_user_ids=[victim_id],
)
elif shape == "beta":
await create_scratch_team(
prisma,
db,
team_id,
organization_id=world.org_b_id,
organization_id=actors.org_b_id,
member_user_ids=[victim_id],
)
else: # pragma: no cover - guard
@ -64,26 +64,26 @@ async def test_team_member_delete_authz_matrix(
actor: Actor,
shape: str,
expected_status: int,
proxy_client,
prisma,
scratch,
world,
gateway_client,
db,
resources,
actors,
):
victim_id = scratch.tag("victim")
await _seed_target(prisma, world, shape, scratch.prefix, victim_id)
caller = world.keys[actor]
victim_id = resources.tag("victim")
await _seed_target(db, actors, shape, resources.prefix, victim_id)
caller = actors.keys[actor]
resp = await proxy_client.post(
resp = await gateway_client.post(
"/team/member_delete",
headers={"Authorization": f"Bearer {caller.cleartext}"},
json={"team_id": scratch.prefix, "user_id": victim_id},
json={"team_id": resources.prefix, "user_id": victim_id},
)
assert (
resp.status_code == expected_status
), f"{actor.value} {shape}: {resp.status_code} {resp.text}"
row = await prisma.db.litellm_teamtable.find_unique(
where={"team_id": scratch.prefix}
row = await db.db.litellm_teamtable.find_unique(
where={"team_id": resources.prefix}
)
assert row is not None
if expected_status == 200:

87
tests/proxy_behavior/management/test_team_member_info_validation.py → tests/integration/management/test_team_member_info_validation.py
View file

@ -13,8 +13,9 @@ user_id is invisible from response-body alone).
import pytest
from .actors import Actor
from .conftest import create_scratch_team, create_scratch_user
from tests.integration._support.actors import Actor
from tests.integration._support.actors import create_scratch_team, create_scratch_user
from tests.integration._support.database import write_rows
pytestmark = pytest.mark.asyncio(loop_scope="session")
@ -24,10 +25,10 @@ pytestmark = pytest.mark.asyncio(loop_scope="session")
# ---------------------------------------------------------------------------
async def test_member_add_both_none_rejected(proxy_client, prisma, scratch, world):
team_id = await create_scratch_team(prisma, team_id=scratch.tag("team"))
seeder = world.keys[Actor.PROXY_ADMIN].cleartext
resp = await proxy_client.post(
async def test_member_add_both_none_rejected(gateway_client, db, resources, actors):
team_id = await create_scratch_team(db, team_id=resources.tag("team"))
seeder = actors.keys[Actor.PROXY_ADMIN].cleartext
resp = await gateway_client.post(
"/team/member_add",
headers={"Authorization": f"Bearer {seeder}"},
json={"team_id": team_id, "member": {"role": "user"}},
@ -36,8 +37,8 @@ async def test_member_add_both_none_rejected(proxy_client, prisma, scratch, worl
# (422). Either shape proves the empty-Member payload is rejected before
# any membership row is written — pin both.
assert resp.status_code in (400, 422), resp.text
rows = await prisma.db.litellm_teammembership.find_many(where={"team_id": team_id})
assert rows == [], "empty-Member payload leaked a membership row"
rows = await db.db.litellm_teammembership.find_many(where={"team_id": team_id})
assert all(row.user_id == "default_user_id" for row in rows), "empty-Member payload leaked a membership row"
# ---------------------------------------------------------------------------
@ -46,17 +47,17 @@ async def test_member_add_both_none_rejected(proxy_client, prisma, scratch, worl
async def test_member_add_email_id_mismatch_rejected(
proxy_client, prisma, scratch, world
gateway_client, db, resources, actors
):
email = f"{scratch.prefix}-mismatch@example.com"
email = f"{resources.prefix}-mismatch@example.com"
real_user_id = await create_scratch_user(
prisma, scratch.prefix, suffix="real", user_email=email
db, resources.prefix, suffix="real", user_email=email
)
other_user_id = await create_scratch_user(prisma, scratch.prefix, suffix="other")
other_user_id = await create_scratch_user(db, resources.prefix, suffix="other")
assert real_user_id != other_user_id # sanity
team_id = await create_scratch_team(prisma, team_id=scratch.tag("team"))
seeder = world.keys[Actor.PROXY_ADMIN].cleartext
resp = await proxy_client.post(
team_id = await create_scratch_team(db, team_id=resources.tag("team"))
seeder = actors.keys[Actor.PROXY_ADMIN].cleartext
resp = await gateway_client.post(
"/team/member_add",
headers={"Authorization": f"Bearer {seeder}"},
json={
@ -70,8 +71,8 @@ async def test_member_add_email_id_mismatch_rejected(
)
assert resp.status_code == 400, resp.text
assert "do not belong to the same user" in resp.text, resp.text
rows = await prisma.db.litellm_teammembership.find_many(where={"team_id": team_id})
assert rows == [], "mismatch payload leaked a membership row"
rows = await db.db.litellm_teammembership.find_many(where={"team_id": team_id})
assert not any(row.user_id in {real_user_id, other_user_id} for row in rows), "mismatch payload leaked a membership row"
# ---------------------------------------------------------------------------
@ -80,15 +81,15 @@ async def test_member_add_email_id_mismatch_rejected(
async def test_member_add_email_only_resolves_user_id(
proxy_client, prisma, scratch, world
gateway_client, db, resources, actors
):
email = f"{scratch.prefix}-resolve@example.com"
email = f"{resources.prefix}-resolve@example.com"
user_id = await create_scratch_user(
prisma, scratch.prefix, suffix="lookup", user_email=email
db, resources.prefix, suffix="lookup", user_email=email
)
team_id = await create_scratch_team(prisma, team_id=scratch.tag("team"))
seeder = world.keys[Actor.PROXY_ADMIN].cleartext
resp = await proxy_client.post(
team_id = await create_scratch_team(db, team_id=resources.tag("team"))
seeder = actors.keys[Actor.PROXY_ADMIN].cleartext
resp = await gateway_client.post(
"/team/member_add",
headers={"Authorization": f"Bearer {seeder}"},
json={
@ -102,7 +103,7 @@ async def test_member_add_email_only_resolves_user_id(
# team's members_with_roles JSON. Re-read that and assert the resolved
# user_id landed — the regression shape is "email resolved to the WRONG
# user_id and was silently written to members_with_roles".
team_row = await prisma.db.litellm_teamtable.find_unique(where={"team_id": team_id})
team_row = await db.db.litellm_teamtable.find_unique(where={"team_id": team_id})
assert team_row is not None
member_user_ids = [m.get("user_id") for m in team_row.members_with_roles]
assert (
@ -116,15 +117,15 @@ async def test_member_add_email_only_resolves_user_id(
async def test_member_add_unknown_user_id_upserted(
proxy_client, prisma, scratch, world
gateway_client, db, resources, actors
):
team_id = await create_scratch_team(prisma, team_id=scratch.tag("team"))
new_user_id = f"{scratch.prefix}-fresh"
team_id = await create_scratch_team(db, team_id=resources.tag("team"))
new_user_id = f"{resources.prefix}-fresh"
# Sanity — user does not exist yet.
pre = await prisma.db.litellm_usertable.find_unique(where={"user_id": new_user_id})
pre = await db.db.litellm_usertable.find_unique(where={"user_id": new_user_id})
assert pre is None
seeder = world.keys[Actor.PROXY_ADMIN].cleartext
resp = await proxy_client.post(
seeder = actors.keys[Actor.PROXY_ADMIN].cleartext
resp = await gateway_client.post(
"/team/member_add",
headers={"Authorization": f"Bearer {seeder}"},
json={
@ -133,7 +134,7 @@ async def test_member_add_unknown_user_id_upserted(
},
)
assert resp.status_code == 200, resp.text
post = await prisma.db.litellm_usertable.find_unique(where={"user_id": new_user_id})
post = await db.db.litellm_usertable.find_unique(where={"user_id": new_user_id})
assert post is not None, "user_id was not upserted"
# The user row was created with NULL email (the helper returned the
# member as-is, no email lookup happened because the user didn't exist).
@ -143,20 +144,24 @@ async def test_member_add_unknown_user_id_upserted(
# ---------------------------------------------------------------------------
# Duplicate-email rejection — two scratch users share an email; email-only
# Duplicate-email rejection — two resources users share an email; email-only
# add → 400 with "Multiple users found" detail.
# ---------------------------------------------------------------------------
async def test_member_add_duplicate_email_rejected(
proxy_client, prisma, scratch, world
gateway_client, db, resources, actors
):
email = f"{scratch.prefix}-dup@example.com"
await create_scratch_user(prisma, scratch.prefix, suffix="dup1", user_email=email)
await create_scratch_user(prisma, scratch.prefix, suffix="dup2", user_email=email)
team_id = await create_scratch_team(prisma, team_id=scratch.tag("team"))
seeder = world.keys[Actor.PROXY_ADMIN].cleartext
resp = await proxy_client.post(
email = f"{resources.prefix}-dup@example.com"
first = await create_scratch_user(db, resources.prefix, suffix="dup1", user_email=email)
write_rows(
'INSERT INTO "LiteLLM_UserTable" (user_id, user_email, user_role, models, metadata) '
"VALUES (%s, %s, 'internal_user', '{}', '{}')",
(resources.tag("dup2"), email),
)
team_id = await create_scratch_team(db, team_id=resources.tag("team"))
seeder = actors.keys[Actor.PROXY_ADMIN].cleartext
resp = await gateway_client.post(
"/team/member_add",
headers={"Authorization": f"Bearer {seeder}"},
json={
@ -166,5 +171,7 @@ async def test_member_add_duplicate_email_rejected(
)
assert resp.status_code == 400, resp.text
assert "Multiple users found" in resp.text, resp.text
rows = await prisma.db.litellm_teammembership.find_many(where={"team_id": team_id})
rows = await db.db.litellm_teammembership.find_many(
where={"team_id": team_id, "user_id": {"in": [first, resources.tag("dup2")]}}
)
assert rows == [], "duplicate-email payload leaked a membership row"

28
tests/proxy_behavior/management/test_team_member_me.py → tests/integration/management/test_team_member_me.py
View file

@ -4,7 +4,7 @@ import pytest
from litellm.proxy.utils import hash_token
from .actors import Actor
from tests.integration._support.actors import Actor
pytestmark = pytest.mark.asyncio(loop_scope="session")
@ -38,16 +38,16 @@ _CASES = [
ids=[c[0] for c in _CASES],
)
async def test_team_member_me_matrix(
actor: Actor, team: str, expected_status: int, proxy_client, world
actor: Actor, team: str, expected_status: int, gateway_client, actors
):
team_id = {
"alpha": world.team_alpha_id,
"beta": world.team_beta_id,
"gamma": world.team_gamma_id,
"alpha": actors.team_alpha_id,
"beta": actors.team_beta_id,
"gamma": actors.team_gamma_id,
}[team]
caller = world.keys[actor]
caller = actors.keys[actor]
resp = await proxy_client.get(
resp = await gateway_client.get(
f"/team/{team_id}/members/me",
headers={"Authorization": f"Bearer {caller.cleartext}"},
)
@ -62,22 +62,22 @@ async def test_team_member_me_matrix(
async def test_team_member_me_team_key_without_user_id_is_400(
proxy_client, prisma, scratch, world
gateway_client, db, resources, actors
):
"""A key with no associated user_id (a team / service-account key) cannot
resolve 'me' — the caller has no identity to look up — so it is 400."""
cleartext = "sk-" + uuid.uuid4().hex
await prisma.db.litellm_verificationtoken.create(
await db.db.litellm_verificationtoken.create(
data={
"token": hash_token(cleartext),
"key_name": f"{scratch.prefix}-teamkey",
"key_alias": f"{scratch.prefix}-teamkey",
"team_id": world.team_alpha_id,
"key_name": f"{resources.prefix}-teamkey",
"key_alias": f"{resources.prefix}-teamkey",
"team_id": actors.team_alpha_id,
"models": [],
}
)
resp = await proxy_client.get(
f"/team/{world.team_alpha_id}/members/me",
resp = await gateway_client.get(
f"/team/{actors.team_alpha_id}/members/me",
headers={"Authorization": f"Bearer {cleartext}"},
)
assert resp.status_code == 400, resp.text

200
tests/integration/management/test_team_member_reset_budget.py Normal file
View file

@ -0,0 +1,200 @@
import uuid
import pytest
from prisma import Json
from tests.integration._support.actors import Actor
from tests.integration._support.actors import create_scratch_team
pytestmark = pytest.mark.asyncio(loop_scope="session")
_SEED_SPEND = 5.0
_TEAM_DEFAULT_MAX_BUDGET = 100.0
_CUSTOM_MAX_BUDGET = 50.0
_MATRIX = [
("alpha/proxy_admin", Actor.PROXY_ADMIN, "alpha", 200),
("alpha/org_admin", Actor.ORG_ADMIN, "alpha", 200),
("alpha/team_admin", Actor.TEAM_ADMIN, "alpha", 200),
("alpha/internal_user", Actor.INTERNAL_USER, "alpha", 403),
("alpha/owner", Actor.OWNER, "alpha", 403),
("alpha/unrelated_same_org", Actor.UNRELATED_SAME_ORG, "alpha", 403),
("alpha/cross_org_user", Actor.CROSS_ORG_USER, "alpha", 403),
("alpha/service_account", Actor.SERVICE_ACCOUNT, "alpha", 403),
("alpha/org_b_admin", Actor.ORG_B_ADMIN, "alpha", 403),
("beta/proxy_admin", Actor.PROXY_ADMIN, "beta", 200),
("beta/org_admin", Actor.ORG_ADMIN, "beta", 403),
("beta/team_admin", Actor.TEAM_ADMIN, "beta", 403),
("beta/org_b_admin", Actor.ORG_B_ADMIN, "beta", 200),
]
async def _seed_budget(db, budget_id: str, max_budget: float) -> str:
await db.db.litellm_budgettable.create(
data={
"budget_id": budget_id,
"max_budget": max_budget,
"created_by": "phase4-resources",
"updated_by": "phase4-resources",
}
)
return budget_id
async def _seed_team_with_default_budget(db, actors, shape: str, team_id: str, resources) -> str:
default_budget_id = await _seed_budget(db, resources.tag("team-default-budget"), _TEAM_DEFAULT_MAX_BUDGET)
if shape == "alpha":
await create_scratch_team(
db,
team_id,
organization_id=actors.org_a_id,
admin_user_ids=[actors.keys[Actor.TEAM_ADMIN].user_id],
)
elif shape == "beta":
await create_scratch_team(db, team_id, organization_id=actors.org_b_id)
else:
pytest.fail(f"unknown shape={shape}")
await db.db.litellm_teamtable.update(
where={"team_id": team_id},
data={"metadata": Json({"team_member_budget_id": default_budget_id})},
)
return default_budget_id
async def _seed_custom_member(db, team_id: str, member_id: str, resources) -> str:
custom_budget_id = await _seed_budget(db, resources.tag("custom-budget"), _CUSTOM_MAX_BUDGET)
await db.db.litellm_teammembership.create(
data={
"user_id": member_id,
"team_id": team_id,
"spend": _SEED_SPEND,
"litellm_budget_table": {"connect": {"budget_id": custom_budget_id}},
}
)
return custom_budget_id
async def _membership(db, team_id: str, member_id: str):
row = await db.db.litellm_teammembership.find_unique(
where={"user_id_team_id": {"user_id": member_id, "team_id": team_id}}
)
assert row is not None
return row
@pytest.mark.parametrize(
"actor,shape,expected_status",
[(a, sh, s) for (_id, a, sh, s) in _MATRIX],
ids=[s[0] for s in _MATRIX],
)
async def test_team_member_reset_budget_authz_matrix(
actor: Actor,
shape: str,
expected_status: int,
gateway_client,
db,
resources,
actors,
):
member_id = resources.tag("member")
default_budget_id = await _seed_team_with_default_budget(db, actors, shape, resources.prefix, resources)
custom_budget_id = await _seed_custom_member(db, resources.prefix, member_id, resources)
caller = actors.keys[actor]
resp = await gateway_client.post(
f"/team/{resources.prefix}/member/{member_id}/reset_budget",
headers={"Authorization": f"Bearer {caller.cleartext}"},
)
assert resp.status_code == expected_status, f"{actor.value} {shape}: {resp.status_code} {resp.text}"
row = await _membership(db, resources.prefix, member_id)
assert row.spend == _SEED_SPEND, "reset_budget must never touch spend"
if expected_status == 200:
assert row.budget_id == default_budget_id
body = resp.json()
assert body["budget_id"] == default_budget_id
assert body["previous_budget_id"] == custom_budget_id
assert body["budget_source"] == "team_default"
else:
assert row.budget_id == custom_budget_id, "denied but budget relinked"
async def test_team_member_reset_budget_leaves_shared_default_row_untouched(gateway_client, db, resources, actors):
member_id = resources.tag("member")
default_budget_id = await _seed_team_with_default_budget(db, actors, "alpha", resources.prefix, resources)
await _seed_custom_member(db, resources.prefix, member_id, resources)
resp = await gateway_client.post(
f"/team/{resources.prefix}/member/{member_id}/reset_budget",
headers={"Authorization": f"Bearer {actors.keys[Actor.PROXY_ADMIN].cleartext}"},
)
assert resp.status_code == 200, resp.text
default_row = await db.db.litellm_budgettable.find_unique(where={"budget_id": default_budget_id})
assert default_row is not None and default_row.max_budget == _TEAM_DEFAULT_MAX_BUDGET
info = await gateway_client.get(
f"/team/info?team_id={resources.prefix}",
headers={"Authorization": f"Bearer {actors.keys[Actor.PROXY_ADMIN].cleartext}"},
)
assert info.status_code == 200, info.text
memberships = {tm["user_id"]: tm for tm in info.json()["team_memberships"]}
assert memberships[member_id]["budget_source"] == "team_default"
assert memberships[member_id]["litellm_budget_table"]["max_budget"] == _TEAM_DEFAULT_MAX_BUDGET
async def test_team_member_reset_budget_without_team_default_detaches_member(gateway_client, db, resources, actors):
member_id = resources.tag("member")
await create_scratch_team(db, resources.prefix, organization_id=actors.org_a_id)
await _seed_custom_member(db, resources.prefix, member_id, resources)
resp = await gateway_client.post(
f"/team/{resources.prefix}/member/{member_id}/reset_budget",
headers={"Authorization": f"Bearer {actors.keys[Actor.PROXY_ADMIN].cleartext}"},
)
assert resp.status_code == 200, resp.text
assert resp.json()["budget_id"] is None
assert resp.json()["budget_source"] == "none"
row = await _membership(db, resources.prefix, member_id)
assert row.budget_id is None
assert row.spend == _SEED_SPEND
async def test_team_member_reset_budget_with_deleted_team_default_detaches_member(gateway_client, db, resources, actors):
member_id = resources.tag("member")
await create_scratch_team(
db,
resources.prefix,
organization_id=actors.org_a_id,
metadata={"team_member_budget_id": resources.tag("deleted-budget")},
)
await _seed_custom_member(db, resources.prefix, member_id, resources)
resp = await gateway_client.post(
f"/team/{resources.prefix}/member/{member_id}/reset_budget",
headers={"Authorization": f"Bearer {actors.keys[Actor.PROXY_ADMIN].cleartext}"},
)
assert resp.status_code == 200, resp.text
assert resp.json()["budget_id"] is None
assert resp.json()["budget_source"] == "none"
row = await _membership(db, resources.prefix, member_id)
assert row.budget_id is None
async def test_team_member_reset_budget_missing_team_is_404(gateway_client, actors):
resp = await gateway_client.post(
f"/team/behavior-pin-no-such-team/member/{uuid.uuid4().hex}/reset_budget",
headers={"Authorization": f"Bearer {actors.keys[Actor.PROXY_ADMIN].cleartext}"},
)
assert resp.status_code == 404, resp.text
async def test_team_member_reset_budget_missing_membership_is_404(gateway_client, db, resources, actors):
await create_scratch_team(db, resources.prefix, organization_id=actors.org_a_id)
resp = await gateway_client.post(
f"/team/{resources.prefix}/member/{uuid.uuid4().hex}/reset_budget",
headers={"Authorization": f"Bearer {actors.keys[Actor.PROXY_ADMIN].cleartext}"},
)
assert resp.status_code == 404, resp.text

94
tests/proxy_behavior/management/test_team_member_reset_spend.py → tests/integration/management/test_team_member_reset_spend.py
View file

@ -2,8 +2,9 @@ import uuid
import pytest
from .actors import Actor
from .conftest import create_scratch_team
from tests.integration._support.actors import Actor
from tests.integration._support.actors import create_scratch_team
from tests.integration._support.database import write_rows
pytestmark = pytest.mark.asyncio(loop_scope="session")
@ -32,19 +33,19 @@ _MATRIX = [
]
async def _seed_target(prisma, world, shape: str, team_id: str, member_id: str) -> None:
async def _seed_target(db, actors, shape: str, team_id: str, member_id: str) -> None:
if shape == "alpha":
await create_scratch_team(
prisma,
db,
team_id,
organization_id=world.org_a_id,
admin_user_ids=[world.keys[Actor.TEAM_ADMIN].user_id],
organization_id=actors.org_a_id,
admin_user_ids=[actors.keys[Actor.TEAM_ADMIN].user_id],
)
elif shape == "beta":
await create_scratch_team(prisma, team_id, organization_id=world.org_b_id)
await create_scratch_team(db, team_id, organization_id=actors.org_b_id)
else: # pragma: no cover - guard
pytest.fail(f"unknown shape={shape}")
await prisma.db.litellm_teammembership.create(
await db.db.litellm_teammembership.create(
data={"user_id": member_id, "team_id": team_id, "spend": _SEED_SPEND}
)
@ -58,17 +59,17 @@ async def test_team_member_reset_spend_authz_matrix(
actor: Actor,
shape: str,
expected_status: int,
proxy_client,
prisma,
scratch,
world,
gateway_client,
db,
resources,
actors,
):
member_id = scratch.tag("member")
await _seed_target(prisma, world, shape, scratch.prefix, member_id)
caller = world.keys[actor]
member_id = resources.tag("member")
await _seed_target(db, actors, shape, resources.prefix, member_id)
caller = actors.keys[actor]
resp = await proxy_client.post(
f"/team/{scratch.prefix}/member/{member_id}/reset_spend",
resp = await gateway_client.post(
f"/team/{resources.prefix}/member/{member_id}/reset_spend",
headers={"Authorization": f"Bearer {caller.cleartext}"},
json={"reset_to": _RESET_TO},
)
@ -76,8 +77,8 @@ async def test_team_member_reset_spend_authz_matrix(
resp.status_code == expected_status
), f"{actor.value} {shape}: {resp.status_code} {resp.text}"
row = await prisma.db.litellm_teammembership.find_unique(
where={"user_id_team_id": {"user_id": member_id, "team_id": scratch.prefix}}
row = await db.db.litellm_teammembership.find_unique(
where={"user_id_team_id": {"user_id": member_id, "team_id": resources.prefix}}
)
assert row is not None
if expected_status == 200:
@ -86,67 +87,68 @@ async def test_team_member_reset_spend_authz_matrix(
assert row.spend == _SEED_SPEND, "denied but spend reset"
async def test_team_member_reset_spend_missing_team_is_404(proxy_client, world):
resp = await proxy_client.post(
async def test_team_member_reset_spend_missing_team_is_404(gateway_client, actors):
resp = await gateway_client.post(
f"/team/behavior-pin-no-such-team/member/{uuid.uuid4().hex}/reset_spend",
headers={"Authorization": f"Bearer {world.keys[Actor.PROXY_ADMIN].cleartext}"},
headers={"Authorization": f"Bearer {actors.keys[Actor.PROXY_ADMIN].cleartext}"},
json={"reset_to": 0.0},
)
assert resp.status_code == 404, resp.text
async def test_team_member_reset_spend_missing_membership_is_404(
proxy_client, prisma, scratch, world
gateway_client, db, resources, actors
):
"""A well-formed team but a user_id with no LiteLLM_TeamMembership row is 404."""
await create_scratch_team(prisma, scratch.prefix, organization_id=world.org_a_id)
resp = await proxy_client.post(
f"/team/{scratch.prefix}/member/{uuid.uuid4().hex}/reset_spend",
headers={"Authorization": f"Bearer {world.keys[Actor.PROXY_ADMIN].cleartext}"},
await create_scratch_team(db, resources.prefix, organization_id=actors.org_a_id)
resp = await gateway_client.post(
f"/team/{resources.prefix}/member/{uuid.uuid4().hex}/reset_spend",
headers={"Authorization": f"Bearer {actors.keys[Actor.PROXY_ADMIN].cleartext}"},
json={"reset_to": 0.0},
)
assert resp.status_code == 404, resp.text
async def test_team_member_reset_spend_above_current_spend_is_400(
proxy_client, prisma, scratch, world
gateway_client, db, resources, actors
):
member_id = scratch.tag("member")
await create_scratch_team(prisma, scratch.prefix, organization_id=world.org_a_id)
await prisma.db.litellm_teammembership.create(
data={"user_id": member_id, "team_id": scratch.prefix, "spend": 1.0}
member_id = resources.tag("member")
await create_scratch_team(db, resources.prefix, organization_id=actors.org_a_id)
await db.db.litellm_teammembership.create(
data={"user_id": member_id, "team_id": resources.prefix, "spend": 1.0}
)
resp = await proxy_client.post(
f"/team/{scratch.prefix}/member/{member_id}/reset_spend",
headers={"Authorization": f"Bearer {world.keys[Actor.PROXY_ADMIN].cleartext}"},
resp = await gateway_client.post(
f"/team/{resources.prefix}/member/{member_id}/reset_spend",
headers={"Authorization": f"Bearer {actors.keys[Actor.PROXY_ADMIN].cleartext}"},
json={"reset_to": 5.0},
)
assert resp.status_code == 400, resp.text
async def test_team_member_reset_spend_team_admin_cannot_reset_own_spend(
proxy_client, prisma, scratch, world
gateway_client, db, resources, actors
):
"""A team admin targeting their own LiteLLM_TeamMembership row is 403: unchecked, an
admin could repeatedly zero their own spend right before it crosses their per-member
cap, consuming the shared team budget without the configured limit ever binding."""
team_admin = world.keys[Actor.TEAM_ADMIN]
team_admin = actors.keys[Actor.TEAM_ADMIN]
await create_scratch_team(
prisma,
scratch.prefix,
organization_id=world.org_a_id,
db,
resources.prefix,
organization_id=actors.org_a_id,
admin_user_ids=[team_admin.user_id],
)
await prisma.db.litellm_teammembership.create(
data={"user_id": team_admin.user_id, "team_id": scratch.prefix, "spend": _SEED_SPEND}
write_rows(
'UPDATE "LiteLLM_TeamMembership" SET spend = %s WHERE user_id = %s AND team_id = %s',
(str(_SEED_SPEND), team_admin.user_id, resources.prefix),
)
resp = await proxy_client.post(
f"/team/{scratch.prefix}/member/{team_admin.user_id}/reset_spend",
resp = await gateway_client.post(
f"/team/{resources.prefix}/member/{team_admin.user_id}/reset_spend",
headers={"Authorization": f"Bearer {team_admin.cleartext}"},
json={"reset_to": 0.0},
)
assert resp.status_code == 403, resp.text
row = await prisma.db.litellm_teammembership.find_unique(
where={"user_id_team_id": {"user_id": team_admin.user_id, "team_id": scratch.prefix}}
row = await db.db.litellm_teammembership.find_unique(
where={"user_id_team_id": {"user_id": team_admin.user_id, "team_id": resources.prefix}}
)
assert row is not None and row.spend == _SEED_SPEND, "denied but spend reset"

40
tests/proxy_behavior/management/test_team_member_update.py → tests/integration/management/test_team_member_update.py
View file

@ -1,12 +1,12 @@
import pytest
from .actors import Actor
from .conftest import create_scratch_team
from tests.integration._support.actors import Actor
from tests.integration._support.actors import create_scratch_team
pytestmark = pytest.mark.asyncio(loop_scope="session")
# POST /team/member_update — actor x team-shape matrix. The scratch team is
# POST /team/member_update — actor x team-shape matrix. The resources team is
# raw-seeded with a "user"-role member; each scenario tries to promote it to
# "admin". PROXY_ADMIN, the team's team admin, or an org admin of the team's
# org may update members; else 403. (The harness forces premium_user, so the
@ -33,20 +33,20 @@ _MATRIX = [
]
async def _seed_target(prisma, world, shape: str, team_id: str, member_id: str) -> None:
async def _seed_target(db, actors, shape: str, team_id: str, member_id: str) -> None:
if shape == "alpha":
await create_scratch_team(
prisma,
db,
team_id,
organization_id=world.org_a_id,
admin_user_ids=[world.keys[Actor.TEAM_ADMIN].user_id],
organization_id=actors.org_a_id,
admin_user_ids=[actors.keys[Actor.TEAM_ADMIN].user_id],
member_user_ids=[member_id],
)
elif shape == "beta":
await create_scratch_team(
prisma,
db,
team_id,
organization_id=world.org_b_id,
organization_id=actors.org_b_id,
member_user_ids=[member_id],
)
else: # pragma: no cover - guard
@ -69,26 +69,26 @@ async def test_team_member_update_authz_matrix(
actor: Actor,
shape: str,
expected_status: int,
proxy_client,
prisma,
scratch,
world,
gateway_client,
db,
resources,
actors,
):
member_id = scratch.tag("member")
await _seed_target(prisma, world, shape, scratch.prefix, member_id)
caller = world.keys[actor]
member_id = resources.tag("member")
await _seed_target(db, actors, shape, resources.prefix, member_id)
caller = actors.keys[actor]
resp = await proxy_client.post(
resp = await gateway_client.post(
"/team/member_update",
headers={"Authorization": f"Bearer {caller.cleartext}"},
json={"team_id": scratch.prefix, "user_id": member_id, "role": "admin"},
json={"team_id": resources.prefix, "user_id": member_id, "role": "admin"},
)
assert (
resp.status_code == expected_status
), f"{actor.value} {shape}: {resp.status_code} {resp.text}"
row = await prisma.db.litellm_teamtable.find_unique(
where={"team_id": scratch.prefix}
row = await db.db.litellm_teamtable.find_unique(
where={"team_id": resources.prefix}
)
assert row is not None
if expected_status == 200:

14
tests/proxy_behavior/management/test_team_metadata_schema.py → tests/integration/management/test_team_metadata_schema.py
View file

@ -1,25 +1,25 @@
"""GET /team/metadata_schema — the behavior world declares no
"""GET /team/metadata_schema — the behavior actors declares no
``general_settings.team_metadata_schema``, so the route is an info route that
returns an empty field list to every authenticated actor and 401s without a key.
"""
import pytest
from .actors import Actor
from tests.integration._support.actors import Actor
pytestmark = pytest.mark.asyncio(loop_scope="session")
@pytest.mark.parametrize("actor", list(Actor), ids=[a.value for a in Actor])
async def test_team_metadata_schema_default_is_empty(actor: Actor, proxy_client, world):
resp = await proxy_client.get(
async def test_team_metadata_schema_default_is_empty(actor: Actor, gateway_client, actors):
resp = await gateway_client.get(
"/team/metadata_schema",
headers={"Authorization": f"Bearer {world.keys[actor].cleartext}"},
headers={"Authorization": f"Bearer {actors.keys[actor].cleartext}"},
)
assert resp.status_code == 200, f"{actor.value}: {resp.status_code} {resp.text}"
assert resp.json() == {"fields": []}
async def test_team_metadata_schema_requires_auth(proxy_client, world):
resp = await proxy_client.get("/team/metadata_schema")
async def test_team_metadata_schema_requires_auth(gateway_client, actors):
resp = await gateway_client.get("/team/metadata_schema", headers={"Authorization": "Bearer invalid"})
assert resp.status_code == 401, resp.text

30
tests/proxy_behavior/management/test_team_model.py → tests/integration/management/test_team_model.py
View file

@ -1,7 +1,7 @@
import pytest
from .actors import Actor
from .conftest import create_scratch_team
from tests.integration._support.actors import Actor
from tests.integration._support.actors import create_scratch_team
pytestmark = pytest.mark.asyncio(loop_scope="session")
@ -37,28 +37,28 @@ async def test_team_model_authz_matrix(
route: str,
actor: Actor,
expected_status: int,
proxy_client,
prisma,
scratch,
world,
gateway_client,
db,
resources,
actors,
):
initial = [] if route == "add" else [_MARKER_MODEL]
await create_scratch_team(
prisma, scratch.prefix, organization_id=world.org_a_id, models=initial
db, resources.prefix, organization_id=actors.org_a_id, models=initial
)
caller = world.keys[actor]
caller = actors.keys[actor]
resp = await proxy_client.post(
resp = await gateway_client.post(
_ROUTE_URL[route],
headers={"Authorization": f"Bearer {caller.cleartext}"},
json={"team_id": scratch.prefix, "models": [_MARKER_MODEL]},
json={"team_id": resources.prefix, "models": [_MARKER_MODEL]},
)
assert (
resp.status_code == expected_status
), f"{route} {actor.value}: {resp.status_code} {resp.text}"
row = await prisma.db.litellm_teamtable.find_unique(
where={"team_id": scratch.prefix}
row = await db.db.litellm_teamtable.find_unique(
where={"team_id": resources.prefix}
)
assert row is not None
if expected_status == 200:
@ -68,11 +68,11 @@ async def test_team_model_authz_matrix(
@pytest.mark.parametrize("route", ["add", "delete"])
async def test_team_model_missing_team_is_404(route: str, proxy_client, world):
async def test_team_model_missing_team_is_404(route: str, gateway_client, actors):
"""A team_id absent from the DB is 404 — the existence check precedes authz."""
resp = await proxy_client.post(
resp = await gateway_client.post(
_ROUTE_URL[route],
headers={"Authorization": f"Bearer {world.keys[Actor.PROXY_ADMIN].cleartext}"},
headers={"Authorization": f"Bearer {actors.keys[Actor.PROXY_ADMIN].cleartext}"},
json={"team_id": "behavior-pin-no-such-team", "models": [_MARKER_MODEL]},
)
assert resp.status_code == 404, resp.text

56
tests/proxy_behavior/management/test_team_new.py → tests/integration/management/test_team_new.py
View file

@ -2,7 +2,7 @@ from typing import Any, Dict
import pytest
from .actors import Actor
from tests.integration._support.actors import Actor
pytestmark = pytest.mark.asyncio(loop_scope="session")
@ -51,30 +51,30 @@ async def test_team_new_authz_matrix(
actor: Actor,
org_target: str,
expected_status: int,
proxy_client,
prisma,
scratch,
world,
gateway_client,
db,
resources,
actors,
):
caller = world.keys[actor]
caller = actors.keys[actor]
org_id = {
"none": None,
"org_a": world.org_a_id,
"org_b": world.org_b_id,
"org_a": actors.org_a_id,
"org_b": actors.org_b_id,
}[org_target]
body: Dict[str, Any] = {"team_id": scratch.prefix, "team_alias": scratch.prefix}
body: Dict[str, Any] = {"team_id": resources.prefix, "team_alias": resources.prefix}
if org_id is not None:
body["organization_id"] = org_id
resp = await proxy_client.post(
resp = await gateway_client.post(
"/team/new",
headers={"Authorization": f"Bearer {caller.cleartext}"},
json=body,
)
assert resp.status_code == expected_status, f"{actor.value} org={org_target}: {resp.status_code} {resp.text}"
row = await prisma.db.litellm_teamtable.find_unique(where={"team_id": scratch.prefix})
row = await db.db.litellm_teamtable.find_unique(where={"team_id": resources.prefix})
if expected_status == 200:
assert row is not None
assert row.organization_id == org_id
@ -82,49 +82,49 @@ async def test_team_new_authz_matrix(
assert row is None, f"{actor.value}: denied but team row leaked"
async def test_team_new_rejects_negative_budget(proxy_client, prisma, scratch, world):
async def test_team_new_rejects_negative_budget(gateway_client, db, resources, actors):
"""Input-validation pin: max_budget < 0 is a 400, no row created."""
resp = await proxy_client.post(
resp = await gateway_client.post(
"/team/new",
headers={"Authorization": f"Bearer {world.keys[Actor.PROXY_ADMIN].cleartext}"},
json={"team_id": scratch.prefix, "max_budget": -1},
headers={"Authorization": f"Bearer {actors.keys[Actor.PROXY_ADMIN].cleartext}"},
json={"team_id": resources.prefix, "max_budget": -1},
)
assert resp.status_code == 400, resp.text
row = await prisma.db.litellm_teamtable.find_unique(where={"team_id": scratch.prefix})
row = await db.db.litellm_teamtable.find_unique(where={"team_id": resources.prefix})
assert row is None
async def test_team_new_rejects_duplicate_team_id(proxy_client, prisma, scratch, world):
async def test_team_new_rejects_duplicate_team_id(gateway_client, db, resources, actors):
"""Input-validation pin: a colliding team_id is a 400 on the second call."""
seeder = world.keys[Actor.PROXY_ADMIN].cleartext
first = await proxy_client.post(
seeder = actors.keys[Actor.PROXY_ADMIN].cleartext
first = await gateway_client.post(
"/team/new",
headers={"Authorization": f"Bearer {seeder}"},
json={"team_id": scratch.prefix, "team_alias": scratch.prefix},
json={"team_id": resources.prefix, "team_alias": resources.prefix},
)
assert first.status_code == 200, first.text
second = await proxy_client.post(
second = await gateway_client.post(
"/team/new",
headers={"Authorization": f"Bearer {seeder}"},
json={"team_id": scratch.prefix, "team_alias": scratch.prefix},
json={"team_id": resources.prefix, "team_alias": resources.prefix},
)
assert second.status_code == 400, second.text
async def test_team_new_unknown_organization_is_400(proxy_client, prisma, scratch, world):
async def test_team_new_unknown_organization_is_400(gateway_client, db, resources, actors):
"""A /team/new with an organization_id that does not exist fails 400:
OrganizationNotFoundError is routed into the handler's own
'Organization not found' guard instead of escaping as a 500."""
resp = await proxy_client.post(
resp = await gateway_client.post(
"/team/new",
headers={"Authorization": f"Bearer {world.keys[Actor.PROXY_ADMIN].cleartext}"},
headers={"Authorization": f"Bearer {actors.keys[Actor.PROXY_ADMIN].cleartext}"},
json={
"team_id": scratch.prefix,
"organization_id": scratch.tag("no-such-org"),
"team_id": resources.prefix,
"organization_id": resources.tag("no-such-org"),
},
)
assert resp.status_code == 400, resp.text
assert "Organization not found" in resp.text
row = await prisma.db.litellm_teamtable.find_unique(where={"team_id": scratch.prefix})
row = await db.db.litellm_teamtable.find_unique(where={"team_id": resources.prefix})
assert row is None

104
tests/proxy_behavior/management/test_team_permissions.py → tests/integration/management/test_team_permissions.py
View file

@ -1,10 +1,9 @@
import litellm
import pytest
from litellm.proxy._types import KeyManagementRoutes
from .actors import Actor
from .conftest import create_scratch_team
from tests.integration._support.actors import Actor
from tests.integration._support.actors import create_scratch_team, owned_settings_client
pytestmark = pytest.mark.asyncio(loop_scope="session")
@ -13,7 +12,7 @@ _PERM = KeyManagementRoutes.KEY_INFO.value
# GET /team/permissions_list and POST /team/permissions_update are self-managed
# routes, so every actor reaches the handler. Both grant access to PROXY_ADMIN,
# the team admin, or an org admin of the team's org. The scratch team is in
# the team admin, or an org admin of the team's org. The resources team is in
# ORG_A with TEAM_ADMIN as its team admin.
_MATRIX = [
("proxy_admin", Actor.PROXY_ADMIN, 200),
@ -28,17 +27,17 @@ _MATRIX = [
]
async def _seed_team(prisma, scratch_prefix, world) -> None:
async def _seed_team(db, scratch_prefix, actors) -> None:
await create_scratch_team(
prisma,
db,
scratch_prefix,
organization_id=world.org_a_id,
admin_user_ids=[world.keys[Actor.TEAM_ADMIN].user_id],
organization_id=actors.org_a_id,
admin_user_ids=[actors.keys[Actor.TEAM_ADMIN].user_id],
member_user_ids=[
world.keys[Actor.INTERNAL_USER].user_id,
world.keys[Actor.OWNER].user_id,
world.keys[Actor.UNRELATED_SAME_ORG].user_id,
world.keys[Actor.SERVICE_ACCOUNT].user_id,
actors.keys[Actor.INTERNAL_USER].user_id,
actors.keys[Actor.OWNER].user_id,
actors.keys[Actor.UNRELATED_SAME_ORG].user_id,
actors.keys[Actor.SERVICE_ACCOUNT].user_id,
],
)
@ -49,18 +48,18 @@ async def _seed_team(prisma, scratch_prefix, world) -> None:
ids=[s[0] for s in _MATRIX],
)
async def test_team_permissions_list_authz_matrix(
actor: Actor, expected_status: int, proxy_client, prisma, scratch, world
actor: Actor, expected_status: int, gateway_client, db, resources, actors
):
await _seed_team(prisma, scratch.prefix, world)
resp = await proxy_client.get(
f"/team/permissions_list?team_id={scratch.prefix}",
headers={"Authorization": f"Bearer {world.keys[actor].cleartext}"},
await _seed_team(db, resources.prefix, actors)
resp = await gateway_client.get(
f"/team/permissions_list?team_id={resources.prefix}",
headers={"Authorization": f"Bearer {actors.keys[actor].cleartext}"},
)
assert (
resp.status_code == expected_status
), f"{actor.value}: {resp.status_code} {resp.text}"
if expected_status == 200:
assert resp.json()["team_id"] == scratch.prefix
assert resp.json()["team_id"] == resources.prefix
@pytest.mark.parametrize(
@ -69,20 +68,20 @@ async def test_team_permissions_list_authz_matrix(
ids=[s[0] for s in _MATRIX],
)
async def test_team_permissions_update_authz_matrix(
actor: Actor, expected_status: int, proxy_client, prisma, scratch, world
actor: Actor, expected_status: int, gateway_client, db, resources, actors
):
await _seed_team(prisma, scratch.prefix, world)
resp = await proxy_client.post(
await _seed_team(db, resources.prefix, actors)
resp = await gateway_client.post(
"/team/permissions_update",
headers={"Authorization": f"Bearer {world.keys[actor].cleartext}"},
json={"team_id": scratch.prefix, "team_member_permissions": [_PERM]},
headers={"Authorization": f"Bearer {actors.keys[actor].cleartext}"},
json={"team_id": resources.prefix, "team_member_permissions": [_PERM]},
)
assert (
resp.status_code == expected_status
), f"{actor.value}: {resp.status_code} {resp.text}"
row = await prisma.db.litellm_teamtable.find_unique(
where={"team_id": scratch.prefix}
row = await db.db.litellm_teamtable.find_unique(
where={"team_id": resources.prefix}
)
assert row is not None
if expected_status == 200:
@ -92,29 +91,24 @@ async def test_team_permissions_update_authz_matrix(
async def test_team_permissions_available_team_self_join_divergence(
proxy_client, prisma, scratch, world, monkeypatch
gateway_client, db, resources, actors, tmp_path
):
"""permissions_list honours the available-team self-join — a non-admin can
READ an available team's permissions — but permissions_update deliberately
does not: the same caller is 403 on update. default_internal_user_params is
module-level litellm.* state, so monkeypatch save/restores it."""
await create_scratch_team(prisma, scratch.prefix, organization_id=world.org_a_id)
monkeypatch.setattr(
litellm, "default_internal_user_params", {"available_teams": [scratch.prefix]}
)
caller = world.keys[Actor.CROSS_ORG_USER] # non-admin, unrelated to the team
await create_scratch_team(db, resources.prefix, organization_id=actors.org_a_id)
caller = actors.keys[Actor.CROSS_ORG_USER]
listed = await proxy_client.get(
f"/team/permissions_list?team_id={scratch.prefix}",
headers={"Authorization": f"Bearer {caller.cleartext}"},
)
with owned_settings_client(
gateway_client, tmp_path, {"default_internal_user_params": {"available_teams": [resources.prefix]}}
) as client:
listed = await client.get(
f"/team/permissions_list?team_id={resources.prefix}",
headers={"Authorization": f"Bearer {caller.cleartext}"},
)
updated = await client.post(
"/team/permissions_update",
headers={"Authorization": f"Bearer {caller.cleartext}"},
json={"team_id": resources.prefix, "team_member_permissions": [_PERM]},
)
assert listed.status_code == 200, listed.text
updated = await proxy_client.post(
"/team/permissions_update",
headers={"Authorization": f"Bearer {caller.cleartext}"},
json={"team_id": scratch.prefix, "team_member_permissions": [_PERM]},
)
assert updated.status_code == 403, updated.text
@ -138,20 +132,20 @@ _BULK_MATRIX = [
ids=[s[0] for s in _BULK_MATRIX],
)
async def test_team_permissions_bulk_update_authz_matrix(
actor: Actor, expected_status: int, proxy_client, prisma, scratch, world
actor: Actor, expected_status: int, gateway_client, db, resources, actors
):
await create_scratch_team(prisma, scratch.prefix, organization_id=world.org_a_id)
resp = await proxy_client.post(
await create_scratch_team(db, resources.prefix, organization_id=actors.org_a_id)
resp = await gateway_client.post(
"/team/permissions_bulk_update",
headers={"Authorization": f"Bearer {world.keys[actor].cleartext}"},
json={"team_ids": [scratch.prefix], "permissions": [_PERM]},
headers={"Authorization": f"Bearer {actors.keys[actor].cleartext}"},
json={"team_ids": [resources.prefix], "permissions": [_PERM]},
)
assert (
resp.status_code == expected_status
), f"{actor.value}: {resp.status_code} {resp.text}"
row = await prisma.db.litellm_teamtable.find_unique(
where={"team_id": scratch.prefix}
row = await db.db.litellm_teamtable.find_unique(
where={"team_id": resources.prefix}
)
assert row is not None
if expected_status == 200:
@ -160,11 +154,11 @@ async def test_team_permissions_bulk_update_authz_matrix(
assert _PERM not in (row.team_member_permissions or []), "denied but mutated"
async def test_team_permissions_bulk_update_no_selector_is_400(proxy_client, world):
async def test_team_permissions_bulk_update_no_selector_is_400(gateway_client, actors):
"""Neither team_ids nor apply_to_all_teams is a 400."""
resp = await proxy_client.post(
resp = await gateway_client.post(
"/team/permissions_bulk_update",
headers={"Authorization": f"Bearer {world.keys[Actor.PROXY_ADMIN].cleartext}"},
headers={"Authorization": f"Bearer {actors.keys[Actor.PROXY_ADMIN].cleartext}"},
json={"permissions": [_PERM]},
)
assert resp.status_code == 400, resp.text

136
tests/proxy_behavior/management/test_team_permissions_bulk_update.py → tests/integration/management/test_team_permissions_bulk_update.py
View file

@ -8,22 +8,20 @@ Pins the three helpers
The route is admin-only — PROXY_ADMIN is the only legal caller. The
contracts under test:
* specific-team list → ONLY listed teams mutate; every other team
(including world teams) is byte-identical on re-read.
* apply_to_all_teams → every team gains the permission, idempotently
(including actors teams) is byte-identical on re-read.
* apply_to_all_teams → every team has the permission, idempotently
merged (re-running with the same permission is a no-op).
* unknown team_id → 404 (the missing_ids guard); no mutation anywhere.
* malformed payload (neither / both selector flags) → 400; no mutation.
The all-teams scenario mutates world teams as a deliberate side effect of
the path under test. The test snapshots every team's
`team_member_permissions` up front and restores them on exit so the
read-world stays immutable for downstream tests.
"""
from typing import Final
import pytest
from .actors import TEAM_ALPHA, TEAM_BETA, TEAM_GAMMA, Actor
from .conftest import create_scratch_team
from tests.integration._support.actors import TEAM_ALPHA, TEAM_BETA, TEAM_GAMMA, Actor
from tests.integration._support.actors import create_scratch_team
pytestmark = pytest.mark.asyncio(loop_scope="session")
@ -34,28 +32,28 @@ pytestmark = pytest.mark.asyncio(loop_scope="session")
async def test_bulk_update_specific_team_only_mutates_listed(
proxy_client, prisma, scratch, world
gateway_client, db, resources, actors
):
target = await create_scratch_team(
prisma,
team_id=scratch.tag("target"),
db,
team_id=resources.tag("target"),
team_member_permissions=[],
)
bystander = await create_scratch_team(
prisma,
team_id=scratch.tag("bystander"),
db,
team_id=resources.tag("bystander"),
team_member_permissions=["pre-existing-perm"],
)
# Snapshot world teams so we can assert byte-equality on re-read.
world_team_ids = [TEAM_ALPHA, TEAM_BETA, TEAM_GAMMA]
# Snapshot actors teams so we can assert byte-equality on re-read.
world_team_ids = [actors.team_alpha_id, actors.team_beta_id, actors.team_gamma_id]
before = {}
for tid in world_team_ids + [bystander]:
row = await prisma.db.litellm_teamtable.find_unique(where={"team_id": tid})
row = await db.db.litellm_teamtable.find_unique(where={"team_id": tid})
before[tid] = list(row.team_member_permissions or [])
seeder = world.keys[Actor.PROXY_ADMIN].cleartext
seeder = actors.keys[Actor.PROXY_ADMIN].cleartext
perm = "/key/info"
resp = await proxy_client.post(
resp = await gateway_client.post(
"/team/permissions_bulk_update",
headers={"Authorization": f"Bearer {seeder}"},
json={"team_ids": [target], "permissions": [perm]},
@ -63,7 +61,7 @@ async def test_bulk_update_specific_team_only_mutates_listed(
assert resp.status_code == 200, resp.text
assert resp.json()["teams_updated"] == 1
target_row = await prisma.db.litellm_teamtable.find_unique(
target_row = await db.db.litellm_teamtable.find_unique(
where={"team_id": target}
)
assert perm in (
@ -71,7 +69,7 @@ async def test_bulk_update_specific_team_only_mutates_listed(
), f"target team did not gain perm; got={target_row.team_member_permissions}"
for tid, before_perms in before.items():
after_row = await prisma.db.litellm_teamtable.find_unique(
after_row = await db.db.litellm_teamtable.find_unique(
where={"team_id": tid}
)
assert list(after_row.team_member_permissions or []) == before_perms, (
@ -87,15 +85,15 @@ async def test_bulk_update_specific_team_only_mutates_listed(
async def test_bulk_update_specific_team_idempotent(
proxy_client, prisma, scratch, world
gateway_client, db, resources, actors
):
target = await create_scratch_team(
prisma,
team_id=scratch.tag("target"),
db,
team_id=resources.tag("target"),
team_member_permissions=["/key/info"],
)
seeder = world.keys[Actor.PROXY_ADMIN].cleartext
resp = await proxy_client.post(
seeder = actors.keys[Actor.PROXY_ADMIN].cleartext
resp = await gateway_client.post(
"/team/permissions_bulk_update",
headers={"Authorization": f"Bearer {seeder}"},
json={
@ -113,25 +111,25 @@ async def test_bulk_update_specific_team_idempotent(
async def test_bulk_update_unknown_team_id_rejected(
proxy_client, prisma, scratch, world
gateway_client, db, resources, actors
):
existing = await create_scratch_team(
prisma,
team_id=scratch.tag("real"),
db,
team_id=resources.tag("real"),
team_member_permissions=[],
)
before = list(
(
await prisma.db.litellm_teamtable.find_unique(where={"team_id": existing})
await db.db.litellm_teamtable.find_unique(where={"team_id": existing})
).team_member_permissions
or []
)
seeder = world.keys[Actor.PROXY_ADMIN].cleartext
resp = await proxy_client.post(
seeder = actors.keys[Actor.PROXY_ADMIN].cleartext
resp = await gateway_client.post(
"/team/permissions_bulk_update",
headers={"Authorization": f"Bearer {seeder}"},
json={
"team_ids": [existing, f"{scratch.prefix}-ghost"],
"team_ids": [existing, f"{resources.prefix}-ghost"],
"permissions": ["/key/info"],
},
)
@ -141,7 +139,7 @@ async def test_bulk_update_unknown_team_id_rejected(
assert "not found" in resp.text.lower() or "ghost" in resp.text, resp.text
# Critical: the partial-success regression shape is "real team got
# mutated before the ghost-id check ran". Re-read and assert it didn't.
after_real = await prisma.db.litellm_teamtable.find_unique(
after_real = await db.db.litellm_teamtable.find_unique(
where={"team_id": existing}
)
assert (
@ -173,10 +171,10 @@ async def test_bulk_update_unknown_team_id_rejected(
ids=["neither_selector", "both_selectors"],
)
async def test_bulk_update_selector_validation(
body, expected_substring: str, proxy_client, world
body, expected_substring: str, gateway_client, actors
):
seeder = world.keys[Actor.PROXY_ADMIN].cleartext
resp = await proxy_client.post(
seeder = actors.keys[Actor.PROXY_ADMIN].cleartext
resp = await gateway_client.post(
"/team/permissions_bulk_update",
headers={"Authorization": f"Bearer {seeder}"},
json=body,
@ -186,60 +184,40 @@ async def test_bulk_update_selector_validation(
# ---------------------------------------------------------------------------
# apply_to_all_teams — mutates every team. Snapshot + restore world teams
# so the read-world contract holds for downstream tests.
# apply_to_all_teams
# ---------------------------------------------------------------------------
async def test_bulk_update_apply_to_all_mutates_every_team(
proxy_client, prisma, scratch, world
gateway_client, db, resources, actors
):
# Two scratch teams so we can assert "every" includes our own targets.
# Two resources teams so we can assert "every" includes our own targets.
a = await create_scratch_team(
prisma, team_id=scratch.tag("a"), team_member_permissions=[]
db, team_id=resources.tag("a"), team_member_permissions=[]
)
b = await create_scratch_team(
prisma, team_id=scratch.tag("b"), team_member_permissions=["other"]
db, team_id=resources.tag("b"), team_member_permissions=["other"]
)
perm = "/key/health" # distinct from the specific-team scenarios above
# Snapshot every team's permission list so we can restore world teams.
all_teams = await prisma.db.litellm_teamtable.find_many()
all_teams = await db.db.litellm_teamtable.find_many()
snapshot = {t.team_id: list(t.team_member_permissions or []) for t in all_teams}
expected_updates: Final = sum(
1 for permissions in snapshot.values() if perm not in permissions
)
try:
seeder = world.keys[Actor.PROXY_ADMIN].cleartext
resp = await proxy_client.post(
"/team/permissions_bulk_update",
headers={"Authorization": f"Bearer {seeder}"},
json={"apply_to_all_teams": True, "permissions": [perm]},
)
assert resp.status_code == 200, resp.text
# `teams_updated` counts teams that didn't already have the perm —
# i.e. every team in the DB at call time.
assert resp.json()["teams_updated"] == len(all_teams)
seeder = actors.keys[Actor.PROXY_ADMIN].cleartext
resp = await gateway_client.post(
"/team/permissions_bulk_update",
headers={"Authorization": f"Bearer {seeder}"},
json={"apply_to_all_teams": True, "permissions": [perm]},
)
assert resp.status_code == 200, resp.text
assert resp.json()["teams_updated"] == expected_updates
post = await prisma.db.litellm_teamtable.find_many()
for team in post:
perms = list(team.team_member_permissions or [])
assert perm in perms, f"team {team.team_id} missing the all-perm: {perms}"
# Existing perms preserved (merge, not replace).
for prior in snapshot.get(team.team_id, []):
assert (
prior in perms
), f"team {team.team_id} lost prior perm {prior!r}: {perms}"
finally:
# Restore every team — scratch teardown handles {a, b}; we must
# explicitly restore world teams (and any other non-scratch teams
# that snuck in) so downstream tests see the immutable world.
for team_id, prior_perms in snapshot.items():
current = await prisma.db.litellm_teamtable.find_unique(
where={"team_id": team_id}
)
if current is None:
continue
if list(current.team_member_permissions or []) != prior_perms:
await prisma.db.litellm_teamtable.update(
where={"team_id": team_id},
data={"team_member_permissions": prior_perms},
)
post = await db.db.litellm_teamtable.find_many()
for team in post:
perms = list(team.team_member_permissions or [])
assert perm in perms, f"team {team.team_id} missing the all-perm: {perms}"
for prior in snapshot.get(team.team_id, []):
assert prior in perms, f"team {team.team_id} lost prior perm {prior!r}: {perms}"

10
tests/proxy_behavior/management/test_team_spend_by_user.py → tests/integration/management/test_team_spend_by_user.py
View file

@ -1,6 +1,6 @@
import pytest
from .actors import Actor
from tests.integration._support.actors import Actor
pytestmark = pytest.mark.asyncio(loop_scope="session")
@ -43,13 +43,13 @@ _DATES = "start_date=2024-01-01&end_date=2024-12-31"
[(a, t, s) for (_id, a, t, s) in _CASES],
ids=[c[0] for c in _CASES],
)
async def test_team_spend_by_user_matrix(actor: Actor, team: str, expected_status: int, proxy_client, world):
team_id = {"alpha": world.team_alpha_id, "beta": world.team_beta_id}.get(team)
async def test_team_spend_by_user_matrix(actor: Actor, team: str, expected_status: int, gateway_client, actors):
team_id = {"alpha": actors.team_alpha_id, "beta": actors.team_beta_id}.get(team)
query = _DATES if team_id is None else f"{_DATES}&team_ids={team_id}"
resp = await proxy_client.get(
resp = await gateway_client.get(
f"/team/spend/by_user?{query}",
headers={"Authorization": f"Bearer {world.keys[actor].cleartext}"},
headers={"Authorization": f"Bearer {actors.keys[actor].cleartext}"},
)
assert resp.status_code == expected_status, f"{actor.value} -> {team}: {resp.status_code} {resp.text}"
if expected_status == 200:

130
tests/proxy_behavior/management/test_team_update.py → tests/integration/management/test_team_update.py
View file

@ -2,8 +2,8 @@ import pytest
from litellm.proxy._types import LitellmUserRoles
from .actors import Actor
from .conftest import create_scratch_actor, create_scratch_team
from tests.integration._support.actors import Actor
from tests.integration._support.actors import create_scratch_actor, create_scratch_team
pytestmark = pytest.mark.asyncio(loop_scope="session")
@ -38,30 +38,30 @@ _MATRIX = [
]
async def _seed_target(prisma, world, shape: str, team_id: str) -> str:
"""Raw-seed the scratch target team; returns its organization_id."""
async def _seed_target(db, actors, shape: str, team_id: str) -> str:
"""Raw-seed the resources target team; returns its organization_id."""
if shape == "alpha":
await create_scratch_team(
prisma,
db,
team_id,
organization_id=world.org_a_id,
admin_user_ids=[world.keys[Actor.TEAM_ADMIN].user_id],
organization_id=actors.org_a_id,
admin_user_ids=[actors.keys[Actor.TEAM_ADMIN].user_id],
member_user_ids=[
world.keys[Actor.INTERNAL_USER].user_id,
world.keys[Actor.OWNER].user_id,
world.keys[Actor.UNRELATED_SAME_ORG].user_id,
world.keys[Actor.SERVICE_ACCOUNT].user_id,
actors.keys[Actor.INTERNAL_USER].user_id,
actors.keys[Actor.OWNER].user_id,
actors.keys[Actor.UNRELATED_SAME_ORG].user_id,
actors.keys[Actor.SERVICE_ACCOUNT].user_id,
],
)
return world.org_a_id
return actors.org_a_id
if shape == "beta":
await create_scratch_team(
prisma,
db,
team_id,
organization_id=world.org_b_id,
member_user_ids=[world.keys[Actor.CROSS_ORG_USER].user_id],
organization_id=actors.org_b_id,
member_user_ids=[actors.keys[Actor.CROSS_ORG_USER].user_id],
)
return world.org_b_id
return actors.org_b_id
pytest.fail(f"unknown shape={shape}") # pragma: no cover
@ -74,19 +74,19 @@ async def test_team_update_authz_matrix(
actor: Actor,
shape: str,
expected_status: int,
proxy_client,
prisma,
scratch,
world,
gateway_client,
db,
resources,
actors,
):
org_id = await _seed_target(prisma, world, shape, scratch.prefix)
caller = world.keys[actor]
org_id = await _seed_target(db, actors, shape, resources.prefix)
caller = actors.keys[actor]
resp = await proxy_client.post(
resp = await gateway_client.post(
"/team/update",
headers={"Authorization": f"Bearer {caller.cleartext}"},
json={
"team_id": scratch.prefix,
"team_id": resources.prefix,
"team_alias": MARKER_ALIAS,
"organization_id": org_id,
},
@ -95,8 +95,8 @@ async def test_team_update_authz_matrix(
resp.status_code == expected_status
), f"{actor.value} {shape}: {resp.status_code} {resp.text}"
row = await prisma.db.litellm_teamtable.find_unique(
where={"team_id": scratch.prefix}
row = await db.db.litellm_teamtable.find_unique(
where={"team_id": resources.prefix}
)
assert row is not None
if expected_status == 200:
@ -106,41 +106,41 @@ async def test_team_update_authz_matrix(
async def test_team_update_org_admin_resolved_from_team_without_org_context(
proxy_client, prisma, scratch, world
gateway_client, db, resources, actors
):
"""With no organization_id in the body the route gate resolves the target
team's org from team_id, so an org admin of the team's own org is allowed
(200), same as PROXY_ADMIN. A team admin of that same team reaches the
handler but is refused (403) until a proxy admin enables fields for team
admins, and the response says so."""
await _seed_target(prisma, world, "alpha", scratch.prefix)
await _seed_target(db, actors, "alpha", resources.prefix)
allowed_org_admin = await proxy_client.post(
allowed_org_admin = await gateway_client.post(
"/team/update",
headers={"Authorization": f"Bearer {world.keys[Actor.ORG_ADMIN].cleartext}"},
json={"team_id": scratch.prefix, "team_alias": MARKER_ALIAS},
headers={"Authorization": f"Bearer {actors.keys[Actor.ORG_ADMIN].cleartext}"},
json={"team_id": resources.prefix, "team_alias": MARKER_ALIAS},
)
assert allowed_org_admin.status_code == 200, allowed_org_admin.text
allowed_proxy_admin = await proxy_client.post(
allowed_proxy_admin = await gateway_client.post(
"/team/update",
headers={"Authorization": f"Bearer {world.keys[Actor.PROXY_ADMIN].cleartext}"},
json={"team_id": scratch.prefix, "team_alias": MARKER_ALIAS},
headers={"Authorization": f"Bearer {actors.keys[Actor.PROXY_ADMIN].cleartext}"},
json={"team_id": resources.prefix, "team_alias": MARKER_ALIAS},
)
assert allowed_proxy_admin.status_code == 200, allowed_proxy_admin.text
denied_team_admin = await proxy_client.post(
denied_team_admin = await gateway_client.post(
"/team/update",
headers={"Authorization": f"Bearer {world.keys[Actor.TEAM_ADMIN].cleartext}"},
json={"team_id": scratch.prefix, "team_alias": MARKER_ALIAS},
headers={"Authorization": f"Bearer {actors.keys[Actor.TEAM_ADMIN].cleartext}"},
json={"team_id": resources.prefix, "team_alias": MARKER_ALIAS},
)
assert denied_team_admin.status_code == 403, denied_team_admin.text
assert "cannot edit team settings" in denied_team_admin.text, denied_team_admin.text
assert "Team admin editable fields" in denied_team_admin.text, denied_team_admin.text
# Relocation gate — moving a team to a different org. The scratch team starts
# in ORG_A; each scenario relocates it to ORG_B. PROXY_ADMIN bypasses;
# Relocation gate — moving a team to a different org. The resources team starts
# in actors.org_a_id; each scenario relocates it to actors.org_b_id. PROXY_ADMIN bypasses;
# ORG_B_ADMIN reaches the handler but holds no role on the source team (403);
# ORG_ADMIN holds the source team but not the destination org (403 from the
# relocation gate); the team admin is refused by the empty field allow-list and
@ -164,31 +164,31 @@ _RELOCATION = [
async def test_team_update_org_relocation_gate(
actor: Actor,
expected_status: int,
proxy_client,
prisma,
scratch,
world,
gateway_client,
db,
resources,
actors,
):
await _seed_target(prisma, world, "alpha", scratch.prefix)
caller = world.keys[actor]
await _seed_target(db, actors, "alpha", resources.prefix)
caller = actors.keys[actor]
resp = await proxy_client.post(
resp = await gateway_client.post(
"/team/update",
headers={"Authorization": f"Bearer {caller.cleartext}"},
json={"team_id": scratch.prefix, "organization_id": world.org_b_id},
json={"team_id": resources.prefix, "organization_id": actors.org_b_id},
)
assert (
resp.status_code == expected_status
), f"{actor.value}: {resp.status_code} {resp.text}"
row = await prisma.db.litellm_teamtable.find_unique(
where={"team_id": scratch.prefix}
row = await db.db.litellm_teamtable.find_unique(
where={"team_id": resources.prefix}
)
assert row is not None
if expected_status == 200:
assert row.organization_id == world.org_b_id
assert row.organization_id == actors.org_b_id
else:
assert row.organization_id == world.org_a_id, "denied but team relocated"
assert row.organization_id == actors.org_a_id, "denied but team relocated"
# Phase 4 F6 — explicit pin on the `team_access_denied` 403 detail string
@ -197,44 +197,44 @@ async def test_team_update_org_relocation_gate(
# status; this guard turns a silent rename of the helper's exception detail
# into a CI red.
async def test_team_update_org_b_admin_relocation_rejection_detail(
proxy_client, prisma, scratch, world
gateway_client, db, resources, actors
):
await _seed_target(prisma, world, "alpha", scratch.prefix)
resp = await proxy_client.post(
await _seed_target(db, actors, "alpha", resources.prefix)
resp = await gateway_client.post(
"/team/update",
headers={"Authorization": f"Bearer {world.keys[Actor.ORG_B_ADMIN].cleartext}"},
json={"team_id": scratch.prefix, "organization_id": world.org_b_id},
headers={"Authorization": f"Bearer {actors.keys[Actor.ORG_B_ADMIN].cleartext}"},
json={"team_id": resources.prefix, "organization_id": actors.org_b_id},
)
assert resp.status_code == 403, resp.text
assert "do not have access to this team" in resp.text, resp.text
async def test_team_update_org_relocation_allowed_for_dual_org_admin(
proxy_client, prisma, scratch, world
gateway_client, db, resources, actors
):
"""Relocation-allowed branch: a caller who is org admin of BOTH the source
and destination org may relocate a team between them. Completes the
_RELOCATION matrix, whose allowed branch PR2 left open — no seeded actor is
a dual-org admin, so one is minted with create_scratch_actor."""
actor = await create_scratch_actor(
prisma,
scratch.prefix,
db,
resources.prefix,
user_role=LitellmUserRoles.ORG_ADMIN.value,
org_admin_of=(world.org_a_id, world.org_b_id),
org_admin_of=(actors.org_a_id, actors.org_b_id),
)
team_id = await create_scratch_team(
prisma, scratch.tag("team"), organization_id=world.org_a_id
db, resources.tag("team"), organization_id=actors.org_a_id
)
resp = await proxy_client.post(
resp = await gateway_client.post(
"/team/update",
headers={"Authorization": f"Bearer {actor.cleartext}"},
json={"team_id": team_id, "organization_id": world.org_b_id},
json={"team_id": team_id, "organization_id": actors.org_b_id},
)
assert resp.status_code == 200, resp.text
row = await prisma.db.litellm_teamtable.find_unique(where={"team_id": team_id})
row = await db.db.litellm_teamtable.find_unique(where={"team_id": team_id})
assert row is not None
assert (
row.organization_id == world.org_b_id
row.organization_id == actors.org_b_id
), "dual-org admin relocation not applied"

70
tests/proxy_behavior/management/test_users_bulk_delete.py → tests/integration/management/test_users_bulk_delete.py
View file

@ -1,7 +1,7 @@
import pytest
from .actors import Actor
from .conftest import create_scratch_team, create_scratch_user
from tests.integration._support.actors import Actor
from tests.integration._support.actors import create_scratch_team, create_scratch_user
pytestmark = pytest.mark.asyncio(loop_scope="session")
@ -25,16 +25,16 @@ def _member_ids(row) -> list:
return [m["user_id"] for m in (row.members_with_roles or [])]
async def _seed_team_members(prisma, scratch, world, member_ids: list, org_id) -> None:
async def _seed_team_members(db, resources, actors, member_ids: list, org_id) -> None:
"""Leave behind what /team/member_add would: roster entry, `teams` array, and org membership."""
await create_scratch_team(prisma, scratch.prefix, organization_id=world.org_a_id, member_user_ids=member_ids)
await prisma.db.litellm_usertable.update_many(
where={"user_id": {"in": member_ids}}, data={"teams": {"set": [scratch.prefix]}}
await create_scratch_team(db, resources.prefix, organization_id=actors.org_a_id, member_user_ids=member_ids)
await db.db.litellm_usertable.update_many(
where={"user_id": {"in": member_ids}}, data={"teams": {"set": [resources.prefix]}}
)
if org_id is None:
return
for uid in member_ids:
await prisma.db.litellm_organizationmembership.create(
await db.db.litellm_organizationmembership.create(
data={"user_id": uid, "organization_id": org_id, "user_role": "internal_user"}
)
@ -49,25 +49,25 @@ async def test_users_bulk_delete_authz_matrix(
org,
expected_status: int,
expect_deleted: bool,
proxy_client,
prisma,
scratch,
world,
gateway_client,
db,
resources,
actors,
):
victims = [await create_scratch_user(prisma, scratch.prefix, suffix=s) for s in ("v1", "v2")]
keep = await create_scratch_user(prisma, scratch.prefix, suffix="keep")
await _seed_team_members(prisma, scratch, world, victims + [keep], world.org_a_id if org == "a" else None)
victims = [await create_scratch_user(db, resources.prefix, suffix=s) for s in ("v1", "v2")]
keep = await create_scratch_user(db, resources.prefix, suffix="keep")
await _seed_team_members(db, resources, actors, victims + [keep], actors.org_a_id if org == "a" else None)
resp = await proxy_client.post(
resp = await gateway_client.post(
_URL,
headers={"Authorization": f"Bearer {world.keys[actor].cleartext}"},
headers={"Authorization": f"Bearer {actors.keys[actor].cleartext}"},
json={"user_ids": victims},
)
assert resp.status_code == expected_status, f"{actor.value}: {resp.status_code} {resp.text}"
team = await prisma.db.litellm_teamtable.find_unique(where={"team_id": scratch.prefix})
team = await db.db.litellm_teamtable.find_unique(where={"team_id": resources.prefix})
assert team is not None and keep in _member_ids(team), "unrelated member removed"
remaining = {u.user_id for u in await prisma.db.litellm_usertable.find_many(where={"user_id": {"in": victims}})}
remaining = {u.user_id for u in await db.db.litellm_usertable.find_many(where={"user_id": {"in": victims}})}
if expected_status == 403:
assert resp.headers["content-type"] == "application/problem+json"
assert resp.json()["type"] == "urn:litellm:error:forbidden"
@ -79,7 +79,7 @@ async def test_users_bulk_delete_authz_matrix(
assert set(body) == {"data"}
rows = [(r["user_id"], r["success"], r["teams_removed"]) for r in body["data"]]
if expect_deleted:
assert rows == [(v, True, [scratch.prefix]) for v in victims]
assert rows == [(v, True, [resources.prefix]) for v in victims]
assert remaining == set()
assert not set(victims) & set(_member_ids(team))
return
@ -89,13 +89,13 @@ async def test_users_bulk_delete_authz_matrix(
assert set(victims) <= set(_member_ids(team))
async def test_users_bulk_delete_reports_each_row_in_order(proxy_client, prisma, scratch, world):
victim = await create_scratch_user(prisma, scratch.prefix, suffix="victim")
ghost = scratch.tag("ghost")
async def test_users_bulk_delete_reports_each_row_in_order(gateway_client, db, resources, actors):
victim = await create_scratch_user(db, resources.prefix, suffix="victim")
ghost = resources.tag("ghost")
resp = await proxy_client.post(
resp = await gateway_client.post(
_URL,
headers={"Authorization": f"Bearer {world.keys[Actor.PROXY_ADMIN].cleartext}"},
headers={"Authorization": f"Bearer {actors.keys[Actor.PROXY_ADMIN].cleartext}"},
json={"user_ids": [ghost, victim, victim]},
)
assert resp.status_code == 200, resp.text
@ -104,34 +104,34 @@ async def test_users_bulk_delete_reports_each_row_in_order(proxy_client, prisma,
(victim, True),
(victim, False),
]
assert await prisma.db.litellm_usertable.find_unique(where={"user_id": victim}) is None
assert await db.db.litellm_usertable.find_unique(where={"user_id": victim}) is None
async def test_users_bulk_delete_unknown_query_param_is_400_problem(proxy_client, prisma, scratch, world):
victim = await create_scratch_user(prisma, scratch.prefix, suffix="victim")
async def test_users_bulk_delete_unknown_query_param_is_400_problem(gateway_client, db, resources, actors):
victim = await create_scratch_user(db, resources.prefix, suffix="victim")
resp = await proxy_client.post(
resp = await gateway_client.post(
f"{_URL}?dry_run=1",
headers={"Authorization": f"Bearer {world.keys[Actor.PROXY_ADMIN].cleartext}"},
headers={"Authorization": f"Bearer {actors.keys[Actor.PROXY_ADMIN].cleartext}"},
json={"user_ids": [victim]},
)
assert resp.status_code == 400, resp.text
assert resp.headers["content-type"] == "application/problem+json"
assert resp.json()["type"] == "urn:litellm:error:unknown-query-parameter"
assert "dry_run" in resp.json()["detail"]
assert await prisma.db.litellm_usertable.find_unique(where={"user_id": victim}) is not None
assert await db.db.litellm_usertable.find_unique(where={"user_id": victim}) is not None
async def test_users_bulk_delete_unknown_body_field_is_422_problem(proxy_client, prisma, scratch, world):
victim = await create_scratch_user(prisma, scratch.prefix, suffix="victim")
async def test_users_bulk_delete_unknown_body_field_is_422_problem(gateway_client, db, resources, actors):
victim = await create_scratch_user(db, resources.prefix, suffix="victim")
resp = await proxy_client.post(
resp = await gateway_client.post(
_URL,
headers={"Authorization": f"Bearer {world.keys[Actor.PROXY_ADMIN].cleartext}"},
headers={"Authorization": f"Bearer {actors.keys[Actor.PROXY_ADMIN].cleartext}"},
json={"user_ids": [victim], "dry_run": True},
)
assert resp.status_code == 422, resp.text
assert resp.headers["content-type"] == "application/problem+json"
assert resp.json()["type"] == "urn:litellm:error:invalid-request-body"
assert "dry_run" in resp.json()["detail"]
assert await prisma.db.litellm_usertable.find_unique(where={"user_id": victim}) is not None
assert await db.db.litellm_usertable.find_unique(where={"user_id": victim}) is not None

Some files were not shown because too many files have changed in this diff Show more