Merge pull request #33308 from BerriAI/litellm_internal_staging

chore(ci): promote internal staging to main
This commit is contained in:
yuneng-jiang 2026-07-14 17:36:19 -07:00 committed by GitHub
commit 6a797f97b2
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
329 changed files with 19507 additions and 5643 deletions

2
.github/CODEOWNERS vendored Normal file
View file

@ -0,0 +1,2 @@
/ui/ @yuneng-jiang @ryan-crabbe-berri
/litellm/proxy/_experimental/out/ @yuneng-jiang @ryan-crabbe-berri

View file

@ -0,0 +1,47 @@
name: "Set up uv with retries"
description: >-
Install uv via astral-sh/setup-uv, retrying on transient failures. Even with
an exact pinned version, the action resolves the artifact URL by fetching
https://raw.githubusercontent.com/astral-sh/versions/main/v1/uv.ndjson in a
single request with no retry, timeout, or fallback, so one connection-level
network error ("fetch failed") fails the whole job before any test runs.
Retrying the full step covers the manifest fetch and the binary download.
inputs:
version:
description: "uv version to install"
required: true
runs:
using: composite
steps:
- name: Set up uv (attempt 1)
id: attempt-1
continue-on-error: true
uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7.6.0
with:
version: ${{ inputs.version }}
- name: Wait before attempt 2
if: steps.attempt-1.outcome == 'failure'
shell: bash
run: sleep 15
- name: Set up uv (attempt 2)
id: attempt-2
if: steps.attempt-1.outcome == 'failure'
continue-on-error: true
uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7.6.0
with:
version: ${{ inputs.version }}
- name: Wait before attempt 3
if: steps.attempt-2.outcome == 'failure'
shell: bash
run: sleep 30
- name: Set up uv (attempt 3)
if: steps.attempt-2.outcome == 'failure'
uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7.6.0
with:
version: ${{ inputs.version }}

View file

@ -41,3 +41,27 @@ If you're seeing a delay in your PR being merged, ping the LiteLLM Team on [Slac
✅ Test
## Changes
## QA runbook
<!-- Only needed when your PR edits tests/e2e; delete this section otherwise
For each e2e test you added or changed, list the manual steps a reviewer can follow to reproduce it by hand against a live proxy, mapping 1:1 to what the test asserts: one top-level bullet per test giving its pytest node id followed by what it proves in plain words, then a nested "- [ ]" checklist where each item is a concrete action (route, request body, expected response) and the final item is the sanity-check step shown in the examples. Note environment prerequisites (provider credentials, config flags) and any nuances a manual run will hit. See PRs #32914 and #32963 for full examples
Example checklists:
- tests/e2e/quota_management/ratelimit/test_rate_limit_e2e.py::TestKeyRateLimits::test_rpm_limit_blocks_over_limit - a key allowed 2 requests a minute serves exactly 2 and refuses the 3rd
- [ ] Generate a limited key: curl -X POST http://localhost:4000/key/generate -H "Authorization: Bearer sk-1234" -d '{"rpm_limit": 2}'
- [ ] Send three /v1/chat/completions requests with that key inside one minute
- [ ] Expect the first two to return 200 and the third to return 429 naming the rpm limit
- [ ] Sanity check: this test makes sense to add and is not hand-wavey (e.g., assert actual expected spend instead of just spend > 0) or potentially flaky
- tests/e2e/management/test_management_e2e.py::TestModelRoutes::test_model_create_appears_in_ui - a deployment created through the API shows up on the Admin UI models page
- [ ] POST /model/new with the master key, a bedrock model, and aws_region_name (needs STORE_MODEL_IN_DB=True and AWS credentials)
- [ ] Open http://localhost:4000/ui/?page=models and expect a deployment row showing the returned model id
- [ ] Sanity check: this test makes sense to add and is not hand-wavey (e.g., assert actual expected spend instead of just spend > 0) or potentially flaky
-->
### Final Attestation
- [ ] The tests check the right things, including the edge cases, and regressions in the respective real-world customer use-cases are not possible after this PR

View file

@ -63,7 +63,7 @@ jobs:
python-version: "3.12"
- name: Set up uv
uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7
uses: ./.github/actions/setup-uv-with-retries
with:
version: "0.10.9"

View file

@ -18,7 +18,7 @@ jobs:
with:
persist-credentials: false
- name: Set up uv
uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7
uses: ./.github/actions/setup-uv-with-retries
with:
version: "0.10.9"
- name: Update JSON Data

View file

@ -31,7 +31,7 @@ jobs:
python-version: "3.12"
- name: Set up uv
uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7
uses: ./.github/actions/setup-uv-with-retries
with:
version: "0.10.9"

View file

@ -37,7 +37,7 @@ jobs:
python-version: "3.12"
- name: Set up uv
uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7
uses: ./.github/actions/setup-uv-with-retries
with:
version: "0.10.9"

View file

@ -39,7 +39,7 @@ jobs:
python-version: "3.12"
- name: Set up uv
uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7
uses: ./.github/actions/setup-uv-with-retries
with:
version: "0.10.9"

View file

@ -35,7 +35,7 @@ jobs:
python-version: "3.12"
- name: Set up uv
uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7
uses: ./.github/actions/setup-uv-with-retries
with:
version: "0.10.9"

View file

@ -38,7 +38,7 @@ jobs:
python-version: "3.12"
- name: Set up uv
uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7
uses: ./.github/actions/setup-uv-with-retries
with:
version: "0.10.9"

View file

@ -33,7 +33,7 @@ jobs:
python-version: "3.12"
- name: Set up uv
uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7
uses: ./.github/actions/setup-uv-with-retries
with:
version: "0.10.9"
@ -172,7 +172,7 @@ jobs:
python-version: "3.12"
- name: Set up uv
uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7
uses: ./.github/actions/setup-uv-with-retries
with:
version: "0.10.9"

View file

@ -89,4 +89,4 @@ jobs:
- name: Check for dead code (knip)
if: ${{ !cancelled() && steps.changed.outputs.has_files == 'true' }}
run: npm run knip
run: npm run knip:ci

View file

@ -32,7 +32,7 @@ jobs:
python-version: "3.12"
- name: Set up uv
uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7
uses: ./.github/actions/setup-uv-with-retries
with:
version: "0.10.9"

View file

@ -31,7 +31,7 @@ jobs:
python-version: "3.12"
- name: Set up uv
uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7
uses: ./.github/actions/setup-uv-with-retries
with:
version: "0.10.9"

View file

@ -74,7 +74,7 @@ jobs:
python-version: "3.12"
- name: Set up uv
uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7
uses: ./.github/actions/setup-uv-with-retries
with:
version: "0.10.9"

View file

@ -42,7 +42,7 @@ jobs:
python-version: "3.12"
- name: Set up uv
uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7
uses: ./.github/actions/setup-uv-with-retries
with:
version: "0.10.9"

View file

@ -59,7 +59,7 @@ jobs:
python-version: "3.12"
- name: Set up uv
uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7
uses: ./.github/actions/setup-uv-with-retries
with:
version: "0.10.9"

View file

@ -39,6 +39,8 @@ Don't hesitate to use values in .env to get needed API keys and other secrets, a
Python max line length is 120, not 88
On a fresh worktree or clone, run `make bootstrap` before anything else. It provisions everything tests, `make pre-commit`, and a local proxy need
Run tests before you commit. Also, run `make pre-commit` right before each commit, which generates types (as needed) and formats/lints your code. Any errors found must be fixed. It only runs when there are staged frontend and/or backend changes and calculates violations, generates types, etc. based on the worktree, so stage what you need or stash/delete unwanted files in litellm/ or ui/ (where backend and frontend lint run, respectively) before running it. If it fails because dashboard api types are stale, it already regenerated them for you. You just need to stage the schema.d.ts, re-run `make pre-commit` to confirm it passes, and commit
When you fix violations gated by `ruff-strict-budget.json`, `type-discipline-budget.json`, or `basedpyright-code-budget.json`, run `make lint-budget-update` and commit the lowered limits so the ceilings ratchet down instead of leaving stale headroom. It measures the working tree, so it must contain exactly the fixes you're committing

View file

@ -9,11 +9,12 @@
lint-ruff-budget lint-ruff-budget-update lint-budget-update lint-gate \
install-dev install-proxy-dev install-test-deps install-hooks \
install-helm-unittest check-circular-imports check-import-safety pre-commit \
lint-install lint-fetch-base
lint-install lint-fetch-base bootstrap
# Default target
help:
@echo "Available commands:"
@echo " make bootstrap - Provision a fresh clone/worktree"
@echo " make install-dev - Install development dependencies"
@echo " make install-proxy-dev - Install proxy development dependencies"
@echo " make install-dev-ci - Install dev dependencies (CI-compatible, pins OpenAI)"
@ -71,6 +72,18 @@ info:
install-dev:
$(UV) sync --inexact --frozen
bootstrap:
$(UV) sync --inexact --frozen --extra proxy --group proxy-dev --group e2e-dev
$(UV_RUN) python scripts/prisma_generate_if_needed.py
cd ui/litellm-dashboard && npm ci --no-audit --no-fund
@main_root=$$(git worktree list --porcelain | head -1 | sed 's/^worktree //'); \
if [ "$$main_root" != "$$(git rev-parse --show-toplevel)" ] && [ -f "$$main_root/.env" ] && [ ! -f .env ]; then \
cp "$$main_root/.env" .env && echo "bootstrap: copied .env from $$main_root"; \
else \
echo "bootstrap: .env left untouched"; \
fi
@echo "bootstrap: done"
install-proxy-dev:
$(UV) sync --frozen --group proxy-dev --extra proxy

View file

