scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-13 08:52:02 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 175

Merge branch 'lts/1.108.x' into litellm_backport_messages_reject_background_lts_1_108_x

Browse source
This commit is contained in:
Yuneng Jiang 2026-11-09 11:57:27 -08:00
commit 7f78209db1
No known key found for this signature in database
135 changed files with 8443 additions and 4809 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
.circleci/config.yml
View file

@ -1291,7 +1291,7 @@ jobs:
name: Install MCP SDK1 peer
command: |
uv venv --python 3.12 .venv-mcp-peer
uv pip install --python .venv-mcp-peer 'mcp==1.28.1' 'langchain-mcp-adapters==0.2.1'
uv pip install --no-config --exclude-newer "3 days" --python .venv-mcp-peer 'mcp==1.28.1' 'langchain-mcp-adapters==0.2.1'
echo "export MCP_TEST_PEER_PYTHON=$PWD/.venv-mcp-peer/bin/python" >> "$BASH_ENV"
- when:
condition:

69
.github/e2e-stack/assert_tests_ran.py vendored
View file

@ -1,69 +0,0 @@
import os
import re
import sys
import xml.etree.ElementTree as ET
from pathlib import Path
from typing import Final
sys.path.insert(0, str(Path(__file__).resolve().parents[2] / "tests/e2e"))
from coverage_registry.management_cases import MANAGEMENT_CASES
def main() -> int:
selected: Final = tuple(sys.argv[2:])
try:
report: Final = ET.parse(Path(sys.argv[1])).getroot()
except (ET.ParseError, OSError):
_ = sys.stdout.write("::error::could not read the test execution report\n")
return 1
cases: Final = tuple(report.iter("testcase"))
expected_count: Final = os.environ.get("E2E_REQUIRED_TEST_COUNT")
passed: Final = frozenset(
case.get("file") for case in cases if all(case.find(tag) is None for tag in ("skipped", "failure", "error"))
)
missing: Final = tuple(path for path in selected if path not in passed)
required_nodes: Final = frozenset(case.node for case in MANAGEMENT_CASES if case.node.split("::", 1)[0] in selected)
passed_nodes: Final = frozenset(
prop.get("value")
for case in cases
if all(case.find(tag) is None for tag in ("skipped", "failure", "error"))
for prop in case.findall("./properties/property")
if prop.get("name") == "management_node"
)
missing_nodes: Final = required_nodes - passed_nodes
for node in sorted(missing_nodes):
_ = sys.stdout.write(f"::error::required management case did not pass: {node}\n")
for path in selected:
collected: Final = sum(case.get("file") == path for case in cases)
skipped: Final = sum(case.get("file") == path and case.find("skipped") is not None for case in cases)
_ = sys.stdout.write(f"{path}: {collected} collected, {skipped} skipped\n")
for case in cases:
if case.get("file") != path or all(case.find(tag) is None for tag in ("failure", "error", "skipped")):
continue
outcome = "skipped" if case.find("skipped") is not None else "failed"
_ = sys.stdout.write(f" {outcome}: {case.get('classname', '')}::{case.get('name', '')}\n")
for prop in case.findall("./properties/property"):
name = prop.get("name", "")
value = prop.get("value", "")
if name in ("oauth_failure_phase", "oauth_exception_type", "oauth_frame") and re.fullmatch(
r"[A-Za-z0-9_.:<>-]{1,240}", value
):
_ = sys.stdout.write(f" {name}: {value}\n")
if expected_count is not None and (
len(cases) != int(expected_count) or any(case.find("skipped") is not None for case in cases)
):
_ = sys.stdout.write("::error::required test count was not met or a required case was skipped\n")
return 1
if (
selected
and not missing
and not missing_nodes
and not any(case.find(tag) is not None for case in cases for tag in ("failure", "error"))
):
return 0
_ = sys.stdout.write("::error::every selected file must execute a passing test, with no failures or errors\n")
return 1
if __name__ == "__main__":
sys.exit(main())

319
.github/scripts/auto_update_price_and_context_window_file.py vendored
View file

@ -1,319 +0,0 @@
import asyncio
import aiohttp
import json
import math
from typing import Any
# Asynchronously fetch data from a given URL
async def fetch_data(url):
try:
# Create an asynchronous session
async with aiohttp.ClientSession() as session:
# Send a GET request to the URL
async with session.get(url) as resp:
# Raise an error if the response status is not OK
resp.raise_for_status()
# Parse the response JSON
resp_json = await resp.json()
print("Fetch the data from URL.")
# Return the 'data' field from the JSON response
return resp_json['data']
except Exception as e:
# Print an error message if fetching data fails
print("Error fetching data from URL:", e)
return None
FRIENDLI_API_URL = "https://api.friendli.ai/serverless/v1/models"
FRIENDLI_PROVIDER = "friendliai"
INHERITABLE_BASE_KEYS = (
"supports_pdf_input",
"supports_assistant_prefill",
"supports_adaptive_thinking",
"supports_output_config",
)
REASONING_EFFORT_LEVEL_ORDER = ("none", "minimal", "low", "medium", "high", "xhigh", "max")
def _find_base_model_entry(base_model: str, local_data: dict) -> str | None:
if not base_model:
return None
bm_tail = base_model.split("/")[-1].lower()
if base_model in local_data:
return base_model
for key in local_data:
if key.startswith("sample_spec") or key == "fallback_generalizations":
continue
if key.split("/")[-1].lower() == bm_tail:
return key
return None
def _reasoning_effort_levels(reasoning_options: list) -> list:
offered = {
val
for opt in reasoning_options or []
if opt.get("type") == "effort"
for val in opt.get("values", [])
}
return [level for level in REASONING_EFFORT_LEVEL_ORDER if level in offered]
def _valid_token_price(value: object) -> bool:
try:
price = float(value) # pyright: ignore[reportArgumentType] # non-numeric values are rejected via the except
except (TypeError, ValueError):
return False
return math.isfinite(price) and price >= 0
def _has_valid_token_prices(pricing: dict | None) -> bool:
prices = pricing or {}
return _valid_token_price(prices.get("input")) and _valid_token_price(prices.get("output"))
def _pricing(pricing: dict) -> dict:
out: dict[str, Any] = {}
if not pricing:
return out
if "input" in pricing:
out["input_cost_per_token"] = float(pricing["input"])
if "output" in pricing:
out["output_cost_per_token"] = float(pricing["output"])
if "input_cache_read" in pricing and pricing["input_cache_read"] is not None:
out["cache_read_input_token_cost"] = float(pricing["input_cache_read"])
return out
def _modality_flags(input_mods: list) -> dict:
mods = input_mods or []
has_image = "image" in mods
return {
"supports_vision": has_image,
"supports_image_input": has_image,
"supports_video_input": "video" in mods,
}
def transform_friendli_data(data: list, local_data: dict) -> dict:
transformed: dict[str, dict] = {}
if not data:
return transformed
for model in data:
# An unpriced row must never wholesale-replace an already priced local entry:
# missing prices cost-calculate as zero, silently zeroing tracked spend
if not _has_valid_token_prices(model.get("pricing")):
continue
model_id = model["id"]
base_model = model.get("base_model") or ""
entry: dict[str, Any] = {
"litellm_provider": FRIENDLI_PROVIDER,
}
base_key = _find_base_model_entry(base_model, local_data)
if base_key:
base_entry = local_data[base_key]
for k in INHERITABLE_BASE_KEYS:
if k in base_entry:
entry[k] = base_entry[k]
ctx = model.get("context_length")
if ctx is not None:
entry["max_input_tokens"] = int(ctx)
max_out = model.get("max_completion_tokens")
if max_out is not None:
entry["max_output_tokens"] = int(max_out)
entry["max_tokens"] = int(max_out)
pricing = _pricing(model.get("pricing", {}))
entry.update(pricing)
entry["supports_prompt_caching"] = "cache_read_input_token_cost" in pricing
reasoning = model.get("reasoning") is True
entry["supports_reasoning"] = reasoning
if reasoning:
entry["reasoning_effort_levels"] = _reasoning_effort_levels(
model.get("reasoning_options", [])
)
func = model.get("functionality", {})
entry["supports_function_calling"] = func.get("tool_call") is True
entry["supports_parallel_function_calling"] = func.get("parallel_tool_call") is True
is_struct = func.get("structured_output") is True
entry["supports_response_schema"] = is_struct
entry["supports_native_structured_output"] = is_struct
entry["supports_system_messages"] = func.get("system_messages") is True
entry["supports_tool_choice"] = func.get("tool_choice") is True
entry.update(_modality_flags(model.get("input_modalities", [])))
entry["mode"] = model.get("mode", "chat")
desc = model.get("description")
if desc:
entry["comment"] = desc
dep = model.get("deprecation_date")
if dep:
entry["deprecation_date"] = dep.split("T")[0]
entry["source"] = FRIENDLI_API_URL
transformed[f"{FRIENDLI_PROVIDER}/{model_id}"] = entry
return transformed
# Synchronize local data with remote data
def sync_local_data_with_remote(local_data, remote_data, replace_keys=frozenset()):
# Update existing keys in local_data with values from remote_data
# (replace_keys entries are swapped wholesale so a field the remote catalog
# dropped, e.g. cache pricing, cannot survive as a stale value)
for key in (set(local_data) & set(remote_data)):
if key in replace_keys:
local_data[key] = remote_data[key]
else:
local_data[key].update(remote_data[key])
# Add new keys from remote_data to local_data
for key in (set(remote_data) - set(local_data)):
local_data[key] = remote_data[key]
# Write data to the json file
def write_to_file(file_path, data):
try:
# Open the file in write mode
with open(file_path, "w") as file:
# Dump the data as JSON into the file
json.dump(data, file, indent=4)
print("Values updated successfully.")
except Exception as e:
# Print an error message if writing to file fails
print("Error updating JSON file:", e)
# Update the existing models and add the missing models for OpenRouter
def transform_openrouter_data(data):
transformed = {}
if not data:
return transformed
for row in data:
# Add the fields 'max_tokens' and 'input_cost_per_token'
obj = {
"max_tokens": row["context_length"],
"input_cost_per_token": float(row["pricing"]["prompt"]),
}
# Add 'max_output_tokens' as a field if it is not None
if "top_provider" in row and "max_completion_tokens" in row["top_provider"] and row["top_provider"]["max_completion_tokens"] is not None:
obj['max_output_tokens'] = int(row["top_provider"]["max_completion_tokens"])
# Add the field 'output_cost_per_token'
obj.update({
"output_cost_per_token": float(row["pricing"]["completion"]),
})
# Add field 'input_cost_per_image' if it exists and is non-zero
if "pricing" in row and "image" in row["pricing"] and float(row["pricing"]["image"]) != 0.0:
obj['input_cost_per_image'] = float(row["pricing"]["image"])
# Add the fields 'litellm_provider' and 'mode'
obj.update({
"litellm_provider": "openrouter",
"mode": "chat"
})
# Add the 'supports_vision' field if the modality is 'multimodal'
if row.get('architecture', {}).get('modality') == 'multimodal':
obj['supports_vision'] = True
# Use a composite key to store the transformed object
transformed[f'openrouter/{row["id"]}'] = obj
return transformed
# Update the existing models and add the missing models for Vercel AI Gateway
def transform_vercel_ai_gateway_data(data):
transformed = {}
if not data:
return transformed
for row in data:
# Rows without token pricing or token limits (video/embedding models) previously KeyError'd the whole sync
if any(row.get(k) is None for k in ("context_window", "max_tokens")) or any(
row.get("pricing", {}).get(k) is None for k in ("input", "output")
):
continue
obj = {
"max_tokens": row["context_window"],
"input_cost_per_token": float(row["pricing"]["input"]),
"output_cost_per_token": float(row["pricing"]["output"]),
'max_output_tokens': row['max_tokens'],
'max_input_tokens': row["context_window"],
}
# Handle cache pricing if available
if "pricing" in row:
if "input_cache_read" in row["pricing"] and row["pricing"]["input_cache_read"] is not None:
obj['cache_read_input_token_cost'] = float(f"{float(row['pricing']['input_cache_read']):e}")
if "input_cache_write" in row["pricing"] and row["pricing"]["input_cache_write"] is not None:
obj['cache_creation_input_token_cost'] = float(f"{float(row['pricing']['input_cache_write']):e}")
mode = "embedding" if "embedding" in row["id"].lower() else "chat"
obj.update({"litellm_provider": "vercel_ai_gateway", "mode": mode})
transformed[f'vercel_ai_gateway/{row["id"]}'] = obj
return transformed
# Load local data from a specified file
def load_local_data(file_path):
try:
# Open the file in read mode
with open(file_path, "r") as file:
# Load and return the JSON data
return json.load(file)
except FileNotFoundError:
# Print an error message if the file is not found
print("File not found:", file_path)
return None
except json.JSONDecodeError as e:
# Print an error message if JSON decoding fails
print("Error decoding JSON:", e)
return None
def main():
local_file_path = "model_prices_and_context_window.json" # Path to the local data file
openrouter_url = "https://openrouter.ai/api/v1/models" # URL to fetch OpenRouter data
vercel_ai_gateway_url = "https://ai-gateway.vercel.sh/v1/models" # URL to fetch Vercel AI Gateway data
# Load local data from file
local_data = load_local_data(local_file_path)
# Fetch OpenRouter data
openrouter_data = asyncio.run(fetch_data(openrouter_url))
# Transform the fetched OpenRouter data
openrouter_data = transform_openrouter_data(openrouter_data)
# Fetch Vercel AI Gateway data
vercel_data = asyncio.run(fetch_data(vercel_ai_gateway_url))
# Transform the fetched Vercel AI Gateway data
vercel_data = transform_vercel_ai_gateway_data(vercel_data)
friendli_data = asyncio.run(fetch_data(FRIENDLI_API_URL))
friendli_data = transform_friendli_data(friendli_data, local_data)
# Combine both datasets
all_remote_data = {**openrouter_data, **vercel_data, **friendli_data}
# If both local and openrouter data are available, synchronize and save
if local_data and all_remote_data:
sync_local_data_with_remote(local_data, all_remote_data, replace_keys=frozenset(friendli_data))
write_to_file(local_file_path, local_data)
else:
print("Failed to fetch model data from either local file or URL.")
# Entry point of the script
if __name__ == "__main__":
main()

39
.github/workflows/auto_update_price_and_context_window.yml vendored
View file

@ -1,39 +0,0 @@
name: Updates model_prices_and_context_window.json and Create Pull Request
on:
schedule:
- cron: "0 0 * * 0" # Run every Sundays at midnight
#- cron: "0 0 * * *" # Run daily at midnight
permissions:
contents: write
pull-requests: write
jobs:
auto_update_price_and_context_window:
if: github.repository == 'BerriAI/litellm'
runs-on: ubuntu-latest
steps:
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
persist-credentials: false
- name: Set up uv
uses: ./.github/actions/setup-uv-with-retries
with:
version: "0.10.9"
- name: Update JSON Data
run: |
uv run --frozen --with 'aiohttp==3.13.3' python ".github/scripts/auto_update_price_and_context_window_file.py"
- name: Regenerate JSON Schema
run: |
uv run --frozen python ci_cd/generate_model_prices_schema.py
- name: Create Pull Request
run: |
git add model_prices_and_context_window.json model_prices_and_context_window.schema.json
git commit -m "Update model_prices_and_context_window.json file: $(date +'%Y-%m-%d')"
gh pr create --title "Update model_prices_and_context_window.json file" \
--body "Automated update for model_prices_and_context_window.json" \
--head auto-update-price-and-context-window-$(date +'%Y-%m-%d') \
--base main
env:
GH_TOKEN: ${{ secrets.GH_TOKEN }}

2
.github/workflows/create-rc-branch.yml vendored
View file

@ -2,7 +2,7 @@ name: Create RC Branch
on:
schedule:
- cron: "0 3 * * 5"
- cron: "0 3 * * 1"
timezone: "America/Los_Angeles"
workflow_dispatch:

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

@ -1,30 +0,0 @@
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

70
.github/workflows/sync-together-ai-models.yml vendored
View file

@ -1,70 +0,0 @@
name: Sync Together AI model registry
on:
schedule:
- cron: "30 6 * * *"
workflow_dispatch:
permissions:
contents: write
pull-requests: write
jobs:
sync_together_ai_models:
if: github.repository == 'BerriAI/litellm'
runs-on: ubuntu-latest
env:
BASE_BRANCH: ${{ github.event.repository.default_branch }}
steps:
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
ref: ${{ env.BASE_BRANCH }}
persist-credentials: false
- name: Set up uv
uses: ./.github/actions/setup-uv-with-retries
with:
version: "0.10.9"
- name: Look for an already-open sync PR
id: existing
run: |
open_pr="$(gh pr list --repo "$GITHUB_REPOSITORY" --state open --limit 1000 --json headRefName \
--jq '[.[].headRefName | select(startswith("litellm_together_registry_sync_"))] | first // empty')"
echo "open_pr=$open_pr" >> "$GITHUB_OUTPUT"
if [ -n "$open_pr" ]; then
echo "An open sync PR already exists on branch $open_pr; skipping this run."
fi
env:
GH_TOKEN: ${{ secrets.GH_TOKEN || github.token }}
- name: Run the sync
if: steps.existing.outputs.open_pr == ''
run: |
uv run --frozen python scripts/sync_together_ai_models.py --write --pr-body-file "$RUNNER_TEMP/pr_body.md"
env:
TOGETHER_API_KEY: ${{ secrets.TOGETHER_API_KEY }}
- name: Regenerate the JSON schema
if: steps.existing.outputs.open_pr == ''
run: |
uv run --frozen python ci_cd/generate_model_prices_schema.py
- name: Create a pull request when the registry changed
if: steps.existing.outputs.open_pr == ''
run: |
if git diff --quiet; then
echo "Registry already in sync; no PR needed."
exit 0
fi
branch="litellm_together_registry_sync_$(date +'%Y-%m-%d')"
git config user.name "github-actions[bot]"
git config user.email "41898282+github-actions[bot]@users.noreply.github.com"
git checkout -b "$branch"
git add model_prices_and_context_window.json \
litellm/model_prices_and_context_window_backup.json \
model_prices_and_context_window.schema.json
git commit -m "feat(models): sync together_ai model registry $(date +'%Y-%m-%d')"
gh auth setup-git
git push origin "$branch"
gh pr create --title "feat(models): sync together_ai model registry" \
--body-file "$RUNNER_TEMP/pr_body.md" \
--head "$branch" \
--base "$BASE_BRANCH"
env:
GH_TOKEN: ${{ secrets.GH_TOKEN || github.token }}

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

@ -376,10 +376,10 @@ jobs:
exit 0
fi
files="$(git diff --name-only "$base" HEAD)"
if grep -Eq '^(litellm/(proxy|types)/|ui/litellm-dashboard/(src/lib/http/schema\.d\.ts|scripts/gen-api-types\.mjs|package(-lock)?\.json)$|\.github/workflows/test-linting\.yml$)' <<< "$files"; then
if grep -Eq '^(litellm/|enterprise/|(pyproject\.toml|uv\.lock)$|ui/litellm-dashboard/(src/lib/http/schema\.d\.ts|scripts/gen-api-types\.mjs|package(-lock)?\.json)$|\.github/workflows/test-linting\.yml$)' <<< "$files"; then
echo "relevant=true" >> "$GITHUB_OUTPUT"
else
echo "No proxy, types or generator changes in this pull request, nothing to verify."
echo "No litellm, enterprise, dependency or generator changes in this pull request, nothing to verify."
echo "relevant=false" >> "$GITHUB_OUTPUT"
fi
@ -539,8 +539,8 @@ jobs:
- name: test_unit_passed_gate
run: uv run --no-sync pytest -q --noconftest -p no:cacheprovider -c /dev/null tests/code_coverage_tests/test_unit_passed_gate.py
- name: test_e2e_changed_gate
run: uv run --no-sync pytest -q --noconftest -p no:cacheprovider -c /dev/null tests/code_coverage_tests/test_e2e_changed_gate.py tests/code_coverage_tests/test_e2e_idp_stack.py
- name: test_e2e_idp_stack
run: uv run --no-sync pytest -q --noconftest -p no:cacheprovider -c /dev/null tests/code_coverage_tests/test_e2e_idp_stack.py
- name: test_e2e_metadata
env:

180
.github/workflows/test-mcp-oauth-e2e.yml vendored
View file

@ -1,180 +0,0 @@
name: MCP OAuth happy path
on:
pull_request:
paths:
- '.github/workflows/test-mcp-oauth-e2e.yml'
- '.github/e2e-stack/**'
- 'tests/e2e/*.py'
- 'tests/e2e/pytest.ini'
- 'tests/e2e/idp_realm.json'
- 'tests/e2e/mcp/**'
- 'litellm/experimental_mcp_client/**'
- 'litellm/proxy/_experimental/mcp_server/**'
- 'litellm/proxy/auth/**'
- 'litellm/proxy/management_endpoints/mcp_management_endpoints.py'
- 'litellm/proxy/_types.py'
- 'litellm/types/mcp_server/mcp_server_manager.py'
- 'litellm/proxy/management_endpoints/*sso*.py'
- 'litellm/proxy/management_endpoints/sso/**'
- 'litellm/proxy/common_utils/encrypt_decrypt_utils.py'
- 'litellm/proxy/proxy_server.py'
- 'litellm/proxy/schema.prisma'
- 'ui/litellm-dashboard/src/app/connect/**'
- 'ui/litellm-dashboard/src/app/mcp/oauth/**'
- 'pyproject.toml'
- 'uv.lock'
workflow_dispatch:
permissions: {}
concurrency:
group: mcp-oauth-${{ github.ref }}
cancel-in-progress: true
jobs:
oauth:
if: github.event_name == 'workflow_dispatch' || github.event.pull_request.head.repo.full_name == github.repository
runs-on: ubuntu-latest
environment: e2e-changed
timeout-minutes: 45
permissions:
contents: read
id-token: write
services:
postgres:
image: postgres:16.6
env:
POSTGRES_USER: litellm
POSTGRES_PASSWORD: dbpassword9090
POSTGRES_DB: litellm
ports:
- 5432:5432
options: >-
--health-cmd "pg_isready -U litellm"
--health-interval 5s
--health-timeout 5s
--health-retries 10
env:
DATABASE_HOST: 127.0.0.1
DATABASE_PORT: '5432'
DATABASE_USER: litellm
DATABASE_PASSWORD: dbpassword9090
DATABASE_NAME: litellm
DATABASE_URL: postgresql://litellm:dbpassword9090@127.0.0.1:5432/litellm
E2E_KEYCLOAK_URL: http://127.0.0.1:8081
E2E_KEYCLOAK_ADMIN_USER: admin
E2E_KEYCLOAK_ADMIN_PASSWORD: e2e-ephemeral-idp-not-a-secret
E2E_FIXTURE_MODE: live
E2E_PROVIDER_CACHE: '0'
E2E_MCP_OAUTH_LIVE: '1'
E2E_REQUIRED_TEST_COUNT: '4'
steps:
- name: Checkout the tested source
uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
ref: ${{ github.event.pull_request.head.sha || github.sha }}
persist-credentials: false
- name: Require and materialize the upstream login
env:
STORAGE_STATE: ${{ secrets.E2E_LINEAR_STORAGE_STATE_B64 }}
run: |
umask 077
python3 - <<'PY'
import base64
import json
import os
import secrets
from pathlib import Path
encoded = os.environ.get("STORAGE_STATE", "")
if not encoded:
raise SystemExit("E2E_LINEAR_STORAGE_STATE_B64 is required; capture and provision a test-account login")
state = json.loads(base64.b64decode(encoded, validate=True))
if not isinstance(state, dict) or not state.get("cookies"):
raise SystemExit("The captured login must contain browser cookies")
directory = Path(os.environ["RUNNER_TEMP"]) / "mcp-oauth-private"
directory.mkdir(mode=0o700)
path = directory / "linear-state.json"
path.write_text(json.dumps(state))
with open(os.environ["GITHUB_ENV"], "a") as output:
output.write(f"E2E_LINEAR_STORAGE_STATE={path}\n")
for name in ("LITELLM_MASTER_KEY", "LITELLM_SALT_KEY"):
value = "sk-e2e-" + secrets.token_hex(24)
print(f"::add-mask::{value}")
output.write(f"{name}={value}\n")
PY
- uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
with:
python-version: '3.13'
- uses: ./.github/actions/setup-uv-with-retries
with:
version: '0.10.9'
- uses: ./.github/actions/cache-cargo-build
- name: Install the frozen E2E environment
run: |
.github/scripts/uv_sync_with_retries.sh --frozen --extra proxy --extra proxy-runtime --extra extra_proxy --group ci --group proxy-dev --group e2e-dev
uv run --no-sync python scripts/prisma_generate_if_needed.py
uv run --no-sync playwright install --with-deps chromium
- name: Configure license access
id: aws
uses: aws-actions/configure-aws-credentials@e7f100cf4c008499ea8adda475de1042d6975c7b # v6.2.0
with:
role-to-assume: ${{ vars.E2E_AWS_ROLE_TO_ASSUME }}
aws-region: us-east-1
role-session-name: mcp-oauth-${{ github.run_id }}
role-duration-seconds: 900
output-env-credentials: false
output-credentials: true
- name: Load the E2E license
env:
AWS_ACCESS_KEY_ID: ${{ steps.aws.outputs.aws-access-key-id }}
AWS_SECRET_ACCESS_KEY: ${{ steps.aws.outputs.aws-secret-access-key }}
AWS_SESSION_TOKEN: ${{ steps.aws.outputs.aws-session-token }}
AWS_DEFAULT_REGION: us-east-1
run: |
license="$(aws secretsmanager get-secret-value --secret-id litellm-e2e-changed-license --query SecretString --output text)"
test -n "${license}"
echo "::add-mask::${license}"
echo "LITELLM_LICENSE=${license}" >> "${GITHUB_ENV}"
- uses: actions/setup-node@49933ea5288caeca8642d1e84afbd3f7d6820020 # v4.4.0
with:
node-version-file: ui/litellm-dashboard/.nvmrc
- name: Build the gateway consent UI at the tested commit
run: |
cd ui/litellm-dashboard
../../scripts/with_dashboard_node.sh npm ci
../../scripts/with_dashboard_node.sh npm run build
mkdir -p ../../litellm/proxy/_experimental/out
cp -r out/. ../../litellm/proxy/_experimental/out/
find ../../litellm/proxy/_experimental/out -name '*.html' ! -name index.html | while read -r page; do
mkdir -p "${page%.html}"
mv "${page}" "${page%.html}/index.html"
done
- name: Prepare the isolated database and IdP
run: |
umask 077
bash .github/e2e-stack/start-idp.sh
uv run --no-sync python migrations/run.py > "${RUNNER_TEMP}/mcp-oauth-private/migrations.log" 2>&1
- name: Run every required OAuth variant without retries
run: |
umask 077
uv run --no-sync pytest -c tests/e2e/pytest.ini tests/e2e/mcp/test_mcp_oauth_happy_path_e2e.py \
--rootdir=. --reruns 0 --tb=short -o junit_family=xunit1 \
--junitxml="${RUNNER_TEMP}/mcp-oauth-private/results.xml" \
> "${RUNNER_TEMP}/mcp-oauth-private/pytest.log" 2>&1
- name: Report JUnit results and reject skipped or missing cases
if: always()
run: |
uv run --no-sync python .github/e2e-stack/assert_tests_ran.py \
"${RUNNER_TEMP}/mcp-oauth-private/results.xml" tests/e2e/mcp/test_mcp_oauth_happy_path_e2e.py
- name: Remove private login and logs
if: always()
run: |
docker rm -f e2e-keycloak >/dev/null 2>&1 || true
rm -rf "${RUNNER_TEMP}/mcp-oauth-private"

1
.github/workflows/test-terraform-provider.yml vendored
View file

@ -12,6 +12,7 @@ on:
paths:
- "terraform/provider/**"
- "litellm/proxy/**"
- "enterprise/**"
- ".github/workflows/test-terraform-provider.yml"
permissions:

85
.github/workflows/weekly_load_anomaly.yml vendored
View file

@ -1,85 +0,0 @@
name: "Weekly Load Anomaly Check"
on:
schedule:
- cron: "0 12 * * 6"
workflow_dispatch:
permissions:
contents: read
jobs:
weekly-load-anomaly:
if: github.event_name != 'schedule' || github.repository == 'BerriAI/litellm'
runs-on: ubuntu-latest
timeout-minutes: 45
services:
postgres:
image: postgres:16.6
env:
POSTGRES_USER: llmproxy
POSTGRES_PASSWORD: dbpassword9090
POSTGRES_DB: litellm
ports:
- 5432:5432
options: >-
--health-cmd "pg_isready -U llmproxy"
--health-interval 5s
--health-timeout 5s
--health-retries 10
env:
DATABASE_URL: postgresql://llmproxy:dbpassword9090@localhost:5432/litellm
LITELLM_MASTER_KEY: sk-weekly-anomaly-check
ANTHROPIC_API_KEY: ${{ secrets.ANTHROPIC_API_KEY }}
AWS_BEARER_TOKEN_BEDROCK: ${{ secrets.AWS_BEARER_TOKEN_BEDROCK }}
steps:
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
persist-credentials: false
- name: Set up Python
uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
with:
python-version: "3.12"
- name: Set up uv
uses: ./.github/actions/setup-uv-with-retries
with:
version: "0.10.9"
- name: Cache the Rust build
uses: ./.github/actions/cache-cargo-build
- name: Install dependencies
run: |
.github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra proxy
- name: Cache Prisma binaries
uses: ./.github/actions/cache-prisma-binaries
- name: Generate Prisma client
run: |
uv run --no-sync prisma generate --schema litellm/proxy/schema.prisma
- name: Start the proxy
run: |
nohup uv run --no-sync litellm --config tests/e2e/perf/weekly_anomaly_config.yml --port 4000 > proxy.log 2>&1 &
for _ in $(seq 1 90); do
if curl -fs http://localhost:4000/health/liveliness > /dev/null; then
exit 0
fi
sleep 2
done
echo "proxy never became live"
tail -n 100 proxy.log
exit 1
- name: Run the weekly session anomaly test
env:
E2E_WEEKLY_ANOMALY: "1"
run: |
uv run --no-sync pytest tests/e2e/perf/test_weekly_session_anomaly_e2e.py -v --tb=short -rA
- name: Show proxy log on failure
if: failure()
run: tail -n 300 proxy.log

19
litellm/_internal_context.py
View file

@ -36,6 +36,8 @@ REDIS_FAMILIES_METADATA_KEY: Final = "families"
_service_caller: Final[ContextVar[str | None]] = ContextVar("service_caller", default=None)
_resent_on: Final[ContextVar[Callable[[Exception], bool] | None]] = ContextVar("resent_on", default=None)
@contextmanager
def post_response_phase() -> Generator[None]:
@ -113,6 +115,23 @@ def current_service_caller() -> str | None:
return _service_caller.get()
@contextmanager
def resent_on(rejection: Callable[[Exception], bool]) -> Generator[None]:
"""An attempt its caller sends again when it fails with ``rejection``. That failure is not the request's
outcome, so it is neither logged as the request's failure nor retried, and its budget reservation and rate-limit
slot stay with the attempt sent again."""
token: Final = _resent_on.set(rejection)
try:
yield
finally:
_restore(_resent_on, token)
def is_resent_failure(error: Exception) -> bool:
rejection: Final = _resent_on.get()
return rejection is not None and rejection(error)
@contextmanager
def pinned_billing_time(moment: datetime) -> Generator[None]:
"""Price every rate lookup inside this block at ``moment`` rather than at each one's own clock read."""

12
litellm/completion_extras/litellm_responses_transformation/handler.py
View file

@ -182,6 +182,9 @@ class ResponsesToCompletionBridgeHandler:
model_response: Final = validated_kwargs["model_response"]
logging_obj: Final = validated_kwargs["logging_obj"]
custom_llm_provider: Final = validated_kwargs["custom_llm_provider"]
from .transformation import serves_anthropic_messages
serves_messages_client: Final = serves_anthropic_messages(logging_obj)
if kwargs.get("stream") is True and "stream" not in optional_params:
optional_params = {**optional_params, "stream": True}
@ -222,6 +225,7 @@ class ResponsesToCompletionBridgeHandler:
encoding=validated_kwargs["encoding"],
api_key=kwargs.get("api_key"),
json_mode=kwargs.get("json_mode"),
keeps_every_reasoning_item=serves_messages_client,
)
elif isinstance(result, ModelResponse):
if not stream:
@ -247,6 +251,7 @@ class ResponsesToCompletionBridgeHandler:
encoding=validated_kwargs["encoding"],
api_key=kwargs.get("api_key"),
json_mode=kwargs.get("json_mode"),
keeps_every_reasoning_item=serves_messages_client,
)
else:
if self._is_preformatted_cached_chat_stream(result):
@ -260,6 +265,7 @@ class ResponsesToCompletionBridgeHandler:
structured_output=is_structured_output_request(
cast("Mapping[str, object]", optional_params) # cast-ok: optional_params is an untyped dict
),
streams_reasoning_signatures=serves_messages_client,
)
streamwrapper: Final = CustomStreamWrapper(
completion_stream=completion_stream,
@ -282,6 +288,9 @@ class ResponsesToCompletionBridgeHandler:
model_response: Final = validated_kwargs["model_response"]
logging_obj: Final = validated_kwargs["logging_obj"]
custom_llm_provider: Final = validated_kwargs["custom_llm_provider"]
from .transformation import serves_anthropic_messages
serves_messages_client: Final = serves_anthropic_messages(logging_obj)
if kwargs.get("stream") is True and "stream" not in optional_params:
optional_params = {**optional_params, "stream": True}
@ -325,6 +334,7 @@ class ResponsesToCompletionBridgeHandler:
encoding=validated_kwargs["encoding"],
api_key=kwargs.get("api_key"),
json_mode=kwargs.get("json_mode"),
keeps_every_reasoning_item=serves_messages_client,
)
elif isinstance(result, ModelResponse):
if not stream:
@ -350,6 +360,7 @@ class ResponsesToCompletionBridgeHandler:
encoding=validated_kwargs["encoding"],
api_key=kwargs.get("api_key"),
json_mode=kwargs.get("json_mode"),
keeps_every_reasoning_item=serves_messages_client,
)
else:
if self._is_preformatted_cached_chat_stream(result):
@ -363,6 +374,7 @@ class ResponsesToCompletionBridgeHandler:
structured_output=is_structured_output_request(
cast("Mapping[str, object]", optional_params) # cast-ok: optional_params is an untyped dict
),
streams_reasoning_signatures=serves_messages_client,
)
streamwrapper: Final = CustomStreamWrapper(
completion_stream=completion_stream,

136
litellm/completion_extras/litellm_responses_transformation/transformation.py
View file

@ -26,9 +26,11 @@ import litellm
from litellm import ModelResponse
from litellm._logging import verbose_logger
from litellm.integrations.anthropic_cache_control_hook import supports_openai_prompt_cache_breakpoint
from litellm.litellm_core_utils.core_helpers import as_str_mapping, attribute_or_none
from litellm.litellm_core_utils.hidden_params import get_hidden_params, get_or_create_hidden_params
from litellm.litellm_core_utils.prompt_templates.common_utils import (
responses_reasoning_items_from_thinking_blocks,
thinking_blocks_from_reasoning_items,
with_prompt_cache_breakpoint,
)
from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator
@ -45,6 +47,8 @@ from litellm.responses.utils import ResponsesAPIRequestUtils, normalize_response
from litellm.types.llms.openai import (
ChatCompletionAnnotation,
ChatCompletionReasoningItem,
ChatCompletionRedactedThinkingBlock,
ChatCompletionThinkingBlock,
ChatCompletionToolCallChunk,
ChatCompletionToolCallFunctionChunk,
ChatCompletionToolParamFunctionChunk,
@ -54,7 +58,7 @@ from litellm.types.llms.openai import (
ResponsesAPIResponse,
ResponsesAPIStreamEvents,
)
from litellm.types.utils import GenericStreamingChunk, ModelResponseStream
from litellm.types.utils import CallTypes, Delta, GenericStreamingChunk, ModelResponseStream, StreamingChoices
if TYPE_CHECKING:
from openai.types.responses import ResponseInputImageParam, ResponseOutputItem
@ -70,8 +74,6 @@ if TYPE_CHECKING:
AllMessageValues,
ChatCompletionFileObject,
ChatCompletionImageObject,
ChatCompletionRedactedThinkingBlock,
ChatCompletionThinkingBlock,
ChatCompletionToolReferenceObject,
OpenAIMessageContentListBlock,
)
@ -225,6 +227,10 @@ class _BuiltReasoningItem(TypedDict):
summary: Sequence[_ReasoningSummaryText]
def serves_anthropic_messages(logging_obj: object) -> bool:
return attribute_or_none(logging_obj, "call_type") == CallTypes.anthropic_messages.value
def _get_reasoning_items(
msg: "AllMessageValues",
) -> list[ChatCompletionReasoningItem]:
@ -443,6 +449,77 @@ def _as_chat_reasoning_items(
return cast(list[ChatCompletionReasoningItem], list(reasoning_items))
def _with_stream_block_index(
block: Mapping[str, object],
index: int,
summary_streamed: bool,
) -> ChatCompletionThinkingBlock | ChatCompletionRedactedThinkingBlock:
if block.get("type") == "redacted_thinking":
data: Final = block.get("data")
return ChatCompletionRedactedThinkingBlock(
type="redacted_thinking",
data=data if isinstance(data, str) else "",
index=index,
)
thinking: Final = block.get("thinking")
signature: Final = block.get("signature")
return ChatCompletionThinkingBlock(
type="thinking",
thinking=thinking if isinstance(thinking, str) and not summary_streamed else "",
signature=signature if isinstance(signature, str) else None,
index=index,
)
def _stream_output_index(chunk: Mapping[str, object]) -> int | None:
output_index: Final = chunk.get("output_index")
return output_index if isinstance(output_index, int) and not isinstance(output_index, bool) else None
def _reasoning_done_thinking_blocks(
item: Mapping[str, object],
output_index: int | None,
summary_streamed: bool,
) -> tuple[ChatCompletionThinkingBlock | ChatCompletionRedactedThinkingBlock, ...]:
index: Final = output_index if output_index is not None else 0
return tuple(
_with_stream_block_index(block, index, summary_streamed)
for block in thinking_blocks_from_reasoning_items((item,))
)
def _reasoning_done_stream_chunk(
chunk: Mapping[str, object],
summary_streamed_outputs: frozenset[int | None],
) -> ModelResponseStream | None:
event_type: Final = chunk.get("type")
event_name: Final = (
event_type.value
if isinstance(event_type, ResponsesAPIStreamEvents)
else event_type
if isinstance(event_type, str)
else None
)
if event_name != ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE:
return None
item: Final = as_str_mapping(chunk.get("item"))
if item is None or item.get("type") != "reasoning":
return None
output_index: Final = _stream_output_index(chunk)
signed_blocks: Final = _reasoning_done_thinking_blocks(item, output_index, output_index in summary_streamed_outputs)
if not signed_blocks:
return None
return ModelResponseStream(
choices=[
StreamingChoices(
index=0,
delta=Delta(thinking_blocks=list(signed_blocks)),
finish_reason=None,
)
]
)
def _shift_annotation(annotation: ChatCompletionAnnotation, offset: int) -> ChatCompletionAnnotation:
shifted: Final = {
key: value + offset if key in ("start_index", "end_index") and isinstance(value, int) else value
@ -476,7 +553,7 @@ def _fold_choices_into_one(
choices: Sequence["Choices"],
tool_calls: Sequence[Mapping[str, object]],
trailing_reasoning_content: str | None,
trailing_reasoning_item: _BuiltReasoningItem | None,
trailing_reasoning_items: Sequence[_BuiltReasoningItem],
tool_search_items: Sequence[Mapping[str, object]],
structured_output: bool = False,
) -> "Choices":
@ -495,7 +572,7 @@ def _fold_choices_into_one(
reasoning_contents: Final = (*(_message_reasoning_content(m) for m in messages), trailing_reasoning_content)
reasoning_items: Final = (
*chain.from_iterable(_message_reasoning_items(m) for m in messages),
*(() if trailing_reasoning_item is None else (trailing_reasoning_item,)),
*trailing_reasoning_items,
)
return Choices(
message=Message(
@ -1149,6 +1226,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
output_items: Sequence[object],
handle_raw_dict_callback: Callable[..., tuple["Choices | None", int]] | None = None,
structured_output: bool = False,
keeps_every_reasoning_item: bool = False,
) -> list["Choices"]:
"""
Convert Responses API output items to chat completion choices.
@ -1158,6 +1236,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
handle_raw_dict_callback: Optional callback for handling raw dict items
structured_output: The chat request carried a json response_format, so only the
first message item becomes content
keeps_every_reasoning_item: Keep each reasoning item before an answer, not only the last
Returns:
List of Choices objects
@ -1180,7 +1259,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
choices: Final[list[Choices]] = []
index = 0
reasoning_content: str | None = None
pending_reasoning_item: _BuiltReasoningItem | None = None
pending_reasoning_items: tuple[_BuiltReasoningItem, ...] = ()
# Collect all tool calls to put them in a single choice
# (Chat Completions API expects all tool calls in one message)
@ -1189,12 +1268,19 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
for item in output_items:
if isinstance(item, ResponseReasoningItem):
pending_reasoning_item = _build_reasoning_item(
item_id=item.id,
encrypted_content=getattr(item, "encrypted_content", None),
summary_raw=item.summary,
pending_reasoning_items = (
*(pending_reasoning_items if keeps_every_reasoning_item else ()),
_build_reasoning_item(
item_id=item.id,
encrypted_content=getattr(item, "encrypted_content", None),
summary_raw=item.summary,
),
)
reasoning_content = " ".join(
s["text"]
for s in chain.from_iterable(built["summary"] for built in pending_reasoning_items)
if s.get("text")
)
reasoning_content = " ".join(s["text"] for s in pending_reasoning_item["summary"] if s.get("text"))
elif isinstance(item, ResponseOutputMessage):
for content in item.content:
@ -1211,7 +1297,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
annotations=annotations,
reasoning_items=cast(
list[ChatCompletionReasoningItem] | None,
([pending_reasoning_item] if pending_reasoning_item is not None else None),
(list(pending_reasoning_items) or None),
),
)
@ -1224,7 +1310,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
)
reasoning_content = None # flush
pending_reasoning_item = None # flush
pending_reasoning_items = () # flush
index += 1
elif isinstance(item, ResponseFunctionToolCall):
@ -1289,7 +1375,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
choices=choices,
tool_calls=accumulated_tool_calls,
trailing_reasoning_content=reasoning_content,
trailing_reasoning_item=pending_reasoning_item,
trailing_reasoning_items=pending_reasoning_items,
tool_search_items=tool_search_items,
structured_output=structured_output,
)
@ -1396,6 +1482,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
encoding: object,
api_key: str | None = None,
json_mode: bool | None = None,
keeps_every_reasoning_item: bool = False,
) -> "ModelResponse":
"""Transform Responses API response to chat completion response"""
from litellm.responses.utils import ResponseAPILoggingUtils
@ -1427,6 +1514,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
structured_output=is_structured_output_request(
cast("Mapping[str, object]", optional_params) # cast-ok: optional_params is an untyped dict
),
keeps_every_reasoning_item=keeps_every_reasoning_item,
)
response_is_incomplete: Final = raw_response.status == "incomplete" or (
@ -1490,9 +1578,14 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
sync_stream: bool,
json_mode: bool | None = False,
structured_output: bool = False,
streams_reasoning_signatures: bool = False,
) -> BaseModelResponseIterator:
return OpenAiResponsesToChatCompletionStreamIterator(
streaming_response, sync_stream, json_mode, structured_output=structured_output
streaming_response,
sync_stream,
json_mode,
structured_output=structured_output,
streams_reasoning_signatures=streams_reasoning_signatures,
)
def _convert_content_str_to_input_text(self, content: str, role: str) -> dict[str, object]:
@ -1871,13 +1964,16 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator):
sync_stream: bool,
json_mode: bool | None = False,
structured_output: bool = False,
streams_reasoning_signatures: bool = False,
):
super().__init__(streaming_response, sync_stream, json_mode)
self._streams_reasoning_signatures: Final = streams_reasoning_signatures
self._chat_completion_id: str | None = None
self._served_service_tier: str | None = None
self._tool_call_index_map: dict[int, int] = {} # mutable-ok: per-stream accumulator state
self._structured_output: Final = structured_output
self._first_message_item_done: bool = False
self._summary_streamed_outputs: frozenset[int | None] = frozenset()
def _handle_string_chunk(
self, str_line: Union[str, "BaseModel"]
@ -2089,7 +2185,6 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator):
)
]
)
elif event_type == "response.output_text.delta":
# Content part added to output
content_part = parsed_chunk.get("delta", None)
@ -2193,15 +2288,24 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator):
"""
verbose_logger.debug("Chat provider: transform_streaming_response called with chunk: %s", chunk)
self._remember_served_service_tier(chunk)
reasoning_chunk: Final = (
_reasoning_done_stream_chunk(chunk, self._summary_streamed_outputs)
if self._streams_reasoning_signatures
else None
)
translated: Final = (
self._empty_content_chunk()
if self._is_text_delta_after_structured_answer(chunk)
else reasoning_chunk
if reasoning_chunk is not None
else OpenAiResponsesToChatCompletionStreamIterator.translate_responses_chunk_to_openai_stream(
chunk, tool_call_index_map=self._tool_call_index_map
)
)
if self._structured_output and _is_message_item_done(chunk):
self._first_message_item_done = True
if _responses_event_type(chunk) == "response.reasoning_summary_text.delta" and chunk.get("delta"):
self._summary_streamed_outputs = self._summary_streamed_outputs | {_stream_output_index(chunk)}
return self._with_served_service_tier(self._with_stream_scoped_id(translated))
def _is_text_delta_after_structured_answer(self, chunk: Mapping[str, object]) -> bool:

3
litellm/constants.py
View file

@ -237,6 +237,9 @@ MCP_PER_USER_TOKEN_EXPIRY_BUFFER_SECONDS: Final = get_env_int("MCP_PER_USER_TOKE
# MCP timeout defaults (seconds). Override via env vars for slow/custom MCP servers.
MCP_CLIENT_TIMEOUT: Final = get_env_float("LITELLM_MCP_CLIENT_TIMEOUT", 60.0)
MCP_MAX_SSE_EVENT_SIZE_BYTES: Final = get_env_int_in_range(
"LITELLM_MCP_MAX_SSE_EVENT_SIZE_BYTES", 64 * 1024 * 1024, 1, sys.maxsize
)
MCP_TOOL_LISTING_TIMEOUT: Final = get_env_float("LITELLM_MCP_TOOL_LISTING_TIMEOUT", 30.0)
MCP_METADATA_TIMEOUT: Final = get_env_float("LITELLM_MCP_METADATA_TIMEOUT", 10.0)
MCP_HEALTH_CHECK_TIMEOUT: Final = get_env_float("LITELLM_MCP_HEALTH_CHECK_TIMEOUT", 10.0)

2
litellm/experimental_mcp_client/client.py
View file

@ -95,6 +95,7 @@ from pydantic import AnyUrl, TypeAdapter
from litellm._logging import verbose_logger
from litellm.constants import (
MCP_CLIENT_TIMEOUT,
MCP_MAX_SSE_EVENT_SIZE_BYTES,
MCP_NPM_CACHE_DIR,
MCP_TOOL_LISTING_MAX_PAGES,
MCP_TOOL_LISTING_TIMEOUT,
@ -631,6 +632,7 @@ class MCPClient:
transport_ctx: Final = streamable_http_client(
url=self.server_url,
http_client=http_client,
max_sse_event_size=MCP_MAX_SSE_EVENT_SIZE_BYTES,
)
return transport_ctx, http_client

132
litellm/integrations/custom_guardrail.py
View file

@ -1,6 +1,7 @@
import contextvars
import copy
import hashlib
import json
import os
import secrets
from collections.abc import Iterable, Mapping, Sequence
@ -9,6 +10,7 @@ from types import MappingProxyType
from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, Optional, cast, get_args
import httpx
from pydantic import BaseModel, TypeAdapter
from litellm._internal_context import with_service_target
from litellm._logging import verbose_logger
@ -71,6 +73,7 @@ GUARDRAIL_SESSIONS_TARGET: Final = "guardrail_sessions"
_PRE_CALL_EXECUTED_TOKEN: Final = secrets.token_hex(16)
_PRE_CALL_RAISED_TOKEN: Final = secrets.token_hex(16)
_DEPLOYMENT_HOOK_RAN_TOKEN: Final = secrets.token_hex(16)
_POST_CALL_CHECKED_TOKEN: Final = secrets.token_hex(16)
_GUARDRAIL_BLOCK_STATUS_CODES: Final = frozenset({400, 403, 422})
@ -239,6 +242,98 @@ def _request_carries_marker(data: Mapping[str, object], marker: str | None) -> b
return False
def _record_guardrail_marker(
data: dict[str, object], # mutable-ok: the request payload the marker is recorded on
marker: str | None,
) -> None:
if marker is None:
return
for meta_key in ("metadata", "litellm_metadata"):
meta = data.get(meta_key)
if isinstance(meta, dict):
executed = meta.get(PRE_CALL_EXECUTED_GUARDRAILS_KEY)
if isinstance(executed, list):
if marker not in executed:
executed.append(marker)
else:
meta[PRE_CALL_EXECUTED_GUARDRAILS_KEY] = [marker]
return
data["metadata"] = {PRE_CALL_EXECUTED_GUARDRAILS_KEY: [marker]}
_RESPONSE_METADATA_KEYS: Final = frozenset(
{
"id",
"created",
"created_at",
"model",
"usage",
"usage_info",
"system_fingerprint",
"_hidden_params",
"_response_ms",
}
)
def _response_fingerprint(response: object) -> str | None:
"""Digest of everything but metadata set between layers, so an unknown field counts as content."""
try:
fields: Final[Mapping[str, object]] = (
response.model_dump(mode="json")
if isinstance(response, BaseModel)
else _REQUEST_FIELDS.validate_python(response)
)
content: Final = {key: value for key, value in fields.items() if key not in _RESPONSE_METADATA_KEYS}
if not content:
return None
return hashlib.sha256(
json.dumps(content, sort_keys=True, separators=(",", ":"), allow_nan=False).encode()
).hexdigest()
except (TypeError, ValueError):
return None
_RESPONSE_GUARDRAIL_CHECKS_KEY: Final = "guardrail_checks"
def _str_keyed_dict(value: object) -> Mapping[str, object] | None:
if not isinstance(value, dict):
return None
return cast("Mapping[str, object]", value) # cast-ok: response bodies and _hidden_params are keyed by field name
def _response_hidden_params(response: object) -> Mapping[str, object] | None:
as_dict: Final = _str_keyed_dict(response)
return _str_keyed_dict(
as_dict.get("_hidden_params") if as_dict is not None else getattr(response, "_hidden_params", None)
)
def _write_response_marker(response: object, marker: str | None) -> None:
"""Record a check on the response itself, so it survives the router copying request metadata per hop."""
if marker is None:
return
hidden: Final = _response_hidden_params(response)
if not isinstance(hidden, dict):
return
hidden[_RESPONSE_GUARDRAIL_CHECKS_KEY] = tuple(
_str_collection(hidden.get(_RESPONSE_GUARDRAIL_CHECKS_KEY)) | {marker}
)
def _response_carries_marker(response: object, marker: str | None) -> bool:
hidden: Final = _response_hidden_params(response)
return (
marker is not None
and hidden is not None
and marker in _str_collection(hidden.get(_RESPONSE_GUARDRAIL_CHECKS_KEY))
)
_REQUEST_FIELDS: Final = TypeAdapter(Mapping[str, object])
class CustomGuardrail(CustomLogger):
# If True, during_call runs async_moderation_hook instead of the unified apply_guardrail path.
use_native_during_call_hook: ClassVar[bool] = False
@ -867,20 +962,7 @@ class CustomGuardrail(CustomLogger):
With ``raised``, records instead that the hook raised as a policy pipeline
step the pipeline then moved past. The deployment-level hook still runs it.
"""
marker: Final = self._pre_call_marker(raised=raised)
if marker is None:
return
for meta_key in ("metadata", "litellm_metadata"):
meta = data.get(meta_key)
if isinstance(meta, dict):
executed = meta.get(PRE_CALL_EXECUTED_GUARDRAILS_KEY)
if isinstance(executed, list):
if marker not in executed:
executed.append(marker)
else:
meta[PRE_CALL_EXECUTED_GUARDRAILS_KEY] = [marker]
return
data["metadata"] = {PRE_CALL_EXECUTED_GUARDRAILS_KEY: [marker]}
_record_guardrail_marker(data, self._pre_call_marker(raised=raised))
def pre_call_hook_ran_or_raised(self, data: Mapping[str, object]) -> bool:
"""Whether this guardrail's ``async_pre_call_hook`` ran for this request, whatever it decided."""
@ -1077,10 +1159,26 @@ class CustomGuardrail(CustomLogger):
cast("dict[str, object]", request_data) # cast-ok: proxy request payloads are str-keyed
)
if not self._is_valid_response_type(result):
return None
returned: Final = result if self._is_valid_response_type(result) else None
if target is self:
return returned
checked: Final = returned if returned is not None else response
marker: Final = self._post_call_check_marker(checked)
_record_guardrail_marker(
cast("dict[str, object]", request_data), # cast-ok: proxy request payloads are str-keyed
marker,
)
_write_response_marker(checked, marker)
return returned
return result
def _post_call_check_marker(self, response: object) -> str | None:
fingerprint: Final = _response_fingerprint(response)
return self._marker(f"{_POST_CALL_CHECKED_TOKEN}:{fingerprint}") if fingerprint is not None else None
def post_call_already_checked(self, data: Mapping[str, object], response: object) -> bool:
"""Whether a deployment-level hook of this request already ran this guardrail on this exact response."""
marker: Final = self._post_call_check_marker(response)
return _request_carries_marker(data, marker) or _response_carries_marker(response, marker)
async def async_logging_hook(
self,

8
litellm/litellm_core_utils/get_llm_provider_logic.py
View file

@ -158,6 +158,14 @@ def declared_authenticating_provider(model: str | None, custom_llm_provider: str
return declared if declared in PROVIDERS_THAT_AUTHENTICATE_ON_PROVIDER_INFO else None
def resolve_model_and_provider_for_metadata(model: str, custom_llm_provider: str | None = None) -> tuple[str, str]:
declared: Final = declared_authenticating_provider(model, custom_llm_provider)
if declared is not None:
return model.removeprefix(f"{declared}/"), declared
resolved_model, resolved_provider, _, _ = get_llm_provider(model=model, custom_llm_provider=custom_llm_provider)
return resolved_model, resolved_provider
def inferred_provider(model: str | None) -> str | None:
if not model:
return None

2
litellm/litellm_core_utils/litellm_logging.py
View file

@ -4593,6 +4593,7 @@ class Logging(LiteLLMLoggingBaseClass):
from litellm.completion_extras.litellm_responses_transformation.transformation import (
LiteLLMResponsesTransformationHandler,
responses_provider_metadata,
serves_anthropic_messages,
)
served_id: Final = _provider_response_id(result)
@ -4607,6 +4608,7 @@ class Logging(LiteLLMLoggingBaseClass):
optional_params={},
litellm_params={},
encoding=litellm.encoding,
keeps_every_reasoning_item=serves_anthropic_messages(self),
)
translated.id = served_id or translated.id
return translated

92
litellm/litellm_core_utils/prompt_templates/common_utils.py
View file

@ -20,6 +20,7 @@ from pydantic import TypeAdapter
import litellm
from litellm import verbose_logger
from litellm.litellm_core_utils.core_helpers import as_str_mapping
from litellm.router_utils.batch_utils import InMemoryFile
from litellm.types.llms.openai import (
AllMessageValues,
@ -2289,6 +2290,89 @@ def encrypted_content_of_block(block: Mapping[str, object]) -> str | None:
return encrypted_content_from_signature(_encrypted_reasoning_field(block))
_REASONING_SUMMARY_SEPARATOR: Final = "\n\n"
_OBJECT_TUPLE: Final = TypeAdapter(tuple[object, ...])
def _object_sequence(value: object) -> tuple[object, ...]:
if not isinstance(value, Sequence) or isinstance(value, (str, bytes)):
return ()
return _OBJECT_TUPLE.validate_python(value)
def _iter_mappings(value: object) -> tuple[Mapping[str, object], ...]:
return tuple(mapping for item in _object_sequence(value) if (mapping := as_str_mapping(item)) is not None)
def _reasoning_summary_part_text(part: object) -> str:
mapping: Final = as_str_mapping(part)
if mapping is not None:
text: Final = mapping.get("text")
return text if isinstance(text, str) else ""
text_attr: Final = getattr(part, "text", None)
return text_attr if isinstance(text_attr, str) else ""
def _reasoning_item_summary_text(item: Mapping[str, object]) -> str:
parts: Final = tuple(_reasoning_summary_part_text(part) for part in _object_sequence(item.get("summary")))
return _REASONING_SUMMARY_SEPARATOR.join(part for part in parts if part)
def thinking_block_from_reasoning_item(
item: object,
) -> ChatCompletionThinkingBlock | ChatCompletionRedactedThinkingBlock | None:
"""The chat thinking block that carries one Responses `reasoning` item back to an Anthropic client.
`encrypted_content` rides the same `litellm_encrypted_reasoning:` signature the Responses
adapter uses. A summary becomes thinking text. Bytes with no summary become `redacted_thinking`,
because an empty thinking string is not something the client can replay.
"""
mapping: Final = as_str_mapping(item)
if mapping is None or mapping.get("type") not in (None, "reasoning"):
return None
encrypted_content: Final = mapping.get("encrypted_content")
if not isinstance(encrypted_content, str) or not encrypted_content:
return None
signature: Final = encrypted_reasoning_signature(encrypted_content)
text: Final = _reasoning_item_summary_text(mapping)
if not text:
return ChatCompletionRedactedThinkingBlock(type="redacted_thinking", data=signature)
return ChatCompletionThinkingBlock(type="thinking", thinking=text, signature=signature)
def thinking_blocks_from_reasoning_items(
items: object,
) -> tuple[ChatCompletionThinkingBlock | ChatCompletionRedactedThinkingBlock, ...]:
return tuple(
block for item in _object_sequence(items) if (block := thinking_block_from_reasoning_item(item)) is not None
)
def _block_encrypted_content(block: object) -> str | None:
mapping: Final = as_str_mapping(block)
if mapping is None:
return None
return encrypted_content_of_block(mapping)
def thinking_blocks_with_encrypted_reasoning(
thinking_blocks: object,
reasoning_items: object,
) -> tuple[Mapping[str, object], ...]:
"""Thinking blocks already on the message, plus any encrypted reasoning item they do not already carry."""
direct: Final = _iter_mappings(thinking_blocks)
represented: Final = frozenset(
content for block in direct if (content := _block_encrypted_content(block)) is not None
)
extra: Final = tuple(
block
for block in thinking_blocks_from_reasoning_items(_iter_mappings(reasoning_items))
if _block_encrypted_content(block) not in represented
)
return (*extra, *direct)
def is_encrypted_reasoning_block(block: object) -> bool:
"""A thinking or redacted_thinking block carrying Responses API encrypted reasoning.
@ -2380,6 +2464,14 @@ def anthropic_content_lists(messages: Sequence[object]) -> Iterator[object]:
)
def _content_carries_encrypted_reasoning(content: object) -> bool:
return any(is_encrypted_reasoning_block(block) for block in _object_sequence(content))
def messages_carry_encrypted_reasoning(messages: Sequence[object]) -> bool:
return any(_content_carries_encrypted_reasoning(content) for content in anthropic_content_lists(messages))
def _strip_encrypted_reasoning_from_blocks(
content: object,
*,

56
litellm/litellm_core_utils/prompt_templates/factory.py
View file

@ -4599,6 +4599,16 @@ def get_assistant_message_block_or_continue_message(
raise ValueError(f"Unsupported content type: {type(content_block)}")
def _as_thinking_block(item: object) -> ChatCompletionThinkingBlock | ChatCompletionRedactedThinkingBlock | None:
match item:
case {"type": "thinking" | "redacted_thinking"}:
return cast( # cast-ok: the mapping pattern matched a thinking block's type tag
ChatCompletionThinkingBlock | ChatCompletionRedactedThinkingBlock, item
)
case _:
return None
class BedrockConverseMessagesProcessor:
@staticmethod
def _initial_message_setup(
@ -4647,11 +4657,15 @@ class BedrockConverseMessagesProcessor:
deployment_model: str | None = None,
base_model: str | None = None,
) -> list[BedrockMessageBlock]:
from litellm.llms.bedrock.common_utils import bedrock_converse_accepts_redacted_thinking
from litellm.llms.bedrock.common_utils import (
bedrock_converse_accepts_redacted_thinking,
bedrock_converse_rejects_signed_thinking,
)
contents: list[BedrockMessageBlock] = []
msg_i = 0
to_claude: Final = bedrock_converse_accepts_redacted_thinking(model, base_model)
drop_signed_thinking: Final = bedrock_converse_rejects_signed_thinking(model, base_model)
messages = BedrockConverseMessagesProcessor._initial_message_setup(
messages, model, llm_provider, user_continue_message
@ -4803,7 +4817,7 @@ class BedrockConverseMessagesProcessor:
if thinking_blocks is not None:
converted_thinking_blocks = (
BedrockConverseMessagesProcessor.translate_thinking_blocks_to_reasoning_content_blocks(
thinking_blocks, to_claude=to_claude
thinking_blocks, to_claude=to_claude, drop_signed=drop_signed_thinking
)
)
assistant_content = BedrockConverseMessagesProcessor.add_thinking_blocks_to_assistant_content(
@ -4817,7 +4831,7 @@ class BedrockConverseMessagesProcessor:
if isinstance(element, dict):
if element["type"] == "thinking" or element["type"] == "redacted_thinking":
thinking_block = BedrockConverseMessagesProcessor.translate_thinking_blocks_to_reasoning_content_blocks(
thinking_blocks=[element], to_claude=to_claude
thinking_blocks=[element], to_claude=to_claude, drop_signed=drop_signed_thinking
)
assistants_parts = (
BedrockConverseMessagesProcessor.add_thinking_blocks_to_assistant_content(
@ -4844,7 +4858,9 @@ class BedrockConverseMessagesProcessor:
model=model,
deployment_model=deployment_model,
)
if _cache_point_block is not None:
if _cache_point_block is not None and BedrockConverseMessagesProcessor.keeps_cache_point(
element, drop_signed=drop_signed_thinking
):
assistants_parts.append(_cache_point_block)
assistant_content.extend(assistants_parts)
elif _assistant_content is not None and isinstance(_assistant_content, str):
@ -4892,6 +4908,7 @@ class BedrockConverseMessagesProcessor:
thinking_blocks: Sequence[ChatCompletionThinkingBlock | ChatCompletionRedactedThinkingBlock],
*,
to_claude: bool = False,
drop_signed: bool = False,
) -> list[BedrockContentBlock]:
return [
BedrockContentBlock(
@ -4901,8 +4918,25 @@ class BedrockConverseMessagesProcessor:
)
for thinking_block in thinking_blocks
if not (to_claude and is_encrypted_reasoning_block(thinking_block))
and not (drop_signed and BedrockConverseMessagesProcessor._is_signed(thinking_block))
]
@staticmethod
def keeps_cache_point(element: object, *, drop_signed: bool) -> bool:
match _as_thinking_block(element):
case None:
return True
case thinking_block:
return not (drop_signed and BedrockConverseMessagesProcessor._is_signed(thinking_block))
@staticmethod
def _is_signed(thinking_block: ChatCompletionThinkingBlock | ChatCompletionRedactedThinkingBlock) -> bool:
match thinking_block:
case {"signature": signature} if signature:
return True
case _:
return False
@staticmethod
def _to_reasoning_content_block(
thinking_block: ChatCompletionThinkingBlock | ChatCompletionRedactedThinkingBlock,
@ -5059,11 +5093,15 @@ def _bedrock_converse_messages_pt(
- Please ensure that function response turn comes immediately after a function call turn
- Conversation blocks and tool result blocks cannot be provided in the same turn. Issue: https://github.com/BerriAI/litellm/issues/6053
"""
from litellm.llms.bedrock.common_utils import bedrock_converse_accepts_redacted_thinking
from litellm.llms.bedrock.common_utils import (
bedrock_converse_accepts_redacted_thinking,
bedrock_converse_rejects_signed_thinking,
)
contents: list[BedrockMessageBlock] = []
msg_i = 0
to_claude: Final = bedrock_converse_accepts_redacted_thinking(model, base_model)
drop_signed_thinking: Final = bedrock_converse_rejects_signed_thinking(model, base_model)
messages = BedrockConverseMessagesProcessor._initial_message_setup(
messages, model, llm_provider, user_continue_message
@ -5216,7 +5254,7 @@ def _bedrock_converse_messages_pt(
if thinking_blocks is not None:
converted_thinking_blocks = (
BedrockConverseMessagesProcessor.translate_thinking_blocks_to_reasoning_content_blocks(
thinking_blocks, to_claude=to_claude
thinking_blocks, to_claude=to_claude, drop_signed=drop_signed_thinking
)
)
assistant_content = BedrockConverseMessagesProcessor.add_thinking_blocks_to_assistant_content(
@ -5231,7 +5269,7 @@ def _bedrock_converse_messages_pt(
if element["type"] == "thinking" or element["type"] == "redacted_thinking":
thinking_block = (
BedrockConverseMessagesProcessor.translate_thinking_blocks_to_reasoning_content_blocks(
thinking_blocks=[element], to_claude=to_claude
thinking_blocks=[element], to_claude=to_claude, drop_signed=drop_signed_thinking
)
)
assistants_parts = (
@ -5260,7 +5298,9 @@ def _bedrock_converse_messages_pt(
model=model,
deployment_model=deployment_model,
)
if _cache_point_block is not None:
if _cache_point_block is not None and BedrockConverseMessagesProcessor.keeps_cache_point(
element, drop_signed=drop_signed_thinking
):
assistants_parts.append(_cache_point_block)
assistant_content.extend(assistants_parts)
elif _assistant_content is not None and isinstance(_assistant_content, str):

1
litellm/litellm_core_utils/streaming_handler.py
View file

@ -912,6 +912,7 @@ class CustomStreamWrapper:
and model_response.choices[0].delta.role is not None
)
or (getattr(model_response.choices[0].delta, "reasoning_items", None) is not None)
or getattr(model_response.choices[0].delta, "thinking_blocks", None)
):
return True
else:

227
litellm/llms/anthropic/pass_through/adapters/handler.py
View file

@ -1,5 +1,8 @@
from collections.abc import AsyncIterator, Coroutine, Iterator, Mapping, Sequence
import re
from collections.abc import AsyncIterator, Callable, Coroutine, Iterator, Mapping, Sequence
from contextlib import AbstractContextManager, nullcontext
from dataclasses import dataclass
from functools import partial
from typing import (
TYPE_CHECKING,
Final,
@ -11,16 +14,25 @@ from pydantic import TypeAdapter
from typing_extensions import TypedDict
import litellm
from litellm._internal_context import resent_on
from litellm._logging import verbose_logger
from litellm.litellm_core_utils.asyncify import run_async_function
from litellm.litellm_core_utils.core_helpers import as_str_mapping, iter_items
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObject
from litellm.litellm_core_utils.model_param_helper import peek_reasoning_summary_aliases
from litellm.litellm_core_utils.prompt_templates.common_utils import (
is_encrypted_reasoning_block,
thinking_blocks_with_encrypted_reasoning,
)
from litellm.litellm_core_utils.tool_fields import (
ROUTES_WITHOUT_MCP_SERVERS,
drops_mcp_servers,
tools_without_mcp_servers,
)
from litellm.llms.anthropic.common_utils import strip_encrypted_reasoning_blocks_from_anthropic_messages
from litellm.llms.anthropic.pass_through.adapters.transformation import (
AnthropicAdapter,
openai_thinking_routes_to_responses,
)
from litellm.llms.anthropic.pass_through.context_management import (
AnthropicContextManagementError,
@ -48,9 +60,9 @@ from litellm.types.llms.anthropic_messages.anthropic_response import (
)
from litellm.types.llms.openai import OpenAIWebSearchOptions
from litellm.types.utils import ModelResponse
from litellm.utils import get_model_info
if TYPE_CHECKING:
from litellm.chat_completions.dispatch import ChatResult
from litellm.router import Router
from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth
@ -68,22 +80,156 @@ class _CompletionKwargs(TypedDict, total=False, extra_items=object):
custom_llm_provider: str
_ANTHROPIC_MESSAGES: Final = TypeAdapter(list[dict[str, object]])
_UNVERIFIED_ENCRYPTED_CONTENT: Final = re.compile(r"encrypted content \S+ could not be verified", re.IGNORECASE)
_PrepareReplay: TypeAlias = "Callable[[_AnthropicMessages], _CompletionKwargs]"
@dataclass(frozen=True, slots=True)
class _ChatTargetReasoning:
effort: str
summary: str | None = None
def _provider_error_code(body: object) -> str | None:
mapping: Final = as_str_mapping(body)
if mapping is None:
return None
nested: Final = as_str_mapping(mapping.get("error"))
source: Final = nested if nested is not None else mapping
code: Final = source.get("code")
return code if isinstance(code, str) else None
def _logged_rejection_code(error: litellm.BadRequestError) -> str:
if error.code == "invalid_encrypted_content" or _provider_error_code(error.body) == "invalid_encrypted_content":
return "invalid_encrypted_content"
if "invalid_encrypted_content" in str(error):
return "invalid_encrypted_content"
return "could_not_be_verified"
def _completion_text(completion_kwargs: _CompletionKwargs, key: str) -> str:
value: Final = completion_kwargs.get(key)
return value if isinstance(value, str) else ""
def _message_carries_encrypted_reasoning(message: object) -> bool:
mapping: Final = as_str_mapping(message)
if mapping is None:
return False
blocks: Final = thinking_blocks_with_encrypted_reasoning(
mapping.get("thinking_blocks"), mapping.get("reasoning_items")
)
return any(is_encrypted_reasoning_block(block) for block in blocks)
def _sends_encrypted_reasoning(completion_kwargs: _CompletionKwargs) -> bool:
messages: Final = completion_kwargs.get("messages")
sent: Final = iter_items(messages) if messages is not None else ()
return any(_message_carries_encrypted_reasoning(message) for message in sent)
def _rejects_encrypted_reasoning(error: litellm.BadRequestError) -> bool:
error_message: Final = str(error)
return (
"invalid_encrypted_content" in error_message or _UNVERIFIED_ENCRYPTED_CONTENT.search(error_message) is not None
)
def _is_rejected_reasoning(error: Exception) -> bool:
return isinstance(error, litellm.BadRequestError) and _rejects_encrypted_reasoning(error)
def resent_without_rejected_reasoning(sent_encrypted_reasoning: bool) -> AbstractContextManager[None]:
return resent_on(_is_rejected_reasoning) if sent_encrypted_reasoning else nullcontext()
def messages_without_rejected_reasoning(
messages: Sequence[Mapping[str, object]],
error: litellm.BadRequestError,
*,
sent_encrypted_reasoning: bool,
provider: str,
model: str,
) -> "_AnthropicMessages | None":
if not sent_encrypted_reasoning or not _rejects_encrypted_reasoning(error):
return None
verbose_logger.warning(
"Backend rejected replayed encrypted reasoning, retrying once without it. "
"Enable the encrypted_content_affinity pre-call check to drop such blocks before the call. "
"provider=%s model=%s status=%s code=%s",
provider,
model,
error.status_code,
_logged_rejection_code(error),
)
return _ANTHROPIC_MESSAGES.validate_python(
strip_encrypted_reasoning_blocks_from_anthropic_messages([dict(message) for message in messages])
)
def _messages_without_rejected_reasoning(
messages: _AnthropicMessages,
error: litellm.BadRequestError,
completion_kwargs: _CompletionKwargs,
sent_encrypted_reasoning: bool,
) -> "_AnthropicMessages | None":
return messages_without_rejected_reasoning(
messages,
error,
sent_encrypted_reasoning=sent_encrypted_reasoning,
provider=_completion_text(completion_kwargs, "custom_llm_provider"),
model=_completion_text(completion_kwargs, "model"),
)
async def _acompletion_without_rejected_reasoning(
completion_kwargs: _CompletionKwargs,
messages: _AnthropicMessages,
prepare_replay: _PrepareReplay,
) -> "ChatResult":
sent_encrypted_reasoning: Final = _sends_encrypted_reasoning(completion_kwargs)
try:
with resent_without_rejected_reasoning(sent_encrypted_reasoning):
return await litellm.acompletion(**completion_kwargs)
except litellm.BadRequestError as error:
replayed_messages: Final = _messages_without_rejected_reasoning(
messages, error, completion_kwargs, sent_encrypted_reasoning
)
if replayed_messages is None:
raise
return await litellm.acompletion(**prepare_replay(replayed_messages))
def _completion_without_rejected_reasoning(
completion_kwargs: _CompletionKwargs,
messages: _AnthropicMessages,
prepare_replay: _PrepareReplay,
) -> "ChatResult | Coroutine[object, object, ChatResult]":
sent_encrypted_reasoning: Final = _sends_encrypted_reasoning(completion_kwargs)
try:
with resent_without_rejected_reasoning(sent_encrypted_reasoning):
return litellm.completion(**completion_kwargs)
except litellm.BadRequestError as error:
replayed_messages: Final = _messages_without_rejected_reasoning(
messages, error, completion_kwargs, sent_encrypted_reasoning
)
if replayed_messages is None:
raise
return litellm.completion(**prepare_replay(replayed_messages))
def _is_compaction_block(block: object) -> bool:
mapping: Final = as_str_mapping(block)
return mapping is not None and mapping.get("type") == "compaction"
def _messages_have_compaction_block(messages: _AnthropicMessages) -> bool:
"""Return True when any message carries a ``compaction`` content block."""
for msg in messages:
content = msg.get("content")
if not isinstance(content, list):
continue
for block in content:
if isinstance(block, dict) and block.get("type") == "compaction":
return True
return False
return any(
any(map(_is_compaction_block, content)) for msg in messages if isinstance(content := msg.get("content"), list)
)
def _proxy_router_fallback() -> "Router | None":
@ -413,34 +559,14 @@ class LiteLLMMessagesToCompletionTransformationHandler:
If the user provides a `summary` field in the thinking dict, it is passed
through to the OpenAI reasoning params (opt-in per OpenAI spec).
"""
custom_llm_provider = completion_kwargs.get("custom_llm_provider")
if custom_llm_provider is None:
try:
_, inferred_provider, _, _ = litellm.utils.get_llm_provider(
model=cast(str, completion_kwargs.get("model"))
)
custom_llm_provider = inferred_provider
except Exception:
custom_llm_provider = None
if custom_llm_provider != "openai":
return
if not isinstance(thinking, dict) or thinking.get("type") != "enabled":
return
model: Final = completion_kwargs.get("model")
try:
model_info: Final = get_model_info(model=cast(str, model), custom_llm_provider=custom_llm_provider)
if model_info and model_info.get("supports_reasoning") is False:
# Model doesn't support reasoning/responses API, don't route
return
except Exception:
pass
if not isinstance(model, str) or thinking is None:
return
if not openai_thinking_routes_to_responses(model, completion_kwargs.get("custom_llm_provider"), thinking):
return
if isinstance(model, str) and model and "responses/" not in model:
local_model: Final = model.removeprefix(f"{custom_llm_provider}/")
completion_kwargs["model"] = f"{custom_llm_provider}/responses/{local_model}"
if "responses/" not in model:
completion_kwargs["model"] = f"openai/responses/{model.removeprefix('openai/')}"
auto_summary: Final = is_reasoning_auto_summary_enabled()
@ -646,12 +772,15 @@ class LiteLLMMessagesToCompletionTransformationHandler:
if "output_config" in extra_kwargs:
request_data["output_config"] = extra_kwargs["output_config"]
api_base: Final = provider_kwargs.get("api_base")
(
openai_request,
tool_name_mapping,
) = ANTHROPIC_ADAPTER.translate_completion_input_params_with_tool_mapping(
request_data,
custom_llm_provider=target_provider,
api_base=api_base if isinstance(api_base, str) else None,
reasoning_summary=peek_reasoning_summary_aliases(provider_kwargs),
)
if openai_request is None:
@ -779,12 +908,9 @@ class LiteLLMMessagesToCompletionTransformationHandler:
parent_trace_id=sub_call_trace_id,
)
(
completion_kwargs,
tool_name_mapping,
) = LiteLLMMessagesToCompletionTransformationHandler._prepare_completion_kwargs(
prepare: Final = partial(
LiteLLMMessagesToCompletionTransformationHandler._prepare_completion_kwargs,
max_tokens=max_tokens,
messages=effective_messages,
model=model,
metadata=metadata,
stop_sequences=stop_sequences,
@ -799,8 +925,10 @@ class LiteLLMMessagesToCompletionTransformationHandler:
output_format=output_format,
extra_kwargs=kwargs,
)
completion_response: Final = await litellm.acompletion(**completion_kwargs)
initial_kwargs, tool_name_mapping = prepare(messages=effective_messages)
completion_response: Final = await _acompletion_without_rejected_reasoning(
initial_kwargs, effective_messages, lambda replayed_messages: prepare(messages=replayed_messages)[0]
)
logging_obj: Final = litellm_logging_obj_from_kwargs(kwargs)
if stream:
@ -924,12 +1052,9 @@ class LiteLLMMessagesToCompletionTransformationHandler:
effective_messages: Final = polyfill_result.messages if polyfill_result is not None else messages
effective_system: Final = polyfill_result.system if polyfill_result is not None else system
(
completion_kwargs,
tool_name_mapping,
) = LiteLLMMessagesToCompletionTransformationHandler._prepare_completion_kwargs(
prepare: Final = partial(
LiteLLMMessagesToCompletionTransformationHandler._prepare_completion_kwargs,
max_tokens=max_tokens,
messages=effective_messages,
model=model,
metadata=metadata,
stop_sequences=stop_sequences,
@ -944,8 +1069,10 @@ class LiteLLMMessagesToCompletionTransformationHandler:
output_format=output_format,
extra_kwargs=kwargs,
)
completion_response: Final = litellm.completion(**completion_kwargs)
initial_kwargs, tool_name_mapping = prepare(messages=effective_messages)
completion_response: Final = _completion_without_rejected_reasoning(
initial_kwargs, effective_messages, lambda replayed_messages: prepare(messages=replayed_messages)[0]
)
logging_obj: Final = litellm_logging_obj_from_kwargs(kwargs)
if stream:

40
litellm/llms/anthropic/pass_through/adapters/streaming_iterator.py
View file

@ -24,6 +24,7 @@ from typing_extensions import assert_never
from litellm._logging import verbose_logger
from litellm._uuid import uuid
from litellm.exceptions import ContextWindowExceededError, MidStreamFallbackError
from litellm.litellm_core_utils.core_helpers import as_str_mapping
from litellm.litellm_core_utils.hidden_params import set_hidden_params
from litellm.llms.anthropic.pass_through.safeguards import StreamedSafeguardResults
from litellm.llms.anthropic.pass_through.utils import anthropic_sse_frame
@ -71,6 +72,21 @@ def _optional_attr_sequence(obj: object, name: str) -> Sequence[object]:
return value if value else ()
def _first_thinking_block(chunk: "ModelResponseStream") -> Mapping[str, object] | None:
delta: Final = _optional_attr(chunk.choices[0], "delta") if chunk.choices else None
blocks: Final = _optional_attr_sequence(delta, "thinking_blocks")
return as_str_mapping(blocks[0]) if blocks else None
def _thinking_block_index(block: Mapping[str, object] | None) -> int | None:
index: Final = block.get("index") if block is not None else None
return index if isinstance(index, int) else None
def _closes_thinking_block(block: Mapping[str, object] | None) -> bool:
return block is not None and _thinking_block_index(block) is not None and bool(block.get("signature"))
def _error_status_and_message(exc: Exception) -> tuple[int, str]:
if isinstance(exc, (BaseLLMException, MidStreamFallbackError, ContextWindowExceededError)):
return exc.status_code, exc.message
@ -458,12 +474,14 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper):
sent_first_chunk: bool = False
sent_content_block_start: bool = False
sent_content_block_finish: bool = False
current_content_block_type: Literal["text", "tool_use", "thinking"] = "text"
current_content_block_type: Literal["text", "tool_use", "thinking", "redacted_thinking"] = "text"
sent_last_message: bool = False
holding_chunk: ContentBlockDelta | None = None
holding_stop_reason_chunk: MessageBlockDelta | None = None
queued_usage_chunk: bool = False
current_content_block_index: int = 0
current_thinking_block_index: int | None = None
current_thinking_block_closed: bool = False
def __init__(
self,
@ -1352,11 +1370,31 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper):
original_name: Final = self.tool_name_mapping.get(truncated_name, truncated_name)
tool_block["name"] = original_name
thinking_block: Final = _first_thinking_block(chunk)
thinking_block_index: Final = _thinking_block_index(thinking_block)
previous_thinking_block_index: Final = self.current_thinking_block_index
follows_closed_thinking_block: Final = self.current_thinking_block_closed and block_type == "thinking"
self.current_thinking_block_closed = _closes_thinking_block(thinking_block)
if (
thinking_block_index is not None
or block_type != self.current_content_block_type
or follows_closed_thinking_block
):
self.current_thinking_block_index = thinking_block_index
if block_type != self.current_content_block_type:
self.current_content_block_type = block_type
self.current_content_block_start = content_block_start
return True
if follows_closed_thinking_block or (
thinking_block_index is not None
and previous_thinking_block_index is not None
and thinking_block_index != previous_thinking_block_index
):
self.current_content_block_start = content_block_start
return True
# For parallel tool calls, we'll necessarily have a new content block
# if we get a function name since it signals a new tool call
if block_type == "tool_use":

173
litellm/llms/anthropic/pass_through/adapters/transformation.py
View file

@ -2,12 +2,14 @@ import copy
import hashlib
import json
from collections.abc import AsyncIterator, Iterator, Mapping, Sequence
from dataclasses import dataclass
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, Literal, TypeAlias, TypeVar, cast
from pydantic import JsonValue, TypeAdapter
import litellm
from litellm.litellm_core_utils.core_helpers import as_str_mapping, attribute_items
from litellm.litellm_core_utils.hidden_params import get_hidden_params
from litellm.llms.anthropic.pass_through.utils import (
is_reasoning_auto_summary_enabled,
@ -122,6 +124,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import (
anthropic_image_source_to_openai_url,
parse_tool_call_arguments,
reasoning_content_from_thinking_blocks,
thinking_blocks_with_encrypted_reasoning,
with_prompt_cache_breakpoint,
)
from litellm.litellm_core_utils.prompt_templates.factory import (
@ -219,6 +222,72 @@ def target_supports_mid_conversation_system(model: str | None, custom_llm_provid
return supports_mid_conversation_system(model=model, custom_llm_provider=custom_llm_provider)
@dataclass(frozen=True, slots=True)
class ChatRouteInputs:
tools: Sequence[Mapping[str, object]] | None = None
reasoning_effort: str | Mapping[str, object] | None = None
api_base: str | None = None
thinking: Mapping[str, object] | None = None
reasoning_summary: object | None = None
_NO_ROUTE_INPUTS: Final = ChatRouteInputs()
def _inferred_provider(model: str) -> str | None:
try:
return litellm.utils.get_llm_provider(model=model)[1]
except Exception: # noqa: BLE001 # an unknown model name has no inferred provider
return None
def openai_thinking_routes_to_responses(
model: str | None,
custom_llm_provider: str | None,
thinking: Mapping[str, object] | None,
) -> bool:
if not model or thinking is None or thinking.get("type") != "enabled":
return False
provider: Final = custom_llm_provider or _inferred_provider(model)
if provider != "openai":
return False
try:
model_info: Final = litellm.get_model_info(model=model, custom_llm_provider=provider)
except Exception: # noqa: BLE001 # a model missing from the cost map still routes, as before
return True
return model_info.get("supports_reasoning") is not False
def _keeps_encrypted_reasoning_for_responses(
model: str | None,
custom_llm_provider: str | None,
route: ChatRouteInputs,
) -> bool:
from litellm.responses.bridge_check import chat_completion_ends_in_responses_api
return openai_thinking_routes_to_responses(
model, custom_llm_provider, route.thinking
) or chat_completion_ends_in_responses_api(
model,
custom_llm_provider,
tools=route.tools,
reasoning_effort=route.reasoning_effort,
reasoning_summary=route.reasoning_summary,
api_base=route.api_base,
)
def _anthropic_thinking_blocks(source: object) -> tuple[Mapping[str, object], ...]:
return thinking_blocks_with_encrypted_reasoning(
getattr(source, "thinking_blocks", None),
getattr(source, "reasoning_items", None),
)
def _streamed_thinking_blocks(delta: object) -> tuple[Mapping[str, object], ...]:
return tuple(block for item in attribute_items(delta, "thinking_blocks") if (block := as_str_mapping(item)))
def _chat_tool_param(function_chunk: ChatCompletionToolParamFunctionChunk, tool: object) -> ChatCompletionToolParam:
eager_input_streaming: Final = eager_input_streaming_flag(tool)
if eager_input_streaming is None:
@ -246,7 +315,12 @@ class AnthropicAdapter:
return result
def translate_completion_input_params_with_tool_mapping(
self, kwargs, *, custom_llm_provider: str | None = None
self,
kwargs,
*,
custom_llm_provider: str | None = None,
api_base: str | None = None,
reasoning_summary: object | None = None,
) -> tuple[ChatCompletionRequest | None, dict[str, str]]:
"""
Translate Anthropic request params to OpenAI format, returning tool name mapping.
@ -280,6 +354,8 @@ class AnthropicAdapter:
) = LiteLLMAnthropicMessagesAdapter().translate_anthropic_to_openai(
anthropic_message_request=request_body,
custom_llm_provider=custom_llm_provider,
api_base=api_base,
reasoning_summary=reasoning_summary,
)
return translated_body, tool_name_mapping
@ -465,9 +541,14 @@ class LiteLLMAnthropicMessagesAdapter:
*,
custom_llm_provider: str | None = None,
preserve_midturn_system: bool = False,
route: ChatRouteInputs = _NO_ROUTE_INPUTS,
) -> list[AllMessageValues]:
new_messages: Final[list[AllMessageValues]] = []
replayable_messages: Final = strip_encrypted_reasoning_blocks_from_anthropic_messages(messages)
replayable_messages: Final = (
messages
if _keeps_encrypted_reasoning_for_responses(model, custom_llm_provider, route)
else strip_encrypted_reasoning_blocks_from_anthropic_messages(messages)
)
leading_count: Final = next(
(i for i, m in enumerate(replayable_messages) if not is_system_role_message(m)),
len(replayable_messages),
@ -1221,6 +1302,8 @@ class LiteLLMAnthropicMessagesAdapter:
anthropic_message_request: AnthropicMessagesRequest,
*,
custom_llm_provider: str | None = None,
api_base: str | None = None,
reasoning_summary: object | None = None,
preserve_midturn_system: bool = False,
) -> tuple[ChatCompletionRequest, dict[str, str]]:
"""
@ -1232,26 +1315,15 @@ class LiteLLMAnthropicMessagesAdapter:
for tools that exceeded OpenAI's 64-char limit
"""
# Debug: Processing Anthropic message request
new_messages: list[AllMessageValues] = []
tool_name_mapping: dict[str, str] = {}
## CONVERT ANTHROPIC MESSAGES TO OPENAI
messages_list: Final[list[AllAnthropicPassThroughMessageValues]] = cast(
list[AllAnthropicPassThroughMessageValues],
anthropic_message_request["messages"],
)
new_messages = self.translate_anthropic_messages_to_openai(
messages=messages_list,
model=anthropic_message_request.get("model"),
custom_llm_provider=custom_llm_provider,
preserve_midturn_system=preserve_midturn_system,
)
## ADD SYSTEM MESSAGE TO MESSAGES
self._add_system_message_to_messages(new_messages, anthropic_message_request)
new_kwargs: Final[ChatCompletionRequest] = {
"model": anthropic_message_request["model"],
"messages": new_messages,
"messages": [],
}
## CONVERT METADATA (user_id + litellm metadata)
self._translate_metadata_to_openai(
@ -1275,6 +1347,23 @@ class LiteLLMAnthropicMessagesAdapter:
new_kwargs=new_kwargs,
custom_llm_provider=custom_llm_provider,
)
## CONVERT ANTHROPIC MESSAGES TO OPENAI, after tools and thinking pick the route
new_messages: Final = self.translate_anthropic_messages_to_openai(
messages=messages_list,
model=anthropic_message_request.get("model"),
custom_llm_provider=custom_llm_provider,
preserve_midturn_system=preserve_midturn_system,
route=ChatRouteInputs(
tools=new_kwargs.get("tools"),
reasoning_effort=new_kwargs.get("reasoning_effort"),
api_base=api_base,
thinking=anthropic_message_request.get("thinking"),
reasoning_summary=reasoning_summary,
),
)
## ADD SYSTEM MESSAGE TO MESSAGES
self._add_system_message_to_messages(new_messages, anthropic_message_request)
new_kwargs["messages"] = new_messages
## CONVERT STOP_SEQUENCES
self._translate_stop_sequences_to_openai(
anthropic_message_request=anthropic_message_request,
@ -1358,9 +1447,9 @@ class LiteLLMAnthropicMessagesAdapter:
for choice in choices
):
new_content.extend(dict(_COMPACTION_BLOCK.validate_python(block)) for block in compaction_blocks)
# Handle thinking blocks first
if hasattr(choice.message, "thinking_blocks") and choice.message.thinking_blocks:
for thinking_block in choice.message.thinking_blocks:
signed_thinking_blocks = _anthropic_thinking_blocks(choice.message)
if signed_thinking_blocks:
for thinking_block in signed_thinking_blocks:
if thinking_block.get("type") == "thinking":
if is_empty_unsigned_thinking_block(thinking_block):
continue
@ -1627,7 +1716,7 @@ class LiteLLMAnthropicMessagesAdapter:
def translate_streaming_openai_chunk_to_anthropic_content_block(
self, choices: list[OpenAIStreamingChoice | StreamingChoices]
) -> tuple[
Literal["text", "tool_use", "thinking"],
Literal["text", "tool_use", "thinking", "redacted_thinking"],
"ContentBlockContentBlockDict",
]:
from litellm._uuid import uuid
@ -1660,20 +1749,24 @@ class LiteLLMAnthropicMessagesAdapter:
choice.delta
) is not None:
return "text", TextBlock(type="text", text="")
elif isinstance(choice, StreamingChoices) and hasattr(choice.delta, "thinking_blocks"):
thinking_blocks = choice.delta.thinking_blocks or []
if len(thinking_blocks) > 0:
thinking_block = thinking_blocks[0]
if thinking_block["type"] == "thinking":
thinking = thinking_block.get("thinking") or ""
signature = thinking_block.get("signature") or ""
elif isinstance(choice, StreamingChoices) and (thinking_blocks := _streamed_thinking_blocks(choice.delta)):
thinking_block = thinking_blocks[0]
if thinking_block["type"] == "redacted_thinking":
redacted_data = thinking_block.get("data")
return "redacted_thinking", ChatCompletionRedactedThinkingBlock(
type="redacted_thinking",
data=redacted_data if isinstance(redacted_data, str) else "",
)
if thinking_block["type"] == "thinking":
thinking = thinking_block.get("thinking") or ""
signature = thinking_block.get("signature") or ""
assert isinstance(thinking, str)
assert isinstance(signature, str)
assert isinstance(thinking, str)
assert isinstance(signature, str)
return "thinking", ChatCompletionThinkingBlock(
type="thinking", thinking=thinking, signature=signature
)
return "thinking", ChatCompletionThinkingBlock(
type="thinking", thinking=thinking, signature=signature
)
# OpenAI-compatible reasoning backends (e.g. vLLM/SGLang reasoning
# parsers) populate ``reasoning_content`` without ``thinking_blocks``.
# ``Delta`` deletes the ``thinking_blocks`` attribute when unset, so the
@ -1706,19 +1799,17 @@ class LiteLLMAnthropicMessagesAdapter:
for tool in choice.delta.tool_calls:
if tool.function is not None and tool.function.arguments is not None:
partial_json = (partial_json or "") + tool.function.arguments
elif isinstance(choice, StreamingChoices) and hasattr(choice.delta, "thinking_blocks"):
thinking_blocks = choice.delta.thinking_blocks or []
if len(thinking_blocks) > 0:
for thinking_block in thinking_blocks:
if thinking_block["type"] == "thinking":
thinking = thinking_block.get("thinking") or ""
signature = thinking_block.get("signature") or ""
elif isinstance(choice, StreamingChoices) and (thinking_blocks := _streamed_thinking_blocks(choice.delta)):
for thinking_block in thinking_blocks:
if thinking_block["type"] == "thinking":
thinking = thinking_block.get("thinking") or ""
signature = thinking_block.get("signature") or ""
assert isinstance(thinking, str)
assert isinstance(signature, str)
assert isinstance(thinking, str)
assert isinstance(signature, str)
reasoning_content += thinking
reasoning_signature += signature
reasoning_content += thinking
reasoning_signature += signature
# Handle reasoning_content when thinking_blocks is not present
# This handles providers like OpenRouter that return reasoning_content
elif isinstance(choice, StreamingChoices) and hasattr(choice.delta, "reasoning_content"):

53
litellm/llms/anthropic/pass_through/messages/handler.py
View file

@ -27,6 +27,7 @@ from litellm.llms.base_llm.anthropic_messages.transformation import (
from litellm.llms.base_llm.chat.transformation import BaseLLMException
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler
from litellm.responses.mcp.tool_precheck import may_name_litellm_mcp_gateway
from litellm.types.llms.anthropic import AllAnthropicToolsValues
from litellm.types.llms.anthropic_messages.anthropic_request import AnthropicMetadata
from litellm.types.llms.anthropic_messages.anthropic_response import (
@ -569,38 +570,34 @@ def anthropic_messages_handler(
# downstream path (native passthrough and both bridges) gets real tools rather than a
# reference the provider cannot resolve. Popped from kwargs so it never reaches the provider.
skip_mcp_handler: Final = kwargs.pop("_skip_mcp_handler", False)
if not skip_mcp_handler and tools:
if not skip_mcp_handler and tools and may_name_litellm_mcp_gateway(tools): # pyright: ignore[reportUnknownArgumentType] # untyped tools list
from litellm.llms.anthropic.pass_through.messages.mcp_handler import (
anthropic_messages_with_mcp,
)
from litellm.responses.mcp.litellm_proxy_mcp_handler import (
LiteLLM_Proxy_MCP_Handler,
)
if LiteLLM_Proxy_MCP_Handler.should_use_litellm_mcp_gateway(tools=tools):
return anthropic_messages_with_mcp(
max_tokens=max_tokens,
messages=messages,
model=model,
metadata=metadata,
stop_sequences=stop_sequences,
stream=stream,
system=cast( # cast-ok: anthropic accepts list system blocks; handler signature is narrower
"str | None", system
),
temperature=temperature,
thinking=thinking,
tool_choice=tool_choice,
tools=tools,
top_k=top_k,
top_p=top_p,
container=container,
api_key=api_key,
api_base=api_base,
client=client,
custom_llm_provider=custom_llm_provider,
**kwargs,
)
return anthropic_messages_with_mcp(
max_tokens=max_tokens,
messages=messages,
model=model,
metadata=metadata,
stop_sequences=stop_sequences,
stream=stream,
system=cast( # cast-ok: anthropic accepts list system blocks; handler signature is narrower
"str | None", system
),
temperature=temperature,
thinking=thinking,
tool_choice=tool_choice,
tools=tools,
top_k=top_k,
top_p=top_p,
container=container,
api_key=api_key,
api_base=api_base,
client=client,
custom_llm_provider=custom_llm_provider,
**kwargs,
)
anthropic_messages_provider_config: Final = _native_messages_config(
model, custom_llm_provider, kwargs.get("model_info")

165
litellm/llms/anthropic/pass_through/responses_adapters/handler.py
View file

@ -14,9 +14,10 @@ from litellm.integrations.anthropic_cache_control_hook import (
AnthropicCacheControlHook,
supports_openai_prompt_cache_breakpoint,
)
from litellm.litellm_core_utils.core_helpers import as_str_mapping, normalize_drop_params
from litellm.litellm_core_utils.core_helpers import as_str_mapping, iter_items, normalize_drop_params
from litellm.litellm_core_utils.prompt_templates.common_utils import cache_control_as_prompt_cache_breakpoint
from litellm.litellm_core_utils.tool_fields import drop_policy, drops_mcp_servers, tools_without_mcp_servers
from litellm.responses.streaming_iterator import BaseResponsesAPIStreamingIterator
from litellm.types.integrations.custom_logger import CONVERTED_STREAM_KEYS
from litellm.types.llms.anthropic import (
AllAnthropicMessageValues,
@ -31,7 +32,11 @@ from litellm.types.llms.anthropic_messages.anthropic_response import (
from litellm.types.llms.openai import ResponsesAPIResponse
from litellm.utils import ProviderConfigManager
from ..adapters.handler import request_safeguards_evaluator
from ..adapters.handler import (
messages_without_rejected_reasoning,
request_safeguards_evaluator,
resent_without_rejected_reasoning,
)
from ..safeguards import with_safeguard_results
from ..utils import litellm_logging_obj_from_kwargs, local_model_name
from .streaming_iterator import AnthropicResponsesStreamWrapper
@ -284,6 +289,78 @@ def _build_responses_kwargs(
return responses_kwargs
_ResponsesKwargsFor: TypeAlias = Callable[[AnthropicRequestMessages], Mapping[str, object]]
_ResponsesResult: TypeAlias = ResponsesAPIResponse | BaseResponsesAPIStreamingIterator
_SyncResponsesResult: TypeAlias = _ResponsesResult | Coroutine[object, object, _ResponsesResult]
def _is_encrypted_reasoning_item(item: object) -> bool:
mapping: Final = as_str_mapping(item)
return mapping is not None and mapping.get("type") == "reasoning" and bool(mapping.get("encrypted_content"))
def _sends_encrypted_reasoning(responses_kwargs: Mapping[str, object]) -> bool:
input_items: Final = responses_kwargs.get("input")
sent: Final = iter_items(input_items) if input_items is not None else ()
return any(map(_is_encrypted_reasoning_item, sent))
def _kwarg_text(responses_kwargs: Mapping[str, object], key: str) -> str:
value: Final = responses_kwargs.get(key)
return value if isinstance(value, str) else ""
def _messages_without_rejected_reasoning(
messages: AnthropicRequestMessages,
error: litellm.BadRequestError,
responses_kwargs: Mapping[str, object],
sent_encrypted_reasoning: bool,
) -> AnthropicRequestMessages | None:
return messages_without_rejected_reasoning(
messages,
error,
sent_encrypted_reasoning=sent_encrypted_reasoning,
provider=_kwarg_text(responses_kwargs, "custom_llm_provider"),
model=_kwarg_text(responses_kwargs, "model"),
)
async def _aresponses_without_rejected_reasoning(
messages: AnthropicRequestMessages, responses_kwargs_for: _ResponsesKwargsFor
) -> tuple[Mapping[str, object], _ResponsesResult]:
responses_kwargs: Final = responses_kwargs_for(messages)
sent_encrypted_reasoning: Final = _sends_encrypted_reasoning(responses_kwargs)
try:
with resent_without_rejected_reasoning(sent_encrypted_reasoning):
return responses_kwargs, await litellm.aresponses(**responses_kwargs)
except litellm.BadRequestError as error:
replayed_messages: Final = _messages_without_rejected_reasoning(
messages, error, responses_kwargs, sent_encrypted_reasoning
)
if replayed_messages is None:
raise
replay_kwargs: Final = responses_kwargs_for(replayed_messages)
return replay_kwargs, await litellm.aresponses(**replay_kwargs)
def _responses_without_rejected_reasoning(
messages: AnthropicRequestMessages, responses_kwargs_for: _ResponsesKwargsFor
) -> tuple[Mapping[str, object], _SyncResponsesResult]:
responses_kwargs: Final = responses_kwargs_for(messages)
sent_encrypted_reasoning: Final = _sends_encrypted_reasoning(responses_kwargs)
try:
with resent_without_rejected_reasoning(sent_encrypted_reasoning):
return responses_kwargs, litellm.responses(**responses_kwargs)
except litellm.BadRequestError as error:
replayed_messages: Final = _messages_without_rejected_reasoning(
messages, error, responses_kwargs, sent_encrypted_reasoning
)
if replayed_messages is None:
raise
replay_kwargs: Final = responses_kwargs_for(replayed_messages)
return replay_kwargs, litellm.responses(**replay_kwargs)
class LiteLLMMessagesToResponsesAPIHandler:
"""
Handles Anthropic /v1/messages requests for OpenAI / Azure models by
@ -311,28 +388,30 @@ class LiteLLMMessagesToResponsesAPIHandler:
output_format: AnthropicOutputSchema | None = None,
**kwargs: object,
) -> AnthropicMessagesResponse | AsyncIterator[bytes]:
responses_kwargs: Final = _build_responses_kwargs(
max_tokens=max_tokens,
messages=messages,
model=model,
context_management=context_management,
metadata=metadata,
output_config=output_config,
stop_sequences=stop_sequences,
stream=stream,
system=system,
temperature=temperature,
thinking=thinking,
tool_choice=tool_choice,
tools=tools,
top_k=top_k,
top_p=top_p,
output_format=output_format,
extra_kwargs=kwargs,
)
def responses_kwargs_for(sent_messages: AnthropicRequestMessages) -> Mapping[str, object]:
return _build_responses_kwargs(
max_tokens=max_tokens,
messages=sent_messages,
model=model,
context_management=context_management,
metadata=metadata,
output_config=output_config,
stop_sequences=stop_sequences,
stream=stream,
system=system,
temperature=temperature,
thinking=thinking,
tool_choice=tool_choice,
tools=tools,
top_k=top_k,
top_p=top_p,
output_format=output_format,
extra_kwargs=kwargs,
)
safeguards_evaluator: Final = request_safeguards_evaluator(kwargs, messages)
result: Final = await litellm.aresponses(**responses_kwargs)
responses_kwargs, result = await _aresponses_without_rejected_reasoning(messages, responses_kwargs_for)
if stream:
wrapper: Final = AnthropicResponsesStreamWrapper(
@ -396,28 +475,28 @@ class LiteLLMMessagesToResponsesAPIHandler:
**kwargs,
)
# Sync path
responses_kwargs: Final = _build_responses_kwargs(
max_tokens=max_tokens,
messages=messages,
model=model,
context_management=context_management,
metadata=metadata,
output_config=output_config,
stop_sequences=stop_sequences,
stream=stream,
system=system,
temperature=temperature,
thinking=thinking,
tool_choice=tool_choice,
tools=tools,
top_k=top_k,
top_p=top_p,
output_format=output_format,
extra_kwargs=kwargs,
)
def responses_kwargs_for(sent_messages: AnthropicRequestMessages) -> Mapping[str, object]:
return _build_responses_kwargs(
max_tokens=max_tokens,
messages=sent_messages,
model=model,
context_management=context_management,
metadata=metadata,
output_config=output_config,
stop_sequences=stop_sequences,
stream=stream,
system=system,
temperature=temperature,
thinking=thinking,
tool_choice=tool_choice,
tools=tools,
top_k=top_k,
top_p=top_p,
output_format=output_format,
extra_kwargs=kwargs,
)
result: Final = litellm.responses(**responses_kwargs)
responses_kwargs, result = _responses_without_rejected_reasoning(messages, responses_kwargs_for)
if stream:
wrapper: Final = AnthropicResponsesStreamWrapper(

32
litellm/llms/bedrock/base_aws_llm.py
View file

@ -1526,28 +1526,16 @@ class BaseAWSLLM(SignsRequestsWithAWS):
aws_region_name: str,
endpoint_type: Literal["runtime", "agent", "agentcore"] | None = "runtime",
) -> tuple[str, str]:
env_aws_bedrock_runtime_endpoint: Final = get_secret("AWS_BEDROCK_RUNTIME_ENDPOINT")
if api_base is not None:
endpoint_url = api_base
elif aws_bedrock_runtime_endpoint is not None and isinstance(aws_bedrock_runtime_endpoint, str):
endpoint_url = aws_bedrock_runtime_endpoint
elif env_aws_bedrock_runtime_endpoint and isinstance(env_aws_bedrock_runtime_endpoint, str):
endpoint_url = env_aws_bedrock_runtime_endpoint
else:
endpoint_url = self._select_default_endpoint_url(
endpoint_type=endpoint_type,
aws_region_name=aws_region_name,
)
# Determine proxy_endpoint_url
if aws_bedrock_runtime_endpoint is not None and isinstance(aws_bedrock_runtime_endpoint, str):
proxy_endpoint_url = aws_bedrock_runtime_endpoint
elif env_aws_bedrock_runtime_endpoint and isinstance(env_aws_bedrock_runtime_endpoint, str):
proxy_endpoint_url = env_aws_bedrock_runtime_endpoint
else:
proxy_endpoint_url = endpoint_url
return endpoint_url, proxy_endpoint_url
configured_endpoint: Final = (
aws_bedrock_runtime_endpoint if isinstance(aws_bedrock_runtime_endpoint, str) else None
)
runtime_endpoint: Final = configured_endpoint or get_secret_str("AWS_BEDROCK_RUNTIME_ENDPOINT")
endpoint_url: Final = (
api_base
or runtime_endpoint
or self._select_default_endpoint_url(endpoint_type=endpoint_type, aws_region_name=aws_region_name)
)
return endpoint_url, runtime_endpoint or endpoint_url
def _select_default_endpoint_url(
self,

18
litellm/llms/bedrock/common_utils.py
View file

@ -1336,15 +1336,29 @@ def bedrock_converse_accepts_redacted_thinking(model: str, base_model: str | Non
its own. An ARN that hides its model, like an application inference profile, falls back to the
deployment's ``base_model`` and counts as non-Claude without one.
"""
return _is_claude(_converse_target_model(model, base_model))
def bedrock_converse_rejects_signed_thinking(model: str, base_model: str | None) -> bool:
if base_model is None and "." not in get_bedrock_base_model(model):
return False
return not _is_claude_behind_any_region_prefix(_converse_target_model(model, base_model))
def _converse_target_model(model: str, base_model: str | None) -> str:
if base_model is None or "." in get_bedrock_base_model(model):
return _is_claude(model)
return _is_claude(base_model)
return model
return base_model
def _is_claude(model: str) -> bool:
return get_bedrock_base_model(model).startswith(("anthropic", "claude"))
def _is_claude_behind_any_region_prefix(model: str) -> bool:
return _is_claude(model) or get_bedrock_base_model(model).partition(".")[2].startswith("anthropic.")
def bedrock_converse_supports_parallel_tool_use_config(model: str) -> bool:
return any(
(litellm.model_cost.get(candidate) or {}).get("supports_parallel_tool_use_config") is True

46
litellm/llms/bedrock/passthrough/transformation.py
View file

@ -13,6 +13,10 @@ from litellm._logging import verbose_logger
from litellm.litellm_core_utils.litellm_logging import Logging
from litellm.llms.base_llm.passthrough.transformation import BasePassthroughConfig, PassthroughStreamCollector
from litellm.llms.bedrock.chat.converse_to_openai import as_json_array, as_json_object
from litellm.llms.bedrock.chat.invoke_transformations.base_invoke_transformation import (
bedrock_invoke_guardrail_headers,
)
from litellm.llms.bedrock.request_metadata import merge_bedrock_invoke_headers
from litellm.types.llms.openai import ChatCompletionSystemMessage, ChatCompletionUserMessage
from litellm.types.utils import ModelResponseStream
@ -38,8 +42,35 @@ def is_bedrock_streaming_endpoint(endpoint: str) -> bool:
_CONVERSE_ACTIONS: Final = frozenset({"converse", "converse-stream"})
def _is_converse_url(url: str) -> bool:
return httpx.URL(url).path.rstrip("/").rsplit("/", 1)[-1] in _CONVERSE_ACTIONS
_INVOKE_ACTIONS: Final = frozenset({"invoke", "invoke-with-response-stream"})
_BEDROCK_INVOKE_GUARDRAIL_HEADERS: Final = frozenset(
{"x-amzn-bedrock-guardrailidentifier", "x-amzn-bedrock-guardrailversion", "x-amzn-bedrock-trace"}
)
def _converse_guardrail_config(action: str, guardrail_config: object) -> object:
fields: Final = as_json_object(guardrail_config)
if action == "converse-stream" or fields is None:
return guardrail_config
return {key: value for key, value in fields.items() if key != "streamProcessingMode"}
def _with_deployment_guardrail(
action: str,
headers: Mapping[str, str],
body: Mapping[str, object],
guardrail_config: object,
) -> tuple[Mapping[str, str], Mapping[str, object]]:
if guardrail_config is None:
return headers, body
if action in _CONVERSE_ACTIONS:
return headers, {**body, "guardrailConfig": _converse_guardrail_config(action, guardrail_config)}
if action in _INVOKE_ACTIONS:
guardrail_headers: Final = tuple(bedrock_invoke_guardrail_headers(guardrail_config).items())
return merge_bedrock_invoke_headers(
dict(headers), (), guardrail_headers, _BEDROCK_INVOKE_GUARDRAIL_HEADERS
), body
return headers, body
def _is_cache_point(entry: object) -> bool:
@ -363,13 +394,20 @@ class BedrockPassthroughConfig(BaseAWSLLM, BedrockModelInfo, BedrockEventStreamD
api_base: str,
model: str | None = None,
) -> tuple[dict, bytes | None]:
action: Final = httpx.URL(api_base).path.rstrip("/").rsplit("/", 1)[-1]
body: Final[Mapping[str, object]] = request_data or {}
optional_params: Final = litellm_params.copy()
guarded_headers, guarded_body = _with_deployment_guardrail(
action,
headers,
_converse_body_to_send(body) if action in _CONVERSE_ACTIONS else body,
litellm_params.get("guardrailConfig"),
)
return self._sign_request(
service_name="bedrock",
headers=headers,
headers=dict(guarded_headers),
optional_params=optional_params,
request_data=dict(_converse_body_to_send(body) if _is_converse_url(api_base) else body),
request_data=dict(guarded_body),
api_base=api_base,
model=model,
api_key=self.get_bedrock_bearer_token(optional_params),

11
litellm/llms/bedrock/request_metadata.py
View file

@ -179,12 +179,11 @@ def merge_bedrock_invoke_headers(
"""
Merge the ``X-Amzn-*`` headers the Invoke paths derive from params.
``caller_owned`` (the guardrail headers) defers to a header the caller already set, which is
the long-standing behaviour for those. ``proxy_owned_names`` are dropped from the caller's
headers unconditionally and re-supplied only from ``proxy_owned``, because those names carry
proxy-authenticated identity into an AWS billing record that the caller must not be able to
write. Names are compared case-insensitively so a caller cannot leave a second spelling in
the dict and let the transport pick the winner.
``caller_owned`` defers to a header the caller already set. ``proxy_owned_names`` are dropped
from the caller's headers unconditionally and re-supplied only from ``proxy_owned``, for names
the caller must not be able to write, such as proxy-authenticated identity headed for an AWS
billing record or a deployment's guardrail. Names are compared case-insensitively so a caller
cannot leave a second spelling in the dict and let the transport pick the winner.
"""
if not caller_owned and not proxy_owned and not proxy_owned_names:
return headers

9
litellm/llms/dashscope/videos/transformation.py
View file

@ -35,6 +35,7 @@ from litellm.types.videos.utils import (
decode_video_id_with_provider,
encode_video_id_with_provider,
extract_original_video_id,
id_safe_name,
)
from litellm.utils import get_model_info
@ -599,17 +600,13 @@ def _stamped_poll_models(metadata: object) -> tuple[str, str | None] | None:
fields: Final = _PARAMETERS_ADAPTER.validate_python(metadata)
except ValidationError:
return None
requested: Final = _id_safe_name(fields.get("model_group"))
served: Final = _id_safe_name(fields.get("deployment_model_name"))
requested: Final = id_safe_name(fields.get("model_group"))
served: Final = id_safe_name(fields.get("deployment_model_name"))
if requested is None:
return None
return requested, None if served is None or served == requested or "*" in served else served
def _id_safe_name(name: object) -> str | None:
return name if isinstance(name, str) and name and ";" not in name else None
class DashScopeVideoConfig(BaseVideoConfig):
def get_supported_openai_params(self, model: str) -> list[str]: # mutable-ok: BaseVideoConfig contract returns list
return [

4
litellm/llms/openai/responses/transformation.py
View file

@ -63,7 +63,9 @@ _NO_TOOL_UPDATE: Final[Mapping[str, object]] = MappingProxyType({})
_MODEL_FAMILIES_REJECTING_TOP_LEVEL_SCHEMA_COMBINATORS: Final = ("gpt-4", "gpt-3.5", "chatgpt-4o", "o1", "o3", "o4")
_PROVIDERS_WITH_OPENAI_SCHEMA_VALIDATOR: Final = frozenset({LlmProviders.AZURE, LlmProviders.OPENAI})
_PROVIDERS_VALIDATING_TOOL_CALL_ITEM_IDS: Final = frozenset({LlmProviders.AZURE, LlmProviders.OPENAI})
_PROVIDERS_REPLAYING_ONLY_THEIR_OWN_REASONING: Final = _PROVIDERS_VALIDATING_TOOL_CALL_ITEM_IDS
_PROVIDERS_REPLAYING_ONLY_THEIR_OWN_REASONING: Final = frozenset(
{LlmProviders.AZURE, LlmProviders.BEDROCK, LlmProviders.BEDROCK_MANTLE, LlmProviders.OPENAI}
)
_ROLES_READING_PROMPT_CACHE_BREAKPOINTS: Final = frozenset({"system", "developer", "user"})
_PART_TYPES_READING_PROMPT_CACHE_BREAKPOINTS: Final = frozenset({"input_text", "input_image", "input_file"})
_INPUT_ITEM: Final[TypeAdapter[Mapping[str, object]]] = TypeAdapter(Mapping[str, object])

34
litellm/llms/vertex_ai/batches/handler.py
View file

@ -14,8 +14,10 @@ from litellm.litellm_core_utils.url_utils import (
)
from litellm.llms.custom_httpx.http_handler import (
HTTPHandler,
MaskedHTTPStatusError,
get_async_httpx_client,
get_httpx_client,
mask_sensitive_info,
)
from litellm.llms.vertex_ai.common_utils import VertexAIError, get_vertex_base_url
from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import VertexLLM
@ -80,6 +82,15 @@ def _vertex_batch_list_payload(response: _VertexBatchListJsonSource) -> dict[str
return response.json()
def _vertex_batch_create_error(error: httpx.HTTPStatusError) -> MaskedHTTPStatusError:
body: Final = mask_sensitive_info(error.response.text)
return MaskedHTTPStatusError(
httpx.HTTPStatusError(f"{error}\n{body}", request=error.request, response=error.response),
message=body,
text=body,
)
class VertexAIBatchPrediction(VertexLLM):
def __init__(self, gcs_bucket_name: str, *args, **kwargs):
super().__init__(*args, **kwargs)
@ -165,11 +176,19 @@ class VertexAIBatchPrediction(VertexLLM):
headers=headers,
)
response: Final = sync_handler.post(
url=api_base,
headers=headers,
data=json.dumps(vertex_batch_request),
)
try:
response: Final = sync_handler.post(
url=api_base,
headers=headers,
data=json.dumps(vertex_batch_request),
)
except httpx.HTTPStatusError as e:
litellm.verbose_logger.error(
"Vertex AI batch create failed: status=%s, body=%s",
e.response.status_code,
e.response.text[:1000],
)
raise _vertex_batch_create_error(e) from e
payload_view: Final[_VertexBatchPayloadView] = {"payload": response.json()}
_json_response: Final = payload_view["payload"]
@ -265,13 +284,12 @@ class VertexAIBatchPrediction(VertexLLM):
data=json.dumps(vertex_batch_request),
)
except httpx.HTTPStatusError as e:
error_body: Final = e.response.text
litellm.verbose_logger.error(
"Vertex AI batch create failed: status=%s, body=%s",
e.response.status_code,
error_body[:1000],
e.response.text[:1000],
)
raise
raise _vertex_batch_create_error(e) from e
payload_view: Final[_VertexBatchPayloadView] = {"payload": response.json()}
_json_response: Final = payload_view["payload"]

67
litellm/llms/vertex_ai/batches/transformation.py
View file

@ -9,14 +9,41 @@ from litellm._logging import verbose_logger
from litellm._uuid import uuid
from litellm.llms.vertex_ai.common_utils import (
VertexAIError,
VertexAIModelRoute,
convert_vertex_datetime_to_openai_datetime,
get_vertex_ai_model_route,
)
from litellm.llms.vertex_ai.vertex_ai_partner_models.anthropic.transformation import VertexAIAnthropicConfig
from litellm.types.llms.openai import BatchJobStatus, CreateBatchRequest
from litellm.types.llms.vertex_ai import *
from litellm.types.llms.vertex_ai import GenerateContentResponseBody
from litellm.types.llms.vertex_ai import GenerateContentResponseBody, VertexPartnerProvider
from litellm.types.utils import LiteLLMBatch, ModelInfo, Usage
_NATIVE_VERTEX_RESPONSE: Final = TypeAdapter(GenerateContentResponseBody)
_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object])
VERTEX_ANTHROPIC_VERSION: Final = "vertex-2023-10-16"
_ANTHROPIC_PUBLISHER_MODEL_PATH: Final = "publishers/anthropic/models/"
def is_vertex_anthropic_batch_model(model: str) -> bool:
"""
Whether a Vertex batch model (bare, `vertex_ai/`-prefixed or a publisher path) is a Claude model, decided the
way Vertex chat routes Claude. Vertex batches Claude only as `publishers/anthropic/models/<model>` with rows in
Anthropic's Messages shape
"""
bare_model: Final = model.rsplit("/", 1)[-1]
return _ANTHROPIC_PUBLISHER_MODEL_PATH in model or (
get_vertex_ai_model_route(bare_model) is VertexAIModelRoute.PARTNER_MODELS
and VertexPartnerProvider.claude.value in bare_model
)
def as_json_object(value: object) -> Mapping[str, object] | None:
return _JSON_OBJECT.validate_python(value) if isinstance(value, dict) else None
def is_vertex_anthropic_batch_output_response(response_body: Mapping[str, object]) -> bool:
return response_body.get("type") == "message" and isinstance(response_body.get("usage"), dict)
def _int_field(mapping: Mapping[str, object], key: str) -> int:
@ -43,8 +70,22 @@ def is_vertex_embedding_batch_output_response(response_body: Mapping[str, object
return isinstance(response_body.get("embedding"), dict)
def is_vertex_anthropic_batch_output_row(vertex_output_row: Mapping[str, object]) -> bool:
"""
Whether a Vertex batch output row came from a Claude (Anthropic Messages) request. Vertex echoes only the
request's top-level `custom_id` on these rows, with an Anthropic `message` response on success and an empty
response plus a `status` on failure, so a Claude input file read back is never mistaken for output
"""
if not ("custom_id" in vertex_output_row and "processed_time" in vertex_output_row):
return False
response: Final = as_json_object(vertex_output_row.get("response"))
return (response is not None and is_vertex_anthropic_batch_output_response(response)) or bool(
vertex_output_row.get("status")
)
def is_native_vertex_batch_output_row(row: Mapping[str, object]) -> bool:
return isinstance(row.get("request"), dict)
return isinstance(row.get("request"), dict) or is_vertex_anthropic_batch_output_row(row)
class NativeVertexBatchCostCalculator(Protocol):
@ -72,6 +113,9 @@ def _native_vertex_row_usage(
response_body: Mapping[str, object],
calculate_usage: Callable[[GenerateContentResponseBody], Usage],
) -> Usage | None:
anthropic_usage: Final = as_json_object(response_body.get("usage"))
if is_vertex_anthropic_batch_output_response(response_body) and anthropic_usage is not None:
return VertexAIAnthropicConfig().calculate_usage(usage_object=anthropic_usage, reasoning_content=None)
if "usageMetadata" not in response_body:
if not is_vertex_embedding_batch_output_response(response_body):
return None
@ -97,21 +141,24 @@ def native_vertex_batch_row_stats(
"""
Usage and cost of one native Vertex predictions.jsonl row, a
`{"request": ..., "response": {"candidates": [...], "usageMetadata": {...}, "modelVersion": ...}}`
generateContent object or a `{"request": ..., "response": {"embedding": {...}, "usageMetadata": {...}}}`
embedding object (an embedding row without `usageMetadata` is billed from its documented `tokenCount`).
`model_name` (the deployment model) prices the row unless it is a wildcard, else its own `modelVersion`
does, else the wildcard name so explicit deployment prices still apply; a row without a response, a
generateContent row without `response.usageMetadata`, and a row whose response fails validation are
generateContent object, a `{"request": ..., "response": {"embedding": {...}, "usageMetadata": {...}}}`
embedding object (an embedding row without `usageMetadata` is billed from its documented `tokenCount`) or a
`{"custom_id": ..., "request": ..., "response": {"type": "message", "model": ..., "usage": {...}}}` Claude object.
`model_name` (the deployment model) prices the row unless it is a wildcard, else its own `modelVersion` (Claude
rows: `model`) does, else the wildcard name so explicit deployment prices still apply; a row without a response,
a generateContent row without `response.usageMetadata`, and a row whose response fails validation are
None (failed).
"""
response_body: Final = row.get("response")
if not isinstance(response_body, dict):
response_body: Final = as_json_object(row.get("response"))
if response_body is None:
return None
usage: Final = _native_vertex_row_usage(response_body, calculate_usage)
if usage is None:
return None
total_tokens: Final = usage.total_tokens or (usage.prompt_tokens + usage.completion_tokens)
model_version: Final = response_body.get("modelVersion")
model_version: Final = response_body.get(
"model" if is_vertex_anthropic_batch_output_response(response_body) else "modelVersion"
)
deployment_model: Final = model_name if model_name and "*" not in model_name else None
model: Final = deployment_model or (model_version if isinstance(model_version, str) else model_name)
if model is None:

258
litellm/llms/vertex_ai/files/transformation.py
View file

@ -19,6 +19,7 @@ from typing_extensions import ReadOnly, Required
import litellm
from litellm._uuid import uuid
from litellm.constants import DEFAULT_CHAT_COMPLETION_PARAM_VALUES, RESPONSE_FORMAT_TOOL_NAME
from litellm.files.types import FileContentStreamingResult
from litellm.files.utils import FilesAPIUtils
from litellm.litellm_core_utils.cloud_storage_security import (
@ -30,11 +31,13 @@ from litellm.litellm_core_utils.cloud_storage_security import (
split_configured_cloud_bucket_name,
validate_managed_cloud_file_id,
)
from litellm.litellm_core_utils.core_helpers import as_str_mapping, iter_items, normalize_drop_params
from litellm.litellm_core_utils.litellm_logging import Logging
from litellm.litellm_core_utils.prompt_templates.common_utils import (
extract_file_data,
extract_file_metadata,
)
from litellm.llms.anthropic.chat.transformation import ANTHROPIC_TOOL_NAME_REVERSE_MAP_KEY
from litellm.llms.base_llm.base_utils import map_developer_role_to_system_role
from litellm.llms.base_llm.chat.transformation import BaseLLMException
from litellm.llms.base_llm.files.batch_records import responses_batch_body_to_chat_body
@ -43,7 +46,14 @@ from litellm.llms.base_llm.files.transformation import (
BaseFileUploadStream,
LiteLLMLoggingObj,
)
from litellm.llms.vertex_ai.batches.transformation import vertex_embedding_prompt_token_count
from litellm.llms.vertex_ai.batches.transformation import (
VERTEX_ANTHROPIC_VERSION,
as_json_object,
is_vertex_anthropic_batch_model,
is_vertex_anthropic_batch_output_response,
is_vertex_anthropic_batch_output_row,
vertex_embedding_prompt_token_count,
)
from litellm.llms.vertex_ai.common_utils import (
convert_vertex_datetime_to_openai_datetime,
get_vertex_ai_fine_tuned_endpoint_id,
@ -55,6 +65,7 @@ from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import (
from litellm.llms.vertex_ai.gemini_embeddings.batch_embed_content_transformation import (
transform_openai_input_gemini_embed_content,
)
from litellm.llms.vertex_ai.vertex_ai_partner_models.anthropic.transformation import VertexAIAnthropicConfig
from litellm.types.files import StreamingMediaUploadConfig
from litellm.types.llms.openai import (
AllMessageValues,
@ -75,6 +86,7 @@ from litellm.types.utils import (
ModelResponse,
Usage,
)
from litellm.utils import get_non_default_completion_params, get_optional_params
from ..common_utils import VertexAIError
from ..vertex_llm_base import VertexBase
@ -83,6 +95,7 @@ _GCP_LABEL_VALUE_MAX_LEN: Final = 63
_CUSTOM_ID_RAW_LABEL_PREFIX: Final = "b32_"
_VERTEX_BATCH_KEY_FIELD: Final = "key"
_MANAGED_GCS_MODEL_PATH_PATTERN: Final = re.compile(r"publishers/[^/]+/models/([^/?]+)")
_PUBLISHER_MODEL_PATH_PATTERN: Final = re.compile(r"publishers/[^/]+/models/")
_EMBED_REQUEST_FIELD_BY_GEMINI_PARAM: Final = (
("outputDimensionality", "output_dimensionality"),
("taskType", "task_type"),
@ -220,10 +233,14 @@ def _get_litellm_batch_custom_id(vertex_output_row: Mapping[str, object]) -> str
"""
Resolve the OpenAI `custom_id` for a Vertex batch output row.
Embedding rows carry it in the top-level `key` field that Vertex echoes back;
`generateContent` rows have no such field, so it is smuggled through request
labels instead (see `_set_litellm_batch_custom_id_labels`).
Claude rows carry it in their own top-level `custom_id` field and embedding rows in
the top-level `key` field that Vertex echoes back; `generateContent` rows have no
such field, so it is smuggled through request labels instead (see
`_set_litellm_batch_custom_id_labels`).
"""
anthropic_custom_id: Final = vertex_output_row.get("custom_id")
if anthropic_custom_id is not None:
return str(anthropic_custom_id)
key = vertex_output_row.get(_VERTEX_BATCH_KEY_FIELD)
if key is not None:
return unquote(str(key))
@ -284,6 +301,12 @@ def _is_vertex_generate_content_batch_output_row(vertex_output_row: Mapping[str,
)
def _is_vertex_chat_batch_output_row(vertex_output_row: Mapping[str, object]) -> bool:
return _is_vertex_generate_content_batch_output_row(vertex_output_row) or is_vertex_anthropic_batch_output_row(
vertex_output_row
)
def _try_parse_vertex_batch_output_row(line: bytes) -> _VertexBatchRow | None:
try:
row: Final = _parse_vertex_batch_output_row(line.decode("utf-8"))
@ -405,6 +428,37 @@ def _openai_batch_output_row(
}
def _vertex_anthropic_batch_output_row_to_openai(
custom_id: str,
response: object,
mock_httpx_response: httpx.Response,
) -> _OpenAIBatchOutputRow:
"""
One OpenAI batch output row from a Vertex Claude row's Anthropic `message` response
"""
anthropic_response: Final = as_json_object(response)
if anthropic_response is None or not is_vertex_anthropic_batch_output_response(anthropic_response):
return _openai_batch_output_row(
custom_id=custom_id,
error_code="vertex_ai_error",
error_message="Claude batch row has no message response",
)
try:
model_response: Final = VertexAIAnthropicConfig().transform_parsed_response(
completion_response=dict(anthropic_response),
raw_response=mock_httpx_response,
model_response=ModelResponse(),
json_mode=True,
)
except Exception as e: # noqa: BLE001 # one malformed row must fail alone, not the whole output file
return _openai_batch_output_row(
custom_id=custom_id,
error_code="transformation_error",
error_message=f"Failed to transform response: {e}",
)
return _openai_batch_output_row(custom_id=custom_id, body=model_response.model_dump())
def _split_vertex_batch_key(vertex_output_row: Mapping[str, object]) -> tuple[str, int, int]:
"""
Resolve `(custom_id, index within that custom_id, group size)` for a Vertex batch
@ -679,21 +733,123 @@ def _openai_batch_jsonl_entry_to_vertex_embeddings_rows(
)
_CLAUDE_BATCH_ROW_FIXED_PARAMS: Final = frozenset(
("model", "messages", "custom_llm_provider", "drop_params", "stream", "stream_options")
)
_CLAUDE_BATCH_OPENAI_PARAMS: Final = frozenset(DEFAULT_CHAT_COMPLETION_PARAM_VALUES) - {"metadata"}
_CLAUDE_BATCH_DEPLOYMENT_PARAMS: Final = (
_CLAUDE_BATCH_OPENAI_PARAMS - _CLAUDE_BATCH_ROW_FIXED_PARAMS - {"extra_headers"}
) | {"drop_params"}
def _claude_batch_deployment_params(litellm_params: Mapping[str, object]) -> Mapping[str, object]:
return MappingProxyType(
{key: value for key, value in litellm_params.items() if key in _CLAUDE_BATCH_DEPLOYMENT_PARAMS}
)
def _names_the_response_format_tool(tool: object) -> bool:
match tool:
case {"function": {"name": str(name)}}:
return name == RESPONSE_FORMAT_TOOL_NAME
case _:
return False
def _openai_batch_jsonl_entry_to_vertex_anthropic_row(
custom_id: object,
messages: list[AllMessageValues],
request_params: Mapping[str, object],
model: str,
deployment_params: Mapping[str, object],
) -> Mapping[str, object]:
"""
Transforms an OpenAI chat batch entry into the Anthropic Messages row Vertex batches Claude with,
`{"custom_id": ..., "request": {"messages": [...], "anthropic_version": "vertex-2023-10-16", "max_tokens": ...}}`
"""
if custom_id is None:
raise ValueError("custom_id is required on Claude lines, since Claude output rows are matched back only by it")
config: Final = VertexAIAnthropicConfig()
bare_model: Final = model.rsplit("/", 1)[-1]
params: Final = {**deployment_params, **request_params}
drop_params: Final = litellm.drop_params is True or normalize_drop_params(params.get("drop_params")) is True
chat_params: Final = {
**get_non_default_completion_params(params),
**{key: value for key, value in params.items() if key in _CLAUDE_BATCH_OPENAI_PARAMS},
}
if any(map(_names_the_response_format_tool, iter_items(chat_params.get("tools") or ()))):
raise ValueError(
f"tool name {RESPONSE_FORMAT_TOOL_NAME!r} is reserved for response_format on Claude, "
"and batch output cannot tell a user tool by that name apart, so rename it"
)
optional_params: Final[Mapping[str, object]] = get_optional_params(
model=bare_model,
custom_llm_provider="vertex_ai",
drop_params=drop_params,
messages=messages,
**{ # pyright: ignore[reportArgumentType] # batch JSONL bodies are untyped; get_optional_params checks each param as chat does
key: value for key, value in chat_params.items() if key not in _CLAUDE_BATCH_ROW_FIXED_PARAMS
},
)
request_state: Final[dict[str, object]] = {"drop_params": drop_params}
anthropic_request: Final[Mapping[str, object]] = config.transform_request(
model=bare_model,
messages=messages,
optional_params={
**{key: value for key, value in optional_params.items() if key not in ("json_mode", "extra_body")},
"anthropic_version": VERTEX_ANTHROPIC_VERSION,
},
litellm_params=request_state,
headers={},
)
renamed_tools: Final = as_str_mapping(request_state.get(ANTHROPIC_TOOL_NAME_REVERSE_MAP_KEY))
if renamed_tools:
raise ValueError(
f"tool names {sorted(map(str, renamed_tools.values()))} must match ^[a-zA-Z0-9_-]{{1,128}}$ for Claude on Vertex, "
"and batch output cannot map a renamed tool back, so rename these tools"
)
extra_body: Final = as_json_object(optional_params.get("extra_body"))
request: Final[Mapping[str, object]] = (
anthropic_request
if extra_body is None
else {
**anthropic_request,
**config.transform_extra_body(
extra_body=extra_body, request=anthropic_request, model=bare_model, litellm_params={}
),
}
)
return {"custom_id": str(custom_id), "request": request}
def _openai_batch_jsonl_entry_to_vertex_rows(
openai_entry: dict[str, Any],
map_openai_to_vertex_params: Callable[[dict[str, object]], dict[str, object]],
batch_model: str | None = None,
deployment_params: Mapping[str, object] = MappingProxyType({}),
) -> tuple[Mapping[str, object], ...]:
"""
Transforms a single OpenAI JSONL batch entry into the Vertex rows it maps to.
`batch_model` is the deployment model the batch runs against and wins over the entry's own `body.model`,
which picks the row shape: Anthropic Messages for Claude, `GenerateContentRequest` otherwise.
jsonl body for vertex is {"request": <request_body>}
Example Vertex jsonl
{"request":{"contents": [{"role": "user", "parts": [{"text": "What is the relation between the following video and image samples?"}, {"fileData": {"fileUri": "gs://cloud-samples-data/generative-ai/video/animals.mp4", "mimeType": "video/mp4"}}, {"fileData": {"fileUri": "gs://cloud-samples-data/generative-ai/image/cricket.jpeg", "mimeType": "image/jpeg"}}]}]}}
"""
entry_body: Final = openai_entry.get("body")
model: Final = batch_model or (str(entry_body.get("model") or "") if isinstance(entry_body, dict) else "")
is_anthropic_model: Final = is_vertex_anthropic_batch_model(model)
if _is_embeddings_batch_entry(openai_entry):
if is_anthropic_model:
raise VertexAIError(
status_code=400,
message=f"Vertex AI batch embeddings are not supported for Claude models ('{model}')",
)
return _openai_batch_jsonl_entry_to_vertex_embeddings_rows(openai_entry)
if not isinstance(openai_entry.get("body"), dict):
if not isinstance(entry_body, dict):
raise TypeError("`body` on batch requests must be a JSON object, but was missing or not an object")
openai_request_body: Final = openai_entry["body"]
chat_request_body: Final = (
@ -701,8 +857,21 @@ def _openai_batch_jsonl_entry_to_vertex_rows(
if _is_responses_batch_entry(openai_entry)
else openai_request_body
)
chat_messages: Final = map_developer_role_to_system_role(chat_request_body.get("messages", []))
if is_anthropic_model:
return (
_openai_batch_jsonl_entry_to_vertex_anthropic_row(
custom_id=openai_entry.get("custom_id"),
messages=chat_messages,
request_params={
key: value for key, value in chat_request_body.items() if key not in ("model", "messages")
},
model=model,
deployment_params=deployment_params,
),
)
vertex_request_body: Final = transform_request_body(
messages=map_developer_role_to_system_role(chat_request_body.get("messages", [])),
messages=chat_messages,
model=chat_request_body.get("model", ""),
optional_params=map_openai_to_vertex_params(chat_request_body),
custom_llm_provider="vertex_ai",
@ -854,15 +1023,24 @@ class _OpenAIToVertexBatchUploadStream(BaseFileUploadStream):
self,
openai_file_content: FileTypes,
map_openai_to_vertex_params: Callable[[dict[str, object]], dict[str, object]],
batch_model: str | None = None,
deployment_params: Mapping[str, object] = MappingProxyType({}),
) -> None:
self._openai_file_content = openai_file_content
self._map_openai_to_vertex_params = map_openai_to_vertex_params
self._batch_model = batch_model
self._deployment_params = deployment_params
def _iter_vertex_jsonl_chunks(self) -> Iterator[bytes]:
first = True
for lineno, entry in _iter_openai_jsonl_entries(self._openai_file_content):
try:
vertex_rows = _openai_batch_jsonl_entry_to_vertex_rows(entry, self._map_openai_to_vertex_params)
vertex_rows = _openai_batch_jsonl_entry_to_vertex_rows(
entry,
self._map_openai_to_vertex_params,
batch_model=self._batch_model,
deployment_params=self._deployment_params,
)
except Exception as e:
raise VertexAIError(
status_code=400,
@ -918,16 +1096,31 @@ class _RawFileUploadStream(BaseFileUploadStream):
return _iter_raw_file_chunks(self._file_content)
def _deployment_batch_model(litellm_params: Mapping[str, object]) -> str | None:
model: Final = litellm_params.get("model")
return model.removeprefix("vertex_ai/") if isinstance(model, str) and model else None
def _publisher_model_path(raw_model: str) -> str:
if _PUBLISHER_MODEL_PATH_PATTERN.search(raw_model):
return raw_model
publisher: Final = "anthropic" if is_vertex_anthropic_batch_model(raw_model) else "google"
return f"publishers/{publisher}/models/{raw_model}"
def _managed_batch_object_name(raw_model: str, *, passthrough: bool) -> str:
endpoint_id: Final = get_vertex_ai_fine_tuned_endpoint_id(raw_model)
model_path: Final = (
f"endpoints/{endpoint_id}"
if endpoint_id is not None
else (raw_model if "publishers/google/models" in raw_model else f"publishers/google/models/{raw_model}")
)
safe_model_path: Final = sanitize_cloud_object_path(model_path, fallback="model")
model_path: Final = f"endpoints/{endpoint_id}" if endpoint_id is not None else _publisher_model_path(raw_model)
prefix: Final = _PASSTHROUGH_MANAGED_GCS_PREFIX if passthrough else VERTEX_AI_MANAGED_GCS_PREFIX
return f"{prefix}{safe_model_path}/{uuid.uuid4()}"
return f"{prefix}{_safe_managed_model_path(model_path)}/{uuid.uuid4()}"
def _safe_managed_model_path(model_path: str) -> str:
if not is_vertex_anthropic_batch_model(model_path):
return sanitize_cloud_object_path(model_path, fallback="model")
return "@".join(
sanitize_cloud_object_path(version_part, fallback="model") for version_part in model_path.split("@")
)
class VertexAIFilesConfig(VertexBase, BaseFilesConfig):
@ -1030,10 +1223,10 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig):
fallback_filename="file",
)
def _get_configured_bucket_name(self, litellm_params: dict) -> str:
def _get_configured_bucket_name(self, litellm_params: Mapping[str, object]) -> str:
configured_bucket_name: Final = litellm_params.get("gcs_bucket_name") or litellm_params.get("bucket_name")
bucket_name: Final = (
litellm_params.get("gcs_bucket_name")
or litellm_params.get("bucket_name")
(configured_bucket_name if isinstance(configured_bucket_name, str) else None)
or os.getenv("GCS_BATCH_BUCKET_NAME")
or os.getenv("GCS_BUCKET_NAME")
)
@ -1047,7 +1240,7 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig):
api_key: str | None,
model: str,
optional_params: dict,
litellm_params: dict,
litellm_params: Mapping[str, object],
data: CreateFileRequest,
) -> str:
"""
@ -1071,11 +1264,10 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig):
raise ValueError("file is required")
if purpose is None:
raise ValueError("purpose is required")
configured_model: Final = litellm_params.get("model")
object_name = self.get_object_name(
file_data,
purpose,
deployment_model=configured_model if isinstance(configured_model, str) else None,
deployment_model=_deployment_batch_model(litellm_params),
passthrough=is_passthrough_batch_upload(data, litellm_params),
)
if object_prefix:
@ -1126,7 +1318,7 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig):
model: str,
create_file_data: CreateFileRequest,
optional_params: dict,
litellm_params: dict,
litellm_params: Mapping[str, object],
) -> bytes | str | dict:
"""
2 Cases:
@ -1157,6 +1349,8 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig):
body_stream=_OpenAIToVertexBatchUploadStream(
file_data,
self._map_openai_to_vertex_params,
batch_model=_deployment_batch_model(litellm_params),
deployment_params=_claude_batch_deployment_params(litellm_params),
),
content_type="application/json",
)
@ -1207,7 +1401,7 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig):
def get_error_class(self, error_message: str, status_code: int, headers: dict | Headers) -> BaseLLMException:
return VertexAIError(status_code=status_code, message=error_message, headers=headers)
def _parse_gcs_uri(self, file_id: str, litellm_params: dict | None = None) -> tuple[str, str]:
def _parse_gcs_uri(self, file_id: str, litellm_params: Mapping[str, object] | None = None) -> tuple[str, str]:
"""
Validate a managed GCS file_id and return (bucket, url-encoded-object-path).
"""
@ -1225,7 +1419,7 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig):
self,
file_id: str,
optional_params: dict,
litellm_params: dict,
litellm_params: Mapping[str, object],
) -> tuple[str, dict]:
bucket, encoded_object = self._parse_gcs_uri(file_id, litellm_params)
url: Final = f"https://storage.googleapis.com/storage/v1/b/{bucket}/o/{encoded_object}"
@ -1255,7 +1449,7 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig):
self,
file_id: str,
optional_params: dict,
litellm_params: dict,
litellm_params: Mapping[str, object],
) -> tuple[str, dict]:
bucket, encoded_object = self._parse_gcs_uri(file_id, litellm_params)
url: Final = f"https://storage.googleapis.com/storage/v1/b/{bucket}/o/{encoded_object}"
@ -1298,7 +1492,7 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig):
self,
file_content_request,
optional_params: dict,
litellm_params: dict,
litellm_params: Mapping[str, object],
) -> tuple[str, dict]:
file_id: Final = file_content_request.get("file_id", "")
bucket, encoded_object = self._parse_gcs_uri(file_id, litellm_params)
@ -1395,7 +1589,7 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig):
headers=MappingProxyType({**headers, "content-length": str(len(transformed_content))}),
)
if not _is_vertex_generate_content_batch_output_row(first_row):
if not _is_vertex_chat_batch_output_row(first_row):
return FileContentStreamingResult(stream_iterator=replayed_stream, headers=headers)
return FileContentStreamingResult(
@ -1479,10 +1673,7 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig):
# first line is not valid UTF-8/JSON) raises and falls through to the
# passthrough below, leaving the content untouched.
first_row: Final = _parse_vertex_batch_output_row(first_line)
if not (
_is_vertex_embeddings_batch_output_row(first_row)
or _is_vertex_generate_content_batch_output_row(first_row)
):
if not (_is_vertex_embeddings_batch_output_row(first_row) or _is_vertex_chat_batch_output_row(first_row)):
return content
context: Final = _new_vertex_batch_output_row_transform_context()
@ -1547,6 +1738,13 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig):
error_message=status,
)
if is_vertex_anthropic_batch_output_row(vertex_output):
return _vertex_anthropic_batch_output_row_to_openai(
custom_id=custom_id,
response=vertex_output.get("response"),
mock_httpx_response=mock_httpx_response,
)
# Transform successful response using existing transformation
vertex_response: Final = vertex_output.get("response", {})

136
litellm/main.py
View file

@ -27,6 +27,7 @@ from functools import partial
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Protocol, Union, cast, get_args
from litellm._internal_context import is_resent_failure
from litellm._logging import redact_string
from litellm._uuid import uuid
@ -135,6 +136,7 @@ from litellm.llms.vertex_ai.common_utils import (
get_vertex_ai_model_route,
)
from litellm.realtime_api.main import realtime_health_check
from litellm.responses.mcp.tool_precheck import may_name_litellm_mcp_gateway
from litellm.secret_managers.main import get_secret_bool, get_secret_str
from litellm.types.completion import (
CompletionDispatchContext,
@ -5366,63 +5368,55 @@ def completion(
control_options: Final = _resolve_control_options(kwargs, model)
skip_mcp_handler: Final = kwargs.pop("_skip_mcp_handler", False)
if not skip_mcp_handler and tools:
if not skip_mcp_handler and tools and may_name_litellm_mcp_gateway(tools): # pyright: ignore[reportUnknownArgumentType] # untyped tools list
from litellm.responses.mcp.chat_completions_handler import acompletion_with_mcp
from litellm.responses.mcp.litellm_proxy_mcp_handler import (
LiteLLM_Proxy_MCP_Handler,
)
from litellm.types.llms.openai import ToolParam
# Check if MCP tools are present (following responses pattern)
# Cast tools to Optional[Iterable[ToolParam]] for type checking
tools_for_mcp: Final = cast(Iterable[ToolParam] | None, tools)
if LiteLLM_Proxy_MCP_Handler.should_use_litellm_mcp_gateway(tools=tools_for_mcp):
return acompletion_with_mcp( # pyright: ignore[reportReturnType] # MCP path returns a coroutine that acompletion() awaits; completion()'s sync return type omits it
model=model,
messages=messages,
functions=functions,
function_call=function_call,
timeout=timeout,
temperature=temperature,
top_p=top_p,
n=n,
stream=stream,
stream_options=stream_options,
stop=stop,
max_tokens=max_tokens,
max_completion_tokens=max_completion_tokens,
modalities=modalities,
prediction=prediction,
audio=audio,
presence_penalty=presence_penalty,
frequency_penalty=frequency_penalty,
logit_bias=logit_bias,
user=user,
response_format=response_format,
seed=seed,
tools=tools,
tool_choice=tool_choice,
parallel_tool_calls=parallel_tool_calls,
logprobs=logprobs,
top_logprobs=top_logprobs,
deployment_id=deployment_id,
reasoning_effort=reasoning_effort,
verbosity=verbosity,
safety_identifier=safety_identifier,
service_tier=service_tier,
store=store,
prompt_cache_key=prompt_cache_key,
base_url=base_url,
api_version=api_version,
api_key=api_key,
model_list=model_list,
extra_headers=extra_headers,
thinking=thinking,
web_search_options=web_search_options,
shared_session=shared_session,
enable_json_schema_validation=enable_json_schema_validation,
**kwargs,
)
return acompletion_with_mcp( # pyright: ignore[reportReturnType] # MCP path returns a coroutine that acompletion() awaits; completion()'s sync return type omits it
model=model,
messages=messages,
functions=functions,
function_call=function_call,
timeout=timeout,
temperature=temperature,
top_p=top_p,
n=n,
stream=stream,
stream_options=stream_options,
stop=stop,
max_tokens=max_tokens,
max_completion_tokens=max_completion_tokens,
modalities=modalities,
prediction=prediction,
audio=audio,
presence_penalty=presence_penalty,
frequency_penalty=frequency_penalty,
logit_bias=logit_bias,
user=user,
response_format=response_format,
seed=seed,
tools=tools,
tool_choice=tool_choice,
parallel_tool_calls=parallel_tool_calls,
logprobs=logprobs,
top_logprobs=top_logprobs,
deployment_id=deployment_id,
reasoning_effort=reasoning_effort,
verbosity=verbosity,
safety_identifier=safety_identifier,
service_tier=service_tier,
store=store,
prompt_cache_key=prompt_cache_key,
base_url=base_url,
api_version=api_version,
api_key=api_key,
model_list=model_list,
extra_headers=extra_headers,
thinking=thinking,
web_search_options=web_search_options,
shared_session=shared_session,
enable_json_schema_validation=enable_json_schema_validation,
**kwargs,
)
api_base = kwargs.get("api_base", None)
mock_response: Final[MOCK_RESPONSE_TYPE | None] = kwargs.get("mock_response", None)
mock_tool_calls: Final = kwargs.get("mock_tool_calls", None)
@ -6242,6 +6236,10 @@ def completion(
)
def _is_not_resent_failure(error: BaseException) -> bool:
return isinstance(error, Exception) and not is_resent_failure(error)
def completion_with_retries(*args, **kwargs):
"""
Executes a litellm.completion() with 3 retries
@ -6261,10 +6259,15 @@ def completion_with_retries(*args, **kwargs):
retryer = tenacity.Retrying(
wait=tenacity.wait_exponential(multiplier=1, max=10),
stop=tenacity.stop_after_attempt(num_retries),
retry=tenacity.retry_if_exception(_is_not_resent_failure),
reraise=True,
)
else:
retryer = tenacity.Retrying(stop=tenacity.stop_after_attempt(num_retries), reraise=True)
retryer = tenacity.Retrying(
stop=tenacity.stop_after_attempt(num_retries),
retry=tenacity.retry_if_exception(_is_not_resent_failure),
reraise=True,
)
return retryer(original_function, *args, **kwargs)
@ -6287,10 +6290,15 @@ async def acompletion_with_retries(*args, **kwargs):
retryer = tenacity.AsyncRetrying(
wait=tenacity.wait_exponential(multiplier=1, max=10),
stop=tenacity.stop_after_attempt(num_retries),
retry=tenacity.retry_if_exception(_is_not_resent_failure),
reraise=True,
)
else:
retryer = tenacity.AsyncRetrying(stop=tenacity.stop_after_attempt(num_retries), reraise=True)
retryer = tenacity.AsyncRetrying(
stop=tenacity.stop_after_attempt(num_retries),
retry=tenacity.retry_if_exception(_is_not_resent_failure),
reraise=True,
)
return await retryer(original_function, *args, **kwargs)
@ -6315,10 +6323,15 @@ def responses_with_retries(*args, **kwargs):
retryer = tenacity.Retrying(
wait=tenacity.wait_exponential(multiplier=1, max=10),
stop=tenacity.stop_after_attempt(num_retries),
retry=tenacity.retry_if_exception(_is_not_resent_failure),
reraise=True,
)
else:
retryer = tenacity.Retrying(stop=tenacity.stop_after_attempt(num_retries), reraise=True)
retryer = tenacity.Retrying(
stop=tenacity.stop_after_attempt(num_retries),
retry=tenacity.retry_if_exception(_is_not_resent_failure),
reraise=True,
)
return retryer(original_function, *args, **kwargs)
@ -6342,10 +6355,15 @@ async def aresponses_with_retries(*args, **kwargs):
retryer = tenacity.AsyncRetrying(
wait=tenacity.wait_exponential(multiplier=1, max=10),
stop=tenacity.stop_after_attempt(num_retries),
retry=tenacity.retry_if_exception(_is_not_resent_failure),
reraise=True,
)
else:
retryer = tenacity.AsyncRetrying(stop=tenacity.stop_after_attempt(num_retries), reraise=True)
retryer = tenacity.AsyncRetrying(
stop=tenacity.stop_after_attempt(num_retries),
retry=tenacity.retry_if_exception(_is_not_resent_failure),
reraise=True,
)
return await retryer(original_function, *args, **kwargs)

3
litellm/passthrough/main.py
View file

@ -437,6 +437,7 @@ def llm_passthrough_route(
params: QueryParamTypes | None = None,
cookies: CookieTypes | None = None,
client: HTTPHandler | AsyncHTTPHandler | None = None,
guardrailConfig: object = None,
**kwargs,
) -> (
httpx.Response
@ -547,7 +548,7 @@ def llm_passthrough_route(
_request_data: dict | None = data if isinstance(data, dict) else (json if isinstance(json, dict) else None)
headers, signed_json_body = provider_config.sign_request(
headers=headers,
litellm_params=litellm_params_dict,
litellm_params={**litellm_params_dict, "guardrailConfig": guardrailConfig},
request_data=_request_data,
api_base=str(updated_url),
model=model,

50
litellm/proxy/auth/auth_utils.py
View file

@ -4,6 +4,7 @@ import re
import sys
from collections.abc import Collection, Iterator, Mapping
from functools import lru_cache
from itertools import accumulate, chain, repeat
from logging import Logger
from types import MappingProxyType
from typing import Any, Final, Protocol
@ -39,6 +40,7 @@ from litellm.types.passthrough_endpoints.pass_through_endpoints import (
from litellm.types.router import CONFIGURABLE_CLIENTSIDE_AUTH_PARAMS, Deployment, server_owned_wif_fields_present
from litellm.types.router import reject_server_owned_wif_params as _reject_server_owned_wif_params
from litellm.types.utils import CustomPricingLiteLLMParams, oauth_token_exchange_litellm_params
from litellm.types.videos.main import DecodedVideoId
def is_invalid_virtual_key_error(exception: BaseException | None) -> bool:
@ -566,6 +568,31 @@ def fallback_target_model_name(target: object) -> str | None:
return None
def _direct_fallback_model_names(model: str, llm_router: Router) -> frozenset[str]:
from litellm.proxy.auth.model_checks import get_all_fallbacks
return frozenset(
name for target in get_all_fallbacks(model, llm_router) if (name := fallback_target_model_name(target))
)
def _fallback_hop(
frontier_and_reached: tuple[frozenset[str], frozenset[str]], llm_router: Router
) -> tuple[frozenset[str], frozenset[str]]:
frontier, reached = frontier_and_reached
next_hop: Final = (
frozenset(chain.from_iterable(_direct_fallback_model_names(model, llm_router) for model in frontier)) - reached
)
return next_hop, reached | next_hop
def router_fallback_model_names(model: str, llm_router: Router) -> frozenset[str]:
"""The groups a failure on `model` can reach through the router's fallbacks, chained and generic (`*`) ones."""
hops: Final = accumulate(repeat(llm_router), _fallback_hop, initial=(frozenset({model}), frozenset({model})))
_, reached = next(hop for hop in hops if not hop[0])
return reached - {model}
def request_fallback_model_names(request_body: Mapping[str, object]) -> tuple[str, ...]:
return tuple(
dict.fromkeys(
@ -1975,10 +2002,10 @@ def _extract_models_from_managed_resource_id(
)
if resource_id_field == "video_id":
model_id = decode_video_id_with_provider(resource_id).get("model_id")
decoded_video: Final = decode_video_id_with_provider(resource_id)
_append_model_candidates(
candidates=candidates,
value=_resolve_model_id_with_router(model_id, llm_router),
value=_video_access_group(decoded_video, llm_router) or decoded_video.get("model_id"),
)
else:
model_id = decode_character_id_with_provider(resource_id).get("model_id")
@ -1992,6 +2019,25 @@ def _extract_models_from_managed_resource_id(
return _dedupe_model_candidates(candidates)
def _video_access_group(decoded_video: DecodedVideoId, llm_router: Router | None) -> str | None:
"""
The group a key must be allowed to call to use a video: the creating deployment's group, or, while router
fallbacks skip model access checks, the group whose router fallback created it.
"""
from litellm.proxy.auth.fallback_model_access import router_fallback_access_check
from litellm.proxy.video_endpoints.utils import video_model_group
creating_group: Final = video_model_group(decoded_video, llm_router)
requested_model: Final = decoded_video.get("requested_model")
if llm_router is None or creating_group is None or requested_model is None:
return creating_group
if router_fallback_access_check.is_enforced():
return creating_group
if creating_group in router_fallback_model_names(requested_model, llm_router):
return requested_model
return creating_group
def _resolve_model_id_with_router(model_id: str | None, llm_router: Router | None) -> str | None:
if model_id is None or llm_router is None:
return model_id

8
litellm/proxy/utils.py
View file

@ -328,6 +328,7 @@ _VIEW_SETUP_POLL_INTERVAL_SECONDS: Final = 5.0
_VIEW_SETUP_DEADLINE_SECONDS: Final = 15 * 60.0
_VIEW_SETUP_GATE_TABLE: Final = '"LiteLLM_SpendLogs"'
_VIEW_SETUP_GATE_PROBE_ROWS: Final = TypeAdapter(tuple[Mapping[str, bool], ...])
_REQUEST_FIELDS: Final = TypeAdapter(Mapping[str, object])
_ViewSetupOutcome: TypeAlias = Literal["ready", "timed_out"]
_ViewSetupAttempt: TypeAlias = Literal["ready", "table_missing"] | Exception
@ -4143,6 +4144,7 @@ class ProxyLogging:
response = pipeline_response # rebind-ok: adopt the pipeline's replacement response, same contract as the callback loops below
pipeline_managed: Final = pipeline_managed_guardrail_names(data, "post_call")
request_view: Final = _REQUEST_FIELDS.validate_python(data)
guardrail_callbacks, other_callbacks = _partition_post_call_callbacks()
try:
# Merge model-level guardrails before checking which guardrails to run
@ -4173,6 +4175,9 @@ class ProxyLogging:
):
continue
if callback.post_call_already_checked(request_view, response) is True:
continue
guardrail_response: Any | None = None
if "apply_guardrail" in type(callback).__dict__ and not callback.use_native_lifecycle_hooks:
@ -4246,10 +4251,13 @@ class ProxyLogging:
immediately before awaiting, and the unified hook pops it before its first
suspension point, so concurrent guardrails never race on that key.
"""
request_view: Final = _REQUEST_FIELDS.validate_python(data)
async def _run_one(callback: CustomGuardrail) -> None:
if callback.should_run_guardrail(data=guardrail_data, event_type=GuardrailEventHooks.post_call) is not True:
return
if callback.post_call_already_checked(request_view, response) is True:
return
if "apply_guardrail" in type(callback).__dict__ and not callback.use_native_lifecycle_hooks:
data["guardrail_to_apply"] = callback
await self._run_guardrail_with_metrics(

92
litellm/proxy/video_endpoints/endpoints.py
View file

@ -20,9 +20,14 @@ from litellm.proxy.common_utils.openai_endpoint_utils import (
)
from litellm.proxy.image_endpoints.endpoints import batch_to_bytesio
from litellm.proxy.video_endpoints.utils import (
VideoRequestRoute,
encode_character_id_in_response,
encode_video_id_in_response,
extract_model_from_target_model_names,
fallback_requested_model,
get_custom_provider_from_data,
routable_model_name,
video_model_group,
video_reference_to_id,
)
from litellm.types.videos.utils import (
@ -84,6 +89,7 @@ async def video_generation(
# Read request body
data: Final = await read_request_body(request=request)
requested: Final = VideoRequestRoute.model_validate(data)
if input_reference is not None:
input_reference_file: Final = await batch_to_bytesio([input_reference])
if input_reference_file:
@ -118,7 +124,13 @@ async def video_generation(
version=version,
)
else:
return generated
requested_model: Final = routable_model_name(requested.model, llm_router, requested.custom_llm_provider)
return encode_video_id_in_response(
response=generated,
fallback_provider=requested.custom_llm_provider,
fallback_model_id=requested_model,
requested_model=fallback_requested_model(requested_model, generated, llm_router),
)
@router.get(
@ -269,12 +281,9 @@ async def video_status(
if custom_llm_provider:
data["custom_llm_provider"] = custom_llm_provider
# Resolve model_name from model_id if available
# This allows the router to automatically inject litellm_params from the model config
if model_id_from_decoded and llm_router:
resolved_model: Final = llm_router.resolve_model_name_from_model_id(model_id_from_decoded, provider_from_id)
if resolved_model:
data["model"] = resolved_model
resolved_model: Final = video_model_group(decoded, llm_router)
if resolved_model:
data["model"] = resolved_model
# Process request using ProxyBaseLLMRequestProcessing
processor: Final = ProxyBaseLLMRequestProcessing(data=data, served_model=decoded.get("served_model"))
@ -305,7 +314,12 @@ async def video_status(
version=version,
)
else:
return status
return encode_video_id_in_response(
response=status,
fallback_provider=custom_llm_provider,
fallback_model_id=model_id_from_decoded,
polled_video=decoded,
)
@router.get(
@ -358,7 +372,6 @@ async def video_content(
decoded: Final = decode_video_id_with_provider(video_id)
provider_from_id: Final = decoded.get("custom_llm_provider")
model_id_from_decoded: Final = decoded.get("model_id")
custom_llm_provider: Final = (
get_custom_llm_provider_from_request_headers(request=request)
@ -369,12 +382,9 @@ async def video_content(
if custom_llm_provider:
data["custom_llm_provider"] = custom_llm_provider
# Resolve model_name from model_id if available
# This allows the router to automatically inject litellm_params from the model config
if model_id_from_decoded and llm_router:
resolved_model: Final = llm_router.resolve_model_name_from_model_id(model_id_from_decoded, provider_from_id)
if resolved_model:
data["model"] = resolved_model
resolved_model: Final = video_model_group(decoded, llm_router)
if resolved_model:
data["model"] = resolved_model
# Process request using ProxyBaseLLMRequestProcessing
processor: Final = ProxyBaseLLMRequestProcessing(data=data, served_model=decoded.get("served_model"))
try:
@ -462,6 +472,7 @@ async def video_remix(
)
data: Final = await read_request_body(request=request)
requested: Final = VideoRequestRoute.model_validate(data)
data["video_id"] = video_id
decoded: Final = decode_video_id_with_provider(video_id)
@ -471,18 +482,15 @@ async def video_remix(
custom_llm_provider: Final = (
get_custom_llm_provider_from_request_headers(request=request)
or get_custom_llm_provider_from_request_query(request=request)
or data.get("custom_llm_provider")
or requested.custom_llm_provider
or provider_from_id
)
if custom_llm_provider:
data["custom_llm_provider"] = custom_llm_provider
# Resolve model_name from model_id if available
# This allows the router to automatically inject litellm_params from the model config
if model_id_from_decoded and llm_router:
resolved_model: Final = llm_router.resolve_model_name_from_model_id(model_id_from_decoded, provider_from_id)
if resolved_model:
data["model"] = resolved_model
resolved_model: Final = video_model_group(decoded, llm_router)
if resolved_model:
data["model"] = resolved_model
# Process request using ProxyBaseLLMRequestProcessing
processor: Final = ProxyBaseLLMRequestProcessing(data=data, served_model=decoded.get("served_model"))
@ -513,7 +521,13 @@ async def video_remix(
version=version,
)
else:
return remixed
return encode_video_id_in_response(
response=remixed,
fallback_provider=custom_llm_provider,
fallback_model_id=model_id_from_decoded
or routable_model_name(requested.model, llm_router, custom_llm_provider),
requested_model=decoded.get("requested_model"),
)
@router.post(
@ -771,6 +785,7 @@ async def video_edit(
)
data: Final = await read_request_body(request=request)
requested: Final = VideoRequestRoute.model_validate(data)
uploaded_video: Final = data.pop("video", None)
if isinstance(uploaded_video, StarletteUploadFile):
video_files: Final = await batch_to_bytesio((uploaded_video,))
@ -793,10 +808,9 @@ async def video_edit(
)
data["custom_llm_provider"] = custom_llm_provider
if model_id_from_decoded and llm_router:
resolved_model: Final = llm_router.resolve_model_name_from_model_id(model_id_from_decoded, provider_from_id)
if resolved_model:
data["model"] = resolved_model
resolved_model: Final = video_model_group(decoded, llm_router)
if resolved_model:
data["model"] = resolved_model
processor: Final = ProxyBaseLLMRequestProcessing(data=data, served_model=decoded.get("served_model"))
try:
@ -826,7 +840,13 @@ async def video_edit(
version=version,
)
else:
return edited
return encode_video_id_in_response(
response=edited,
fallback_provider=custom_llm_provider,
fallback_model_id=model_id_from_decoded
or routable_model_name(requested.model, llm_router, custom_llm_provider),
requested_model=decoded.get("requested_model"),
)
@router.post(
@ -875,6 +895,7 @@ async def video_extension(
)
data: Final = await read_request_body(request=request)
requested: Final = VideoRequestRoute.model_validate(data)
data["video_id"] = video_reference_to_id(data.pop("video", None))
decoded: Final = decode_video_id_with_provider(data["video_id"])
@ -890,10 +911,9 @@ async def video_extension(
)
data["custom_llm_provider"] = custom_llm_provider
if model_id_from_decoded and llm_router:
resolved_model: Final = llm_router.resolve_model_name_from_model_id(model_id_from_decoded, provider_from_id)
if resolved_model:
data["model"] = resolved_model
resolved_model: Final = video_model_group(decoded, llm_router)
if resolved_model:
data["model"] = resolved_model
processor: Final = ProxyBaseLLMRequestProcessing(data=data, served_model=decoded.get("served_model"))
try:
@ -923,4 +943,10 @@ async def video_extension(
version=version,
)
else:
return extended
return encode_video_id_in_response(
response=extended,
fallback_provider=custom_llm_provider,
fallback_model_id=model_id_from_decoded
or routable_model_name(requested.model, llm_router, custom_llm_provider),
requested_model=decoded.get("requested_model"),
)

115
litellm/proxy/video_endpoints/utils.py
View file

@ -1,8 +1,18 @@
from typing import Any, Final
from typing import TYPE_CHECKING, Any, Final
import orjson
from pydantic import BaseModel, ConfigDict
from litellm.types.videos.utils import encode_character_id_with_provider
from litellm.types.videos.main import DecodedVideoId, VideoObject
from litellm.types.videos.utils import (
decode_video_id_with_provider,
encode_character_id_with_provider,
encode_video_id_with_provider,
id_safe_name,
)
if TYPE_CHECKING:
from litellm.router import Router
def extract_model_from_target_model_names(target_model_names: Any) -> str | None:
@ -65,3 +75,104 @@ def encode_character_id_in_response(response: Any, custom_llm_provider: str, mod
model_id=model_id,
)
return response
class VideoRequestRoute(BaseModel):
model_config = ConfigDict(frozen=True)
model: object = None
custom_llm_provider: object = None
def routable_model_name(model: object, llm_router: "Router | None", custom_llm_provider: object) -> str | None:
name: Final = id_safe_name(model)
if name is None or llm_router is None:
return name
provider: Final = custom_llm_provider if isinstance(custom_llm_provider, str) else None
return name if llm_router.resolve_model_name_from_model_id(name, provider) else None
def video_model_group(decoded: DecodedVideoId, llm_router: "Router | None") -> str | None:
"""
The group a video id names. A provider model that several groups share resolves to the group of the
deployment that created the video, under its team's public name for a team model, while a name the router
routes by itself (the id's model, bare or provider-prefixed), or one that routes to that deployment's group
(an alias of it), is kept as written.
"""
if llm_router is None:
return None
model_id: Final = decoded.get("model_id")
provider: Final = decoded.get("custom_llm_provider")
resolved: Final = llm_router.resolve_model_name_from_model_id(model_id, provider) if model_id else None
deployment_id: Final = decoded.get("deployment_id")
if not deployment_id or resolved in (model_id, f"{provider}/{model_id}"):
return resolved
deployment: Final = llm_router.get_deployment(model_id=deployment_id)
if deployment is None:
return resolved
creating_group: Final = deployment.model_info.team_public_model_name or deployment.model_name
if resolved is not None and llm_router.resolve_served_model_name(resolved, creating_group) == creating_group:
return resolved
return creating_group
def fallback_requested_model(requested_model: str | None, response: object, llm_router: "Router | None") -> str | None:
"""The group a create asked for when one of that group's router fallbacks served it."""
from litellm.proxy.auth.auth_utils import router_fallback_model_names
if requested_model is None or llm_router is None or not isinstance(response, VideoObject):
return None
deployment_id: Final = id_safe_name(response.hidden_params.get("model_id"))
deployment: Final = llm_router.get_deployment(model_id=deployment_id) if deployment_id else None
if deployment is None:
return None
creating_group: Final = deployment.model_info.team_public_model_name or deployment.model_name
return requested_model if creating_group in router_fallback_model_names(requested_model, llm_router) else None
def encode_video_id_in_response(
response: object,
fallback_provider: object,
fallback_model_id: object,
polled_video: DecodedVideoId | None = None,
requested_model: str | None = None,
) -> object:
"""
Stamps the provider, model and the deployment that answered into a returned video id. A poll of an existing
video (`polled_video`) answered without a routed deployment keeps the deployment its id already named, and
every returned id keeps the group a router fallback was serving (`requested_model`).
"""
if not isinstance(response, VideoObject) or not response.id:
return response
decoded: Final = decode_video_id_with_provider(response.id)
hidden_params: Final = response.hidden_params
routed_deployment_id: Final = id_safe_name(hidden_params.get("model_id"))
deployment_id: Final = routed_deployment_id or (polled_video or decoded).get("deployment_id")
kept_requested_model: Final = (
id_safe_name(requested_model) or decoded.get("requested_model") or (polled_video or {}).get("requested_model")
)
model_id: Final = decoded.get("model_id") or id_safe_name(fallback_model_id) or routed_deployment_id
provider: Final = (
id_safe_name(decoded.get("custom_llm_provider"))
or id_safe_name(hidden_params.get("custom_llm_provider"))
or id_safe_name(fallback_provider)
)
if not model_id or not provider:
return response
if (
model_id == decoded.get("model_id")
and deployment_id == decoded.get("deployment_id")
and kept_requested_model == decoded.get("requested_model")
):
return response
restamped_id: Final = encode_video_id_with_provider(
video_id=decoded.get("video_id") or response.id,
provider=provider,
model_id=model_id,
served_model=decoded.get("served_model"),
deployment_id=deployment_id,
requested_model=kept_requested_model,
)
return response.model_copy(update={"id": restamped_id})

42
litellm/responses/bridge_check.py
View file

@ -4,7 +4,10 @@ from urllib.parse import urlsplit
import litellm
from litellm._logging import verbose_logger
from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider
from litellm.litellm_core_utils.get_llm_provider_logic import (
get_llm_provider,
resolve_model_and_provider_for_metadata,
)
from litellm.llms.azure_ai.chat.transformation import AzureAIGPT5Config
from litellm.llms.azure_ai.common_utils import (
azure_ai_supports_native_responses,
@ -69,7 +72,7 @@ def responses_api_bridge_check(
reasoning_effort: str | Mapping[str, object] | None = None,
reasoning_summary: object | None = None,
api_base: str | None = None,
) -> tuple[dict, str]:
) -> tuple[dict[str, object], str]:
model_info: dict[str, object] = {}
# Global flag: route ALL OpenAI chat completions through Responses API.
@ -186,3 +189,38 @@ def responses_api_bridge_check(
model = model.replace("responses/", "")
return model_info, model
def chat_completion_ends_in_responses_api(
model: str | None,
custom_llm_provider: str | None,
*,
tools: Sequence[Mapping[str, object]] | None = None,
reasoning_effort: str | Mapping[str, object] | None = None,
reasoning_summary: object | None = None,
api_base: str | None = None,
) -> bool:
"""Whether chat completion for this model is served by a Responses API.
Encrypted reasoning can be replayed only there. A normal chat provider cannot
verify the block, so the messages adapter strips it. ``chatgpt``, a ``responses/``
model, and a request whose tools or reasoning effort the bridge sends to Responses keep it.
"""
if model is None and custom_llm_provider is None:
return False
try:
resolved_model, provider = resolve_model_and_provider_for_metadata(model or "", custom_llm_provider)
except Exception as error: # noqa: BLE001 # unknown model names must not fail message translation
verbose_logger.debug("responses replay check failed for %s: %s", model, error)
return custom_llm_provider == LlmProviders.CHATGPT.value or (
isinstance(model, str) and (model == "chatgpt" or model.startswith("chatgpt/"))
)
model_info, _resolved = responses_api_bridge_check(
model=resolved_model,
custom_llm_provider=provider,
tools=tools,
reasoning_effort=reasoning_effort,
reasoning_summary=reasoning_summary,
api_base=api_base,
)
return model_info.get("mode") == "responses"

7
litellm/responses/main.py
View file

@ -41,6 +41,7 @@ from litellm.responses.litellm_completion_transformation.handler import (
LiteLLMCompletionTransformationHandler,
)
from litellm.responses.mcp.request_context import MCPRequestContext
from litellm.responses.mcp.tool_precheck import may_name_litellm_mcp_gateway
from litellm.responses.utils import ResponsesAPIRequestUtils
from litellm.types.llms.openai import (
PromptObject,
@ -1112,11 +1113,7 @@ def _responses_try_dispatch_mcp_gateway(
| None
):
"""Return a response when MCP gateway handles the call; otherwise None."""
from litellm.responses.mcp.litellm_proxy_mcp_handler import (
LiteLLM_Proxy_MCP_Handler,
)
if skip_mcp_handler or not LiteLLM_Proxy_MCP_Handler.should_use_litellm_mcp_gateway(tools=tools):
if skip_mcp_handler or not may_name_litellm_mcp_gateway(tools):
return None
mcp_call_kwargs: Final = {
"model": model,

44
litellm/responses/mcp/litellm_proxy_mcp_handler.py
View file

@ -1,4 +1,3 @@
import re
import traceback
from collections.abc import Awaitable, Callable, Collection, Iterable, Mapping, Sequence
from datetime import datetime
@ -20,6 +19,12 @@ from litellm.proxy._experimental.mcp_server.utils import (
strip_known_server_prefix,
)
from litellm.responses.main import aresponses # noqa: TID251 # inner call must skip the MCP gateway that invoked it
from litellm.responses.mcp.tool_precheck import (
LITELLM_PROXY_MCP_SERVER_URL_PREFIX,
may_name_litellm_mcp_gateway,
names_gateway_explicitly,
proxy_path_mcp_name,
)
from litellm.responses.streaming_iterator import BaseResponsesAPIStreamingIterator
from litellm.types.llms.openai import (
ResponseInputParam,
@ -65,29 +70,6 @@ class _AresponsesOptions(ExtraItemsTypedDict, total=False, extra_items=object):
pass
LITELLM_PROXY_MCP_SERVER_URL: Final = "litellm_proxy"
LITELLM_PROXY_MCP_SERVER_URL_PREFIX: Final = f"{LITELLM_PROXY_MCP_SERVER_URL}/mcp/"
_PROXY_MCP_PATH_RE: Final = re.compile(r"^https?://.+/mcp/([^/]+)$")
def _mcp_server_url(tool: ToolParam) -> str | None:
if not isinstance(tool, dict) or tool.get("type") != "mcp":
return None
server_url: Final = tool.get("server_url")
return server_url if isinstance(server_url, str) else None
def _names_gateway_explicitly(tool: ToolParam) -> bool:
return (_mcp_server_url(tool) or "").startswith(LITELLM_PROXY_MCP_SERVER_URL)
def _proxy_path_mcp_name(tool: ToolParam) -> str | None:
server_url: Final = _mcp_server_url(tool)
match: Final = None if server_url is None else _PROXY_MCP_PATH_RE.match(server_url)
return None if match is None else match.group(1)
def _registered_mcp_servers() -> Collection[MCPServer]:
from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
global_mcp_server_manager,
@ -133,7 +115,7 @@ async def _gateway_served_names(
async def _served_mcp_path_names(
tools: Collection[ToolParam], served_names: Callable[[Collection[str]], Awaitable[frozenset[str]]]
) -> frozenset[str]:
names: Final = frozenset(name for name in map(_proxy_path_mcp_name, tools) if name is not None)
names: Final = frozenset(name for name in map(proxy_path_mcp_name, tools) if name is not None)
return await served_names(names) if names else frozenset[str]()
@ -162,17 +144,15 @@ class LiteLLM_Proxy_MCP_Handler:
@staticmethod
def should_use_litellm_mcp_gateway(tools: Iterable[ToolParam] | None) -> bool:
"""True when a tool may name this gateway: server_url "litellm_proxy..." or an http(s) URL ending in
/mcp/<name>. `_split_mcp_tools` then settles which of the latter the gateway actually serves."""
return any(_names_gateway_explicitly(tool) or _proxy_path_mcp_name(tool) is not None for tool in tools or ())
return may_name_litellm_mcp_gateway(tools)
_should_use_litellm_mcp_gateway = should_use_litellm_mcp_gateway
@staticmethod
def _parse_mcp_tools(tools: Iterable[Mapping[str, object]] | None) -> SplitTools:
items: Final = tuple(tools or ())
gateway_tools: Final[list[ToolParam]] = [tool for tool in items if _names_gateway_explicitly(tool)]
other_tools: Final[list[ToolParam]] = [tool for tool in items if not _names_gateway_explicitly(tool)]
gateway_tools: Final[list[ToolParam]] = [tool for tool in items if names_gateway_explicitly(tool)]
other_tools: Final[list[ToolParam]] = [tool for tool in items if not names_gateway_explicitly(tool)]
return gateway_tools, other_tools
@staticmethod
@ -185,7 +165,7 @@ class LiteLLM_Proxy_MCP_Handler:
return LiteLLM_Proxy_MCP_Handler._parse_mcp_tools(
[
{**tool, "server_url": f"{LITELLM_PROXY_MCP_SERVER_URL_PREFIX}{name}"}
if (name := _proxy_path_mcp_name(tool)) in served
if (name := proxy_path_mcp_name(tool)) in served
else tool
for tool in items
]
@ -201,7 +181,7 @@ class LiteLLM_Proxy_MCP_Handler:
) -> tuple[bool, ...]:
items: Final = tuple(tools or ())
served: Final = await _served_mcp_path_names(items, served_names)
return tuple(_names_gateway_explicitly(tool) or _proxy_path_mcp_name(tool) in served for tool in items)
return tuple(names_gateway_explicitly(tool) or proxy_path_mcp_name(tool) in served for tool in items)
@staticmethod
async def _apply_toolset_permissions(

43
litellm/responses/mcp/tool_precheck.py Normal file
View file

@ -0,0 +1,43 @@
import importlib
import re
from collections.abc import Iterable
from typing import Final
LITELLM_PROXY_MCP_SERVER_URL: Final = "litellm_proxy"
LITELLM_PROXY_MCP_SERVER_URL_PREFIX: Final = f"{LITELLM_PROXY_MCP_SERVER_URL}/mcp/"
_PROXY_MCP_PATH_RE: Final = re.compile(r"^https?://.+/mcp/([^/]+)$")
_GATEWAY_REGISTRY_MODULE: Final = "litellm.proxy._experimental.mcp_server.mcp_server_manager"
def _mcp_server_url(tool: object) -> str | None:
match tool:
case {"type": "mcp", "server_url": str() as server_url} if isinstance(tool, dict):
return server_url
case _:
return None
def names_gateway_explicitly(tool: object) -> bool:
return (_mcp_server_url(tool) or "").startswith(LITELLM_PROXY_MCP_SERVER_URL)
def proxy_path_mcp_name(tool: object) -> str | None:
server_url: Final = _mcp_server_url(tool)
match: Final = None if server_url is None else _PROXY_MCP_PATH_RE.match(server_url)
return None if match is None else match.group(1)
def _gateway_registry_loads() -> bool:
try:
importlib.import_module(_GATEWAY_REGISTRY_MODULE)
except ImportError:
return False
return True
def may_name_litellm_mcp_gateway(tools: Iterable[object] | None) -> bool:
items: Final = tuple(tools or ())
if any(map(names_gateway_explicitly, items)):
return True
return any(proxy_path_mcp_name(tool) is not None for tool in items) and _gateway_registry_loads()

107
litellm/router.py
View file

@ -34,6 +34,7 @@ from collections.abc import (
MutableMapping,
Sequence,
)
from contextvars import ContextVar
from dataclasses import dataclass
from datetime import datetime, timezone
from functools import lru_cache, partial
@ -223,6 +224,7 @@ from litellm.router_utils.cooldown_handlers import (
is_background_response_poll_not_found,
is_caller_scoped_auth_failure,
is_caller_timeout_408,
is_video_lookup_not_found,
set_cooldown_deployments,
)
from litellm.router_utils.deployment_pricing import deployment_pricing_models
@ -370,6 +372,7 @@ from litellm.types.utils import (
shared_backend_model_info,
)
from litellm.types.utils import ModelInfo as ModelMapInfo
from litellm.types.videos.utils import VIDEO_ID_CALL_TYPES, decode_video_id_with_provider
from litellm.utils import (
CustomStreamWrapper,
EmbeddingResponse,
@ -488,6 +491,7 @@ def model_info_is_active_for_environment(model_info: Mapping[str, object] | None
_PreRoutingStrategyT = TypeVar("_PreRoutingStrategyT")
_CallbackT = TypeVar("_CallbackT")
_DeploymentT = TypeVar("_DeploymentT", bound=Mapping[str, object])
_ALIAS_PARAMS_NEVER_FORWARDED: Final = frozenset({"model", "api_base", "api_key", "api_version"})
_ALIAS_MARKER_FORWARDED_PARAMS_KWARG: Final = "_alias_marker_forwarded_params"
@ -496,6 +500,7 @@ _CLAUDE_CODE_SESSION_ROUTER_TTL_SECONDS: Final = 3600
CLAUDE_CODE_SESSION_ROUTER_BINDING_TARGET: Final = "claude_code_session_router_binding"
_PRE_CALL_CHECK_FILTERED_PARAMS: Final = frozenset({"response_format"})
_video_call_deployment_id: Final[ContextVar[str | None]] = ContextVar("video_call_deployment_id", default=None)
_RUNTIME_TOGGLEABLE_PRE_CALL_CHECKS: Final[Mapping[str, type[CustomLogger]]] = MappingProxyType(
{
"prompt_caching": PromptCachingDeploymentCheck,
@ -5717,17 +5722,19 @@ class Router:
passthrough_on_no_deployment: Final = kwargs.pop("passthrough_on_no_deployment", False)
function_name: Final = "_ageneric_api_call_with_fallbacks"
handler_name: Final[str | None] = getattr(original_generic_function, "__name__", None)
deployment = None # rebind-ok: pre-init so the except block can stamp a failure with no deployment picked
try:
parent_otel_span: Final = get_parent_otel_span_from_kwargs(kwargs)
try:
deployment = await self.async_get_available_deployment( # rebind-ok: set on success, see pre-init above
model=model,
request_kwargs=kwargs,
messages=kwargs.get("messages", None),
input=kwargs.get("input", None),
specific_deployment=kwargs.pop("specific_deployment", None),
)
with self._video_call_pin(handler_name, kwargs):
deployment = await self.async_get_available_deployment( # rebind-ok: set on success, see pre-init
model=model,
request_kwargs=kwargs,
messages=kwargs.get("messages", None),
input=kwargs.get("input", None),
specific_deployment=kwargs.pop("specific_deployment", None),
)
except Exception as e:
if passthrough_on_no_deployment:
return await original_generic_function(model=model, **kwargs)
@ -6409,6 +6416,37 @@ class Router:
)
return await self._ageneric_api_call_with_fallbacks(original_function=original_function, **kwargs)
@staticmethod
@contextlib.contextmanager
def _video_call_pin(handler_name: str | None, request_kwargs: Mapping[str, object]) -> Generator[None]:
"""
Scopes the deployment a video id names to the routing of one call on that video (status, content, remix,
edit, extension). Any other call routes as if no id were sent, whatever its body carries.
"""
video_id: Final = request_kwargs.get("video_id")
deployment_id: Final = (
decode_video_id_with_provider(video_id).get("deployment_id")
if handler_name in VIDEO_ID_CALL_TYPES and isinstance(video_id, str)
else None
)
token: Final = _video_call_deployment_id.set(deployment_id)
try:
yield
finally:
_video_call_deployment_id.reset(token)
@staticmethod
def _narrow_to_video_deployment(healthy_deployments: Sequence[_DeploymentT]) -> Sequence[_DeploymentT]:
"""
A video lives on the provider account of the deployment that created it, so a call on an id naming a
deployment that every routing filter kept goes to that deployment, else to whichever one routing picks.
"""
deployment_id: Final = _video_call_deployment_id.get()
creating: Final = tuple(
deployment for deployment in healthy_deployments if deployment_id in Router._deployment_ids((deployment,))
)
return creating or healthy_deployments
def _generic_api_call_with_fallbacks(self, model: str, original_function: Callable, **kwargs):
"""
Make a generic LLM API call through the router, this allows you to use retries/fallbacks with litellm router
@ -6430,13 +6468,14 @@ class Router:
kwargs=kwargs,
metadata_variable_name=metadata_variable_name,
)
deployment: Final = self.get_available_deployment(
model=model,
messages=kwargs.get("messages", None),
input=kwargs.get("input", None),
specific_deployment=kwargs.pop("specific_deployment", None),
request_kwargs=kwargs,
)
with self._video_call_pin(handler_name, kwargs):
deployment: Final = self.get_available_deployment(
model=model,
messages=kwargs.get("messages", None),
input=kwargs.get("input", None),
specific_deployment=kwargs.pop("specific_deployment", None),
request_kwargs=kwargs,
)
self._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs, function_name="generic_api_call")
data: Final = deployment["litellm_params"].copy()
@ -8914,6 +8953,13 @@ class Router:
)
return False
if is_video_lookup_not_found(exception, kwargs):
verbose_router_logger.debug(
"Router: Exiting 'deployment_callback_on_failure' without cooldown. "
"Provider 404 says the video is not on this deployment's account, not that the deployment is down."
)
return False
exception_status: Final = getattr(exception, "status_code", "")
if is_caller_timeout_408(kwargs, exception_status):
@ -12337,7 +12383,8 @@ class Router:
1. First, check if model_id directly matches a model_name or deployment ID
2. If not, search through router's model_list to find a match by litellm_params.model
3. If not, keep model_id when the router routes that name itself (an alias, a routing group, or a wildcard
whose deployment serves `custom_llm_provider`)
whose deployment serves `custom_llm_provider`), else name it `custom_llm_provider/model_id` when only that
form routes to such a wildcard
4. Return the model_name if found, None otherwise
Args:
@ -12385,8 +12432,15 @@ class Router:
if self.is_recognized_model(model_id):
return model_id
if custom_llm_provider is not None and self.pattern_router.routes_to_provider(model_id, custom_llm_provider):
if custom_llm_provider is None:
return None
if self.pattern_router.routes_to_provider(model_id, custom_llm_provider):
return model_id
provider_prefixed: Final = f"{custom_llm_provider}/{model_id}"
if not model_id.startswith(f"{custom_llm_provider}/") and self.pattern_router.routes_to_provider(
provider_prefixed, custom_llm_provider
):
return provider_prefixed
return None
def resolve_served_model_name(self, model_name: str, served_model: str) -> str:
@ -12915,13 +12969,13 @@ class Router:
return [{**DeploymentTypedDict(**m), "model_name": model_name} for m in matches]
def wildcard_route_deployments(self, model: str, team_id: str | None) -> Sequence[Mapping[str, object]]:
"""The deployments routing serves ``model`` from through a wildcard: the most specific
proxy-wide pattern matching it, else ``team_id``'s own pattern, never both."""
proxy_wide: Final[Sequence[Mapping[str, object]]] = self.pattern_router.get_deployments_by_pattern(model=model)
"""The deployments routing serves ``model`` from through a wildcard: ``team_id``'s own pattern
matching it, else the most specific proxy-wide pattern, never both."""
team_router: Final = self.team_pattern_routers.get(team_id) if team_id is not None else None
if proxy_wide or team_router is None:
return proxy_wide
return team_router.get_deployments_by_pattern(model=model)
team_owned: Final[Sequence[Mapping[str, object]]] = (
team_router.get_deployments_by_pattern(model=model) if team_router is not None else ()
)
return team_owned or self.pattern_router.get_deployments_by_pattern(model=model)
def team_ptu_ceiling(self, team_id: str, model_group: str) -> PTUTeamCeiling | None:
"""The per-minute ceiling ``team_id``'s PTU shares set on the group serving ``model_group``,
@ -13897,10 +13951,10 @@ class Router:
def _filter_deployments_by_model_access_groups(
self,
model: str,
healthy_deployments: list,
healthy_deployments: list[DeploymentTypedDict],
request_kwargs: Mapping[str, object] | None,
request_team_id: str | None,
) -> list:
) -> list[DeploymentTypedDict]:
"""
Restrict candidate deployments to caller-authorized model access groups.
@ -14084,7 +14138,8 @@ class Router:
## ORDER FILTERING ## -> if user set 'order' in deployments, return deployments with lowest order (e.g. order=1 > order=2)
_target_order: Final = (request_kwargs or {}).pop("_target_order", None)
healthy_deployments = litellm.utils.get_order_filtered_deployments(
cast(list[dict], healthy_deployments), target_order=_target_order
list(self._narrow_to_video_deployment(cast(list[dict[str, object]], healthy_deployments))),
target_order=_target_order,
)
## WEIGHTED FAILOVER EXCLUSION ## -> drop deployments already tried in
@ -15271,7 +15326,7 @@ class Router:
## ORDER FILTERING ## -> if user set 'order' in deployments, return deployments with lowest order (e.g. order=1 > order=2)
_target_order: Final = (request_kwargs or {}).pop("_target_order", None)
healthy_deployments = litellm.utils.get_order_filtered_deployments(
healthy_deployments, target_order=_target_order
list(self._narrow_to_video_deployment(healthy_deployments)), target_order=_target_order
)
## WEIGHTED FAILOVER EXCLUSION ## -> drop deployments already tried in

6
litellm/router_utils/cooldown_handlers.py
View file

@ -35,6 +35,7 @@ from litellm.types.utils import (
BACKGROUND_RESPONSE_STATUS_POLL_CALL_ORIGIN,
InternalCallOrigin,
)
from litellm.types.videos.utils import VIDEO_ID_CALL_TYPES
from .router_callbacks.track_deployment_metrics import (
get_deployment_failures_for_current_minute,
@ -89,6 +90,11 @@ def is_background_response_poll_not_found(exception: Exception, litellm_params:
)
def is_video_lookup_not_found(exception: object, failure_kwargs: Mapping[str, object]) -> bool:
"""Whether a call on an existing video id failed with a provider 404: the video is not on that account."""
return getattr(exception, "status_code", None) == 404 and failure_kwargs.get("call_type") in VIDEO_ID_CALL_TYPES
_EXCEPTION_POLICY_FIELDS: Final[tuple[tuple[type, str], ...]] = (
# ContentPolicyViolationError subclasses BadRequestError, so it must be checked first.
(litellm.ContentPolicyViolationError, "ContentPolicyViolationErrorAllowedFails"),

2
litellm/router_utils/pattern_match_deployments.py
View file

@ -173,7 +173,7 @@ class PatternMatchRouter:
return None # No matching pattern found
def routes_to_provider(self, request: str, custom_llm_provider: str) -> bool:
deployments: Final = _ROUTED_DEPLOYMENTS.validate_python(self.route(request) or ())
deployments: Final = _ROUTED_DEPLOYMENTS.validate_python(self.get_pattern(request) or ())
return any(
(deployment.litellm_params.custom_llm_provider or inferred_provider(deployment.litellm_params.model))
== custom_llm_provider

4
litellm/types/llms/anthropic.py
View file

@ -545,7 +545,9 @@ class ContentBlockStartText(TypedDict):
content_block: TextBlock
ContentBlockContentBlockDict = ToolUseBlock | TextBlock | ChatCompletionThinkingBlock
ContentBlockContentBlockDict = (
ToolUseBlock | TextBlock | ChatCompletionThinkingBlock | ChatCompletionRedactedThinkingBlock
)
ContentBlockStart = ContentBlockStartToolUse | ContentBlockStartText

2
litellm/types/videos/main.py
View file

@ -113,6 +113,8 @@ class DecodedVideoId(TypedDict, total=False):
model_id: str | None
video_id: str
served_model: ReadOnly[str | None]
deployment_id: ReadOnly[str | None]
requested_model: ReadOnly[str | None]
class CharacterObject(LiteLLMBaseModel):

111
litellm/types/videos/utils.py
View file

@ -6,6 +6,7 @@ Format: vid_{base64_encoded_string}
"""
import base64
from collections.abc import Sequence
from typing import Final
from litellm._logging import verbose_logger
@ -14,6 +15,22 @@ from litellm.types.videos.main import DecodedVideoId
VIDEO_ID_PREFIX: Final = "video_"
VIDEO_ID_SERVED_MODEL_PREFIX: Final = "served_model:"
VIDEO_ID_DEPLOYMENT_PREFIX: Final = "deployment_id:"
VIDEO_ID_REQUESTED_MODEL_PREFIX: Final = "requested_model:"
VIDEO_ID_CALL_TYPES: Final = frozenset(
{
"avideo_status",
"video_status",
"avideo_content",
"video_content",
"avideo_remix",
"video_remix",
"avideo_edit",
"video_edit",
"avideo_extension",
"video_extension",
}
)
CHARACTER_ID_PREFIX: Final = "character_"
CHARACTER_ID_TEMPLATE: Final = "litellm:custom_llm_provider:{};model_id:{};character_id:{}"
@ -37,7 +54,12 @@ def _add_base64_padding(value: str) -> str:
def encode_video_id_with_provider(
video_id: str, provider: str, model_id: str | None = None, served_model: str | None = None
video_id: str,
provider: str,
model_id: str | None = None,
served_model: str | None = None,
deployment_id: str | None = None,
requested_model: str | None = None,
) -> str:
"""Encode provider and model_id into video_id using base64."""
if not provider or not video_id:
@ -54,12 +76,24 @@ def encode_video_id_with_provider(
# ID is not encoded (even if it starts with video_), so encode it
assembled_id = str(SpecialEnums.LITELLM_MANAGED_VIDEO_COMPLETE_STR.value).format(provider, model_id or "", video_id)
served_part: Final = f";{VIDEO_ID_SERVED_MODEL_PREFIX}{served_model}" if served_model else ""
deployment_part: Final = f";{VIDEO_ID_DEPLOYMENT_PREFIX}{deployment_id}" if deployment_id else ""
requested_part: Final = f";{VIDEO_ID_REQUESTED_MODEL_PREFIX}{requested_model}" if requested_model else ""
base64_encoded_id: Final[str] = base64.b64encode((assembled_id + served_part).encode("utf-8")).decode("utf-8")
base64_encoded_id: Final[str] = base64.b64encode(
(assembled_id + served_part + deployment_part + requested_part).encode("utf-8")
).decode("utf-8")
return f"{VIDEO_ID_PREFIX}{base64_encoded_id}"
def id_safe_name(name: object) -> str | None:
return name if isinstance(name, str) and name and ";" not in name else None
def _trailing_part(parts: Sequence[str], prefix: str) -> str | None:
return next((part.removeprefix(prefix) or None for part in parts[3:] if part.startswith(prefix)), None)
def decode_video_id_with_provider(encoded_video_id: str) -> DecodedVideoId:
"""Decode provider and model_id from encoded video_id."""
if not encoded_video_id:
@ -80,43 +114,6 @@ def decode_video_id_with_provider(encoded_video_id: str) -> DecodedVideoId:
cleaned_id = encoded_video_id.replace(VIDEO_ID_PREFIX, "")
cleaned_id = _add_base64_padding(cleaned_id)
decoded_id: Final = base64.b64decode(cleaned_id.encode("utf-8")).decode("utf-8")
if ";" not in decoded_id:
return DecodedVideoId(
custom_llm_provider=None,
model_id=None,
video_id=encoded_video_id,
)
parts: Final = decoded_id.split(";")
custom_llm_provider = None
model_id = None
decoded_video_id = encoded_video_id
if len(parts) >= 3:
custom_llm_provider_part: Final = parts[0]
model_id_part: Final = parts[1]
video_id_part: Final = parts[2]
custom_llm_provider = custom_llm_provider_part.replace("litellm:custom_llm_provider:", "")
model_id = model_id_part.replace("model_id:", "")
decoded_video_id = video_id_part.replace("video_id:", "")
served_model: Final = (
parts[3].removeprefix(VIDEO_ID_SERVED_MODEL_PREFIX) or None
if len(parts) >= 4 and parts[3].startswith(VIDEO_ID_SERVED_MODEL_PREFIX)
else None
)
if served_model is None:
return DecodedVideoId(custom_llm_provider=custom_llm_provider, model_id=model_id, video_id=decoded_video_id)
return DecodedVideoId(
custom_llm_provider=custom_llm_provider,
model_id=model_id,
video_id=decoded_video_id,
served_model=served_model,
)
except Exception as e:
verbose_logger.debug("Error decoding video_id '%s': %s", encoded_video_id, e)
return DecodedVideoId(
@ -125,6 +122,42 @@ def decode_video_id_with_provider(encoded_video_id: str) -> DecodedVideoId:
video_id=encoded_video_id,
)
if ";" not in decoded_id:
return DecodedVideoId(
custom_llm_provider=None,
model_id=None,
video_id=encoded_video_id,
)
parts: Final = decoded_id.split(";")
custom_llm_provider = None
model_id = None
decoded_video_id = encoded_video_id
if len(parts) >= 3:
custom_llm_provider_part: Final = parts[0]
model_id_part: Final = parts[1]
video_id_part: Final = parts[2]
custom_llm_provider = custom_llm_provider_part.replace("litellm:custom_llm_provider:", "")
model_id = model_id_part.replace("model_id:", "")
decoded_video_id = video_id_part.replace("video_id:", "")
served_model: Final = _trailing_part(parts, VIDEO_ID_SERVED_MODEL_PREFIX)
deployment_id: Final = _trailing_part(parts, VIDEO_ID_DEPLOYMENT_PREFIX)
requested_model: Final = _trailing_part(parts, VIDEO_ID_REQUESTED_MODEL_PREFIX)
decoded_parts: Final = DecodedVideoId(
custom_llm_provider=custom_llm_provider, model_id=model_id, video_id=decoded_video_id
)
with_served: Final[DecodedVideoId] = (
decoded_parts if served_model is None else {**decoded_parts, "served_model": served_model}
)
with_deployment: Final[DecodedVideoId] = (
with_served if deployment_id is None else {**with_served, "deployment_id": deployment_id}
)
return with_deployment if requested_model is None else {**with_deployment, "requested_model": requested_model}
def extract_original_video_id(encoded_video_id: str) -> str:
"""Extract original video ID without encoding."""

36
litellm/utils.py
View file

@ -51,7 +51,7 @@ import litellm.litellm_core_utils
# audio_utils.utils is lazy-loaded - only imported when needed for transcription calls
import litellm.litellm_core_utils.json_validation_rule
from litellm._internal_context import is_internal_call
from litellm._internal_context import is_internal_call, is_resent_failure
from litellm._lazy_imports import (
_get_modified_max_tokens,
get_default_encoding,
@ -1980,12 +1980,13 @@ def client(original_function):
retry_policy=kwargs.get("retry_policy"),
)
kwargs["retry_policy"] = reset_retry_policy() # prevent infinite loops
litellm.num_retries = None # set retries to None to prevent infinite loops
if not is_resent_failure(e):
litellm.num_retries = None # set retries to None to prevent infinite loops
context_window_fallback_dict: Final = kwargs.get("context_window_fallback_dict", {})
is_completion_litellm_router_call: Final = _is_litellm_router_call(kwargs, is_async=False)
if (
num_retries and not is_completion_litellm_router_call
num_retries and not is_completion_litellm_router_call and not is_resent_failure(e)
): # only enter this if call is not from litellm router/proxy. router has it's own logic for retrying
if (
isinstance(e, openai.APIError)
@ -2015,11 +2016,12 @@ def client(original_function):
retry_policy=kwargs.get("retry_policy"),
)
kwargs["retry_policy"] = reset_retry_policy() # prevent infinite loops
litellm.num_retries = None # set retries to None to prevent infinite loops
if not is_resent_failure(e):
litellm.num_retries = None # set retries to None to prevent infinite loops
is_responses_litellm_router_call: Final = _is_litellm_router_call(kwargs, is_async=False)
if (
num_retries and not is_responses_litellm_router_call
num_retries and not is_responses_litellm_router_call and not is_resent_failure(e)
): # only enter this if call is not from litellm router/proxy. router has it's own logic for retrying
if (
isinstance(e, openai.APIError)
@ -2032,7 +2034,7 @@ def client(original_function):
end_time = datetime.datetime.now()
# LOG FAILURE - handle streaming failure logging in the _next_ object, remove `handle_failure` once it's deprecated
if logging_obj:
if logging_obj and not is_resent_failure(e):
logging_obj.failure_handler(
e, traceback_exception, start_time, end_time
) # DO NOT MAKE THREADED - router retry fallback relies on this!
@ -2279,7 +2281,7 @@ def client(original_function):
# Reuse the timestamp taken right when the deployment call itself failed, before
# the failure hook ran, so a slow callback doesn't inflate the reported duration.
end_time = _deployment_call_end_time if _deployment_call_end_time is not None else datetime.datetime.now() # noqa: DTZ005 # matches the naive datetimes this whole function already times start_time/end_time with
if logging_obj and not _is_litellm_internal_call:
if logging_obj and not _is_litellm_internal_call and not is_resent_failure(e):
unbind_budget_reservation_from_callbacks(logging_obj.litellm_params)
try:
logging_obj.failure_handler(
@ -2300,7 +2302,10 @@ def client(original_function):
is_acompletion_litellm_router_call: Final = _is_litellm_router_call(kwargs, is_async=True)
if (
num_retries and not is_acompletion_litellm_router_call and not isinstance(e, ImportError)
num_retries
and not is_acompletion_litellm_router_call
and not isinstance(e, ImportError)
and not is_resent_failure(e)
): # only enter this if call is not from litellm router/proxy. router has it's own logic for retrying
try:
litellm.num_retries = None # set retries to None to prevent infinite loops
@ -2311,8 +2316,9 @@ def client(original_function):
elif isinstance(e, openai.APIError): # generic api error
kwargs["retry_strategy"] = "constant_retry"
result = await litellm.acompletion_with_retries(*args, **kwargs)
except Exception:
pass
except Exception as retry_error:
if is_resent_failure(retry_error):
raise
else:
return result
elif (
@ -2331,7 +2337,10 @@ def client(original_function):
is_aresponses_litellm_router_call: Final = _is_litellm_router_call(kwargs, is_async=True)
if (
num_retries and not is_aresponses_litellm_router_call and not isinstance(e, ImportError)
num_retries
and not is_aresponses_litellm_router_call
and not isinstance(e, ImportError)
and not is_resent_failure(e)
): # only enter this if call is not from litellm router/proxy. router has it's own logic for retrying
try:
litellm.num_retries = None # set retries to None to prevent infinite loops
@ -2342,8 +2351,9 @@ def client(original_function):
elif isinstance(e, openai.APIError): # generic api error
kwargs["retry_strategy"] = "constant_retry"
result = await litellm.aresponses_with_retries(*args, **kwargs)
except Exception:
pass
except Exception as retry_error:
if is_resent_failure(retry_error):
raise
else:
return result

2
packaging/litellm-core/pyproject.toml
View file

@ -11,7 +11,7 @@ authors = [
]
dependencies = [
"fastuuid>=0.14.0,<1.0",
"filelock>=3.16.1,<4.0",
"filelock>=3.16.1,<5",
"httpx[http2]>=0.28.0,<1.0",
"openai>=2.20.0,<4.0.0",
"python-dateutil>=2.8.2,<3.0",

19
pyproject.toml
View file

@ -1,6 +1,6 @@
[project]
name = "litellm"
version = "1.108.2"
version = "1.108.3"
description = "Library to easily interface with LLM API providers"
readme = "README.md"
requires-python = ">=3.10, <3.15"
@ -15,7 +15,7 @@ dependencies = [
# When changing a floor, verify it installs + imports on every supported
# Python with: `uv pip install --resolution=lowest-direct .`
"fastuuid>=0.14.0,<1.0",
"filelock>=3.16.1,<4.0",
"filelock>=3.16.1,<5",
"httpx[http2]>=0.28.0,<1.0",
"openai>=2.20.0,<4.0.0",
"python-dotenv>=1.0.0,<2.0",
@ -75,8 +75,8 @@ proxy = [
"boto3>=1.43.117,<2.0",
"azure-identity>=1.25.2,<2.0",
"azure-storage-blob>=12.28.0,<13.0",
"mcp>=2.2.0,<3",
"httpx2>=2.5.0,<3",
"mcp>=2.3.0,<3",
"httpx2>=2.10.0,<3",
"pydantic>=2.12.0,<3",
"litellm-proxy-extras==0.4.110",
"litellm-enterprise==0.1.79.post1",
@ -124,7 +124,7 @@ utils = [
"numpydoc>=1.8.0,<2.0",
]
caching = ["diskcache>=5.6.3,<6.0"]
mcp = ["mcp>=2.2.0,<3", "httpx2>=2.5.0,<3", "pydantic>=2.12.0,<3"]
mcp = ["mcp>=2.3.0,<3", "httpx2>=2.10.0,<3", "pydantic>=2.12.0,<3"]
# Driver for the MongoDB Atlas vector store; Atlas Vector Search has no HTTP query API.
# The floor is 4.9 because that is the release AsyncMongoClient landed in.
# SAML SSO for the admin UI. python3-saml pulls in xmlsec/lxml, whose wheels
@ -242,7 +242,7 @@ e2e-dev = [
"locust==2.45.0",
"anthropic==0.84.0",
"psutil==7.2.2",
"mcp>=2.2.0,<3",
"mcp>=2.3.0,<3",
]
admin-mcp = [
"litellm-admin-mcp @ https://github.com/BerriAI/liteadmin-mcp/archive/d35ec9c19c117d4c50cc1eccf6ce1296aac25a1a.tar.gz#sha256=86b930f6706fb2da10d53d1798ee7fd14e44fb0afdef0da122cf0e6bb88fa2c3 ; python_version >= '3.12'",
@ -302,7 +302,7 @@ healthcheck = [
benchmarks = [
"pytest==9.0.3",
"pytest-codspeed==4.3.0",
"mcp>=2.2.0,<3",
"mcp>=2.3.0,<3",
"a2a-sdk==1.1.0",
]
@ -358,11 +358,14 @@ constraint-dependencies = [
"httplib2>=0.32.0",
"setuptools>=84.1.0",
"dnspython>=2.9.0; python_version >= '3.11'",
"filelock<3.20.4; python_version < '3.11'",
"filelock>=4.2.1; python_version >= '3.11'",
]
override-dependencies = [
# a2a-sdk 1.x requires packaging>=24.0; lunary 1.4.x still caps at <24.0.
"packaging>=24.0",
"cryptography>=50.0.0,<51.0",
"mcp>=2.3.0,<3",
]
default-groups = ["dev"]
required-version = ">=0.10.9"
@ -393,7 +396,7 @@ litellm-enterprise = { workspace = true }
members = ["enterprise", "litellm-proxy-extras"]
[tool.commitizen]
version = "1.108.2"
version = "1.108.3"
version_files = [
"pyproject.toml:^version",
]

10
scripts/pre_commit_lint.sh
View file

@ -18,7 +18,8 @@
# -> ruff over ruff-tests.toml + `make lint-test-quality` (test-linting.yml's
# test-tree ruff and test-quality gate steps)
# - dashboard -> prettier + eslint + lint budgets (test-linting.yml's ui job)
# - proxy/types -> regenerate the lazy OpenAPI snapshot and dashboard API types, fail on drift (test-linting.yml's ui-api-types job)
# - spec_pattern -> regenerate the lazy OpenAPI snapshot and dashboard API types, fail on drift
# (test-linting.yml's ui-api-types job)
#
# Each block is skipped when no matching files are in scope, so unrelated commits
# stay fast. This is intentionally not auto-installed as a git hook (see
@ -98,7 +99,7 @@ existing_files() {
litellm_py_pattern='^litellm/.*\.py$'
e2e_py_pattern='^tests/e2e(_harness)?/.*\.py$'
test_tree_pattern='^(tests/.*\.py|ruff-tests\.toml|scripts/(check_test_quality|test_quality_gate)\.py)$'
spec_pattern='^(litellm/(proxy|types)/.*|ui/litellm-dashboard/(scripts/gen-api-types\.mjs|package\.json|package-lock\.json|src/lib/http/schema\.d\.ts))$'
spec_pattern='^(litellm/.*|enterprise/.*|pyproject\.toml|uv\.lock|ui/litellm-dashboard/(scripts/gen-api-types\.mjs|package\.json|package-lock\.json|src/lib/http/schema\.d\.ts))$'
ui_prettier_pattern='^ui/litellm-dashboard/.*\.(js|jsx|ts|tsx|mjs|cjs|json|css|scss|md|mdx|yml|yaml|html)$'
ui_eslint_pattern='^ui/litellm-dashboard/.*\.(js|jsx|ts|tsx|mjs|cjs)$'
@ -110,9 +111,6 @@ e2e_py_files=$(scope_match "$e2e_py_pattern")
test_tree_files=$(scope_match "$test_tree_pattern")
# ruff format (and CI's format step) skip enterprise; the rest of make lint covers it.
fmt_files=$(printf '%s\n' "$litellm_py_files" | grep -v '^litellm/enterprise/' | existing_files)
# test-linting.yml triggers on any file under litellm/proxy or litellm/types
# (Prisma schema and configs included, not just Python) plus the generator and its
# lockfiles, so match that whole trigger set rather than a Python subset.
spec_files=$(scope_match "$spec_pattern")
# CI's ui job runs prettier over a wider extension set than eslint; keep that
# split so this flags exactly what the job would.
@ -337,7 +335,7 @@ summary_item "tests/e2e checks (basedpyright + raw HTTP client ban)" "$e2e_py_fi
summary_item "test-tree lint (ruff-tests.toml + test-quality gate)" "$test_tree_files" \
"no tests/ Python files or test-tree lint inputs in scope"
summary_item "dashboard lint (prettier + eslint + lint budgets)" "$ui_prettier_changed$ui_eslint_changed" "no dashboard files in scope"
summary_item "dashboard API-type sync (npm run gen:api)" "$spec_files" "no litellm/proxy, litellm/types, or generator files in scope"
summary_item "dashboard API-type sync (npm run gen:api)" "$spec_files" "no litellm/, enterprise/, dependency, or generator files in scope"
if [ -z "$litellm_py_files$e2e_py_files$test_tree_files$ui_prettier_changed$ui_eslint_changed$spec_files" ]; then
echo "check: NOTE - no gating lint check matches the files in scope, so nothing ran:" >&2

547
scripts/sync_together_ai_models.py
View file

@ -1,547 +0,0 @@
"""Sync the together_ai entries of model_prices_and_context_window.json with Together's live serverless catalog.
Pulls ``GET https://api.together.ai/v1/models?serverless`` plus the deprecations doc, maps API fields onto
registry fields, merges the reviewed capability rules below for everything the API cannot express, and diffs
the result against the registry. Dry run (the default) prints the diff summary and the generated PR body;
``--write`` applies the changes to the root cost map and its ``litellm/`` backup copy.
Policy highlights:
- Prices arrive per 1M tokens with float artifacts and are normalized to clean per-token values.
- A registry entry absent from the serverless catalog is marked with ``deprecation_date`` from the docs
deprecation table, never deleted; absences with no docs date are surfaced for a human call.
- Availability comes from the API: a model the docs list as removed but the API still serves stays live,
with the conflict surfaced as a warning.
- Manually curated values the API cannot express (``metadata.successor``, ``max_output_tokens`` on existing
entries, capability flags no rule covers) are never overwritten; conflicts are surfaced instead.
"""
import argparse
import json
import os
import re
import sys
from collections.abc import Mapping, Sequence
from dataclasses import dataclass, field
from pathlib import Path
from types import MappingProxyType
from typing import Final
import httpx
from pydantic import BaseModel, TypeAdapter, ValidationError
MODELS_URL: Final = "https://api.together.ai/v1/models?serverless"
DEPRECATIONS_URL: Final = "https://docs.together.ai/docs/deprecations.md"
PROVIDER: Final = "together_ai"
PREFIX: Final = "together_ai/"
SOURCE_URL: Final = "https://docs.together.ai/docs/serverless-models"
COST_MAP_RELPATHS: Final = (
"model_prices_and_context_window.json",
"litellm/model_prices_and_context_window_backup.json",
)
TYPE_TO_MODE: Final = MappingProxyType({"chat": "chat", "embedding": "embedding", "moderation": "chat"})
class SyncError(RuntimeError):
pass
class CatalogPricing(BaseModel):
input: float
output: float
cached_input: float | None = None
class CatalogModel(BaseModel):
id: str
type: str
context_length: int | None = None
pricing: CatalogPricing
CATALOG_ADAPTER: Final = TypeAdapter(list[CatalogModel])
RegistryEntry = dict[str, object]
CostMap = dict[str, object]
@dataclass(frozen=True, slots=True)
class CapabilityRule:
model_id: str
fields: Mapping[str, bool | int]
provenance: str
def _rule(model_id: str, provenance: str, **fields: bool | int) -> CapabilityRule:
return CapabilityRule(model_id=model_id, fields=MappingProxyType(dict(fields)), provenance=provenance)
_TOOLS: Final = MappingProxyType(
{
"supports_function_calling": True,
"supports_parallel_function_calling": True,
"supports_response_schema": True,
"supports_tool_choice": True,
}
)
CAPABILITY_RULES: Final = (
_rule(
"MiniMaxAI/MiniMax-M3",
"reviewed for the LIT-5968 backfill against https://www.together.ai/models/minimax-m3",
**_TOOLS,
supports_reasoning=True,
supports_vision=True,
),
_rule("Prism-ML/Ternary-Bonsai-27B", "reviewed for the LIT-5968 backfill; no tool or vision support documented"),
_rule(
"Qwen/Qwen3.5-9B",
"reviewed for the LIT-5968 backfill against https://www.together.ai/models/qwen3-5-9b",
**_TOOLS,
supports_reasoning=True,
supports_vision=True,
),
_rule(
"Qwen/Qwen3.6-Plus",
"reviewed for the LIT-5968 backfill; hybrid reasoning model without a documented tools contract",
supports_reasoning=True,
),
_rule("Qwen/Qwen3.7-Max", "reviewed for the LIT-5968 backfill; no tool or vision support documented"),
_rule("Qwen/Qwen3.7-Plus", "reviewed for the LIT-5968 backfill; no tool or vision support documented"),
_rule("Qwen/Qwen3.8-2.4T-A95B", "reviewed for the LIT-5968 backfill; no tool or vision support documented"),
_rule("arize-ai/qwen-2-1.5b-instruct", "reviewed for the LIT-5968 backfill; no tool or vision support documented"),
_rule(
"deepseek-ai/DeepSeek-V4-Flash-0731",
"reviewed for the LIT-5968 backfill against https://www.together.ai/models/deepseek-v4-flash",
**_TOOLS,
),
_rule(
"deepseek-ai/DeepSeek-V4-Pro",
"reviewed for the LIT-5968 backfill against https://www.together.ai/models/deepseek-v4-pro",
**_TOOLS,
supports_reasoning=True,
),
_rule(
"deepseek-ai/DeepSeek-V4-Pro-0813",
"reviewed for the LIT-5968 backfill against https://www.together.ai/models/deepseek-v4-pro",
**_TOOLS,
),
_rule("google/gemma-3n-E4B-it", "reviewed for the LIT-5968 backfill; no tool or vision support documented"),
_rule(
"google/gemma-4-31B-it",
"reviewed for the LIT-5968 backfill against https://www.together.ai/models/gemma-4-31b-it",
**_TOOLS,
supports_vision=True,
),
_rule(
"intfloat/multilingual-e5-large-instruct",
"embedding dims per https://huggingface.co/intfloat/multilingual-e5-large-instruct",
output_vector_size=1024,
),
_rule(
"meta-llama/Llama-3.3-70B-Instruct-Turbo",
"reviewed for the LIT-5968 backfill against https://docs.together.ai/docs/function-calling",
**_TOOLS,
),
_rule(
"meta-llama/Llama-Guard-4-12B",
"moderation classifier with a chat-shaped API; no tools per the LIT-5968 backfill review",
),
_rule("meta-models/Muse-Glimmer-30B", "reviewed for the LIT-5968 backfill; no tool or vision support documented"),
_rule(
"moonshotai/Kimi-K2.7-Code",
"reviewed for the LIT-5968 backfill against https://www.together.ai/models/kimi-k2-7-code",
**_TOOLS,
supports_vision=True,
),
_rule(
"moonshotai/Kimi-K3",
"reviewed for the LIT-5968 backfill against https://www.together.ai/models/kimi-k3",
**_TOOLS,
supports_reasoning=True,
supports_vision=True,
),
_rule(
"nvidia/nemotron-3-ultra-550b-a55b",
"reviewed for the LIT-5968 backfill against https://www.together.ai/models/nemotron-3-ultra",
**_TOOLS,
supports_reasoning=True,
),
_rule(
"openai/gpt-oss-120b",
"reviewed for the LIT-5968 backfill against https://www.together.ai/models/gpt-oss-120b",
**_TOOLS,
supports_reasoning=True,
),
_rule(
"openai/gpt-oss-20b",
"reviewed for the LIT-5968 backfill against https://www.together.ai/models/gpt-oss-20b",
**_TOOLS,
),
_rule("pearl-ai/gemma-4-31b-it", "reviewed for the LIT-5968 backfill; no tool or vision support documented"),
_rule(
"thinkingmachines/Inkling",
"reviewed for the LIT-5968 backfill against https://www.together.ai/models/inkling",
**_TOOLS,
),
_rule("thinkingmachines/Inkling-Small", "reviewed for the LIT-5968 backfill; no tool or vision support documented"),
_rule(
"zai-org/GLM-5.2",
"reviewed for the LIT-5968 backfill against https://www.together.ai/models/glm-5-2;"
" 128K output ceiling per https://docs.z.ai/guides/llm/glm-5.2",
**_TOOLS,
supports_reasoning=True,
max_output_tokens=128000,
max_tokens=128000,
),
_rule(
"zai-org/GLM-5.3-Flash",
"reviewed for LIT-6489 against https://www.together.ai/models/glm-5-3-flash;"
" 128K output ceiling per https://docs.z.ai/guides/llm/glm-5.3",
**_TOOLS,
supports_reasoning=True,
supports_vision=True,
max_output_tokens=128000,
max_tokens=128000,
),
)
RULES_BY_ID: Final = MappingProxyType({rule.model_id: rule for rule in CAPABILITY_RULES})
@dataclass(frozen=True, slots=True)
class DeprecationDoc:
removal_dates: Mapping[str, str]
redirects: Mapping[str, str]
_REDIRECT_ROW: Final = re.compile(r"^\|\s*`([^`]+)`\s*\|\s*`([^`]+)`\s*\|")
_REMOVAL_ROW: Final = re.compile(r"^\|\s*(\d{4}-\d{2}-\d{2})\s*\|\s*`([^`]+)`\s*\|")
def _section(markdown: str, heading: str) -> str:
level: Final = heading.split(" ", 1)[0]
start: Final = markdown.find(f"\n{heading}\n")
if start < 0:
return ""
body: Final = markdown[start + 1 + len(heading) :]
next_heading: Final = re.search(rf"^{re.escape(level)} ", body, flags=re.MULTILINE)
return body[: next_heading.start()] if next_heading else body
def parse_deprecations(markdown: str) -> DeprecationDoc:
redirect_rows: Final = tuple(
m.groups()
for m in (_REDIRECT_ROW.match(line) for line in _section(markdown, "## Active model redirects").splitlines())
if m
)
inference: Final = _section(_section(markdown, "## Deprecation history"), "### Inference")
removal_rows: Final = tuple(m.groups() for m in (_REMOVAL_ROW.match(line) for line in inference.splitlines()) if m)
if not redirect_rows or not removal_rows:
raise SyncError(
"deprecations doc parsed to zero redirect or removal rows; the table format at "
f"{DEPRECATIONS_URL} changed and the parser needs updating"
)
removal_dates: Final = {model: date for date, model in reversed(removal_rows)}
return DeprecationDoc(
removal_dates=MappingProxyType(dict(reversed(removal_dates.items()))),
redirects=MappingProxyType({original: target for original, target in redirect_rows}),
)
def per_token(price_per_million: float) -> float:
return float(f"{price_per_million / 1e6:.6g}")
def _resolve_name(name: str, universe: frozenset[str]) -> str | None:
if name in universe:
return name
suffix_matches: Final = tuple(candidate for candidate in universe if candidate.endswith(f"/{name}"))
return suffix_matches[0] if len(suffix_matches) == 1 else None
def resolve_successor(model_id: str, doc: DeprecationDoc, live_ids: frozenset[str]) -> str | None:
canonical: Final = live_ids | frozenset(doc.removal_dates)
redirects: Final = {
(_resolve_name(raw_source, canonical) or raw_source): (_resolve_name(raw_target, canonical) or raw_target)
for raw_source, raw_target in doc.redirects.items()
}
seen: Final = set()
current = model_id # rebind-ok: walks the redirect chain
while current in redirects and current not in seen:
seen.add(current)
current = redirects[current] # rebind-ok: walks the redirect chain
return current if current != model_id and current in live_ids else None
@dataclass(frozen=True, slots=True)
class SyncOutcome:
cost_map: CostMap
added: tuple[str, ...] = ()
updated: tuple[str, ...] = ()
deprecated: tuple[str, ...] = ()
reappeared: tuple[str, ...] = ()
warnings: tuple[str, ...] = ()
skipped_types: Mapping[str, int] = field(default_factory=dict)
@property
def has_changes(self) -> bool:
return bool(self.added or self.updated or self.deprecated or self.reappeared)
def _api_fields(model: CatalogModel) -> RegistryEntry:
cached: Final = model.pricing.cached_input
return {
"input_cost_per_token": per_token(model.pricing.input),
"output_cost_per_token": per_token(model.pricing.output),
**({"cache_read_input_token_cost": per_token(cached), "supports_prompt_caching": True} if cached else {}),
**({"max_input_tokens": model.context_length} if model.context_length is not None else {}),
}
def _new_entry(model: CatalogModel, mode: str) -> RegistryEntry:
rule: Final = RULES_BY_ID.get(model.id)
legacy_ceiling: Final = {} if model.context_length is None else {"max_tokens": model.context_length}
merged: Final = {
**_api_fields(model),
**legacy_ceiling,
"litellm_provider": PROVIDER,
"mode": mode,
"source": SOURCE_URL,
**(dict(rule.fields) if rule else {}),
}
return dict(sorted(merged.items()))
def _updated_entry(entry: RegistryEntry, model: CatalogModel) -> tuple[RegistryEntry, tuple[str, ...]]:
rule: Final = RULES_BY_ID.get(model.id)
desired: Final = {**_api_fields(model), **(dict(rule.fields) if rule else {})}
dropped: Final = () if model.pricing.cached_input else ("cache_read_input_token_cost", "supports_prompt_caching")
changes: Final = tuple(
f"{name}: {entry.get(name)!r} -> {value!r}" for name, value in desired.items() if entry.get(name) != value
) + tuple(
f"{name}: {entry[name]!r} removed (no longer in the catalog pricing)" for name in dropped if name in entry
)
merged: Final = {name: value for name, value in {**entry, **desired}.items() if name not in dropped}
return dict(sorted(merged.items())), changes
def _with_new_keys_in_block(original: CostMap, result: CostMap, new_keys: Sequence[str]) -> CostMap:
provider_keys: Final = tuple(key for key in original if key.startswith(PREFIX))
if not new_keys or not provider_keys:
return result
block_end: Final = provider_keys[-1]
return {
key: value
for existing in original
for key, value in (
(existing, result[existing]),
*((new, result[new]) for new in sorted(new_keys) if existing == block_end),
)
}
def compute_sync(cost_map: CostMap, catalog: Sequence[CatalogModel], doc: DeprecationDoc) -> SyncOutcome:
live_ids: Final = frozenset(model.id for model in catalog)
token_models: Final = {model.id: model for model in catalog if model.type in TYPE_TO_MODE}
skipped: Final = {
model.type: sum(1 for m in catalog if m.type == model.type)
for model in catalog
if model.type not in TYPE_TO_MODE
}
registry_ids: Final = {key.removeprefix(PREFIX): key for key in cost_map if key.startswith(PREFIX)}
added: Final[list[str]] = []
updated: Final[list[str]] = []
deprecated: Final[list[str]] = []
reappeared: Final[list[str]] = []
warnings: Final[list[str]] = []
result: Final[CostMap] = dict(cost_map)
for model_id, model in sorted(token_models.items()):
mode: Final = TYPE_TO_MODE[model.type]
key: Final = f"{PREFIX}{model_id}"
if model_id in doc.removal_dates:
warnings.append(
f"`{key}` is listed as removed on {doc.removal_dates[model_id]} in the docs but the serverless "
"catalog still serves it; availability kept from the API"
)
entry = result.get(key)
if not isinstance(entry, dict):
result[key] = _new_entry(model, mode)
added.append(key)
if model.type == "chat" and model_id not in RULES_BY_ID:
warnings.append(
f"`{key}` added without a capability rule; review its tools/vision/reasoning support and add one"
)
continue
if entry.get("mode") != mode:
warnings.append(
f"`{key}` has curated mode {entry.get('mode')!r} but the catalog maps to {mode!r}; left unchanged"
)
new_entry, changes = _updated_entry(entry, model)
if "deprecation_date" in new_entry:
new_entry.pop("deprecation_date")
reappeared.append(key)
if changes:
updated.append(f"{key}: " + "; ".join(changes))
if changes or key in reappeared:
result[key] = new_entry
for model_id, key in sorted(registry_ids.items()):
if model_id in token_models:
continue
entry = result.get(key)
if not isinstance(entry, dict):
continue
removal_date: Final = doc.removal_dates.get(model_id)
successor: Final = resolve_successor(model_id, doc, live_ids)
metadata = entry.get("metadata")
curated_successor: Final = metadata.get("successor") if isinstance(metadata, dict) else None
new_entry = dict(entry)
if removal_date is not None and entry.get("deprecation_date") != removal_date:
if "deprecation_date" in entry:
warnings.append(
f"`{key}` has curated deprecation_date {entry.get('deprecation_date')!r} but the docs list "
f"{removal_date!r}; left unchanged"
)
else:
new_entry["deprecation_date"] = removal_date
if removal_date is None and "deprecation_date" not in entry:
warnings.append(
f"`{key}` is absent from the serverless catalog with no removal date in the docs; "
"needs a human deprecation call"
)
if successor is not None:
desired_successor: Final = f"{PREFIX}{successor}"
if curated_successor is None:
new_entry["metadata"] = dict(
sorted({**(metadata if isinstance(metadata, dict) else {}), "successor": desired_successor}.items())
)
elif curated_successor != desired_successor:
warnings.append(
f"`{key}` has curated successor {curated_successor!r} but the docs redirects resolve to "
f"{desired_successor!r}; left unchanged"
)
if new_entry != entry:
result[key] = dict(sorted(new_entry.items()))
deprecated.append(f"{key}: " + ", ".join(sorted(set(new_entry) - set(entry)) or ["updated"]))
return SyncOutcome(
cost_map=_with_new_keys_in_block(cost_map, result, tuple(added)),
added=tuple(added),
updated=tuple(updated),
deprecated=tuple(deprecated),
reappeared=tuple(reappeared),
warnings=tuple(warnings),
skipped_types=MappingProxyType(skipped),
)
def _section_block(title: str, lines: Sequence[str], backtick: bool) -> str:
bullets: Final = "\n".join(f"- `{line}`" if backtick else f"- {line}" for line in lines) or "- none"
return f"### {title} ({len(lines)})\n{bullets}\n"
def render_pr_body(outcome: SyncOutcome) -> str:
skipped: Final = ", ".join(f"{kind} ({count})" for kind, count in sorted(outcome.skipped_types.items())) or "none"
return (
"Automated daily sync of the together_ai entries in model_prices_and_context_window.json against "
f"`GET {MODELS_URL}` and {DEPRECATIONS_URL} by scripts/sync_together_ai_models.py.\n"
"\n"
f"{_section_block('Added', outcome.added, backtick=True)}"
"\n"
f"{_section_block('Updated', outcome.updated, backtick=True)}"
"\n"
f"{_section_block('Marked deprecated', outcome.deprecated, backtick=True)}"
"\n"
f"{_section_block('Returned to the catalog', outcome.reappeared, backtick=True)}"
"\n"
f"{_section_block('Warnings needing a human call', outcome.warnings, backtick=False)}"
"\n"
f"Catalog model types outside the sync's token-pricing scope, skipped: {skipped}\n"
)
def render_summary(outcome: SyncOutcome) -> str:
return (
f"added={len(outcome.added)} updated={len(outcome.updated)} deprecated={len(outcome.deprecated)} "
f"reappeared={len(outcome.reappeared)} warnings={len(outcome.warnings)}"
)
def load_catalog(raw: bytes) -> list[CatalogModel]:
parsed: Final = json.loads(raw)
entries: Final = parsed.get("data") if isinstance(parsed, dict) else parsed
try:
catalog: Final = CATALOG_ADAPTER.validate_python(entries)
except ValidationError as error:
raise SyncError(f"the catalog response no longer matches the expected shape: {error}") from error
if not any(model.type in TYPE_TO_MODE for model in catalog):
raise SyncError(
"the catalog response contains no token-priced models; refusing to mark the whole registry deprecated"
)
return catalog
def _fetch(url: str, headers: Mapping[str, str]) -> bytes:
response: Final = httpx.get(url, headers=dict(headers), timeout=30, follow_redirects=True)
if response.status_code != 200:
raise SyncError(f"GET {url} returned {response.status_code}")
return response.content
def _serialize(cost_map: CostMap) -> str:
return json.dumps(cost_map, indent=4, ensure_ascii=False) + "\n"
def main(argv: Sequence[str]) -> int:
parser: Final = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--write", action="store_true", help="apply the sync to the cost map files (default: dry run)")
parser.add_argument("--models-json", type=Path, help="recorded catalog response to use instead of the live API")
parser.add_argument(
"--deprecations-md", type=Path, help="recorded deprecations doc to use instead of the live docs"
)
parser.add_argument("--pr-body-file", type=Path, help="write the generated PR body to this path")
parser.add_argument("--repo-root", type=Path, default=Path(__file__).resolve().parent.parent)
args: Final = parser.parse_args(argv)
if args.models_json is not None:
catalog_raw: Final = args.models_json.read_bytes()
else:
api_key: Final = os.environ.get("TOGETHER_API_KEY")
if not api_key:
raise SyncError("TOGETHER_API_KEY is not set and --models-json was not given")
catalog_raw = _fetch(MODELS_URL, {"Authorization": f"Bearer {api_key}"}) # rebind-ok: branch-dependent source
catalog: Final = load_catalog(catalog_raw)
markdown: Final = (
args.deprecations_md.read_text() if args.deprecations_md is not None else _fetch(DEPRECATIONS_URL, {}).decode()
)
doc: Final = parse_deprecations(markdown)
cost_map_path: Final = args.repo_root / COST_MAP_RELPATHS[0]
cost_map: Final = json.loads(cost_map_path.read_text())
outcome: Final = compute_sync(cost_map, catalog, doc)
body: Final = render_pr_body(outcome)
if args.pr_body_file is not None:
args.pr_body_file.write_text(body)
if args.write and outcome.has_changes:
for relpath in COST_MAP_RELPATHS:
(args.repo_root / relpath).write_text(_serialize(outcome.cost_map))
print(render_summary(outcome))
print()
print(body)
if not args.write:
print("dry run: no files were touched")
elif not outcome.has_changes:
print("registry already in sync: no files were touched")
return 0
if __name__ == "__main__":
try:
raise SystemExit(main(sys.argv[1:]))
except SyncError as error:
print(f"SYNC FAILED: {error}", file=sys.stderr)
raise SystemExit(1) from error

24
tests/base_sdk_tests/check_base_sdk_install.py
View file

@ -102,6 +102,29 @@ def check_completion() -> str:
return "mock completion round-trips"
def check_completion_with_function_tool() -> str:
import litellm
response: Final = litellm.completion(
model="gpt-4o",
messages=[{"role": "user", "content": "ping"}],
tools=[
{
"type": "function",
"function": {
"name": "noop",
"description": "No operation",
"parameters": {"type": "object", "properties": {}},
},
}
],
mock_response="pong",
)
content: Final = response.choices[0].message.content
_require(content == "pong", f"mock completion with a function tool returned {content!r}")
return "mock completion with a function tool round-trips"
def check_mcp_install_guidance() -> str:
try:
import litellm.experimental_mcp_client
@ -209,6 +232,7 @@ CHECKS: tuple[tuple[str, Callable[[], str]], ...] = (
("import litellm", check_import),
("optional MCP installation guidance", check_mcp_install_guidance),
("chat completion", check_completion),
("chat completion with a function tool", check_completion_with_function_tool),
("embedding", check_embedding),
("bundled model metadata", check_bundled_model_metadata),
("token counter", check_token_counter),

55
tests/base_sdk_tests/check_sdk_http.py
View file

@ -89,7 +89,6 @@ def payload(response: object) -> Mapping[str, JsonValue]:
return JSON_OBJECT.validate_json(response.model_dump_json(exclude_none=True))
@skipIf(importlib.util.find_spec("mcp") is None, "BUG: LIT-9333 completion(tools=...) imports optional mcp")
def check_tools(base: str, asynchronous: bool = False) -> None:
tools: Final = [
{
@ -122,7 +121,6 @@ def check_tools(base: str, asynchronous: bool = False) -> None:
assert json.loads(body) == {"model": "gpt-6", "messages": MESSAGES, "tools": tools}
@skipIf(importlib.util.find_spec("mcp") is None, "BUG: LIT-9333 responses() imports optional mcp")
def check_responses(base: str, asynchronous: bool = False) -> None:
output: Final[Sequence[JsonValue]] = [
{
@ -144,14 +142,62 @@ def check_responses(base: str, asynchronous: bool = False) -> None:
"usage": {"input_tokens": 3, "output_tokens": 2, "total_tokens": 5},
}
)
arguments: Final = dict(model="openai/gpt-6", input="ping", api_key="test-key", api_base=base + "/v1")
tools: Final = [
{"type": "mcp", "server_label": "docs", "server_url": "https://mcp.example.com/sse"},
{"type": "mcp", "server_label": "repos", "server_url": "https://mcp.example.com/mcp/repos"},
]
arguments: Final = dict(model="openai/gpt-6", input="ping", tools=tools, api_key="test-key", api_base=base + "/v1")
response: Final = asyncio.run(litellm.aresponses(**arguments)) if asynchronous else litellm.responses(**arguments)
assert payload(response)["output"] == output
assert payload(response)["usage"] == {"input_tokens": 3, "output_tokens": 2, "total_tokens": 5}
path, headers, body = REQUESTS.get(timeout=10)
assert path == "/v1/responses"
assert headers["authorization"] == "Bearer test-key"
assert json.loads(body) == {"model": "gpt-6", "input": "ping"}
assert json.loads(body) == {"model": "gpt-6", "input": "ping", "tools": tools}
def check_messages_tools(base: str) -> None:
tools: Final = [
{
"name": "lookup",
"description": "Look up the forecast for a city",
"input_schema": {"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]},
}
]
content: Final = [{"type": "tool_use", "id": "toolu-http-check", "name": "lookup", "input": {"city": "Paris"}}]
enqueue(
{
"id": "msg-tools-http-check",
"type": "message",
"role": "assistant",
"model": "claude-sonnet-5-5",
"content": content,
"stop_reason": "tool_use",
"stop_sequence": None,
"usage": {"input_tokens": 3, "output_tokens": 2},
}
)
async def send() -> object:
create: Final = cast( # cast-ok: SDK facade exposes untyped kwargs
Callable[..., Awaitable[object]], litellm.anthropic.messages.acreate
)
return await create(
model="anthropic/claude-sonnet-5-5",
messages=MESSAGES,
tools=tools,
max_tokens=16,
api_key="test-key",
api_base=base,
)
response: Final = JSON_OBJECT.validate_python(asyncio.run(send()))
assert response["content"] == content, response
path, headers, body = REQUESTS.get(timeout=10)
assert path == "/v1/messages"
assert headers["x-api-key"] == "test-key"
sent: Final = JSON_OBJECT.validate_json(body)
assert (sent["model"], sent["messages"], sent["tools"]) == ("claude-sonnet-5-5", MESSAGES, tools), sent
@skipIf(importlib.util.find_spec("fastapi") is None, "BUG: LIT-9333 cache hits import proxy-only fastapi")
@ -392,6 +438,7 @@ def main() -> None:
partial(check_tools, asynchronous=True),
check_responses,
partial(check_responses, asynchronous=True),
check_messages_tools,
check_cache_hit,
check_messages_stream_and_callback,
):

1
tests/code_coverage_tests/check_e2e_no_raw_requests.py
View file

@ -25,7 +25,6 @@ ALLOWED_RAW_CLIENT_FILES = {
"e2e/e2e_http.py": ("requests",),
"e2e/conftest.py": ("requests",),
"e2e/claude_code/pr_gate_version_resolver.py": ("urllib.request",),
"e2e/mcp/oauth_chat_client.py": ("httpx",),
}
EXCEPTION_ONLY_NAMES = frozenset({"RequestException", "ConnectionError", "Timeout", "HTTPError"})

2
tests/code_coverage_tests/router_code_coverage.py
View file

@ -99,6 +99,8 @@ ignored_function_names = [
"_embedding",
"_aembedding",
"_stream_pre_content_error", # Tested through the non-retriable retry error tests in test_router.py
"_narrow_to_video_deployment", # Tested through video_status and avideo_status polls in test_router_provider_endpoints.py
"_video_call_pin", # Tested through the video_status pin tests and the chat video_id test in test_router_provider_endpoints.py
"_deployment_num_retries", # Tested through the deployment num_retries mid-stream budget test in test_router.py
"_request_fallback_list", # Tested through every mid-stream retry test in test_router.py
"_request_model_group", # Tested through test_anthropic_messages_retry_budget_precedence_direct_call

156
tests/code_coverage_tests/test_e2e_changed_gate.py
View file

@ -1,156 +0,0 @@
import os
import subprocess
import sys
import xml.etree.ElementTree as ET
from pathlib import Path
from typing import Final
import pytest
GATE: Final = Path(__file__).resolve().parents[2] / ".github/e2e-stack/assert_tests_ran.py"
SELECTED: Final = ("tests/e2e/access_control/test_a.py", "tests/e2e/access_control/test_b.py")
@pytest.mark.parametrize(
("second_outcome", "expected_status"),
(("passed", 0), ("skipped", 1), ("failure", 1), ("error", 1), ("deselected", 1)),
)
def test_each_changed_file_must_run(tmp_path: Path, second_outcome: str, expected_status: int) -> None:
suite: Final = ET.Element("testsuite")
_ = ET.SubElement(suite, "testcase", file=SELECTED[0])
if second_outcome != "deselected":
second: Final = ET.SubElement(suite, "testcase", file=SELECTED[1])
if second_outcome != "passed":
_ = ET.SubElement(second, second_outcome)
report: Final = tmp_path / "report.xml"
ET.ElementTree(suite).write(report)
result: Final = subprocess.run([sys.executable, str(GATE), str(report), *SELECTED], capture_output=True, text=True)
assert result.returncode == expected_status, result.stdout
@pytest.mark.parametrize("outcome", ("failure", "error"))
def test_passing_case_does_not_hide_a_failure_in_the_same_file(tmp_path: Path, outcome: str) -> None:
suite: Final = ET.Element("testsuite")
_ = ET.SubElement(suite, "testcase", file=SELECTED[0])
failed: Final = ET.SubElement(suite, "testcase", file=SELECTED[0])
_ = ET.SubElement(failed, outcome)
report: Final = tmp_path / "report.xml"
ET.ElementTree(suite).write(report)
result: Final = subprocess.run(
[sys.executable, str(GATE), str(report), SELECTED[0]], capture_output=True, text=True
)
assert result.returncode == 1
def test_failed_cases_are_named_per_selected_file(tmp_path: Path) -> None:
suite: Final = ET.Element("testsuite")
_ = ET.SubElement(suite, "testcase", file=SELECTED[0], classname="tests.e2e.access_control.test_a", name="test_ok")
failed: Final = ET.SubElement(
suite, "testcase", file=SELECTED[0], classname="tests.e2e.access_control.test_a", name="test_boom"
)
_ = ET.SubElement(failed, "failure", message="secret-bearing message")
errored: Final = ET.SubElement(
suite, "testcase", file=SELECTED[1], classname="tests.e2e.access_control.test_b", name="test_setup"
)
_ = ET.SubElement(errored, "error")
report: Final = tmp_path / "report.xml"
ET.ElementTree(suite).write(report)
result: Final = subprocess.run([sys.executable, str(GATE), str(report), *SELECTED], capture_output=True, text=True)
assert result.returncode == 1
assert " failed: tests.e2e.access_control.test_a::test_boom\n" in result.stdout
assert " failed: tests.e2e.access_control.test_b::test_setup\n" in result.stdout
assert "test_ok" not in result.stdout
assert "secret-bearing message" not in result.stdout
@pytest.mark.parametrize("contents", ("<testsuite/>", "<testsuite", '<testsuite><testcase name="a"/></testsuite>'))
def test_missing_execution_evidence_fails(tmp_path: Path, contents: str) -> None:
report: Final = tmp_path / "report.xml"
_ = report.write_text(contents)
result: Final = subprocess.run([sys.executable, str(GATE), str(report), *SELECTED], capture_output=True, text=True)
assert result.returncode == 1
@pytest.mark.parametrize("omitted_role", ("proxy_admin", "team_member", "internal_user_viewer"))
def test_one_passing_management_case_cannot_hide_a_missing_actor(tmp_path: Path, omitted_role: str) -> None:
suite: Final = ET.Element("testsuite")
path: Final = "tests/e2e/management/test_jwt_management_e2e.py"
case: Final = ET.SubElement(suite, "testcase", file=path)
properties: Final = ET.SubElement(case, "properties")
_ = ET.SubElement(
properties,
"property",
name="management_node",
value=f"{path}::TestJwtManagement::test_actor_subject_and_database_role[proxy_admin_viewer]",
)
report: Final = tmp_path / "report.xml"
ET.ElementTree(suite).write(report)
result: Final = subprocess.run([sys.executable, str(GATE), str(report), path], capture_output=True, text=True)
assert result.returncode == 1
assert f"test_actor_subject_and_database_role[{omitted_role}]" in result.stdout
@pytest.mark.parametrize("phase", ("setup", "call", "teardown"))
@pytest.mark.parametrize("required_count", ("1", "4"))
def test_oauth_failure_diagnostics_do_not_publish_private_payloads(
tmp_path: Path, phase: str, required_count: str
) -> None:
suite: Final = ET.Element("testsuite")
case: Final = ET.SubElement(suite, "testcase", file=SELECTED[0])
private: Final = "private-token-in-exception-message"
failure: Final = ET.SubElement(case, "failure", message=private)
failure.text = private
properties: Final = ET.SubElement(case, "properties")
for name, value in (
("oauth_failure_phase", phase),
("oauth_exception_type", "AssertionError"),
("oauth_frame", "oauth_gateway.py:120:start"),
("oauth_frame", f"injected\\n{private}"),
("unrelated_property", private),
):
_ = ET.SubElement(properties, "property", name=name, value=value)
report: Final = tmp_path / "report.xml"
ET.ElementTree(suite).write(report)
result: Final = subprocess.run(
[sys.executable, "-I", str(GATE), str(report), SELECTED[0]],
capture_output=True,
text=True,
env={**os.environ, "E2E_REQUIRED_TEST_COUNT": required_count},
)
assert result.returncode == 1
assert f"oauth_failure_phase: {phase}" in result.stdout
assert "oauth_exception_type: AssertionError" in result.stdout
assert "oauth_frame: oauth_gateway.py:120:start" in result.stdout
assert private not in result.stdout + result.stderr
@pytest.mark.parametrize(
("count", "skip", "expected"), ((0, False, 1), (3, False, 1), (4, False, 0), (5, False, 1), (4, True, 1))
)
def test_required_count_reports_cases_before_rejecting(tmp_path: Path, count: int, skip: bool, expected: int) -> None:
suite = ET.Element("testsuite")
for index in range(count):
case = ET.SubElement(suite, "testcase", file=SELECTED[0], classname="OAuth", name=f"variant{index}")
if skip and index == 0:
ET.SubElement(case, "skipped", message="private-skip-reason")
report = tmp_path / "report.xml"
ET.ElementTree(suite).write(report)
result = subprocess.run(
[sys.executable, "-I", str(GATE), str(report), SELECTED[0]],
env={**os.environ, "E2E_REQUIRED_TEST_COUNT": "4"},
capture_output=True,
text=True,
)
assert result.returncode == expected
assert f"{count} collected, {int(skip)} skipped" in result.stdout
if skip:
assert "skipped: OAuth::variant0" in result.stdout
assert "private-skip-reason" not in result.stdout + result.stderr

16
tests/code_coverage_tests/test_e2e_junit_report.py
View file

@ -172,12 +172,7 @@ def open_consent() -> None:
raise RuntimeError("consent page timed out")
@pytest.mark.mcp_oauth_live
def test_oauth_dies_on_consent() -> None:
open_consent()
def test_plain_dies_on_consent() -> None:
def test_dies_on_consent() -> None:
open_consent()
"""
@ -353,16 +348,13 @@ class TestStepsReachTheReport:
def test_a_failed_phase_s_own_report_carries_the_steps(tmp_path: Path) -> None:
"""Plugins that read the failed setup or call report, not the teardown one
junitxml writes from, see where the test died too, oauth-live or not."""
junitxml writes from, see where the test died too."""
write_suite(tmp_path, {"test_consent.py": FAILED_PHASE_SUITE, "report_spy.py": REPORT_SPY_PLUGIN})
child: Final = run_child_pytest(tmp_path, "-p", "report_spy", env={"E2E_MCP_OAUTH_LIVE": "1"})
child: Final = run_child_pytest(tmp_path, "-p", "report_spy")
seen_path: Final = tmp_path / "failed-reports.jsonl"
assert seen_path.exists(), f"no failed report reached the spy:\n{child.stdout}\n{child.stderr}"
seen: Final = dict(map(FailedReport.validate_json, seen_path.read_text().splitlines()))
assert seen == {
"test_oauth_dies_on_consent": ("open the consent page",),
"test_plain_dies_on_consent": ("open the consent page",),
}, child.stdout
assert seen == {"test_dies_on_consent": ("open the consent page",)}, child.stdout
class TestDeclaredPropertiesReachTheReport:

14
tests/e2e/AGENTS.md
View file

File diff suppressed because one or more lines are too long

52
tests/e2e/CONTRIBUTING.md
View file

@ -250,55 +250,3 @@ 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 `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
`test_mcp_oauth_happy_path_e2e.py` runs one shared scenario with four variants:
aggregate gateway SSO and explicitly configured per-server JWT, each directly
against Linear and through the live provider edge. The edge forwards to real
Linear without replay and compares the forwarded bearer to the encrypted
canonical user/server credential. This observes the forwarding boundary, not
Linear's internal logs. Direct variants independently exercise discovery
Use the existing database preparation, Prisma generation and Keycloak setup.
Build and stage the dashboard from the tested checkout as in the UI runner.
Provide `DATABASE_URL`, `LITELLM_MASTER_KEY`, `LITELLM_SALT_KEY`, `LITELLM_LICENSE`,
and the `E2E_KEYCLOAK_*` settings. Capture a test-account Linear login using
`mcp/linear_session_capture.py` and set `E2E_LINEAR_STORAGE_STATE` to that private
file. The test workspace must contain a team. Do not publish browser state or
raw test/proxy output
```bash
E2E_MCP_OAUTH_LIVE=1 E2E_FIXTURE_MODE=live E2E_PROVIDER_CACHE=0 \
uv run --no-sync pytest tests/e2e/mcp/test_mcp_oauth_happy_path_e2e.py \
--rootdir=. --reruns 0
```
The test starts and restarts its own source-built proxy on a free loopback port,
retaining its database and SSO client but no Redis or process-local cache. It
does not restart an existing proxy or clear shared databases. Gateway login,
consent, immediate list/call and post-restart reconnect must all succeed. The
aggregate client never injects a gateway header; the explicitly labeled JWT
variant configures `x-litellm-api-key` for the first consent and reconnects with
only its gateway JWT after restart
`.github/workflows/test-mcp-oauth-e2e.yml` automatically requests a run for
same-repository pull requests changing MCP, gateway authentication/SSO, consent
UI, dependencies or the relevant E2E harness/workflow paths. It retains manual
`workflow_dispatch` for targeted verification. The four cases run in the
protected `e2e-changed` environment after its normal deployment approval;
reviewers should approve and inspect this separate OAuth check when it appears.
Fork pull requests do not run this credentialed job; use a reviewed
same-repository branch for their verification. The workflow's path-filtered
check is not configured here as a globally required branch-protection check.
Provision `E2E_LINEAR_STORAGE_STATE_B64` as a secret there and retain the existing E2E license/AWS role configuration. A missing or
expired session fails the job; collection, deselection and skips are not passes.
No LLM call is needed
Coverage remains limited to authorization-code OAuth over HTTP. M2M, OBO,
PKCE passthrough, static/BYOK, ID-JAG, forwarding, SigV4 and stdio are outside this
scenario; consult the registry and LIT-3559 for their existing coverage and gaps.
LIT-4506 owns broader isolation/failure regressions. LIT-7737 retains ownership
of dependency/Python compatibility and its matrix; this test reuses its delivered
environment and does not change dependency constraints or compatibility gates

35
tests/e2e/claude_code/conftest.py
View file

@ -597,14 +597,21 @@ def _build_control_plane_client(proxy_config: ProxyConfig):
)
def _register_deployment(proxy, deployment: CompatDeployment) -> str:
"""Register one deployment and return its proxy-assigned model_id
once it is servable on the data plane. The aliases are shared by every
cell and, under xdist, by every worker, so no call to them belongs to
one test and none is cached: the matrix exists to reach real providers."""
return proxy.create_model(
deployment.model_name,
deployment.litellm_params,
def _register_deployments(proxy, deployments: tuple[CompatDeployment, ...]):
"""The aliases are shared by every cell and, under xdist, by every worker,
so no call to them belongs to one test and none is cached: the matrix
exists to reach real providers."""
from models import ModelInfoBody, ModelNewBody
return proxy.register_models(
tuple(
ModelNewBody(
model_name=deployment.model_name,
litellm_params=deployment.litellm_params,
model_info=ModelInfoBody(),
)
for deployment in deployments
),
provider_live=True,
)
@ -632,18 +639,14 @@ def _compat_models_registered() -> Any:
yield
return
from proxy_client import ModelNotRegistered, ModelRegistered
from requests import RequestException
proxy = _build_control_plane_client(proxy_config)
registered_ids: list[str] = []
failures: list[tuple[str, str]] = []
registrations = _register_deployments(proxy, load_all_deployments())
registered_ids = tuple(r.model_id for r in registrations if isinstance(r, ModelRegistered))
failures = tuple((r.model_name, r.reason) for r in registrations if isinstance(r, ModelNotRegistered))
try:
for deployment in load_all_deployments():
try:
model_id = _register_deployment(proxy, deployment)
registered_ids.append(model_id)
except (AssertionError, RequestException) as exc:
failures.append((deployment.model_name, str(exc)))
if failures:
summary = "\n".join(
f" - {name}: {reason}" for name, reason in failures

37
tests/e2e/conftest.py
View file

@ -17,7 +17,6 @@ import functools
import os
from collections.abc import Generator, Iterator
from datetime import datetime, timezone
from pathlib import Path
from types import MappingProxyType
from typing import Final
@ -29,7 +28,6 @@ from e2e_config import (
FIXTURE_DIR,
FIXTURE_MODE_RAW,
MANAGED_FILES_OPT_IN_ENV,
MCP_OAUTH_LIVE_OPT_IN_ENV,
OTEL_TLS_OPT_IN_ENV,
OTEL_V2_OPT_IN_ENV,
OWNED_GATEWAY_OPT_IN_ENV,
@ -37,6 +35,7 @@ from e2e_config import (
PROMPT_CACHING_OPT_IN_ENV,
PROVIDER_EDGE_HOST_OPT_IN_ENV,
PROXY_BASE_URL,
PROXY_WIDE_CONFIG_OPT_IN_ENV,
SAGEMAKER_NOVA2_LITE_OPT_IN_ENV,
SAGEMAKER_NOVA_OPT_IN_ENV,
SECRET_MANAGER_OPT_IN_ENV,
@ -56,7 +55,6 @@ from models import TeamNewBody, UserNewBody, UserNewResponse
from provider_cache_routing import LIVE_PROVIDER_REQUIRED
from provider_edge import replay_leftover_error
from proxy_client import ProxyClient, build_proxy_client
from stack_lock import stack_lock
_E2E_TEST_RAN = pytest.StashKey[bool]()
_CALL_PASSED = pytest.StashKey[bool]()
@ -72,7 +70,6 @@ OPT_IN_MARKERS: Final = MappingProxyType(
"redis_chaos": PERF_TESTS_OPT_IN_ENV,
"perf": PERF_TESTS_OPT_IN_ENV,
"cli_determinism": CLI_DETERMINISM_OPT_IN_ENV,
"mcp_oauth_live": MCP_OAUTH_LIVE_OPT_IN_ENV,
"provider_edge_host": PROVIDER_EDGE_HOST_OPT_IN_ENV,
"owned_gateway": OWNED_GATEWAY_OPT_IN_ENV,
"otel_v2": OTEL_V2_OPT_IN_ENV,
@ -80,6 +77,7 @@ OPT_IN_MARKERS: Final = MappingProxyType(
"secret_manager": SECRET_MANAGER_OPT_IN_ENV,
"sagemaker_nova": SAGEMAKER_NOVA_OPT_IN_ENV,
"sagemaker_nova2_lite": SAGEMAKER_NOVA2_LITE_OPT_IN_ENV,
"proxy_wide_config": PROXY_WIDE_CONFIG_OPT_IN_ENV,
}
)
@ -165,16 +163,6 @@ def pytest_configure(config: pytest.Config) -> None:
"redis_chaos: load test that pauses the proxy's Redis outright mid-run; needs a proxy booted from "
"gateway/redis_chaos_ci_config.yml on the same host, and is deselected unless E2E_PERF_TESTS is set",
)
config.addinivalue_line(
"markers",
"quiet_stack: measures the proxy itself, so it runs while no other test on this host is hitting the stack; "
"every other test waits for it to finish",
)
config.addinivalue_line(
"markers",
"mcp_oauth_live: real Linear OAuth consent via a captured browser session; deselected unless "
"E2E_MCP_OAUTH_LIVE is set",
)
config.addinivalue_line(
"markers",
"provider_edge_host: routes provider traffic through the pytest host's edge in every fixture mode, so the "
@ -208,6 +196,11 @@ def pytest_configure(config: pytest.Config) -> None:
"sagemaker_nova2_lite: needs a deployed SageMaker Nova 2 Lite endpoint the proxy's AWS credentials can invoke; "
"deselected unless SAGEMAKER_NOVA2_LITE_ENDPOINT names it",
)
config.addinivalue_line(
"markers",
"proxy_wide_config: turns on a proxy-wide setting through /config/update, so it runs on a stack no other suite "
"is driving; deselected unless E2E_PROXY_WIDE_CONFIG_STACK is set",
)
def pytest_sessionstart(session: pytest.Session) -> None:
@ -277,7 +270,7 @@ def pytest_collection_finish(session: pytest.Session) -> None:
ahead of xdist's own collection-finish report, and the controller schedules no
test until every worker has reported, so this is the idle footprint of a stack
that just passed its readiness gate. The fixture hands the capture to the
idle-budget test in router/test_reliability_memory_e2e.py."""
idle-budget test in perf/test_reliability_memory_e2e.py."""
if session.config.getoption("collectonly") or not any(_uses_idle_rss(item) for item in session.items):
return
session.config.stash[_IDLE_RSS] = read_rss_everywhere(build_proxy_client(), timeout=IDLE_RSS_READ_TIMEOUT_SECONDS)
@ -307,12 +300,6 @@ def _proxy_fail_reason() -> str | None:
return None
@pytest.hookimpl(wrapper=True)
def pytest_runtest_protocol(item: pytest.Item, nextitem: pytest.Item | None) -> Generator[None, object, object]:
with stack_lock(exclusive=item.get_closest_marker("quiet_stack") is not None):
return (yield)
@pytest.hookimpl(tryfirst=True)
def pytest_runtest_setup(item: pytest.Item) -> None:
"""Hard-fail `e2e`-marked tests unless a proxy answers its liveness probe.
@ -331,8 +318,6 @@ def pytest_runtest_setup(item: pytest.Item) -> None:
item.user_properties.extend(item.config.stash[_IDLE_RSS].junit_properties)
if not _reaches_proxy(item):
return
if isinstance(item, pytest.Function) and "oauth_gateway" in item.fixturenames:
return
reason = _proxy_fail_reason()
if reason is not None:
pytest.fail(reason)
@ -376,12 +361,6 @@ def pytest_runtest_makereport(
report = yield
if report.when in ("setup", "call"):
attach_step_properties(item)
if item.get_closest_marker("mcp_oauth_live") is not None and call.excinfo is not None:
# Publish code locations only, never exception messages, source text or locals.
item.user_properties.append(("oauth_failure_phase", report.when))
item.user_properties.append(("oauth_exception_type", call.excinfo.type.__name__))
for entry in call.excinfo.traceback:
item.user_properties.append(("oauth_frame", f"{Path(entry.path).name}:{entry.lineno + 1}:{entry.name}"))
if report.when == "call":
item.stash[_CALL_PASSED] = report.passed
report.user_properties = list(item.user_properties)

8
tests/e2e/coverage_registry/mcp.yaml
View file

@ -79,14 +79,6 @@
assertions: [succeeds]
source: "db.py user_oauth_credential lookup"
rationale: OAuth2 token passthrough; per-user credential storage
- id: mcp.call_tool.oauth.persists_across_processes
module: mcp
tier: P1
operation: call_tool
auth_family: oauth
assertions: [persists_across_processes]
source: "outbound_credentials/per_user_oauth_store.py V2PerUserTokenStore"
rationale: Stored per-user token survives a verified restart of an owned gateway with no Redis cache
- id: mcp.list_tools.none.succeeds
module: mcp
tier: P1

2
tests/e2e/coverage_registry/quota_management.yaml
View file

@ -44,8 +44,6 @@
- {id: quota_management.spend_tracking.messages_bridge.logs_cost, module: quota_management, tier: P1, behavior: spend_tracking, variant: messages_bridge, assertions: [logs_cost], exercised_on: [messages], source: "llms/anthropic/pass_through/responses_adapters/handler.py", rationale: "A streaming /v1/messages request served by an openai-provider model is bridged through the anthropic-messages -> Responses adapter and must aggregate the consumed SSE stream into one spend row with nonzero cost and token counts, attributed to custom_llm_provider openai under call_type anthropic_messages"}
- {id: quota_management.spend_tracking.embeddings.logs_cost, module: quota_management, tier: P1, behavior: spend_tracking, variant: embeddings, assertions: [logs_cost], exercised_on: [embeddings], source: "proxy/spend_tracking/spend_tracking_utils.py", rationale: "Embedding calls write nonzero spend rows"}
- {id: quota_management.spend_tracking.cache_hit.zero_cost, module: quota_management, tier: P1, behavior: spend_tracking, variant: cache_hit, assertions: [zero_cost], exercised_on: [chat_completions], source: "proxy/spend_tracking/spend_tracking_utils.py", rationale: "A response-cache hit logs at zero cost with the cache-hit marker"}
- {id: quota_management.spend_tracking.discount_config.logs_cost, module: quota_management, tier: P1, behavior: spend_tracking, variant: discount_config, assertions: [logs_cost], exercised_on: [chat_completions], source: "proxy/management_endpoints/cost_tracking_settings.py + cost_calculator.py", rationale: "A persisted provider discount changes the real request's spend row"}
- {id: quota_management.spend_tracking.margin_config.logs_cost, module: quota_management, tier: P1, behavior: spend_tracking, variant: margin_config, assertions: [logs_cost], exercised_on: [chat_completions], source: "proxy/management_endpoints/cost_tracking_settings.py + cost_calculator.py", rationale: "A persisted provider margin is added to the real request's spend row"}
- {id: quota_management.spend_tracking.guardrail_cost.logs_cost, module: quota_management, tier: P1, behavior: spend_tracking, variant: guardrail_cost, assertions: [logs_cost], exercised_on: [chat_completions], source: "litellm_core_utils/llm_cost_calc/guardrail_cost.py + proxy/spend_tracking/spend_tracking_utils.py", rationale: "A priced Bedrock guardrail charge reaches the persisted spend row"}
- {id: quota_management.spend_tracking.key_rollup.matches_sum_of_logs, module: quota_management, tier: P1, behavior: spend_tracking, variant: key_rollup, assertions: [matches_sum_of_logs], exercised_on: [chat_completions], source: "proxy/db/db_spend_update_writer.py", rationale: "A key's rolled-up spend equals the sum of its log rows"}
- {id: quota_management.spend_tracking.concurrent_burst.loses_no_spend, module: quota_management, tier: P1, behavior: spend_tracking, variant: concurrent_burst, assertions: [loses_no_spend], exercised_on: [chat_completions], source: "proxy/db/db_spend_update_writer.py", rationale: "Concurrent calls all land as spend; no row lost to write contention"}

7
tests/e2e/e2e_config.py
View file

@ -111,10 +111,6 @@ CHEAP_ANTHROPIC_MODEL = os.environ.get("E2E_CHEAP_ANTHROPIC_MODEL", "claude-haik
CHEAP_OPENAI_MODEL = os.environ.get("E2E_CHEAP_OPENAI_MODEL", "gpt-5.5")
S3_PARTITION_GRANULARITY = os.environ.get("E2E_S3_PARTITION_GRANULARITY", "day")
LINEAR_MCP_URL = os.environ.get("E2E_LINEAR_MCP_URL", "https://mcp.linear.app/mcp")
LINEAR_STORAGE_STATE = os.environ.get("E2E_LINEAR_STORAGE_STATE", "")
LINEAR_READONLY_TOOL: Final = "list_teams" # as listed by tools/list on mcp.linear.app when PR #33787 landed
# Jaeger query API of the compose stack's OTEL trace destination (the `jaeger`
# service in docker-compose.yml maps it to host 16686). Trace-completeness tests
# read exported spans back through it.
@ -207,7 +203,6 @@ MANAGED_FILES_OPT_IN_ENV = "E2E_MANAGED_FILES_STACK"
PROMPT_CACHING_OPT_IN_ENV = "E2E_PROMPT_CACHING_STACK"
PERF_TESTS_OPT_IN_ENV: Final = "E2E_PERF_TESTS"
CLI_DETERMINISM_OPT_IN_ENV = "E2E_CLI_DETERMINISM"
MCP_OAUTH_LIVE_OPT_IN_ENV: Final = "E2E_MCP_OAUTH_LIVE"
PROVIDER_EDGE_HOST_OPT_IN_ENV: Final = "E2E_PROVIDER_EDGE_HOST_REACHABLE"
OWNED_GATEWAY_OPT_IN_ENV: Final = "E2E_OWNED_GATEWAY"
OTEL_V2_OPT_IN_ENV: Final = "E2E_OTEL_V2"
@ -215,6 +210,7 @@ OTEL_TLS_OPT_IN_ENV: Final = "E2E_OTEL_EXPORTER_ENDPOINT"
SECRET_MANAGER_OPT_IN_ENV: Final = "E2E_SECRET_MANAGER"
SAGEMAKER_NOVA_OPT_IN_ENV: Final = "SAGEMAKER_NOVA_ENDPOINT"
SAGEMAKER_NOVA2_LITE_OPT_IN_ENV: Final = "SAGEMAKER_NOVA2_LITE_ENDPOINT"
PROXY_WIDE_CONFIG_OPT_IN_ENV: Final = "E2E_PROXY_WIDE_CONFIG_STACK"
ANOMALY_SESSIONS = int(os.environ.get("E2E_ANOMALY_SESSIONS", "6"))
ANOMALY_TURNS_PER_SESSION = int(os.environ.get("E2E_ANOMALY_TURNS_PER_SESSION", "6"))
ANOMALY_TURN_ATTEMPTS = int(os.environ.get("E2E_ANOMALY_TURN_ATTEMPTS", "3"))
@ -239,7 +235,6 @@ MEMORY_RSS_SETTLE_SAMPLES = int(os.environ.get("E2E_MEMORY_RSS_SETTLE_SAMPLES",
MEMORY_RSS_SAMPLE_INTERVAL_SECONDS = float(os.environ.get("E2E_MEMORY_RSS_SAMPLE_INTERVAL_SECONDS", "1"))
MEMORY_RSS_CEILING_MB = float(os.environ.get("E2E_MEMORY_RSS_CEILING_MB", "200"))
MEMORY_IDLE_RSS_BUDGET_MB = float(os.environ.get("E2E_MEMORY_IDLE_RSS_BUDGET_MB", "768"))
MEMORY_STORED_REQUEST_BUDGET_KB = float(os.environ.get("E2E_MEMORY_STORED_REQUEST_BUDGET_KB", "64"))
def ws_base_url() -> str:

6
tests/e2e/llm_translation/test_together_ai_e2e.py
View file

@ -7,9 +7,9 @@ backends are pinned because the registry has no flag for what they prove: ``enab
and the ``{"reasoning": {"enabled": false}}`` toggle that ``reasoning_effort="none"`` maps
to are Qwen hybrid-model contracts, and the structured-output case runs on that hybrid
model with reasoning off, since a reasoning-only model can spend the whole token budget
thinking and return no content. MiniMax-M3 is the serverless model whose
thinking and return no content. MiniMax-M3.1-Flash is the serverless model whose
template renders a replayed ``reasoning_content`` back into the prompt (Qwen and
DeepSeek silently drop it). MiniMax-M3 honors that replayed field on nearly every call, not
DeepSeek silently drop it). MiniMax-M3.1-Flash honors that replayed field on nearly every call, not
every call (one miss in dozens of otherwise identical calls), so the replay case asks
up to ``REPLAY_ATTEMPTS`` times and fails only when no answer carries the secret, which
a proxy that strips the field guarantees. Requires TOGETHER_API_KEY on the proxy; no
@ -56,7 +56,7 @@ from pydantic import BaseModel
pytestmark = pytest.mark.e2e
HYBRID_REASONING_BACKEND = "together_ai/Qwen/Qwen3.5-9B"
REASONING_REPLAY_BACKEND = "together_ai/MiniMaxAI/MiniMax-M3"
REASONING_REPLAY_BACKEND = "together_ai/MiniMaxAI/MiniMax-M3.1-Flash"
SECRET_PROMPT = "Remember this for later and reply with just OK."
SECRET_REASONING = "The user told me their favorite color is chartreuse. I must remember it."

58
tests/e2e/mcp/linear_session_capture.py
View file

@ -1,58 +0,0 @@
"""One-time helper to capture a logged-in Linear browser session for the
real-Linear MCP e2e test.
The real-Linear test drives the genuine gateway-managed authorization_code
dance against ``mcp.linear.app``. The only step that cannot be scripted is
Linear's login (magic link / SSO), so a human authenticates once here and the
resulting session (cookies + local storage) is persisted to disk. The e2e test
then loads that session in a headless Playwright context and clicks Approve on
Linear's consent screen every run, with no human and no login automation.
Run it with the e2e venv, log into Linear in the window that opens, then return
to the terminal and press Enter:
LITELLM=~/litellm-mcpe2e
"$LITELLM"/.venv/bin/python "$LITELLM"/tests/e2e/mcp/linear_session_capture.py
The session is written to ``E2E_LINEAR_STORAGE_STATE`` (default
``~/.litellm-e2e/linear_storage_state.json``), outside the repo. It is a
secret: never commit it. Re-run this whenever Linear expires the session.
"""
from __future__ import annotations
import os
from pathlib import Path
from playwright.sync_api import sync_playwright
DEFAULT_STATE_PATH = Path.home() / ".litellm-e2e" / "linear_storage_state.json"
def capture(state_path: Path) -> None:
"""Open a headed browser at Linear, wait for the human to log in, then save
the authenticated session to ``state_path``."""
state_path.parent.mkdir(parents=True, exist_ok=True)
with sync_playwright() as playwright:
browser = playwright.chromium.launch(headless=False)
context = browser.new_context()
page = context.new_page()
page.goto("https://linear.app/login", wait_until="domcontentloaded")
print("\n" + "=" * 72)
print("Log into Linear in the browser window that just opened.")
print("If Linear emails you a magic link, paste the link into THIS window's")
print("address bar (opening it in your default browser won't capture the")
print("session). Google SSO works too as long as you complete it here.")
print("When your Linear workspace has loaded, come back and press Enter.")
print("=" * 72)
input("Press Enter once you are logged in... ")
page.goto("https://mcp.linear.app/", wait_until="domcontentloaded")
context.storage_state(path=str(state_path))
browser.close()
print(f"\nSaved Linear session to {state_path}")
print("Point the e2e test at it with:")
print(f' export E2E_LINEAR_STORAGE_STATE="{state_path}"')
if __name__ == "__main__":
capture(Path(os.environ.get("E2E_LINEAR_STORAGE_STATE", str(DEFAULT_STATE_PATH))))

442
tests/e2e/mcp/oauth_chat_client.py
View file

@ -1,442 +0,0 @@
"""Client for the mcp chat-completion OAuth e2e suite.
Registers a gateway-managed OAuth (authorization_code) MCP server, seeds the
per-user upstream token by driving the interactive authorize dance with the
official mcp SDK's OAuthClientProvider (the browser leg is a headless Chromium
primed with a human's saved Linear session), then exercises the server through
/chat/completions, where the gateway lists and executes its tools with the
stored per-user token.
Management routes (/v1/mcp/server CRUD, /chat/completions) go through the
shared ProxyClient transport. The MCP protocol used to seed the token goes through
the mcp SDK, the same library production MCP hosts run.
"""
from __future__ import annotations
import asyncio
import re
import time
from dataclasses import dataclass
from typing import TYPE_CHECKING, Final
from urllib.parse import parse_qsl
import httpx
import httpx2
import pytest
from e2e_config import PROXY_BASE_URL, REQUEST_TIMEOUT
from e2e_http import AuthHeaders, NoBody, unwrap
from e2e_metadata import step
from idp import Identity
from mcp import ClientSession
from mcp.client.auth import OAuthClientProvider
from mcp.client.streamable_http import streamable_http_client
from mcp.shared.auth import AuthorizationCodeResult, OAuthClientInformationFull, OAuthClientMetadata, OAuthToken
from mcp.types import TextContent
from models import (
ChatBody,
ChatResponse,
McpOauthUserCredentialStatus,
McpServerCreateBody,
McpServerInfo,
McpServerUserCredentialListResponse,
McpServerUserCredentialRow,
)
from proxy_client import ProxyClient
if TYPE_CHECKING:
from playwright.async_api import Route
# Where the "browser" lands at the end of the authorize dance. Nothing listens
# here: the route interceptor short-circuits the final redirect and reads the
# code/state off its query string, exactly like a desktop MCP host intercepting
# its loopback redirect.
OAUTH_CLIENT_REDIRECT_URI = "http://127.0.0.1:53682/e2e/callback"
BROWSER_CONSENT_TIMEOUT = 60.0
def _mcp_url(alias: str, base_url: str = PROXY_BASE_URL) -> str:
return f"{base_url.rstrip('/')}/{alias}/mcp"
class InMemoryTokenStorage:
"""The mcp SDK's TokenStorage protocol, in memory for one dance: the
DCR-registered client and the gateway tokens minted for it."""
def __init__(self) -> None:
self._tokens: OAuthToken | None = None
self._client_info: OAuthClientInformationFull | None = None
async def get_tokens(self) -> OAuthToken | None:
return self._tokens
async def set_tokens(self, tokens: OAuthToken) -> None:
self._tokens = tokens
async def get_client_info(self) -> OAuthClientInformationFull | None:
return self._client_info
async def set_client_info(self, client_info: OAuthClientInformationFull) -> None:
self._client_info = client_info
async def _browser_follow_authorize(
start_url: str,
storage_state_path: str,
identity: Identity | None = None,
server_alias: str | None = None,
allow_upstream_consent: bool = True,
) -> tuple[str, str | None]:
"""Play the browser's role for a real upstream whose authorize endpoint
serves an interactive consent page (Linear). A headless Chromium primed
with a human's saved Linear session opens the gateway authorize URL and
clicks through Linear's consent screens (the mcp.linear.app Approve form,
then the linear.app workspace-selection page), riding the rest of the chain
(Linear -> gateway callback -> host redirect_uri). The final hop is
intercepted and short-circuited, since nothing listens there, and its
code/state are read off the query string."""
from playwright.async_api import async_playwright
captured: dict[str, str] = {} # mutable-ok: hand-off from the request listener
trail: list[str] = [] # mutable-ok: navigation diagnostics for a failed dance
def _note_request(request: object) -> None:
url = getattr(request, "url", "")
host = httpx.URL(url).host
if not allow_upstream_consent and (host == "linear.app" or host.endswith(".linear.app")):
captured["upstream_consent"] = "seen"
if url.startswith(OAUTH_CLIENT_REDIRECT_URI) and "url" not in captured:
captured["url"] = url
async def _swallow_redirect(route: Route) -> None:
await route.fulfill(status=200, content_type="text/plain", body="ok")
async with async_playwright() as playwright:
browser = await playwright.chromium.launch(headless=True)
context = await browser.new_context(storage_state=storage_state_path)
await context.route(re.compile(re.escape(OAUTH_CLIENT_REDIRECT_URI) + r".*"), _swallow_redirect)
page = await context.new_page()
context.on("request", _note_request)
page.on("framenavigated", lambda frame: trail.append(frame.url.split("?", 1)[0]))
await page.goto(start_url, wait_until="domcontentloaded")
deadline = time.monotonic() + BROWSER_CONSENT_TIMEOUT
while "url" not in captured and time.monotonic() < deadline:
try:
await page.wait_for_load_state("networkidle", timeout=8000)
except Exception: # noqa: BLE001 - a busy consent page never idles; fall through and try to advance it
pass
if "upstream_consent" in captured or "url" in captured:
break
if await page.locator("#username").count() and identity is not None:
await page.locator("#username").fill(identity.username)
await page.locator("#password").fill(identity.password)
await page.locator("#kc-login").click()
continue
if "/ui/connect" in page.url and server_alias is not None:
card = page.locator("div.cursor-pointer").filter(has=page.get_by_text(server_alias, exact=True))
if await card.count() != 1:
await asyncio.sleep(0.5)
continue
connect = card.get_by_text("Connect", exact=True)
if await connect.count():
await connect.click()
continue
if not await card.locator("svg.text-success").count():
await asyncio.sleep(0.5)
continue
finish = page.get_by_role("button", name="Finish connecting", exact=True)
if await finish.count() and await finish.is_enabled():
await finish.click()
continue
control = page.locator(
'button[name="action"][value="approve"], button:has-text("Authorize"), '
'button:has-text("Allow"), button:has-text("@"), a:has-text("@")'
).first
try:
await control.click(timeout=5000)
except Exception: # noqa: BLE001 - nothing to advance yet; loop and re-check
await asyncio.sleep(0.5)
final_url = page.url
await browser.close()
# A redirect chain can finish inside goto/networkidle before the loop checks the page.
assert "upstream_consent" not in captured, "cold reconnect required upstream consent"
landing = captured.get("url")
assert landing is not None, (
f"consent flow never reached {OAUTH_CLIENT_REDIRECT_URI}; "
f"final={final_url.split('?', 1)[0]!r}; trail={trail[-6:]}"
)
params = dict(parse_qsl(httpx.URL(landing).query.decode()))
assert "code" in params, "client redirect_uri carried no authorization code"
return params["code"], params.get("state")
def _oauth_provider(
url: str,
storage: InMemoryTokenStorage,
storage_state_path: str | None,
identity: Identity | None = None,
server_alias: str | None = None,
allow_upstream_consent: bool = True,
) -> OAuthClientProvider:
"""The SDK's real OAuth machinery (RFC 9728/8414 discovery, RFC 7591 DCR,
PKCE, token exchange) with the browser leg driven by Playwright against the
upstream's consent screen."""
code_holder: dict[str, str | None] = {} # mutable-ok: hand-off between the two SDK callbacks
async def _reject_redirect(_: str) -> None:
raise AssertionError("gateway demanded a fresh upstream consent; stored per-user token was not reused")
async def _follow_redirect(authorize_url: str) -> None:
assert storage_state_path is not None
code, state = await _browser_follow_authorize(
authorize_url, storage_state_path, identity, server_alias, allow_upstream_consent
)
code_holder["code"] = code
code_holder["state"] = state
redirect_handler: Final = _reject_redirect if storage_state_path is None else _follow_redirect
async def callback_handler() -> AuthorizationCodeResult:
code = code_holder.get("code")
assert code is not None, "callback_handler ran before the authorize redirect completed"
return AuthorizationCodeResult(code=code, state=code_holder.get("state"))
return OAuthClientProvider(
server_url=url,
client_metadata=OAuthClientMetadata.model_validate(
{
"redirect_uris": [OAUTH_CLIENT_REDIRECT_URI],
"token_endpoint_auth_method": "none",
"grant_types": ["authorization_code", "refresh_token"],
"response_types": ["code"],
"client_name": "e2e-mcp-host",
}
),
storage=storage,
redirect_handler=redirect_handler,
callback_handler=callback_handler,
)
class _HeaderInjectingTransport(httpx2.AsyncBaseTransport):
"""Adds the caller's LiteLLM key header to every outgoing SDK request
(discovery, DCR, token exchange), so the gateway resolves which user to
store the upstream token for from the key on the token exchange, exactly
like a production MCP host configured with a LiteLLM key header."""
def __init__(self, inner: httpx2.AsyncBaseTransport, headers: dict[str, str], gateway_url: str) -> None:
self._inner = inner
self._headers = headers
self._gateway_url = httpx2.URL(gateway_url)
@staticmethod
def _port(url: httpx2.URL) -> int | None:
if url.port is not None:
return url.port
return {"http": 80, "https": 443}.get(url.scheme)
async def handle_async_request(self, request: httpx2.Request) -> httpx2.Response:
same_origin: Final = (
request.url.scheme == self._gateway_url.scheme
and request.url.host == self._gateway_url.host
and self._port(request.url) == self._port(self._gateway_url)
)
if same_origin:
for name, value in self._headers.items():
if name not in request.headers:
request.headers[name] = value
else:
for name, value in self._headers.items():
if request.headers.get(name) == value:
del request.headers[name]
return await self._inner.handle_async_request(request)
async def aclose(self) -> None:
await self._inner.aclose()
def _oauth_http_client(
headers: dict[str, str], auth: OAuthClientProvider, gateway_url: str = PROXY_BASE_URL
) -> httpx2.AsyncClient:
return httpx2.AsyncClient(
auth=auth,
timeout=httpx2.Timeout(REQUEST_TIMEOUT),
follow_redirects=True,
transport=_HeaderInjectingTransport(httpx2.AsyncHTTPTransport(), headers, gateway_url),
)
async def _seed_via_dance(
url: str, headers: dict[str, str], storage: InMemoryTokenStorage, storage_state_path: str
) -> tuple[str, ...]:
async with _oauth_http_client(headers, _oauth_provider(url, storage, storage_state_path)) as http_client:
async with streamable_http_client(url, http_client=http_client) as (read, write):
async with ClientSession(read, write) as session:
await session.initialize()
listed = await session.list_tools()
return tuple(sorted(tool.name for tool in listed.tools))
@dataclass(frozen=True, slots=True)
class OauthToolRun:
tools: tuple[str, ...]
is_error: bool
text: str
async def _list_and_call(
url: str,
headers: dict[str, str],
storage: InMemoryTokenStorage,
storage_state_path: str | None,
tool: str,
arguments: dict[str, str],
gateway_url: str = PROXY_BASE_URL,
identity: Identity | None = None,
server_alias: str | None = None,
allow_upstream_consent: bool = True,
) -> OauthToolRun:
async with _oauth_http_client(
headers,
_oauth_provider(url, storage, storage_state_path, identity, server_alias, allow_upstream_consent),
gateway_url,
) as http_client:
async with streamable_http_client(url, http_client=http_client) as (read, write):
async with ClientSession(read, write) as session:
await session.initialize()
listed: Final = await session.list_tools()
result: Final = await session.call_tool(tool, arguments)
text: Final = "".join(content.text for content in result.content if isinstance(content, TextContent))
return OauthToolRun(
tools=tuple(sorted(tool_item.name for tool_item in listed.tools)),
is_error=result.is_error,
text=text,
)
@dataclass(frozen=True, slots=True)
class ChatMcpClient:
proxy: ProxyClient
@step("Register the MCP server with the alias {body.alias}")
def create_server(self, body: McpServerCreateBody) -> McpServerInfo:
return unwrap(
self.proxy.transport.post(
"/v1/mcp/server",
headers=self.proxy.transport.master,
json=body,
response_type=McpServerInfo,
)
)
@step("Read the MCP server back from /v1/mcp/server")
def server_info(self, server_id: str) -> McpServerInfo:
return unwrap(
self.proxy.transport.get(
f"/v1/mcp/server/{server_id}",
headers=self.proxy.transport.master,
params=NoBody(),
response_type=McpServerInfo,
)
)
@step("Delete the MCP server")
def delete_server(self, server_id: str) -> None:
_ = self.proxy.transport.delete(
f"/v1/mcp/server/{server_id}",
headers=self.proxy.transport.master,
json=NoBody(),
response_type=NoBody,
)
@step("Sign the key's user in to the MCP server {alias} through the OAuth consent flow")
def seed_user_token(self, alias: str, key: str, storage_state_path: str) -> tuple[str, ...]:
"""Drive the interactive authorize dance for `key`'s user so the gateway
stores their upstream token, retried to the shared deadline since the
just-created server and key propagate asynchronously. The LiteLLM key
rides x-litellm-api-key so the gateway binds the token to that user.
Returns the upstream tool names the dance listed, proof the token works."""
headers = {"x-litellm-api-key": f"Bearer {key}"}
storage = InMemoryTokenStorage()
deadline = time.monotonic() + self.proxy.poll_timeout
last_error: Exception | None = None
while time.monotonic() < deadline:
try:
return asyncio.run(_seed_via_dance(_mcp_url(alias), headers, storage, storage_state_path))
except Exception as exc: # noqa: BLE001 - retried to the deadline; the last error surfaces below
last_error = exc
time.sleep(self.proxy.poll_interval)
pytest.fail(
f"authorize dance for {alias!r} never completed within {self.proxy.poll_timeout}s; "
f"last error: {last_error!r}"
)
@step("List the tools on the MCP server {alias} and call {tool} over the MCP protocol")
def list_and_call(
self,
alias: str,
headers: dict[str, str],
storage: InMemoryTokenStorage,
storage_state_path: str | None,
tool: str,
arguments: dict[str, str],
base_url: str = PROXY_BASE_URL,
identity: Identity | None = None,
allow_upstream_consent: bool = True,
) -> OauthToolRun:
return asyncio.run(
_list_and_call(
f"{base_url.rstrip('/')}/mcp" if identity is not None else _mcp_url(alias, base_url),
headers,
storage,
storage_state_path,
tool,
arguments,
base_url,
identity,
alias,
allow_upstream_consent,
)
)
@step("List the users with a stored OAuth token for the MCP server")
def server_user_credentials(self, server_id: str) -> tuple[McpServerUserCredentialRow, ...]:
return unwrap(
self.proxy.transport.get(
f"/v1/mcp/server/{server_id}/user-credentials",
headers=self.proxy.transport.master,
params=NoBody(),
response_type=McpServerUserCredentialListResponse,
)
).root
@step("Revoke the user's stored OAuth token for the MCP server")
def revoke_user_token(self, server_id: str, headers: AuthHeaders) -> None:
_ = unwrap(
self.proxy.transport.delete(
f"/v1/mcp/server/{server_id}/oauth-user-credential",
headers=headers,
json=NoBody(),
response_type=McpOauthUserCredentialStatus,
)
)
@step("Send a /chat/completions request to {body.model} with an MCP server attached as a tool")
def chat_with_mcp(self, headers: AuthHeaders, body: ChatBody) -> ChatResponse:
"""POST /chat/completions carrying the LiteLLM key in `headers` (either
ingress form) with an MCP server attached in `body.tools`. The gateway
resolves the user from the key and lists/executes the server's tools
with that user's stored upstream token."""
return unwrap(
self.proxy.transport.post(
"/chat/completions",
headers=headers,
json=body,
response_type=ChatResponse,
)
)
def build_chat_client(proxy: ProxyClient) -> ChatMcpClient:
return ChatMcpClient(proxy=proxy)

197
tests/e2e/mcp/oauth_gateway.py
View file

@ -1,197 +0,0 @@
"""An owned, source-built OAuth gateway with cold restarts and credential observations.
Only this child process is restarted. Its database and SSO client survive while
its process-local caches do not; Redis is deliberately absent from its config.
The optional live edge measures headers without recording credentials or bodies.
"""
from __future__ import annotations
import os
import subprocess
import sys
import threading
import time
from collections.abc import Callable, Mapping
from contextlib import ExitStack
from dataclasses import dataclass, field
from pathlib import Path
from typing import Final
import psycopg
from e2e_config import INHERITED_ENV_PREFIXES, available_port
from e2e_http import NoBody
from e2e_metadata import step
from idp import Keycloak, stop_process_group
from proxy_client import ProxyClient, build_proxy_client
from psycopg.rows import class_row
from pydantic import BaseModel, SecretStr, ValidationError
class StoredOAuth(BaseModel):
type: str
access_token: SecretStr
@dataclass(frozen=True, slots=True)
class CredentialRow:
credential_b64: str = field(repr=False)
@step("Read the user's stored OAuth credential for the MCP server from the database and decrypt it")
def stored_oauth(user_id: str, server_id: str) -> StoredOAuth:
"""Read the encrypted credential because management APIs omit the plaintext token."""
from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value_helper
with psycopg.Connection[CredentialRow].connect(
os.environ["DATABASE_URL"], row_factory=class_row(CredentialRow)
) as conn:
row: Final = conn.execute(
'SELECT credential_b64 FROM "LiteLLM_MCPUserCredentials" WHERE user_id = %s AND server_id = %s',
(user_id, server_id),
).fetchone()
assert row is not None, "canonical user/server has no persisted credential"
plaintext: Final = decrypt_value_helper(
row.credential_b64, "e2e_mcp_oauth", exception_type="debug", return_original_value=False
)
assert plaintext is not None, "persisted credential must decrypt with the gateway salt"
assert plaintext != row.credential_b64, "persisted credential must be encrypted"
try:
credential: Final = StoredOAuth.model_validate_json(plaintext)
except ValidationError:
raise AssertionError("decrypted credential is not an OAuth payload") from None
assert credential.type == "oauth2"
assert bool(credential.access_token.get_secret_value()), "stored upstream token is empty"
return credential
class RpcMethod(BaseModel):
method: str = ""
@dataclass(slots=True)
class OAuthObservation:
gateway_token: str = field(default="", repr=False)
_seen: tuple[tuple[str, str, bool], ...] = field(default=(), init=False, repr=False)
_lock: threading.Lock = field(default_factory=threading.Lock, init=False, repr=False)
def observe(self, url: str, headers: Mapping[str, str], body: bytes | None) -> None:
if body is None or not url.endswith("/mcp"):
return
try:
operation: Final = RpcMethod.model_validate_json(body).method
except ValidationError:
return
if operation not in ("tools/list", "tools/call"):
return
received: Final = headers.get("authorization", "")
gateway_leaked: Final = any(
value in (self.gateway_token, f"Bearer {self.gateway_token}") for value in headers.values()
)
with self._lock:
self._seen = (*self._seen, (operation, received, gateway_leaked))
def assert_forwarded(self, expected: StoredOAuth) -> None:
with self._lock:
snapshot: Final = self._seen
self._seen = ()
assert {item[0] for item in snapshot} == {"tools/list", "tools/call"}, "missing upstream observations"
expected_header: Final = f"Bearer {expected.access_token.get_secret_value()}"
assert all(item[1] == expected_header for item in snapshot), "upstream bearer did not match the stored token"
assert all(not item[2] for item in snapshot), "gateway bearer leaked to the upstream"
@dataclass(slots=True)
class OAuthGateway:
base_url: str
proxy: ProxyClient
_environment: Mapping[str, str] = field(repr=False)
_command: tuple[str, ...] = field(repr=False)
_log_path: Path
_child: subprocess.Popen[bytes] | None = field(default=None, init=False, repr=False)
@step("Start the separate LiteLLM proxy and wait for /health/liveliness")
def start(self) -> None:
with self._log_path.open("ab") as log:
self._child = subprocess.Popen(
self._command,
env=self._environment,
stdout=log,
stderr=log,
start_new_session=True,
)
deadline: Final = time.monotonic() + 120
while time.monotonic() < deadline:
assert self._child.poll() is None, "owned OAuth gateway exited; inspect its private log"
result = self.proxy.transport.probe("/health/liveliness", params=NoBody())
if result.status_code == 200:
return
time.sleep(0.5)
raise AssertionError("owned OAuth gateway did not become ready")
@step("Stop the separate LiteLLM proxy")
def stop(self) -> None:
if self._child is not None:
stop_process_group(self._child)
assert self._child.poll() is not None, "old gateway process is still alive"
@step("Restart the separate LiteLLM proxy process so its in-memory caches start empty")
def restart(self) -> None:
assert self._child is not None
previous: Final = self._child.pid
self.stop()
self.start()
assert self._child.pid != previous, "gateway restart did not create a new process"
@step("Start a separate LiteLLM proxy from source with JWT auth against Keycloak")
def owned_gateway(idp: Keycloak, directory: Path, cleanup: ExitStack) -> OAuthGateway:
for name in ("DATABASE_URL", "LITELLM_LICENSE", "LITELLM_SALT_KEY", "LITELLM_MASTER_KEY"):
assert os.environ.get(name), f"{name} is required for the owned OAuth gateway"
port: Final = available_port()
base_url: Final = f"http://127.0.0.1:{port}"
def defer(callback: Callable[[], object]) -> None:
cleanup.callback(callback)
browser: Final = idp.browser_client(callback_url=f"{base_url}/sso/callback", defer=defer)
config: Final = directory / "oauth-gateway.yaml"
config.write_text(
"model_list: []\n"
"general_settings:\n"
" master_key: os.environ/LITELLM_MASTER_KEY\n"
" database_url: os.environ/DATABASE_URL\n"
" enable_jwt_auth: true\n"
" litellm_jwtauth:\n"
" user_id_jwt_field: sub\n"
" user_email_jwt_field: email\n"
" team_ids_jwt_field: groups\n"
" user_id_upsert: true\n"
)
environment: Final = {
**{key: value for key, value in os.environ.items() if not key.startswith(INHERITED_ENV_PREFIXES)},
**browser.environment(idp.discovery()),
"PROXY_BASE_URL": base_url,
"JWT_PUBLIC_KEY_URL": idp.jwks_url,
"JWT_ISSUER": idp.issuer,
"JWT_AUDIENCE": "litellm-e2e",
"DISABLE_SCHEMA_UPDATE": "true",
"STORE_MODEL_IN_DB": "True",
"PYTHONPATH": str(Path(__file__).resolve().parents[3]),
}
gateway: Final = OAuthGateway(
base_url=base_url,
proxy=build_proxy_client(
base_url=base_url,
control_plane_base_url=base_url,
replica_urls=(base_url,),
control_replica_urls=(base_url,),
master_key=os.environ["LITELLM_MASTER_KEY"],
),
_environment=environment,
_command=(sys.executable, "-m", "litellm.proxy.proxy_cli", "--config", str(config), "--port", str(port)),
_log_path=directory / "oauth-gateway.log",
)
cleanup.callback(gateway.stop)
gateway.start()
return gateway

218
tests/e2e/mcp/test_mcp_chat_completion_oauth_e2e.py
View file

@ -1,218 +0,0 @@
"""On-demand e2e: a chat completion drives a gateway-managed OAuth MCP server.
The real end-user flow for MCP over an OAuth server: a user registers a Linear
authorization_code server, authorizes it once so the gateway stores their
upstream token, then sends a normal /chat/completions request with the Linear
MCP attached. The gateway resolves the user from the LiteLLM key, lists Linear's
tools with the stored per-user token, lets the model call one, executes it
upstream with that token, and returns the answer. This is proven against the
real Linear MCP server (mcp.linear.app) and a real Anthropic model, once per
documented ingress header (x-litellm-api-key and Authorization).
The authorize dance is seeded through the mcp SDK's OAuthClientProvider; the one
step Linear cannot auto-approve is the human consent, so it is captured once out
of band (mcp/linear_session_capture.py) into a saved browser session and a
headless Chromium clicks Approve every run. The test therefore skips unless
E2E_LINEAR_STORAGE_STATE points at that session, so it never runs on the per-PR
CI path; it is a nightly/on-demand real-server smoke test.
Fail-before-fix: without the stored per-user token the gateway lists no Linear
tools, so mcp_list_tools comes back empty, nothing is called, and the
assertions fail; a served, called, non-empty Linear tool proves the gateway
pulled and used the user's token.
"""
from __future__ import annotations
import os
import pytest
from e2e_config import CHEAP_ANTHROPIC_MODEL, LINEAR_MCP_URL, LINEAR_STORAGE_STATE, unique_marker
from e2e_http import AuthHeaders
from e2e_metadata import Capability, Domain, Mode, Provider, Route, Subject, meta
from lifecycle import ResourceManager
from models import ChatBody, ChatMessage, KeyGenerateBody, McpChatTool, McpServerCreateBody, ObjectPermission
from proxy_client import ProxyClient
pytest.importorskip("mcp", reason="mcp SDK not installed; run `uv sync --inexact --group e2e-dev`")
pytest.importorskip(
"playwright.async_api",
reason="playwright not installed; run `uv pip install playwright` and `playwright install chromium`",
)
from oauth_chat_client import ChatMcpClient, build_chat_client # noqa: E402 # imports follow the importorskip guards
pytestmark = [
pytest.mark.e2e,
pytest.mark.skipif(
not LINEAR_STORAGE_STATE or not os.path.exists(LINEAR_STORAGE_STATE),
reason="set E2E_LINEAR_STORAGE_STATE to a Linear session captured via mcp/linear_session_capture.py",
),
]
# Pinned from a live dance during verification (never guessed); the gateway
# prefixes every upstream tool name with the server alias. list_teams is a
# read-only Linear tool that takes no arguments and returns the caller's teams.
LINEAR_READONLY_TOOL = "list_teams"
LINEAR_PROMPT = "Use the list_teams tool to list my Linear teams, then reply with the name of one of them."
@pytest.fixture(scope="session")
def chat_client(proxy: ProxyClient) -> ChatMcpClient:
return build_chat_client(proxy)
class TestMcpChatCompletionOauth:
"""A scoped internal-user key on a real Linear authorization_code server,
used through /chat/completions once per ingress header: the gateway pulls
the user's stored upstream token, lists and executes Linear's tools during
the completion, and returns the answer."""
@pytest.mark.covers("mcp.list_tools.oauth.succeeds")
@pytest.mark.covers("mcp.call_tool.oauth.succeeds")
@meta(
Subject(
domain=Domain.MCP,
route=Route.CHAT_COMPLETIONS,
providers=(Provider.ANTHROPIC,),
models=(CHEAP_ANTHROPIC_MODEL,),
capabilities=(Capability.FUNCTION_CALLING,),
mode=Mode.NONSTREAM,
)
)
def test_chat_completion_uses_linear_with_x_litellm_api_key_header(
self, chat_client: ChatMcpClient, resources: ResourceManager
) -> None:
marker = unique_marker()
alias = f"e2elinear{marker}"
created = chat_client.create_server(
McpServerCreateBody(
alias=alias,
url=LINEAR_MCP_URL,
allow_all_keys=False,
auth_type="oauth2",
oauth2_flow="authorization_code",
)
)
resources.defer(lambda: chat_client.delete_server(created.server_id))
stored = chat_client.server_info(created.server_id)
assert stored.auth_type == "oauth2"
assert stored.oauth2_flow == "authorization_code"
assert stored.allow_all_keys is False
key = chat_client.proxy.generate_key(
KeyGenerateBody(
user_id="e2e-test-user",
object_permission=ObjectPermission(mcp_servers=[created.server_id]),
)
)
resources.defer(lambda: chat_client.proxy.delete_key(key))
seeded = chat_client.seed_user_token(alias, key, LINEAR_STORAGE_STATE)
assert f"{alias}-{LINEAR_READONLY_TOOL}" in seeded, (
f"the authorize dance listed {seeded}, expected it to include {alias}-{LINEAR_READONLY_TOOL}"
)
response = chat_client.chat_with_mcp(
AuthHeaders.model_validate({"x-litellm-api-key": f"Bearer {key}"}),
ChatBody(
model=CHEAP_ANTHROPIC_MODEL,
messages=[ChatMessage(role="user", content=LINEAR_PROMPT)],
tools=[
McpChatTool(
server_url=f"litellm_proxy/mcp/{alias}",
server_label=alias,
require_approval="never",
)
],
),
)
message = response.choices[0].message
assert message is not None and message.content, f"completion returned no answer: {response}"
meta = message.provider_specific_fields
assert meta is not None, f"no MCP metadata on the completion: {response}"
listed = {t.function.name for t in (meta.mcp_list_tools or []) if t.function}
assert f"{alias}-{LINEAR_READONLY_TOOL}" in listed, (
f"the gateway listed {sorted(listed)}, expected the stored token to surface {alias}-{LINEAR_READONLY_TOOL}"
)
results = [r for r in (meta.mcp_call_results or []) if r.name == f"{alias}-{LINEAR_READONLY_TOOL}"]
assert results and results[0].result, (
f"Linear tool {alias}-{LINEAR_READONLY_TOOL} was not executed with a result: {meta.mcp_call_results}"
)
@pytest.mark.covers("mcp.list_tools.oauth.succeeds")
@pytest.mark.covers("mcp.call_tool.oauth.succeeds")
@meta(
Subject(
domain=Domain.MCP,
route=Route.CHAT_COMPLETIONS,
providers=(Provider.ANTHROPIC,),
models=(CHEAP_ANTHROPIC_MODEL,),
capabilities=(Capability.FUNCTION_CALLING,),
mode=Mode.NONSTREAM,
)
)
def test_chat_completion_uses_linear_with_authorization_bearer_header(
self, chat_client: ChatMcpClient, resources: ResourceManager
) -> None:
marker = unique_marker()
alias = f"e2elinear{marker}"
created = chat_client.create_server(
McpServerCreateBody(
alias=alias,
url=LINEAR_MCP_URL,
allow_all_keys=False,
auth_type="oauth2",
oauth2_flow="authorization_code",
)
)
resources.defer(lambda: chat_client.delete_server(created.server_id))
stored = chat_client.server_info(created.server_id)
assert stored.auth_type == "oauth2"
assert stored.oauth2_flow == "authorization_code"
assert stored.allow_all_keys is False
key = chat_client.proxy.generate_key(
KeyGenerateBody(
user_id="e2e-test-user",
object_permission=ObjectPermission(mcp_servers=[created.server_id]),
)
)
resources.defer(lambda: chat_client.proxy.delete_key(key))
seeded = chat_client.seed_user_token(alias, key, LINEAR_STORAGE_STATE)
assert f"{alias}-{LINEAR_READONLY_TOOL}" in seeded, (
f"the authorize dance listed {seeded}, expected it to include {alias}-{LINEAR_READONLY_TOOL}"
)
response = chat_client.chat_with_mcp(
AuthHeaders.model_validate({"authorization": f"Bearer {key}"}),
ChatBody(
model=CHEAP_ANTHROPIC_MODEL,
messages=[ChatMessage(role="user", content=LINEAR_PROMPT)],
tools=[
McpChatTool(
server_url=f"litellm_proxy/mcp/{alias}",
server_label=alias,
require_approval="never",
)
],
),
)
message = response.choices[0].message
assert message is not None and message.content, f"completion returned no answer: {response}"
meta = message.provider_specific_fields
assert meta is not None, f"no MCP metadata on the completion: {response}"
listed = {t.function.name for t in (meta.mcp_list_tools or []) if t.function}
assert f"{alias}-{LINEAR_READONLY_TOOL}" in listed, (
f"the gateway listed {sorted(listed)}, expected the stored token to surface {alias}-{LINEAR_READONLY_TOOL}"
)
results = [r for r in (meta.mcp_call_results or []) if r.name == f"{alias}-{LINEAR_READONLY_TOOL}"]
assert results and results[0].result, (
f"Linear tool {alias}-{LINEAR_READONLY_TOOL} was not executed with a result: {meta.mcp_call_results}"
)

216
tests/e2e/mcp/test_mcp_oauth_happy_path_e2e.py
View file

@ -1,216 +0,0 @@
"""Real OAuth consent, immediate MCP operations and cold-restart persistence.
Aggregate SSO uses the SDK's normal authentication. The per-server variant is
explicitly a configured two-header client, not an Authorization-only OAuth host.
The observed variants forward to the same real Linear upstream and compare its
bearer at the forwarding boundary; direct variants retain unmodified discovery.
"""
from __future__ import annotations
import os
from collections.abc import Iterator
from contextlib import ExitStack
from pathlib import Path
from types import MappingProxyType
from typing import Final, Literal
import pytest
from e2e_config import LINEAR_MCP_URL, LINEAR_READONLY_TOOL, LINEAR_STORAGE_STATE, unique_marker
from e2e_http import AuthHeaders, NoBody, get_external, unwrap
from e2e_metadata import Domain, Route, Subject, meta
from idp import Identity, Keycloak
from lifecycle import ResourceManager
from models import (
McpOauthCredentials,
McpServerCreateBody,
ObjectPermission,
TeamMemberAddBody,
TeamMemberEntry,
TeamUpdateBody,
)
from oauth_chat_client import ChatMcpClient, InMemoryTokenStorage, OauthToolRun, build_chat_client
from oauth_gateway import OAuthGateway, OAuthObservation, owned_gateway, stored_oauth
from provider_edge import LiveEdge, start_provider_edge
from proxy_client import ProxyClient
from pydantic import BaseModel, ValidationError
pytestmark = [pytest.mark.e2e, pytest.mark.mcp_oauth_live, pytest.mark.provider_live]
class OAuthMetadata(BaseModel):
issuer: str
authorization_endpoint: str
token_endpoint: str
registration_endpoint: str
class LinearTeam(BaseModel):
id: str
name: str
class LinearTeams(BaseModel):
teams: tuple[LinearTeam, ...]
def assert_tool_result(run: OauthToolRun, tool: str) -> None:
assert tool in run.tools
assert run.is_error is False
try:
result: Final = LinearTeams.model_validate_json(run.text)
except ValidationError:
raise AssertionError("list_teams did not return the expected teams payload") from None
assert result.teams, "the test workspace must contain at least one team"
assert all(team.id and team.name for team in result.teams), "team results must contain identifiers and names"
@pytest.fixture(scope="module")
def oauth_gateway(idp: Keycloak, tmp_path_factory: pytest.TempPathFactory) -> Iterator[OAuthGateway]:
assert LINEAR_STORAGE_STATE and Path(LINEAR_STORAGE_STATE).is_file(), (
"E2E_LINEAR_STORAGE_STATE must name a captured Linear login; see mcp/linear_session_capture.py"
)
assert os.environ.get("E2E_FIXTURE_MODE", "live") == "live", "OAuth acceptance cannot use replay"
with ExitStack() as cleanup:
yield owned_gateway(idp, tmp_path_factory.mktemp("mcp-oauth"), cleanup)
@pytest.fixture(scope="module")
def proxy(oauth_gateway: OAuthGateway) -> ProxyClient:
return oauth_gateway.proxy
@pytest.fixture(scope="module")
def client(proxy: ProxyClient) -> ChatMcpClient:
return build_chat_client(proxy)
class TestMcpOauthHappyPath:
@pytest.mark.covers("mcp.list_tools.oauth.succeeds")
@pytest.mark.covers("mcp.call_tool.oauth.succeeds")
@pytest.mark.covers("mcp.call_tool.oauth.persists_across_processes")
@meta(
Subject(
domain=Domain.MCP,
route=Route.MCP,
)
)
@pytest.mark.parametrize("route", ("aggregate_sso", "explicit_header_jwt"))
@pytest.mark.parametrize("observed", (False, True), ids=("direct", "observed"))
def test_consent_list_call_and_cold_restart(
self,
client: ChatMcpClient,
resources: ResourceManager,
jwt_identity: Identity,
idp: Keycloak,
oauth_gateway: OAuthGateway,
route: Literal["aggregate_sso", "explicit_header_jwt"],
observed: bool,
) -> None:
alias: Final = f"e2elinear{unique_marker()}"
tool: Final = f"{alias}-{LINEAR_READONLY_TOOL}"
token: Final = idp.access_token(jwt_identity)
observation: Final = OAuthObservation(gateway_token=token)
edge: Final = (
start_provider_edge(
LiveEdge(observe_request=observation.observe),
mounts=MappingProxyType(
{"linear": "https://mcp.linear.app", ".well-known": "https://mcp.linear.app/.well-known"}
),
)
if observed
else None
)
if edge is not None:
resources.defer(edge.shutdown)
metadata: Final = (
unwrap(
get_external(
"https://mcp.linear.app/.well-known/oauth-authorization-server",
response_type=OAuthMetadata,
)
)
if observed
else None
)
created: Final = client.create_server(
McpServerCreateBody(
alias=alias,
server_name=alias,
url=f"{edge.edge.api_base('linear')}/mcp" if edge is not None else LINEAR_MCP_URL,
transport="http",
allow_all_keys=False,
auth_type="oauth2",
oauth2_flow="authorization_code",
per_server_oauth_discovery=route == "explicit_header_jwt",
issuer=metadata.issuer if metadata else None,
authorization_url=metadata.authorization_endpoint if metadata else None,
token_url=metadata.token_endpoint if metadata else None,
registration_url=metadata.registration_endpoint if metadata else None,
credentials=McpOauthCredentials(upstream_resource=LINEAR_MCP_URL) if observed else None,
)
)
resources.defer(lambda: client.delete_server(created.server_id))
assert client.server_user_credentials(created.server_id) == (), (
"scenario must start without upstream credentials"
)
unwrap(
client.proxy.transport.post(
"/team/member_add",
headers=client.proxy.transport.master,
json=TeamMemberAddBody(
team_id=jwt_identity.group, member=TeamMemberEntry(user_id=jwt_identity.user_id, role="user")
),
response_type=NoBody,
)
)
client.proxy.update_team(
TeamUpdateBody(
team_id=jwt_identity.group,
object_permission=ObjectPermission(mcp_servers=[created.server_id]),
)
)
headers: Final = {"x-litellm-api-key": f"Bearer {token}"} if route == "explicit_header_jwt" else {}
resources.defer(
lambda: client.revoke_user_token(
created.server_id,
AuthHeaders(authorization=f"Bearer {idp.access_token(jwt_identity)}"),
)
)
identity: Final = jwt_identity if route == "aggregate_sso" else None
first: Final = client.list_and_call(
alias,
headers,
InMemoryTokenStorage(),
LINEAR_STORAGE_STATE,
tool,
{},
base_url=oauth_gateway.base_url,
identity=identity,
)
assert_tool_result(first, tool)
credentials: Final = client.server_user_credentials(created.server_id)
assert len(credentials) == 1
assert credentials[0].user_id == jwt_identity.user_id
assert credentials[0].credential_type == "oauth2"
first_stored_oauth: Final = stored_oauth(jwt_identity.user_id, created.server_id)
if observed:
observation.assert_forwarded(first_stored_oauth)
oauth_gateway.restart()
fresh_token: Final = idp.access_token(jwt_identity)
observation.gateway_token = fresh_token
second: Final = client.list_and_call(
alias,
{"Authorization": f"Bearer {fresh_token}"} if identity is None else {},
InMemoryTokenStorage(),
LINEAR_STORAGE_STATE if identity is not None else None,
tool,
{},
base_url=oauth_gateway.base_url,
identity=identity,
allow_upstream_consent=False,
)
assert_tool_result(second, tool)
second_stored_oauth: Final = stored_oauth(jwt_identity.user_id, created.server_id)
if observed:
observation.assert_forwarded(second_stored_oauth)

5
tests/e2e/mcp/test_mcp_semantic_tool_filter_e2e.py
View file

@ -5,7 +5,8 @@ embedding deployment, and turns on `litellm_settings.mcp_semantic_tool_filter` t
/config/update. A chat completion that attaches the server and asks to search logs must report,
in the x-litellm-semantic-filter headers, that the catalog was cut to at most top_k tools, all
of them tools the key lists. Requires DD_API_KEY + DD_APP_KEY and the proxy's
OPENAI_API_KEY for the embedding deployment.
OPENAI_API_KEY for the embedding deployment. The filter is proxy-wide, so the test is
marked `proxy_wide_config` and runs on a stack the main suite is not driving.
"""
from __future__ import annotations
@ -23,7 +24,7 @@ from mcp_client import McpClient
from models import ChatBody, ChatMessage, LiteLLMParamsBody, McpChatTool
from pydantic import BaseModel
pytestmark = [pytest.mark.e2e, pytest.mark.quiet_stack]
pytestmark = [pytest.mark.e2e, pytest.mark.proxy_wide_config]
TOP_K: Final = 3

59
tests/e2e/models.py
View file

@ -730,28 +730,14 @@ class McpInfo(BaseModel):
logo_url: str | None = None
class McpOauthCredentials(BaseModel):
upstream_resource: str
class McpServerCreateBody(BaseModel):
"""POST /v1/mcp/server. For a gateway-managed OAuth server, `auth_type` is
`oauth2` and `oauth2_flow` is `authorization_code`; the upstream endpoints
are discovered and registered via DCR when left unset. `allow_all_keys`
false scopes the server to keys granted it through object_permission."""
"""POST /v1/mcp/server. `allow_all_keys` false scopes the server to keys
granted it through object_permission."""
alias: str
url: str
transport: str = "http"
allow_all_keys: bool = True
auth_type: str | None = None
oauth2_flow: Literal["client_credentials", "authorization_code"] | None = None
per_server_oauth_discovery: bool | None = None
issuer: str | None = None
authorization_url: str | None = None
token_url: str | None = None
registration_url: str | None = None
credentials: McpOauthCredentials | None = Field(default=None, repr=False)
server_name: str | None = None
description: str | None = None
mcp_info: McpInfo | None = None
@ -791,26 +777,6 @@ class McpServerListResponse(RootModel[list[McpServerRow]]):
"""GET /v1/mcp/server answers with a bare array of servers."""
class McpServerUserCredentialRow(BaseModel):
user_id: str
credential_type: Literal["oauth2", "byok"]
expires_at: str | None = None
connected_at: str | None = None
updated_at: str
class McpServerUserCredentialListResponse(RootModel[tuple[McpServerUserCredentialRow, ...]]):
"""GET /v1/mcp/server/{server_id}/user-credentials answers with a bare array."""
class McpOauthUserCredentialStatus(BaseModel):
server_id: str
has_credential: bool
expires_at: str | None = None
is_expired: bool = False
connected_at: str | None = None
class ToolsetTool(BaseModel):
server_id: str
tool_name: str
@ -1295,27 +1261,6 @@ class CostMap(RootModel[dict[str, CostMapEntry]]):
pass
class ConfigPatchResponse(BaseModel):
status: str
values: dict[str, float | dict[str, float]]
class CostDiscountConfig(RootModel[dict[str, float]]):
pass
class CostMarginConfig(RootModel[dict[str, float | dict[str, float]]]):
pass
class CostDiscountConfigResponse(BaseModel):
values: dict[str, float]
class CostMarginConfigResponse(BaseModel):
values: dict[str, float | dict[str, float]]
class BedrockGuardrailParams(BaseModel):
guardrail: Literal["bedrock"] = "bedrock"
mode: Literal["pre_call"] = "pre_call"

2
tests/e2e/perf/test_perf_cases.py
View file

@ -32,7 +32,7 @@ from setup.podcpu import gateway_cpu_sampler
from setup.provision import MOCK_PROVIDERS, build_setup, guardrail_url, mock_url
from setup.run import run_plan
pytestmark = [pytest.mark.perf, pytest.mark.quiet_stack]
pytestmark = pytest.mark.perf
_SPEND_SAMPLE: Final = 40

111
tests/e2e/router/test_reliability_memory_e2e.py → tests/e2e/perf/test_reliability_memory_e2e.py
View file

@ -1,17 +1,19 @@
"""Live e2e: a few hundred requests that fail before any provider answers must not
grow the proxy's resident memory past a loose ceiling once the proxy is warm.
grow the proxy's resident memory past a loose ceiling once the proxy is warm. It reads
the proxy's own RSS, so it runs in the perf lane, where nothing else is hitting the stack.
The regression this guards shipped in v1.100.0: every retry breadcrumb copied the
whole request and the copies nested into one router-global list, so a proxy under
retry-heavy failing traffic grew until it was OOM-killed. The traffic here has that
shape: a model group whose deployments refuse at the socket (an unreachable base
URL) with cooldown_time 0 so the router keeps retrying them, per-request retries,
and a fallback group that refuses the same way, each request carrying a long chat
URL, with a literal key since the perf stack carries no provider keys) with
cooldown_time 0 so the router keeps retrying them, per-request retries, and a
fallback group that refuses the same way, each request carrying a long chat
transcript so every whole-request copy costs hundreds of containers instead of a
handful. Under the stack's cooldown policy a
deployment that fails a handful of times in a row is benched (a bad-credential 401
included), the router answers "No deployments available" without retrying, and the
retry loop that leaks stops running; cooldown_time 0 keeps it running.
handful. Under the stack's cooldown policy a deployment that fails a handful of
times in a row is benched (a bad-credential 401 included), the router answers "No
deployments available" without retrying, and the retry loop that leaks stops running;
cooldown_time 0 keeps it running.
Four identical batches run back to back. The first two are the warmup that grows
the proxy's caches and allocator arenas, the last two are the measured window, and
@ -35,19 +37,11 @@ single load-balanced address reaches the workers behind it one answer at a time,
a worker that answered only one checkpoint would otherwise drop out of the
comparison, which is where a leaking worker could hide.
RSS alone is a coarse gauge: on the release stack (spend logs storing prompts,
json logs, prometheus and otel callbacks) the same v1.100.0 breadcrumbs grew RSS
by only about 15 MB per 300 failing requests, while every failing request's stored
request snapshot carried a copy of the request per failed attempt, over 100 KB on
the first call and a couple of MB once the copies nested, against tens of KB with
the fix. So the first check sends one failing request before the phases, reads its
spend log back through /spend/logs, and holds the stored request body to a fixed
size budget: the deterministic catch for a breadcrumb that copies the whole
request. It runs before the phases because the leaking writer drops its own rows
under the phases' traffic (a queue budget hit, a recursion limit on the nested
copies), which would turn the size check into a missing-row check.
The stored-request half of that regression (every failing request's spend log carrying a
copy of the request per failed attempt) is a deterministic size check, so it lives in
tests/integration/spend/test_failing_request_stored_snapshot.py.
A second, cheaper check holds the idle footprint: every worker's RSS as the harness
A cheaper check holds the idle footprint: every worker's RSS as the harness
read it at collection time, before this pytest process sent any traffic (see
conftest.pytest_collection_finish), must sit under a fixed budget. On the
release gate that is a fresh stack right after its readiness gate, one worker per
@ -74,7 +68,6 @@ from typing import Final
import pytest
from complexity_router_client import ComplexityRouterClient
from e2e_config import (
MEMORY_CONCURRENCY,
MEMORY_IDLE_RSS_BUDGET_MB,
@ -83,7 +76,6 @@ from e2e_config import (
MEMORY_RSS_CEILING_MB,
MEMORY_RSS_SAMPLE_INTERVAL_SECONDS,
MEMORY_RSS_SETTLE_SAMPLES,
MEMORY_STORED_REQUEST_BUDGET_KB,
MEMORY_TRANSCRIPT_TURNS,
PROXY_REPLICA_URLS,
unique_marker,
@ -91,11 +83,12 @@ from e2e_config import (
from e2e_metadata import Domain, Mode, Provider, Subject, meta
from lifecycle import ResourceManager
from memory_readings import RssCapture, RssReading, WorkerKey, read_rss_everywhere
from models import ChatMessage, RouterSettingsOverride, SpendLogRow
from models import ChatMessage, LiteLLMParamsBody, ReliabilityChatBody, RouterSettingsOverride
from proxy_client import ProxyClient
from reliability_support import REAL_MODEL, chat_override, create_never_benched_refusing_deployment
pytestmark = [pytest.mark.e2e, pytest.mark.quiet_stack]
pytestmark = [pytest.mark.e2e, pytest.mark.perf]
REFUSING_MODEL: Final = "openai/gpt-5.5"
DEPLOYMENTS_PER_GROUP: Final = 2
RSS_SAMPLE_CAP: Final = 4 * MEMORY_RSS_SETTLE_SAMPLES
@ -110,7 +103,6 @@ class FailedCall:
status_code: int
seconds: float
body_head: str
call_id: str | None
@dataclass(frozen=True, slots=True)
@ -123,8 +115,17 @@ class WorkerGrowth:
return self.after.ram_usage_mb - self.warm.ram_usage_mb
def _create_never_benched_refusing_deployment(proxy: ProxyClient, name: str) -> str:
return proxy.create_model(
name,
LiteLLMParamsBody(
model=REFUSING_MODEL, api_key="sk-unreachable", api_base="http://127.0.0.1:9/v1", cooldown_time=0
),
)
def _register_refusing_group(proxy: ProxyClient, resources: ResourceManager, name: str) -> None:
for model_id in tuple(create_never_benched_refusing_deployment(proxy, name) for _ in range(DEPLOYMENTS_PER_GROUP)):
for model_id in tuple(_create_never_benched_refusing_deployment(proxy, name) for _ in range(DEPLOYMENTS_PER_GROUP)):
resources.defer(lambda model_id=model_id: proxy.delete_model(model_id))
@ -141,10 +142,18 @@ TRANSCRIPT: Final = _transcript(MEMORY_TRANSCRIPT_TURNS)
def _fail_once(proxy: ProxyClient, key: str, model: str, override: RouterSettingsOverride) -> FailedCall:
started: Final = time.perf_counter()
resp: Final = chat_override(
proxy, key, model, f"memory regression {unique_marker()}", override=override, history=TRANSCRIPT
resp: Final = proxy.transport.send(
"/chat/completions",
headers=proxy.transport.bearer(key),
json=ReliabilityChatBody(
model=model,
messages=[*TRANSCRIPT, ChatMessage(role="user", content=f"memory regression {unique_marker()}")],
max_tokens=512,
router_settings_override=override,
cache={"no-cache": True},
),
)
return FailedCall(resp.status_code, time.perf_counter() - started, resp.body[:300], resp.call_id)
return FailedCall(resp.status_code, time.perf_counter() - started, resp.body[:300])
def _fail_many(proxy: ProxyClient, key: str, model: str, override: RouterSettingsOverride) -> tuple[FailedCall, ...]:
@ -259,24 +268,6 @@ def _assert_every_call_failed_through_fallback(calls: Sequence[FailedCall], fall
)
def _stored_request_kb(proxy: ProxyClient, call: FailedCall) -> float:
assert call.call_id, (
f"the failing call carried no x-litellm-call-id header, so its spend log cannot be read back: {call.body_head}"
)
rows: Final[Sequence[SpendLogRow]] = proxy.poll_logs_for_request_id(call.call_id)
assert rows, (
f"no spend log row appeared for failing call {call.call_id} within the poll window: either the stack "
"writes no spend logs or its writer dropped the row, which the v1.100.0 one did once the stored "
"request outgrew the writer's queue budget"
)
snapshot: Final = rows[0].proxy_server_request
assert snapshot, (
f"spend log {call.call_id} stored no request body, so the stack is not running with "
"general_settings.store_prompts_in_spend_logs and the stored-request check would pass vacuously"
)
return len(json.dumps(snapshot).encode()) / 1024
class TestReliabilityMemory:
@pytest.mark.covers("reliability.perf.idle_memory.under_slo")
@meta(
@ -309,13 +300,13 @@ class TestReliabilityMemory:
Subject(
domain=Domain.ROUTING,
providers=(Provider.OPENAI,),
models=(REAL_MODEL,),
models=(REFUSING_MODEL,),
mode=Mode.NONSTREAM,
)
)
def test_failing_requests_do_not_grow_rss_or_stored_request(
def test_failing_requests_do_not_grow_rss(
self,
client: ComplexityRouterClient,
proxy: ProxyClient,
resources: ResourceManager,
scoped_key: str,
record_property: Callable[[str, object], None],
@ -323,35 +314,25 @@ class TestReliabilityMemory:
marker: Final = unique_marker()
primary: Final = f"reliability-memory-{marker}"
fallback: Final = f"reliability-memory-fb-{marker}"
_register_refusing_group(client.proxy, resources, primary)
_register_refusing_group(client.proxy, resources, fallback)
_register_refusing_group(proxy, resources, primary)
_register_refusing_group(proxy, resources, fallback)
override: Final = RouterSettingsOverride(
num_retries=MEMORY_RETRIES_PER_REQUEST, fallbacks=[{primary: [fallback]}]
)
probe: Final = _fail_once(client.proxy, scoped_key, primary, override)
_assert_every_call_failed_through_fallback((probe,), fallback)
stored_kb: Final = _stored_request_kb(client.proxy, probe)
assert stored_kb <= MEMORY_STORED_REQUEST_BUDGET_KB, (
f"the spend log of one failing request stored a {stored_kb:.0f} KB request body, past the "
f"{MEMORY_STORED_REQUEST_BUDGET_KB:.0f} KB budget for a {len(TRANSCRIPT)}-message transcript with "
f"{MEMORY_RETRIES_PER_REQUEST} retries and a fallback; the retry breadcrumbs are copying the whole "
f"request into the stored snapshot the way the v1.100.0 ones did"
)
warm: Final = _settled_rss_after_batches(
client.proxy,
proxy,
scoped_key,
primary,
override,
fallback,
record_property,
_settled_rss_per_worker(client.proxy),
_settled_rss_per_worker(proxy),
1,
WARMUP_BATCHES,
)
after: Final = _settled_rss_after_batches(
client.proxy,
proxy,
scoped_key,
primary,
override,

96
tests/e2e/proxy_client.py
View file

@ -11,7 +11,7 @@ from __future__ import annotations
import os
import time
import warnings
from collections.abc import Callable, Mapping
from collections.abc import Callable, Mapping, Sequence
from dataclasses import dataclass, field, replace
from datetime import datetime
from functools import reduce
@ -155,6 +155,21 @@ class NotServableOn:
last_result: Result[ModelsListResponse] | None
@dataclass(frozen=True, slots=True)
class ModelRegistered:
model_name: str
model_id: str
@dataclass(frozen=True, slots=True)
class ModelNotRegistered:
model_name: str
reason: str
ModelRegistration = ModelRegistered | ModelNotRegistered
def await_servable(
list_models: Callable[[float], Result[ModelsListResponse]],
*,
@ -461,6 +476,7 @@ class ProxyClient:
model_servable_db_sync_seconds: float = MODEL_SERVABLE_DB_SYNC_SECONDS
model_servable_interval: float = MODEL_SERVABLE_INTERVAL
model_servable_request_timeout: float = MODEL_SERVABLE_REQUEST_TIMEOUT
settle: Callable[[float], None] = settle_propagation
def with_caller(self, caller: Caller) -> ProxyClient:
return replace(self, caller=caller)
@ -694,31 +710,59 @@ class ProxyClient:
balancer address is configured (every request opens a fresh connection, so
the caller's next request re-rolls), so waiting out PROPAGATION_TIMEOUT is
what makes the model safe to use anywhere."""
model_id = unwrap(
self.transport.post(
"/model/new",
headers=self.management_headers(),
json=body.model_copy(update={"litellm_params": route_cache_model(
body.litellm_params, provider_edge_base,
enabled=os.environ.get("E2E_PROVIDER_CACHE", "0") == "1" and not provider_live,
mode=body.model_info.mode,
)}),
response_type=ModelNewResponse,
)
).model_id
model_id = unwrap(self._post_model(body, provider_live=provider_live)).model_id
written_at = time.monotonic()
try:
self._await_model_servable(body.model_name, listed_for)
except BaseException:
self.delete_model(model_id)
raise
settle_propagation(written_at)
self.settle(written_at)
return model_id
@step("Add a batch of deployments, then wait once for all of them to propagate")
def register_models(
self, bodies: Sequence[ModelNewBody], *, provider_live: bool = False
) -> tuple[ModelRegistration, ...]:
posted: Final = tuple((body.model_name, self._post_model(body, provider_live=provider_live)) for body in bodies)
written_at: Final = time.monotonic()
registrations: Final = tuple(self._servable_registration(name, result) for name, result in posted)
self.settle(written_at)
return registrations
def _post_model(self, body: ModelNewBody, *, provider_live: bool) -> Result[ModelNewResponse]:
return self.transport.post(
"/model/new",
headers=self.management_headers(),
json=body.model_copy(update={"litellm_params": route_cache_model(
body.litellm_params, provider_edge_base,
enabled=os.environ.get("E2E_PROVIDER_CACHE", "0") == "1" and not provider_live,
mode=body.model_info.mode,
)}),
response_type=ModelNewResponse,
)
def _servable_registration(self, model_name: str, posted: Result[ModelNewResponse]) -> ModelRegistration:
if not isinstance(posted, Success):
return ModelNotRegistered(model_name=model_name, reason=f"/model/new failed: {posted!r}")
match self._servable_outcome(model_name, None):
case Servable():
return ModelRegistered(model_name=model_name, model_id=posted.data.model_id)
case NotServableOn() as missing:
self.delete_model(posted.data.model_id)
return ModelNotRegistered(model_name=model_name, reason=self._not_servable_message(model_name, missing))
def _await_model_servable(self, model_name: str, listed_for: str | None = None) -> None:
"""Block until every replica lists `model_name`, or fail at model_servable_timeout."""
match self._servable_outcome(model_name, listed_for):
case Servable():
return
case NotServableOn() as missing:
raise AssertionError(self._not_servable_message(model_name, missing))
def _servable_outcome(self, model_name: str, listed_for: str | None) -> Servable | NotServableOn:
headers: Final = self.management_headers(listed_for)
outcome: Final = await_servable_everywhere(
return await_servable_everywhere(
{url: self._models_poller(transport, headers) for url, transport in self.replicas.items()},
model_name=model_name,
timeout=self.model_servable_timeout,
@ -728,19 +772,15 @@ class ProxyClient:
now=time.monotonic,
sleep=time.sleep,
)
match outcome:
case Servable():
return
case NotServableOn(replica=replica, last_result=last_result):
raise AssertionError(
servable_timeout_message(
model_name=model_name,
replica=replica,
timeout=self.model_servable_timeout,
db_sync_seconds=self.model_servable_db_sync_seconds,
last_result=last_result,
)
)
def _not_servable_message(self, model_name: str, missing: NotServableOn) -> str:
return servable_timeout_message(
model_name=model_name,
replica=missing.replica,
timeout=self.model_servable_timeout,
db_sync_seconds=self.model_servable_db_sync_seconds,
last_result=missing.last_result,
)
@staticmethod
def _models_poller(transport: Transport, headers: AuthHeaders) -> ModelsPoller:

3
tests/e2e/pytest.ini
View file

@ -13,8 +13,6 @@ markers =
prompt_caching_stack: needs a proxy running with router_settings.optional_pre_call_checks including prompt_caching; deselected unless E2E_PROMPT_CACHING_STACK is set
cli_determinism: drives the real claude CLI for several seconds; deselected unless E2E_CLI_DETERMINISM is set
redis_chaos: load test that pauses the proxy's Redis outright mid-run; needs a proxy booted from gateway/redis_chaos_ci_config.yml on the same host, and is deselected unless E2E_PERF_TESTS is set
quiet_stack: measures the proxy itself, so it runs while no other test on this host is hitting the stack; every other test waits for it to finish
mcp_oauth_live: real Linear OAuth consent via a captured browser session; deselected unless E2E_MCP_OAUTH_LIVE is set
provider_edge_host: routes provider traffic through the pytest host's edge in every fixture mode, so the gateway must reach the pytest host; deselected unless E2E_PROVIDER_EDGE_HOST_REACHABLE is set
owned_gateway: boots its own proxy from source against the stack's Postgres, so it needs DATABASE_URL on the pytest host; deselected unless E2E_OWNED_GATEWAY is set
otel_v2: needs a proxy running with LITELLM_OTEL_V2=true; deselected unless E2E_OTEL_V2 is set
@ -22,3 +20,4 @@ markers =
secret_manager: needs a proxy booted from gateway/secret_manager_<system>_ci_config.yml against that live secret manager; deselected unless E2E_SECRET_MANAGER names the backend (see secret_manager/secret_backends.py)
sagemaker_nova: needs a deployed SageMaker Nova endpoint the proxy's AWS credentials can invoke; deselected unless SAGEMAKER_NOVA_ENDPOINT names it
sagemaker_nova2_lite: needs a deployed SageMaker Nova 2 Lite endpoint the proxy's AWS credentials can invoke; deselected unless SAGEMAKER_NOVA2_LITE_ENDPOINT names it
proxy_wide_config: turns on a proxy-wide setting through /config/update, so it runs on a stack no other suite is driving; deselected unless E2E_PROXY_WIDE_CONFIG_STACK is set

152
tests/e2e/quota_management/spend_tracking/test_pricing_config_spend_e2e.py → tests/e2e/quota_management/spend_tracking/test_guardrail_cost_spend_e2e.py
View file

@ -17,11 +17,6 @@ from models import (
ChatBody,
ChatMessage,
ChatResponse,
ConfigPatchResponse,
CostDiscountConfig,
CostDiscountConfigResponse,
CostMarginConfig,
CostMarginConfigResponse,
GuardrailCreateBody,
GuardrailCreateResponse,
KeyDeleteBody,
@ -31,14 +26,11 @@ from models import (
)
from spend_e2e_client import SpendClient
pytestmark = [pytest.mark.e2e, pytest.mark.quiet_stack]
pytestmark = pytest.mark.e2e
BACKEND_MODEL: Final = "openai/gpt-6.1-sol"
INPUT_RATE: Final = 0.00005
OUTPUT_RATE: Final = 0.0001
DISCOUNT: Final = 0.25
MARGIN_PERCENT: Final = 0.1
MARGIN_FIXED: Final = 0.0005
def _register_bedrock_guardrail(
@ -92,52 +84,6 @@ def _guarded_chat(client: SpendClient, key: str, model: str, name: str) -> Strea
)
def _set_discount(client: SpendClient, values: dict[str, float]) -> float:
unwrap(
client.proxy.transport.patch(
"/config/cost_discount_config",
headers=client.proxy.transport.master,
json=CostDiscountConfig(values),
response_type=ConfigPatchResponse,
)
)
return time.monotonic()
def _set_margin(client: SpendClient, values: dict[str, float | dict[str, float]]) -> float:
unwrap(
client.proxy.transport.patch(
"/config/cost_margin_config",
headers=client.proxy.transport.master,
json=CostMarginConfig(values),
response_type=ConfigPatchResponse,
)
)
return time.monotonic()
def _get_discount(client: SpendClient) -> dict[str, float]:
return unwrap(
client.proxy.transport.get(
"/config/cost_discount_config",
headers=client.proxy.transport.master,
params=NoBody(),
response_type=CostDiscountConfigResponse,
)
).values
def _get_margin(client: SpendClient) -> dict[str, float | dict[str, float]]:
return unwrap(
client.proxy.transport.get(
"/config/cost_margin_config",
headers=client.proxy.transport.master,
params=NoBody(),
response_type=CostMarginConfigResponse,
)
).values
def _register_model(client: SpendClient, resources: ResourceManager, prefix: str) -> str:
model: Final = f"{prefix}-{unique_marker()}"
model_id: Final = client.proxy.create_model(
@ -199,100 +145,7 @@ def scoped_key(client: SpendClient, strict_resources: ResourceManager) -> str:
return key
@pytest.fixture
def restored_pricing_config(client: SpendClient) -> Iterator[None]:
discount: Final = _get_discount(client)
margin: Final = _get_margin(client)
try:
_set_discount(client, {})
reset_at: Final = _set_margin(client, {})
settle_propagation(reset_at)
yield
finally:
try:
_set_discount(client, discount)
finally:
restored_at: Final = _set_margin(client, margin)
settle_propagation(restored_at)
class TestPricingConfigSpend:
@pytest.mark.covers(
"quota_management.spend_tracking.discount_config.logs_cost",
exercised_on=["chat_completions"],
)
@meta(
Subject(
domain=Domain.SPEND_BUDGETS,
route=Route.SPEND_REPORTING,
providers=(Provider.OPENAI,),
models=(BACKEND_MODEL,),
mode=Mode.NONSTREAM,
)
)
def test_configured_discount_reaches_persisted_spend_row(
self,
client: SpendClient,
strict_resources: ResourceManager,
scoped_key: str,
restored_pricing_config: None,
) -> None:
discount_set_at: Final = _set_discount(client, {"openai": DISCOUNT})
settle_propagation(discount_set_at)
model: Final = _register_model(client, strict_resources, "discount-priced")
chat: Final = unwrap(client.chat(scoped_key, model, f"reply with one word {unique_marker()}", max_tokens=16))
assert chat.id and chat.usage
assert chat.usage.prompt_tokens is not None and chat.usage.prompt_tokens > 0
assert chat.usage.completion_tokens is not None and chat.usage.completion_tokens >= 0
row: Final = poll_cost_row(client.proxy, chat.id)
assert row is not None
base_cost: Final = _base_cost(row, chat.usage.prompt_tokens, chat.usage.completion_tokens)
breakdown: Final = row.breakdown
assert breakdown.discount_percent is not None and approx_equal(breakdown.discount_percent, DISCOUNT)
assert breakdown.discount_amount is not None and approx_equal(breakdown.discount_amount, base_cost * DISCOUNT)
assert row.spend is not None and approx_equal(row.spend, base_cost * (1 - DISCOUNT))
@pytest.mark.covers(
"quota_management.spend_tracking.margin_config.logs_cost",
exercised_on=["chat_completions"],
)
@meta(
Subject(
domain=Domain.SPEND_BUDGETS,
route=Route.SPEND_REPORTING,
providers=(Provider.OPENAI,),
models=(BACKEND_MODEL,),
mode=Mode.NONSTREAM,
)
)
def test_configured_margin_reaches_persisted_spend_row(
self,
client: SpendClient,
strict_resources: ResourceManager,
scoped_key: str,
restored_pricing_config: None,
) -> None:
margin_set_at: Final = _set_margin(
client, {"openai": {"percentage": MARGIN_PERCENT, "fixed_amount": MARGIN_FIXED}}
)
settle_propagation(margin_set_at)
model: Final = _register_model(client, strict_resources, "margin-priced")
chat: Final = unwrap(client.chat(scoped_key, model, f"reply with one word {unique_marker()}", max_tokens=16))
assert chat.id and chat.usage
assert chat.usage.prompt_tokens is not None and chat.usage.prompt_tokens > 0
assert chat.usage.completion_tokens is not None and chat.usage.completion_tokens >= 0
row: Final = poll_cost_row(client.proxy, chat.id)
assert row is not None
base_cost: Final = _base_cost(row, chat.usage.prompt_tokens, chat.usage.completion_tokens)
breakdown: Final = row.breakdown
expected_margin: Final = base_cost * MARGIN_PERCENT + MARGIN_FIXED
assert breakdown.margin_percent is not None and approx_equal(breakdown.margin_percent, MARGIN_PERCENT)
assert breakdown.margin_fixed_amount is not None and approx_equal(breakdown.margin_fixed_amount, MARGIN_FIXED)
assert breakdown.margin_total_amount is not None and approx_equal(breakdown.margin_total_amount, expected_margin)
assert row.spend is not None and approx_equal(row.spend, base_cost + expected_margin)
class TestGuardrailCostSpend:
@pytest.mark.covers(
"quota_management.spend_tracking.guardrail_cost.logs_cost",
exercised_on=["chat_completions"],
@ -311,7 +164,6 @@ class TestPricingConfigSpend:
client: SpendClient,
strict_resources: ResourceManager,
scoped_key: str,
restored_pricing_config: None,
) -> None:
identifier: Final = os.environ["BEDROCK_GUARDRAIL_IDENTIFIER"]
version: Final = os.environ["BEDROCK_GUARDRAIL_VERSION"]

45
tests/e2e/stack_lock.py
View file

@ -1,45 +0,0 @@
"""Cross-process reader/writer lock over the proxy stack every xdist worker shares.
Every collected test holds it shared, marker or not, since the Claude Code cells and
other unmarked suites drive the same stack; a `quiet_stack` test holds it exclusive,
and the `gate` file makes a waiting exclusive holder win over readers that arrive
after it."""
from __future__ import annotations
import fcntl
import hashlib
import tempfile
from collections.abc import Generator
from contextlib import ExitStack, contextmanager
from pathlib import Path
from typing import Final
from e2e_config import PROXY_BASE_URL
STACK_DIGEST: Final = hashlib.sha256(PROXY_BASE_URL.encode()).hexdigest()[:12]
LOCK_DIR: Final = Path(tempfile.gettempdir()) / f"litellm-e2e-stack-{STACK_DIGEST}"
GATE_FILE: Final = LOCK_DIR / "gate"
STACK_FILE: Final = LOCK_DIR / "stack"
@contextmanager
def _flock(path: Path, operation: int) -> Generator[None]:
with path.open("a") as handle:
fcntl.flock(handle, operation)
try:
yield
finally:
fcntl.flock(handle, fcntl.LOCK_UN)
@contextmanager
def stack_lock(exclusive: bool) -> Generator[None]:
LOCK_DIR.mkdir(parents=True, exist_ok=True)
if exclusive:
with _flock(GATE_FILE, fcntl.LOCK_EX), _flock(STACK_FILE, fcntl.LOCK_EX):
yield
return
with ExitStack() as held:
with _flock(GATE_FILE, fcntl.LOCK_SH):
held.enter_context(_flock(STACK_FILE, fcntl.LOCK_SH))
yield

2
tests/e2e_harness/AGENTS.md
View file

@ -1,6 +1,6 @@
# e2e harness tests
Tests of the harness under `tests/e2e/` (the transport, the clients, the fixture bundle and replay edge, the stack lock, the IdP launcher, the coverage collector, the JUnit properties and the Claude Code driver), not of the product. They live outside `tests/e2e/` because the Buildkite e2e run copies that folder into the runner image and runs every test in it, so a harness test in there counts as a product test in the nightly numbers. Nothing here needs a proxy, provider keys or the network
Tests of the harness under `tests/e2e/` (the transport, the clients, the fixture bundle and replay edge, the IdP launcher, the coverage collector, the JUnit properties and the Claude Code driver), not of the product. They live outside `tests/e2e/` because the Buildkite e2e run copies that folder into the runner image and runs every test in it, so a harness test in there counts as a product test in the nightly numbers. Nothing here needs a proxy, provider keys or the network
The layout mirrors `tests/e2e/`: `test_e2e_http.py` covers `tests/e2e/e2e_http.py`, `logging/test_datadog_reader.py` covers `tests/e2e/logging/datadog_reader.py`, and `claude_code/` covers the driver, builder, probe and version resolver. Put a new harness test under the folder that mirrors the suite folder whose module it covers

90
tests/e2e_harness/test_proxy_client.py
View file

@ -12,10 +12,11 @@ injected clock, so nothing here monkeypatches anything.
from __future__ import annotations
import json
import time
from builtins import ExceptionGroup
from collections.abc import Callable, Generator, Iterable, Mapping
from contextlib import contextmanager
from dataclasses import dataclass
from dataclasses import dataclass, replace
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from itertools import chain, repeat
from queue import SimpleQueue
@ -40,7 +41,10 @@ from models import (
LiteLLMParamsBody,
McpServerCreateBody,
McpServerUpdateBody,
ModelDeleteBody,
ModelInfoBody,
ModelListEntry,
ModelNewBody,
ModelsListResponse,
OrgNewBody,
OrgUpdateBody,
@ -59,6 +63,8 @@ from proxy_client import (
ConvergeOutcome,
CredentialKind,
EverywhereConverged,
ModelNotRegistered,
ModelRegistered,
ModelsPoller,
NeverConvergedOn,
NotConverged,
@ -691,3 +697,85 @@ def test_partial_updates_preserve_explicit_null_at_the_http_boundary(operation:
)
assert json.loads(bodies.get_nowait()) == expected
assert bodies.empty()
@contextmanager
def model_registry(
listed: frozenset[str], refused: frozenset[str]
) -> Generator[tuple[ProxyClient, SimpleQueue[float], SimpleQueue[str]]]:
writes: Final[SimpleQueue[float]] = SimpleQueue()
deleted: Final[SimpleQueue[str]] = SimpleQueue()
class Handler(BaseHTTPRequestHandler):
def log_message(self, format: str, *args: object) -> None:
pass
def reply(self, status: int, body: object) -> None:
self.send_response(status)
self.end_headers()
self.wfile.write(json.dumps(body).encode())
def do_GET(self) -> None:
self.reply(200, {"data": [{"id": name} for name in sorted(listed)]})
def do_POST(self) -> None:
raw: Final = self.rfile.read(int(self.headers.get("Content-Length", "0")))
if self.path == "/model/delete":
deleted.put(ModelDeleteBody.model_validate_json(raw).id)
self.reply(200, {})
return
writes.put(time.monotonic())
name: Final = ModelNewBody.model_validate_json(raw).model_name
if name in refused:
self.reply(500, {"error": "refused"})
return
self.reply(200, {"model_id": f"id-{name}"})
server: Final = ThreadingHTTPServer(("127.0.0.1", 0), Handler)
thread: Final = Thread(target=lambda: server.serve_forever(poll_interval=0.01), daemon=True)
thread.start()
url: Final = f"http://127.0.0.1:{server.server_port}"
proxy: Final = build_proxy_client(
base_url=url, control_plane_base_url=url, replica_urls=(url,), control_replica_urls=(url,), master_key="bootstrap"
)
try:
yield replace(proxy, model_servable_timeout=0.3, model_servable_interval=0.01), writes, deleted
finally:
server.shutdown()
server.server_close()
thread.join(timeout=5)
class TestRegisterModels:
def test_a_batch_waits_once_after_its_last_write_and_reports_each_model(self) -> None:
names: Final = ("first", "refused", "ghost", "last")
settled: Final[SimpleQueue[float]] = SimpleQueue()
with (
model_registry(listed=frozenset({"first", "last"}), refused=frozenset({"refused"})) as (
proxy,
writes,
deleted,
),
without_retries(),
):
registrations: Final = replace(proxy, settle=settled.put).register_models(
tuple(
ModelNewBody(
model_name=name, litellm_params=LiteLLMParamsBody(model="synthetic"), model_info=ModelInfoBody()
)
for name in names
)
)
last_write: Final = max(writes.get_nowait() for _ in names)
assert writes.empty()
assert settled.get_nowait() >= last_write
assert settled.empty(), "a batch must wait out propagation once, not once per model"
assert tuple((type(r), r.model_name) for r in registrations) == (
(ModelRegistered, "first"),
(ModelNotRegistered, "refused"),
(ModelNotRegistered, "ghost"),
(ModelRegistered, "last"),
)
assert tuple(r.model_id for r in registrations if isinstance(r, ModelRegistered)) == ("id-first", "id-last")
assert deleted.get_nowait() == "id-ghost"
assert deleted.empty()

117
tests/e2e_harness/test_stack_lock.py
View file

@ -1,117 +0,0 @@
"""Cross-process behavior of the stack lock: readers share it, an exclusive holder waits for
every reader and keeps them out, and a reader arriving behind a waiting exclusive holder
queues behind it instead of starving it."""
from __future__ import annotations
import fcntl
import inspect
import os
import subprocess
import sys
import time
from contextlib import ExitStack
from pathlib import Path
from typing import Final
import pytest
import stack_lock
HARNESS_DIR: Final = Path(inspect.getfile(stack_lock)).resolve().parent
DEADLINE_SECONDS: Final = 30.0
SETTLE_SECONDS: Final = 0.5
HOLDER_SCRIPT: Final = """
import sys, time
from pathlib import Path
from stack_lock import stack_lock
name, mode, release_path, log_path = sys.argv[1:]
def record(event):
with Path(log_path).open("a") as log:
log.write(f"{name} {event}\\n")
record("waiting")
with stack_lock(exclusive=mode == "exclusive"):
record("enter")
while not Path(release_path).exists():
time.sleep(0.02)
record("exit")
"""
def _events(log_path: Path) -> tuple[str, ...]:
return tuple(log_path.read_text().splitlines()) if log_path.exists() else ()
def _wait_for_event(log_path: Path, event: str) -> None:
deadline: Final = time.monotonic() + DEADLINE_SECONDS
while event not in _events(log_path):
if time.monotonic() > deadline:
pytest.fail(f"{event!r} never appeared; events so far: {_events(log_path)}")
time.sleep(0.02)
def _wait_until_gate_is_held_exclusively(gate_path: Path) -> None:
deadline: Final = time.monotonic() + DEADLINE_SECONDS
with gate_path.open("a") as handle:
while True:
try:
fcntl.flock(handle, fcntl.LOCK_SH | fcntl.LOCK_NB)
except BlockingIOError:
return
fcntl.flock(handle, fcntl.LOCK_UN)
if time.monotonic() > deadline:
pytest.fail("no exclusive holder ever took the gate")
time.sleep(0.02)
def _start_holder(held: ExitStack, tmp_path: Path, name: str, mode: str) -> subprocess.Popen[bytes]:
holder: Final = held.enter_context(
subprocess.Popen(
(
sys.executable,
"-P",
"-c",
HOLDER_SCRIPT,
name,
mode,
str(tmp_path / f"release-{name}"),
str(tmp_path / "events"),
),
cwd=HARNESS_DIR,
env={**os.environ, "TMPDIR": str(tmp_path), "PYTHONPATH": str(HARNESS_DIR)},
)
)
held.callback(holder.kill)
return holder
def test_readers_share_exclusive_waits_and_a_waiting_exclusive_beats_later_readers(tmp_path: Path) -> None:
lock_dir: Final = tmp_path / f"litellm-e2e-stack-{stack_lock.STACK_DIGEST}"
lock_dir.mkdir()
log_path: Final = tmp_path / "events"
with ExitStack() as held:
first_reader: Final = _start_holder(held, tmp_path, "A", "shared")
_wait_for_event(log_path, "A enter")
second_reader: Final = _start_holder(held, tmp_path, "R", "shared")
_wait_for_event(log_path, "R enter")
(tmp_path / "release-R").touch()
_wait_for_event(log_path, "R exit")
writer: Final = _start_holder(held, tmp_path, "W", "exclusive")
_wait_until_gate_is_held_exclusively(lock_dir / "gate")
late_reader: Final = _start_holder(held, tmp_path, "B", "shared")
_wait_for_event(log_path, "B waiting")
time.sleep(SETTLE_SECONDS)
(tmp_path / "release-A").touch()
_wait_for_event(log_path, "W enter")
(tmp_path / "release-W").touch()
_wait_for_event(log_path, "B enter")
(tmp_path / "release-B").touch()
for holder in (first_reader, second_reader, writer, late_reader):
assert holder.wait(timeout=DEADLINE_SECONDS) == 0
events: Final = _events(log_path)
assert events.index("R enter") < events.index("A exit")
assert events.index("W enter") > events.index("A exit")
assert events.index("B enter") > events.index("W exit")

14
tests/integration/_support/client.py
View file

@ -95,6 +95,10 @@ class Gateway:
client: httpx.Client
key: str
upstream_url: str
database_url: str | None = None
def rows(self, query: str, parameters: tuple[str, ...]) -> list[dict[str, JsonValue]]:
return read_rows(query, parameters, database_url=self.database_url)
def request(
self,
@ -267,12 +271,14 @@ class Scenario:
)
assert read_rows('SELECT model_id FROM "LiteLLM_ProxyModelTable" WHERE model_id = %s', (identity,)) == []
def model(self, *, model_info: Mapping[str, JsonValue] | None = None, **parameters: JsonValue) -> str:
name: Final = f"integration-{uuid.uuid4().hex}"
def model(
self, *, name: str | None = None, model_info: Mapping[str, JsonValue] | None = None, **parameters: JsonValue
) -> str:
model_name: Final = name if name is not None else f"integration-{uuid.uuid4().hex}"
created: Final = self.gateway.post(
"/model/new",
{
"model_name": name,
"model_name": model_name,
"litellm_params": {
"model": "openai/gpt-4o-mini",
"api_key": "integration-provider-key",
@ -284,7 +290,7 @@ class Scenario:
)
identity: Final = string_value(object_value(created["model_info"])["id"])
self.cleanups.callback(self.delete_model, identity)
return name
return model_name
@contextmanager

6
tests/integration/_support/mcp.py
View file

@ -282,6 +282,7 @@ class ScriptedTool:
respond: Callable[[JsonRpc], Reply | JsonRpc]
description: str | Callable[[Mapping[str, str]], str] | None = None
input_schema: JsonRpc = field(default_factory=lambda: {"type": "object"})
event_stream: bool = False
def listing(self, headers: Mapping[str, str]) -> JsonRpc:
described: Final = self.description(headers) if callable(self.description) else self.description
@ -334,7 +335,10 @@ def scripted_peer(*tools: ScriptedTool) -> Iterator[McpPeer]:
if tool is None:
return jsonrpc_error(identity, -32602, "unknown tool")
produced: Final = tool.respond(body["params"])
return produced if isinstance(produced, Reply) else jsonrpc_reply(identity, produced)
reply: Final = produced if isinstance(produced, Reply) else jsonrpc_reply(identity, produced)
if not tool.event_stream:
return reply
return Reply(content_type="text/event-stream", body=b"event: message\ndata: " + reply.body + b"\n\n")
with wire_server(provider) as wire:
yield McpPeer(wire.url + "/mcp", observed)

55
tests/integration/mcp/test_mcp_llm_endpoints.py
View file

@ -468,6 +468,61 @@ def test_server_scoped_gateway_url_exposes_only_that_servers_tools(gateway: Gate
assert [call["body"]["params"]["name"] for call in _peer_add_calls(rig.peer)] == ["add"]
@pytest.mark.parametrize("surface", ("chat", "responses", "messages"))
def test_proxy_path_url_of_a_served_server_is_listed_and_executed_by_the_gateway(
gateway: Gateway, surface: Surface
) -> None:
with _rig(gateway, surface) as rig:
server_url: Final = f"{str(gateway.client.base_url).rstrip('/')}/mcp/{rig.alias}"
response: Final = rig.send(_granted_key(rig), [{**AUTO, "server_url": server_url}])
assert response.status_code == 200, response.text
requests: Final = rig.upstream_tools()
assert len(requests) == 2 and all(rig.tool in names for names in requests), requests
peer_calls: Final = tuple(JSON_OBJECT.validate_python(call["body"]) for call in _peer_add_calls(rig.peer))
assert [object_value(call["params"])["name"] for call in peer_calls] == ["add"], peer_calls
assert rig.final_text(JSON_OBJECT.validate_json(response.content)) == ANSWER, response.text
LOOKUP_PARAMETERS: Final = {"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]}
CALLER_FUNCTION_TOOLS: Final[Mapping[Surface, Mapping[str, object]]] = {
"chat": {"type": "function", "function": {"name": "lookup", "parameters": LOOKUP_PARAMETERS}},
"responses": {"type": "function", "name": "lookup", "parameters": LOOKUP_PARAMETERS},
"messages": {"name": "lookup", "input_schema": LOOKUP_PARAMETERS},
}
@pytest.mark.parametrize("surface", ("chat", "responses", "messages"))
def test_tools_that_do_not_name_the_gateway_reach_the_provider_unchanged_and_never_the_peer(
gateway: Gateway, surface: Surface
) -> None:
with _rig(gateway, surface) as rig:
remote_mcp_tools: Final = [
{
"type": "mcp",
"server_label": "docs",
"server_url": "https://mcp.example.com/sse",
"require_approval": "never",
},
{
"type": "mcp",
"server_label": "repos",
"server_url": f"https://mcp.example.com/mcp/unserved{uuid.uuid4().hex[:8]}",
"require_approval": "never",
},
]
response: Final = rig.send(_granted_key(rig), [CALLER_FUNCTION_TOOLS[surface], *remote_mcp_tools])
assert response.status_code == 200, response.text
bodies: Final = tuple(
JSON_OBJECT.validate_json(request.body) for request in rig.wire.drain() if request.method == "POST"
)
assert len(bodies) == 1, bodies
forwarded: Final = bodies[0].get("tools")
assert isinstance(forwarded, list) and len(forwarded) == 3, forwarded
assert [tool for tool in forwarded if isinstance(tool, dict) and tool.get("type") == "mcp"] == remote_mcp_tools
assert "lookup" in json.dumps([tool for tool in forwarded if tool not in remote_mcp_tools]), forwarded
assert tool_calls(rig.peer.drain()) == (), "a tool that does not name the gateway reached the peer"
def test_streaming_chat_executes_the_tool_once_and_streams_the_follow_up(gateway: Gateway) -> None:
with _rig(gateway, "chat") as rig:
key: Final = _granted_key(rig)

48
tests/integration/mcp/test_mcp_oauth_flows.py
View file

@ -697,6 +697,54 @@ def test_per_user_authorization_code_with_pkce_binds_the_token_to_the_authorizin
)
@pytest.mark.timeout(600)
def test_per_user_authorization_code_token_is_read_back_from_the_database_by_a_fresh_proxy_without_redis(
gateway: Gateway, tmp_path: Path
) -> None:
with mcp_peer() as peer, oauth_server() as auth, gateway.scenario() as scenario:
alias: Final = "cold" + uuid.uuid4().hex[:8]
identity: Final = _register_oauth(
scenario,
peer,
auth,
alias,
auth_type="oauth2",
oauth2_flow="authorization_code",
credentials={"client_id": "cold-client", "client_secret": "cold-secret", "scopes": ["tools.call"]},
)
owner_id: Final = scenario.user()
owner: Final = scenario.key(user_id=owner_id, object_permission={"mcp_servers": [identity]})
registered: Final = gateway.client.post(
f"/{alias}/register", json={"redirect_uris": [CLIENT_REDIRECT], "client_name": "integration"}
)
assert registered.status_code in (200, 201), registered.text
client_id: Final = str(_response_object(registered)["client_id"])
pkce: Final = _Pkce(secrets.token_urlsafe(32))
code, _ = _authorize_through_gateway(gateway, auth, alias, owner, client_id, pkce)
redeemed: Final = _redeem(gateway, alias, owner, client_id, code, pkce)
assert redeemed.status_code == 200, redeemed.text
token: Final = _issued_token(redeemed.json())
stored: Final = read_rows(
'SELECT credential_b64 FROM "LiteLLM_MCPUserCredentials" WHERE user_id = %s AND server_id = %s',
(owner_id, identity),
)
assert len(stored) == 1, stored
stored_bytes: Final = base64.urlsafe_b64decode(str(stored[0]["credential_b64"]))
assert token.encode() not in stored_bytes, "stored token is plaintext"
config: Final = tmp_path / "no_redis.yaml"
config.write_text(
"general_settings:\n"
" master_key: os.environ/LITELLM_MASTER_KEY\n"
" database_url: os.environ/DATABASE_URL\n"
" store_model_in_db: true\n"
)
peer.drain()
with owned_proxy(gateway, tmp_path, {}, config=config, remove_environment=("REDIS_HOST", "REDIS_PORT")) as cold:
called: Final = call_tool(cold, owner, identity, f"{alias}-add", ADD)
assert called.status_code == 200, called.text
assert _authorizations(peer) == (f"Bearer {token}".encode(),)
def test_authorization_request_without_pkce_is_refused_before_reaching_the_authorization_server(
gateway: Gateway,
) -> None:

78
tests/integration/mcp/test_mcp_transports.py
View file

@ -2,10 +2,12 @@ import builtins
import json
import uuid
from collections.abc import Mapping
from pathlib import Path
from typing import Final
import httpx
import pytest
import yaml
from mcp.shared.exceptions import MCPError
from integration._support.client import Gateway
from integration._support.database import read_rows
@ -17,15 +19,21 @@ from integration._support.mcp import (
McpCaller,
Outcome,
PeerKind,
ScriptedTool,
_outcome_from_rpc,
_parse_rpc_body,
call_tool,
listed_tools,
mcp_peer,
official_client_outcomes,
peer_of,
register_mcp,
scripted_peer,
text_result,
tool_calls,
)
from integration._support.mcp_grants import create_toolset
from integration._support.process import owned_proxy
ADD: Final = {"http": "add", "sse": "add", "stdio": "add", "openapi": "getpet"}
ARGUMENTS: Final = {"add": {"a": 3, "b": 4}, "getpet": {"petId": "7"}}
@ -670,7 +678,7 @@ def test_sdk_progress_and_cancellation_reach_upstream_without_losing_session(
from starlette.routing import Mount
from integration._support.asgi import asgi_server
from integration._support.mcp import _drain_sse_streams
from integration._support.mcp import _acknowledge_notifications_after_dispatch, _drain_sse_streams
from integration._support.process import owned_proxy
upstream: Final = MCPServer("progress-cancellation-peer")
@ -762,7 +770,9 @@ def test_sdk_progress_and_cancellation_reach_upstream_without_losing_session(
watchers = tuple(task for task in asyncio.all_tasks() if "_shutdown_watcher" in repr(task.get_coro()))
await asyncio.wait_for(asyncio.gather(*watchers), 2)
hosted: Final = Starlette(routes=[Mount("/", app=app)], lifespan=lifespan)
hosted: Final = Starlette(
routes=[Mount("/", app=_acknowledge_notifications_after_dispatch(app))], lifespan=lifespan
)
with asgi_server(hosted) as upstream_url, httpx.Client() as client:
config: Final = tmp_path / "proxy.yaml"
config.write_text(
@ -817,3 +827,67 @@ def test_sdk_progress_and_cancellation_reach_upstream_without_losing_session(
),
) as gateway:
asyncio.run(exercise(gateway))
def test_tool_result_larger_than_one_mib_in_a_single_sse_event_reaches_the_caller(gateway: Gateway) -> None:
payload: Final = "x" * (3 * 1024 * 1024)
tool: Final = ScriptedTool("dump", lambda _: text_result(payload), event_stream=True)
with scripted_peer(tool) as peer, gateway.scenario() as scenario:
alias: Final = "bigsse" + uuid.uuid4().hex[:8]
identity: Final = register_mcp(scenario, peer, alias)
key: Final = scenario.key(object_permission={"mcp_servers": [identity]})
name: Final = next(full for full in listed_tools(gateway, key, identity) if full.endswith("dump"))
response: Final = call_tool(gateway, key, identity, name, {})
assert response.status_code == 200, response.text[:500]
body: Final = response.json()
assert body["isError"] is False, response.text[:500]
assert body["content"][0]["text"] == payload
def test_configured_sse_limit_rejects_large_results_and_allows_the_next_call(tmp_path: Path) -> None:
limit: Final = 4096
payload: Final = "x" * 512
small: Final = ScriptedTool("small", lambda _: text_result(payload), event_stream=True)
large: Final = ScriptedTool("large", lambda _: text_result("x" * (2 * limit)), event_stream=True)
key: Final = "sk-local-mcp-sse-limit-regression-key"
with scripted_peer(small, large) as peer, httpx.Client() as client:
config: Final = tmp_path / "sse-limit.yaml"
config.write_text(
yaml.safe_dump(
{
"general_settings": {"master_key": key, "store_model_in_db": False},
"mcp_servers": {"limited": peer.registration()},
}
)
)
seed: Final = Gateway(client, key, peer.url)
with owned_proxy(
seed,
tmp_path / "proxy",
{
"STORE_MODEL_IN_DB": "False",
"DISABLE_SCHEMA_UPDATE": "true",
"LITELLM_LOCAL_MODEL_COST_MAP": "True",
"LITELLM_MCP_MAX_SSE_EVENT_SIZE_BYTES": str(limit),
},
config=config,
database_setup=(),
remove_environment=(
"DATABASE_URL",
"DATABASE_URL_READ_REPLICA",
"DIRECT_URL",
"INTEGRATION_PROXY_DATABASE_URL",
"INTEGRATION_PROXY_READ_REPLICA_URL",
),
) as gateway:
assert set(listed_tools(gateway, key)) == {"small", "large"}
for name in ("small", "large", "small"):
response: Final = call_tool(gateway, key, "limited", name, {})
assert response.status_code == 200, response.text[:500]
body: Final = response.json()
assert body["isError"] is (name == "large"), response.text[:500]
if name == "large":
assert str(limit) in body["content"][0]["text"], response.text[:500]
else:
assert body["content"][0]["text"] == payload, response.text[:500]

32
tests/integration/pricing/test_cost_margin_discount.py
View file

@ -7,7 +7,7 @@ import httpx
import pytest
from integration._support.client import Gateway, eventually, object_value, string_value
from integration._support.database import read_rows
from pydantic import BaseModel, JsonValue
from pydantic import BaseModel, ConfigDict, JsonValue
class CostEstimate(BaseModel):
@ -110,6 +110,25 @@ def _assert_logged_spends(expectations: tuple[tuple[str, float], ...]) -> None:
assert actual == pytest.approx(dict(expectations)), rows
class LoggedBreakdown(BaseModel):
model_config = ConfigDict(frozen=True)
original_cost: float
discount_percent: float | None = None
discount_amount: float | None = None
margin_percent: float | None = None
margin_fixed_amount: float | None = None
margin_total_amount: float | None = None
def _logged_breakdown(request_id: str) -> LoggedBreakdown:
rows: Final = read_rows('SELECT metadata FROM "LiteLLM_SpendLogs" WHERE request_id = %s', (request_id,))
assert len(rows) == 1, rows
metadata: Final = rows[0]["metadata"]
parsed: Final = object_value(json.loads(metadata) if isinstance(metadata, str) else metadata)
return LoggedBreakdown.model_validate(parsed["cost_breakdown"])
def _estimated(gateway: Gateway, model: str) -> CostEstimate:
response: Final = gateway.request(
"POST", "/cost/estimate", {"model": model, "input_tokens": 20, "output_tokens": 20}
@ -236,6 +255,17 @@ def test_discount_and_margin_reprice_a_request_and_survive_rejection(gateway: Ga
(streamed_global_request_id, 0.12),
)
)
discounted: Final = _logged_breakdown(discounted_request_id)
assert discounted.original_cost == pytest.approx(0.06), discounted
assert discounted.discount_percent == pytest.approx(0.05), discounted
assert discounted.discount_amount == pytest.approx(0.06 * 0.05), discounted
margined: Final = _logged_breakdown(provider_margin_request_id)
assert margined.original_cost == pytest.approx(0.06), margined
assert margined.discount_percent == pytest.approx(0.05), margined
assert margined.discount_amount == pytest.approx(0.06 * 0.05), margined
assert margined.margin_percent == pytest.approx(0.1), margined
assert margined.margin_fixed_amount == pytest.approx(0.001), margined
assert margined.margin_total_amount == pytest.approx(0.057 * 0.1 + 0.001), margined
rejections: Final = (
(

71
tests/integration/providers/test_bedrock_converse_redacted_thinking_wire.py
View file

@ -414,3 +414,74 @@ def test_converse_redacted_thinking_replays_to_the_anthropic_api_as_claude_nativ
{"type": "redacted_thinking", "data": NATIVE_REDACTED},
{"type": "text", "text": FIRST_ANSWER},
)
_SIGNED_THINKING: Final[dict[str, JsonValue]] = {
"type": "thinking",
"thinking": "Checking the forecast.",
"signature": "scripted-signature-" + "s" * 32,
}
_CONVERSE_SIGNED_PART: Final[dict[str, JsonValue]] = {
"reasoningContent": {
"reasoningText": {"text": _SIGNED_THINKING["thinking"], "signature": _SIGNED_THINKING["signature"]}
}
}
@pytest.mark.parametrize(
("deployment", "replayed_thinking"),
[
pytest.param({"model": f"bedrock/converse/{NOVA_MODEL_ID}"}, (), id="nova"),
pytest.param(
{
"model": f"bedrock/converse/{APPLICATION_PROFILE_ARN}",
"model_info": {"base_model": f"bedrock/{NOVA_MODEL_ID}"},
},
(),
id="application-profile-arn-with-nova-base-model",
),
pytest.param(
{"model": f"bedrock/converse/{APPLICATION_PROFILE_ARN}"},
(_CONVERSE_SIGNED_PART,),
id="application-profile-arn-without-base-model",
),
pytest.param(
{
"model": f"bedrock/converse/{APPLICATION_PROFILE_ARN}",
"model_info": {"base_model": f"bedrock/{CLAUDE_MODEL_ID}"},
},
(_CONVERSE_SIGNED_PART,),
id="application-profile-arn-with-claude-base-model",
),
pytest.param(_CLAUDE, (_CONVERSE_SIGNED_PART,), id="claude"),
],
)
def test_signed_history_keeps_its_thinking_unless_the_converse_target_is_known_not_to_be_claude(
gateway: Gateway, deployment: dict[str, JsonValue], replayed_thinking: tuple[dict[str, JsonValue], ...]
) -> None:
model_info: Final = deployment.get("model_info")
parameters: Final = {key: value for key, value in deployment.items() if key != "model_info"}
with wire_server(_bedrock_peer()) as wire, gateway.scenario() as scenario:
model: Final = scenario.model(
model_info=model_info if isinstance(model_info, dict) else None,
api_key=TOKEN,
aws_region_name="us-east-1",
aws_bedrock_runtime_endpoint=wire.url,
**parameters,
)
reply: Final = _chat(
gateway,
{
"model": model,
"max_tokens": 256,
"messages": [
{"role": "user", "content": QUESTION},
{"role": "assistant", "content": FIRST_ANSWER, "thinking_blocks": [_SIGNED_THINKING]},
{"role": "user", "content": FOLLOW_UP},
],
},
)
assert reply.content == FINAL_ANSWER
assert _replayed_assistant_turn(wire).content == (*replayed_thinking, {"text": FIRST_ANSWER})

531
tests/integration/providers/test_bedrock_passthrough_guardrail_config_wire.py Normal file
View file

@ -0,0 +1,531 @@
import base64
import json
import uuid
from collections.abc import Iterator, Mapping
from dataclasses import dataclass
from pathlib import Path
from types import MappingProxyType
from typing import Final
import boto3
import pytest
import yaml
from botocore.config import Config
from botocore.exceptions import ClientError
from integration._support.budget import case_timeout_seconds
from integration._support.client import Gateway
from integration._support.process import graceful_stop_seconds, owned_proxy
from integration._support.upstream import _aws_event_frame
from integration._support.wire import Reply, Request, Wire, wire_server
from pydantic import JsonValue
pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 180)
_MODEL_ID: Final = "anthropic.claude-sonnet-5-v1:0"
_EVENT_STREAM: Final = "application/vnd.amazon.eventstream"
_BLOCKED_WORD: Final = "zzzforbidden"
_BLOCKED_TEXT: Final = "Blocked by guardrail"
_VERSIONLESS_GUARDRAIL_MESSAGE: Final = "guardrailIdentifier and guardrailVersion must be provided together."
_ACTIONS: Final = ("converse", "converse-stream", "invoke", "invoke-with-response-stream")
_DEPLOYMENT_GUARDRAIL: Final[Mapping[str, str]] = {
"guardrailIdentifier": "deployment-guardrail",
"guardrailVersion": "3",
"trace": "enabled",
"streamProcessingMode": "sync",
}
_CALLER_GUARDRAIL: Final[Mapping[str, str]] = {
"guardrailIdentifier": "caller-guardrail",
"guardrailVersion": "9",
"trace": "disabled",
}
_MANTLE_KEY: Final = "scripted-mantle-key"
_CALLER_INVOKE_HEADERS: Final[Mapping[str, str]] = {
"x-pass-x-amzn-bedrock-guardrailidentifier": "caller-guardrail",
"x-pass-x-amzn-bedrock-guardrailversion": "9",
"x-pass-x-amzn-bedrock-trace": "DISABLED",
}
@dataclass(frozen=True, slots=True)
class Applied:
"""The guardrail a request asked Bedrock for, read from where Bedrock reads it for that action."""
identifier: str
version: str | None
trace: str | None
stream_processing_mode: str | None
@dataclass(frozen=True, slots=True)
class Outcome:
"""What the AWS SDK caller reads back."""
text: str
intervened: bool
traced_guardrail: str | None
def _action(request: Request) -> str:
return request.target.partition("?")[0].rsplit("/", 1)[-1]
def _applied(request: Request) -> Applied | None:
if _action(request).startswith("converse"):
config: Final = json.loads(request.body).get("guardrailConfig")
if config is None:
return None
return Applied(
config["guardrailIdentifier"],
config.get("guardrailVersion"),
config.get("trace"),
config.get("streamProcessingMode"),
)
identifier: Final = request.headers.get("x-amzn-bedrock-guardrailidentifier")
if identifier is None:
return None
return Applied(
identifier,
request.headers.get("x-amzn-bedrock-guardrailversion"),
request.headers.get("x-amzn-bedrock-trace"),
None,
)
def _identifier(request: Request) -> str | None:
applied: Final = _applied(request)
return None if applied is None else applied.identifier
def _assessment(applied: Applied) -> Mapping[str, JsonValue]:
return {
applied.identifier: {
"wordPolicy": {"customWords": [{"match": _BLOCKED_WORD, "action": "BLOCKED", "detected": True}]},
"appliedGuardrailDetails": {"guardrailId": applied.identifier, "guardrailVersion": applied.version},
}
}
def _bedrock_with_guardrails(request: Request) -> Reply:
"""Bedrock's runtime as the caller sees it: the guardrail the request names blocks a prompt holding the blocked
word, answering with its message and, when the request enabled tracing, a trace naming that guardrail. A guardrail
named without a version is refused with a 400 ValidationException, as real Bedrock in us-east-1 refused it on
every one of these actions on 2026-10-19."""
applied: Final = _applied(request)
if applied is not None and applied.version is None:
return Reply(
status=400,
body=json.dumps({"message": _VERSIONLESS_GUARDRAIL_MESSAGE}).encode(),
headers={"x-amzn-errortype": "ValidationException"},
)
blocked: Final = applied is not None and _BLOCKED_WORD.encode() in request.body
traced: Final = blocked and applied is not None and (applied.trace or "").lower() == "enabled"
text: Final = _BLOCKED_TEXT if blocked else "ok"
stop_reason: Final = "guardrail_intervened" if blocked else "end_turn"
usage: Final = {"inputTokens": 5, "outputTokens": 1, "totalTokens": 6}
converse_trace: Final = {"trace": {"guardrail": {"inputAssessment": _assessment(applied)}}} if traced else {}
action: Final = _action(request)
if action == "converse":
return Reply(
body=json.dumps(
{
"output": {"message": {"role": "assistant", "content": [{"text": text}]}},
"stopReason": stop_reason,
"usage": usage,
"metrics": {"latencyMs": 3},
**converse_trace,
}
).encode()
)
if action == "converse-stream":
return Reply(
chunks=(
_aws_event_frame("messageStart", {"role": "assistant"}, "sc", "u"),
_aws_event_frame("contentBlockDelta", {"delta": {"text": text}, "contentBlockIndex": 0}, "sc", "u"),
_aws_event_frame("contentBlockStop", {"contentBlockIndex": 0}, "sc", "u"),
_aws_event_frame("messageStop", {"stopReason": stop_reason}, "sc", "u"),
_aws_event_frame(
"metadata", {"usage": usage, "metrics": {"latencyMs": 3}, **converse_trace}, "sc", "u"
),
),
content_type=_EVENT_STREAM,
)
message: Final = {
"id": "msg_bedrock",
"type": "message",
"role": "assistant",
"model": "claude",
"content": [{"type": "text", "text": text}],
"stop_reason": "end_turn",
"stop_sequence": None,
"usage": {"input_tokens": 5, "output_tokens": 1},
**({"amazon-bedrock-guardrailAction": "INTERVENED"} if blocked else {}),
**({"amazon-bedrock-trace": {"guardrail": {"input": _assessment(applied)}}} if traced else {}),
}
if action == "invoke":
return Reply(body=json.dumps(message).encode())
chunk: Final = {"bytes": base64.b64encode(json.dumps(message).encode()).decode()}
return Reply(chunks=(_aws_event_frame("chunk", chunk, "sc", "u"),), content_type=_EVENT_STREAM)
def _bedrock_outage(request: Request) -> Reply:
return Reply(
status=500,
body=json.dumps({"message": "scripted outage"}).encode(),
headers={"x-amzn-errortype": "InternalServerException"},
)
def _traced_guardrail(assessment: object) -> str | None:
return None if not isinstance(assessment, dict) else next(iter(assessment))
def _send(runtime: object, action: str, model: str, prompt: str, caller_guardrail: bool) -> Outcome:
messages: Final = [{"role": "user", "content": [{"text": prompt}]}]
invoke_body: Final = json.dumps(
{"anthropic_version": "bedrock-2023-05-31", "max_tokens": 16, "messages": [{"role": "user", "content": prompt}]}
)
invoke_guardrail: Final = (
{"guardrailIdentifier": "caller-guardrail", "guardrailVersion": "9", "trace": "DISABLED"}
if caller_guardrail
else {}
)
if action == "converse":
response: Final = getattr(runtime, "converse")(
modelId=model,
messages=messages,
**({"guardrailConfig": dict(_CALLER_GUARDRAIL)} if caller_guardrail else {}),
)
return Outcome(
text=response["output"]["message"]["content"][0]["text"],
intervened=response["stopReason"] == "guardrail_intervened",
traced_guardrail=_traced_guardrail(response.get("trace", {}).get("guardrail", {}).get("inputAssessment")),
)
if action == "converse-stream":
events: Final = list(
getattr(runtime, "converse_stream")(
modelId=model,
messages=messages,
**(
{"guardrailConfig": {**_CALLER_GUARDRAIL, "streamProcessingMode": "async"}}
if caller_guardrail
else {}
),
)["stream"]
)
(stop,) = tuple(event["messageStop"]["stopReason"] for event in events if "messageStop" in event)
(metadata,) = tuple(event["metadata"] for event in events if "metadata" in event)
return Outcome(
text="".join(
event["contentBlockDelta"]["delta"]["text"] for event in events if "contentBlockDelta" in event
),
intervened=stop == "guardrail_intervened",
traced_guardrail=_traced_guardrail(metadata.get("trace", {}).get("guardrail", {}).get("inputAssessment")),
)
if action == "invoke":
message: Final = json.loads(
getattr(runtime, "invoke_model")(modelId=model, body=invoke_body, **invoke_guardrail)["body"].read()
)
else:
(message,) = tuple(
json.loads(event["chunk"]["bytes"])
for event in getattr(runtime, "invoke_model_with_response_stream")(
modelId=model, body=invoke_body, **invoke_guardrail
)["body"]
)
return Outcome(
text=message["content"][0]["text"],
intervened=message.get("amazon-bedrock-guardrailAction") == "INTERVENED",
traced_guardrail=_traced_guardrail(message.get("amazon-bedrock-trace", {}).get("guardrail", {}).get("input")),
)
def _signed_header_names(request: Request) -> frozenset[str]:
return frozenset(request.headers["authorization"].split("SignedHeaders=")[1].split(",")[0].split(";"))
def _deployment(wire: Wire, guardrail_config: Mapping[str, str] | None = None) -> Mapping[str, JsonValue]:
return {
"model": f"bedrock/{_MODEL_ID}",
"api_base": wire.url,
"aws_access_key_id": "AKIASCRIPTEDPROVIDER",
"aws_secret_access_key": "scripted-secret",
"aws_region_name": "us-east-1",
"use_in_pass_through": True,
**({} if guardrail_config is None else {"guardrailConfig": dict(guardrail_config)}),
}
@dataclass(frozen=True, slots=True)
class Rig:
proxy: Gateway
wire: Wire
guarded_model: str
unguarded_model: str
balanced_model: str
balanced_wires: Mapping[str, Wire]
failing_model: str
failing_wire: Wire
fallback_wire: Wire
mantle_model: str
versionless_model: str
key: str
def runtime( # noqa: ANN201 # boto3 clients are untyped
self, monkeypatch: pytest.MonkeyPatch, extra_headers: Mapping[str, str] = MappingProxyType({})
):
monkeypatch.setenv("AWS_BEARER_TOKEN_BEDROCK", self.key)
client: Final = boto3.client(
"bedrock-runtime",
endpoint_url=f"{str(self.proxy.client.base_url).rstrip('/')}/bedrock",
region_name="us-east-1",
config=Config(read_timeout=case_timeout_seconds(), retries={"total_max_attempts": 1}),
)
client.meta.events.register(
"before-send.bedrock-runtime.*",
lambda request, **_: request.headers.update(extra_headers),
)
return client
def received(self, wire: Wire, marker: str) -> tuple[Request, ...]:
return tuple(request for request in wire.drain() if marker.encode() in request.body)
@pytest.fixture(scope="module")
def rig(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]:
suffix: Final = uuid.uuid4().hex[:12]
guarded_model: Final = f"bedrock-guarded-{suffix}"
unguarded_model: Final = f"bedrock-unguarded-{suffix}"
balanced_model: Final = f"bedrock-balanced-{suffix}"
failing_model: Final = f"bedrock-failing-{suffix}"
fallback_model: Final = f"bedrock-fallback-{suffix}"
mantle_model: Final = f"bedrock-mantle-{suffix}"
versionless_model: Final = f"bedrock-versionless-{suffix}"
directory: Final = tmp_path_factory.mktemp("bedrock-guardrail-config")
base_config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
with (
wire_server(_bedrock_with_guardrails) as wire,
wire_server(_bedrock_with_guardrails) as wire_a,
wire_server(_bedrock_with_guardrails) as wire_b,
wire_server(_bedrock_outage) as failing_wire,
wire_server(_bedrock_with_guardrails) as fallback_wire,
):
config: Final = {
**base_config,
"model_list": [
*base_config.get("model_list", []),
{
"model_name": guarded_model,
"litellm_params": _deployment(wire, _DEPLOYMENT_GUARDRAIL),
},
{"model_name": unguarded_model, "litellm_params": _deployment(wire)},
{
"model_name": balanced_model,
"litellm_params": _deployment(
wire_a, {"guardrailIdentifier": "guardrail-a", "guardrailVersion": "1"}
),
},
{
"model_name": balanced_model,
"litellm_params": _deployment(
wire_b, {"guardrailIdentifier": "guardrail-b", "guardrailVersion": "2"}
),
},
{
"model_name": failing_model,
"litellm_params": _deployment(
failing_wire,
{"guardrailIdentifier": "failing-guardrail", "guardrailVersion": "5", "trace": "enabled"},
),
},
{
"model_name": fallback_model,
"litellm_params": _deployment(
fallback_wire,
{"guardrailIdentifier": "fallback-guardrail", "guardrailVersion": "6", "trace": "enabled"},
),
},
{
"model_name": versionless_model,
"litellm_params": _deployment(
wire, {"guardrailIdentifier": "versionless-guardrail", "trace": "enabled"}
),
},
{
"model_name": mantle_model,
"litellm_params": {
"model": "bedrock_mantle/anthropic.claude-sonnet-5",
"api_base": wire.url,
"api_key": _MANTLE_KEY,
"aws_region_name": "us-east-1",
"use_in_pass_through": True,
"guardrailConfig": {
"guardrailIdentifier": "mantle-guardrail",
"guardrailVersion": "4",
"trace": "enabled",
},
},
},
],
"router_settings": {
**base_config.get("router_settings", {}),
"fallbacks": [{failing_model: [fallback_model]}],
},
}
config_path: Final = directory / "bedrock-guardrail-config.yaml"
config_path.write_text(yaml.safe_dump(config))
with (
owned_proxy(session_gateway, directory, {}, config=config_path, workers=2) as proxy,
proxy.scenario() as scenario,
):
yield Rig(
proxy=proxy,
wire=wire,
guarded_model=guarded_model,
unguarded_model=unguarded_model,
balanced_model=balanced_model,
balanced_wires={"guardrail-a": wire_a, "guardrail-b": wire_b},
failing_model=failing_model,
failing_wire=failing_wire,
fallback_wire=fallback_wire,
mantle_model=mantle_model,
versionless_model=versionless_model,
key=scenario.key(key_alias=f"guardrail-key-{suffix}"),
)
def _expected_on_the_wire(action: str, request: Request) -> Applied:
if action == "converse":
return Applied("deployment-guardrail", "3", "enabled", None)
if action == "converse-stream":
return Applied("deployment-guardrail", "3", "enabled", "sync")
assert {"x-amzn-bedrock-guardrailidentifier", "x-amzn-bedrock-guardrailversion", "x-amzn-bedrock-trace"} <= (
_signed_header_names(request)
), request.headers["authorization"]
return Applied("deployment-guardrail", "3", "ENABLED", None)
@pytest.mark.parametrize("action", _ACTIONS)
def test_passthrough_on_a_guardrailed_deployment_blocks_with_the_deployment_guardrail(
rig: Rig, monkeypatch: pytest.MonkeyPatch, action: str
) -> None:
marker: Final = uuid.uuid4().hex
runtime: Final = rig.runtime(monkeypatch)
blocked: Final = _send(runtime, action, rig.guarded_model, f"{marker} say {_BLOCKED_WORD}", caller_guardrail=False)
allowed: Final = _send(runtime, action, rig.guarded_model, f"{marker} say hello", caller_guardrail=False)
assert blocked == Outcome(_BLOCKED_TEXT, intervened=True, traced_guardrail="deployment-guardrail"), blocked
assert allowed == Outcome("ok", intervened=False, traced_guardrail=None), allowed
received: Final = rig.received(rig.wire, marker)
assert [_action(request) for request in received] == [action, action], received
assert [_applied(request) for request in received] == [
_expected_on_the_wire(action, request) for request in received
], [(request.headers, request.body) for request in received]
@pytest.mark.parametrize("action", _ACTIONS)
def test_caller_guardrail_is_replaced_by_the_deployment_guardrail(
rig: Rig, monkeypatch: pytest.MonkeyPatch, action: str
) -> None:
marker: Final = uuid.uuid4().hex
runtime: Final = rig.runtime(monkeypatch, _CALLER_INVOKE_HEADERS)
outcome: Final = _send(runtime, action, rig.guarded_model, f"{marker} say {_BLOCKED_WORD}", caller_guardrail=True)
assert outcome == Outcome(_BLOCKED_TEXT, intervened=True, traced_guardrail="deployment-guardrail"), outcome
(request,) = rig.received(rig.wire, marker)
assert _applied(request) == _expected_on_the_wire(action, request), (request.headers, request.body)
assert "caller-guardrail" not in json.dumps({**request.headers, "body": request.body.decode()}), request
@pytest.mark.parametrize("action", _ACTIONS)
def test_mantle_deployment_passthrough_sends_its_guardrail_with_bearer_auth(
rig: Rig, monkeypatch: pytest.MonkeyPatch, action: str
) -> None:
marker: Final = uuid.uuid4().hex
runtime: Final = rig.runtime(monkeypatch)
outcome: Final = _send(runtime, action, rig.mantle_model, f"{marker} say {_BLOCKED_WORD}", caller_guardrail=False)
assert outcome == Outcome(_BLOCKED_TEXT, intervened=True, traced_guardrail="mantle-guardrail"), outcome
(request,) = rig.received(rig.wire, marker)
assert request.headers["authorization"] == f"Bearer {_MANTLE_KEY}", request.headers
assert _applied(request) == Applied(
"mantle-guardrail", "4", "ENABLED" if action.startswith("invoke") else "enabled", None
), (request.headers, request.body)
@pytest.mark.parametrize("action", _ACTIONS)
def test_deployment_guardrail_without_a_version_returns_the_bedrock_rejection(
rig: Rig, monkeypatch: pytest.MonkeyPatch, action: str
) -> None:
marker: Final = uuid.uuid4().hex
runtime: Final = rig.runtime(monkeypatch)
with pytest.raises(ClientError) as raised:
_send(runtime, action, rig.versionless_model, f"{marker} say hello", caller_guardrail=False)
error: Final = raised.value.response.get("Error", {})
assert (error.get("Code"), error.get("Message")) == ("ValidationException", _VERSIONLESS_GUARDRAIL_MESSAGE), error
received: Final = rig.received(rig.wire, marker)
assert {_applied(request) for request in received} == {
Applied("versionless-guardrail", None, "ENABLED" if action.startswith("invoke") else "enabled", None)
}, [(request.headers, request.body) for request in received]
@pytest.mark.parametrize("action", _ACTIONS)
def test_deployment_without_a_guardrail_leaves_the_caller_guardrail_as_sent(
rig: Rig, monkeypatch: pytest.MonkeyPatch, action: str
) -> None:
if action.startswith("invoke"):
pytest.skip("BUG: the router-model invoke pass-through drops the caller's X-Amzn-Bedrock-Guardrail* headers")
marker: Final = uuid.uuid4().hex
runtime: Final = rig.runtime(monkeypatch)
outcome: Final = _send(runtime, action, rig.unguarded_model, f"{marker} say {_BLOCKED_WORD}", caller_guardrail=True)
assert outcome == Outcome(_BLOCKED_TEXT, intervened=True, traced_guardrail=None), outcome
(request,) = rig.received(rig.wire, marker)
assert _applied(request) == Applied(
"caller-guardrail", "9", "disabled", "async" if action == "converse-stream" else None
), request.body
def test_each_load_balanced_request_carries_the_guardrail_of_the_deployment_it_reached(
rig: Rig, monkeypatch: pytest.MonkeyPatch
) -> None:
marker: Final = uuid.uuid4().hex
runtime: Final = rig.runtime(monkeypatch)
outcomes: Final = tuple(
_send(runtime, action, rig.balanced_model, f"{marker}-{attempt} say {_BLOCKED_WORD}", caller_guardrail=False)
for attempt, action in enumerate(_ACTIONS * 6)
)
assert all(outcome.intervened for outcome in outcomes), outcomes
reached: Final = {name: rig.received(wire, marker) for name, wire in rig.balanced_wires.items()}
assert sum(len(requests) for requests in reached.values()) == len(outcomes), reached
assert all(reached.values()), {name: len(requests) for name, requests in reached.items()}
assert {name: {_identifier(request) for request in requests} for name, requests in reached.items()} == {
"guardrail-a": {"guardrail-a"},
"guardrail-b": {"guardrail-b"},
}, reached
def test_a_fallback_sends_the_guardrail_of_the_deployment_it_fell_back_to(
rig: Rig, monkeypatch: pytest.MonkeyPatch
) -> None:
marker: Final = uuid.uuid4().hex
runtime: Final = rig.runtime(monkeypatch)
outcomes: Final = tuple(
_send(runtime, action, rig.failing_model, f"{marker} say {_BLOCKED_WORD}", caller_guardrail=False)
for action in _ACTIONS
)
assert outcomes == (Outcome(_BLOCKED_TEXT, intervened=True, traced_guardrail="fallback-guardrail"),) * 4, outcomes
failed: Final = rig.received(rig.failing_wire, marker)
served: Final = rig.received(rig.fallback_wire, marker)
assert [_action(request) for request in served] == list(_ACTIONS), served
assert {_identifier(request) for request in failed} == {"failing-guardrail"}, failed
assert {_identifier(request) for request in served} == {"fallback-guardrail"}, served

22
tests/integration/providers/test_responses_minted_reasoning_replay_wire.py
View file

@ -77,7 +77,18 @@ _AZURE_AI_OPENAI_HOST: Final = _Deployment(
"/openai/v1/responses?api-version=preview",
strips_message_status=True,
)
_DROPPING: Final = (_OPENAI, _AZURE, _AZURE_AI_OPENAI_HOST)
_BEDROCK_MANTLE: Final = _Deployment(
"bedrock_mantle", "bedrock_mantle/openai.gpt-oss-120b", "synthetic-mantle-key", "/v1/responses"
)
_BEDROCK: Final = _Deployment(
"bedrock",
"bedrock/openai.gpt-oss-120b-1:0",
"synthetic-bedrock-key",
"/openai/v1/responses",
MappingProxyType({"aws_region_name": "us-east-1"}),
model_info=MappingProxyType({"supported_endpoints": ["/v1/responses"]}),
)
_DROPPING: Final = (_OPENAI, _AZURE, _AZURE_AI_OPENAI_HOST, _BEDROCK_MANTLE, _BEDROCK)
_KEEPING: Final = (
_Deployment("litellm_proxy", f"litellm_proxy/{_GPT}", "synthetic-proxy-key", "/responses"),
_Deployment("databricks", "databricks/gpt-5.6", "synthetic-databricks-key", "/responses"),
@ -95,15 +106,6 @@ _KEEPING: Final = (
"/v1/responses",
types_untyped_items_as_messages=True,
),
_Deployment("bedrock_mantle", "bedrock_mantle/openai.gpt-oss-120b", "synthetic-mantle-key", "/v1/responses"),
_Deployment(
"bedrock",
"bedrock/openai.gpt-oss-120b-1:0",
"synthetic-bedrock-key",
"/openai/v1/responses",
MappingProxyType({"aws_region_name": "us-east-1"}),
model_info=MappingProxyType({"supported_endpoints": ["/v1/responses"]}),
),
*(
_Deployment(slug, f"{slug}/{model}", f"synthetic-{slug}-key", "/responses")
for slug, model in (

4
tests/integration/routing/test_video_poll_routing_wire.py
View file

@ -338,9 +338,9 @@ def test_an_id_without_a_served_group_still_polls_its_own_deployment(rig: _Rig)
video_id: Final = created_id(create(rig.gateway, {"model": rig.member_a, "seconds": "5", "size": "1280x720"}))
assert decoded(video_id).get("served_model") is None
legacy: Final = forged_id("dashscope", rig.member_a, task_of(video_id), None)
assert legacy == video_id
assert {**decoded(legacy), "deployment_id": decoded(video_id).get("deployment_id")} == decoded(video_id)
rig.drain()
_assert_served(_poll(rig, legacy), legacy)
_assert_served(_poll(rig, legacy), video_id)
download: Final = _download(rig, legacy)
assert download.status_code == 200 and download.content == MP4, download.text
received: Final = rig.a.drain()

84
tests/integration/spend/test_failing_request_stored_snapshot.py Normal file
View file

@ -0,0 +1,84 @@
import json
import uuid
from collections.abc import Iterator
from pathlib import Path
from typing import Final
import pytest
import yaml
from integration._support.client import Gateway, eventually, gateway_from_environment, object_value
from integration._support.process import owned_proxy
from pydantic import JsonValue
UNREACHABLE_BASE: Final = "http://127.0.0.1:9/v1"
TRANSCRIPT_TURNS: Final = 40
RETRIES: Final = 2
STORED_REQUEST_BUDGET_KB: Final = 64
def _config_storing_prompts(directory: Path) -> Path:
config: Final = object_value(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()))
general_settings: Final = object_value(config["general_settings"])
path: Final = directory / "store-prompts.yaml"
path.write_text(
yaml.safe_dump({**config, "general_settings": {**general_settings, "store_prompts_in_spend_logs": True}})
)
return path
@pytest.fixture(scope="module")
def gateway(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Gateway]:
with gateway_from_environment() as parent:
directory: Final = tmp_path_factory.mktemp("failing_request_stored_snapshot")
with owned_proxy(parent, directory, {}, config=_config_storing_prompts(directory)) as owned:
yield owned
def _transcript(marker: str) -> list[JsonValue]:
turns: Final[list[JsonValue]] = [
{"role": role, "content": f"turn {turn // 2} {role}"}
for turn, role in enumerate(("user", "assistant") * TRANSCRIPT_TURNS)
]
return [*turns, {"role": "user", "content": marker}]
def _stored_request(gateway: Gateway, call_id: str) -> JsonValue:
rows: Final = eventually(
lambda: gateway.rows('SELECT proxy_server_request FROM "LiteLLM_SpendLogs" WHERE request_id = %s', (call_id,)),
lambda values: len(values) == 1,
seconds=40,
)
stored: Final = rows[0]["proxy_server_request"]
return json.loads(stored) if isinstance(stored, str) else stored
def test_failing_request_with_retries_and_fallback_stores_one_copy_of_the_request(gateway: Gateway) -> None:
with gateway.scenario() as scenario:
primary: Final = f"stored-snapshot-{uuid.uuid4().hex}"
fallback: Final = f"stored-snapshot-fb-{uuid.uuid4().hex}"
for name in (primary, primary, fallback, fallback):
scenario.model(name=name, api_base=UNREACHABLE_BASE, cooldown_time=0)
key: Final = scenario.key(models=[primary, fallback])
marker: Final = f"stored snapshot {uuid.uuid4().hex}"
response: Final = gateway.request(
"POST",
"/v1/chat/completions",
{
"model": primary,
"max_tokens": 16,
"messages": _transcript(marker),
"router_settings_override": {"num_retries": RETRIES, "fallbacks": [{primary: [fallback]}]},
},
key=key,
)
assert response.status_code != 200, response.text
assert fallback in response.text, response.text
serialized: Final = json.dumps(_stored_request(gateway, response.headers["x-litellm-call-id"]))
stored_kb: Final = len(serialized.encode()) / 1024
assert marker in serialized, serialized[:400]
assert stored_kb <= STORED_REQUEST_BUDGET_KB, (
f"one failing request stored a {stored_kb:.0f} KB request snapshot for a "
f"{2 * TRANSCRIPT_TURNS + 1}-message transcript with {RETRIES} retries and a fallback"
)

30
tests/unit/batches/test_batch_utils.py
View file

@ -2296,6 +2296,36 @@ async def test_native_vertex_rows_route_to_vertex_cost_path_without_flag(monkeyp
assert {(call["model"], call["custom_llm_provider"]) for call in calls} == {("gemini-2.5-flash", "vertex_ai")}
@pytest.mark.asyncio
async def test_raw_vertex_claude_rows_are_priced_like_with_the_native_output_flag(monkeypatch):
claude_message: Final = {
"id": "msg_1",
"type": "message",
"role": "assistant",
"model": "claude-sonnet-4-6",
"content": [{"type": "text", "text": "Blue."}],
"stop_reason": "end_turn",
"usage": {"input_tokens": 12, "output_tokens": 5},
}
rows: Final = [
{"custom_id": "ok", "processed_time": "2026-10-29T03:00:00Z", "response": claude_message, "status": ""},
{"custom_id": "bad", "processed_time": "2026-10-29T03:00:00Z", "response": {}, "status": "Bad Request"},
]
async def cost_with_flag(flag: bool) -> bu.BatchCostUsageResult:
monkeypatch.setattr(litellm, "disable_vertex_batch_output_transformation", flag, raising=False)
return await bu.calculate_batch_cost_and_usage(
file_content_dictionary=rows, custom_llm_provider="vertex_ai", model_name="vertex_ai/claude-sonnet-4-6"
)
unflagged: Final = await cost_with_flag(False)
flagged: Final = await cost_with_flag(True)
assert unflagged.cost > 0
assert (unflagged.successful_requests, unflagged.failed_requests) == (1, 1)
assert (unflagged.cost, unflagged.usage.total_tokens) == (flagged.cost, flagged.usage.total_tokens)
@pytest.mark.asyncio
async def test_openai_shaped_vertex_rows_keep_the_generic_path_without_flag(monkeypatch):
monkeypatch.setattr(litellm, "disable_vertex_batch_output_transformation", False, raising=False)

113
tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_handler.py
View file

@ -1,7 +1,12 @@
import json
from datetime import datetime
from pathlib import Path
from typing import Final
from unittest.mock import AsyncMock, MagicMock, patch
import httpx
import pytest
import respx
import litellm
@ -11,6 +16,7 @@ from litellm.completion_extras.litellm_responses_transformation.handler import (
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging
from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
from litellm.types.router import GenericLiteLLMParams
from litellm.litellm_core_utils.prompt_templates.common_utils import encrypted_reasoning_signature
from litellm.types.utils import ModelResponse
@ -335,3 +341,110 @@ async def test_acompletion_keeps_provider_native_model_id_through_responses(
handed_model = responses_call.call_args.kwargs["model"]
assert _upstream_model_for(handed_model, custom_llm_provider) == expected_upstream_model
_CHATGPT_RESPONSES_URL: Final = "https://chatgpt.com/backend-api/codex/responses"
def _reasoning_item(item_id: str) -> dict[str, object]:
return {"type": "reasoning", "id": item_id, "summary": [], "encrypted_content": f"gAAAA_{item_id}"}
def _weather_call(index: int, city: str) -> dict[str, object]:
return {
"type": "function_call",
"id": f"fc_{index}",
"status": "completed",
"arguments": json.dumps({"city": city}),
"call_id": f"call_{index}",
"name": "get_weather",
}
def _two_reasoning_tool_turn_stream() -> httpx.Response:
turn: Final = {
"id": "resp_two_reasoning",
"object": "response",
"created_at": 1,
"status": "completed",
"model": "gpt-5.6-sol",
"output": [
_reasoning_item("rs_first"),
_weather_call(1, "Oslo"),
_reasoning_item("rs_second"),
_weather_call(2, "Lyon"),
],
"usage": {"input_tokens": 10, "output_tokens": 20, "total_tokens": 30},
}
body: Final = f"data: {json.dumps({'type': 'response.completed', 'response': turn})}\n\ndata: [DONE]\n\n"
return httpx.Response(200, headers={"content-type": "text/event-stream"}, text=body)
_FAR_FUTURE_EXPIRY: Final = 4_102_444_800
@pytest.fixture
def _chatgpt_backend(tmp_path: Path, monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter) -> None:
auth: Final = {"access_token": "chatgpt-test-token", "expires_at": _FAR_FUTURE_EXPIRY, "account_id": "acct_test"}
(tmp_path / "auth.json").write_text(json.dumps(auth))
monkeypatch.setenv("CHATGPT_TOKEN_DIR", str(tmp_path))
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
litellm.in_memory_llm_clients_cache.flush_cache()
respx_mock.post(_CHATGPT_RESPONSES_URL).mock(side_effect=lambda _request: _two_reasoning_tool_turn_stream())
_WEATHER_QUESTION: Final = "Weather in Oslo and Lyon?"
_WEATHER_TOOL_INPUT: Final = {"type": "object", "properties": {"city": {"type": "string"}}}
_EVERY_REASONING_SIGNATURE: Final = [
encrypted_reasoning_signature("gAAAA_rs_first"),
encrypted_reasoning_signature("gAAAA_rs_second"),
]
def _thinking_signatures(reply: object) -> list[object]:
assert isinstance(reply, dict)
return [block.get("data") for block in reply["content"] if block["type"] == "redacted_thinking"]
@pytest.mark.asyncio
@pytest.mark.usefixtures("_chatgpt_backend")
async def test_async_messages_client_gets_every_reasoning_item_of_a_non_streamed_turn() -> None:
reply: Final = await litellm.anthropic.messages.acreate(
model="chatgpt/gpt-5.6-sol",
max_tokens=64,
messages=[{"role": "user", "content": _WEATHER_QUESTION}],
tools=[{"name": "get_weather", "input_schema": _WEATHER_TOOL_INPUT}],
)
assert _thinking_signatures(reply) == _EVERY_REASONING_SIGNATURE
_CHAT_WEATHER_TOOL: Final = {"type": "function", "function": {"name": "get_weather", "parameters": _WEATHER_TOOL_INPUT}}
def _reasoning_item_ids(response: object) -> list[object]:
assert isinstance(response, ModelResponse)
return [item["id"] for item in response.choices[0].message.reasoning_items]
@pytest.mark.asyncio
@pytest.mark.usefixtures("_chatgpt_backend")
async def test_async_chat_client_keeps_only_the_last_reasoning_item_of_a_turn() -> None:
response: Final = await litellm.acompletion(
model="chatgpt/gpt-5.6-sol",
messages=[{"role": "user", "content": _WEATHER_QUESTION}],
tools=[_CHAT_WEATHER_TOOL],
)
assert _reasoning_item_ids(response) == ["rs_second"]
@pytest.mark.usefixtures("_chatgpt_backend")
def test_sync_chat_client_keeps_only_the_last_reasoning_item_of_a_turn() -> None:
response: Final = litellm.completion(
model="chatgpt/gpt-5.6-sol",
messages=[{"role": "user", "content": _WEATHER_QUESTION}],
tools=[_CHAT_WEATHER_TOOL],
)
assert _reasoning_item_ids(response) == ["rs_second"]

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