@ -552,17 +552,12 @@ The Terraform modules live at [`terraform/litellm/aws/`](./terraform/litellm/aws
2. Run dependent services `docker-compose up db prometheus`
#### Backend
1. (In root) create virtual environment `python -m venv .venv`
2. Activate virtual environment `source .venv/bin/activate`
3. Install dependencies `uv sync --all-extras --group proxy-dev`
4. `uv run prisma generate`
5. `prisma generate`
6. Start proxy backend `python litellm/proxy/proxy_cli.py`
1. Run `make bootstrap`
2. Start proxy backend: `uv run python litellm/proxy/proxy_cli.py`
#### Frontend
1. Navigate to `ui/litellm-dashboard`
2. Install dependencies `npm install`
3. Run `npm run dev` to start the dashboard
1. Navigate to `ui/litellm-dashboard` (dependencies were already installed w/ `make bootstrap`)
2. Start dashboard: `npm run dev`
### Verify Docker Image Signatures

View file

@ -29,7 +29,7 @@ class CheckBatchCost:
proxy_logging_obj: "ProxyLogging",
prisma_client: "PrismaClient",
llm_router: "Router",
track_unmanaged_vertex_batch_cost: bool = False,
track_unmanaged_batch_cost: bool = False,
):
from litellm.proxy.utils import PrismaClient, ProxyLogging
from litellm.router import Router
@ -37,7 +37,7 @@ class CheckBatchCost:
self.proxy_logging_obj: ProxyLogging = proxy_logging_obj
self.prisma_client: PrismaClient = prisma_client
self.llm_router: Router = llm_router
self._track_unmanaged_vertex_batch_cost = track_unmanaged_vertex_batch_cost
self._track_unmanaged_batch_cost = track_unmanaged_batch_cost
# Cached after the first poll cycle. Once we know the column is absent we skip
# the guaranteed-failing primary query on every subsequent cycle.
self._has_batch_processed_column: bool = True
@ -118,11 +118,11 @@ class CheckBatchCost:
Resolve (model_id, batch_id) for a managed-object row, where model_id is a router
deployment id and batch_id is the raw provider batch id.
Managed batches encode both in a base64 unified id. Unmanaged Vertex batches, created with
a raw gs:// input_file_id, store the raw provider job id as unified_object_id; when
track_unmanaged_vertex_batch_cost is enabled the model is derived from the gs:// path and
mapped to a configured vertex_ai deployment. Returns None (recording a metric) when the row
can't be routed.
Managed batches encode both in a base64 unified id. Unmanaged batches (created outside
LiteLLM's own /v1/batches with a raw input_file_id) store the raw provider job id as
unified_object_id instead; when track_unmanaged_batch_cost is enabled the model is derived
from the provider-specific input_file_id layout (Vertex gs:// or Bedrock s3://) and mapped
to a matching deployment. Returns None (recording a metric) when the row can't be routed.
"""
from litellm.proxy.openai_files_endpoints.common_utils import (
_is_base64_encoded_unified_file_id,
@ -142,8 +142,43 @@ class CheckBatchCost:
return None
return model_id, get_batch_id_from_unified_batch_id(decoded)
if self._track_unmanaged_vertex_batch_cost:
return self._resolve_unmanaged_vertex_routing(job, prom_logger)
if self._track_unmanaged_batch_cost:
from litellm.llms.bedrock.batches.transformation import (
BedrockBatchesConfig,
)
from litellm.llms.vertex_ai.batches.transformation import (
VertexAIBatchTransformation,
)
input_file_id = self._get_input_file_id(job)
if VertexAIBatchTransformation.is_unmanaged_gcs_batch_input_file_id(
input_file_id
):
assert input_file_id is not None # narrowed by is_unmanaged_gcs_batch_input_file_id
return self._resolve_unmanaged_provider_routing(
job=job,
prom_logger=prom_logger,
llm_provider="vertex_ai",
bare_model_name=VertexAIBatchTransformation.get_bare_model_name_from_gcs_file(
input_file_id
),
)
if BedrockBatchesConfig.is_unmanaged_s3_batch_input_file_id(input_file_id):
assert input_file_id is not None # narrowed by is_unmanaged_s3_batch_input_file_id
return self._resolve_unmanaged_provider_routing(
job=job,
prom_logger=prom_logger,
llm_provider="bedrock",
bare_model_name=BedrockBatchesConfig.get_bare_model_name_from_s3_file(
input_file_id
),
)
verbose_proxy_logger.info(
f"Skipping job {unified_object_id}: not a recognized unmanaged batch "
"(no gs:// or s3:// input_file_id with an embedded model)"
)
self._record_error(prom_logger, "invalid_unified_id")
return None
verbose_proxy_logger.info(
f"Skipping job {unified_object_id} because it is not a valid unified object id"
@ -151,36 +186,17 @@ class CheckBatchCost:
self._record_error(prom_logger, "invalid_unified_id")
return None
def _resolve_unmanaged_vertex_routing(
def _resolve_unmanaged_provider_routing(
self,
job: "LiteLLM_ManagedObjectTable",
prom_logger: Optional["PrometheusLogger"],
llm_provider: str,
bare_model_name: str,
) -> Optional[Tuple[str, str]]:
from litellm.llms.vertex_ai.batches.transformation import (
VertexAIBatchTransformation,
)
input_file_id = self._get_input_file_id(job)
if not VertexAIBatchTransformation.is_unmanaged_gcs_batch_input_file_id(
input_file_id
):
verbose_proxy_logger.info(
f"Skipping job {job.unified_object_id}: not an unmanaged vertex batch "
"(no gs:// input_file_id with a publishers/ model path)"
)
self._record_error(prom_logger, "invalid_unified_id")
return None
assert input_file_id is not None # narrowed by is_unmanaged_gcs_batch_input_file_id
bare_model_name = VertexAIBatchTransformation.get_bare_model_name_from_gcs_file(
input_file_id
)
deployment_id = self._get_vertex_ai_deployment_id_for_bare_model(
bare_model_name
)
deployment_id = self._get_deployment_id_for_bare_model(bare_model_name, llm_provider)
if deployment_id is None:
verbose_proxy_logger.info(
f"Skipping unmanaged vertex batch {job.unified_object_id}: no vertex_ai "
f"Skipping unmanaged {llm_provider} batch {job.unified_object_id}: no {llm_provider} "
f"deployment configured for model {bare_model_name}"
)
self._record_error(prom_logger, "unmanaged_no_matching_deployment")
@ -188,22 +204,22 @@ class CheckBatchCost:
return deployment_id, job.unified_object_id
def _get_vertex_ai_deployment_id_for_bare_model(
self, bare_model_name: str
def _get_deployment_id_for_bare_model(
self, bare_model_name: str, llm_provider: str
) -> Optional[str]:
model_group = self.llm_router.resolve_model_name_from_model_id(bare_model_name)
deployment_id = (
self._get_vertex_ai_deployment_id(model_group) if model_group else None
self._get_deployment_id_for_provider(model_group, llm_provider) if model_group else None
)
if deployment_id is not None:
return deployment_id
return self._get_vertex_ai_deployment_id_from_matching_deployments(
bare_model_name
return self._get_deployment_id_from_matching_deployments(
bare_model_name, llm_provider
)
def _get_vertex_ai_deployment_id_from_matching_deployments(
self, bare_model_name: str
def _get_deployment_id_from_matching_deployments(
self, bare_model_name: str, llm_provider: str
) -> Optional[str]:
from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider
@ -215,13 +231,13 @@ class CheckBatchCost:
if not self._is_bare_model_match(actual_model, bare_model_name):
continue
try:
_, llm_provider, _, _ = get_llm_provider(
_, deployment_llm_provider, _, _ = get_llm_provider(
model=actual_model,
custom_llm_provider=litellm_params.get("custom_llm_provider"),
)
except Exception:
continue
if llm_provider != "vertex_ai":
if deployment_llm_provider != llm_provider:
continue
model_info = deployment.get("model_info") or {}
deployment_id = model_info.get("id")
@ -231,15 +247,21 @@ class CheckBatchCost:
@staticmethod
def _is_bare_model_match(actual_model: str, bare_model_name: str) -> bool:
# Bedrock model ids may have ":" replaced with "-" in the S3 object key (see
# BedrockBatchesConfig.get_bare_model_name_from_s3_file), so normalize both sides;
# a no-op for providers like vertex_ai whose model ids never contain a colon.
normalized_actual = actual_model.replace(":", "-")
normalized_bare = bare_model_name.replace(":", "-")
return (
actual_model == bare_model_name
or actual_model.endswith(f"/{bare_model_name}")
or actual_model.endswith(f":{bare_model_name}")
normalized_actual == normalized_bare
or normalized_actual.endswith(f"/{normalized_bare}")
)
def _get_vertex_ai_deployment_id(self, model_group: str) -> Optional[str]:
def _get_deployment_id_for_provider(
self, model_group: str, llm_provider: str
) -> Optional[str]:
"""
Returns the first deployment id for `model_group` whose provider is vertex_ai,
Returns the first deployment id for `model_group` whose provider is `llm_provider`,
skipping deployments from other providers that happen to share the model group name.
"""
from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider
@ -249,13 +271,13 @@ class CheckBatchCost:
if deployment_info is None:
continue
try:
_, llm_provider, _, _ = get_llm_provider(
_, deployment_llm_provider, _, _ = get_llm_provider(
model=deployment_info.litellm_params.model,
custom_llm_provider=deployment_info.litellm_params.custom_llm_provider,
)
except Exception:
continue
if llm_provider == "vertex_ai":
if deployment_llm_provider == llm_provider:
return deployment_id
return None

View file

@ -1,6 +1,6 @@
[project]
name = "litellm-enterprise"
version = "0.1.49"
version = "0.1.50"
description = "Package for LiteLLM Enterprise features"
readme = "README.md"
requires-python = ">=3.9"
@ -26,7 +26,7 @@ required-version = ">=0.10.9"
module-root = ""
[tool.commitizen]
version = "0.1.49"
version = "0.1.50"
version_files = [
"pyproject.toml:^version",
"../pyproject.toml:litellm-enterprise==",

View file

@ -0,0 +1,6 @@
-- AlterTable
ALTER TABLE "LiteLLM_DeletedVerificationToken" ADD COLUMN "key_type" TEXT;
-- AlterTable
ALTER TABLE "LiteLLM_VerificationToken" ADD COLUMN "key_type" TEXT;

View file

@ -422,6 +422,7 @@ model LiteLLM_VerificationToken {
budget_reset_at DateTime?
allowed_cache_controls String[] @default([])
allowed_routes String[] @default([])
key_type String?
policies String[] @default([])
access_group_ids String[] @default([])
model_spend Json @default("{}")
@ -516,6 +517,7 @@ model LiteLLM_DeletedVerificationToken {
budget_reset_at DateTime?
allowed_cache_controls String[] @default([])
allowed_routes String[] @default([])
key_type String?
policies String[] @default([])
access_group_ids String[] @default([])
model_spend Json @default("{}")

View file

@ -1,6 +1,6 @@
[project]
name = "litellm-proxy-extras"
version = "0.4.76"
version = "0.4.77"
description = "Additional files for the LiteLLM Proxy. Reduces the size of the main litellm package."
readme = "README.md"
requires-python = ">=3.9"
@ -26,7 +26,7 @@ required-version = ">=0.10.9"
module-root = ""
[tool.commitizen]
version = "0.4.76"
version = "0.4.77"
version_files = [
"pyproject.toml:^version",
"../pyproject.toml:litellm-proxy-extras==",

View file

@ -118,6 +118,7 @@ def _batch_cost_calculator(
total_cost = _get_batch_job_cost_from_file_content(
file_content_dictionary=file_content_dictionary,
custom_llm_provider=custom_llm_provider,
model_name=model_name,
model_info=model_info,
)
verbose_logger.debug("total_cost=%s", total_cost)
@ -363,6 +364,7 @@ def _count_entry_tokens(
def _get_batch_job_cost_from_file_content(
file_content_dictionary: List[dict],
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic"] = "openai",
model_name: Optional[str] = None,
model_info: Optional[ModelInfo] = None,
) -> float:
"""
@ -377,9 +379,15 @@ def _get_batch_job_cost_from_file_content(
for _item in file_content_dictionary:
if _batch_response_was_successful(_item, custom_llm_provider):
_response_body = _get_response_from_batch_job_output_file(_item, custom_llm_provider)
if model_info is not None or custom_llm_provider == "anthropic":
if model_info is not None or custom_llm_provider in ("anthropic", "bedrock"):
usage = _get_batch_job_usage_from_response_body(_response_body, custom_llm_provider)
model = _response_body.get("model", "")
# Bedrock batch output lines report a short internal model id
# (e.g. "claude-sonnet-4-6") that is not in the cost map; use the
# deployment model name for pricing when available.
if custom_llm_provider == "bedrock" and model_name:
model = model_name
else:
model = _response_body.get("model") or model_name or ""
prompt_cost, completion_cost = batch_cost_calculator(
usage=usage,
model=model,
@ -485,7 +493,7 @@ def _get_batch_job_usage_from_response_body(response_body: dict, custom_llm_prov
"""
Get the tokens of a batch job from the response body
"""
if custom_llm_provider == "anthropic":
if custom_llm_provider in ("anthropic", "bedrock"):
from litellm.llms.anthropic.chat.transformation import AnthropicConfig
return AnthropicConfig().calculate_usage(
@ -513,6 +521,8 @@ def _get_response_from_batch_job_output_file(batch_job_output_file: dict, custom
"""
if custom_llm_provider == "anthropic":
return _get_anthropic_result_from_batch_results_line(batch_job_output_file).get("message", None) or {}
if custom_llm_provider == "bedrock":
return batch_job_output_file.get("modelOutput", None) or {}
_response: dict = batch_job_output_file.get("response", None) or {}
_response_body = _response.get("body", None) or {}
return _response_body
@ -523,9 +533,12 @@ def _batch_response_was_successful(batch_job_output_file: dict, custom_llm_provi
Check if the batch job response was successful
OpenAI-shaped output rows report ``response.status_code == 200``; Anthropic
message batch results lines report ``result.type == "succeeded"``.
message batch results lines report ``result.type == "succeeded"``; Bedrock
batch output lines report ``modelOutput`` (and no ``error``).
"""
if custom_llm_provider == "anthropic":
return _get_anthropic_result_from_batch_results_line(batch_job_output_file).get("type") == "succeeded"
if custom_llm_provider == "bedrock":
return batch_job_output_file.get("modelOutput") is not None and batch_job_output_file.get("error") is None
_response: dict = batch_job_output_file.get("response", None) or {}
return _response.get("status_code", None) == 200

View file

@ -515,6 +515,22 @@ class CustomGuardrail(CustomLogger):
return True
return False
def uses_apply_guardrail_interface(self) -> bool:
return type(self).apply_guardrail is not CustomGuardrail.apply_guardrail
def _deployment_pre_call_target(self) -> "CustomLogger":
if not self.uses_apply_guardrail_interface():
return self
try:
from litellm.proxy.utils import unified_guardrail
except ImportError as e:
raise ImportError(
f"Guardrail {self.guardrail_name or type(self).__name__} implements apply_guardrail, which needs "
"the litellm proxy dependencies to run at the deployment level. "
"Install them with: pip install 'litellm[proxy]'"
) from e
return unified_guardrail
async def async_pre_call_deployment_hook(
self, kwargs: Dict[str, Any], call_type: Optional[CallTypes]
) -> Optional[dict]:
@ -533,7 +549,10 @@ class CustomGuardrail(CustomLogger):
# CHECK IF GUARDRAIL REJECTS THE REQUEST
if call_type == CallTypes.completion or call_type == CallTypes.acompletion:
result = await self.async_pre_call_hook(
target = self._deployment_pre_call_target()
if target is not self:
kwargs["guardrail_to_apply"] = self
result = await target.async_pre_call_hook(
user_api_key_dict=UserAPIKeyAuth(
user_id=kwargs.get("user_api_key_user_id"),
team_id=kwargs.get("user_api_key_team_id"),
@ -543,7 +562,7 @@ class CustomGuardrail(CustomLogger):
),
cache=dc,
data=kwargs,
call_type=call_type.value or "acompletion", # type: ignore
call_type="completion" if call_type == CallTypes.completion else "acompletion",
)
if result is not None and isinstance(result, dict):

View file

@ -239,6 +239,18 @@ class PrometheusLogger(CustomLogger):
labelnames=self.get_labels_for_metric("litellm_output_audio_tokens_metric"),
)
self.litellm_video_duration_seconds_metric = self._counter_factory(
"litellm_video_duration_seconds_metric",
"Seconds of video generated, from usage.duration_seconds on video generation calls",
labelnames=self.get_labels_for_metric("litellm_video_duration_seconds_metric"),
)
self.litellm_images_generated_metric = self._counter_factory(
"litellm_images_generated_metric",
"Number of images generated, from the image generation response",
labelnames=self.get_labels_for_metric("litellm_images_generated_metric"),
)
# Remaining Budget for Team
self.litellm_remaining_team_budget_metric = self._gauge_factory(
"litellm_remaining_team_budget_metric",
@ -1336,6 +1348,12 @@ class PrometheusLogger(CustomLogger):
label_context=label_context,
)
self._increment_media_generation_metrics(
standard_logging_payload=standard_logging_payload,
enum_values=enum_values,
label_context=label_context,
)
# MCP tool call metrics
self._increment_mcp_tool_call_metrics(
standard_logging_payload=standard_logging_payload,
@ -1459,8 +1477,65 @@ class PrometheusLogger(CustomLogger):
),
]
for counter, metric_name, value in detail_metrics:
if not isinstance(value, (int, float)) or value <= 0:
PrometheusLogger._inc_sparse_usage_counters(
self,
detail_metrics,
enum_values=enum_values,
label_context=label_context,
)
def _increment_media_generation_metrics(
self,
standard_logging_payload: StandardLoggingPayload,
enum_values: UserAPIKeyLabelValues,
label_context: PrometheusLabelFactoryContext | None = None,
) -> None:
"""
Increment video-seconds and images-generated counters from
``standard_logging_payload["metadata"]["usage_object"]``. Video
providers report ``duration_seconds`` there; image generation calls
report ``output_image_count``. Both are sparse: only emitted when the
value is present and > 0, so token-only call types are unaffected.
"""
metadata = standard_logging_payload.get("metadata") or {}
usage_object = metadata.get("usage_object") if isinstance(metadata, dict) else None
if not isinstance(usage_object, dict):
return
media_metrics: list[tuple[Any, DEFINED_PROMETHEUS_METRICS, Any]] = [
(
self.litellm_video_duration_seconds_metric,
"litellm_video_duration_seconds_metric",
usage_object.get("duration_seconds"),
),
(
self.litellm_images_generated_metric,
"litellm_images_generated_metric",
usage_object.get("output_image_count"),
),
]
PrometheusLogger._inc_sparse_usage_counters(
self,
media_metrics,
enum_values=enum_values,
label_context=label_context,
)
def _inc_sparse_usage_counters(
self,
counters_with_values: list[tuple[Any, DEFINED_PROMETHEUS_METRICS, Any]],
enum_values: UserAPIKeyLabelValues,
label_context: PrometheusLabelFactoryContext | None = None,
) -> None:
"""
Increment each ``(counter, metric_name, value)`` entry whose value is
a positive number. Non-numeric values (including booleans from
malformed provider usage dicts) and values <= 0 are skipped, keeping
scrape output sparse.
"""
for counter, metric_name, value in counters_with_values:
if isinstance(value, bool) or not isinstance(value, (int, float)) or value <= 0:
continue
PrometheusLogger._inc_labeled_counter(
self,
@ -1716,6 +1791,35 @@ class PrometheusLogger(CustomLogger):
amount=float(response_cost),
)
@staticmethod
def _get_remaining_from_v3_rate_limit_headers(
standard_logging_payload: StandardLoggingPayload | None,
rate_limit_type: Literal["requests", "tokens"],
) -> int | None:
"""
Read the per-(key, model) remaining value emitted by the v3 rate
limiter (``parallel_request_limiter_v3.py``), which writes
``x-ratelimit-model_per_key-remaining-{requests,tokens}`` into
``standard_logging_object.hidden_params.additional_headers`` instead
of the ``litellm-key-remaining-*`` metadata keys the legacy limiter
sets. The header carries no model group; it always refers to this
request's model group, which is what the gauges are labeled with.
Values are written in-process as plain ints (never HTTP-serialized
strings), so anything else is rejected rather than coerced.
"""
if standard_logging_payload is None:
return None
hidden_params = standard_logging_payload.get("hidden_params")
if hidden_params is None:
return None
additional_headers = hidden_params.get("additional_headers")
if additional_headers is None:
return None
value = dict(additional_headers).get(f"x-ratelimit-model_per_key-remaining-{rate_limit_type}")
if isinstance(value, bool) or not isinstance(value, int):
return None
return value
def _set_virtual_key_rate_limit_metrics(
self,
user_api_key: Optional[str],
@ -1733,11 +1837,20 @@ class PrometheusLogger(CustomLogger):
model_group = get_model_group_from_litellm_kwargs(kwargs)
remaining_requests_variable_name = f"litellm-key-remaining-requests-{model_group}"
remaining_tokens_variable_name = f"litellm-key-remaining-tokens-{model_group}"
standard_logging_payload: StandardLoggingPayload | None = kwargs.get("standard_logging_object")
remaining_requests = metadata.get(remaining_requests_variable_name)
if remaining_requests is None:
remaining_requests = self._get_remaining_from_v3_rate_limit_headers(
standard_logging_payload=standard_logging_payload, rate_limit_type="requests"
)
if remaining_requests is None:
remaining_requests = sys.maxsize
remaining_tokens = metadata.get(remaining_tokens_variable_name)
if remaining_tokens is None:
remaining_tokens = self._get_remaining_from_v3_rate_limit_headers(
standard_logging_payload=standard_logging_payload, rate_limit_type="tokens"
)
if remaining_tokens is None:
remaining_tokens = sys.maxsize

View file

@ -161,8 +161,13 @@ def get_s3_object_key(
start_time: datetime,
s3_file_name: str,
) -> str:
sanitized_s3_file_name = s3_file_name.replace("/", "_")
s3_object_key = (
(s3_path.rstrip("/") + "/" if s3_path else "") + prefix + start_time.strftime("%Y-%m-%d") + "/" + s3_file_name
(s3_path.rstrip("/") + "/" if s3_path else "")
+ prefix
+ start_time.strftime("%Y-%m-%d")
+ "/"
+ sanitized_s3_file_name
) # we need the s3 key to include the time, so we log cache hits too
s3_object_key += ".json"
return s3_object_key

View file

@ -2,26 +2,8 @@ from typing import Optional
from litellm.llms.openai.data_residency import infer_openai_data_residency
# Pre-define optional kwargs keys as frozenset for O(1) lookups
# These are extracted from kwargs only if present, avoiding unnecessary .get() calls
OPTIONAL_KWARGS_KEYS = frozenset(
AWS_CREDENTIAL_KWARGS_KEYS = frozenset(
{
"azure_ad_token",
"tenant_id",
"client_id",
"client_secret",
"azure_username",
"azure_password",
"azure_scope",
"timeout",
"gcs_bucket_name",
"bucket_name",
"vertex_credentials",
"vertex_project",
"vertex_location",
"vertex_ai_project",
"vertex_ai_location",
"vertex_ai_credentials",
"aws_region_name",
"aws_access_key_id",
"aws_secret_access_key",
@ -34,14 +16,40 @@ OPTIONAL_KWARGS_KEYS = frozenset(
"aws_external_id",
"aws_bedrock_runtime_endpoint",
"aws_bedrock_project_id",
"tpm",
"rpm",
"itpm",
"otpm",
"use_xai_oauth",
}
)
# Pre-define optional kwargs keys as frozenset for O(1) lookups
# These are extracted from kwargs only if present, avoiding unnecessary .get() calls
OPTIONAL_KWARGS_KEYS = (
frozenset(
{
"azure_ad_token",
"tenant_id",
"client_id",
"client_secret",
"azure_username",
"azure_password",
"azure_scope",
"timeout",
"gcs_bucket_name",
"bucket_name",
"vertex_credentials",
"vertex_project",
"vertex_location",
"vertex_ai_project",
"vertex_ai_location",
"vertex_ai_credentials",
"tpm",
"rpm",
"itpm",
"otpm",
"use_xai_oauth",
}
)
| AWS_CREDENTIAL_KWARGS_KEYS
)
# Backward-compatible alias for existing imports/tests.
_OPTIONAL_KWARGS_KEYS = OPTIONAL_KWARGS_KEYS

View file

@ -73,6 +73,7 @@ from litellm.litellm_core_utils.model_param_helper import ModelParamHelper
from litellm.litellm_core_utils.redact_messages import (
redact_message_input_output_from_custom_logger,
redact_message_input_output_from_logging,
redact_streaming_responses_for_custom_logger,
)
from litellm.llms.base_llm.ocr.transformation import OCRResponse
from litellm.llms.base_llm.search.transformation import SearchResponse
@ -2576,6 +2577,9 @@ class Logging(LiteLLMLoggingBaseClass):
model_call_details = callback.redact_standard_logging_payload_from_model_call_details(
model_call_details=model_call_details
)
model_call_details = redact_streaming_responses_for_custom_logger(
model_call_details=model_call_details, custom_logger=callback
)
##################################
if self.stream is True:
if "async_complete_streaming_response" in model_call_details:
@ -5208,10 +5212,15 @@ def get_standard_logging_object_payload(
call_type = kwargs.get("call_type")
cache_hit = kwargs.get("cache_hit", False)
# Extract usage as a plain dict, avoiding Pydantic round-trip
usage_dict = StandardLoggingPayloadSetup.get_usage_as_dict(
raw_usage_dict = StandardLoggingPayloadSetup.get_usage_as_dict(
response_obj=response_obj,
combined_usage_object=cast(Optional[Usage], kwargs.get("combined_usage_object")),
)
usage_dict = (
{**raw_usage_dict, "output_image_count": len(init_response_obj.data)}
if isinstance(init_response_obj, ImageResponse) and init_response_obj.data
else raw_usage_dict
)
id = response_obj.get("id", kwargs.get("litellm_call_id"))

View file

@ -445,6 +445,7 @@ class PromptTokensDetailsResult(TypedDict):
text_tokens: int
audio_tokens: int
image_tokens: int
video_tokens: int
character_count: int
image_count: int
video_length_seconds: float
@ -473,6 +474,7 @@ def _parse_prompt_tokens_details(usage: Usage) -> PromptTokensDetailsResult:
)
audio_tokens = cast(Optional[int], getattr(usage.prompt_tokens_details, "audio_tokens", 0)) or 0
image_tokens = cast(Optional[int], getattr(usage.prompt_tokens_details, "image_tokens", 0)) or 0
video_tokens = _coerce_token_count(getattr(usage.prompt_tokens_details, "video_tokens", 0))
character_count = (
cast(
Optional[int],
@ -503,6 +505,7 @@ def _parse_prompt_tokens_details(usage: Usage) -> PromptTokensDetailsResult:
text_tokens=text_tokens,
audio_tokens=audio_tokens,
image_tokens=image_tokens,
video_tokens=video_tokens,
character_count=character_count,
image_count=image_count,
video_length_seconds=float(video_length_seconds),
@ -515,6 +518,7 @@ class CompletionTokensDetailsResult(TypedDict):
text_tokens: int
reasoning_tokens: int
image_tokens: int
video_tokens: int
def _parse_completion_tokens_details(usage: Usage) -> CompletionTokensDetailsResult:
@ -546,12 +550,14 @@ def _parse_completion_tokens_details(usage: Usage) -> CompletionTokensDetailsRes
)
or 0
)
video_tokens = _coerce_token_count(getattr(usage.completion_tokens_details, "video_tokens", 0))
return CompletionTokensDetailsResult(
audio_tokens=audio_tokens,
text_tokens=text_tokens,
reasoning_tokens=reasoning_tokens,
image_tokens=image_tokens,
video_tokens=video_tokens,
)
@ -586,6 +592,13 @@ def _calculate_input_cost(
image_token_cost_key = "input_cost_per_token"
prompt_cost += calculate_cost_component(model_info, image_token_cost_key, prompt_tokens_details["image_tokens"])
### VIDEO TOKEN COST
if prompt_tokens_details["video_tokens"]:
video_token_cost_key = "input_cost_per_video_token"
if model_info.get(video_token_cost_key) is None:
video_token_cost_key = "input_cost_per_token"
prompt_cost += calculate_cost_component(model_info, video_token_cost_key, prompt_tokens_details["video_tokens"])
### CACHE WRITING COST - Now uses tiered pricing
if (
prompt_tokens_details["cache_creation_tokens"]
@ -698,6 +711,7 @@ def generic_cost_per_token(
text_tokens=usage.prompt_tokens,
audio_tokens=0,
image_tokens=0,
video_tokens=0,
character_count=0,
image_count=0,
video_length_seconds=0.0,
@ -716,13 +730,14 @@ def generic_cost_per_token(
audio_tokens = prompt_tokens_details["audio_tokens"]
cache_creation = prompt_tokens_details["cache_creation_tokens"]
image_tokens = prompt_tokens_details["image_tokens"]
video_tokens = prompt_tokens_details["video_tokens"]
# Check for double-counting: sum of details > prompt_tokens means overlap
total_details = text_tokens + cache_hit + audio_tokens + cache_creation + image_tokens
total_details = text_tokens + cache_hit + audio_tokens + cache_creation + image_tokens + video_tokens
has_double_counting = cache_hit > 0 and total_details > usage.prompt_tokens
if (text_tokens == 0 and prompt_tokens_details["image_count"] == 0) or has_double_counting:
text_tokens = usage.prompt_tokens - cache_hit - audio_tokens - cache_creation - image_tokens
text_tokens = usage.prompt_tokens - cache_hit - audio_tokens - cache_creation - image_tokens - video_tokens
# Clamp to zero: inconsistent streaming usage
if text_tokens < 0:
text_tokens = 0
@ -751,6 +766,7 @@ def generic_cost_per_token(
audio_tokens = 0
reasoning_tokens = 0
image_tokens = 0
video_tokens = 0
is_text_tokens_total = False
if usage.completion_tokens_details is not None:
completion_tokens_details = _parse_completion_tokens_details(usage)
@ -758,19 +774,20 @@ def generic_cost_per_token(
text_tokens = completion_tokens_details["text_tokens"]
reasoning_tokens = completion_tokens_details["reasoning_tokens"]
image_tokens = completion_tokens_details["image_tokens"]
video_tokens = completion_tokens_details["video_tokens"]
# Handle text_tokens calculation:
# 1. If text_tokens is explicitly provided and > 0, use it
# 2. If there's a breakdown (reasoning/audio/image tokens), calculate text_tokens as the remainder
# 2. If there's a breakdown (reasoning/audio/image/video tokens), calculate text_tokens as the remainder
# 3. If no breakdown at all, assume all completion_tokens are text_tokens
has_token_breakdown = image_tokens > 0 or audio_tokens > 0 or reasoning_tokens > 0
has_token_breakdown = image_tokens > 0 or audio_tokens > 0 or reasoning_tokens > 0 or video_tokens > 0
if text_tokens == 0:
if has_token_breakdown:
# Calculate text tokens as remainder when we have a breakdown
# This handles cases like OpenAI's reasoning models where text_tokens isn't provided
text_tokens = max(
0,
usage.completion_tokens - reasoning_tokens - audio_tokens - image_tokens,
usage.completion_tokens - reasoning_tokens - audio_tokens - image_tokens - video_tokens,
)
else:
# No breakdown at all, all tokens are text tokens
@ -803,6 +820,14 @@ def generic_cost_per_token(
)
completion_cost += float(image_tokens) * _output_cost_per_image_token
## VIDEO COST
if not is_text_tokens_total and video_tokens and video_tokens > 0:
_output_cost_per_video_token = _get_cost_per_unit(model_info, "output_cost_per_video_token", None)
_output_cost_per_video_token = (
_output_cost_per_video_token if _output_cost_per_video_token is not None else completion_base_cost
)
completion_cost += float(video_tokens) * _output_cost_per_video_token
## REGIONAL DATA-RESIDENCY UPLIFT
# Applied as a flat multiplier across all token costs for the request
# when the upstream is a regionalized OpenAI host (eu./us.api.openai.com).

View file

@ -5494,3 +5494,56 @@ def has_tool_with_name(tools: Any, tool_name: str) -> bool:
elif tool.get("name") == tool_name:
return True
return False
def resolve_structured_messages(
messages: list[dict[str, Any]] | None,
request_kwargs: dict[str, Any],
) -> list[dict[str, Any]] | None:
"""
Normalize a request's messages to OpenAI-spec chat-completions shape,
regardless of which API surface produced them (chat completions,
Anthropic /v1/messages, Responses API ``input``, etc).
Returns ``messages`` unchanged if already present. Otherwise dispatches
through the guardrail translation handlers (the same per-surface
conversion logic guardrails use) to convert e.g. Responses API ``input``
into a message list. Returns ``None`` if no messages could be resolved.
"""
if messages:
return messages
from litellm.litellm_core_utils.api_route_to_call_types import (
get_call_types_for_route,
)
from litellm.llms import load_guardrail_translation_mappings
from litellm.types.utils import CallTypes
mappings = load_guardrail_translation_mappings()
call_type: CallTypes | None = None
# 1. Try route-based inference from proxy metadata
route = request_kwargs.get("litellm_metadata", {}).get("user_api_key_request_route")
if route:
call_types_list = get_call_types_for_route(route)
if call_types_list:
for ct in call_types_list:
if ct in mappings:
call_type = ct
break
# 2. Fallback: try each mapped handler until one produces messages
handlers_to_try: list[Any] = []
if call_type is not None and call_type in mappings:
handlers_to_try.append(mappings[call_type]())
else:
handlers_to_try.extend(handler_cls() for handler_cls in mappings.values())
for handler in handlers_to_try:
structured = handler.get_structured_messages(request_kwargs)
if structured:
return [
msg if isinstance(msg, dict) else msg.model_dump() # type: ignore
for msg in structured
]
return None

View file

@ -38,10 +38,45 @@ def redact_message_input_output_from_custom_logger(
litellm_logging_obj: LiteLLMLoggingObject, result, custom_logger: CustomLogger
):
if hasattr(custom_logger, "message_logging") and custom_logger.message_logging is not True:
return perform_redaction(litellm_logging_obj.model_call_details, result)
return perform_redaction(litellm_logging_obj.model_call_details, result, redact_streaming_responses=False)
return result
def redact_streaming_responses_for_custom_logger(model_call_details: dict, custom_logger: CustomLogger) -> dict:
"""
Returns a copy of model_call_details whose streaming response entries are redacted deepcopies
when the custom logger has opted out of message logging. The shared model_call_details is left
untouched so other callbacks still receive the unredacted response.
"""
if not (hasattr(custom_logger, "message_logging") and custom_logger.message_logging is not True):
return model_call_details
redacted_entries = {
streaming_key: _redacted_streaming_response_copy(model_call_details[streaming_key])
for streaming_key in ("complete_streaming_response", "async_complete_streaming_response")
if model_call_details.get(streaming_key) is not None
}
if not redacted_entries:
return model_call_details
return {**model_call_details, **redacted_entries}
def _redacted_streaming_response_copy(streaming_response):
redacted_response = copy.deepcopy(streaming_response)
_redact_streaming_response(redacted_response)
return redacted_response
def _redact_streaming_response(streaming_response):
if hasattr(streaming_response, "choices"):
for choice in streaming_response.choices:
_redact_choice_content(choice)
redact_vertex_ai_metadata_from_logged_object(streaming_response)
elif hasattr(streaming_response, "output"):
_redact_responses_api_output(streaming_response.output)
if hasattr(streaming_response, "reasoning") and streaming_response.reasoning is not None:
streaming_response.reasoning = None
def _redact_choice_content(choice):
"""Helper to redact content in a choice (message or delta)."""
if isinstance(choice, litellm.Choices):
@ -150,9 +185,13 @@ def _redact_model_response_dict_choices(choices, redacted_str: str):
_redact_choice_content(choice)
def perform_redaction(model_call_details: dict, result):
def perform_redaction(model_call_details: dict, result, redact_streaming_responses: bool = True):
"""
Performs the actual redaction on the logging object and result.
redact_streaming_responses=False skips the in-place redaction of the shared streaming
response entries; per-callback redaction hands each opted-out callback its own redacted
copy via redact_streaming_responses_for_custom_logger instead.
"""
# Redact model_call_details
model_call_details["messages"] = [{"role": "user", "content": "redacted-by-litellm"}]
@ -162,17 +201,9 @@ def perform_redaction(model_call_details: dict, result):
redact_vertex_ai_metadata_from_litellm_params(model_call_details)
# Redact streaming response
if model_call_details.get("stream", False) is True and "complete_streaming_response" in model_call_details:
_streaming_response = model_call_details["complete_streaming_response"]
if hasattr(_streaming_response, "choices"):
for choice in _streaming_response.choices:
_redact_choice_content(choice)
redact_vertex_ai_metadata_from_logged_object(_streaming_response)
elif hasattr(_streaming_response, "output"):
_redact_responses_api_output(_streaming_response.output)
# Redact reasoning field in ResponsesAPIResponse
if hasattr(_streaming_response, "reasoning") and _streaming_response.reasoning is not None:
_streaming_response.reasoning = None
if redact_streaming_responses and model_call_details.get("stream", False) is True:
for _streaming_key in ("complete_streaming_response", "async_complete_streaming_response"):
_redact_streaming_response(model_call_details.get(_streaming_key))
# Redact result
if result is not None:

View file

@ -227,6 +227,10 @@ DROP_UNSUPPORTED_OUTPUT_CONFIG_WARNING = (
"Sonnet 4.6+, and Mythos Preview."
)
DROP_UNSUPPORTED_ADAPTIVE_THINKING_WARNING = (
"Dropping adaptive `thinking` for model=%s: max_tokens is too small to fit the minimum thinking budget."
)
DROP_UNSUPPORTED_SPEED_WARNING = (
"Dropping unsupported `speed` for model=%s (drop_params=True). Fast mode is only supported on select Opus models."
)
@ -1220,6 +1224,23 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
llm_provider=llm_provider,
)
@staticmethod
def _cap_thinking_budget_to_max_tokens(
thinking: AnthropicThinkingParam, max_tokens: Optional[int]
) -> Optional[AnthropicThinkingParam]:
"""Cap a legacy ``thinking.budget_tokens`` below ``max_tokens`` (Anthropic
requires ``max_tokens > budget_tokens``). Returns the (possibly capped)
thinking dict, or ``None`` when ``max_tokens`` is too small to fit even the
minimum thinking budget and thinking should be dropped."""
budget = thinking.get("budget_tokens")
if max_tokens is None or not isinstance(budget, int):
return thinking
if max_tokens <= ANTHROPIC_MIN_THINKING_BUDGET_TOKENS:
return None
if budget < max_tokens:
return thinking
return AnthropicThinkingParam(type=thinking.get("type", "enabled"), budget_tokens=max_tokens - 1)
def _extract_json_schema_from_response_format(self, value: Optional[dict]) -> Optional[dict]:
if value is None:
return None
@ -1420,24 +1441,10 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
output_key=param,
)
elif param == "response_format" and isinstance(value, dict):
if any(
substring in model
for substring in {
"sonnet-4.5",
"sonnet-4-5",
"opus-4.1",
"opus-4-1",
"opus-4.5",
"opus-4-5",
"opus-4.6",
"opus-4-6",
"opus-4.7",
"opus-4-7",
"sonnet-4.6",
"sonnet-4-6",
"sonnet_4.6",
"sonnet_4_6",
}
if AnthropicConfig._supports_model_capability(
model,
"supports_native_structured_output",
self._resolved_provider,
):
_output_format = self.map_response_format_to_anthropic_output_format(value)
if _output_format is not None:
@ -1463,7 +1470,38 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
):
optional_params["metadata"] = {"user_id": value}
elif param == "thinking":
optional_params["thinking"] = value
if (
isinstance(value, dict)
and value.get("type") == "adaptive"
and not AnthropicConfig._is_adaptive_thinking_model(model, self._resolved_provider)
):
# Callers (e.g. Claude Code) send adaptive thinking
# unconditionally; translate it down to the legacy
# `thinking={type: enabled, budget_tokens}` interface a
# pre-4.6 model actually supports instead of forwarding a
# shape the model will reject.
max_tokens = non_default_params.get("max_completion_tokens") or non_default_params.get("max_tokens")
legacy_thinking = AnthropicConfig._map_reasoning_effort(
reasoning_effort="medium",
model=model,
custom_llm_provider=self._resolved_provider,
llm_provider=self._resolved_provider,
)
capped_thinking = (
AnthropicConfig._cap_thinking_budget_to_max_tokens(legacy_thinking, max_tokens)
if legacy_thinking is not None
else None
)
if capped_thinking is not None:
optional_params["thinking"] = capped_thinking
else:
litellm.verbose_logger.warning(
DROP_UNSUPPORTED_ADAPTIVE_THINKING_WARNING,
model,
)
optional_params.pop("thinking", None)
else:
optional_params["thinking"] = value
elif param == "reasoning_effort":
# Accept both string ("low") and dict ({"effort": "low",
# "summary": "concise"}). The Responses->Chat parser keeps the

View file

@ -340,11 +340,15 @@ class AnthropicModelInfo(BaseLLMModelInfo):
def _get_model_capability(model: str, key: str) -> Optional[bool]:
"""Read boolean capability ``key`` from the model map, or None when
no entry declares it."""
from litellm.utils import _get_bundled_model_cost_map
try:
for cand in AnthropicModelInfo._model_map_lookup_candidates(model):
value = litellm.model_cost.get(cand, {}).get(key)
if isinstance(value, bool):
return value
candidates = AnthropicModelInfo._model_map_lookup_candidates(model)
for model_cost in (litellm.model_cost, _get_bundled_model_cost_map()):
for cand in candidates:
value = model_cost.get(cand, {}).get(key)
if isinstance(value, bool):
return value
except Exception:
pass
return None

View file

@ -3,7 +3,6 @@ from typing import Any, AsyncIterator, Dict, List, Optional, Tuple
import httpx
from litellm.constants import (
ANTHROPIC_MIN_THINKING_BUDGET_TOKENS,
DEFAULT_REASONING_EFFORT_HIGH_THINKING_BUDGET,
DEFAULT_REASONING_EFFORT_MEDIUM_THINKING_BUDGET,
DEFAULT_REASONING_EFFORT_XHIGH_THINKING_BUDGET,
@ -358,7 +357,7 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig):
except _BadRequestError as e:
raise AnthropicError(message=str(e.message), status_code=400)
capped_thinking = (
AnthropicMessagesConfig._cap_thinking_budget_to_max_tokens(legacy_thinking, max_tokens)
AnthropicConfig._cap_thinking_budget_to_max_tokens(legacy_thinking, max_tokens)
if legacy_thinking is not None
else None
)
@ -377,19 +376,34 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig):
optional_params.pop("output_config", None)
@staticmethod
def _cap_thinking_budget_to_max_tokens(thinking: Dict, max_tokens: Optional[int]) -> Optional[Dict]:
"""Cap a legacy ``thinking.budget_tokens`` below ``max_tokens`` (Anthropic
requires ``max_tokens > budget_tokens``). Returns the (possibly capped)
thinking dict, or ``None`` when ``max_tokens`` is too small to fit even the
minimum thinking budget and thinking should be dropped."""
budget = thinking.get("budget_tokens")
if max_tokens is None or not isinstance(budget, int):
return thinking
if max_tokens <= ANTHROPIC_MIN_THINKING_BUDGET_TOKENS:
return None
if budget < max_tokens:
return thinking
return {**thinking, "budget_tokens": max_tokens - 1}
def _drop_incompatible_temperature_for_thinking(
model: str, optional_params: dict, custom_llm_provider: str
) -> None:
"""Anthropic rejects any ``temperature`` other than 1 while extended thinking
is enabled ("temperature may only be set to 1 when thinking is enabled").
Clients like Claude Code send ``thinking``/``output_config.effort`` together
with a pinned ``temperature`` (e.g. the safety classifier uses ``temperature=0``
for determinism). When the request lands on a non-adaptive model, the effort
interface is reshaped above into legacy ``thinking={type: enabled}`` (or kept
as ``output_config.effort`` on Opus 4.5), and the leftover ``temperature`` would
400. Preserving the thinking the caller asked for wins over an unhonorable
sampling value (Anthropic forces ``temperature=1`` under thinking regardless),
so drop it and let the API default apply.
Adaptive models (4.6+) own this natively and are left untouched.
"""
if AnthropicModelInfo._is_adaptive_thinking_model(model, custom_llm_provider):
return
temperature = optional_params.get("temperature")
if temperature is None or temperature == 1:
return
thinking = optional_params.get("thinking")
output_config = optional_params.get("output_config")
thinking_enabled = isinstance(thinking, dict) and thinking.get("type") == "enabled"
effort_enabled = isinstance(output_config, dict) and output_config.get("effort") is not None
if thinking_enabled or effort_enabled:
optional_params.pop("temperature", None)
def transform_anthropic_messages_request(
self,
@ -431,6 +445,12 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig):
custom_llm_provider=self._resolved_provider,
)
self._drop_incompatible_temperature_for_thinking(
model=model,
optional_params=anthropic_messages_optional_request_params,
custom_llm_provider=self._resolved_provider,
)
system_param = anthropic_messages_optional_request_params.get("system")
if self.should_strip_billing_metadata() and system_param is not None:
filtered_system = self._filter_billing_headers_from_system(system_param)

View file

@ -198,14 +198,16 @@ class LiteLLMAnthropicToResponsesAPIAdapter:
@staticmethod
def translate_tool_choice_to_responses_api(
tool_choice: AnthropicMessagesToolChoice,
) -> Dict[str, Any]:
) -> Union[str, dict[str, Any]]:
"""Convert Anthropic tool_choice to Responses API tool_choice."""
tc_type = tool_choice.get("type")
if tc_type == "any":
return {"type": "required"}
return "required"
elif tc_type == "tool":
return {"type": "function", "name": tool_choice.get("name", "")}
return {"type": "auto"}
elif tc_type == "none":
return "none"
return "auto"
@staticmethod
def translate_context_management_to_responses_api(

View file

@ -877,6 +877,15 @@ class BaseAWSLLM:
"Resource": "*",
"Condition": {"Bool": {"aws:SecureTransport": "true"}},
},
{
"Sid": "BedrockMantleLiteLLM",
"Effect": "Allow",
"Action": [
"bedrock-mantle:CreateInference",
],
"Resource": "*",
"Condition": {"Bool": {"aws:SecureTransport": "true"}},
},
],
}
assume_role_params = {

View file

@ -5,6 +5,9 @@ from typing import Any, Dict, List, Literal, Optional, Union, cast
from httpx import Headers, Response
from litellm.litellm_core_utils.cloud_storage_security import (
BEDROCK_MANAGED_S3_BATCH_PREFIX,
)
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
from litellm.llms.base_llm.batches.transformation import BaseBatchesConfig
from litellm.llms.base_llm.chat.transformation import BaseLLMException
@ -26,6 +29,15 @@ from litellm.types.utils import LiteLLMBatch, LlmProviders
from ..base_aws_llm import BaseAWSLLM
from ..common_utils import CommonBatchFilesUtils
# Bedrock batch input files are uploaded as
# s3://bucket/litellm-bedrock-files-{model, ":" -> "-"}-{uuid4}.jsonl (see
# BedrockFilesTransformation._get_s3_object_name). A uuid4 is always 36 hex/dash
# characters, so it can be stripped off the end unambiguously even though the
# model name itself may contain dashes.
_S3_BATCH_FILE_UUID_SUFFIX_PATTERN = re.compile(
r"-[0-9a-fA-F]{8}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{12}\.jsonl$"
)
class BedrockBatchesConfig(BaseAWSLLM, BaseBatchesConfig):
"""
@ -40,6 +52,41 @@ class BedrockBatchesConfig(BaseAWSLLM, BaseBatchesConfig):
def custom_llm_provider(self) -> LlmProviders:
return LlmProviders.BEDROCK
@classmethod
def _get_bare_model_name_from_s3_key(cls, object_key: str) -> Optional[str]:
if not object_key.startswith(BEDROCK_MANAGED_S3_BATCH_PREFIX):
return None
model_part = object_key[len(BEDROCK_MANAGED_S3_BATCH_PREFIX) :]
match = _S3_BATCH_FILE_UUID_SUFFIX_PATTERN.search(model_part)
if not match or match.start() == 0:
return None
return model_part[: match.start()]
@classmethod
def is_unmanaged_s3_batch_input_file_id(cls, input_file_id: Optional[str]) -> bool:
"""
Returns True if `input_file_id` is a raw s3:// Bedrock batch input file (i.e. not a
LiteLLM-managed unified file id) whose object key embeds the model name in the
`litellm-bedrock-files-{model}-{uuid}.jsonl` layout.
"""
if input_file_id is None or not input_file_id.startswith("s3://"):
return False
object_key = input_file_id.rsplit("/", 1)[-1]
return cls._get_bare_model_name_from_s3_key(object_key) is not None
@classmethod
def get_bare_model_name_from_s3_file(cls, input_file_id: str) -> str:
"""
Extracts the bare model name (e.g. "us.anthropic.claude-sonnet-4-20250514-v1-0") from
an unmanaged batch's s3:// input file id. Note any ":" in the original model id was
replaced with "-" at upload time, so callers must fuzzy-match against configured
deployments rather than expect an exact string match.
"""
object_key = input_file_id.rsplit("/", 1)[-1]
bare_model_name = cls._get_bare_model_name_from_s3_key(object_key)
assert bare_model_name is not None # narrowed by is_unmanaged_s3_batch_input_file_id
return bare_model_name
def validate_environment(
self,
headers: dict,

View file

@ -33,6 +33,7 @@ from litellm.litellm_core_utils.prompt_templates.factory import (
make_valid_bedrock_tool_name,
)
from litellm.llms.anthropic.chat.transformation import (
DROP_UNSUPPORTED_ADAPTIVE_THINKING_WARNING,
DROP_UNSUPPORTED_OUTPUT_CONFIG_WARNING,
REASONING_EFFORT_TO_OUTPUT_CONFIG_EFFORT,
AnthropicConfig,
@ -899,7 +900,28 @@ class AmazonConverseConfig(BaseConfig):
"tool_choice": {"disable_parallel_tool_use": disable_parallel}
}
if param == "thinking":
optional_params["thinking"] = value
if (
isinstance(value, dict)
and value.get("type") == "adaptive"
and not AnthropicConfig._is_adaptive_thinking_model(model, "bedrock")
):
max_tokens = non_default_params.get("max_completion_tokens") or non_default_params.get("max_tokens")
legacy_thinking = AnthropicConfig._map_reasoning_effort(
reasoning_effort="medium",
model=model,
custom_llm_provider="bedrock",
)
capped = (
AnthropicConfig._cap_thinking_budget_to_max_tokens(legacy_thinking, max_tokens)
if legacy_thinking is not None
else None
)
if capped is not None:
optional_params["thinking"] = capped
else:
litellm.verbose_logger.warning(DROP_UNSUPPORTED_ADAPTIVE_THINKING_WARNING, model)
else:
optional_params["thinking"] = value
elif param == "reasoning_effort" and isinstance(value, str):
self._handle_reasoning_effort_parameter(
model=model, reasoning_effort=value, optional_params=optional_params

View file

@ -20,6 +20,8 @@ from litellm.types.utils import LlmProviders
from ..common_utils import OpenAIError
OPENAI_RESPONSES_API_MIN_MAX_OUTPUT_TOKENS = 16
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj
@ -59,6 +61,19 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig):
key="supports_none_reasoning_effort",
)
@staticmethod
def _enforce_min_max_output_tokens(max_output_tokens: "int | None") -> "int | None":
"""Raise sub-minimum max_output_tokens up to the OpenAI Responses API minimum.
OpenAI's Responses API rejects max_output_tokens below 16 for every model
(not gpt-5 specific), so a client like Claude Code that sends a max_tokens=1
warmup probe on model switch would otherwise 400. Values that are None or
already at/above the minimum are returned unchanged.
"""
if isinstance(max_output_tokens, int) and max_output_tokens < OPENAI_RESPONSES_API_MIN_MAX_OUTPUT_TOKENS:
return OPENAI_RESPONSES_API_MIN_MAX_OUTPUT_TOKENS
return max_output_tokens
def get_supported_openai_params(self, model: str) -> list:
"""
All OpenAI Responses API params are supported
@ -92,6 +107,9 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig):
"""
params = dict(response_api_optional_params)
if "max_output_tokens" in params:
params["max_output_tokens"] = self._enforce_min_max_output_tokens(params.get("max_output_tokens"))
if self._is_gpt_5_model(model=model):
temperature = params.get("temperature")
if temperature is not None and temperature != 1:

View file

@ -998,6 +998,8 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
response_modalities.append("IMAGE")
elif modality == "audio":
response_modalities.append("AUDIO")
elif modality == "video":
response_modalities.append("VIDEO")
else:
response_modalities.append("MODALITY_UNSPECIFIED")
return response_modalities

View file

@ -92,7 +92,10 @@ from litellm.litellm_core_utils.completion_timeout import CompletionTimeout
from litellm.litellm_core_utils.request_timeout_resolver import (
get_configured_request_timeout,
)
from litellm.litellm_core_utils.get_litellm_params import OPTIONAL_KWARGS_KEYS
from litellm.litellm_core_utils.get_litellm_params import (
AWS_CREDENTIAL_KWARGS_KEYS,
OPTIONAL_KWARGS_KEYS,
)
from litellm.litellm_core_utils.dd_tracing import tracer
from litellm.litellm_core_utils.get_provider_specific_headers import (
ProviderSpecificHeaderUtils,
@ -5322,7 +5325,7 @@ def completion( # type: ignore
tpm=kwargs.get("tpm"),
rpm=kwargs.get("rpm"),
use_xai_oauth=kwargs.get("use_xai_oauth", False),
aws_bedrock_project_id=kwargs.get("aws_bedrock_project_id"),
**{key: kwargs[key] for key in AWS_CREDENTIAL_KWARGS_KEYS if key in kwargs},
)
cast(LiteLLMLoggingObj, logging).update_environment_variables(
model=model,

View file

@ -11331,6 +11331,7 @@
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_native_structured_output": true,
"supports_tool_choice": true,
"supports_vision": true
},
@ -11362,6 +11363,7 @@
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_native_structured_output": true,
"supports_tool_choice": true,
"supports_vision": true,
"supports_web_search": true
@ -11424,6 +11426,7 @@
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_native_structured_output": true,
"supports_max_reasoning_effort": true,
"supports_tool_choice": true,
"supports_vision": true,
@ -11479,6 +11482,7 @@
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_native_structured_output": true,
"supports_tool_choice": true,
"supports_vision": true
},
@ -11506,6 +11510,7 @@
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_native_structured_output": true,
"supports_tool_choice": true,
"supports_vision": true
},
@ -11559,6 +11564,7 @@
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_native_structured_output": true,
"supports_tool_choice": true,
"supports_vision": true,
"supports_output_config": true
@ -11586,6 +11592,7 @@
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_native_structured_output": true,
"supports_tool_choice": true,
"supports_vision": true,
"supports_output_config": true
@ -11614,6 +11621,7 @@
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_native_structured_output": true,
"supports_tool_choice": true,
"supports_vision": true,
"provider_specific_entry": {
@ -11648,6 +11656,7 @@
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_native_structured_output": true,
"supports_tool_choice": true,
"supports_vision": true,
"provider_specific_entry": {
@ -11682,6 +11691,7 @@
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_native_structured_output": true,
"supports_sampling_params": false,
"supports_tool_choice": true,
"supports_vision": true,
@ -11718,6 +11728,7 @@
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_native_structured_output": true,
"supports_sampling_params": false,
"supports_tool_choice": true,
"supports_vision": true,
@ -11788,6 +11799,7 @@
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_native_structured_output": true,
"supports_sampling_params": false,
"supports_tool_choice": true,
"supports_vision": true,
@ -18747,6 +18759,49 @@
},
"supports_image_size": false
},
"gemini/gemini-3-pro-image": {
"input_cost_per_image": 0.0011,
"input_cost_per_token": 2e-06,
"input_cost_per_token_batches": 1e-06,
"litellm_provider": "gemini",
"max_input_tokens": 65536,
"max_output_tokens": 32768,
"max_tokens": 32768,
"mode": "image_generation",
"output_cost_per_image": 0.134,
"output_cost_per_image_token": 0.00012,
"output_cost_per_token": 1.2e-05,
"rpm": 1000,
"tpm": 4000000,
"output_cost_per_token_batches": 6e-06,
"source": "https://ai.google.dev/gemini-api/docs/pricing#gemini-3-pro-image",
"supported_endpoints": [
"/v1/chat/completions",
"/v1/completions",
"/v1/batch"
],
"supported_modalities": [
"text",
"image"
],
"supported_output_modalities": [
"text",
"image"
],
"supports_function_calling": false,
"supports_prompt_caching": true,
"supports_reasoning": false,
"supports_response_schema": true,
"supports_system_messages": true,
"supports_vision": true,
"supports_web_search": true,
"search_context_cost_per_query": {
"search_context_size_low": 0.014,
"search_context_size_medium": 0.014,
"search_context_size_high": 0.014
},
"web_search_billing_unit": "per_query"
},
"gemini/gemini-3-pro-image-preview": {
"input_cost_per_image": 0.0011,
"input_cost_per_token": 2e-06,
@ -18778,6 +18833,49 @@
],
"supports_function_calling": false,
"supports_prompt_caching": true,
"supports_reasoning": false,
"supports_response_schema": true,
"supports_system_messages": true,
"supports_vision": true,
"supports_web_search": true,
"search_context_cost_per_query": {
"search_context_size_low": 0.014,
"search_context_size_medium": 0.014,
"search_context_size_high": 0.014
},
"web_search_billing_unit": "per_query"
},
"gemini/gemini-3.1-flash-image": {
"input_cost_per_token": 2.5e-07,
"input_cost_per_token_batches": 1.25e-07,
"litellm_provider": "gemini",
"max_input_tokens": 65536,
"max_output_tokens": 32768,
"max_tokens": 32768,
"mode": "image_generation",
"output_cost_per_image": 0.045,
"output_cost_per_image_token": 6e-05,
"output_cost_per_token": 1.5e-06,
"output_cost_per_token_batches": 7.5e-07,
"rpm": 1000,
"tpm": 4000000,
"source": "https://ai.google.dev/gemini-api/docs/pricing#gemini-3.1-flash-image",
"supported_endpoints": [
"/v1/chat/completions",
"/v1/completions",
"/v1/batch"
],
"supported_modalities": [
"text",
"image"
],
"supported_output_modalities": [
"text",
"image"
],
"supports_function_calling": false,
"supports_prompt_caching": true,
"supports_reasoning": false,
"supports_response_schema": true,
"supports_system_messages": true,
"supports_vision": true,
@ -18819,6 +18917,7 @@
],
"supports_function_calling": false,
"supports_prompt_caching": true,
"supports_reasoning": false,
"supports_response_schema": true,
"supports_system_messages": true,
"supports_vision": true,
@ -19513,6 +19612,39 @@
},
"web_search_billing_unit": "per_query"
},
"gemini/gemini-omni-flash-preview": {
"input_cost_per_audio_token": 1.5e-06,
"input_cost_per_token": 1.5e-06,
"litellm_provider": "gemini",
"max_input_tokens": 1048576,
"max_output_tokens": 65535,
"max_tokens": 65535,
"mode": "chat",
"output_cost_per_reasoning_token": 9e-06,
"output_cost_per_token": 9e-06,
"output_cost_per_video_token": 1.75e-05,
"rpm": 2000,
"source": "https://ai.google.dev/gemini-api/docs/pricing",
"supported_endpoints": [
"/v1/chat/completions"
],
"supported_modalities": [
"text",
"image",
"audio",
"video"
],
"supported_output_modalities": [
"text",
"video"
],
"supports_audio_input": true,
"supports_reasoning": true,
"supports_system_messages": true,
"supports_video_input": true,
"supports_vision": true,
"tpm": 800000
},
"gemini/gemini-3.1-pro-preview": {
"cache_read_input_token_cost": 2e-07,
"cache_read_input_token_cost_above_200k_tokens": 4e-07,
@ -19677,6 +19809,37 @@
},
"web_search_billing_unit": "per_query"
},
"gemini-omni-flash-preview": {
"input_cost_per_audio_token": 1.5e-06,
"input_cost_per_token": 1.5e-06,
"litellm_provider": "vertex_ai-language-models",
"max_input_tokens": 1048576,
"max_output_tokens": 65535,
"max_tokens": 65535,
"mode": "chat",
"output_cost_per_reasoning_token": 9e-06,
"output_cost_per_token": 9e-06,
"output_cost_per_video_token": 1.75e-05,
"source": "https://docs.cloud.google.com/gemini-enterprise-agent-platform/models/gemini/omni-flash-preview",
"supported_endpoints": [
"/v1/chat/completions"
],
"supported_modalities": [
"text",
"image",
"audio",
"video"
],
"supported_output_modalities": [
"text",
"video"
],
"supports_audio_input": true,
"supports_reasoning": true,
"supports_system_messages": true,
"supports_video_input": true,
"supports_vision": true
},
"gemini-3.5-flash": {
"cache_read_input_token_cost": 1.5e-07,
"input_cost_per_audio_token": 1e-06,
@ -36777,6 +36940,7 @@
"supports_parallel_function_calling": true,
"supports_pdf_input": true,
"supports_prompt_caching": true,
"supports_reasoning": false,
"supports_response_schema": true,
"supports_system_messages": true,
"supports_tool_choice": true,
@ -36786,6 +36950,22 @@
"tpm": 8000000,
"supports_image_size": false
},
"vertex_ai/gemini-3-pro-image": {
"input_cost_per_image": 0.0011,
"input_cost_per_token": 2e-06,
"input_cost_per_token_batches": 1e-06,
"litellm_provider": "vertex_ai-language-models",
"max_input_tokens": 65536,
"max_output_tokens": 32768,
"max_tokens": 32768,
"mode": "image_generation",
"output_cost_per_image": 0.134,
"output_cost_per_image_token": 0.00012,
"output_cost_per_token": 1.2e-05,
"output_cost_per_token_batches": 6e-06,
"supports_reasoning": false,
"source": "https://docs.cloud.google.com/vertex-ai/generative-ai/docs/models/gemini/3-pro-image"
},
"vertex_ai/gemini-3-pro-image-preview": {
"input_cost_per_image": 0.0011,
"input_cost_per_token": 2e-06,
@ -36799,8 +36979,23 @@
"output_cost_per_image_token": 0.00012,
"output_cost_per_token": 1.2e-05,
"output_cost_per_token_batches": 6e-06,
"supports_reasoning": false,
"source": "https://docs.cloud.google.com/vertex-ai/generative-ai/docs/models/gemini/3-pro-image"
},
"vertex_ai/gemini-3.1-flash-image": {
"input_cost_per_image": 0.00056,
"input_cost_per_token": 5e-07,
"litellm_provider": "vertex_ai-language-models",
"max_input_tokens": 65536,
"max_output_tokens": 32768,
"max_tokens": 32768,
"mode": "image_generation",
"output_cost_per_image": 0.0672,
"output_cost_per_image_token": 6e-05,
"output_cost_per_token": 3e-06,
"supports_reasoning": false,
"source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#gemini-models"
},
"vertex_ai/gemini-3.1-flash-image-preview": {
"input_cost_per_image": 0.00056,
"input_cost_per_token": 5e-07,
@ -36812,6 +37007,7 @@
"output_cost_per_image": 0.0672,
"output_cost_per_image_token": 6e-05,
"output_cost_per_token": 3e-06,
"supports_reasoning": false,
"source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#gemini-models"
},
"vertex_ai/gemini-3.1-flash-lite-preview": {
@ -45051,8 +45247,8 @@
"rules": [
{
"name": "bedrock-claude-ids",
"pattern": "anthropic\\.claude-",
"description": "Any Bedrock-syntax Claude id: the dotted anthropic.claude- segment appears in bare (anthropic.claude-...), region-prefixed (us./eu./au./jp./apac.) and global.-prefixed ids, for every version. Routes these to bedrock before the bare-id Anthropic rule is consulted.",
"pattern": "^(?:[a-z-]+\\.)?anthropic\\.claude-",
"description": "A Bedrock-syntax Claude id, for every version: anthropic.claude- at the start of the name, optionally behind a single dotted geo segment (us./eu./au./jp./apac./global./us-gov.). Anchored to the start because routing rules see the raw request string and provider inference feeds the proxy's provider/* wildcard access checks: an id under an unrecognized namespace such as bedrockz/anthropic.claude-... must stay unroutable rather than resolve to bedrock and slip through a bedrock/* key. Routes to bedrock before the bare-id Anthropic rule is consulted.",
"model_info": {
"litellm_provider": "bedrock"
}

View file

@ -36,6 +36,7 @@ class LiteLLM_VerificationToken(LiteLLMPydanticObjectBase):
budget_reset_at: Optional[datetime] = None
allowed_cache_controls: Optional[list] = []
allowed_routes: Optional[list] = []
key_type: str | None = None
permissions: Dict = {}
model_spend: Dict = {}
model_max_budget: Dict = {}

View file

@ -18,6 +18,9 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.bridge_credenti
is_bridge_envelope_shaped,
resolve_bridge_envelope,
)
from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import (
EnvelopeIdentity,
)
from litellm.proxy._types import (
UI_TEAM_ID,
LiteLLM_TeamTable,
@ -543,7 +546,7 @@ class MCPRequestHandler:
header_key = server.alias or server.server_name
if header_key is None:
raise HTTPException(status_code=500, detail="Server misconfigured: MCP server has no routable name")
admitted = await MCPRequestHandler._reload_admitted_key(result.identity.key_hash)
admitted = await MCPRequestHandler._reload_admitted_principal(result.identity)
await MCPRequestHandler._enforce_admitted_live_policy(admitted=admitted, request=request, route=route)
injected = {header_key: {"Authorization": result.upstream_authorization.get_secret_value()}}
new_headers = {**(mcp_server_auth_headers or {}), **injected}
@ -572,6 +575,89 @@ class MCPRequestHandler:
route=route,
)
@staticmethod
async def _reload_admitted_principal(identity: EnvelopeIdentity) -> UserAPIKeyAuth:
"""Reload the live litellm record the envelope's subject references.
Dispatches on the sealed subject type: a ``key_hash`` reloads the virtual key that
minted the envelope (the scripted two-header client that presents a litellm key at the
token endpoint), a ``user_id`` reloads the user that authenticated interactively (the
DCR client, whose SSO login at the bridged authorize yields a user, not a key). Both
return a ``UserAPIKeyAuth`` the caller runs through the centralized policy gate, so
team/project/org/budget/SCIM enforcement is identical to the principal presenting
itself directly."""
match identity.subject_type:
case "key_hash":
return await MCPRequestHandler._reload_admitted_key(identity.subject)
case "user_id":
return await MCPRequestHandler._reload_admitted_user(identity.subject)
case _:
assert_never(identity.subject_type)
@staticmethod
async def _reload_admitted_user(user_id: str) -> UserAPIKeyAuth:
"""Reload the live user an interactively-minted envelope references and admit them as
themselves.
The DCR client authenticates via SSO at the bridged authorize, which yields a user
subject rather than a virtual key, so the envelope admits under the user's own
identity: the reloaded ``user_id`` and the user's own MCP object permission ride on the
returned ``UserAPIKeyAuth``, and the SAME ``get_allowed_mcp_servers`` the key path uses then
computes which servers the user may reach, so the user's litellm MCP grants and access groups
gate the request exactly as a key's do. Only the user's OWN object permission is bound: a
``UserAPIKeyAuth`` carries a single ``team_id`` while a user may belong to many teams, so
team-inherited MCP grants for a user are a follow-up (they need a many-teams union
``get_allowed_mcp_servers`` does not do off one auth object). The caller's centralized policy
gate enforces the user's live budget and org state, and a SCIM-deactivated owner fails closed.
Error handling mirrors the key path's retryable-503 contract, but ``get_user_object`` defeats a
type-based check: where ``get_key_object`` raises a typed ``ProxyException`` for a missing key
and lets a DB outage propagate raw, ``get_user_object`` catches every DB failure and re-raises a
bare ``ValueError``, so a missing user and a real outage look identical and the original error
survives only as ``__context__``. ``_raise_503_if_db_unavailable`` therefore walks the cause
chain: a transient DB outage still surfaces as a retryable 503, while a missing user, or any
other non-outage resolution failure, fails closed as a 401 rather than an opaque 500. The
object-permission load shares this one boundary, so an outage there is classified the same
way (``get_object_permission`` itself swallows a failed load to ``None``, matching how
``get_key_object`` best-effort-loads a key's object permission)."""
from litellm.proxy.auth.auth_checks import get_object_permission, get_user_object
from litellm.proxy.proxy_server import prisma_client, user_api_key_cache
if prisma_client is None:
raise HTTPException(status_code=500, detail="Server misconfigured: no database connection")
try:
user_object = await get_user_object(
user_id=user_id,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
user_id_upsert=False,
)
# Resolve the user's own MCP object permission (get_user_object does not load it) so the shared
# get_allowed_mcp_servers can grant the user their litellm-granted servers. Reuses the same
# get_object_permission resolver the key and team paths use; no permission logic is duplicated.
object_permission = user_object.object_permission if user_object is not None else None
if user_object is not None and object_permission is None and user_object.object_permission_id:
object_permission = await get_object_permission(
object_permission_id=user_object.object_permission_id,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
)
except (ProxyException, HTTPException):
raise HTTPException(status_code=401, detail="Invalid or expired credential") from None
except Exception as e: # noqa: BLE001 # a DB outage anywhere in the resolution is a retryable 503, not an opaque 500; anything else fails closed as 401
MCPRequestHandler._raise_503_if_db_unavailable(e)
raise HTTPException(status_code=401, detail="Invalid or expired credential") from None
if user_object is None:
raise HTTPException(status_code=401, detail="Invalid or expired credential")
if isinstance(user_object.metadata, dict) and user_object.metadata.get("scim_active") is False:
raise HTTPException(status_code=401, detail="Invalid or expired credential")
return UserAPIKeyAuth(
user_id=user_object.user_id,
user_role=user_object.user_role,
object_permission=object_permission,
object_permission_id=user_object.object_permission_id,
)
@staticmethod
async def _reload_admitted_key(key_hash: str) -> UserAPIKeyAuth:
"""Reload the live key record an admitted envelope references and re-check live policy.
@ -615,10 +701,14 @@ class MCPRequestHandler:
"""Raise a retryable 503 when ``e`` means the auth database is unreachable, else return so the
caller applies its own fail-closed mapping. A DB outage must not masquerade as an auth failure
(401) or surface as an opaque 500; the caller retries. Mirrors ``UserAPIKeyAuthExceptionHandler``,
which renders a service-unavailable database error as 503 on the standard pipeline."""
which renders a service-unavailable database error as 503 on the standard pipeline.
Classifies across the ``__cause__``/``__context__`` chain, not just ``e`` itself: ``get_user_object``
re-raises every DB failure as a bare ``ValueError``, so a type-based check on the top exception
would miss a real outage wrapped inside it."""
from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler
if PrismaDBExceptionHandler.is_database_service_unavailable_error(e):
if PrismaDBExceptionHandler.is_database_service_unavailable_error_in_chain(e):
raise HTTPException(
status_code=503,
detail="Service Unavailable, the authentication database is temporarily unreachable. Please retry shortly.",

View file

@ -0,0 +1,694 @@
"""Bridge token flow: litellm identity resolution and the DCR-bridge oauth_delegate mint/refresh pipeline."""
import math
from dataclasses import dataclass
from datetime import datetime, timezone
from typing import TYPE_CHECKING, Literal, Optional
from fastapi import HTTPException, Request
from fastapi.responses import JSONResponse
from pydantic import SecretStr
from typing_extensions import assert_never
from litellm._logging import verbose_logger
from litellm.proxy._experimental.mcp_server.oauth_utils import TOKEN_NO_CACHE_HEADERS
from litellm.types.mcp_server.mcp_server_manager import MCPServer
if TYPE_CHECKING:
from litellm.proxy._experimental.mcp_server.discoverable_endpoints import _BridgeAuthorizationCode
from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import (
EnvelopeIdentity,
EnvelopeKeys,
RefreshCredential,
UpstreamTokenGrant,
)
from litellm.proxy._types import UserAPIKeyAuth
def _litellm_key_from_request(request: Request) -> Optional[str]:
"""Return the LiteLLM API key presented on the request, or ``None``.
Accepts the key from ``x-litellm-api-key`` (what MCP clients such as Claude Desktop/Code
send) as well as ``Authorization``; either may carry a bare token or ``Bearer <token>``.
``x-litellm-api-key`` wins when both are present, since ``Authorization`` may instead carry
an OAuth/upstream bearer.
"""
for header_value in (
request.headers.get("x-litellm-api-key"),
request.headers.get("Authorization") or request.headers.get("authorization"),
):
if not header_value:
continue
value = header_value.strip()
if value.lower().startswith("bearer "):
value = value[7:].strip()
if value:
return value
return None
def _key_is_active(key_obj: "UserAPIKeyAuth") -> bool:
"""``True`` when the presented key is neither blocked nor past its expiry.
The OAuth token endpoint is unauthenticated, so the presented key is validated here before it is
trusted; a revoked or expired key must not mint a bridge envelope or write a stored credential.
``get_key_object`` resolves a row without these checks (the main ``user_api_key_auth`` pipeline
enforces them downstream, which this endpoint bypasses), so they are applied here. Deleted keys
are already rejected upstream, where ``get_key_object`` raises on a row that no longer exists.
This is an active-state gate only; it deliberately does not require a ``user_id``. A valid
team-scoped or service-account key has no ``user_id`` yet is a legitimate credential, so gating
on ``user_id`` presence would wrongly reject it. Callers that need the user (the per-user token
store) derive it separately via :func:`_active_key_user_id`.
Total by design: ``expires`` is typed ``str | datetime``, and an unparseable string would make
``datetime.fromisoformat`` raise. Since the callers run this outside their key-resolution
``try``, an uncaught parse error would surface as a 500 instead of the endpoint's fail-closed
behavior, so a malformed expiry is treated as inactive (return ``False``) rather than raising.
"""
if key_obj.blocked is True:
return False
expires = key_obj.expires
if expires is not None:
if isinstance(expires, datetime):
expiry = expires
else:
try:
expiry = datetime.fromisoformat(expires)
except (ValueError, TypeError):
return False
if expiry.tzinfo is None or expiry.tzinfo.utcoffset(expiry) is None:
expiry = expiry.replace(tzinfo=timezone.utc)
if expiry < datetime.now(timezone.utc):
return False
return True
def _active_key_user_id(key_obj: "UserAPIKeyAuth") -> str | None:
"""The active key's ``user_id``, or ``None`` when the key is blocked/expired or simply has no
``user_id`` (a team-scoped or service-account key). Used only by the per-user token store, which
needs a user to key the stored credential; the bridge mint uses the key hash and does not."""
return key_obj.user_id if _key_is_active(key_obj) else None
@dataclass(frozen=True, slots=True)
class _ResolvedKey:
"""An active litellm key resolved from the token request: its hash (the value ``get_key_object``
and the cache/DB layer key the record by) and the live record."""
key_hash: str
key: "UserAPIKeyAuth"
_KeyResolutionFailure = Literal["no_active_key", "unavailable", "unresolvable"]
"""Why a token request yielded no active litellm key, kept distinct so a caller statuses each truthfully
instead of blaming the client for a gateway problem:
- ``no_active_key``: none was presented, or the presented key is unknown / blocked / expired (the
caller's request is at fault)
- ``unavailable``: the auth database was transiently unreachable while resolving (retryable)
- ``unresolvable``: the gateway cannot resolve identity right now (no DB connection, or an unexpected
error) -- a gateway fault, not the caller's
The classification mirrors admission's ``_reload_admitted_key`` so the mint (ingress) and admission
(egress) never disagree on the status of the same outage."""
async def _resolve_active_litellm_key(request: Request) -> "_ResolvedKey | _KeyResolutionFailure":
"""Resolve the presented litellm key to an active key record, or say precisely why not.
Single resolution path the OAuth token endpoint reuses, resolving authoritatively via
``get_key_object`` (cache first, then DB). The failure is a value, not a bare ``None``, so a caller
can tell "the client sent no usable credential" (a request error) apart from "the gateway could not
check" (an infrastructure error) and status each truthfully; collapsing both to ``None`` is what let
a DB outage read as a 400. A resolved key is still gated by ``_key_is_active``, so a blocked or
expired key is ``no_active_key`` while a valid team-scoped or service-account key (no ``user_id``)
resolves. Classification mirrors admission's ``_reload_admitted_key``: no DB connection is a gateway
fault, a ``ProxyException`` / ``HTTPException`` from ``get_key_object`` is an unknown or invalid key,
a database-service-unavailable error is a retryable outage, and anything else is an unexpected
gateway fault."""
token = _litellm_key_from_request(request)
if not token:
return "no_active_key"
from litellm.proxy._types import hash_token # noqa: PLC0415 # inline import avoids a module-load circular import
return await _reload_active_key_by_hash(hash_token(token))
async def _reload_active_key_by_hash(key_hash: str) -> "_ResolvedKey | _KeyResolutionFailure":
"""Reload the live key record for ``key_hash`` (cache first, then DB) and gate it on active state,
returning the resolved key or a precise failure. Shared by the token request's presented-key
resolution (:func:`_resolve_active_litellm_key`, which hashes the presented key) and the refresh
path (which already holds the hash sealed in the refresh envelope), so both re-validate identity
through one active-key gate and one failure classification. Classification mirrors admission's
``_reload_admitted_key``: no DB connection is a gateway fault, a ``ProxyException`` / ``HTTPException``
from ``get_key_object`` is an unknown or invalid key, a database-service-unavailable error is a
retryable outage, and anything else is an unexpected gateway fault. A blocked or expired key is
``no_active_key``, so a revoked key can neither mint nor refresh a bridge envelope."""
from litellm.proxy._types import (
ProxyException, # noqa: PLC0415 # inline import avoids a module-load circular import
)
from litellm.proxy.auth.auth_checks import ( # noqa: PLC0415 # inline import avoids a module-load circular import
get_key_object,
)
from litellm.proxy.db.exception_handler import ( # noqa: PLC0415 # inline import avoids a module-load circular import
PrismaDBExceptionHandler,
)
from litellm.proxy.proxy_server import ( # noqa: PLC0415 # inline import avoids a module-load circular import
prisma_client,
user_api_key_cache,
)
if prisma_client is None:
return "unresolvable"
try:
key_obj = await get_key_object(
hashed_token=key_hash,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
)
except (ProxyException, HTTPException):
return "no_active_key"
except Exception as exc: # noqa: BLE001 # classify: a DB outage is retryable, anything else is an opaque gateway fault
if PrismaDBExceptionHandler.is_database_service_unavailable_error(exc):
return "unavailable"
verbose_logger.debug(
"_reload_active_key_by_hash: unexpected key-resolution error (%s)",
type(exc).__name__,
)
return "unresolvable"
if not _key_is_active(key_obj):
return "no_active_key"
return _ResolvedKey(key_hash=key_hash, key=key_obj)
async def _reload_active_user_by_id(user_id: str) -> "_KeyResolutionFailure | None":
"""Re-validate a live litellm user by id, returning ``None`` when the user is active or a precise
failure otherwise. The interactive DCR client authenticates via SSO, so its refresh envelope seals a
user subject; renewing it must re-check the user is still live (present and not SCIM-deactivated) so a
deactivated user cannot keep refreshing, mirroring how admission re-validates the same user subject on
the egress side. No DB connection is a gateway fault (``unresolvable``) and a
database-service-unavailable error is a retryable outage (``unavailable``). Everything else fails
closed as ``no_active_key`` (the caller maps it to invalid_grant): a ``ProxyException`` /
``HTTPException``, a SCIM-deactivated user, and, unlike the key path, a missing user. ``get_user_object``
catches every DB failure and re-raises a bare ``ValueError`` (a deleted user and a real outage look
identical, the original error surviving only as ``__context__``), so the outage check walks the cause
chain, and a missing user falls through to ``no_active_key`` rather than an opaque gateway fault."""
from litellm.proxy._types import (
ProxyException, # noqa: PLC0415 # inline import avoids a module-load circular import
)
from litellm.proxy.auth.auth_checks import ( # noqa: PLC0415 # inline import avoids a module-load circular import
get_user_object,
)
from litellm.proxy.db.exception_handler import ( # noqa: PLC0415 # inline import avoids a module-load circular import
PrismaDBExceptionHandler,
)
from litellm.proxy.proxy_server import ( # noqa: PLC0415 # inline import avoids a module-load circular import
prisma_client,
user_api_key_cache,
)
if prisma_client is None:
return "unresolvable"
try:
user_object = await get_user_object(
user_id=user_id,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
user_id_upsert=False,
)
except (ProxyException, HTTPException):
return "no_active_key"
except Exception as exc: # noqa: BLE001 # a DB outage is retryable; a missing user (get_user_object's wrapped ValueError) or any other resolution failure fails closed as no_active_key, never a 500
if PrismaDBExceptionHandler.is_database_service_unavailable_error_in_chain(exc):
return "unavailable"
verbose_logger.debug("_reload_active_user_by_id: user-resolution error (%s)", type(exc).__name__)
return "no_active_key"
if user_object is None:
return "no_active_key"
if isinstance(user_object.metadata, dict) and user_object.metadata.get("scim_active") is False:
return "no_active_key"
return None
async def _key_owner_scim_deactivated(key: "UserAPIKeyAuth") -> bool:
"""True only when the key's owning user was explicitly SCIM-deactivated, so a refresh revokes an
offboarded owner's key exactly as admission does via ``_reject_if_admitted_owner_scim_deactivated``.
A key with no owner, a missing owner record, or a failed lookup fails OPEN (returns ``False``),
matching admission and the standard builder: a key may outlive its owner record, and a transient DB
blip must not revoke a live key. Only an explicit ``scim_active`` of ``False`` gates renewal."""
if key.user_id is None:
return False
from litellm.proxy.auth.auth_checks import ( # noqa: PLC0415 # inline import avoids a module-load circular import
get_user_object,
)
from litellm.proxy.proxy_server import ( # noqa: PLC0415 # inline import avoids a module-load circular import
prisma_client,
user_api_key_cache,
)
if prisma_client is None:
return False
try:
owner = await get_user_object(
user_id=key.user_id,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
user_id_upsert=False,
)
except Exception as exc: # noqa: BLE001 # fail open: a missing owner (get_user_object's wrapped ValueError) or a DB blip must not revoke a live key
verbose_logger.debug("refresh: key-owner SCIM lookup failed, not revoking (%s)", type(exc).__name__)
return False
return owner is not None and isinstance(owner.metadata, dict) and owner.metadata.get("scim_active") is False
async def _revalidate_active_subject(identity: "EnvelopeIdentity") -> "_KeyResolutionFailure | None":
"""Re-validate that the subject sealed in a refresh envelope is still live, dispatching on its type:
a key_hash reloads the virtual key, a user_id reloads the user. Returns ``None`` when the subject is
active or a precise failure otherwise, so revocation gates renewal for either identity source the same
way admission gates the egress: a blocked or expired key, a SCIM-deactivated key owner (mirroring
admission's owner check, so an offboarded user cannot keep renewing a still-active key), and a
deactivated or deleted user all fail closed to ``no_active_key``."""
match identity.subject_type:
case "key_hash":
reloaded = await _reload_active_key_by_hash(identity.subject)
if not isinstance(reloaded, _ResolvedKey):
return reloaded
if await _key_owner_scim_deactivated(reloaded.key):
return "no_active_key"
return None
case "user_id":
return await _reload_active_user_by_id(identity.subject)
case _:
assert_never(identity.subject_type)
async def _extract_user_id_from_request(request: Request) -> str | None:
"""The litellm ``user_id`` for the token request, so a per-user token is stored under the same
identity the egress later reads it by. Storage is best-effort, so every non-resolved outcome
(including a transient DB outage) collapses to ``None`` here and the caller simply skips the store;
the bridge mint, which must status those outcomes differently, consumes
:func:`_resolve_active_litellm_key` directly."""
resolved = await _resolve_active_litellm_key(request)
if not isinstance(resolved, _ResolvedKey):
return None
return _active_key_user_id(resolved.key)
_UpstreamGrantRejection = Literal["no_access_token", "expired_lifetime"]
"""Why an upstream token response cannot back a bridge envelope:
- ``no_access_token``: the response carries no usable ``access_token``
- ``expired_lifetime``: the response reports a parseable, non-positive ``expires_in``, i.e. an upstream
token that is already dead, so sealing it would forward a bearer the edge cannot use
An absent or unparseable ``expires_in`` is NOT a rejection; the lifetime is merely unknown and the
envelope caps it, the by-design behaviour for an upstream that omits the field."""
def _classify_upstream_lifetime(raw_expires_in: object) -> "int | Literal['unspecified', 'expired']":
"""Classify an upstream ``expires_in`` into a positive number of seconds, ``"unspecified"`` (absent
or unparseable, so the envelope caps it), or ``"expired"`` (a non-positive value the upstream reports
as already elapsed). Telling "we do not know the lifetime" apart from "the upstream says it is
already dead" is what stops an explicitly-expired token from silently receiving the envelope's 1h
cap. The expired decision is made on the parsed numeric value, not on ``int(...)`` of it, so a
positive sub-second lifetime in ``(0, 1)`` is not truncated to ``0`` and misread as elapsed; the
envelope works in whole seconds, so such a lifetime clamps up to its 1s floor. ``bool`` is excluded
(an ``int`` subclass but never a real lifetime), and the conversions can raise on ``NaN`` /
``Infinity`` / oversized input, which reads as unparseable rather than surfacing as a 500."""
if raw_expires_in is None or isinstance(raw_expires_in, bool) or not isinstance(raw_expires_in, (int, float, str)):
return "unspecified"
try:
numeric = float(raw_expires_in)
seconds = int(numeric)
except (ValueError, TypeError, OverflowError):
return "unspecified"
if numeric <= 0:
return "expired"
return max(1, seconds)
def _bridge_grant_from_token_response(token_response: object) -> "UpstreamTokenGrant | _UpstreamGrantRejection":
"""Validate an upstream OAuth token response into a typed grant, or say why it cannot back an
envelope. Each field is isinstance-checked so nothing untyped from ``response.json()`` reaches the
grant. ``expires_in`` is read three ways (see :func:`_classify_upstream_lifetime`): an unknown
lifetime leaves the grant ``expires_in`` ``None`` for the envelope to cap, a positive value is
honoured, and an explicit already-elapsed value is a rejection rather than a silent fall-through to
the cap."""
from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import ( # noqa: PLC0415 # inline import avoids a module-load circular import
UpstreamTokenGrant,
)
if not isinstance(token_response, dict):
return "no_access_token"
access = token_response.get("access_token")
if not isinstance(access, str) or not access:
return "no_access_token"
lifetime = _classify_upstream_lifetime(token_response.get("expires_in"))
if lifetime == "expired":
return "expired_lifetime"
token_type = token_response.get("token_type")
scope = token_response.get("scope")
return UpstreamTokenGrant(
access_token=SecretStr(access),
token_type=token_type if isinstance(token_type, str) and token_type else "Bearer",
# The upstream refresh_token is deliberately NOT sealed: the edge never consumes it (it forwards
# only token_type + access_token), so it would be dead weight embedding a long-lived upstream
# credential in the client-held bearer, and it enlarges the envelope. Refresh support is a
# follow-up (a dedicated refresh-envelope); the client re-runs authorization_code at the cap.
refresh_token=None,
scope=scope if isinstance(scope, str) and scope else None,
expires_in=lifetime if isinstance(lifetime, int) else None,
)
# ---------------------------------------------------------------------------
# DCR-bridge oauth_delegate mint: a three-phase pipeline whose failures are values.
#
# prepare (before the upstream exchange) -> validate every precondition and resolve identity+keys
# exchange (the single-use upstream code is consumed here, in exchange_token_with_server)
# finish (after the exchange) -> seal the upstream grant into the client-held envelope
#
# Every precondition lives in ``prepare``, which runs BEFORE the exchange, so no failure can burn the
# single-use code or rotate a refresh token, for either grant type -- that whole class of bug is gone
# by construction rather than guarded case by case. Failures are values mapped to an OAuth-shaped
# response in one place (``_bridge_mint_error_response``), so status codes and the RFC 6749 §5.2 body
# shape are uniform. Adding a failure mode is a new literal plus a match arm the type checker forces.
# ---------------------------------------------------------------------------
_BridgeMintError = Literal[
"no_identity",
"invalid_refresh",
"identity_unavailable",
"identity_unresolvable",
"not_configured",
"no_upstream_token",
"upstream_token_expired",
"too_large",
]
@dataclass(frozen=True, slots=True)
class _BridgeMintReady:
"""Everything the seal needs, resolved once before the exchange: the identity to bind the envelope
to and the master-key-derived envelope keys. The identity is a key_hash subject for the scripted
two-header client (resolved from the litellm key it presents) or a user_id subject for the
interactive SSO client (the user recovered from the gateway authorization code), so one phase-3 seal
serves both. Resolving identity here means ``_finish_bridge_mint`` has no preconditions left to
fail."""
identity: "EnvelopeIdentity"
keys: "EnvelopeKeys"
def _bridge_mint_error_response(error: _BridgeMintError) -> JSONResponse:
"""Map a bridge-mint failure value to its token-endpoint response: one place, RFC 6749 §5.2 shape
(top-level ``error``, no-store headers) for every case, with a status truthful about where the
failure is. The caller's request is 400, a transient gateway outage is 503, a gateway
misconfiguration is 500, and an upstream problem is 502. The identity-resolution statuses match how
admission statuses the same conditions on the egress side, so mint and admit never disagree under
one outage."""
match error:
case "no_identity":
status, code, desc = (
400,
"invalid_request",
"this server issues a gateway-bound credential; complete the interactive sign-in, or "
"send a litellm credential (x-litellm-api-key or Authorization) on the token request",
)
case "invalid_refresh":
status, code, desc = (
400,
"invalid_grant",
"the refresh credential is not a valid, live refresh envelope for this server; "
"re-run authorization_code to obtain a new one",
)
case "identity_unavailable":
status, code, desc = (
503,
"temporarily_unavailable",
"the authentication database is temporarily unreachable; retry shortly",
)
case "identity_unresolvable":
status, code, desc = (
500,
"server_error",
"the gateway could not resolve the litellm identity for this request",
)
case "not_configured":
status, code, desc = (
500,
"server_error",
"the gateway is not configured to mint a gateway-bound credential (master_key is not set)",
)
case "no_upstream_token":
status, code, desc = (
502,
"server_error",
"the upstream token response has no usable access_token",
)
case "upstream_token_expired":
status, code, desc = (
502,
"server_error",
"the upstream token response reports an already-expired lifetime",
)
case "too_large":
status, code, desc = (
502,
"server_error",
"the upstream token is too large to seal into a gateway-bound credential",
)
case _:
assert_never(error)
return JSONResponse(
status_code=status, content={"error": code, "error_description": desc}, headers=TOKEN_NO_CACHE_HEADERS
)
def _key_resolution_failure_to_mint_error(failure: _KeyResolutionFailure) -> _BridgeMintError:
"""Lift an identity-resolution failure into the mint taxonomy, preserving origin so the status stays
truthful: the caller's missing credential is 400, a transient DB outage is 503, and a gateway that
cannot resolve identity is 500."""
match failure:
case "no_active_key":
return "no_identity"
case "unavailable":
return "identity_unavailable"
case "unresolvable":
return "identity_unresolvable"
case _:
assert_never(failure)
def _upstream_rejection_to_mint_error(rejection: _UpstreamGrantRejection) -> _BridgeMintError:
"""Lift an upstream-response rejection into the mint taxonomy; both are upstream faults (502)."""
match rejection:
case "no_access_token":
return "no_upstream_token"
case "expired_lifetime":
return "upstream_token_expired"
case _:
assert_never(rejection)
async def _prepare_bridge_mint(
request: Request,
mcp_server: MCPServer,
bridge_identity: "_BridgeAuthorizationCode | None" = None,
) -> "_BridgeMintReady | _BridgeMintError":
"""Phase 1 for the authorization_code grant, BEFORE the upstream exchange: confirm the gateway can
mint (master_key set), resolve the litellm identity, and derive the envelope keys. Returns a ready
context or a precise failure value. Running before the exchange is what makes every failure here fail
closed without consuming the single-use code.
Two identity sources, one envelope. The interactive DCR client authenticates via SSO at the bridged
authorize, so its identity arrives as ``bridge_identity`` (the user recovered from the gateway
authorization code) and mints a user subject. The scripted two-header client presents a litellm key
on the token request instead, so its identity is the active key's hash and mints a key_hash subject.
A missing or invalid presented key keeps its resolution origin so the mapper statuses it truthfully;
neither source present is ``no_identity``. The refresh_token grant has its own phase-1
(:func:`_prepare_bridge_refresh`), which recovers identity from the presented refresh envelope."""
from litellm.proxy._experimental.mcp_server.outbound_credentials.bridge_credentials import ( # noqa: PLC0415 # inline import avoids a module-load circular import
envelope_keys_from_master_key,
)
from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import ( # noqa: PLC0415 # inline import avoids a module-load circular import
key_hash_identity,
user_identity,
)
from litellm.proxy.proxy_server import ( # noqa: PLC0415 # inline import avoids a module-load circular import
master_key,
)
if not master_key:
return "not_configured"
keys = envelope_keys_from_master_key(master_key)
if bridge_identity is not None:
identity = user_identity(server_id=mcp_server.server_id, user_id=bridge_identity.litellm_user_id)
return _BridgeMintReady(identity=identity, keys=keys)
resolved = await _resolve_active_litellm_key(request)
if not isinstance(resolved, _ResolvedKey):
return _key_resolution_failure_to_mint_error(resolved)
identity = key_hash_identity(server_id=mcp_server.server_id, key_hash=resolved.key_hash)
return _BridgeMintReady(identity=identity, keys=keys)
@dataclass(frozen=True, slots=True)
class _BridgeRefreshReady:
"""A validated refresh request: the identity+keys to mint the renewed pair under, the upstream refresh
token (unwrapped from the client's refresh envelope) to exchange with the upstream IdP, and the scope
sealed alongside it at mint. The upstream refresh token is a ``SecretStr`` like every other credential
in this layer, so a repr or a traceback that captures this value never exposes the raw upstream refresh
token in plaintext. ``upstream_scope`` carries the originally-granted scope so the renewal re-requests
it when the client (a DCR/MCP client that typically omits scope on refresh) sends none, keeping the
renewed token's scope stable against an upstream that would otherwise narrow or drop it."""
ready: "_BridgeMintReady"
upstream_refresh_token: SecretStr
upstream_scope: str | None = None
def _refresh_key_failure_to_mint_error(failure: _KeyResolutionFailure) -> _BridgeMintError:
"""Lift an identity-resolution failure on the refresh path into the mint taxonomy. Unlike the mint
path, a resolved-but-inactive (or unknown) key is ``invalid_grant`` rather than ``invalid_request``:
the client did present an identity (sealed in the refresh envelope), but it is no longer live, so the
refresh is invalid and the client must re-authenticate. A transient outage is still 503 and a gateway
fault still 500, matching the mint path and admission."""
match failure:
case "no_active_key":
return "invalid_refresh"
case "unavailable":
return "identity_unavailable"
case "unresolvable":
return "identity_unresolvable"
case _:
assert_never(failure)
async def _prepare_bridge_refresh(
mcp_server: MCPServer, refresh_value: str | None
) -> "_BridgeRefreshReady | _BridgeMintError":
"""Phase 1 for the refresh_token grant, BEFORE the upstream exchange: open the client's refresh
envelope, re-validate the sealed litellm identity so a revoked key cannot keep refreshing, and
recover the upstream refresh token to exchange. Identity comes entirely from the sealed envelope, not
the HTTP request, so the request object is not needed here. The client presents a refresh envelope,
never a raw upstream refresh token, so a missing value, a non-envelope, an unopenable envelope, or one
minted for another server is ``invalid_grant``. Running before the exchange means a rejected refresh
never consumes or rotates the upstream refresh token."""
from litellm.proxy._experimental.mcp_server.outbound_credentials.bridge_credentials import ( # noqa: PLC0415 # inline import avoids a module-load circular import
BridgeRefreshOpened,
envelope_keys_from_master_key,
open_bridge_refresh_envelope,
)
from litellm.proxy.proxy_server import ( # noqa: PLC0415 # inline import avoids a module-load circular import
master_key,
)
if not master_key:
return "not_configured"
if not refresh_value:
return "invalid_refresh"
keys = envelope_keys_from_master_key(master_key)
opened = open_bridge_refresh_envelope(refresh_value, keys, datetime.now(timezone.utc), mcp_server.server_id)
if not isinstance(opened, BridgeRefreshOpened):
return "invalid_refresh"
failure = await _revalidate_active_subject(opened.identity)
if failure is not None:
return _refresh_key_failure_to_mint_error(failure)
return _BridgeRefreshReady(
ready=_BridgeMintReady(identity=opened.identity, keys=keys),
upstream_refresh_token=opened.refresh.refresh_token,
upstream_scope=opened.refresh.scope,
)
def _finish_bridge_mint(
ready: "_BridgeMintReady", mcp_server: MCPServer, token_response: object, now: datetime
) -> "JSONResponse | _BridgeMintError":
"""Phase 3, AFTER the upstream exchange: seal the upstream grant into the client-held access envelope
using the pre-resolved identity and keys, and, when the upstream returned a refresh token, seal a
long-lived refresh envelope alongside it so the client can renew without re-authenticating. Shared by
the authorization_code and refresh_token paths, so a renewal that the upstream rotates re-issues a
fresh refresh envelope. The only hard failures here are properties of the upstream access token (no
usable token, an already-expired lifetime, or a token too large to seal); a refresh token that cannot
be sealed degrades to an access-only response rather than failing the whole exchange."""
from litellm.proxy._experimental.mcp_server.outbound_credentials.bridge_credentials import ( # noqa: PLC0415 # inline import avoids a module-load circular import
build_bridge_token_response,
)
from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import ( # noqa: PLC0415 # inline import avoids a module-load circular import
SealedEnvelope,
UpstreamTokenGrant,
)
grant = _bridge_grant_from_token_response(token_response)
if not isinstance(grant, UpstreamTokenGrant):
return _upstream_rejection_to_mint_error(grant)
sealed = build_bridge_token_response(ready.identity, grant, ready.keys, now)
if not isinstance(sealed, SealedEnvelope):
return "too_large"
# Report expires_in from the JWT's own second-truncated exp, rounding the elapsed portion up, so the
# client is never told the bearer lives past the point admission (which uses that exp) rejects it.
expires_in = max(0, int(sealed.expires_at.timestamp()) - math.ceil(now.timestamp()))
refresh_envelope = _mint_refresh_envelope_value(ready.identity, token_response, ready.keys, now, mcp_server)
body = {
"access_token": sealed.token.get_secret_value(),
"token_type": "Bearer",
"expires_in": expires_in,
# A refresh envelope rides along only when the upstream returned a refresh token to seal; when it
# rotates on renewal, the client receives the new one and the old envelope's upstream token dies.
**({"refresh_token": refresh_envelope} if refresh_envelope is not None else {}),
}
return JSONResponse(body, headers=TOKEN_NO_CACHE_HEADERS)
def _upstream_refresh_credential(token_response: object) -> "RefreshCredential | None":
"""Extract the upstream refresh grant from a token response, or ``None`` when there is none to seal.
Each field is isinstance-checked so nothing untyped reaches the refresh envelope; ``refresh_expires_in``
(the refresh token's own lifetime, when the upstream reports it) is classified like ``expires_in`` and
bounds the refresh envelope's TTL. An upstream that reports the refresh token itself as already elapsed
(``refresh_expires_in`` non-positive) yields ``None`` rather than a refresh envelope: sealing a dead
token would hand the client a full-TTL-capped envelope the IdP will reject, so the exchange degrades to
an access-only response (the client re-authenticates at access expiry), mirroring how
:func:`_bridge_grant_from_token_response` refuses an already-elapsed access token instead of capping it."""
from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import ( # noqa: PLC0415 # inline import avoids a module-load circular import
RefreshCredential,
)
if not isinstance(token_response, dict):
return None
refresh = token_response.get("refresh_token")
if not isinstance(refresh, str) or not refresh:
return None
lifetime = _classify_upstream_lifetime(token_response.get("refresh_expires_in"))
if lifetime == "expired":
return None
scope = token_response.get("scope")
return RefreshCredential(
refresh_token=SecretStr(refresh),
scope=scope if isinstance(scope, str) and scope else None,
expires_in=lifetime if isinstance(lifetime, int) else None,
)
def _mint_refresh_envelope_value(
identity: "EnvelopeIdentity", token_response: object, keys: "EnvelopeKeys", now: datetime, mcp_server: MCPServer
) -> str | None:
"""Seal the upstream refresh grant (if any) into a refresh envelope and return its bearer string, or
``None`` when the upstream returned no refresh token or the refresh token is too large to seal. A
too-large refresh token degrades to an access-only response (logged) rather than failing an exchange
that already succeeded upstream: the client simply re-authenticates when the access envelope expires."""
from litellm.proxy._experimental.mcp_server.outbound_credentials.bridge_credentials import ( # noqa: PLC0415 # inline import avoids a module-load circular import
build_bridge_refresh_token_response,
)
from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import ( # noqa: PLC0415 # inline import avoids a module-load circular import
SealedEnvelope,
)
refresh_credential = _upstream_refresh_credential(token_response)
if refresh_credential is None:
return None
sealed = build_bridge_refresh_token_response(identity, refresh_credential, keys, now)
if isinstance(sealed, SealedEnvelope):
return sealed.token.get_secret_value()
verbose_logger.warning(
"bridge mint: the upstream refresh token is too large to seal into a refresh envelope for "
"server=%s; issuing an access-only response, so the client re-authenticates at access expiry",
mcp_server.server_id,
)
return None

View file

@ -10,7 +10,7 @@ from urllib.parse import parse_qsl, urlencode, urlparse, urlunparse
import httpx
from fastapi import APIRouter, Form, HTTPException, Request
from fastapi.responses import HTMLResponse, JSONResponse, RedirectResponse, Response
from pydantic import BaseModel, ValidationError
from pydantic import BaseModel, ConfigDict, Field, SecretStr, ValidationError
from litellm._logging import verbose_logger
from litellm.llms.custom_httpx.http_handler import (
@ -21,6 +21,24 @@ from litellm.proxy._experimental.mcp_server.auth.token_endpoint_auth import (
TokenEndpointAuthConfigError,
build_token_endpoint_client_auth,
)
from litellm.proxy._experimental.mcp_server.bridge_token_flow import (
_bridge_mint_error_response,
_BridgeMintReady,
_BridgeRefreshReady,
_extract_user_id_from_request,
_finish_bridge_mint,
_prepare_bridge_mint,
_prepare_bridge_refresh,
)
from litellm.proxy._experimental.mcp_server.faults import (
CallerRejected,
CredentialSource,
UpstreamProtocolFault,
classify_upstream_dcr_rejection,
classify_upstream_token_rejection,
dcr_fault_detail,
render_token_fault,
)
from litellm.proxy._experimental.mcp_server.oauth_utils import (
TOKEN_NO_CACHE_HEADERS,
get_request_base_url,
@ -37,7 +55,7 @@ from litellm.types.mcp import MCPAuth, MCPCredentials
from litellm.types.mcp_server.mcp_server_manager import MCPServer
if TYPE_CHECKING:
from litellm.proxy._types import LiteLLM_MCPServerTable, UserAPIKeyAuth
from litellm.proxy._types import LiteLLM_MCPServerTable
# TTL cache for upstream OAuth metadata fetched from pass-through MCP servers.
# Keeps us from hammering the upstream IdP on each discovery request.
@ -91,6 +109,8 @@ def encode_state_with_base_url(
code_challenge: Optional[str] = None,
code_challenge_method: Optional[str] = None,
client_redirect_uri: Optional[str] = None,
litellm_user_id: str | None = None,
mcp_server_id: str | None = None,
) -> str:
"""
Encode the base_url, original state, and PKCE parameters using encryption.
@ -101,6 +121,11 @@ def encode_state_with_base_url(
code_challenge: PKCE code challenge from client
code_challenge_method: PKCE code challenge method from client
client_redirect_uri: Original redirect_uri from client
litellm_user_id: The SSO-authenticated litellm user captured at the bridge authorize
(interactive dcr_bridge oauth_delegate only); the callback seals it into the gateway
authorization code so the token mint can bind the envelope to this user
mcp_server_id: The bridge server the interactive flow targets, sealed alongside
litellm_user_id so the gateway code cannot be replayed against another server
Returns:
An encrypted string that encodes all values
@ -111,6 +136,8 @@ def encode_state_with_base_url(
"code_challenge": code_challenge,
"code_challenge_method": code_challenge_method,
"client_redirect_uri": client_redirect_uri,
"litellm_user_id": litellm_user_id,
"mcp_server_id": mcp_server_id,
}
state_json = json.dumps(state_data, sort_keys=True)
encrypted_state = encrypt_value_helper(state_json)
@ -138,6 +165,68 @@ def decode_state_hash(encrypted_state: str) -> dict:
return state_data
_BRIDGE_AUTH_CODE_PREFIX = "llm_bcode_"
class _BridgeAuthorizationCode(BaseModel):
"""The identity and upstream code the gateway seals into the authorization code it hands a DCR
client for an interactive dcr_bridge oauth_delegate sign-in, recovered at the token endpoint."""
model_config = ConfigDict(frozen=True)
upstream_code: str = Field(min_length=1)
litellm_user_id: str = Field(min_length=1)
mcp_server_id: str = Field(min_length=1)
def is_bridge_authorization_code(code: str) -> bool:
"""Cheap prefix check that ``code`` is a gateway-sealed bridge authorization code rather than a
raw upstream code, so the token endpoint can route without decrypting."""
return code.startswith(_BRIDGE_AUTH_CODE_PREFIX)
def seal_bridge_authorization_code(upstream_code: str, litellm_user_id: str, mcp_server_id: str) -> str:
"""Seal the upstream authorization code and the SSO-captured litellm user into a gateway
authorization code. The DCR client only echoes this opaque value back at the token endpoint; the
gateway decrypts it there to recover the user (to bind the envelope) and the upstream code (to
exchange with the upstream), so a litellm identity captured in the browser at authorize survives
to the back-channel token call with nothing stored server-side. Encrypted with the repo's
authenticated symmetric helper (the same family the OAuth state uses), so the client can neither
read nor forge it."""
payload = json.dumps(
{"upstream_code": upstream_code, "litellm_user_id": litellm_user_id, "mcp_server_id": mcp_server_id},
sort_keys=True,
)
return _BRIDGE_AUTH_CODE_PREFIX + encrypt_value_helper(payload)
def open_bridge_authorization_code(code: str) -> _BridgeAuthorizationCode | None:
"""Recover the sealed identity and upstream code, or ``None`` when ``code`` is not a gateway
bridge code or does not decrypt / validate. Total over hostile input: a raw upstream code (the
scripted two-header path) returns ``None`` and the caller falls through to the existing
behavior."""
if not is_bridge_authorization_code(code):
return None
decrypted = decrypt_value_helper(
code[len(_BRIDGE_AUTH_CODE_PREFIX) :], "bridge_authorization_code", return_original_value=False
)
if not isinstance(decrypted, str):
return None
try:
return _BridgeAuthorizationCode.model_validate_json(decrypted)
except ValidationError:
return None
def _redirect_to_litellm_login(request: Request) -> RedirectResponse:
"""Send an unauthenticated browser through litellm login before the interactive bridge authorize
can capture its identity. The bridge oauth_delegate flow seals the SSO user into the gateway code,
so a session is required; without one there is nothing to bind. After login the user re-initiates
the connection, which then finds the session cookie (the seamless return-to round-trip, which is
origin-validated against the control-plane URL, is a follow-up)."""
base_url = get_request_base_url(request)
return RedirectResponse(f"{base_url}/sso/key/generate")
# LIT-4197: some upstream authorization servers reject an over-long ``state``
# (the encrypted OAuth session blob routinely exceeds their limit). The upstream
# only needs an opaque value it echoes back on ``/callback``, so we forward a
@ -304,90 +393,6 @@ def _validate_token_response(
)
def _litellm_key_from_request(request: Request) -> Optional[str]:
"""Return the LiteLLM API key presented on the request, or ``None``.
Accepts the key from ``x-litellm-api-key`` (what MCP clients such as Claude Desktop/Code
send) as well as ``Authorization``; either may carry a bare token or ``Bearer <token>``.
``x-litellm-api-key`` wins when both are present, since ``Authorization`` may instead carry
an OAuth/upstream bearer.
"""
for header_value in (
request.headers.get("x-litellm-api-key"),
request.headers.get("Authorization") or request.headers.get("authorization"),
):
if not header_value:
continue
value = header_value.strip()
if value.lower().startswith("bearer "):
value = value[7:].strip()
if value:
return value
return None
def _active_key_user_id(key_obj: "UserAPIKeyAuth") -> Optional[str]:
"""The key's ``user_id``, or ``None`` if the key is blocked or expired.
The OAuth token endpoint is unauthenticated, so the presented key is validated here before its
identity is trusted to key a stored credential; a revoked or expired key must not be able to
write or overwrite the per-user OAuth token. ``get_key_object`` resolves a row without these
checks (the main ``user_api_key_auth`` pipeline enforces them downstream, which this endpoint
bypasses), so they are applied here. Deleted keys are already rejected upstream, where
``get_key_object`` raises on a row that no longer exists.
"""
if key_obj.blocked is True:
return None
expires = key_obj.expires
if expires is not None:
expiry = expires if isinstance(expires, datetime) else datetime.fromisoformat(expires)
if expiry.tzinfo is None or expiry.tzinfo.utcoffset(expiry) is None:
expiry = expiry.replace(tzinfo=timezone.utc)
if expiry < datetime.now(timezone.utc):
return None
return key_obj.user_id
async def _extract_user_id_from_request(request: Request) -> Optional[str]:
"""Resolve the LiteLLM ``user_id`` at the OAuth token endpoint so a per-user token is stored
under the same identity the egress later reads it by (``user_api_key_auth.user_id``).
Resolves authoritatively via ``get_key_object`` (cache first, then DB) instead of a raw cache
peek. On a multi-replica gateway the token-exchange request can land on a worker whose in-memory
cache never saw the key, and a cross-replica Redis hit deserializes to a plain ``dict`` rather
than a ``UserAPIKeyAuth``; the previous code read only ``Authorization`` and did
``getattr(cached, "user_id")`` with no ``model_type`` rehydration and no DB fallback, so it
silently returned ``None`` and the token was never persisted, which makes the egress 401 on every
reconnect. The resolved key is validated (``_active_key_user_id``) before its identity is trusted,
so a blocked or expired key cannot write. Returns ``None`` when no key is present, the key cannot
be resolved, or it is blocked/expired.
"""
token = _litellm_key_from_request(request)
if not token:
return None
try:
from litellm.proxy._types import hash_token # noqa: PLC0415
from litellm.proxy.auth.auth_checks import get_key_object # noqa: PLC0415
from litellm.proxy.proxy_server import ( # noqa: PLC0415
prisma_client,
user_api_key_cache,
)
key_obj = await get_key_object(
hashed_token=hash_token(token),
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
)
return _active_key_user_id(key_obj)
except Exception as exc:
verbose_logger.debug(
"_extract_user_id_from_request: could not resolve a LiteLLM user_id for the presented "
"key (%s); per-user token will not be stored server-side.",
type(exc).__name__,
)
return None
async def _store_per_user_token_server_side(
server: MCPServer,
user_id: str,
@ -620,12 +625,31 @@ async def authorize_with_server(
parsed = urlparse(redirect_uri)
base_url = urlunparse(parsed._replace(query=""))
request_base_url = get_request_base_url(request)
# Interactive dcr_bridge oauth_delegate sign-in: this arm runs the gateway /callback and /token in
# the loop, so the gateway can capture the litellm user here (from the browser's UI session) and
# carry it to the back-channel token mint. Seal the SSO user and the target server into the state;
# the callback reads them back to mint the gateway authorization code. A DCR client cannot present a
# litellm key, so the browser session is the only identity source; without one there is nothing to
# bind, so send the user through login first. Every other oauth2 server keeps the identity-less state.
litellm_user_id: str | None = None
if mcp_server.is_dcr_bridge and mcp_server.is_oauth_delegate:
from litellm.proxy._experimental.mcp_server.byok_oauth_endpoints import ( # noqa: PLC0415 # inline import avoids a module-load circular import
_user_id_from_session_cookie,
)
litellm_user_id = _user_id_from_session_cookie(request)
if litellm_user_id is None:
return _redirect_to_litellm_login(request)
encoded_state = encode_state_with_base_url(
base_url=base_url,
original_state=state,
code_challenge=code_challenge,
code_challenge_method=code_challenge_method,
client_redirect_uri=redirect_uri,
litellm_user_id=litellm_user_id,
mcp_server_id=mcp_server.server_id if litellm_user_id else None,
)
relay_state = secrets.token_urlsafe(_OAUTH_STATE_HANDLE_BYTES)
@ -654,6 +678,13 @@ async def authorize_with_server(
return response
def _token_credential_source(mcp_server: MCPServer) -> CredentialSource:
"""Mirrors the resolved-client rule in :func:`exchange_token_with_server`: when the server has a
stored client_id the gateway presents its own credentials upstream, so a credential rejection is
the operator's fault, not the caller's."""
return "gateway_stored" if mcp_server.client_id else "caller_supplied"
async def exchange_token_with_server(
request: Request,
mcp_server: MCPServer,
@ -688,25 +719,61 @@ async def exchange_token_with_server(
except TokenEndpointAuthConfigError as exc:
raise HTTPException(status_code=400, detail=str(exc)) from exc
bridge_identity: _BridgeAuthorizationCode | None = None
bridge_mint_ready: _BridgeMintReady | None = None
bridge_upstream_refresh: SecretStr | None = None
bridge_upstream_scope: str | None = None
refresh_request_scope: str | None = None
is_bridge = mcp_server.is_oauth_delegate and mcp_server.is_dcr_bridge
if grant_type == "refresh_token":
if not refresh_token:
# Phase 1 for a bridge refresh: open the client's refresh envelope, re-validate the sealed
# identity, and unwrap the real upstream refresh token BEFORE building token_data, so the exchange
# sends the upstream token and never the envelope. A failure returns without touching the upstream.
if is_bridge:
prepared_refresh = await _prepare_bridge_refresh(mcp_server, refresh_token)
if not isinstance(prepared_refresh, _BridgeRefreshReady):
return _bridge_mint_error_response(prepared_refresh)
bridge_mint_ready = prepared_refresh.ready
bridge_upstream_refresh = prepared_refresh.upstream_refresh_token
bridge_upstream_scope = prepared_refresh.upstream_scope
# A bridge server sends the unwrapped upstream refresh token recovered from the client's refresh
# envelope above; every other server sends the client's own refresh token verbatim.
upstream_refresh_token = (
bridge_upstream_refresh.get_secret_value() if bridge_upstream_refresh is not None else refresh_token
)
if not upstream_refresh_token:
raise HTTPException(
status_code=400,
detail="refresh_token is required for refresh_token grant",
)
token_data: dict = {
"grant_type": "refresh_token",
"refresh_token": refresh_token,
"refresh_token": upstream_refresh_token,
**client_auth.body,
}
if scope:
token_data["scope"] = scope
refresh_request_scope = scope or bridge_upstream_scope
if refresh_request_scope:
token_data["scope"] = refresh_request_scope
else:
if not code:
raise HTTPException(
status_code=400,
detail="code is required for authorization_code grant",
)
# Interactive dcr_bridge oauth_delegate: the client presents the gateway authorization code the
# callback sealed. Recover the SSO user and the real upstream code from it; the upstream exchange
# below uses the upstream code, and the mint binds the envelope to the recovered user. Bind the
# sealed server to this request so a code minted for one bridge server cannot be spent at another.
# A raw upstream code (scripted path) opens to None and the code is used as-is.
bridge_identity = open_bridge_authorization_code(code)
if bridge_identity is not None:
if bridge_identity.mcp_server_id != mcp_server.server_id:
raise HTTPException(
status_code=400,
detail="Authorization code was issued for a different MCP server",
)
code = bridge_identity.upstream_code
bridge_token_relay = _dcr_bridge_relays_client_registration(mcp_server)
if bridge_token_relay and not redirect_uri:
raise HTTPException(
@ -726,32 +793,49 @@ async def exchange_token_with_server(
}
if code_verifier:
token_data["code_verifier"] = code_verifier
# Phase 1 for a bridge authorization_code mint: resolve identity (the SSO user recovered above, or
# the presented litellm key) and the envelope keys BEFORE the exchange consumes the single-use code.
if is_bridge:
prepared = await _prepare_bridge_mint(request, mcp_server, bridge_identity)
if not isinstance(prepared, _BridgeMintReady):
return _bridge_mint_error_response(prepared)
bridge_mint_ready = prepared
async_client = get_async_httpx_client(llm_provider=httpxSpecialProvider.Oauth2Check)
response = await async_client.post(
mcp_server.token_url,
headers={"Accept": "application/json", **client_auth.headers},
data=token_data,
)
try:
response = await async_client.post(
mcp_server.token_url,
headers={"Accept": "application/json", **client_auth.headers},
data=token_data,
)
if response is not None:
response.raise_for_status()
except httpx.HTTPStatusError as exc:
fault = classify_upstream_token_rejection(
exc.response,
credential_source=_token_credential_source(mcp_server),
log_context=mcp_server.server_id,
)
upstream_rejected_bridge_refresh = (
is_bridge
and grant_type == "refresh_token"
and isinstance(fault, CallerRejected)
and fault.code == "invalid_grant"
)
if upstream_rejected_bridge_refresh:
verbose_logger.info(
"bridge refresh: the upstream rejected the sealed refresh token for server=%s with "
"invalid_grant (revoked or expired at the IdP); returning invalid_grant so the client "
"re-runs authorization_code rather than an opaque upstream error",
mcp_server.server_id,
)
return _bridge_mint_error_response("invalid_refresh")
return render_token_fault(fault)
if response is None:
raise HTTPException(
status_code=502,
detail="MCP upstream token endpoint returned no response",
)
try:
response.raise_for_status()
except httpx.HTTPStatusError as exc:
if "invalid_target" in exc.response.text:
verbose_logger.warning(
"MCP server %s: the upstream authorization server rejected the token request with "
"invalid_target; it may require RFC 8707 resource indicators, which the gateway "
"does not send yet (tracked as LIT-4339)",
mcp_server.server_id,
)
raise
token_response = response.json()
access_token = token_response["access_token"]
# Validate token response against server-configured rules before any storage.
# This rejects tokens from wrong Slack workspaces, Atlassian orgs, etc.
@ -791,8 +875,23 @@ async def exchange_token_with_server(
mcp_server.server_id,
)
# A DCR-bridge oauth_delegate server hands the client a gateway-bound envelope (identity plus the
# upstream token) instead of the raw upstream token, so the one bearer both admits the caller and
# forwards the upstream credential. Only this mode mints; every other server returns the raw token.
if bridge_mint_ready is not None:
if refresh_request_scope and isinstance(token_response, dict) and not token_response.get("scope"):
token_response = {**token_response, "scope": refresh_request_scope}
# Phase 3: seal the upstream grant into the client-held envelope; failures map through the same
# OAuth-shaped response as the phase-1 preconditions.
minted = _finish_bridge_mint(bridge_mint_ready, mcp_server, token_response, datetime.now(timezone.utc))
return minted if isinstance(minted, JSONResponse) else _bridge_mint_error_response(minted)
raw_access_token = token_response.get("access_token") if isinstance(token_response, dict) else None
if not isinstance(raw_access_token, str) or not raw_access_token:
return render_token_fault(UpstreamProtocolFault(note="the upstream token response has no usable access_token"))
result = {
"access_token": access_token,
"access_token": raw_access_token,
"token_type": token_response.get("token_type", "Bearer"),
}
@ -1048,21 +1147,6 @@ async def _persist_dcr_client_registration(
return "failed"
_MAX_UPSTREAM_ERROR_CHARS = 500
def _safe_upstream_error_detail(response: httpx.Response) -> str:
"""Bounded plaintext summary of an upstream registration failure for the client.
RFC 7591 error bodies are small JSON objects (``error`` / ``error_description``); relaying the
text lets the client read the real reason instead of a bare 500, and the length bound keeps a
hostile or oversized upstream body from bloating the gateway response."""
body = response.text
if not body:
return response.reason_phrase or "upstream registration failed"
return body[:_MAX_UPSTREAM_ERROR_CHARS]
async def register_client_with_server(
request: Request,
mcp_server: MCPServer,
@ -1122,19 +1206,24 @@ async def register_client_with_server(
}
async_client = get_async_httpx_client(llm_provider=httpxSpecialProvider.Oauth2Register)
response = await async_client.post(
mcp_server.registration_url,
headers=headers,
json=register_data,
)
try:
response = await async_client.post(
mcp_server.registration_url,
headers=headers,
json=register_data,
)
if response is not None:
response.raise_for_status()
except httpx.HTTPStatusError as exc:
status_code, detail = dcr_fault_detail(
classify_upstream_dcr_rejection(exc.response, log_context=mcp_server.server_id)
)
raise HTTPException(status_code=status_code, detail=detail) from exc
if response is None:
raise HTTPException(
status_code=502,
detail="MCP upstream registration endpoint returned no response",
)
if bridge_relay and response.status_code >= 400:
raise HTTPException(status_code=response.status_code, detail=_safe_upstream_error_detail(response))
response.raise_for_status()
token_response = response.json()
@ -1362,7 +1451,20 @@ async def callback(
# states while permitting same-origin / allowlisted clients.
redirect_uri = _get_validated_client_redirect_uri(request, state_data)
params = {"code": code, "state": original_state}
# Interactive dcr_bridge oauth_delegate: the state carries the litellm user the authorize step
# captured. Instead of forwarding the raw upstream code (which the client would present at the
# token endpoint with no way to prove who signed in), seal the user and the upstream code into a
# gateway authorization code and forward THAT. The token endpoint decrypts it to bind the
# envelope to this user. Every other flow forwards the raw code unchanged.
litellm_user_id = state_data.get("litellm_user_id")
mcp_server_id = state_data.get("mcp_server_id")
forwarded_code = code
if isinstance(litellm_user_id, str) and litellm_user_id and isinstance(mcp_server_id, str) and mcp_server_id:
forwarded_code = seal_bridge_authorization_code(
upstream_code=code, litellm_user_id=litellm_user_id, mcp_server_id=mcp_server_id
)
params = {"code": forwarded_code, "state": original_state}
complete_returned_url = _append_query_params(redirect_uri, params)
response = RedirectResponse(url=complete_returned_url, status_code=302)
_clear_oauth_state_cookie(response, request, state)

View file

@ -0,0 +1,38 @@
"""Typed fault values for upstream OAuth/DCR failures (phase 1 of the MCP error-handling framework).
The invariant this package exists to enforce: an upstream failure is classified ONCE into a single
fault value, and the response status, wire error code, and prose are all derived from that value.
Deriving all three from one classification makes contradictory pairings (a caller-fault error code on
a server-fault status) unrepresentable, and gives the trust-boundary rule one enforcement point:
spec-defined machine fields may cross to callers, upstream prose and raw bodies go to server logs.
"""
from litellm.proxy._experimental.mcp_server.faults.classify import (
classify_upstream_dcr_rejection,
classify_upstream_token_rejection,
)
from litellm.proxy._experimental.mcp_server.faults.render_oauth import (
dcr_fault_detail,
render_token_fault,
)
from litellm.proxy._experimental.mcp_server.faults.types import (
CallerRejected,
CredentialSource,
GatewayRejected,
UpstreamOAuthFault,
UpstreamProtocolFault,
UpstreamReportedFault,
)
__all__ = [
"CallerRejected",
"CredentialSource",
"GatewayRejected",
"UpstreamOAuthFault",
"UpstreamProtocolFault",
"UpstreamReportedFault",
"classify_upstream_dcr_rejection",
"classify_upstream_token_rejection",
"dcr_fault_detail",
"render_token_fault",
]

View file

@ -0,0 +1,133 @@
"""The single place that reads upstream OAuth/DCR failure responses.
Every accessor here is total: an upstream that lies about its content encoding, sends an undecodable
body, or omits the spec fields yields a classified fault, never an exception. Nothing outside this
module should touch a failed upstream response's body.
"""
from __future__ import annotations
import httpx
from litellm._logging import verbose_logger
from litellm.proxy._experimental.mcp_server.faults.types import (
GATEWAY_CAPABILITY_CODES,
GATEWAY_CREDENTIAL_CODES,
MAX_WIRE_FIELD_CHARS,
CallerRejected,
CredentialSource,
GatewayRejected,
UpstreamOAuthFault,
UpstreamProtocolFault,
UpstreamReportedFault,
)
def _safe_text(response: httpx.Response) -> str:
try:
return response.text
except Exception:
return ""
def _safe_json(response: httpx.Response) -> object:
try:
return response.json()
except Exception:
return None
def _bounded_field(value: object) -> str | None:
if not isinstance(value, str) or not value:
return None
return value[:MAX_WIRE_FIELD_CHARS]
def _log_out_of_contract(endpoint_kind: str, response: httpx.Response, log_context: str) -> None:
verbose_logger.warning(
"MCP upstream %s endpoint (%s) returned HTTP %s outside the OAuth error contract (first %s chars): %s",
endpoint_kind,
log_context,
response.status_code,
MAX_WIRE_FIELD_CHARS,
_safe_text(response)[:MAX_WIRE_FIELD_CHARS],
)
def _classify_oauth_error_code(
code: str,
description: str | None,
error_uri: str | None,
credential_source: CredentialSource,
log_context: str,
) -> UpstreamOAuthFault:
"""Blame assignment for a contract-conformant OAuth error code, shared by the token and DCR
classifiers. Codes by which the upstream blames itself keep that blame; ``invalid_target`` is a
gateway capability gap (RFC 8707 resource indicators, LIT-4339) no matter whose credentials were
presented; credential-indicting codes follow the credential source; everything else, including
codes we do not recognize, is the caller's to act on. The upstream's HTTP status is deliberately
never consulted: status derives from this classification at render time, which is what keeps
status and code from contradicting each other."""
if code == "server_error" or code == "temporarily_unavailable":
return UpstreamReportedFault(code=code)
if code in GATEWAY_CAPABILITY_CODES:
verbose_logger.warning(
"MCP server %s: the upstream authorization server rejected the request with "
"invalid_target; it may require RFC 8707 resource indicators, which the gateway "
"does not send yet (tracked as LIT-4339)",
log_context,
)
return GatewayRejected(code=code)
if credential_source == "gateway_stored" and code in GATEWAY_CREDENTIAL_CODES:
verbose_logger.warning(
"MCP server %s: upstream authorization server rejected the gateway's configured client "
"credentials (%s): %s",
log_context,
code,
description or "<no description>",
)
return GatewayRejected(code=code)
return CallerRejected(code=code, description=description, error_uri=error_uri)
def classify_upstream_token_rejection(
response: httpx.Response,
credential_source: CredentialSource,
log_context: str,
) -> UpstreamOAuthFault:
"""Classify a token-endpoint rejection into exactly one fault: a body with an RFC 6749 §5.2
``error`` field goes through blame assignment (:func:`_classify_oauth_error_code`); anything
without a usable ``error`` field is an upstream protocol fault."""
parsed = _safe_json(response)
fields = parsed if isinstance(parsed, dict) else {}
code = _bounded_field(fields.get("error"))
if code is None:
_log_out_of_contract("token", response, log_context)
return UpstreamProtocolFault(note=f"upstream token endpoint returned HTTP {response.status_code}")
return _classify_oauth_error_code(
code,
description=_bounded_field(fields.get("error_description")),
error_uri=_bounded_field(fields.get("error_uri")),
credential_source=credential_source,
log_context=log_context,
)
def classify_upstream_dcr_rejection(response: httpx.Response, log_context: str) -> UpstreamOAuthFault:
"""Classify a dynamic-client-registration rejection. RFC 7591 §3.2.2 errors carry
``error`` / ``error_description`` and go through the same blame assignment as token errors
(registration sends no client credentials, so credential codes stay caller-actionable); anything
without a usable ``error`` field is an upstream protocol fault."""
parsed = _safe_json(response)
fields = parsed if isinstance(parsed, dict) else {}
code = _bounded_field(fields.get("error"))
if code is None:
_log_out_of_contract("registration", response, log_context)
return UpstreamProtocolFault(note=f"upstream registration failed with HTTP {response.status_code}")
return _classify_oauth_error_code(
code,
description=_bounded_field(fields.get("error_description")),
error_uri=None,
credential_source="caller_supplied",
log_context=log_context,
)

View file

@ -0,0 +1,89 @@
"""Render upstream OAuth/DCR faults onto the wire. The only place that chooses statuses and bodies
for these faults, so every consumer emits the same contract: RFC 6749 §5.2-shaped JSON with the §5.1
no-store headers on token endpoints, HTTPException details on registration. Status, code, and prose
all derive from the fault tag; exhaustive matches keep a new fault arm from shipping unrendered.
"""
from __future__ import annotations
from fastapi.responses import JSONResponse
from typing_extensions import assert_never
from litellm.proxy._experimental.mcp_server.faults.types import UpstreamOAuthFault
from litellm.proxy._experimental.mcp_server.oauth_utils import TOKEN_NO_CACHE_HEADERS
def _gateway_rejected_description(code: str) -> str:
if code == "invalid_target":
return (
"the upstream authorization server rejected the request (invalid_target); "
"it may require RFC 8707 resource indicators, which the gateway does not send yet"
)
return (
f"the upstream authorization server rejected the gateway's configured client credentials "
f"({code}); verify the MCP server's client_id and client_secret"
)
def _upstream_reported_status_and_description(code: str) -> tuple[int, str]:
if code == "temporarily_unavailable":
return 503, "the upstream authorization server is temporarily unavailable; retry shortly"
return 502, "the upstream authorization server reported an internal error"
def render_token_fault(fault: UpstreamOAuthFault) -> JSONResponse:
"""RFC 6749 §5.2 response for a token-endpoint fault. Caller-actionable rejections relay the
upstream's code on the status that code implies (401 for invalid_client per §5.2, else 400);
gateway-side faults are 502 ``server_error`` with gateway-authored prose so a caller is never
blamed for, or shown the internals of, a failure only the operator can fix."""
match fault.tag:
case "caller_rejected":
content = {
"error": fault.code,
**({"error_description": fault.description} if fault.description else {}),
**({"error_uri": fault.error_uri} if fault.error_uri else {}),
}
status_code = 401 if fault.code == "invalid_client" else 400
return JSONResponse(status_code=status_code, content=content, headers=TOKEN_NO_CACHE_HEADERS)
case "gateway_rejected":
return JSONResponse(
status_code=502,
content={
"error": "server_error",
"error_description": _gateway_rejected_description(fault.code),
},
headers=TOKEN_NO_CACHE_HEADERS,
)
case "upstream_reported_fault":
status_code, description = _upstream_reported_status_and_description(fault.code)
return JSONResponse(
status_code=status_code,
content={"error": fault.code, "error_description": description},
headers=TOKEN_NO_CACHE_HEADERS,
)
case "upstream_protocol_fault":
return JSONResponse(
status_code=502,
content={"error": "server_error", "error_description": fault.note},
headers=TOKEN_NO_CACHE_HEADERS,
)
case _:
assert_never(fault.tag)
def dcr_fault_detail(fault: UpstreamOAuthFault) -> tuple[int, str]:
"""Status and detail string for a registration fault, raised as HTTPException by the caller.
RFC 7591 §3.2.2 defines registration errors as 400, so a contract-conformant rejection is 400
regardless of the status the upstream chose; everything else is a 502 upstream fault."""
match fault.tag:
case "caller_rejected":
detail = f"{fault.code}: {fault.description}" if fault.description else fault.code
return 400, detail
case "gateway_rejected":
return 502, _gateway_rejected_description(fault.code)
case "upstream_reported_fault":
return _upstream_reported_status_and_description(fault.code)
case "upstream_protocol_fault":
return 502, fault.note
case _:
assert_never(fault.tag)

View file

@ -0,0 +1,79 @@
"""Fault taxonomy for upstream OAuth token and DCR registration failures.
Each fault is a frozen model on a ``tag`` literal. The tag alone decides the HTTP status, the wire
error code, and whose prose the caller sees, so those three facts can never disagree the way they can
when an upstream's status and error code are relayed independently.
"""
from __future__ import annotations
from typing import Literal, TypeAlias
from pydantic import BaseModel, ConfigDict
MAX_WIRE_FIELD_CHARS = 500
"""Bound on every upstream-derived string that crosses to a caller or into a log line."""
CredentialSource: TypeAlias = Literal["gateway_stored", "caller_supplied"]
"""Whose client credentials the gateway presented upstream: the MCP server's stored configuration or
credentials the caller supplied on the request. Decides whether a credential rejection is the
caller's problem to fix or the gateway operator's."""
GATEWAY_CREDENTIAL_CODES: frozenset[str] = frozenset({"invalid_client", "unauthorized_client"})
"""RFC 6749 error codes that indict the OAuth client's credentials or grant authorization. When the
gateway presented its own stored credentials, these are gateway-side faults the caller cannot act on;
when the caller supplied the credentials, they are the caller's to fix."""
GATEWAY_CAPABILITY_CODES: frozenset[str] = frozenset({"invalid_target"})
"""Codes that indict a gateway capability regardless of whose credentials were presented:
``invalid_target`` means the upstream wants RFC 8707 resource indicators, which the gateway does not
send yet (LIT-4339). Never the caller's fault."""
UPSTREAM_FAULT_CODES: frozenset[str] = frozenset({"server_error", "temporarily_unavailable"})
"""Codes by which the upstream blames itself. Relaying them as caller faults would invert blame, so
they classify as upstream-reported faults and render on the 5xx their meaning implies."""
class CallerRejected(BaseModel):
"""The upstream spoke the OAuth error contract and the failure is actionable by our caller
(e.g. ``invalid_grant``: re-run authorization). The code and its bounded prose relay on the
4xx status the code itself implies."""
model_config = ConfigDict(frozen=True)
tag: Literal["caller_rejected"] = "caller_rejected"
code: str
description: str | None = None
error_uri: str | None = None
class GatewayRejected(BaseModel):
"""The upstream rejected the request for a cause only the gateway operator can address: the
server's stored client credentials or a gateway capability gap. Not actionable by the caller:
rendered as 502 with gateway-authored prose naming the code; the upstream's prose goes to
server logs only."""
model_config = ConfigDict(frozen=True)
tag: Literal["gateway_rejected"] = "gateway_rejected"
code: str
class UpstreamReportedFault(BaseModel):
"""The upstream blamed itself in the OAuth vocabulary. Rendered on the 5xx the code implies
(``server_error`` 502, ``temporarily_unavailable`` 503) so blame and status agree."""
model_config = ConfigDict(frozen=True)
tag: Literal["upstream_reported_fault"] = "upstream_reported_fault"
code: Literal["server_error", "temporarily_unavailable"]
class UpstreamProtocolFault(BaseModel):
"""The upstream broke the error contract: no JSON ``error`` field, an undecodable body, or a
success response without a usable token. Rendered as 502 with a gateway-authored note; the
upstream body never crosses to the caller."""
model_config = ConfigDict(frozen=True)
tag: Literal["upstream_protocol_fault"] = "upstream_protocol_fault"
note: str
UpstreamOAuthFault: TypeAlias = CallerRejected | GatewayRejected | UpstreamReportedFault | UpstreamProtocolFault

View file

@ -21,11 +21,16 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import
EnvelopeKeys,
EnvelopeMintError,
OpenedEnvelope,
OpenedRefreshEnvelope,
RefreshCredential,
SealedEnvelope,
UpstreamTokenGrant,
is_envelope,
is_refresh_envelope,
mint_envelope,
mint_refresh_envelope,
open_envelope,
open_refresh_envelope,
)
_SIGNING_KEY_DOMAIN = b"litellm-mcp-bridge:envelope-signing:"
@ -92,6 +97,67 @@ def build_bridge_token_response(
return mint_envelope(identity, grant, keys, now)
def build_bridge_refresh_token_response(
identity: EnvelopeIdentity,
refresh: RefreshCredential,
keys: EnvelopeKeys,
now: datetime,
) -> SealedEnvelope | EnvelopeMintError:
"""Seal ``refresh`` for ``identity`` into the long-lived refresh envelope the token endpoint returns
alongside the access envelope, so the client can renew without re-authenticating. A thin, pure
wrapper over :func:`mint_refresh_envelope`; returns the mint error as a value for the caller to map.
"""
return mint_refresh_envelope(identity, refresh, keys, now)
class BridgeRefreshOpened(BaseModel):
"""A valid refresh envelope presented to the token endpoint: the identity to re-validate and renew
under, and the upstream refresh grant to exchange."""
model_config = ConfigDict(frozen=True)
tag: Literal["opened"] = "opened"
identity: EnvelopeIdentity
refresh: RefreshCredential
class BridgeRefreshInvalid(BaseModel):
"""The presented refresh grant is not a valid refresh envelope for this server (not refresh-shaped,
will not open, or minted for a different server); the token endpoint fails the refresh closed."""
model_config = ConfigDict(frozen=True)
tag: Literal["invalid"] = "invalid"
BridgeRefreshResult: TypeAlias = BridgeRefreshOpened | BridgeRefreshInvalid
def open_bridge_refresh_envelope(
refresh_value: str,
keys: EnvelopeKeys,
now: datetime,
expected_server_id: str,
) -> BridgeRefreshResult:
"""Open a refresh envelope a bridge ``oauth_delegate`` client presented on a refresh_token grant.
The token-endpoint mirror of :func:`resolve_bridge_envelope`: strips an optional ``Bearer`` scheme,
then returns ``BridgeRefreshOpened`` with the recovered identity and upstream refresh grant, or
``BridgeRefreshInvalid`` for anything that is not a valid refresh envelope for this server. Never
raises; total over hostile input via :func:`open_refresh_envelope`. ``expected_server_id`` binds the
envelope to the server the request targets, so a refresh envelope minted for one server cannot renew
against another. A raw upstream refresh token (not envelope-shaped) is ``BridgeRefreshInvalid``: this
mode never hands the client a bare upstream refresh token, so it must never accept one.
"""
candidate = _strip_bearer(refresh_value)
if not is_refresh_envelope(candidate):
return BridgeRefreshInvalid()
opened = open_refresh_envelope(candidate, keys, now)
if not isinstance(opened, OpenedRefreshEnvelope):
return BridgeRefreshInvalid()
if opened.identity.server_id != expected_server_id:
return BridgeRefreshInvalid()
return BridgeRefreshOpened(identity=opened.identity, refresh=opened.refresh)
class NotBridgeEnvelope(BaseModel):
"""The bearer is not an envelope; admission continues on its normal path."""
@ -128,10 +194,12 @@ def _strip_bearer(value: str) -> str:
def is_bridge_envelope_shaped(authorization_value: str) -> bool:
"""Cheap, keyless test that an ``Authorization`` value carries an envelope (optional
``Bearer`` scheme stripped). The admission edge engages the bridge arm only for an
envelope, so a plain upstream bearer falls through to normal oauth2 admission."""
return is_envelope(_strip_bearer(authorization_value))
"""Cheap, keyless test that an ``Authorization`` value carries an envelope of either kind (optional
``Bearer`` scheme stripped). The admission edge engages the bridge arm for an access envelope (to
admit) and for a refresh envelope (to reject it explicitly, since a refresh credential is never
usable at the tool-call edge); a plain upstream bearer falls through to normal oauth2 admission."""
candidate = _strip_bearer(authorization_value)
return is_envelope(candidate) or is_refresh_envelope(candidate)
def resolve_bridge_envelope(
@ -148,6 +216,10 @@ def resolve_bridge_envelope(
envelope, and ``BridgeEnvelopeInvalid`` for an envelope-shaped bearer that will not
open. Never raises: it is total over hostile input via :func:`open_envelope`.
A refresh envelope is ``BridgeEnvelopeInvalid`` here: it is a valid gateway credential but only ever
presented back to the token endpoint, never usable to authenticate a tool call, so admission must
fail it closed rather than let it fall through to another arm.
``expected_server_id`` is the ``server_id`` of the MCP server the request targets; an
opened envelope whose sealed ``server_id`` does not match is rejected as
``BridgeEnvelopeInvalid``. Binding here (rather than leaving it to the caller) prevents
@ -157,6 +229,8 @@ def resolve_bridge_envelope(
unlike ``hmac.compare_digest`` on ``str``, does not raise on a non-ASCII server_id.
"""
candidate = _strip_bearer(authorization_value)
if is_refresh_envelope(candidate):
return BridgeEnvelopeInvalid()
if not is_envelope(candidate):
return NotBridgeEnvelope()
opened = open_envelope(candidate, keys, now)

View file

@ -44,18 +44,33 @@ from pydantic import BaseModel, ConfigDict, Field, SecretStr, ValidationError
from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value, encrypt_value
ENVELOPE_PREFIX = "llm_env_"
"""Marker prefix on every serialized envelope so the edge can cheaply tell an envelope
"""Marker prefix on every serialized ACCESS envelope so the edge can cheaply tell an envelope
from a raw upstream token before doing any cryptography."""
REFRESH_ENVELOPE_PREFIX = "llm_refresh_"
"""Marker prefix on every serialized REFRESH envelope. A distinct prefix keeps the two credentials
routable without crypto and, together with the signed ``kind`` claim, stops one from being presented
where the other is expected: a refresh envelope carries a long-lived upstream refresh token and is only
ever presented back to the token endpoint, never forwarded upstream on a tool call."""
ENVELOPE_ISSUER = "litellm-mcp-bridge"
"""``iss`` claim stamped into every envelope and required back on open."""
MAX_ENVELOPE_TTL_SECONDS = 3600
"""Hard ceiling on envelope lifetime. ``exp`` is ``min(upstream expires_in, this cap)``
"""Hard ceiling on ACCESS envelope lifetime. ``exp`` is ``min(upstream expires_in, this cap)``
(the cap alone when the upstream omits ``expires_in``), matching the 1h lifetime of the
BYOK session bearer this module's signing approach is borrowed from: a client-held
credential should never outlive a bounded window even when the upstream token does."""
MAX_REFRESH_ENVELOPE_TTL_SECONDS = 1209600
"""Hard ceiling on REFRESH envelope lifetime (14 days). A refresh envelope only renews the short-lived
access envelope, and each renewal re-validates the sealed litellm key (revocation gates it) and is
re-minted with a fresh window, so the practical bound is idle time, not a fixed session. ``exp`` is
``min(upstream refresh_expires_in, this cap)`` (the cap alone when the upstream omits it); if the
upstream refresh token dies first, the next renewal simply fails at the upstream and the client
re-authenticates. The value is deliberately far shorter than a typical upstream refresh-token lifetime
so a leaked refresh envelope is bounded even if the upstream would have honoured it for longer."""
MAX_ENVELOPE_BYTES = 12288
"""Size cap on the final serialized envelope (prefix + JWT, in bytes). Upstream JWTs
commonly run 2-4KB; base64 plus encryption overhead roughly doubles that inside the
@ -66,21 +81,48 @@ typed error, never truncated."""
_ENVELOPE_JWT_ALGORITHM = "HS256"
EnvelopeKind = Literal["access", "refresh"]
"""Which credential an envelope is. Stamped into the signed claims and required to match on open, so a
signature-valid envelope of one kind cannot be replayed as the other even if its wire prefix is swapped
(the prefix is not part of the signed payload; this claim is)."""
EnvelopeSubjectType: TypeAlias = Literal["key_hash", "user_id"]
"""Discriminator for what litellm principal the envelope binds the grant to.
``key_hash`` is a hashed virtual key (the scripted two-header client mints under the key it
presents at the token endpoint); ``user_id`` is a litellm user subject (the interactive DCR
client mints under the SSO-authenticated user, which is the only identity that browser login
yields). Admission reloads a key record for the first and a user record for the second, then
runs both through the same live-policy gate, so team/org/budget/revocation enforcement is
identical either way."""
class EnvelopeIdentity(BaseModel):
"""The litellm identity the envelope binds the inner grant to.
"""The litellm principal the envelope binds the inner grant to.
``key_hash`` is the hashed litellm key that authorized the mint, never a raw
credential (and the edge rejects a bare hash presented as a bearer). Admission
reloads the live key record by it, so the key's current team/org/object-permission
restrictions and its revocation state are enforced at use time rather than frozen at
mint time. ``server_id`` binds the envelope to one MCP server so it cannot be replayed
across a server boundary.
``subject`` is the principal identifier and ``subject_type`` says how to resolve it: a
hashed litellm key (``key_hash``) or a litellm user id (``user_id``), never a raw
credential (and the edge rejects a bare hash or id presented as a bearer). Admission
reloads the live record by it, so the principal's current team/org restrictions and its
revocation state are enforced at use time rather than frozen at mint time. ``server_id``
binds the envelope to one MCP server so it cannot be replayed across a server boundary.
"""
model_config = ConfigDict(frozen=True)
server_id: str = Field(min_length=1)
key_hash: str = Field(min_length=1)
subject_type: EnvelopeSubjectType
subject: str = Field(min_length=1)
def key_hash_identity(server_id: str, key_hash: str) -> EnvelopeIdentity:
"""The identity for the scripted client that mints under a presented virtual key."""
return EnvelopeIdentity(server_id=server_id, subject_type="key_hash", subject=key_hash)
def user_identity(server_id: str, user_id: str) -> EnvelopeIdentity:
"""The identity for the interactive DCR client that mints under its SSO user subject."""
return EnvelopeIdentity(server_id=server_id, subject_type="user_id", subject=user_id)
class UpstreamTokenGrant(BaseModel):
@ -99,6 +141,21 @@ class UpstreamTokenGrant(BaseModel):
expires_in: int | None = Field(default=None, gt=0)
class RefreshCredential(BaseModel):
"""The upstream refresh grant sealed inside a refresh envelope.
Only the refresh token (plus the scope to re-request and the refresh token's own lifetime, when the
upstream reports it) is sealed; the access token is never in a refresh envelope. ``refresh_token`` is
a ``SecretStr`` so reprs never leak it, and ``expires_in`` (the refresh token's lifetime, not the
access token's) must be positive when present.
"""
model_config = ConfigDict(frozen=True)
refresh_token: SecretStr = Field(min_length=1)
scope: str | None = None
expires_in: int | None = Field(default=None, gt=0)
class EnvelopeKeys(BaseModel):
"""Injected key material: the HS256 signing key and the symmetric encryption key.
@ -121,13 +178,21 @@ class SealedEnvelope(BaseModel):
class OpenedEnvelope(BaseModel):
"""A validated envelope: the identity it was minted for and the recovered grant."""
"""A validated access envelope: the identity it was minted for and the recovered grant."""
model_config = ConfigDict(frozen=True)
identity: EnvelopeIdentity
grant: UpstreamTokenGrant
class OpenedRefreshEnvelope(BaseModel):
"""A validated refresh envelope: the identity it was minted for and the recovered refresh grant."""
model_config = ConfigDict(frozen=True)
identity: EnvelopeIdentity
refresh: RefreshCredential
class EnvelopeTooLarge(BaseModel):
"""The serialized envelope exceeded ``MAX_ENVELOPE_BYTES``; carries sizes only."""
@ -199,8 +264,10 @@ class _EnvelopeClaims(BaseModel):
iss: str
iat: int
exp: int
kind: EnvelopeKind
server_id: str = Field(min_length=1)
key_hash: str = Field(min_length=1)
subject_type: EnvelopeSubjectType
subject: str = Field(min_length=1)
grant: str = Field(min_length=1)
@ -213,11 +280,25 @@ class _GrantWire(BaseModel):
expires_in: int | None = None
class _RefreshWire(BaseModel):
model_config = ConfigDict(frozen=True)
refresh_token: str
scope: str | None = None
expires_in: int | None = None
def is_envelope(candidate: str) -> bool:
"""Cheap prefix check so the edge can route envelopes vs raw tokens without crypto."""
"""Cheap prefix check for an ACCESS envelope so the edge can route envelopes vs raw tokens without
crypto. A refresh envelope has a different prefix and is not an access envelope."""
return candidate.startswith(ENVELOPE_PREFIX)
def is_refresh_envelope(candidate: str) -> bool:
"""Cheap prefix check for a REFRESH envelope so the token endpoint can route a refresh grant that
carries an envelope vs a raw upstream refresh token without crypto."""
return candidate.startswith(REFRESH_ENVELOPE_PREFIX)
def mint_envelope(
identity: EnvelopeIdentity,
grant: UpstreamTokenGrant,
@ -231,23 +312,15 @@ def mint_envelope(
serialized envelope exceeds ``MAX_ENVELOPE_BYTES``.
"""
expires_at = now + timedelta(seconds=_envelope_ttl_seconds(grant.expires_in))
claims = _EnvelopeClaims(
iss=ENVELOPE_ISSUER,
iat=int(now.timestamp()),
exp=int(expires_at.timestamp()),
server_id=identity.server_id,
key_hash=identity.key_hash,
grant=_encrypt_grant_blob(_grant_plaintext(grant), keys.encryption_key),
return _seal(
kind="access",
prefix=ENVELOPE_PREFIX,
identity=identity,
grant_blob=_encrypt_grant_blob(_grant_plaintext(grant), keys.encryption_key),
expires_at=expires_at,
signing_key=keys.signing_key,
now=now,
)
token = ENVELOPE_PREFIX + jwt.encode(
claims.model_dump(),
keys.signing_key.get_secret_value(),
algorithm=_ENVELOPE_JWT_ALGORITHM,
)
size_bytes = len(token.encode("utf-8"))
if size_bytes > MAX_ENVELOPE_BYTES:
return EnvelopeTooLarge(size_bytes=size_bytes, max_bytes=MAX_ENVELOPE_BYTES)
return SealedEnvelope(token=SecretStr(token), expires_at=expires_at)
def open_envelope(
@ -263,35 +336,136 @@ def open_envelope(
re-derived, so it is stale by up to the envelope's lifetime; callers that need a
live remaining lifetime should use ``now`` against the upstream, not this field.
"""
if not is_envelope(candidate):
return NotAnEnvelope()
# UTF-8 byte length is never below character length, so a character count already over the
# cap rejects an oversize candidate in O(1) without encoding it; the exact byte check then
# runs only on candidates already bounded to <= MAX_ENVELOPE_BYTES characters.
if len(candidate) > MAX_ENVELOPE_BYTES:
return MalformedPayload()
if len(candidate.encode("utf-8", "surrogatepass")) > MAX_ENVELOPE_BYTES:
return MalformedPayload()
claims = _decode_claims(candidate.removeprefix(ENVELOPE_PREFIX), keys.signing_key)
claims = _open_claims(candidate, prefix=ENVELOPE_PREFIX, expected_kind="access", keys=keys, now=now)
if not isinstance(claims, _EnvelopeClaims):
return claims
if now.timestamp() >= claims.exp:
return Expired()
grant = _decrypt_grant(claims.grant, keys.encryption_key)
if not isinstance(grant, UpstreamTokenGrant):
return grant
return OpenedEnvelope(
identity=EnvelopeIdentity(server_id=claims.server_id, key_hash=claims.key_hash),
identity=EnvelopeIdentity(server_id=claims.server_id, subject_type=claims.subject_type, subject=claims.subject),
grant=grant,
)
def mint_refresh_envelope(
identity: EnvelopeIdentity,
refresh: RefreshCredential,
keys: EnvelopeKeys,
now: datetime,
) -> SealedEnvelope | EnvelopeMintError:
"""Seal ``refresh`` for ``identity`` into a long-lived, client-held refresh envelope.
``exp`` is ``min(refresh.expires_in, MAX_REFRESH_ENVELOPE_TTL_SECONDS)`` seconds from ``now`` (the
cap alone when the upstream omits the refresh lifetime). Sealing a distinct ``kind="refresh"`` claim
is what keeps a refresh envelope from ever opening as an access credential at the MCP edge. Returns
``EnvelopeTooLarge`` when the serialized envelope exceeds ``MAX_ENVELOPE_BYTES``.
"""
expires_at = now + timedelta(seconds=_refresh_ttl_seconds(refresh.expires_in))
return _seal(
kind="refresh",
prefix=REFRESH_ENVELOPE_PREFIX,
identity=identity,
grant_blob=_encrypt_grant_blob(_refresh_plaintext(refresh), keys.encryption_key),
expires_at=expires_at,
signing_key=keys.signing_key,
now=now,
)
def open_refresh_envelope(
candidate: str,
keys: EnvelopeKeys,
now: datetime,
) -> OpenedRefreshEnvelope | EnvelopeOpenError:
"""Validate a refresh ``candidate`` and recover the identity and inner refresh grant.
Total over hostile input exactly like :func:`open_envelope`: every invalid, expired, tampered,
wrong-kind, or undecryptable candidate maps to a distinct ``EnvelopeOpenError`` variant, never a
raise. The ``kind="refresh"`` claim is required, so an access envelope re-prefixed as a refresh one
is rejected as ``MalformedPayload``.
"""
claims = _open_claims(candidate, prefix=REFRESH_ENVELOPE_PREFIX, expected_kind="refresh", keys=keys, now=now)
if not isinstance(claims, _EnvelopeClaims):
return claims
refresh = _decrypt_refresh(claims.grant, keys.encryption_key)
if not isinstance(refresh, RefreshCredential):
return refresh
return OpenedRefreshEnvelope(
identity=EnvelopeIdentity(server_id=claims.server_id, subject_type=claims.subject_type, subject=claims.subject),
refresh=refresh,
)
def _seal(
kind: EnvelopeKind,
prefix: str,
identity: EnvelopeIdentity,
grant_blob: str,
expires_at: datetime,
signing_key: SecretStr,
now: datetime,
) -> SealedEnvelope | EnvelopeTooLarge:
"""Sign the claims for either envelope kind and enforce the size cap. Shared by both mints so the
JWT shape, issuer, and size guard cannot drift between access and refresh envelopes."""
claims = _EnvelopeClaims(
iss=ENVELOPE_ISSUER,
iat=int(now.timestamp()),
exp=int(expires_at.timestamp()),
kind=kind,
server_id=identity.server_id,
subject_type=identity.subject_type,
subject=identity.subject,
grant=grant_blob,
)
token = prefix + jwt.encode(claims.model_dump(), signing_key.get_secret_value(), algorithm=_ENVELOPE_JWT_ALGORITHM)
size_bytes = len(token.encode("utf-8"))
if size_bytes > MAX_ENVELOPE_BYTES:
return EnvelopeTooLarge(size_bytes=size_bytes, max_bytes=MAX_ENVELOPE_BYTES)
return SealedEnvelope(token=SecretStr(token), expires_at=expires_at)
def _open_claims(
candidate: str,
prefix: str,
expected_kind: EnvelopeKind,
keys: EnvelopeKeys,
now: datetime,
) -> _EnvelopeClaims | EnvelopeOpenError:
"""Prefix-route, size-bound, signature-verify, kind-check, and expiry-check an attacker-controlled
candidate, shared by both openers so the security gate is identical for access and refresh. Returns
the validated claims or a distinct ``EnvelopeOpenError``; never raises."""
if not candidate.startswith(prefix):
return NotAnEnvelope()
# UTF-8 byte length is never below character length, so a character count already over the cap
# rejects an oversize candidate in O(1) without encoding it; the exact byte check then runs only on
# candidates already bounded to <= MAX_ENVELOPE_BYTES characters.
if len(candidate) > MAX_ENVELOPE_BYTES:
return MalformedPayload()
if len(candidate.encode("utf-8", "surrogatepass")) > MAX_ENVELOPE_BYTES:
return MalformedPayload()
claims = _decode_claims(candidate.removeprefix(prefix), keys.signing_key)
if not isinstance(claims, _EnvelopeClaims):
return claims
if claims.kind != expected_kind:
return MalformedPayload()
if now.timestamp() >= claims.exp:
return Expired()
return claims
def _envelope_ttl_seconds(upstream_expires_in: int | None) -> int:
if upstream_expires_in is None:
return MAX_ENVELOPE_TTL_SECONDS
return min(upstream_expires_in, MAX_ENVELOPE_TTL_SECONDS)
def _refresh_ttl_seconds(upstream_refresh_expires_in: int | None) -> int:
if upstream_refresh_expires_in is None:
return MAX_REFRESH_ENVELOPE_TTL_SECONDS
return min(upstream_refresh_expires_in, MAX_REFRESH_ENVELOPE_TTL_SECONDS)
def _grant_plaintext(grant: UpstreamTokenGrant) -> str:
wire = _GrantWire(
access_token=grant.access_token.get_secret_value(),
@ -303,6 +477,15 @@ def _grant_plaintext(grant: UpstreamTokenGrant) -> str:
return wire.model_dump_json(exclude_none=True)
def _refresh_plaintext(refresh: RefreshCredential) -> str:
wire = _RefreshWire(
refresh_token=refresh.refresh_token.get_secret_value(),
scope=refresh.scope,
expires_in=refresh.expires_in,
)
return wire.model_dump_json(exclude_none=True)
def _decode_claims(
compact: str,
signing_key: SecretStr,
@ -364,3 +547,22 @@ def _decrypt_grant(
return UpstreamTokenGrant.model_validate_json(plaintext)
except ValidationError:
return MalformedPayload()
def _decrypt_refresh(
blob: str,
encryption_key: SecretStr,
) -> RefreshCredential | DecryptFailed | MalformedPayload:
from nacl.exceptions import CryptoError
try:
plaintext = decrypt_value(
value=base64.urlsafe_b64decode(blob),
signing_key=encryption_key.get_secret_value(),
)
except (CryptoError, ValueError):
return DecryptFailed()
try:
return RefreshCredential.model_validate_json(plaintext)
except ValidationError:
return MalformedPayload()

View file

@ -3719,8 +3719,15 @@ if MCP_AVAILABLE:
headers={"www-authenticate": upstream_www_authenticate},
)
def _get_authorization_header_from_scope(scope: Scope) -> Optional[str]:
"""First ``Authorization`` header value in the ASGI scope, or None."""
for key, value in scope.get("headers", []):
if key.lower() == b"authorization":
return value.decode("latin-1")
return None
def _scope_has_authorization_header(scope: Scope) -> bool:
return any(key.lower() == b"authorization" for key, _ in scope.get("headers", []))
return _get_authorization_header_from_scope(scope) is not None
def _get_forwarded_auth_from_scope(scope: Scope) -> Optional[str]:
"""Return the upstream-bound ``Authorization`` header value, or None.
@ -3733,17 +3740,24 @@ if MCP_AVAILABLE:
``MCPRequestHandler.process_mcp_request``), and forwarding it upstream
would leak the proxy key to a third-party MCP server.
"""
authorization = None
has_litellm_key_header = False
for key, value in scope.get("headers", []):
key_lower = key.lower()
if key_lower == b"authorization":
authorization = value.decode("latin-1")
elif key_lower == b"x-litellm-api-key":
has_litellm_key_header = True
has_litellm_key_header = any(key.lower() == b"x-litellm-api-key" for key, _ in scope.get("headers", []))
if not has_litellm_key_header:
return None
return authorization
return _get_authorization_header_from_scope(scope)
def _is_delegate_upstream_probe_target(server: MCPServer) -> bool:
"""Whether ``server`` is an interactive delegate-auth server whose client-supplied
token should be preflighted upstream.
Mirrors the anonymous-delegate gate in ``get_allowed_mcp_servers``: the flow is
resolved via ``effective_oauth2_flow`` so an unstamped M2M-shape row fails closed
(its stored client credentials drive egress; the caller's bearer is irrelevant).
"""
return (
server.auth_type == MCPAuth.oauth2
and server.delegate_auth_to_upstream is True
and MCPServerManager.effective_oauth2_flow(server) != "client_credentials"
)
async def _probe_upstream_auth(
url: str,
@ -3805,7 +3819,7 @@ if MCP_AVAILABLE:
mcp_servers: Optional[List[str]],
client_ip: Optional[str],
) -> None:
"""Probe pass-through upstream servers in parallel before the MCP session starts.
"""Probe pass-through and delegate-auth upstream servers in parallel before the MCP session starts.
Only servers the caller's key is already authorized to reach are probed —
the list is derived from _get_allowed_mcp_servers so that a user cannot
@ -3813,11 +3827,42 @@ if MCP_AVAILABLE:
The MCP SDK commits HTTP 200 headers before invoking handlers, so a 401
can only be returned before that point. This function raises HTTPException(401)
with a WWW-Authenticate header if any upstream rejects the client token.
with a WWW-Authenticate header if any upstream rejects the client token, or 403
if the upstream accepts it but forbids the caller.
Fails-open: network errors are logged and the request is allowed through.
Delegate-auth servers (``auth_type=oauth2`` + ``delegate_auth_to_upstream``)
are probed with the caller's bare ``Authorization`` bearer. That bearer is only
an upstream token (never a LiteLLM key) when admission took the delegate bypass,
so the delegate target is resolved through ``get_mcp_server_by_name`` -- the same
resolver admission used -- rather than the wider allowed-server prefix/access-group
matching. A name that only reaches a delegate server via server_id or an access
group would have been admitted as a real LiteLLM key, so probing it would leak that
key upstream; requiring the admission-resolver match closes that gap. Without the
probe a rejected token is absorbed by the tools/list handler and masked as an empty
tool list. Gated to single-server routes so one rejected token cannot 401 a
multi-server aggregate connect, matching the OBO preflight gating; the challenge
echoes the requested name so aliased routes get the same resource_metadata URL as
the tokenless preemptive challenge.
"""
forwarded_auth = _get_forwarded_auth_from_scope(scope)
if not forwarded_auth:
requested_single_target = mcp_servers[0] if mcp_servers is not None and len(mcp_servers) == 1 else None
# The bare Authorization header (no x-litellm-api-key) is a valid upstream token
# only when admission classified it as one, i.e. the single requested name resolves
# to a delegate server under admission's own resolver. Resolve it the same way here
# so a server_id- or access-group-named delegate (which admission would have treated
# as a LiteLLM key) is never probed with that key.
delegate_server = (
global_mcp_server_manager.get_mcp_server_by_name(requested_single_target, client_ip=client_ip)
if requested_single_target
else None
)
delegate_auth = (
_get_authorization_header_from_scope(scope)
if delegate_server is not None and _is_delegate_upstream_probe_target(delegate_server)
else None
)
if not forwarded_auth and not delegate_auth:
return
# Use the authorized server set, not the raw user-supplied names, so that
@ -3827,33 +3872,49 @@ if MCP_AVAILABLE:
mcp_servers=mcp_servers,
client_ip=client_ip,
)
passthrough_servers = [
srv
for srv in allowed_servers
# Restrict to genuine OAuth pass-through servers (auth_type none +
# Authorization in extra_headers). Gateway-managed OAuth2 servers
# must not receive the ``resource_metadata=`` challenge emitted
# below — they require ``authorization_uri=`` pointing at the
# gateway AS metadata. ``is_oauth_passthrough`` already requires
# ``auth_type in (None, MCPAuth.none)``, which is mutually
# exclusive with ``has_client_credentials`` (oauth2 + M2M flow),
# so M2M servers are implicitly excluded here.
if srv.is_oauth_passthrough
]
if not passthrough_servers:
passthrough_targets: Tuple[Tuple[MCPServer, str, str], ...] = (
tuple(
(srv, forwarded_auth, srv.name)
for srv in allowed_servers
# Restrict to genuine OAuth pass-through servers (auth_type none +
# Authorization in extra_headers). Gateway-managed OAuth2 servers
# must not receive the ``resource_metadata=`` challenge emitted
# below — they require ``authorization_uri=`` pointing at the
# gateway AS metadata. ``is_oauth_passthrough`` already requires
# ``auth_type in (None, MCPAuth.none)``, which is mutually
# exclusive with ``has_client_credentials`` (oauth2 + M2M flow),
# so M2M servers are implicitly excluded here.
if srv.is_oauth_passthrough
)
if forwarded_auth
else ()
)
# Probe the admission-resolved delegate server only when the caller is actually
# authorized for it (present in the IP-filtered allowed set), keyed by server_id.
delegate_targets: Tuple[Tuple[MCPServer, str, str], ...] = (
tuple(
(srv, delegate_auth, requested_single_target)
for srv in allowed_servers
if delegate_server is not None and srv.server_id == delegate_server.server_id
)
if delegate_auth and requested_single_target
else ()
)
probe_targets = passthrough_targets + delegate_targets
if not probe_targets:
return
probe_results = await asyncio.gather(
*[_probe_upstream_auth(srv.url or "", forwarded_auth) for srv in passthrough_servers]
*[_probe_upstream_auth(srv.url or "", auth_header) for srv, auth_header, _ in probe_targets]
)
for srv, (probe_status, _) in zip(passthrough_servers, probe_results):
for (srv, _, challenge_server_name), (probe_status, _) in zip(probe_targets, probe_results):
if probe_status == 401:
# Token is missing or expired: keep pass-through clients on the
# protected-resource discovery flow so they re-authorize against
# the upstream IdP metadata proxied by LiteLLM.
www_authenticate = _get_passthrough_www_authenticate(
scope=scope,
server_name=srv.name,
server_name=challenge_server_name,
invalid_token=True,
)
raise HTTPException(

View file

@ -1118,6 +1118,7 @@ class GenerateKeyRequest(KeyRequestBase):
class GenerateKeyResponse(KeyRequestBase):
key: str # type: ignore
key_name: Optional[str] = None
key_type: str | None = None
expires: Optional[datetime] = None
user_id: Optional[str] = None
token_id: Optional[str] = None
@ -2421,6 +2422,16 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase):
"is active as a reminder that hard enforcement is relaxed."
),
)
skip_user_budget_on_team_key: bool | None = Field(
None,
description=(
"If True, restores the legacy behavior where a user's personal "
"max_budget is NOT enforced when their key belongs to a team; only "
"the team (and team-member) budgets apply. Defaults to False, meaning "
"the user's personal max_budget is always enforced regardless of "
"whether the key belongs to a team (see GitHub issue #12905)."
),
)
user_url_validation: Optional[bool] = Field(
None,
description=(

View file

@ -626,26 +626,29 @@ async def common_checks(
)
async def _user_max_budget_check() -> None:
# 4.1 personal budget, if personal key
if (
(team_object is None or team_object.team_id is None)
and user_object is not None
and user_object.max_budget is not None
):
from litellm.proxy.proxy_server import get_current_spend
if user_object is None or user_object.max_budget is None:
return
skip_for_team = (
general_settings.get("skip_user_budget_on_team_key") is True
and team_object is not None
and team_object.team_id is not None
)
if skip_for_team:
return
from litellm.proxy.proxy_server import get_current_spend
user_budget = user_object.max_budget
user_spend = await get_current_spend(
counter_key=f"spend:user:{user_object.user_id}",
fallback_spend=user_object.spend or 0.0,
user_budget = user_object.max_budget
user_spend = await get_current_spend(
counter_key=f"spend:user:{user_object.user_id}",
fallback_spend=user_object.spend or 0.0,
max_budget=user_budget,
)
if math.isfinite(user_budget) and user_spend >= user_budget:
raise litellm.BudgetExceededError(
current_cost=user_spend,
max_budget=user_budget,
message=f"ExceededBudget: User={user_object.user_id} over budget. Spend={user_spend}, Budget={user_budget}",
)
if math.isfinite(user_budget) and user_spend >= user_budget:
raise litellm.BudgetExceededError(
current_cost=user_spend,
max_budget=user_budget,
message=f"ExceededBudget: User={user_object.user_id} over budget. Spend={user_spend}, Budget={user_budget}",
)
# Each scope reads a distinct counter key with no cross-scope ordering
# dependency, so the per-scope Redis-first reads run concurrently instead
@ -4383,14 +4386,23 @@ def _model_custom_llm_provider_matches_wildcard_pattern(model: str, allowed_mode
or
- `model=claude-3-5-sonnet-20240620`
- `allowed_model_pattern=anthropic/*`
A model that already carries a namespace get_llm_provider did not consume
(e.g. `bedrockz/anthropic.claude-...`) is never granted here: its provider was
inferred from a fragment of the full string, so rebuilding
`{provider}/{model}` would produce `bedrock/bedrockz/...` and slip an
unrecognized namespace through a `bedrock/*` key.
"""
try:
model, custom_llm_provider, _, _ = get_llm_provider(model=model)
stripped_model, custom_llm_provider, _, _ = get_llm_provider(model=model)
except Exception:
return False
if stripped_model == model and "/" in model:
return False
return is_model_allowed_by_pattern(
model=f"{custom_llm_provider}/{model}",
model=f"{custom_llm_provider}/{stripped_model}",
allowed_model_pattern=allowed_model_pattern,
)

View file

@ -1191,13 +1191,15 @@ async def _user_api_key_auth_builder(
return await handle_oauth2_proxy_request(request=request)
if general_settings.get("enable_jwt_auth", False) is True:
from litellm.proxy.proxy_server import premium_user
if premium_user is not True:
raise ValueError(f"JWT Auth is an enterprise only feature. {CommonProxyErrors.not_premium_user.value}")
is_jwt = jwt_handler.is_jwt(token=api_key)
verbose_proxy_logger.debug("is_jwt: %s", is_jwt)
if is_jwt:
from litellm.proxy.proxy_server import premium_user
if premium_user is not True:
raise ValueError(
f"JWT Auth is an enterprise only feature. {CommonProxyErrors.not_premium_user.value}"
)
# Try JWT-to-Virtual-Key mapping first to avoid
# unnecessary DB queries in auth_builder
do_standard_jwt_auth = True
@ -2442,6 +2444,7 @@ async def _reserve_budget_after_common_checks(
proxy_logging_obj=proxy_logging_obj,
end_user_id=end_user_id,
end_user_object=end_user_object,
skip_user_budget_on_team_key=general_settings.get("skip_user_budget_on_team_key") is True,
)

View file

@ -8,6 +8,10 @@ from litellm.proxy._types import (
)
from litellm.secret_managers.main import str_to_bool
# Bounds the __cause__/__context__ walk in is_database_service_unavailable_error_in_chain.
# Real exception chains are a few links deep; the cap also makes the walk cycle-safe.
_MAX_EXCEPTION_CHAIN_DEPTH = 20
class PrismaDBExceptionHandler:
"""
@ -218,6 +222,32 @@ class PrismaDBExceptionHandler:
),
)
@staticmethod
def is_database_service_unavailable_error_in_chain(e: BaseException) -> bool:
"""Like ``is_database_service_unavailable_error`` but also walks the
``__cause__`` / ``__context__`` chain.
``is_database_service_unavailable_error`` classifies a single exception
by type, which a caller that catches a raw DB failure and re-raises a
domain exception of a different type defeats. ``get_user_object`` in
``litellm/proxy/auth/auth_checks.py`` is the concrete case: it wraps
every DB error, a genuine outage included, in a bare ``ValueError``
whose original error survives only as ``__context__``. A type check on
the ``ValueError`` misses the outage, so the caller would mistake an
infrastructure fault for an auth failure. Walking the chain recovers the
real signal, which is the PEP 3134 way to inspect a wrapped cause.
The walk is depth-bounded, which also makes it cycle-safe.
"""
current: BaseException | None = e
for _ in range(_MAX_EXCEPTION_CHAIN_DEPTH):
if not isinstance(current, Exception):
return False
if PrismaDBExceptionHandler.is_database_service_unavailable_error(current):
return True
current = current.__cause__ or current.__context__
return False
@staticmethod
def handle_db_exception(e: Exception):
"""

View file

@ -199,9 +199,10 @@ model_list:
general_settings:
master_key: sk-1234
# Opt-in: let CheckBatchCost track cost for unmanaged Vertex batches created with a raw gs:// input_file_id.
# Requires a vertex_ai deployment configured for the batched model. Defaults to false.
# track_unmanaged_vertex_batch_cost: true
# Opt-in: let CheckBatchCost track cost for unmanaged batches created with a raw
# gs:// (Vertex) or s3:// (Bedrock) input_file_id. Requires a matching deployment
# configured for the batched model. Defaults to false.
# track_unmanaged_batch_cost: true
sandbox_tools:
- sandbox_tool_name: e2b_sandbox

View file

@ -34,6 +34,12 @@ def is_text_content_call_type(call_type: str) -> bool:
TEXT_PART_TYPES: FrozenSet[str] = frozenset({"text", "input_text", "output_text"})
# Responses-API item types whose ``output`` field carries user/tool text
# that guardrails should inspect. ``function_call_output`` is the
# built-in shape; ``custom_tool_call_output`` is the custom-tool
# counterpart (see ``ChatCompletionCustomToolCallOutput``).
_OUTPUT_ITEM_TYPES: frozenset[str] = frozenset({"function_call_output", "custom_tool_call_output"})
def _iter_text_parts_in_content(content: Any) -> Iterator[str]:
"""Yield text fragments from a ``message.content`` value (string or
@ -72,7 +78,7 @@ def _coerce_input_to_messages(input_value: Any) -> List[Dict[str, Any]]:
messages.append({"role": item.get("role") or "user", "content": [item]})
elif "content" in item:
messages.append({"role": item.get("role") or "user", "content": item["content"]})
elif item.get("type") == "function_call_output" and "output" in item:
elif item.get("type") in _OUTPUT_ITEM_TYPES and "output" in item:
messages.append({"role": item.get("role") or "tool", "content": item["output"]})
return messages
@ -157,7 +163,7 @@ def walk_user_text(data: Dict[str, Any], visit: Callable[[str], str]) -> int:
input_value[idx] = {**item, "text": visit(item["text"])}
elif "content" in item:
item["content"] = _rewrite_content(item["content"])
elif item.get("type") == "function_call_output" and "output" in item:
elif item.get("type") in _OUTPUT_ITEM_TYPES and "output" in item:
item["output"] = _rewrite_content(item["output"])
return visited

View file

@ -772,7 +772,13 @@ class LassoGuardrail(CustomGuardrail):
data: Request data (used for conversation_id generation and tools extraction)
cache: Cache instance for storing conversation_id (optional for post-call)
"""
payload: Dict[str, Any] = {"messages": messages, "messageType": message_type}
payload: Dict[str, Any] = {
"messages": messages,
"messageType": message_type,
# Drives the "Used By" badge on Lasso Application API Keys: every call from this
# integration is attributed as "litellm" on the keys list.
"source": {"type": "litellm"},
}
# Add optional parameters if available
if self.user_id:

View file

@ -44,12 +44,18 @@ from litellm.integrations.custom_guardrail import (
CustomGuardrail,
log_guardrail_information,
)
from litellm.litellm_core_utils.core_helpers import redact_nested_match_and_regex_keys
from litellm.litellm_core_utils.sensitive_data_masker import mask_credentials_in_payload
from litellm.llms.custom_httpx.http_handler import (
get_async_httpx_client,
httpxSpecialProvider,
)
from litellm.types.guardrails import GuardrailEventHooks
from litellm.types.utils import GenericGuardrailAPIInputs, GuardrailStatus
from litellm.types.utils import (
GenericGuardrailAPIInputs,
GuardrailStatus,
StandardLoggingGuardrailInformation,
)
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import (
@ -60,6 +66,13 @@ if TYPE_CHECKING:
)
def _sanitize_scan_result_for_logging(scan_result: dict) -> dict:
without_secrets = {key: value for key, value in scan_result.items() if key != "secret_fields"}
redacted = redact_nested_match_and_regex_keys(without_secrets)
masked = mask_credentials_in_payload(redacted if isinstance(redacted, dict) else without_secrets)
return masked if isinstance(masked, dict) else without_secrets
_DEFAULT_API_BASE = "https://api-xecguard.cycraft.ai"
_SCAN_ENDPOINT = "/xecguard/v1/scan"
_GROUNDING_ENDPOINT = "/xecguard/v1/grounding"
@ -246,16 +259,21 @@ class XecGuardGuardrail(CustomGuardrail):
"guardrail_intervened" if scan_result.get("decision") == "UNSAFE" else "success"
)
end_time = datetime.now()
kwargs["standard_logging_object"]["guardrail_information"] = {
"duration": (end_time - start_time).total_seconds(),
"end_time": end_time.timestamp(),
"guardrail_mode": "logging_only",
"guardrail_name": "xecguard",
"guardrail_response": scan_result,
"guardrail_status": guardrail_status,
"masked_entity_count": None,
"start_time": start_time.timestamp(),
}
slg = StandardLoggingGuardrailInformation(
guardrail_name=self.guardrail_name or "xecguard",
guardrail_mode=GuardrailEventHooks.logging_only,
guardrail_response=_sanitize_scan_result_for_logging(scan_result),
guardrail_status=guardrail_status,
start_time=start_time.timestamp(),
end_time=end_time.timestamp(),
duration=(end_time - start_time).total_seconds(),
masked_entity_count=None,
)
existing = kwargs["standard_logging_object"].get("guardrail_information")
if isinstance(existing, list):
existing.append(slg)
else:
kwargs["standard_logging_object"]["guardrail_information"] = [slg]
except Exception as exc:
verbose_proxy_logger.debug(

View file

@ -28,11 +28,20 @@ from litellm.proxy.spend_tracking.spend_tracking_utils import (
)
from litellm.proxy.utils import ProxyUpdateSpend
from litellm.types.utils import (
CallTypes,
StandardLoggingPayload,
StandardLoggingPayloadErrorInformation,
)
from litellm.utils import get_end_user_id_for_cost_tracking
_PASS_THROUGH_CALL_TYPES: frozenset[str] = frozenset(
{
CallTypes.pass_through.value,
CallTypes.llm_passthrough_route.value,
CallTypes.allm_passthrough_route.value,
}
)
class _ProxyDBLogger(CustomLogger):
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
@ -219,11 +228,13 @@ class _ProxyDBLogger(CustomLogger):
verbose_proxy_logger.debug(
f"user_api_key {user_api_key}, user_id {user_id}, team_id {team_id}, end_user_id {end_user_id}"
)
call_type: Optional[str] = kwargs.get("call_type")
if _should_track_cost_callback(
user_api_key=user_api_key,
user_id=user_id,
team_id=team_id,
end_user_id=end_user_id,
call_type=call_type,
):
## UPDATE DATABASE
await _update_database_and_spend_counters(
@ -412,9 +423,15 @@ def _should_track_cost_callback(
user_id: Optional[str],
team_id: Optional[str],
end_user_id: Optional[str],
call_type: Optional[str] = None,
) -> bool:
"""
Determine if the cost callback should be tracked based on the kwargs
Pass-through endpoints can be configured with ``auth=false``, which leaves
the request with no key/user/team/end-user to attribute spend to. Those
requests still forward real provider traffic that operators expect to see
in request/usage logs, so they are tracked even when unauthenticated.
"""
# don't run track cost callback if user opted into disabling spend
@ -423,7 +440,7 @@ def _should_track_cost_callback(
if user_api_key is not None or user_id is not None or team_id is not None or end_user_id is not None:
return True
return False
return call_type in _PASS_THROUGH_CALL_TYPES
def _get_budget_reservation_from_metadata(metadata: dict) -> Optional[dict]:

View file

@ -42,7 +42,7 @@ from litellm.proxy._experimental.mcp_server.db import (
rotate_mcp_user_env_vars_master_key,
)
from litellm.proxy._types import *
from litellm.proxy._types import LiteLLM_VerificationToken
from litellm.proxy._types import LiteLLM_VerificationToken, hash_token
from litellm.proxy.auth.auth_checks import (
_delete_cache_key_object,
can_team_access_model,
@ -468,7 +468,10 @@ def handle_key_type(data: GenerateKeyRequest, data_json: dict) -> dict:
Handle the key type.
"""
key_type = data.key_type
data_json.pop("key_type", None)
if key_type is None:
data_json.pop("key_type", None)
return data_json
data_json["key_type"] = key_type.value
if key_type == LiteLLMKeyType.LLM_API:
data_json["allowed_routes"] = ["llm_api_routes"]
elif key_type == LiteLLMKeyType.MANAGEMENT:
@ -3566,6 +3569,7 @@ async def generate_key_helper_fn(
created_by: Optional[str] = None,
updated_by: Optional[str] = None,
allowed_routes: Optional[list] = None,
key_type: str | None = None,
sso_user_id: Optional[str] = None,
object_permission_id: Optional[str] = None, # object_permission_id <-> LiteLLM_ObjectPermissionTable
object_permission: Optional[LiteLLM_ObjectPermissionBase] = None,
@ -3706,6 +3710,7 @@ async def generate_key_helper_fn(
"created_by": created_by,
"updated_by": updated_by,
"allowed_routes": allowed_routes or [],
"key_type": key_type,
"object_permission_id": object_permission_id,
"router_settings": router_settings_json,
"access_group_ids": access_group_ids or [],
@ -3772,7 +3777,10 @@ async def generate_key_helper_fn(
return user_data
## CREATE KEY
verbose_proxy_logger.debug("prisma_client: Creating Key= %s", key_data)
verbose_proxy_logger.debug(
"prisma_client: Creating Key= %s",
{**key_data, "token": hash_token(token=token)},
)
create_key_response = await prisma_client.insert_data(data=key_data, table_name="key")
key_data["token_id"] = getattr(create_key_response, "token", None)

View file

@ -3550,7 +3550,7 @@ async def team_info(
try:
team_info: Optional[BaseModel] = await TeamRepository(prisma_client).table.find_unique(
where={"team_id": team_id},
include={"object_permission": True},
include={"litellm_model_table": True, "object_permission": True},
)
if team_info is None:
raise Exception

View file

@ -4034,30 +4034,41 @@ class MicrosoftSSOHandler:
base_url = MicrosoftSSOHandler.get_graph_api_base_url()
# Endpoint to get app role assignments for the given service principal
endpoint = f"/servicePrincipals/{service_principal_id}/appRoleAssignedTo"
url = base_url + endpoint
next_link: str | None = base_url + endpoint
headers = {
"Authorization": f"Bearer {access_token}",
"Content-Type": "application/json",
}
response = await async_client.get(url, headers=headers)
response_json = response.json()
verbose_proxy_logger.debug(f"Response from service principal app role assigned to: {response_json}")
group_ids: List[str] = []
service_principal_teams: List[MicrosoftServicePrincipalTeam] = []
page_count = 0
for _object in response_json.get("value", []):
if _object.get("principalType") == "Group":
# Append the group ID to the list
group_ids.append(_object.get("principalId"))
# Append the service principal team to the list
service_principal_teams.append(
MicrosoftServicePrincipalTeam(
principalDisplayName=_object.get("principalDisplayName"),
principalId=_object.get("principalId"),
while next_link is not None and page_count < MicrosoftSSOHandler.MAX_GRAPH_API_PAGES:
response = await async_client.get(next_link, headers=headers)
response_json = response.json()
verbose_proxy_logger.debug(f"Response from service principal app role assigned to: {response_json}")
for _object in response_json.get("value", []):
if _object.get("principalType") == "Group":
# Append the group ID to the list
group_ids.append(_object.get("principalId"))
# Append the service principal team to the list
service_principal_teams.append(
MicrosoftServicePrincipalTeam(
principalDisplayName=_object.get("principalDisplayName"),
principalId=_object.get("principalId"),
)
)
)
next_link = response_json.get("@odata.nextLink")
page_count += 1
if next_link is not None and page_count >= MicrosoftSSOHandler.MAX_GRAPH_API_PAGES:
verbose_proxy_logger.warning(
f"Reached maximum page limit of {MicrosoftSSOHandler.MAX_GRAPH_API_PAGES}. Some service principal group assignments may not be included."
)
return group_ids, service_principal_teams

View file

@ -7882,7 +7882,7 @@ class ProxyStartupEvent:
proxy_logging_obj=proxy_logging_obj,
prisma_client=prisma_client,
llm_router=llm_router,
track_unmanaged_vertex_batch_cost=general_settings.get("track_unmanaged_vertex_batch_cost", False),
track_unmanaged_batch_cost=general_settings.get("track_unmanaged_batch_cost", False),
)
scheduler.add_job(
check_batch_cost_job.check_batch_cost,
@ -14805,6 +14805,7 @@ async def get_config_list(
"forward_client_headers_to_llm_api": {"type": "Boolean"},
"mcp_required_fields": {"type": "List"},
"cancel_on_disconnect": {"type": "Boolean"},
"skip_user_budget_on_team_key": {"type": "Boolean"},
}
return_val = []

View file

@ -422,6 +422,7 @@ model LiteLLM_VerificationToken {
budget_reset_at DateTime?
allowed_cache_controls String[] @default([])
allowed_routes String[] @default([])
key_type String?
policies String[] @default([])
access_group_ids String[] @default([])
model_spend Json @default("{}")
@ -516,6 +517,7 @@ model LiteLLM_DeletedVerificationToken {
budget_reset_at DateTime?
allowed_cache_controls String[] @default([])
allowed_routes String[] @default([])
key_type String?
policies String[] @default([])
access_group_ids String[] @default([])
model_spend Json @default("{}")

View file

@ -123,6 +123,7 @@ async def reserve_budget_for_request(
proxy_logging_obj: ProxyLogging,
end_user_id: Optional[str] = None,
end_user_object: Optional[Any] = None,
skip_user_budget_on_team_key: bool = False,
) -> Optional[dict]:
if valid_token is None or not RouteChecks.is_llm_api_route(route=route):
return None
@ -141,6 +142,7 @@ async def reserve_budget_for_request(
proxy_logging_obj=proxy_logging_obj,
end_user_id=end_user_id,
end_user_object=end_user_object,
skip_user_budget_on_team_key=skip_user_budget_on_team_key,
)
if not counters:
return None
@ -296,6 +298,7 @@ async def _get_budget_counters(
proxy_logging_obj: ProxyLogging,
end_user_id: Optional[str] = None,
end_user_object: Optional[Any] = None,
skip_user_budget_on_team_key: bool = False,
) -> List[_BudgetCounter]:
counters: List[_BudgetCounter] = []
@ -344,8 +347,9 @@ async def _get_budget_counters(
)
)
is_team_key = team_object is not None and team_object.team_id is not None
if (
(team_object is None or team_object.team_id is None)
not (is_team_key and skip_user_budget_on_team_key)
and user_object is not None
and user_object.user_id is not None
and user_object.max_budget is not None

View file

@ -3592,7 +3592,10 @@ class PrismaClient:
"""
start_time = time.time()
try:
verbose_proxy_logger.debug("PrismaClient: insert_data: %s", data)
verbose_proxy_logger.debug(
"PrismaClient: insert_data: %s",
{**data, "token": self.hash_token(token=data["token"])} if data.get("token") is not None else data,
)
if table_name == "key":
token = data["token"]
hashed_token = self.hash_token(token=token)

View file

@ -261,7 +261,7 @@ async def aresponses_api_with_mcp(
pre_processed_mcp_tools=original_mcp_tools,
)
return LiteLLM_Proxy_MCP_Handler._create_mcp_streaming_response(
mcp_streaming_response = LiteLLM_Proxy_MCP_Handler._create_mcp_streaming_response(
input=input,
model=model,
all_tools=all_tools,
@ -272,6 +272,10 @@ async def aresponses_api_with_mcp(
tool_server_map=tool_server_map,
**kwargs,
)
await mcp_streaming_response._create_initial_response_iterator()
if mcp_streaming_response._initial_creation_error is not None:
raise mcp_streaming_response._initial_creation_error
return mcp_streaming_response
# Determine if we should auto-execute tools
should_auto_execute = bool(mcp_tools_with_litellm_proxy) and LiteLLM_Proxy_MCP_Handler._should_auto_execute_tools(

View file

@ -5,6 +5,8 @@ from litellm._uuid import uuid
from litellm.responses.streaming_iterator import BaseResponsesAPIStreamingIterator
from litellm.types.llms.openai import (
BaseLiteLLMOpenAIResponseObject,
ErrorEvent,
ErrorEventError,
MCPCallArgumentsDeltaEvent,
MCPCallArgumentsDoneEvent,
MCPCallCompletedEvent,
@ -316,6 +318,11 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator):
# Cache the response ID to ensure consistency across all events
self._cached_response_id: Optional[str] = None
self._initial_creation_error: Exception | None = None
self._stream_error: Exception | None = None
self._error_event_emitted = False
self._last_sequence_number = 0
def _extract_mcp_headers_from_params(self) -> None:
"""Extract MCP headers from original request params to pass to tool calls"""
from typing import Dict, Optional
@ -380,10 +387,31 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator):
return LiteLLM_Proxy_MCP_Handler._should_auto_execute_tools(self.mcp_tools_with_litellm_proxy)
def _make_stream_error_event(self) -> ResponsesAPIStreamingResponse:
err = self._stream_error
status_code = getattr(err, "status_code", None)
return ErrorEvent(
type=ResponsesAPIStreamEvents.ERROR,
sequence_number=self._last_sequence_number + 1,
error=ErrorEventError(
type="mcp_gateway_error",
code=str(status_code) if status_code is not None else "internal_error",
message=str(err) if err is not None else "MCP gateway stream failed",
param=None,
),
)
def __aiter__(self):
return self
async def __anext__(self) -> ResponsesAPIStreamingResponse:
chunk = await self._anext_impl()
sequence_number = getattr(chunk, "sequence_number", None)
if isinstance(sequence_number, int) and sequence_number > self._last_sequence_number:
self._last_sequence_number = sequence_number
return chunk
async def _anext_impl(self) -> ResponsesAPIStreamingResponse:
"""
Phase-based streaming:
1. initial_response - Stream the first LLM response (includes response.created, response.in_progress, response.output_item.added)
@ -438,10 +466,16 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator):
self.phase = "continue_initial_response"
return await self.__anext__()
self.phase = "finished"
if self._stream_error is not None and not self._error_event_emitted:
self._error_event_emitted = True
return self._make_stream_error_event()
raise StopAsyncIteration
# Phase 6: Finished
if self.phase == "finished":
if self._stream_error is not None and not self._error_event_emitted:
self._error_event_emitted = True
return self._make_stream_error_event()
raise StopAsyncIteration
# Should not reach here
@ -530,6 +564,12 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator):
chunk = await cast(Any, self.base_iterator).__anext__() # type: ignore[attr-defined]
if self._cached_response_id is None and hasattr(chunk, "response"):
new_response = getattr(chunk, "response", None)
new_response_id = getattr(new_response, "id", None) if new_response is not None else None
if new_response_id:
self._cached_response_id = new_response_id
# Ensure response ID consistency - update chunk if needed
if self._cached_response_id and hasattr(chunk, "response"):
response_obj = getattr(chunk, "response", None)
@ -589,6 +629,8 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator):
traceback.print_exc()
self.base_iterator = None
self._initial_creation_error = e
self._stream_error = e
# Don't set phase to "finished" here — let __anext__ emit any
# pre-generated MCP discovery events before ending the iteration.
@ -761,6 +803,7 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator):
if hasattr(follow_up_response, "__aiter__"):
self.base_iterator = follow_up_response
self.collected_response = None
self._cached_response_id = None
except Exception as e:
verbose_logger.error(f"Error creating follow-up iterator: {e}")
@ -768,6 +811,7 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator):
traceback.print_exc()
self.base_iterator = None
self._stream_error = e
def __iter__(self):
return self

View file

@ -179,7 +179,9 @@ from litellm.types.router import (
RouterModelGroupAliasItem,
RouterRateLimitError,
RouterRateLimitErrorBasic,
RoutingContext,
RoutingGroup,
RoutingPlugin,
RoutingStrategy,
SearchToolTypedDict,
)
@ -299,6 +301,7 @@ class Router:
enable_pre_call_checks: bool = False,
enable_tag_filtering: bool = False,
tag_filtering_match_any: bool = True,
plugins: list[RoutingPlugin] | None = None,
retry_after: int = 0, # min time to wait before retrying a failed request
retry_policy: Optional[Union[RetryPolicy, dict]] = None, # set custom retries for different exceptions
model_group_retry_policy: Dict[str, RetryPolicy] = {}, # set custom retry policies based on model group
@ -477,6 +480,7 @@ class Router:
self.complexity_routers: Dict[str, "ComplexityRouter"] = {}
self.adaptive_routers: Dict[str, "AdaptiveRouter"] = {}
self.quality_routers: Dict[str, "QualityRouter"] = {}
self.routing_plugins: list[RoutingPlugin] = list(plugins) if plugins else []
# Initialize model_group_alias early since it's used in set_model_list
self.model_group_alias: Dict[str, Union[str, RouterModelGroupAliasItem]] = (
@ -7552,7 +7556,11 @@ class Router:
if default_model is None and complexity_router_config:
tiers = complexity_router_config.get("tiers", {})
# Use MEDIUM tier as fallback default
default_model = tiers.get("MEDIUM") or tiers.get("SIMPLE")
medium = tiers.get("MEDIUM") or tiers.get("SIMPLE")
if isinstance(medium, list):
default_model = medium[0] if medium else None
else:
default_model = medium
if default_model is None:
raise ValueError(
@ -7589,15 +7597,6 @@ class Router:
AdaptiveRouterPostCallHook,
)
for _cb_list in (
litellm.callbacks,
litellm.success_callback,
litellm.failure_callback,
litellm._async_success_callback,
litellm._async_failure_callback,
):
litellm.logging_callback_manager.remove_callbacks_by_type(_cb_list, AdaptiveRouterPostCallHook)
for entry in self.model_list or []:
lp = entry.get("litellm_params") if isinstance(entry, dict) else entry.litellm_params
lp_model = (lp.get("model") if isinstance(lp, dict) else lp.model) if lp else None
@ -7606,15 +7605,29 @@ class Router:
model_name = entry.get("model_name") if isinstance(entry, dict) else entry.model_name
if not model_name or not lp:
continue
if model_name in self.adaptive_routers:
continue
deployment = Deployment(
model_name=model_name,
litellm_params=(lp if not isinstance(lp, dict) else LiteLLM_Params(**lp)),
model_info=(entry.get("model_info") if isinstance(entry, dict) else entry.model_info),
)
if model_name in self.adaptive_routers:
continue
self.init_adaptive_router_deployment(deployment=deployment)
for model_name, complexity_router in self.complexity_routers.items():
if not complexity_router.config.adaptive or model_name in self.adaptive_routers:
continue
adaptive_router = complexity_router._ensure_adaptive_router()
if adaptive_router is not None:
self.adaptive_routers[model_name] = adaptive_router
for callback in litellm.logging_callback_manager.get_custom_loggers_for_type(AdaptiveRouterPostCallHook):
litellm.logging_callback_manager.remove_callback_from_all_lists(callback)
for adaptive_router in self.adaptive_routers.values():
litellm.logging_callback_manager.add_litellm_callback(
AdaptiveRouterPostCallHook(adaptive_router=adaptive_router)
)
def init_adaptive_router_deployment(self, deployment: Deployment) -> None:
"""
Build an AdaptiveRouter instance for this deployment and register its
@ -10321,6 +10334,12 @@ class Router:
metadata_variable_name=self._get_metadata_variable_name_from_kwargs(request_kwargs),
)
# narrow to whatever `self.routing_plugins` left in candidate_models
healthy_deployments = self._filter_by_routing_plugin_candidates(
healthy_deployments=healthy_deployments,
request_kwargs=request_kwargs,
)
## ORDER FILTERING ## -> if user set 'order' in deployments, return deployments with lowest order (e.g. order=1 > order=2)
_target_order = (request_kwargs or {}).pop("_target_order", None)
healthy_deployments = litellm.utils._get_order_filtered_deployments(
@ -10596,6 +10615,76 @@ class Router:
)
raise e
async def _run_routing_plugins(
self,
model: str,
request_kwargs: dict,
messages: list[dict[str, Any]] | None,
) -> RoutingContext:
"""
Build a RoutingContext for `model`, run it through `self.routing_plugins`
in order, then stash the narrowed candidate list and accumulated signals
onto `request_kwargs["metadata"]` so `_filter_by_routing_plugin_candidates`
(called later, during healthy-deployment filtering) can consume them.
"""
from litellm.litellm_core_utils.prompt_templates.factory import (
resolve_structured_messages,
)
deployments = self.get_model_list(model_name=model) or []
candidate_models = [
d["litellm_params"]["model"] for d in deployments if d.get("litellm_params", {}).get("model")
]
metadata_key = self._get_metadata_variable_name_from_kwargs(request_kwargs)
metadata = request_kwargs.setdefault(metadata_key, {})
context = RoutingContext(
raw_messages=messages or [],
structured_messages=resolve_structured_messages(messages=messages, request_kwargs=request_kwargs) or [],
candidate_models=candidate_models,
metadata=metadata,
)
for plugin in self.routing_plugins:
context = await plugin.run(context)
metadata["routing_plugin_signals"] = context.signals
if len(context.candidate_models) < len(candidate_models):
metadata["_routing_plugin_candidate_models"] = context.candidate_models
return context
def _filter_by_routing_plugin_candidates(
self,
healthy_deployments: Union[list[dict], dict],
request_kwargs: dict,
) -> Union[list[dict], dict]:
"""
Narrow `healthy_deployments` to whatever `self.routing_plugins` left in
`context.candidate_models`. Raises rather than silently falling back to
the unfiltered pool -- a plugin narrowing to nothing is a policy decision
(e.g. no model this tenant's budget allows), not something to bypass.
"""
if not self.routing_plugins or not isinstance(healthy_deployments, list):
return healthy_deployments
metadata_key = self._get_metadata_variable_name_from_kwargs(request_kwargs)
candidate_models = (request_kwargs.get(metadata_key) or {}).get("_routing_plugin_candidate_models")
# `is None` (not falsy-check): a plugin narrowing to an empty list must
# still hit the "no deployments left" raise below, not be treated the
# same as "no plugin ever set this key".
if candidate_models is None:
return healthy_deployments
candidate_set = set(candidate_models)
filtered = [d for d in healthy_deployments if d.get("litellm_params", {}).get("model") in candidate_set]
if not filtered:
raise ValueError(f"No deployments left after routing-plugin filtering. candidate_models={candidate_models}")
return filtered
async def async_pre_routing_hook(
self,
model: str,
@ -10610,55 +10699,47 @@ class Router:
Used for the litellm auto-router to modify the request before the routing decision is made.
"""
#########################################################
# Check if any auto-router should be used
# Run the routing-plugin pipeline, if any plugins are configured.
# Plugins narrow the candidate deployment pool (consumed later by
# `_filter_by_routing_plugin_candidates`) and may attach signals for
# downstream strategies (auto-router, complexity-router, ...) to read.
#########################################################
if model in self.auto_routers:
return await self.auto_routers[model].async_pre_routing_hook(
model=model,
request_kwargs=request_kwargs,
messages=messages,
input=input,
specific_deployment=specific_deployment,
)
if self.routing_plugins:
await self._run_routing_plugins(model=model, request_kwargs=request_kwargs, messages=messages)
#########################################################
# Check if any complexity-router should be used
#########################################################
if model in self.complexity_routers:
return await self.complexity_routers[model].async_pre_routing_hook(
model=model,
request_kwargs=request_kwargs,
messages=messages,
input=input,
specific_deployment=specific_deployment,
)
router_strategy = (
self.auto_routers.get(model)
or self.complexity_routers.get(model)
or self.adaptive_routers.get(model)
or self.quality_routers.get(model)
)
if router_strategy is None:
return None
#########################################################
# Check if an adaptive-router should be used
#########################################################
adaptive_router = self.adaptive_routers.get(model)
if adaptive_router is not None:
return await adaptive_router.async_pre_routing_hook(
model=model,
request_kwargs=request_kwargs,
messages=messages,
input=input,
specific_deployment=specific_deployment,
)
pre_routing_hook_response = await router_strategy.async_pre_routing_hook(
model=model,
request_kwargs=request_kwargs,
messages=messages,
input=input,
specific_deployment=specific_deployment,
)
#########################################################
# Check if any quality-router should be used
#########################################################
if model in self.quality_routers:
return await self.quality_routers[model].async_pre_routing_hook(
model=model,
request_kwargs=request_kwargs,
messages=messages,
input=input,
specific_deployment=specific_deployment,
)
# `model` (the alias, e.g. "smart-router") is never the deployment actually
# called - apply the alias's own litellm_params (besides `model` itself,
# which is just the alias marker) to the request, since the tier/route
# deployment the hook selected won't have them. Router-only fields
# (tpm, rpm, weight, complexity_router_config, ...) are excluded from the
# actual outbound LLM call downstream by litellm.types.utils.all_litellm_params,
# not here.
if pre_routing_hook_response is not None:
alias_index = self.model_name_to_deployment_indices.get(model, [])
if alias_index:
alias_litellm_params = self.model_list[alias_index[0]].get("litellm_params", {})
for key, value in alias_litellm_params.items():
if key != "model" and value is not None:
request_kwargs.setdefault(key, value)
return None
return pre_routing_hook_response
def get_available_deployment(
self,
@ -10671,6 +10752,18 @@ class Router:
"""
Returns the deployment based on routing strategy
"""
if self.routing_plugins:
raise ValueError(
"Router(plugins=[...]) is configured, but this call resolved to the synchronous "
"deployment-selection path, which never runs the routing-plugin pipeline. This "
"happens for sync Router methods (e.g. Router.completion()) and for async calls "
"with a routing_strategy that has no async-native selector (e.g. legacy "
"'usage-based-routing', v1). Silently skipping "
"configured plugins would let a policy plugin (e.g. a deny-all rule) be bypassed. "
"Use an async Router method with a supported routing_strategy (simple-shuffle, "
"usage-based-routing-v2, cost-based-routing, latency-based-routing, least-busy), "
"or remove `plugins` from the Router config."
)
# users need to explicitly call a specific deployment, by setting `specific_deployment = True` as completion()/embedding() kwarg
# When this was no explicit we had several issues with fallbacks timing out

View file

@ -56,11 +56,10 @@ Callers may pass header `x-litellm-min-quality-tier: 3` (or metadata key
- **Per-request decision.** Sample once per eligible model, score with
`quality_weight·sample + cost_weight·normalized_cost`, pick the argmax.
Routing is stateless per-turn — no sticky lookup. Each call resamples.
- **Owner-cache attribution.** Post-call, the conversation's first picked
model claims an "owner slot" for `OWNER_CACHE_TTL_SECONDS` (24h). Later
turns of the same conversation only fire bandit/state updates if the
same model handled them — mismatches are dropped (no attribution) and
counted in `skipped_updates_total`. Conversation identity is the
- **Previous-response attribution.** Post-call, feedback from the current user
message is attributed to the model that produced the previous response, while
response signals are attributed to the current model. Contexts expire after
24 hours and the in-memory cache is capped at 1,024 sessions. Conversation identity is the
client-supplied `litellm_session_id` if present, otherwise a sha256 over
caller identity (api key hash, team, user, end-user) + the first message.
- **Per-turn updates.** `satisfaction → +α`. `misalignment, stagnation,
@ -76,12 +75,6 @@ Callers may pass header `x-litellm-min-quality-tier: 3` (or metadata key
model can still be picked.
- **Hard sample cap at 200.** Once `α + β > 200`, deltas are silently dropped.
No rescaling — drift is a v1 concern.
- **24h owner-cache TTL.** No explicit eviction below TTL. The in-memory map
can grow if traffic patterns produce many one-shot sessions.
- **Owner-recovery skew.** If model A "owns" a conversation but is then
dethroned in the bandit, later turns served by model B are dropped — so
bandit updates for that conversation flatline until A's TTL expires.
Tracked via `skipped_updates_total`.
- **Signals are regex + tool-call only.** No LLM-judge, no embedding similarity,
no exemplar storage. Signals are best-effort and biased toward English.
- **One AdaptiveRouter per `Router`.** Multiple `adaptive_router/*` deployments

View file

@ -3,25 +3,21 @@ Main adaptive router strategy. See README.md for design overview.
One AdaptiveRouter instance per router_name. Holds in-memory caches:
- _cells: Beta(alpha, beta) bandit posteriors per (request_type, model)
- _owner_cache: session_key -> (owner_model, expires_at) the first model
picked for a conversation owns its bandit-update slot
- _session_states: (session_key, model) -> SessionState for incremental signal updates
Owns the AdaptiveRouterUpdateQueue used by the proxy's flusher to persist
state and session snapshots back to Postgres.
Routing is stateless per-turn (Thompson sample fresh on every call). The
owner cache is consulted only at post-call time to decide whether a turn's
signals should fire a bandit update turns served by a different model than
the conversation's owner are skipped to avoid cross-model misattribution.
Routing is stateless per-turn (Thompson sample fresh on every call).
"""
from __future__ import annotations
import asyncio
import time
from dataclasses import asdict
from typing import Any, Dict, List, Optional, Tuple, Union, cast
from collections import OrderedDict
from dataclasses import asdict, dataclass
from typing import Any, Union, cast
from litellm._logging import verbose_router_logger
from litellm.litellm_core_utils.prompt_templates.common_utils import (
@ -38,13 +34,18 @@ from litellm.router_strategy.adaptive_router.config import (
ADAPTIVE_ROUTER_CHOSEN_MODEL_KEY,
MIN_QUALITY_TIER_HEADER,
MIN_QUALITY_TIER_METADATA_KEY,
MIN_TURNS_FOR_CLEAN_CREDIT,
OWNER_CACHE_TTL_SECONDS,
)
from litellm.router_strategy.adaptive_router.signals import (
SessionState,
SignalDelta,
Turn,
apply_turn,
advance_session_state,
apply_signal_delta,
detect_response_signals,
detect_user_feedback,
merge_signal_deltas,
)
from litellm.router_strategy.adaptive_router.update_queue import (
AdaptiveRouterUpdateQueue,
@ -53,8 +54,7 @@ from litellm.router_strategy.adaptive_router.update_queue import (
# Sweep session-state cache when it exceeds this many live entries. Expired
# entries are dropped in bulk; amortizes to O(1) per insert.
_SESSION_STATE_SWEEP_THRESHOLD: int = 1024
# Same pattern for the owner cache.
_OWNER_CACHE_SWEEP_THRESHOLD: int = 1024
_FEEDBACK_CONTEXT_MAX_ENTRIES: int = 1024
from litellm.repositories.table_repositories import AdaptiveRouterStateRepository
from litellm.types.llms.openai import AllMessageValues
from litellm.types.router import (
@ -70,6 +70,17 @@ def _default_prefs() -> AdaptiveRouterPreferences:
return AdaptiveRouterPreferences(quality_tier=2, strengths=[])
@dataclass(frozen=True, slots=True)
class _FeedbackContext:
model_name: str
request_type: RequestType
user_content: str | None
assistant_content: str | None
turn_count: int
clean_credit_awarded: bool
expires_at: float
class AdaptiveRouter:
"""One instance per router_name. Holds in-memory caches + the update queue."""
@ -77,8 +88,8 @@ class AdaptiveRouter:
self,
router_name: str,
config: AdaptiveRouterConfig,
model_to_prefs: Dict[str, AdaptiveRouterPreferences],
model_to_cost: Dict[str, float],
model_to_prefs: dict[str, AdaptiveRouterPreferences],
model_to_cost: dict[str, float],
) -> None:
self.router_name = router_name
self.config = config
@ -86,13 +97,14 @@ class AdaptiveRouter:
self.model_to_cost = model_to_cost
self.queue = AdaptiveRouterUpdateQueue()
self._cells: Dict[Tuple[RequestType, str], BanditCell] = {}
self._owner_cache: Dict[str, Tuple[str, float]] = {}
self._session_states: Dict[Tuple[str, str], SessionState] = {}
# Parallel expiry map for _session_states, same TTL as _owner_cache.
# Evicted opportunistically in `get_or_create_session_state`.
self._session_states_expiry: Dict[Tuple[str, str], float] = {}
self._skipped_updates_total: int = 0
self._cells: dict[tuple[RequestType, str], BanditCell] = {}
self._session_states: dict[tuple[str, str], SessionState] = {}
self._feedback_contexts: OrderedDict[str, _FeedbackContext] = OrderedDict()
self._session_states_expiry: dict[tuple[str, str], float] = {}
self._feedback_attributed_total: int = 0
self._feedback_without_context_total: int = 0
self._cross_model_feedback_total: int = 0
self._response_signal_updates_total: int = 0
# Set to True once the proxy flusher has loaded persisted priors from
# Postgres. Checked to support lazy-load on hot-reloaded routers.
self._state_loaded: bool = False
@ -145,11 +157,11 @@ class AdaptiveRouter:
async def async_pre_routing_hook(
self,
model: str,
request_kwargs: Dict[str, Any],
messages: Optional[List[Dict[str, Any]]] = None,
input: Optional[Union[str, List]] = None,
specific_deployment: Optional[bool] = False,
) -> Optional[PreRoutingHookResponse]:
request_kwargs: dict[str, Any],
messages: list[dict[str, Any]] | None = None,
input: Union[str, list] | None = None,
specific_deployment: bool | None = False,
) -> PreRoutingHookResponse | None:
"""
Plugin entry point invoked by `Router.async_pre_routing_hook` when the
inbound `model` matches this adaptive router's `router_name`.
@ -159,11 +171,9 @@ class AdaptiveRouter:
post-call hook can surface it as a response header.
Routing is stateless per-turn: every call Thompson-samples fresh,
regardless of any prior pick for the same session. Cross-turn
attribution is enforced post-call via the owner cache (see
`claim_or_check_owner`).
regardless of any prior pick for the same session.
"""
user_text = get_last_user_message(cast(List[AllMessageValues], messages or [])) or ""
user_text = get_last_user_message(cast(list[AllMessageValues], messages or [])) or ""
request_type = classify_prompt(user_text)
min_quality_tier = self._extract_min_quality_tier(request_kwargs)
@ -190,7 +200,7 @@ class AdaptiveRouter:
async def pick_model(
self,
request_type: RequestType,
min_quality_tier: Optional[int] = None,
min_quality_tier: int | None = None,
) -> str:
"""Thompson-sample across eligible models. Stateless per-turn."""
eligible = self._eligible_models(min_quality_tier)
@ -206,44 +216,7 @@ class AdaptiveRouter:
cost_weight=self.config.weights.cost,
)
def claim_or_check_owner(self, session_key: str, current_model: str) -> bool:
"""Resolve attribution for a turn under stateless routing.
Returns True iff this turn should fire a bandit/state update. The
first call for a `session_key` claims ownership for `current_model`
and returns True. Subsequent calls return True only if the owner is
still live AND matches `current_model`. Mismatches (a different
model handled this turn) and expired owners both increment
`_skipped_updates_total` and return False no attribution.
"""
now = time.time()
existing = self._owner_cache.get(session_key)
if existing is not None and existing[1] > now:
owner_model, _ = existing
if owner_model == current_model:
return True
self._skipped_updates_total += 1
return False
# Opportunistic bulk sweep — sessions that never come back would
# otherwise pile up here forever. Same threshold pattern as the
# session-state cache.
if len(self._owner_cache) >= _OWNER_CACHE_SWEEP_THRESHOLD:
self._evict_expired_owner_cache(now)
# No live owner -> claim for current_model.
self._owner_cache[session_key] = (
current_model,
now + OWNER_CACHE_TTL_SECONDS,
)
return True
def _evict_expired_owner_cache(self, now: float) -> None:
expired = [k for k, (_, exp) in self._owner_cache.items() if exp <= now]
for k in expired:
self._owner_cache.pop(k, None)
async def get_state_snapshot(self) -> Dict[str, Any]:
async def get_state_snapshot(self) -> dict[str, Any]:
"""In-memory snapshot for the introspection endpoint. Cheap; no DB hit."""
cells = []
for (rt, model), cell in sorted(self._cells.items(), key=lambda kv: (kv[0][0].value, kv[0][1])):
@ -264,7 +237,7 @@ class AdaptiveRouter:
)
queue = await self.queue.queue_size()
now = time.time()
owner_cache_live = sum(1 for _, exp in self._owner_cache.values() if exp > now)
feedback_contexts_live = sum(1 for context in self._feedback_contexts.values() if context.expires_at > now)
return {
"router_name": self.router_name,
"available_models": list(self.config.available_models),
@ -274,15 +247,18 @@ class AdaptiveRouter:
},
"model_costs": dict(self.model_to_cost),
"cells": cells,
"owner_cache_live": owner_cache_live,
"skipped_updates_total": self._skipped_updates_total,
"feedback_contexts_live": feedback_contexts_live,
"feedback_attributed_total": self._feedback_attributed_total,
"feedback_without_context_total": self._feedback_without_context_total,
"cross_model_feedback_total": self._cross_model_feedback_total,
"response_signal_updates_total": self._response_signal_updates_total,
"queue": queue,
}
@staticmethod
def _extract_min_quality_tier(
request_kwargs: Dict[str, Any],
) -> Optional[int]:
request_kwargs: dict[str, Any],
) -> int | None:
"""Pull `min_quality_tier` from request headers or metadata.
Precedence: headers (`x-litellm-min-quality-tier`) over metadata
@ -310,7 +286,7 @@ class AdaptiveRouter:
return None
return None
def _eligible_models(self, min_quality_tier: Optional[int]) -> List[str]:
def _eligible_models(self, min_quality_tier: int | None) -> list[str]:
if min_quality_tier is None:
return list(self.config.available_models)
return [
@ -363,17 +339,131 @@ class AdaptiveRouter:
request_type: RequestType,
turn: Turn,
) -> SignalDelta:
"""Apply one turn, push session snapshot + bandit deltas to the queue."""
state = self.get_or_create_session_state(session_id, model_name, request_type)
delta = apply_turn(state, turn)
verbose_router_logger.debug("AdaptiveRouter[%s]: record_turn delta=%s", self.router_name, delta)
"""Attribute feedback to the previous response and response signals to the current model."""
async with self._lock:
now = time.time()
while self._feedback_contexts:
oldest_context = next(iter(self._feedback_contexts.values()))
if oldest_context.expires_at > now:
break
self._feedback_contexts.popitem(last=False)
previous = self._feedback_contexts.pop(session_id, None)
# Strip the raw conversation content before persisting. The
# last_user/assistant_content and tool_call_history fields are only
# needed in-memory for the next turn's incremental signal detection;
# writing user prompts and tool payloads to the DB would store PII
# for every adaptive-router conversation. Counts + bookkeeping is
# all the persisted row needs.
effective_request_type = (
previous.request_type if previous is not None and request_type == RequestType.GENERAL else request_type
)
current_state = self.get_or_create_session_state(
session_id,
model_name,
effective_request_type,
)
feedback_delta = detect_user_feedback(
previous.user_content if previous else None,
turn.user_content,
turn.tool_results,
allow_satisfaction=(
previous is not None
and not previous.clean_credit_awarded
and previous.turn_count + 1 >= MIN_TURNS_FOR_CLEAN_CREDIT
),
)
previous_assistant = previous.assistant_content if previous else None
response_delta = detect_response_signals(
previous_assistant,
turn.assistant_content,
current_state.tool_call_history,
turn.tool_calls,
turn.tool_results,
turn.response_status,
)
states_to_persist: dict[str, SessionState] = {model_name: current_state}
bandit_deltas: dict[tuple[RequestType, str], SignalDelta] = {}
if previous is not None:
feedback_state = self.get_or_create_session_state(
session_id,
previous.model_name,
previous.request_type,
)
apply_signal_delta(feedback_state, feedback_delta)
if feedback_delta.satisfaction:
feedback_state.clean_credit_awarded = True
states_to_persist[previous.model_name] = feedback_state
if feedback_delta.any_fired():
self._feedback_attributed_total += 1
if previous.model_name != model_name:
self._cross_model_feedback_total += 1
bandit_deltas[(previous.request_type, previous.model_name)] = feedback_delta
else:
if feedback_delta.any_fired():
self._feedback_without_context_total += 1
initial_failure = SignalDelta(failure=feedback_delta.failure)
apply_signal_delta(current_state, initial_failure)
bandit_deltas[(effective_request_type, model_name)] = initial_failure
apply_signal_delta(current_state, response_delta)
if self._compute_bandit_delta(response_delta) != (0.0, 0.0):
self._response_signal_updates_total += 1
current_key = (effective_request_type, model_name)
bandit_deltas[current_key] = merge_signal_deltas(
bandit_deltas.get(current_key, SignalDelta()),
response_delta,
)
advance_session_state(current_state, turn)
next_turn_count = (previous.turn_count if previous else 0) + 1
clean_credit_awarded = bool((previous and previous.clean_credit_awarded) or feedback_delta.satisfaction)
if len(self._feedback_contexts) >= _FEEDBACK_CONTEXT_MAX_ENTRIES:
self._feedback_contexts.popitem(last=False)
self._feedback_contexts[session_id] = _FeedbackContext(
model_name=model_name,
request_type=effective_request_type,
user_content=turn.user_content,
assistant_content=turn.assistant_content,
turn_count=next_turn_count,
clean_credit_awarded=clean_credit_awarded,
expires_at=now + OWNER_CACHE_TTL_SECONDS,
)
for state_model, state in states_to_persist.items():
await self.queue.add_session_state(
session_id,
self.router_name,
state_model,
self._persistable_session_snapshot(state),
)
combined_delta = SignalDelta()
for (attribution_type, target_model), delta in bandit_deltas.items():
combined_delta = merge_signal_deltas(combined_delta, delta)
d_alpha, d_beta = self._compute_bandit_delta(delta)
if d_alpha == 0 and d_beta == 0:
continue
cell_key = (attribution_type, target_model)
self._cells[cell_key] = apply_delta(
self._cells[cell_key],
d_alpha,
d_beta,
)
await self.queue.add_state_delta(
self.router_name,
attribution_type.value,
target_model,
d_alpha,
d_beta,
)
verbose_router_logger.debug(
"AdaptiveRouter[%s]: feedback_target=%s current_model=%s delta=%s",
self.router_name,
previous.model_name if previous else None,
model_name,
combined_delta,
)
return combined_delta
@staticmethod
def _persistable_session_snapshot(state: SessionState) -> dict[str, Any]:
snapshot = asdict(state)
for sensitive in (
"last_user_content",
@ -382,38 +472,10 @@ class AdaptiveRouter:
"pending_tool_calls",
):
snapshot.pop(sensitive, None)
await self.queue.add_session_state(session_id, self.router_name, model_name, snapshot)
d_alpha, d_beta = self._compute_bandit_delta(delta)
verbose_router_logger.debug(
"AdaptiveRouter[%s]: bandit delta alpha=%.2f beta=%.2f",
self.router_name,
d_alpha,
d_beta,
)
if d_alpha != 0 or d_beta != 0:
# For non-GENERAL turns, attribute to the current-turn classification
# so genuine mid-session topic shifts (e.g. code → math) update the
# correct cell. For GENERAL turns ("thanks!", "ok", "sounds good"), fall
# back to the session's original type so closing pleasantries don't
# misattribute the reward.
attribution_type = (
request_type if request_type != RequestType.GENERAL else RequestType(state.classified_type)
)
cell_key = (attribution_type, model_name)
self._cells[cell_key] = apply_delta(self._cells[cell_key], d_alpha, d_beta)
await self.queue.add_state_delta(
self.router_name,
attribution_type.value,
model_name,
d_alpha,
d_beta,
)
return delta
return snapshot
@staticmethod
def _compute_bandit_delta(delta: SignalDelta) -> Tuple[float, float]:
def _compute_bandit_delta(delta: SignalDelta) -> tuple[float, float]:
"""
Translate per-turn signal deltas into bandit-cell deltas.

View file

@ -214,10 +214,6 @@ class AdaptiveRouterPostCallHook(CustomLogger):
) -> None:
try:
messages = kwargs.get("messages") or []
if len(messages) < SIGNAL_GATE_MIN_MESSAGES:
# Too few turns for any signal to be meaningful — skip.
return
session_key = _resolve_session_key(kwargs)
if not session_key:
return
@ -233,10 +229,6 @@ class AdaptiveRouterPostCallHook(CustomLogger):
if not current_model:
return
if not self.adaptive_router.claim_or_check_owner(session_key, current_model):
# A different model owns this conversation — skip attribution.
return
user_text = _last_user_content(messages)
assistant_text, tool_calls = _assistant_content_and_tool_calls(response_obj)
tool_results = _recent_tool_results(messages)

View file

@ -14,7 +14,7 @@ from __future__ import annotations
import re
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional, Set
from typing import Any
from litellm.router_strategy.adaptive_router.config import (
LOOP_REPEAT_THRESHOLD,
@ -74,26 +74,26 @@ class SessionState:
loop_count: int = 0
exhaustion_count: int = 0
last_user_content: Optional[str] = None
last_assistant_content: Optional[str] = None
tool_call_history: List[str] = field(default_factory=list)
pending_tool_calls: Dict[str, str] = field(default_factory=dict)
last_user_content: str | None = None
last_assistant_content: str | None = None
tool_call_history: list[str] = field(default_factory=list)
pending_tool_calls: dict[str, str] = field(default_factory=dict)
turn_count: int = 0
last_processed_turn: int = -1
clean_credit_awarded: bool = False
terminal_status: Optional[int] = None
terminal_status: int | None = None
@dataclass
class Turn:
"""One turn of input. Caller assembles this from the request/response."""
user_content: Optional[str] = None
assistant_content: Optional[str] = None
tool_calls: List[Dict[str, Any]] = field(default_factory=list)
tool_results: List[Dict[str, Any]] = field(default_factory=list)
response_status: Optional[int] = None
user_content: str | None = None
assistant_content: str | None = None
tool_calls: list[dict[str, Any]] = field(default_factory=list)
tool_results: list[dict[str, Any]] = field(default_factory=list)
response_status: int | None = None
# ---- Detection helpers ----------------------------------------------------
@ -101,13 +101,13 @@ class Turn:
_TOKEN_RE = re.compile(r"[A-Za-z0-9]+")
def _tokens(text: Optional[str]) -> Set[str]:
def _tokens(text: str | None) -> set[str]:
if not text:
return set()
return {t.lower() for t in _TOKEN_RE.findall(text)}
def _jaccard(a: Set[str], b: Set[str]) -> float:
def _jaccard(a: set[str], b: set[str]) -> float:
union = a | b
if not union:
return 0.0
@ -130,7 +130,7 @@ _SATISFACTION_PATTERNS = [
]
def _detect_misalignment(prev_user: Optional[str], curr_user: Optional[str]) -> bool:
def _detect_misalignment(prev_user: str | None, curr_user: str | None) -> bool:
"""Fires when consecutive user messages share *some* topic (jaccard > 0)
but are sufficiently different (jaccard < threshold) i.e. user is
rephrasing, not changing topic, not repeating."""
@ -140,7 +140,7 @@ def _detect_misalignment(prev_user: Optional[str], curr_user: Optional[str]) ->
return 0.0 < j < MISALIGNMENT_JACCARD_THRESHOLD
def _detect_stagnation(prev_asst: Optional[str], curr_asst: Optional[str]) -> bool:
def _detect_stagnation(prev_asst: str | None, curr_asst: str | None) -> bool:
"""Fires when consecutive assistant messages are near-duplicates."""
if not prev_asst or not curr_asst:
return False
@ -148,19 +148,19 @@ def _detect_stagnation(prev_asst: Optional[str], curr_asst: Optional[str]) -> bo
return j >= STAGNATION_JACCARD_NEAR_DUP
def _detect_disengagement(curr_user: Optional[str]) -> bool:
def _detect_disengagement(curr_user: str | None) -> bool:
if not curr_user:
return False
return any(p.search(curr_user) for p in _DISENGAGEMENT_PATTERNS)
def _detect_satisfaction(curr_user: Optional[str]) -> bool:
def _detect_satisfaction(curr_user: str | None) -> bool:
if not curr_user:
return False
return any(p.search(curr_user) for p in _SATISFACTION_PATTERNS)
def _detect_failure(tool_results: List[Dict[str, Any]]) -> bool:
def _detect_failure(tool_results: list[dict[str, Any]]) -> bool:
"""Any tool result explicitly flagged as an error.
We do NOT treat empty content as failure many tools legitimately return
@ -173,7 +173,7 @@ def _detect_failure(tool_results: List[Dict[str, Any]]) -> bool:
return False
def _signature(call: Dict[str, Any]) -> str:
def _signature(call: dict[str, Any]) -> str:
"""Stable signature for loop detection: name + sorted JSON-ish args."""
name = call.get("name") or call.get("function", {}).get("name", "")
call_args = call.get("arguments")
@ -184,7 +184,7 @@ def _signature(call: Dict[str, Any]) -> str:
return f"{name}({call_args})"
def _detect_loop(history: List[str], new_calls: List[Dict[str, Any]]) -> bool:
def _detect_loop(history: list[str], new_calls: list[dict[str, Any]]) -> bool:
"""Fires if any new call's signature appears >= LOOP_REPEAT_THRESHOLD-1 times
in recent history (so this call would be the Nth)."""
if not new_calls:
@ -209,7 +209,7 @@ _EXHAUSTION_KEYWORDS = (
)
def _detect_exhaustion(status: Optional[int], tool_results: List[Dict[str, Any]]) -> bool:
def _detect_exhaustion(status: int | None, tool_results: list[dict[str, Any]]) -> bool:
if status is not None and status in _EXHAUSTION_STATUSES:
return True
for r in tool_results:
@ -219,39 +219,53 @@ def _detect_exhaustion(status: Optional[int], tool_results: List[Dict[str, Any]]
return False
# ---- Public entrypoint ----------------------------------------------------
def detect_user_feedback(
previous_user_content: str | None,
current_user_content: str | None,
tool_results: list[dict[str, Any]],
allow_satisfaction: bool,
) -> SignalDelta:
return SignalDelta(
misalignment=int(_detect_misalignment(previous_user_content, current_user_content)),
disengagement=int(_detect_disengagement(current_user_content)),
satisfaction=int(allow_satisfaction and _detect_satisfaction(current_user_content)),
failure=int(_detect_failure(tool_results)),
)
def apply_turn(state: SessionState, turn: Turn) -> SignalDelta:
"""
Detect signals on this turn, mutate state, return the delta.
def detect_response_signals(
previous_assistant_content: str | None,
current_assistant_content: str | None,
tool_call_history: list[str],
tool_calls: list[dict[str, Any]],
tool_results: list[dict[str, Any]],
response_status: int | None,
) -> SignalDelta:
return SignalDelta(
stagnation=int(
_detect_stagnation(
previous_assistant_content,
current_assistant_content,
)
),
loop=int(_detect_loop(tool_call_history, tool_calls)),
exhaustion=int(_detect_exhaustion(response_status, tool_results)),
)
O(1) per turn (no full-history rescan). Only inspects last_*, recent tool history
(which is bounded at TOOL_CALL_HISTORY_MAX), and the new turn payload.
"""
delta = SignalDelta()
if _detect_misalignment(state.last_user_content, turn.user_content):
delta.misalignment = 1
if _detect_stagnation(state.last_assistant_content, turn.assistant_content):
delta.stagnation = 1
if _detect_disengagement(turn.user_content):
delta.disengagement = 1
if _detect_satisfaction(turn.user_content):
# Gate: only award satisfaction credit once per session, and only
# after MIN_TURNS_FOR_CLEAN_CREDIT turns of context. Early "thanks"
# on turn 1-2 is noise, not a validated quality signal.
current_turn_index = state.turn_count + 1
if not state.clean_credit_awarded and current_turn_index >= MIN_TURNS_FOR_CLEAN_CREDIT:
delta.satisfaction = 1
state.clean_credit_awarded = True
if _detect_failure(turn.tool_results):
delta.failure = 1
if _detect_loop(state.tool_call_history, turn.tool_calls):
delta.loop = 1
if _detect_exhaustion(turn.response_status, turn.tool_results):
delta.exhaustion = 1
def merge_signal_deltas(*deltas: SignalDelta) -> SignalDelta:
return SignalDelta(
misalignment=sum(delta.misalignment for delta in deltas),
stagnation=sum(delta.stagnation for delta in deltas),
disengagement=sum(delta.disengagement for delta in deltas),
satisfaction=sum(delta.satisfaction for delta in deltas),
failure=sum(delta.failure for delta in deltas),
loop=sum(delta.loop for delta in deltas),
exhaustion=sum(delta.exhaustion for delta in deltas),
)
def apply_signal_delta(state: SessionState, delta: SignalDelta) -> None:
state.misalignment_count += delta.misalignment
state.stagnation_count += delta.stagnation
state.disengagement_count += delta.disengagement
@ -260,6 +274,8 @@ def apply_turn(state: SessionState, turn: Turn) -> SignalDelta:
state.loop_count += delta.loop
state.exhaustion_count += delta.exhaustion
def advance_session_state(state: SessionState, turn: Turn) -> None:
if turn.user_content:
state.last_user_content = turn.user_content
if turn.assistant_content:
@ -276,4 +292,38 @@ def apply_turn(state: SessionState, turn: Turn) -> SignalDelta:
state.turn_count += 1
state.last_processed_turn = state.turn_count
# ---- Public entrypoint ----------------------------------------------------
def apply_turn(state: SessionState, turn: Turn) -> SignalDelta:
"""
Detect signals on this turn, mutate state, return the delta.
O(1) per turn (no full-history rescan). Only inspects last_*, recent tool history
(which is bounded at TOOL_CALL_HISTORY_MAX), and the new turn payload.
"""
feedback_delta = detect_user_feedback(
state.last_user_content,
turn.user_content,
turn.tool_results,
allow_satisfaction=(not state.clean_credit_awarded and state.turn_count + 1 >= MIN_TURNS_FOR_CLEAN_CREDIT),
)
response_delta = detect_response_signals(
state.last_assistant_content,
turn.assistant_content,
state.tool_call_history,
turn.tool_calls,
turn.tool_results,
turn.response_status,
)
delta = merge_signal_deltas(
feedback_delta,
response_delta,
)
apply_signal_delta(state, delta)
if delta.satisfaction:
state.clean_credit_awarded = True
advance_session_state(state, turn)
return delta

View file

@ -13,9 +13,12 @@ evaluated before either classification strategy and force a tier outright when m
Inspired by ClawRouter: https://github.com/BlockRunAI/ClawRouter
"""
from __future__ import annotations
import asyncio
import random
import re
from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Tuple, Union, cast
from typing import TYPE_CHECKING, Any, Literal, Union, cast
from pydantic import BaseModel
@ -37,6 +40,7 @@ if TYPE_CHECKING:
from semantic_router.routers import SemanticRouter
from litellm.router import Router
from litellm.router_strategy.adaptive_router.adaptive_router import AdaptiveRouter
from litellm.types.router import PreRoutingHookResponse
else:
Router = Any
@ -62,7 +66,7 @@ Tiers:
{prompt}"""
def _append_custom_keywords(base_keywords: list[str], custom_keywords: Optional[list[str]]) -> list[str]:
def _append_custom_keywords(base_keywords: list[str], custom_keywords: list[str] | None) -> list[str]:
if not custom_keywords:
return base_keywords
base_lowered = frozenset(keyword.lower() for keyword in base_keywords)
@ -94,7 +98,7 @@ def _sanitize_user_api_key_auth(auth: Any) -> Any:
return auth
def _classifier_call_metadata(metadata: Optional[dict[str, Any]]) -> Optional[dict[str, Any]]:
def _classifier_call_metadata(metadata: dict[str, Any] | None) -> dict[str, Any] | None:
if not metadata:
return metadata
return {
@ -109,7 +113,7 @@ class DimensionScore:
__slots__ = ("name", "score", "signal")
def __init__(self, name: str, score: float, signal: Optional[str] = None):
def __init__(self, name: str, score: float, signal: str | None = None):
self.name = name
self.score = score
self.signal = signal
@ -133,9 +137,9 @@ class ComplexityRouter(CustomLogger):
def __init__(
self,
model_name: str,
litellm_router_instance: "Router",
complexity_router_config: Optional[Dict[str, Any]] = None,
default_model: Optional[str] = None,
litellm_router_instance: Router,
complexity_router_config: dict[str, Any] | None = None,
default_model: str | None = None,
):
"""
Initialize ComplexityRouter.
@ -172,7 +176,7 @@ class ComplexityRouter(CustomLogger):
# embeddings are static, only the prompt is embedded per request). The lock
# serializes the one-time build so concurrent cold-start requests don't each
# construct the index and fire duplicate embedding calls.
self._semantic_routelayer: Optional[SemanticRouter] = None
self._semantic_routelayer: SemanticRouter | None = None
self._semantic_routelayer_lock = asyncio.Lock()
# Pre-compile regex patterns for efficiency
@ -184,6 +188,10 @@ class ComplexityRouter(CustomLogger):
re.compile(r"[a-z]\)\s", re.IGNORECASE),
]
self.adaptive_router: AdaptiveRouter | None = None
self._model_tiers: dict[str, tuple[ComplexityTier, ...]] = {}
self._adaptive_init_attempted = False
verbose_router_logger.debug(f"ComplexityRouter initialized for {model_name} with tiers: {self.config.tiers}")
def _estimate_tokens(self, text: str) -> int:
@ -227,12 +235,12 @@ class ComplexityRouter(CustomLogger):
def _score_keyword_match(
self,
text: str,
keywords: List[str],
keywords: list[str],
name: str,
signal_label: str,
thresholds: Tuple[int, int], # (low, high)
scores: Tuple[float, float, float], # (none, low, high)
) -> Tuple[DimensionScore, int]:
thresholds: tuple[int, int], # (low, high)
scores: tuple[float, float, float], # (none, low, high)
) -> tuple[DimensionScore, int]:
"""Score based on keyword matches using word boundary matching.
Returns:
@ -270,7 +278,7 @@ class ComplexityRouter(CustomLogger):
return DimensionScore("questionComplexity", 0.5, f"{count} questions")
return DimensionScore("questionComplexity", 0, None)
def classify(self, prompt: str, system_prompt: Optional[str] = None) -> Tuple[ComplexityTier, float, List[str]]:
def classify(self, prompt: str, system_prompt: str | None = None) -> tuple[ComplexityTier, float, list[str]]:
"""
Classify a prompt by complexity.
@ -329,7 +337,7 @@ class ComplexityRouter(CustomLogger):
(0, -1.0, -1.0),
)
dimensions: List[DimensionScore] = [
dimensions: list[DimensionScore] = [
self._score_token_count(estimated_tokens),
code_score,
reasoning_score,
@ -371,8 +379,8 @@ class ComplexityRouter(CustomLogger):
async def aclassify(
self,
prompt: str,
system_prompt: Optional[str] = None,
request_kwargs: Optional[dict[str, Any]] = None,
system_prompt: str | None = None,
request_kwargs: dict[str, Any] | None = None,
) -> tuple[ComplexityTier, float, list[str]]:
"""
Classify a prompt by complexity, using the LLM classifier when configured.
@ -395,8 +403,8 @@ class ComplexityRouter(CustomLogger):
async def _classify_with_llm(
self,
prompt: str,
system_prompt: Optional[str] = None,
request_kwargs: Optional[dict[str, Any]] = None,
system_prompt: str | None = None,
request_kwargs: dict[str, Any] | None = None,
) -> ComplexityTier:
"""Call the configured classifier model and parse its structured tier response."""
llm_config = self.config.classifier_llm_config
@ -437,23 +445,196 @@ class ComplexityRouter(CustomLogger):
"""
tier_key = tier.value if isinstance(tier, ComplexityTier) else tier
# Check config tiers mapping
model = self.config.tiers.get(tier_key)
if model:
return model
if tier_key in self.config.tiers:
return self._pick_from_tier_value(self.config.tiers[tier_key], tier_key)
# Fallback to default model if configured
if self.config.default_model:
return self.config.default_model
# Last resort: return MEDIUM tier model or error
medium_model = self.config.tiers.get(ComplexityTier.MEDIUM.value)
if medium_model:
return medium_model
medium_key = ComplexityTier.MEDIUM.value
if medium_key in self.config.tiers:
return self._pick_from_tier_value(self.config.tiers[medium_key], medium_key)
raise ValueError(f"No model configured for tier {tier_key} and no default_model set")
def _lexical_tier_override(self, user_message: str) -> Optional[ComplexityTier]:
@staticmethod
def _pick_from_tier_value(model: str | list[str], tier_key: str) -> str:
if isinstance(model, str):
return model
if not model:
raise ValueError(f"Empty model pool for tier {tier_key}")
return random.choice(model)
def _tier_pools(self) -> dict[str, list[str]]:
return {tier: (models if isinstance(models, list) else [models]) for tier, models in self.config.tiers.items()}
def _ensure_adaptive_router(self) -> Any | None:
if not self.config.adaptive:
return None
if self.adaptive_router is not None:
return self.adaptive_router
if self._adaptive_init_attempted:
return self.adaptive_router
self._adaptive_init_attempted = True
from litellm.router_strategy.adaptive_router.adaptive_router import (
AdaptiveRouter,
)
from litellm.router_strategy.adaptive_router.config import (
ADAPTIVE_ROUTER_CHOSEN_MODEL_KEY,
)
from litellm.types.router import (
AdaptiveRouterConfig,
AdaptiveRouterPreferences,
)
pools = self._tier_pools()
available_models = list(dict.fromkeys(model for models in pools.values() for model in models))
self._model_tiers = {
model: tuple(ComplexityTier(tier_name) for tier_name, models in pools.items() if model in models)
for model in available_models
}
model_to_prefs: dict[str, AdaptiveRouterPreferences] = {}
model_to_cost: dict[str, float] = {}
model_list = getattr(self.litellm_router_instance, "model_list", None) or []
name_to_indices = getattr(self.litellm_router_instance, "model_name_to_deployment_indices", {}) or {}
for name in available_models:
indices = name_to_indices.get(name, [])
if not indices:
model_to_prefs[name] = AdaptiveRouterPreferences(quality_tier=2, strengths=[])
model_to_cost[name] = 0.0
continue
deployment = model_list[indices[0]]
mi = deployment.get("model_info") if isinstance(deployment, dict) else deployment.model_info
mi_dict: dict[str, Any] = mi if isinstance(mi, dict) else (mi.model_dump() if mi else {})
prefs_raw = mi_dict.get("adaptive_router_preferences")
if prefs_raw is not None:
model_to_prefs[name] = AdaptiveRouterPreferences(**prefs_raw)
else:
model_to_prefs[name] = AdaptiveRouterPreferences(quality_tier=2, strengths=[])
lp = deployment.get("litellm_params") if isinstance(deployment, dict) else deployment.litellm_params
lp_dict: dict[str, Any] = lp if isinstance(lp, dict) else (lp.model_dump() if lp else {})
cost = lp_dict.get("input_cost_per_token")
model_to_cost[name] = float(cost) if cost is not None else 0.0
self.adaptive_router = AdaptiveRouter(
router_name=self.model_name,
config=AdaptiveRouterConfig(
available_models=available_models,
weights=self.config.adaptive_weights,
),
model_to_prefs=model_to_prefs,
model_to_cost=model_to_cost,
)
self._adaptive_chosen_model_key = ADAPTIVE_ROUTER_CHOSEN_MODEL_KEY
return self.adaptive_router
def _soft_floor_pick(
self,
classified_tier: ComplexityTier,
user_message: str,
request_kwargs: dict[str, Any] | None = None,
) -> str:
from litellm.router_strategy.adaptive_router.bandit import (
normalized_cost,
thompson_sample,
)
from litellm.router_strategy.adaptive_router.classifier import classify_prompt
adaptive = self._ensure_adaptive_router()
if adaptive is None:
return self.get_model_for_tier(classified_tier)
request_type = classify_prompt(user_message)
classified_idx = TIER_SEVERITY_ORDER.index(classified_tier)
pools = self._tier_pools()
classified_candidates = tuple(pools.get(classified_tier.value, ()))
cold_start_candidates = tuple(
model for model in classified_candidates if adaptive._cells[(request_type, model)].total_samples == 0
)
if cold_start_candidates:
chosen_model = random.choice(cold_start_candidates)
if request_kwargs is not None:
metadata = request_kwargs.setdefault("metadata", {})
if isinstance(metadata, dict):
metadata["adaptive_router_decision"] = {
"phase": "cold_start",
"classified_tier": classified_tier.value,
"request_type": request_type.value,
"eligible_mode": "classified_tier",
"quality_weight": self.config.adaptive_weights.quality,
"cost_weight": self.config.adaptive_weights.cost,
"tier_distance_penalty": self.config.tier_distance_penalty,
"chosen_model": chosen_model,
"candidates": [
{
"model": model,
"total_samples": adaptive._cells[(request_type, model)].total_samples,
}
for model in cold_start_candidates
],
}
return chosen_model
if self.config.adaptive_eligible == "classified_tier":
candidates = list(classified_candidates)
if not candidates:
return self.get_model_for_tier(classified_tier)
else:
candidates = list(adaptive.config.available_models)
all_costs = [adaptive.model_to_cost.get(m, 0.0) for m in candidates]
quality_weight = self.config.adaptive_weights.quality
cost_weight = self.config.adaptive_weights.cost
penalty_weight = self.config.tier_distance_penalty
best_model: str | None = None
best_score = float("-inf")
candidate_scores: list[dict[str, Any]] = []
for model in candidates:
cell = adaptive._cells[(request_type, model)]
quality_sample = thompson_sample(cell)
cost_score = normalized_cost(adaptive.model_to_cost.get(model, 0.0), all_costs)
if self.config.adaptive_eligible == "classified_tier":
distance = 0
else:
model_tiers = self._model_tiers.get(model, (classified_tier,))
distance = min(
abs(TIER_SEVERITY_ORDER.index(model_tier) - classified_idx) for model_tier in model_tiers
)
score = quality_weight * quality_sample + cost_weight * cost_score - penalty_weight * distance
candidate_scores.append(
{
"model": model,
"quality_sample": quality_sample,
"cost_score": cost_score,
"tier_distance": distance,
"score": score,
}
)
if score > best_score:
best_score = score
best_model = model
if best_model is None:
return self.get_model_for_tier(classified_tier)
if request_kwargs is not None:
metadata = request_kwargs.setdefault("metadata", {})
if isinstance(metadata, dict):
metadata["adaptive_router_decision"] = {
"phase": "adaptive",
"classified_tier": classified_tier.value,
"request_type": request_type.value,
"eligible_mode": self.config.adaptive_eligible,
"quality_weight": quality_weight,
"cost_weight": cost_weight,
"tier_distance_penalty": penalty_weight,
"chosen_model": best_model,
"candidates": candidate_scores,
}
return best_model
def _lexical_tier_override(self, user_message: str) -> ComplexityTier | None:
"""When keyword_tier_rules match literally, the most-severe matched tier wins.
Escalating to the highest tier (rather than the first rule in the list) keeps
@ -471,7 +652,7 @@ class ComplexityRouter(CustomLogger):
return None
return max(matched_tiers, key=TIER_SEVERITY_ORDER.index)
def _get_or_create_semantic_routelayer(self) -> "SemanticRouter":
def _get_or_create_semantic_routelayer(self) -> SemanticRouter:
"""Build (once) a SemanticRouter with one route per tier, utterances = that tier's keywords."""
if self._semantic_routelayer is not None:
return self._semantic_routelayer
@ -510,7 +691,7 @@ class ComplexityRouter(CustomLogger):
self._semantic_routelayer = routelayer
return routelayer
async def _ensure_semantic_routelayer(self) -> "SemanticRouter":
async def _ensure_semantic_routelayer(self) -> SemanticRouter:
"""Return the cached route layer, building it once under a lock if needed.
The build embeds the static route utterances via the encoder's synchronous path,
@ -526,7 +707,7 @@ class ComplexityRouter(CustomLogger):
routelayer = await asyncio.to_thread(self._get_or_create_semantic_routelayer)
return routelayer
async def _semantic_tier_override(self, user_message: str, request_kwargs: Dict) -> Optional[ComplexityTier]:
async def _semantic_tier_override(self, user_message: str, request_kwargs: dict) -> ComplexityTier | None:
"""Match the prompt against keyword_tier_rules by embedding similarity.
Embeds the query ourselves (instead of letting SemanticRouter.acall embed it
@ -566,7 +747,7 @@ class ComplexityRouter(CustomLogger):
except ValueError:
return None
async def _resolve_keyword_tier_override(self, user_message: str, request_kwargs: Dict) -> Optional[ComplexityTier]:
async def _resolve_keyword_tier_override(self, user_message: str, request_kwargs: dict) -> ComplexityTier | None:
"""Resolve a keyword_tier_rule override, semantically or lexically per config.
Returns None (no override -> fall through to the scorer) not only when no rule
@ -587,60 +768,28 @@ class ComplexityRouter(CustomLogger):
def _resolve_messages(
self,
messages: Optional[List[Dict[str, Any]]],
request_kwargs: Dict,
) -> Optional[List[Dict[str, Any]]]:
messages: list[dict[str, Any]] | None,
request_kwargs: dict,
) -> list[dict[str, Any]] | None:
"""
Resolve messages from the request, converting from other formats if needed.
Uses the guardrail translation handler dispatch to convert Responses API
``input`` (or other non-chat-completions formats) into OpenAI-spec messages.
"""
if messages:
return messages
from litellm.litellm_core_utils.api_route_to_call_types import (
get_call_types_for_route,
from litellm.litellm_core_utils.prompt_templates.factory import (
resolve_structured_messages,
)
from litellm.llms import load_guardrail_translation_mappings
from litellm.types.utils import CallTypes
mappings = load_guardrail_translation_mappings()
call_type: Optional[CallTypes] = None
# 1. Try route-based inference from proxy metadata
route = request_kwargs.get("litellm_metadata", {}).get("user_api_key_request_route")
if route:
call_types_list = get_call_types_for_route(route)
if call_types_list:
for ct in call_types_list:
if ct in mappings:
call_type = ct
break
# 2. Fallback: try each mapped handler until one produces messages
handlers_to_try: List[Any] = []
if call_type is not None and call_type in mappings:
handlers_to_try.append(mappings[call_type]())
else:
handlers_to_try.extend(handler_cls() for handler_cls in mappings.values())
for handler in handlers_to_try:
structured = handler.get_structured_messages(request_kwargs)
if structured:
return [
msg if isinstance(msg, dict) else msg.model_dump() # type: ignore
for msg in structured
]
return None
return resolve_structured_messages(messages=messages, request_kwargs=request_kwargs)
@staticmethod
def _extract_user_message_and_system_prompt(
messages: List[Dict[str, Any]],
) -> Tuple[Optional[str], Optional[str]]:
messages: list[dict[str, Any]],
) -> tuple[str | None, str | None]:
"""Extract the last user message text and last system prompt from messages."""
user_message: Optional[str] = None
system_prompt: Optional[str] = None
user_message: str | None = None
system_prompt: str | None = None
for msg in reversed(messages):
role = msg.get("role", "")
@ -660,17 +809,115 @@ class ComplexityRouter(CustomLogger):
return user_message, system_prompt
@staticmethod
def _iter_metadata_dicts(request_kwargs: dict) -> list[dict]:
"""Metadata may land on `metadata` or `litellm_metadata` depending on the
endpoint, mirroring DeploymentAffinityCheck's precedence."""
return [
metadata
for metadata_key in ("litellm_metadata", "metadata")
if isinstance(metadata := request_kwargs.get(metadata_key), dict)
]
@staticmethod
def _get_session_id_from_request_kwargs(request_kwargs: dict) -> str | None:
"""Resolve a client-supplied session_id."""
for metadata in ComplexityRouter._iter_metadata_dicts(request_kwargs):
session_id = metadata.get("session_id")
if session_id is not None:
return str(session_id)
return None
@staticmethod
def _get_user_api_key_hash_from_request_kwargs(request_kwargs: dict) -> str | None:
"""Resolve the proxy-derived API key hash, the same trust boundary
DeploymentAffinityCheck uses for its own key-based affinity (not the
client-supplied OpenAI `user` param, which isn't authenticated)."""
for metadata in ComplexityRouter._iter_metadata_dicts(request_kwargs):
user_key = metadata.get("user_api_key_hash")
if user_key is not None:
return str(user_key)
return None
def _get_session_affinity_cache_key(self, session_id: str, request_kwargs: dict) -> str:
# Namespace by the caller's API key hash so two different callers reusing the
# same client-supplied session_id can't poison each other's routing pin. Falls
# back to "unscoped" only when there's no authenticated caller to scope by
# (e.g. direct Router usage without the proxy layer).
caller_scope = self._get_user_api_key_hash_from_request_kwargs(request_kwargs) or "unscoped"
return f"complexity_router_session_affinity:v1:{self.model_name}:{caller_scope}:{session_id}"
async def async_pre_routing_hook(
self,
model: str,
request_kwargs: Dict,
messages: Optional[List[Dict[str, Any]]] = None,
input: Optional[Union[str, List]] = None,
specific_deployment: Optional[bool] = False,
) -> Optional["PreRoutingHookResponse"]:
request_kwargs: dict,
messages: list[dict[str, Any]] | None = None,
input: Union[str, list] | None = None,
specific_deployment: bool | None = False,
) -> PreRoutingHookResponse | None:
"""
Pre-routing hook called before the routing decision.
When `session_affinity` is enabled and a session_id is resolvable on the request,
pins the model chosen on the session's first turn and reuses it for every later
turn, skipping classification entirely. Otherwise delegates to `_classify_and_route`.
"""
from litellm.types.router import PreRoutingHookResponse
session_id = self._get_session_id_from_request_kwargs(request_kwargs) if self.config.session_affinity else None
cache_key = self._get_session_affinity_cache_key(session_id, request_kwargs) if session_id is not None else None
if cache_key is not None:
pinned_model = await self.litellm_router_instance.cache.async_get_cache(key=cache_key)
if isinstance(pinned_model, str):
# Refresh the TTL on every hit so an active session doesn't lose its
# pin mid-conversation just because it outlives the original write.
await self.litellm_router_instance.cache.async_set_cache(
key=cache_key,
value=pinned_model,
ttl=self.config.session_affinity_ttl_seconds,
)
if self.config.adaptive:
from litellm.router_strategy.adaptive_router.config import (
ADAPTIVE_ROUTER_CHOSEN_MODEL_KEY,
)
kwargs_metadata = request_kwargs.setdefault("metadata", {})
if isinstance(kwargs_metadata, dict):
kwargs_metadata[ADAPTIVE_ROUTER_CHOSEN_MODEL_KEY] = pinned_model
verbose_router_logger.info(
f"ComplexityRouter: routing decision cause=session_affinity_pin, routed_model={pinned_model}"
)
has_original_messages = messages is not None and len(messages) > 0
return PreRoutingHookResponse(
model=pinned_model,
messages=messages if has_original_messages else None,
)
response = await self._classify_and_route(
model=model,
request_kwargs=request_kwargs,
messages=messages,
input=input,
specific_deployment=specific_deployment,
)
if cache_key is not None and response is not None:
await self.litellm_router_instance.cache.async_set_cache(
key=cache_key,
value=response.model,
ttl=self.config.session_affinity_ttl_seconds,
)
return response
async def _classify_and_route(
self,
model: str,
request_kwargs: dict,
messages: list[dict[str, Any]] | None = None,
input: Union[str, list] | None = None,
specific_deployment: bool | None = False,
) -> PreRoutingHookResponse | None:
"""
Classifies the request by complexity and returns the appropriate model.
Supports chat completions (messages), Responses API (input), and other
formats via the guardrail translation handler dispatch.
@ -719,12 +966,25 @@ class ComplexityRouter(CustomLogger):
)
tier, score, signals = await self.aclassify(user_message, system_prompt, request_kwargs)
routed_model = self.get_model_for_tier(tier)
verbose_router_logger.info(
f"ComplexityRouter: routing decision cause=complexity_scorer, tier={tier.value}, "
f"score={score:.3f}, signals={signals}, routed_model={routed_model}"
)
if self.config.adaptive:
routed_model = self._soft_floor_pick(tier, user_message, request_kwargs)
adaptive = self._ensure_adaptive_router()
if adaptive is not None:
kwargs_metadata = request_kwargs.setdefault("metadata", {})
if isinstance(kwargs_metadata, dict):
chosen_key = getattr(self, "_adaptive_chosen_model_key", "adaptive_router_chosen_model")
kwargs_metadata[chosen_key] = routed_model
verbose_router_logger.info(
f"ComplexityRouter[adaptive]: routing decision cause=complexity_scorer, "
f"tier={tier.value}, score={score:.3f}, "
f"signals={signals}, routed_model={routed_model}"
)
else:
routed_model = self.get_model_for_tier(tier)
verbose_router_logger.info(
f"ComplexityRouter: routing decision cause=complexity_scorer, tier={tier.value}, "
f"score={score:.3f}, signals={signals}, routed_model={routed_model}"
)
return PreRoutingHookResponse(
model=routed_model,

View file

@ -6,9 +6,11 @@ All values are configurable via proxy config.yaml.
"""
from enum import Enum
from typing import Dict, List, Literal, Optional
from typing import Literal
from pydantic import BaseModel, ConfigDict, Field, model_validator
from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator
from litellm.types.router import AdaptiveRouterWeights
class ComplexityTier(str, Enum):
@ -27,11 +29,13 @@ TIER_SEVERITY_ORDER: tuple[ComplexityTier, ...] = (
ComplexityTier.REASONING,
)
DEFAULT_TIER_DISTANCE_PENALTY: float = 0.5
class KeywordTierRule(BaseModel):
"""A deterministic override: if any keyword matches, route to this tier."""
keywords: List[str] = Field(
keywords: list[str] = Field(
min_length=1,
description="Keywords/phrases that trigger this rule (lexical or semantic match)",
)
@ -56,7 +60,7 @@ class KeywordTierRule(BaseModel):
# Note: Keywords should be full words/phrases to avoid substring false positives.
# The matching logic uses word boundary detection for single-word keywords.
DEFAULT_CODE_KEYWORDS: List[str] = [
DEFAULT_CODE_KEYWORDS: list[str] = [
"function",
"class",
"def",
@ -104,7 +108,7 @@ DEFAULT_CODE_KEYWORDS: List[str] = [
"pull request",
]
DEFAULT_REASONING_KEYWORDS: List[str] = [
DEFAULT_REASONING_KEYWORDS: list[str] = [
"step by step",
"think through",
"let's think",
@ -126,7 +130,7 @@ DEFAULT_REASONING_KEYWORDS: List[str] = [
"conclude",
]
DEFAULT_TECHNICAL_KEYWORDS: List[str] = [
DEFAULT_TECHNICAL_KEYWORDS: list[str] = [
"architecture",
"distributed",
"scalable",
@ -158,7 +162,7 @@ DEFAULT_TECHNICAL_KEYWORDS: List[str] = [
# Note: "async", "kubernetes", "docker" are in DEFAULT_CODE_KEYWORDS
]
DEFAULT_SIMPLE_KEYWORDS: List[str] = [
DEFAULT_SIMPLE_KEYWORDS: list[str] = [
"what is",
"what's",
"define",
@ -191,7 +195,7 @@ DEFAULT_SIMPLE_KEYWORDS: List[str] = [
# ─── Default Dimension Weights ───
DEFAULT_DIMENSION_WEIGHTS: Dict[str, float] = {
DEFAULT_DIMENSION_WEIGHTS: dict[str, float] = {
"tokenCount": 0.10, # Reduced - length is less important than content
"codePresence": 0.30, # High - code requests need capable models
"reasoningMarkers": 0.25, # High - explicit reasoning requests
@ -204,7 +208,7 @@ DEFAULT_DIMENSION_WEIGHTS: Dict[str, float] = {
# ─── Default Tier Boundaries ───
DEFAULT_TIER_BOUNDARIES: Dict[str, float] = {
DEFAULT_TIER_BOUNDARIES: dict[str, float] = {
"simple_medium": 0.15, # Lower threshold to catch more MEDIUM cases
"medium_complex": 0.35, # Lower threshold to catch technical COMPLEX cases
"complex_reasoning": 0.60, # Reasoning tier reserved for explicit reasoning markers
@ -213,7 +217,7 @@ DEFAULT_TIER_BOUNDARIES: Dict[str, float] = {
# ─── Default Token Thresholds ───
DEFAULT_TOKEN_THRESHOLDS: Dict[str, int] = {
DEFAULT_TOKEN_THRESHOLDS: dict[str, int] = {
"simple": 15, # Only very short prompts (<15 tokens) are penalized
"complex": 400, # Long prompts (>400 tokens) get complexity boost
}
@ -221,7 +225,7 @@ DEFAULT_TOKEN_THRESHOLDS: Dict[str, int] = {
# ─── Default Tier to Model Mapping ───
DEFAULT_TIER_MODELS: Dict[str, str] = {
DEFAULT_TIER_MODELS: dict[str, str] = {
"SIMPLE": "gpt-4o-mini",
"MEDIUM": "gpt-4o",
"COMPLEX": "claude-sonnet-4-20250514",
@ -244,44 +248,47 @@ class ClassifierLLMConfig(BaseModel):
class ComplexityRouterConfig(BaseModel):
"""Configuration for the ComplexityRouter."""
# Tier to model mapping
tiers: Dict[str, str] = Field(
# string = pin; list = random pick when adaptive=False, soft-floor home pool when adaptive=True
tiers: dict[str, str | list[str]] = Field(
default_factory=lambda: DEFAULT_TIER_MODELS.copy(),
description="Mapping of complexity tiers to model names",
description=(
"Mapping of complexity tiers to a model or model pool. "
"A list is randomly picked from when adaptive=False, and used as a soft-floor home pool when adaptive=True"
),
)
# Tier boundaries (normalized scores)
tier_boundaries: Dict[str, float] = Field(
tier_boundaries: dict[str, float] = Field(
default_factory=lambda: DEFAULT_TIER_BOUNDARIES.copy(),
description="Score boundaries between tiers",
)
# Token count thresholds
token_thresholds: Dict[str, int] = Field(
token_thresholds: dict[str, int] = Field(
default_factory=lambda: DEFAULT_TOKEN_THRESHOLDS.copy(),
description="Token count thresholds for simple/complex classification",
)
# Dimension weights
dimension_weights: Dict[str, float] = Field(
dimension_weights: dict[str, float] = Field(
default_factory=lambda: DEFAULT_DIMENSION_WEIGHTS.copy(),
description="Weights for each scoring dimension",
)
# Keyword lists (overridable)
code_keywords: Optional[List[str]] = Field(
code_keywords: list[str] | None = Field(
default=None,
description="Keywords indicating code-related content",
)
reasoning_keywords: Optional[List[str]] = Field(
reasoning_keywords: list[str] | None = Field(
default=None,
description="Keywords indicating reasoning-required content",
)
technical_keywords: Optional[List[str]] = Field(
technical_keywords: list[str] | None = Field(
default=None,
description="Keywords indicating technical content",
)
custom_technical_keywords: Optional[list[str]] = Field(
custom_technical_keywords: list[str] | None = Field(
default=None,
description=(
"Domain-specific technical keywords appended to the effective base list "
@ -290,13 +297,13 @@ class ComplexityRouterConfig(BaseModel):
"the base list and within this list."
),
)
simple_keywords: Optional[List[str]] = Field(
simple_keywords: list[str] | None = Field(
default=None,
description="Keywords indicating simple/basic queries",
)
# Default model if scoring fails
default_model: Optional[str] = Field(
default_model: str | None = Field(
default=None,
description="Default model to use if tier cannot be determined",
)
@ -306,13 +313,34 @@ class ComplexityRouterConfig(BaseModel):
default="heuristic",
description="Classification strategy: local regex/keyword scoring, or an LLM call",
)
classifier_llm_config: Optional[ClassifierLLMConfig] = Field(
classifier_llm_config: ClassifierLLMConfig | None = Field(
default=None,
description="Configuration for the LLM classifier; required when classifier_type is 'llm'",
)
adaptive: bool = Field(
default=False,
description="Enable adaptive bandit selection with soft complexity floors",
)
adaptive_weights: AdaptiveRouterWeights = Field(
default_factory=lambda: AdaptiveRouterWeights(quality=0.3, cost=0.7),
description="Quality vs cost weights for adaptive selection (used when adaptive=True)",
)
tier_distance_penalty: float = Field(
default=DEFAULT_TIER_DISTANCE_PENALTY,
ge=0.0,
description="Score penalty per tier-step away from the classified tier when adaptive=True",
)
adaptive_eligible: Literal["all", "classified_tier"] = Field(
default="all",
description=(
"When adaptive=True: 'all' scores every pool model with a tier-distance penalty (soft floors); "
"'classified_tier' Thompson-samples only inside the classified tier's pool"
),
)
# Deterministic keyword -> tier overrides, evaluated before weighted scoring
keyword_tier_rules: Optional[List[KeywordTierRule]] = Field(
keyword_tier_rules: list[KeywordTierRule] | None = Field(
default=None,
description="Rules that force a specific tier when their keywords match the prompt",
)
@ -322,7 +350,7 @@ class ComplexityRouterConfig(BaseModel):
default=False,
description="Match keyword_tier_rules by embedding similarity instead of literal text",
)
embedding_model: Optional[str] = Field(
embedding_model: str | None = Field(
default=None,
description="Embedding model (LiteLLM model name) used when semantic_keyword_matching is enabled",
)
@ -333,14 +361,56 @@ class ComplexityRouterConfig(BaseModel):
description="Minimum cosine similarity for a semantic keyword match",
)
# Session affinity: pin the first turn's routed model for the rest of the session
session_affinity: bool = Field(
default=False,
description=(
"When True and a session_id is resolvable on the request, pin the model chosen on the "
"session's first turn and reuse it for every later turn, skipping re-classification."
),
)
session_affinity_ttl_seconds: int = Field(
default=3600,
gt=0,
description="TTL for the session affinity pin; refreshed on every cache hit",
)
model_config = ConfigDict(extra="allow") # Allow additional fields
@field_validator("tiers", mode="before")
@classmethod
def _coerce_tier_values(cls, value: object) -> object:
if not isinstance(value, dict):
return value
coerced: dict[str, object] = {}
for key, item in value.items():
if isinstance(item, str):
coerced[key] = item
elif isinstance(item, (list, tuple)):
coerced[key] = list(item)
else:
coerced[key] = item
return coerced
@model_validator(mode="after")
def _validate_llm_classifier_config(self) -> "ComplexityRouterConfig":
if self.classifier_type == "llm" and self.classifier_llm_config is None:
raise ValueError("classifier_llm_config is required when classifier_type is 'llm'")
return self
@model_validator(mode="after")
def _validate_adaptive_pools(self) -> "ComplexityRouterConfig":
if not self.adaptive:
return self
normalized = {tier: (models if isinstance(models, list) else [models]) for tier, models in self.tiers.items()}
if not any(normalized.values()):
raise ValueError("adaptive=True requires at least one non-empty tier pool")
empty = [tier for tier, models in normalized.items() if not models]
if empty:
raise ValueError(f"adaptive=True tier pools must be non-empty; empty tiers: {empty}")
self.tiers = normalized
return self
@model_validator(mode="after")
def _validate_semantic_matching(self) -> "ComplexityRouterConfig":
if not self.semantic_keyword_matching:

View file

@ -213,6 +213,8 @@ DEFINED_PROMETHEUS_METRICS = Literal[
"litellm_input_audio_tokens_metric",
"litellm_output_reasoning_tokens_metric",
"litellm_output_audio_tokens_metric",
"litellm_video_duration_seconds_metric",
"litellm_images_generated_metric",
"litellm_deployment_successful_fallbacks",
"litellm_deployment_failed_fallbacks",
"litellm_remaining_team_budget_metric",
@ -506,6 +508,9 @@ class PrometheusMetricLabels:
litellm_output_reasoning_tokens_metric = litellm_output_tokens_metric
litellm_output_audio_tokens_metric = litellm_output_tokens_metric
litellm_video_duration_seconds_metric = litellm_output_tokens_metric
litellm_images_generated_metric = litellm_output_tokens_metric
litellm_deployment_state = [
UserAPIKeyLabelNames.v2_LITELLM_MODEL_NAME.value,
UserAPIKeyLabelNames.MODEL_ID.value,
@ -717,6 +722,8 @@ class PrometheusMetricLabels:
"litellm_input_tokens_metric",
"litellm_total_tokens_metric",
"litellm_output_tokens_metric",
"litellm_video_duration_seconds_metric",
"litellm_images_generated_metric",
}
)
# Managed batch metrics

View file

@ -9,7 +9,7 @@ from typing import Any, Dict, List, Literal, Optional, Tuple, Union, get_type_hi
import httpx
from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator
from typing_extensions import Required, TypedDict
from typing_extensions import Protocol, Required, TypedDict
from litellm._uuid import uuid
@ -829,6 +829,35 @@ class PreRoutingHookResponse(BaseModel):
messages: Optional[List[Dict[str, Any]]]
class RoutingContext(BaseModel):
"""
Passed through a Router's `plugins` pipeline before the routing decision is made.
Each plugin reads and mutates this object; the next plugin sees the previous
plugin's changes. `candidate_models` narrows as the pipeline runs -- Router
only selects a deployment whose `litellm_params.model` survives the pipeline.
`raw_messages` and `structured_messages` mirror the pattern
`CustomGuardrail.apply_guardrail` uses: the message shape differs by API
surface (chat completions, Anthropic /v1/messages, Responses API `input`,
...), so plugins that need a stable, provider-agnostic shape should read
`structured_messages` (normalized to OpenAI chat-completions format);
plugins that need the exact original payload can read `raw_messages`.
"""
raw_messages: list[dict[str, Any]]
structured_messages: list[dict[str, Any]]
candidate_models: list[str]
metadata: dict[str, Any] = Field(default_factory=dict)
signals: dict[str, Any] = Field(default_factory=dict)
class RoutingPlugin(Protocol):
"""Interface a custom routing plugin must implement to run in `Router(plugins=[...])`."""
async def run(self, context: RoutingContext) -> RoutingContext: ...
class RequestType(str, enum.Enum):
"""Fixed v0 taxonomy. User-extensible types come in v1."""

View file

@ -209,6 +209,7 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False):
input_cost_per_query: Optional[float] # only for rerank models
input_cost_per_image: Optional[float] # only for vertex ai models
input_cost_per_image_token: Optional[float] # for gpt-image-1 and similar models
input_cost_per_video_token: Optional[float] # for gemini omni models with video input
input_cost_per_audio_per_second: Optional[float] # only for vertex ai models
input_cost_per_video_per_second: Optional[float] # only for vertex ai models
input_cost_per_second: Optional[float] # for OpenAI Speech models
@ -234,6 +235,7 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False):
output_cost_per_character_above_128k_tokens: Optional[float] # only for vertex ai models
output_cost_per_image: Optional[float]
output_cost_per_image_token: Optional[float]
output_cost_per_video_token: Optional[float] # for gemini omni models with video output
output_vector_size: Optional[int]
output_cost_per_reasoning_token: Optional[float]
output_cost_per_video_per_second: Optional[float] # only for vertex ai models
@ -3046,6 +3048,7 @@ class CustomPricingLiteLLMParams(BaseModel):
output_cost_per_character_above_128k_tokens: Optional[float] = None
output_cost_per_image: Optional[float] = None
output_cost_per_image_token: Optional[float] = None
output_cost_per_video_token: Optional[float] = None
output_cost_per_reasoning_token: Optional[float] = None
output_cost_per_video_per_second: Optional[float] = None
output_cost_per_audio_per_second: Optional[float] = None
@ -3055,6 +3058,7 @@ class CustomPricingLiteLLMParams(BaseModel):
cache_read_input_token_cost_above_272k_tokens: Optional[float] = None
cache_read_input_token_cost_above_512k_tokens: Optional[float] = None
input_cost_per_image_token: Optional[float] = None
input_cost_per_video_token: Optional[float] = None
input_cost_per_token_above_272k_tokens: Optional[float] = None
input_cost_per_token_above_512k_tokens: Optional[float] = None
output_cost_per_token_above_272k_tokens: Optional[float] = None
@ -3210,6 +3214,16 @@ all_litellm_params = (
"_litellm_tpm_reserved_model",
"_litellm_tpm_reserved_scopes",
"_litellm_tpm_reservation_released",
"auto_router_config_path",
"auto_router_config",
"auto_router_default_model",
"auto_router_embedding_model",
"complexity_router_config",
"complexity_router_default_model",
"adaptive_router_config",
"adaptive_router_default_model",
"quality_router_config",
"quality_router_default_model",
]
+ list(StandardCallbackDynamicParams.__annotations__.keys())
+ list(CustomPricingLiteLLMParams.model_fields.keys())

View file

@ -5437,6 +5437,7 @@ def _get_model_info_helper(
input_cost_per_second=_model_info.get("input_cost_per_second", None),
input_cost_per_audio_token=_model_info.get("input_cost_per_audio_token", None),
input_cost_per_image_token=_model_info.get("input_cost_per_image_token", None),
input_cost_per_video_token=_model_info.get("input_cost_per_video_token", None),
input_cost_per_image=_model_info.get("input_cost_per_image", None),
input_cost_per_audio_per_second=_model_info.get("input_cost_per_audio_per_second", None),
input_cost_per_video_per_second=_model_info.get("input_cost_per_video_per_second", None),
@ -5480,6 +5481,7 @@ def _get_model_info_helper(
output_cost_per_video_per_second=_model_info.get("output_cost_per_video_per_second", None),
output_cost_per_image=_model_info.get("output_cost_per_image", None),
output_cost_per_image_token=_model_info.get("output_cost_per_image_token", None),
output_cost_per_video_token=_model_info.get("output_cost_per_video_token", None),
output_vector_size=_model_info.get("output_vector_size", None),
citation_cost_per_token=_model_info.get("citation_cost_per_token", None),
tiered_pricing=_model_info.get("tiered_pricing", None),

View file

@ -11331,6 +11331,7 @@
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_native_structured_output": true,
"supports_tool_choice": true,
"supports_vision": true
},
@ -11362,6 +11363,7 @@
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_native_structured_output": true,
"supports_tool_choice": true,
"supports_vision": true,
"supports_web_search": true
@ -11424,6 +11426,7 @@
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_native_structured_output": true,
"supports_max_reasoning_effort": true,
"supports_tool_choice": true,
"supports_vision": true,
@ -11479,6 +11482,7 @@
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_native_structured_output": true,
"supports_tool_choice": true,
"supports_vision": true
},
@ -11506,6 +11510,7 @@
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_native_structured_output": true,
"supports_tool_choice": true,
"supports_vision": true
},
@ -11559,6 +11564,7 @@
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_native_structured_output": true,
"supports_tool_choice": true,
"supports_vision": true,
"supports_output_config": true
@ -11586,6 +11592,7 @@
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_native_structured_output": true,
"supports_tool_choice": true,
"supports_vision": true,
"supports_output_config": true
@ -11614,6 +11621,7 @@
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_native_structured_output": true,
"supports_tool_choice": true,
"supports_vision": true,
"provider_specific_entry": {
@ -11648,6 +11656,7 @@
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_native_structured_output": true,
"supports_tool_choice": true,
"supports_vision": true,
"provider_specific_entry": {
@ -11682,6 +11691,7 @@
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_native_structured_output": true,
"supports_sampling_params": false,
"supports_tool_choice": true,
"supports_vision": true,
@ -11718,6 +11728,7 @@
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_native_structured_output": true,
"supports_sampling_params": false,
"supports_tool_choice": true,
"supports_vision": true,
@ -11788,6 +11799,7 @@
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_native_structured_output": true,
"supports_sampling_params": false,
"supports_tool_choice": true,
"supports_vision": true,
@ -18865,7 +18877,8 @@
"search_context_size_medium": 0.014,
"search_context_size_high": 0.014
},
"web_search_billing_unit": "per_query"
"web_search_billing_unit": "per_query",
"supports_reasoning": false
},
"gemini/gemini-3-pro-image-preview": {
"input_cost_per_image": 0.0011,
@ -18898,6 +18911,7 @@
],
"supports_function_calling": false,
"supports_prompt_caching": true,
"supports_reasoning": false,
"supports_response_schema": true,
"supports_system_messages": true,
"supports_vision": true,
@ -18939,6 +18953,7 @@
],
"supports_function_calling": false,
"supports_prompt_caching": true,
"supports_reasoning": false,
"supports_response_schema": true,
"supports_system_messages": true,
"supports_vision": true,
@ -18980,6 +18995,7 @@
],
"supports_function_calling": false,
"supports_prompt_caching": true,
"supports_reasoning": false,
"supports_response_schema": true,
"supports_system_messages": true,
"supports_vision": true,
@ -19674,6 +19690,39 @@
},
"web_search_billing_unit": "per_query"
},
"gemini/gemini-omni-flash-preview": {
"input_cost_per_audio_token": 1.5e-06,
"input_cost_per_token": 1.5e-06,
"litellm_provider": "gemini",
"max_input_tokens": 1048576,
"max_output_tokens": 65535,
"max_tokens": 65535,
"mode": "chat",
"output_cost_per_reasoning_token": 9e-06,
"output_cost_per_token": 9e-06,
"output_cost_per_video_token": 1.75e-05,
"rpm": 2000,
"source": "https://ai.google.dev/gemini-api/docs/pricing",
"supported_endpoints": [
"/v1/chat/completions"
],
"supported_modalities": [
"text",
"image",
"audio",
"video"
],
"supported_output_modalities": [
"text",
"video"
],
"supports_audio_input": true,
"supports_reasoning": true,
"supports_system_messages": true,
"supports_video_input": true,
"supports_vision": true,
"tpm": 800000
},
"gemini/gemini-3.1-pro-preview": {
"cache_read_input_token_cost": 2e-07,
"cache_read_input_token_cost_above_200k_tokens": 4e-07,
@ -19838,6 +19887,37 @@
},
"web_search_billing_unit": "per_query"
},
"gemini-omni-flash-preview": {
"input_cost_per_audio_token": 1.5e-06,
"input_cost_per_token": 1.5e-06,
"litellm_provider": "vertex_ai-language-models",
"max_input_tokens": 1048576,
"max_output_tokens": 65535,
"max_tokens": 65535,
"mode": "chat",
"output_cost_per_reasoning_token": 9e-06,
"output_cost_per_token": 9e-06,
"output_cost_per_video_token": 1.75e-05,
"source": "https://docs.cloud.google.com/gemini-enterprise-agent-platform/models/gemini/omni-flash-preview",
"supported_endpoints": [
"/v1/chat/completions"
],
"supported_modalities": [
"text",
"image",
"audio",
"video"
],
"supported_output_modalities": [
"text",
"video"
],
"supports_audio_input": true,
"supports_reasoning": true,
"supports_system_messages": true,
"supports_video_input": true,
"supports_vision": true
},
"gemini-3.5-flash": {
"cache_read_input_token_cost": 1.5e-07,
"input_cost_per_audio_token": 1e-06,
@ -36951,6 +37031,7 @@
"supports_parallel_function_calling": true,
"supports_pdf_input": true,
"supports_prompt_caching": true,
"supports_reasoning": false,
"supports_response_schema": true,
"supports_system_messages": true,
"supports_tool_choice": true,
@ -36973,6 +37054,7 @@
"output_cost_per_image_token": 0.00012,
"output_cost_per_token": 1.2e-05,
"output_cost_per_token_batches": 6e-06,
"supports_reasoning": false,
"source": "https://docs.cloud.google.com/vertex-ai/generative-ai/docs/models/gemini/3-pro-image"
},
"vertex_ai/gemini-3-pro-image-preview": {
@ -36988,6 +37070,7 @@
"output_cost_per_image_token": 0.00012,
"output_cost_per_token": 1.2e-05,
"output_cost_per_token_batches": 6e-06,
"supports_reasoning": false,
"source": "https://docs.cloud.google.com/vertex-ai/generative-ai/docs/models/gemini/3-pro-image"
},
"vertex_ai/gemini-3.1-flash-image": {
@ -37001,6 +37084,7 @@
"output_cost_per_image": 0.0672,
"output_cost_per_image_token": 6e-05,
"output_cost_per_token": 3e-06,
"supports_reasoning": false,
"source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#gemini-models"
},
"vertex_ai/gemini-3.1-flash-image-preview": {
@ -37014,6 +37098,7 @@
"output_cost_per_image": 0.0672,
"output_cost_per_image_token": 6e-05,
"output_cost_per_token": 3e-06,
"supports_reasoning": false,
"source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#gemini-models"
},
"vertex_ai/gemini-3.1-flash-lite-preview": {
@ -45284,8 +45369,8 @@
"rules": [
{
"name": "bedrock-claude-ids",
"pattern": "anthropic\\.claude-",
"description": "Any Bedrock-syntax Claude id: the dotted anthropic.claude- segment appears in bare (anthropic.claude-...), region-prefixed (us./eu./au./jp./apac.) and global.-prefixed ids, for every version. Routes these to bedrock before the bare-id Anthropic rule is consulted.",
"pattern": "^(?:[a-z-]+\\.)?anthropic\\.claude-",
"description": "A Bedrock-syntax Claude id, for every version: anthropic.claude- at the start of the name, optionally behind a single dotted geo segment (us./eu./au./jp./apac./global./us-gov.). Anchored to the start because routing rules see the raw request string and provider inference feeds the proxy's provider/* wildcard access checks: an id under an unrecognized namespace such as bedrockz/anthropic.claude-... must stay unroutable rather than resolve to bedrock and slip through a bedrock/* key. Routes to bedrock before the bare-id Anthropic rule is consulted.",
"model_info": {
"litellm_provider": "bedrock"
}

View file

@ -1,6 +1,6 @@
[project]
name = "litellm"
version = "1.93.0"
version = "1.94.0"
description = "Library to easily interface with LLM API providers"
readme = "README.md"
requires-python = ">=3.10, <3.14"
@ -62,8 +62,8 @@ proxy = [
"azure-identity>=1.25.2,<2.0",
"azure-storage-blob>=12.28.0,<13.0",
"mcp>=1.26.0,<2.0",
"litellm-proxy-extras==0.4.76",
"litellm-enterprise==0.1.49",
"litellm-proxy-extras==0.4.77",
"litellm-enterprise==0.1.50",
"RestrictedPython>=8.1,<9.0",
"rich>=13.9.4,<14.0",
"polars>=1.38.1,<2.0",
@ -205,7 +205,7 @@ ci = [
# protobuf, Pillow is a compiled C extension).
"tenacity==8.5.0",
"google-generativeai==0.8.6",
"Pillow==12.2.0",
"Pillow==12.3.0",
# Azure batch E2E tests still import psycopg2 directly.
"psycopg2-binary==2.9.11",
"pytest-codspeed==4.3.0",
@ -264,6 +264,8 @@ constraint-dependencies = [
"aiohttp>=3.14.1,<4.0",
"packaging>=24.0",
"soupsieve>=2.8.4",
"httplib2>=0.32.0",
"setuptools>=83.0.0",
]
override-dependencies = [
# a2a-sdk 1.x requires packaging>=24.0; lunary 1.4.x still caps at <24.0.
@ -284,7 +286,7 @@ members = ["enterprise", "litellm-proxy-extras"]
profile = "black"
[tool.commitizen]
version = "1.93.0"
version = "1.94.0"
version_files = [
"pyproject.toml:^version",
]

View file

@ -306,7 +306,7 @@
"limit": 9
},
"TID251": {
"limit": 2710
"limit": 2701
},
"TRY002": {
"limit": 548
@ -324,7 +324,7 @@
"limit": 883
},
"UP006": {
"limit": 12869
"limit": 12792
},
"UP007": {
"limit": 2570
@ -354,7 +354,7 @@
"limit": 4
},
"UP035": {
"limit": 2295
"limit": 2284
},
"UP036": {
"limit": 4
@ -363,6 +363,6 @@
"limit": 105
},
"UP045": {
"limit": 18517
"limit": 18462
}
}

View file

@ -422,6 +422,7 @@ model LiteLLM_VerificationToken {
budget_reset_at DateTime?
allowed_cache_controls String[] @default([])
allowed_routes String[] @default([])
key_type String?
policies String[] @default([])
access_group_ids String[] @default([])
model_spend Json @default("{}")
@ -516,6 +517,7 @@ model LiteLLM_DeletedVerificationToken {
budget_reset_at DateTime?
allowed_cache_controls String[] @default([])
allowed_routes String[] @default([])
key_type String?
policies String[] @default([])
access_group_ids String[] @default([])
model_spend Json @default("{}")

View file

@ -89,6 +89,11 @@ EOF
status=0
bootstrap_hint() {
echo " This checkout looks unprovisioned (fresh worktree or clone)." >&2
echo " Fix: make bootstrap" >&2
}
if [ -n "$litellm_py_files" ]; then
echo "pre-commit: linting Python (make lint)"
make lint || { echo "✗ Python lint failed. Fix the reds above, then re-run make pre-commit." >&2; status=1; }
@ -109,7 +114,13 @@ fi
if [ -n "$ui_prettier_files" ] || [ -n "$ui_eslint_files" ]; then
echo "pre-commit: linting dashboard (prettier + eslint + lint budgets)"
lint_dashboard || { echo "✗ Dashboard lint failed. See above; format with: (cd ui/litellm-dashboard && npm run format)." >&2; status=1; }
if [ ! -d ui/litellm-dashboard/node_modules ]; then
echo "✗ ui/litellm-dashboard/node_modules is missing; dashboard lint cannot run." >&2
bootstrap_hint
status=1
else
lint_dashboard || { echo "✗ Dashboard lint failed. See above; format with: (cd ui/litellm-dashboard && npm run format)." >&2; status=1; }
fi
fi
if [ -n "$spec_files" ]; then
@ -118,7 +129,15 @@ if [ -n "$spec_files" ]; then
# and an up-to-date Prisma client; check-ui-api-types.yml installs those and runs
# prisma generate before gen:api, so mirror that here or a stale client can mask
# drift that CI will still flag.
if ! uv run --no-sync python scripts/prisma_generate_if_needed.py; then
if [ ! -d ui/litellm-dashboard/node_modules ]; then
echo "✗ ui/litellm-dashboard/node_modules is missing; the gen:api sync check cannot run." >&2
bootstrap_hint
status=1
elif ! uv run --no-sync python -c "import orjson, prisma" 2>/dev/null; then
echo "✗ The Python env lacks the proxy deps (orjson/prisma) that gen:api needs." >&2
bootstrap_hint
status=1
elif ! uv run --no-sync python scripts/prisma_generate_if_needed.py; then
echo "✗ Could not regenerate Prisma client (prisma generate failed)." >&2
status=1
elif ( cd ui/litellm-dashboard && LITELLM_PYTHON="uv run --no-sync python" npm run gen:api ); then

View file

@ -11,7 +11,7 @@ Each subdirectory under `tests/e2e/` is one suite, scoped to an endpoint family
- `embeddings/` - the `/embeddings` endpoint across providers
- `batches/` - the `/batches` endpoint (placeholder until the first test lands)
- `realtime/` - realtime websocket sessions, including the pipecat audio path
- `quota_management/` - quota enforcement and accounting, one subfolder per behavior: `budgets/` (budget definition, enforcement, and reset windows: key, team, tag, soft, multi-window) and `spend_tracking/` (spend logging and cost attribution on `/spend/*`)
- `quota_management/` - quota enforcement and accounting, one subfolder per behavior: `ratelimit/` (rpm/tpm blocks, window reset, pacing headers on live traffic), `budgets/` (budget definition, enforcement, and reset windows: key, team, tag, soft, multi-window), and `spend_tracking/` (spend logging and cost attribution on `/spend/*`)
- `management/` - key/team/user/organization management routes: create/update/delete persistence via the info routes, team membership, and llm-only-key route denials; also the dashboard UI behavior on top of them, driven through the proxy-served UI at /ui with playwright (optional dep behind importorskip)
- `logging/` - logging-integration delivery (datadog and friends)
- `security/` - secret handling and log-leak protection

View file

@ -17,6 +17,7 @@ from __future__ import annotations
import json
import time
from datetime import datetime, timedelta, timezone
from typing import Callable
import pytest
@ -49,7 +50,7 @@ from e2e_http import (
unwrap,
)
from lifecycle import ResourceManager
from models import KeyGenerateBody, SpendLogRow, SpendLogsParams
from models import KeyGenerateBody, SpendLogRow
pytestmark = pytest.mark.e2e
@ -349,6 +350,10 @@ def test_rate_limited_batch_create_leaves_no_unattributed_spend_row(
the file-read path fires while the batch itself is not blocked.
``resources.key()`` cannot set limits, so the key is minted on the gateway
directly and its delete deferred.
Snapshots read /spend/logs/v2 over a bounded window around the test instead
of the unpaginated /spend/logs whole-table read, which grows with the
environment and OOMed the e2e runner on stage.
"""
user_id = f"e2e-batch-rl-{unique_marker()}"
key = client.gateway.generate_key(
@ -356,8 +361,13 @@ def test_rate_limited_batch_create_leaves_no_unattributed_spend_row(
)
resources.defer(lambda: client.gateway.delete_key(key))
window_start = datetime.now(timezone.utc) - timedelta(hours=1)
window_end = window_start + timedelta(hours=2)
before = frozenset(
row.request_id for row in unattributed_rows(client.gateway.spend_logs(SpendLogsParams()))
row.request_id
for row in unattributed_rows(
client.gateway.spend_logs_window(start=window_start, end=window_end)
)
)
file = unwrap(
@ -379,7 +389,9 @@ def test_rate_limited_batch_create_leaves_no_unattributed_spend_row(
new_orphans = [
row
for row in unattributed_rows(client.gateway.spend_logs(SpendLogsParams()))
for row in unattributed_rows(
client.gateway.spend_logs_window(start=window_start, end=window_end)
)
if row.request_id not in before
]
assert not new_orphans, (

View file

@ -8,7 +8,7 @@
- {id: logging.datadog.success.exports_metric, module: logging, tier: P0, event: success, assertions: [exports_metric], exercised_on: [chat_completions, messages, embeddings], source: "integrations/datadog/datadog.py", rationale: "Powers dashboards/alerts; cardinality regressions common"}
- {id: logging.datadog.failure.exports_metric, module: logging, tier: P0, event: failure, assertions: [exports_metric], exercised_on: [chat_completions], source: "integrations/datadog/datadog.py", rationale: "Failure metrics for alerting/SLO"}
- {id: logging.prometheus.success.exports_metric, module: logging, tier: P0, event: success, assertions: [exports_metric], exercised_on: [chat_completions, messages, embeddings], source: "integrations/prometheus.py", rationale: "Standard OSS metrics; per-key cardinality (existing e2e)"}
- {id: logging.otel.success.exports_metric, module: logging, tier: P0, event: success, assertions: [exports_metric], exercised_on: [chat_completions, messages, embeddings], source: "integrations/otel/logger.py", rationale: "OTEL spans on every call path"}
- {id: logging.otel.success.exports_metric, module: logging, tier: P0, event: success, assertions: [exports_metric], exercised_on: [chat_completions, messages, responses, embeddings], source: "integrations/otel/logger.py", rationale: "OTEL spans on every call path"}
- {id: logging.otel.failure.exports_metric, module: logging, tier: P0, event: failure, assertions: [exports_metric], exercised_on: [chat_completions, messages], source: "integrations/otel/logger.py", rationale: "Error spans for observability continuity"}
- {id: logging.braintrust.success.logs_spend, module: logging, tier: P1, event: success, assertions: [logs_spend], exercised_on: [chat_completions, messages], source: "integrations/braintrust_logging.py", rationale: "Evals platform spend"}
- {id: logging.langsmith.success.logs_spend, module: logging, tier: P1, event: success, assertions: [logs_spend], exercised_on: [chat_completions, messages], source: "integrations/langsmith.py", rationale: "LangChain ecosystem"}

View file

@ -18,6 +18,12 @@ configs:
type: redis
host: redis
port: 6379
# OTEL v2 trace destination for the logging suite's trace-completeness
# tests: the arize_phoenix preset is OTLP with a configurable endpoint
# (PHOENIX_COLLECTOR_HTTP_ENDPOINT below points it at the jaeger service),
# so gen-AI spans export through a preset-owned provider - the code path
# where trace splits actually happen - with no cloud credentials needed.
callbacks: ["arize_phoenix"]
router_settings:
routing_strategy: simple-shuffle
@ -68,9 +74,14 @@ services:
condition: service_healthy
redis:
condition: service_healthy
jaeger:
condition: service_healthy
env_file: .env
environment:
LITELLM_MASTER_KEY: sk-1234
LITELLM_OTEL_V2: "true"
PHOENIX_COLLECTOR_HTTP_ENDPOINT: http://jaeger:4318/v1/traces
PHOENIX_API_KEY: local-jaeger-noauth
DATABASE_URL: postgresql://litellm:litellm@db:5432/litellm
UI_USERNAME: admin
UI_PASSWORD: sk-1234
@ -114,3 +125,15 @@ services:
interval: 3s
timeout: 3s
retries: 20
# throwaway OTEL trace destination (OTLP ingest on 4318 inside the network,
# query API on host 16686 for test read-back; see E2E_OTEL_QUERY_URL)
jaeger:
image: jaegertracing/all-in-one:1.62.0
ports:
- "16686:16686"
healthcheck:
test: ["CMD", "wget", "-qO-", "http://localhost:14269/"]
interval: 3s
timeout: 3s
retries: 20

View file

@ -24,6 +24,14 @@ CONTROL_PLANE_BASE_URL = os.environ.get(
UI_USERNAME = os.environ.get("E2E_UI_USERNAME", "admin")
UI_PASSWORD = os.environ.get("E2E_UI_PASSWORD", MASTER_KEY)
CHEAP_ANTHROPIC_MODEL = os.environ.get("E2E_CHEAP_ANTHROPIC_MODEL", "claude-haiku-4-5")
CHEAP_OPENAI_MODEL = os.environ.get("E2E_CHEAP_OPENAI_MODEL", "gpt-5.5")
# 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.
OTEL_QUERY_URL = os.environ.get("E2E_OTEL_QUERY_URL", "http://localhost:16686").rstrip("/")
# Writes on the proxy are eventually consistent (e.g. spend rows flush on
# proxy_batch_write_at, ~60s). Read-backs poll to this deadline, never sleep-once.
POLL_TIMEOUT = float(os.environ.get("E2E_POLL_TIMEOUT", "120"))

View file

@ -12,6 +12,7 @@ import time
import warnings
from collections.abc import Callable
from dataclasses import dataclass
from datetime import datetime
from e2e_http import (
NoBody,
@ -50,6 +51,8 @@ from models import (
OcrResponse,
SpendLogRow,
SpendLogs,
SpendLogsPage,
SpendLogsPageParams,
SpendLogsParams,
)
from e2e_config import (
@ -255,6 +258,28 @@ class Gateway:
case _:
return []
def spend_logs_window(self, *, start: datetime, end: datetime) -> list[SpendLogRow]:
def fetch(page: int) -> SpendLogsPage:
return unwrap(
self.transport.get(
"/spend/logs/v2",
headers=self.transport.master,
params=SpendLogsPageParams(
start_date=start.strftime("%Y-%m-%d %H:%M:%S"),
end_date=end.strftime("%Y-%m-%d %H:%M:%S"),
page=page,
page_size=100,
),
response_type=SpendLogsPage,
)
)
first = fetch(1)
return [
*first.data,
*(row for page in range(2, first.total_pages + 1) for row in fetch(page).data),
]
def poll_logs_for_key(
self, key: str, *, min_rows: int = 1, predicate: RowsPredicate | None = None
) -> list[SpendLogRow]:

View file

@ -109,14 +109,16 @@ class StreamingResponse(BaseModel):
"""Raw outcome for calls whose body is provider-native or streamed: status, the
x-litellm-call-id header, the x-litellm-response-cost header (StandardLogging
response_cost), the content-type (which tells streaming `text/event-stream` from
non-streaming `application/json`), and the body. SpendLogs.request_id is the
completion body id, not call_id. Used by passthrough and streaming, where one
validated JSON model does not fit."""
non-streaming `application/json`), the response headers (lowercased names, e.g.
the x-ratelimit-* pacing headers and retry-after on a 429), and the body.
SpendLogs.request_id is the completion body id, not call_id. Used by passthrough
and streaming, where one validated JSON model does not fit."""
status_code: int
call_id: str | None = None # x-litellm-call-id header
response_cost: float | None = None # x-litellm-response-cost header
content_type: str | None = None
headers: dict[str, str] = {}
body: str
chunks: int = 0 # streamed events (0 for non-streaming)
@ -276,12 +278,14 @@ def _streaming_outcome(resp: requests.Response, stream: bool) -> StreamingRespon
call_id = _hdr(resp, "x-litellm-call-id")
response_cost = _parse_response_cost(resp)
content_type = _hdr(resp, "content-type")
headers = {name.lower(): value for name, value in resp.headers.items()}
if not stream or not (200 <= resp.status_code < 300):
return StreamingResponse(
status_code=resp.status_code,
call_id=call_id,
response_cost=response_cost,
content_type=content_type,
headers=headers,
body=resp.text,
)
lines = cast("Iterator[bytes]", resp.iter_lines())
@ -291,6 +295,7 @@ def _streaming_outcome(resp: requests.Response, stream: bool) -> StreamingRespon
call_id=call_id,
response_cost=response_cost,
content_type=content_type,
headers=headers,
body="<streamed>",
chunks=chunks,
)

View file

@ -11,6 +11,7 @@ import os
import pytest
from logging_client import LangfuseCreds, LoggingClient, build_logging_client, load_langfuse_creds
from otel_client import OtelReader, build_otel_reader
def pytest_configure(config: pytest.Config) -> None:
@ -28,6 +29,12 @@ def client() -> LoggingClient:
return build_logging_client()
@pytest.fixture(scope="session")
def otel_reader() -> OtelReader:
"""Read-back client for the compose stack's Jaeger trace destination."""
return build_otel_reader()
@pytest.fixture
def datadog_creds() -> None:
"""Require Datadog shipping credentials. Hard-fail when absent; never skip."""

View file

@ -35,6 +35,7 @@ from e2e_http import (
unwrap,
)
from models import (
AnthropicMessagesBody,
ChatBody,
ChatMessage,
ChatResponse,
@ -75,6 +76,14 @@ WEATHER_TOOL = ChatTool(
)
class ResponsesRequestBody(BaseModel):
"""OpenAI Responses API /v1/responses request (non-streaming)."""
model: str
input: str
max_output_tokens: int
class TeamCallbackBody(BaseModel):
callback_name: Literal["langfuse_otel", "langfuse", "langsmith", "gcs"]
callback_type: Literal["success", "failure", "success_and_failure"]
@ -455,6 +464,32 @@ class LoggingClient:
json=body,
)
def messages_raw(self, key: str, model: str, text: str, *, max_tokens: int = 16) -> StreamingResponse:
"""Non-streaming POST /v1/messages (Anthropic-native body): raw outcome
judged by status/body/headers, for tests that need x-litellm-call-id."""
return self.gateway.transport.send(
"/v1/messages",
headers=self.gateway.transport.bearer(key),
json=AnthropicMessagesBody(
model=model,
max_tokens=max_tokens,
messages=[ChatMessage(role="user", content=text)],
),
)
def responses_raw(
self, key: str, model: str, text: str, *, max_output_tokens: int = 64
) -> StreamingResponse:
"""Non-streaming POST /v1/responses (OpenAI Responses API): raw outcome
judged by status/body/headers, for tests that need x-litellm-call-id.
max_output_tokens caps reasoning-model output cost; a capped response is
still a 200 and still exports the trace."""
return self.gateway.transport.send(
"/v1/responses",
headers=self.gateway.transport.bearer(key),
json=ResponsesRequestBody(model=model, input=text, max_output_tokens=max_output_tokens),
)
def scrape_metrics(self) -> str:
return self.gateway.probe("/metrics", params=NoBody()).body

View file

@ -0,0 +1,138 @@
"""Jaeger read-back for the OTEL trace-completeness tests: typed models over the
Jaeger query API (the destination's own API - completeness is judged on what the
backend actually holds, never on "export succeeded" proxy-side).
Traces are fetched server-side by the ``litellm.call_id`` tag the gen-AI span
carries (the request's x-litellm-call-id response header), so read-back is
immune to the query page filling up with unrelated traffic (background jobs,
other suites sharing the stack). Jaeger returns every span of a matching trace,
so the completeness assertions see the whole tree. A failed query is a hard
failure, never an empty result - an unreachable destination must not read as
"the trace never arrived".
External reads go through ``e2e_http`` (the only module allowed to call
``requests.*``).
"""
from __future__ import annotations
import json
import time
from dataclasses import dataclass
import pytest
from pydantic import BaseModel, ConfigDict, Field
from e2e_config import OTEL_QUERY_URL, POLL_INTERVAL, POLL_TIMEOUT
from e2e_http import URL, NoBody, Success, get
#: OTEL resource service.name the proxy exports under (OTEL_SERVICE_NAME default).
JAEGER_SERVICE = "litellm"
#: Span tag carrying the request's x-litellm-call-id (stamped on the gen-AI span).
CALL_ID_TAG = "litellm.call_id"
class JaegerTag(BaseModel):
model_config = ConfigDict(extra="ignore")
key: str
value: str | int | float | bool | None = None
class JaegerReference(BaseModel):
model_config = ConfigDict(extra="ignore", populate_by_name=True)
ref_type: str = Field(alias="refType")
trace_id: str = Field(alias="traceID")
span_id: str = Field(alias="spanID")
class JaegerSpan(BaseModel):
model_config = ConfigDict(extra="ignore", populate_by_name=True)
span_id: str = Field(alias="spanID")
operation_name: str = Field(alias="operationName")
start_time: int = Field(default=0, alias="startTime")
references: list[JaegerReference] = []
tags: list[JaegerTag] = []
@property
def kind(self) -> str:
for tag in self.tags:
if tag.key == "span.kind":
return str(tag.value)
return ""
class JaegerTrace(BaseModel):
model_config = ConfigDict(extra="ignore", populate_by_name=True)
trace_id: str = Field(alias="traceID")
spans: list[JaegerSpan] = []
def span_names(self) -> list[str]:
return sorted(span.operation_name for span in self.spans)
class JaegerTracesPage(BaseModel):
model_config = ConfigDict(extra="ignore")
data: list[JaegerTrace] = []
class _TracesQuery(BaseModel):
service: str
tags: str
limit: int = 20
lookback: str = "1h"
def _settled(trace: JaegerTrace, names: set[str], prefixes: set[str]) -> bool:
present = set(trace.span_names())
return names.issubset(present) and all(
any(name.startswith(prefix) for name in present) for prefix in prefixes
)
@dataclass(frozen=True, slots=True)
class OtelReader:
query_url: str
def traces_for_call(self, call_id: str) -> list[JaegerTrace]:
"""Every trace holding a span tagged with this call id. Jaeger matches
spans server-side and returns their full traces; more than one hit for
one call IS the split-trace bug, so this never collapses to one."""
result = get(
URL(f"{self.query_url}/api/traces"),
headers=NoBody(),
params=_TracesQuery(service=JAEGER_SERVICE, tags=json.dumps({CALL_ID_TAG: call_id})),
response_type=JaegerTracesPage,
timeout=30.0,
)
match result:
case Success(data=page):
return page.data
case failure:
pytest.fail(f"Jaeger query API at {self.query_url} failed: {failure}")
def poll_traces_for_call(
self, *, call_id: str, settled_names: set[str], settled_prefixes: set[str]
) -> list[JaegerTrace]:
"""Poll until exactly one trace holds the call and it carries every span
name in ``settled_names`` plus at least one name per prefix in
``settled_prefixes`` (spans flush in batches, the cost write lands after
the response), then return the hits. At the deadline the last hits are
returned as-is so the caller's assertions report the real final state -
on a split trace this never settles and the orphan comes back."""
deadline = time.monotonic() + POLL_TIMEOUT
hits: list[JaegerTrace] = []
while time.monotonic() < deadline:
hits = self.traces_for_call(call_id)
if len(hits) == 1 and _settled(hits[0], settled_names, settled_prefixes):
return hits
time.sleep(POLL_INTERVAL)
return hits
def build_otel_reader() -> OtelReader:
return OtelReader(query_url=OTEL_QUERY_URL)

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