scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-17 23:51:30 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 139

Merge remote-tracking branch 'origin/litellm_internal_staging' into litellm_fix_guardrail_info_client_metadata

Browse source
This commit is contained in:
Devin AI 2026-07-18 03:35:35 +00:00
commit c5a5ea1efa
626 changed files with 41527 additions and 15344 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
.github/CODEOWNERS vendored
View file

@ -1,2 +1,3 @@
/ui/ @yuneng-jiang @ryan-crabbe-berri
/litellm/proxy/_experimental/out/ @yuneng-jiang @ryan-crabbe-berri
/ui/litellm-dashboard/src/lib/http/schema.d.ts

2
.github/workflows/test-unit-proxy-db.yml vendored
View file

@ -5,6 +5,8 @@ on:
branches:
- main
- litellm_internal_staging
- litellm_oss_staging
- "litellm_**"
permissions:
contents: read

6
.github/workflows/zizmor.yml vendored
View file

@ -4,7 +4,11 @@ on:
push:
branches: [main, litellm_internal_staging]
pull_request:
branches: [main, litellm_internal_staging]
branches:
- main
- litellm_internal_staging
- litellm_oss_staging
- "litellm_**"
concurrency:
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}

3
.gitignore vendored
View file

@ -15,6 +15,9 @@ litellm/rust_bridge/_native*.so
litellm/rust_bridge/_native*.pyd
litellm-rust/target/
# Python package build output
dist/
bun.lockb
**/.DS_Store
.aider*

1
Dockerfile
View file

@ -114,6 +114,7 @@ COPY --from=builder /app/litellm/proxy/prisma_migration.py /app/litellm/proxy/pr
# working directory on sys.path; litellm/proxy/hooks resolves
# enterprise.enterprise_hooks from it)
COPY --from=builder /app/enterprise /app/enterprise
COPY --from=builder /app/litellm-proxy-extras /app/litellm-proxy-extras
# Prisma binaries live in $HOME/.cache (default prisma-python location),
# which is /root/.cache here. Copy only the Prisma subdirs — copying the
# whole /root/.cache drags in the uv build cache (~660 MB, includes a

BIN
dist/litellm-1.79.1.tar.gz vendored
View file

Binary file not shown.

1
docker/Dockerfile.database
View file

@ -111,6 +111,7 @@ COPY --from=builder /app/litellm/proxy/prisma_migration.py /app/litellm/proxy/pr
# working directory on sys.path; litellm/proxy/hooks resolves
# enterprise.enterprise_hooks from it)
COPY --from=builder /app/enterprise /app/enterprise
COPY --from=builder /app/litellm-proxy-extras /app/litellm-proxy-extras
# Prisma binaries live in $HOME/.cache (default prisma-python location),
# which is /root/.cache here. Copy them from the builder so they survive
# deployments that volume-mount /app/.cache (e.g. readOnlyRootFilesystem

1
docker/Dockerfile.non_root
View file

@ -137,6 +137,7 @@ COPY --from=builder /app/litellm/proxy/prisma_migration.py /app/litellm/proxy/pr
# working directory on sys.path; litellm/proxy/hooks resolves
# enterprise.enterprise_hooks from it)
COPY --from=builder /app/enterprise /app/enterprise
COPY --from=builder /app/litellm-proxy-extras /app/litellm-proxy-extras
COPY --from=builder /app/.cache /app/.cache
COPY --from=builder /var/lib/litellm/ui /var/lib/litellm/ui
COPY --from=builder /var/lib/litellm/assets /var/lib/litellm/assets

2
docker/build_from_pip/Dockerfile.build_from_pip
View file

@ -36,7 +36,7 @@ RUN uv venv --python python && \
"opentelemetry-api==1.28.0" \
"opentelemetry-sdk==1.28.0" \
"opentelemetry-exporter-otlp==1.28.0" \
"ddtrace==2.19.0" \
"ddtrace==4.11.0" \
"sentry-sdk==2.21.0" \
"mangum==0.17.0" \
"azure-ai-contentsafety==1.0.0" \

141
enterprise/litellm_enterprise/enterprise_callbacks/llm_guard.py
View file

@ -7,7 +7,8 @@
# Thank you users! We ❤️ you! - Krrish & Ishaan
## This provides an LLM Guard Integration for content moderation on the proxy
from typing import Literal, Optional
import asyncio
from typing import Optional
import aiohttp
from fastapi import HTTPException
@ -18,7 +19,6 @@ from litellm.integrations.custom_logger import CustomLogger
from litellm.proxy._types import UserAPIKeyAuth
from litellm.secret_managers.main import get_secret_str
from litellm.types.utils import CallTypesLiteral
from litellm.utils import get_formatted_prompt
class _ENTERPRISE_LLMGuard(CustomLogger):
@ -46,45 +46,44 @@ class _ENTERPRISE_LLMGuard(CustomLogger):
except Exception:
pass
async def moderation_check(self, text: str):
async def moderation_check(self, text: str) -> str:
"""
Runs the LLM Guard moderation check on ``text``.
Raises an HTTPException when the content violates the safety policy;
otherwise returns the sanitized prompt from LLM Guard, falling back to
the original text when the API does not provide one.
[TODO] make this more performant for high-throughput scenario
"""
try:
async with aiohttp.ClientSession() as session:
if self.mock_redacted_text is not None:
redacted_text = self.mock_redacted_text
else:
# Make the first request to /analyze
analyze_url = f"{self.llm_guard_api_base}analyze/prompt"
verbose_proxy_logger.debug("Making request to: %s", analyze_url)
analyze_payload = {"prompt": text}
redacted_text = None
if self.mock_redacted_text is not None:
redacted_text = self.mock_redacted_text
else:
analyze_url = f"{self.llm_guard_api_base}analyze/prompt"
verbose_proxy_logger.debug("Making request to: %s", analyze_url)
async with aiohttp.ClientSession() as session:
async with session.post(
analyze_url, json=analyze_payload
analyze_url, json={"prompt": text}
) as response:
redacted_text = await response.json()
verbose_proxy_logger.debug(
f"LLM Guard: Received response - {redacted_text}"
verbose_proxy_logger.debug(
f"LLM Guard: Received response - {redacted_text}"
)
if redacted_text is None:
raise HTTPException(
status_code=500,
detail={
"error": f"Invalid content moderation response: {redacted_text}"
},
)
if redacted_text is not None:
if (
redacted_text.get("is_valid", None) is not None
and redacted_text["is_valid"] is False
):
raise HTTPException(
status_code=400,
detail={"error": "Violated content safety policy"},
)
else:
pass
else:
raise HTTPException(
status_code=500,
detail={
"error": f"Invalid content moderation response: {redacted_text}"
},
)
if redacted_text.get("is_valid", None) is False:
raise HTTPException(
status_code=400,
detail={"error": "Violated content safety policy"},
)
sanitized_prompt = redacted_text.get("sanitized_prompt")
return sanitized_prompt if isinstance(sanitized_prompt, str) else text
except Exception as e:
verbose_proxy_logger.exception(
"litellm.enterprise.enterprise_hooks.llm_guard::moderation_check - Exception occurred - {}".format(
@ -138,23 +137,75 @@ class _ENTERPRISE_LLMGuard(CustomLogger):
return
self.print_verbose("Makes LLM Guard Check")
try:
assert call_type in [
"completion",
"embeddings",
"image_generation",
"moderation",
"audio_transcription",
]
except Exception:
if call_type not in [
"completion",
"embeddings",
"image_generation",
"moderation",
"audio_transcription",
]:
self.print_verbose(
f"Call Type - {call_type}, not in accepted list - ['completion','embeddings','image_generation','moderation','audio_transcription']"
)
return data
formatted_prompt = get_formatted_prompt(data=data, call_type=call_type) # type: ignore
self.print_verbose(f"LLM Guard, formatted_prompt: {formatted_prompt}")
return await self.moderation_check(text=formatted_prompt)
return await self._moderate_request(data=data)
async def _moderate_request(self, data: dict) -> dict:
"""
Sanitizes the request in place using the prompt returned by LLM Guard so
the provider-bound request carries the redacted content, then returns it.
"""
messages = data.get("messages")
if messages is not None:
data["messages"] = list(
await asyncio.gather(
*(self._moderate_message(message) for message in messages)
)
)
return data
input_ = data.get("input")
if input_ is not None:
data["input"] = await self._moderate_input(input_)
return data
prompt = data.get("prompt")
if isinstance(prompt, str):
data["prompt"] = await self.moderation_check(text=prompt)
return data
async def _moderate_message(self, message: dict) -> dict:
content = message.get("content")
if isinstance(content, str):
return {**message, "content": await self.moderation_check(text=content)}
if isinstance(content, list):
return {
**message,
"content": list(
await asyncio.gather(
*(self._moderate_content_part(part) for part in content)
)
),
}
return message
async def _moderate_content_part(self, part: dict) -> dict:
if part.get("type") == "text" and isinstance(part.get("text"), str):
return {**part, "text": await self.moderation_check(text=part["text"])}
return part
async def _moderate_input(self, input_: object) -> object:
if isinstance(input_, str):
return await self.moderation_check(text=input_)
if isinstance(input_, list):
return [
await self.moderation_check(text=item)
if isinstance(item, str)
else item
for item in input_
]
return input_
async def async_post_call_streaming_hook(
self, user_api_key_dict: UserAPIKeyAuth, response: str

8
enterprise/litellm_enterprise/enterprise_callbacks/pagerduty/pagerduty.py
View file

@ -113,6 +113,10 @@ class PagerDutyAlerting(SlackAlerting):
user_api_key_spend=_meta.get("user_api_key_spend"),
user_api_key_max_budget=_meta.get("user_api_key_max_budget"),
user_api_key_budget_reset_at=_meta.get("user_api_key_budget_reset_at"),
user_api_key_user_spend=_meta.get("user_api_key_user_spend"),
user_api_key_user_max_budget=_meta.get("user_api_key_user_max_budget"),
user_api_key_team_spend=_meta.get("user_api_key_team_spend"),
user_api_key_team_max_budget=_meta.get("user_api_key_team_max_budget"),
user_api_key_org_id=_meta.get("user_api_key_org_id"),
user_api_key_org_alias=_meta.get("user_api_key_org_alias"),
user_api_key_team_id=_meta.get("user_api_key_team_id"),
@ -196,6 +200,10 @@ class PagerDutyAlerting(SlackAlerting):
if user_api_key_dict.budget_reset_at
else None
),
user_api_key_user_spend=user_api_key_dict.user_spend,
user_api_key_user_max_budget=user_api_key_dict.user_max_budget,
user_api_key_team_spend=user_api_key_dict.team_spend,
user_api_key_team_max_budget=user_api_key_dict.team_max_budget,
user_api_key_org_id=user_api_key_dict.org_id,
user_api_key_org_alias=user_api_key_dict.organization_alias,
user_api_key_team_id=user_api_key_dict.team_id,

4
enterprise/pyproject.toml
View file

@ -1,6 +1,6 @@
[project]
name = "litellm-enterprise"
version = "0.1.50"
version = "0.1.51"
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.50"
version = "0.1.51"
version_files = [
"pyproject.toml:^version",
"../pyproject.toml:litellm-enterprise==",

21
helm/litellm-helm/README.md
View file

@ -54,6 +54,12 @@ If `db.useStackgresOperator` is used (not yet implemented):
| `pdb.annotations` | Extra metadata annotations to add to the PDB | `{}` |
| `pdb.labels` | Extra metadata labels to add to the PDB | `{}` |
| `billingMetrics.enabled` | Enable enterprise billable-request metering. Requires an enterprise license. | `false` |
| `billingMetrics.endpoint` | Collector that the billable-request counter is pushed to. | `https://telemetry.litellm.ai` |
| `billingMetrics.secretName` | Name of an existing Secret holding the mTLS client certificate, under the keys `tls.crt` and `tls.key`. | `litellm-billing-metrics-mtls` |
| `billingMetrics.caSecretName` | Name of an existing Secret holding a CA bundle under the key `ca.crt`. Only needed for a private or test collector whose server certificate is not on the public web PKI. | `""` |
| `billingMetrics.exportIntervalMs` | How often the counter is pushed, in milliseconds. The proxy defaults to `60000` when unset. | `""` |
#### Example `proxy_config` ConfigMap from values (default):
```
@ -94,6 +100,21 @@ data:
type: Opaque
```
#### Enterprise billable-request metering
Enterprise licenses meter billable requests by pushing a counter to LiteLLM's collector over mutual TLS. The chart does not create the client certificate; it mounts one you already hold, read-only, so the private key is never exposed through the environment. Create the Secret under the name the chart expects, then turn the block on:
```
kubectl create secret tls litellm-billing-metrics-mtls --cert=client.crt --key=client.key
```
```
billingMetrics:
enabled: true
```
Set `billingMetrics.caSecretName` only when the collector is a private or test one whose server certificate is not on the public web PKI; the production collector needs no CA override. The chart fails the render rather than deploying a proxy that silently never exports, so a missing `secretName` or an emptied `endpoint` surfaces at `helm install` time.
### Database Settings
| Name | Description | Value |

47
helm/litellm-helm/templates/_helpers.tpl
View file

@ -50,6 +50,53 @@ app.kubernetes.io/name: {{ include "litellm.name" . }}
app.kubernetes.io/instance: {{ .Release.Name }}
{{- end }}
{{/*
Enterprise billable-request metering. The client certificate identifies the
deployment to LiteLLM's collector, so it is mounted read-only from an existing
Secret rather than passed through the environment.
*/}}
{{- define "litellm.billingMetrics.certDir" -}}/etc/litellm/billing-mtls{{- end -}}
{{- define "litellm.billingMetrics.caDir" -}}/etc/litellm/billing-mtls-ca{{- end -}}
{{- define "litellm.billingMetricsEnv" -}}
- name: LITELLM_BILLING_METRICS_ENDPOINT
value: {{ required "billingMetrics.endpoint is required when billingMetrics.enabled is true" .Values.billingMetrics.endpoint | quote }}
- name: LITELLM_BILLING_METRICS_CLIENT_CERT
value: {{ printf "%s/tls.crt" (include "litellm.billingMetrics.certDir" .) | quote }}
- name: LITELLM_BILLING_METRICS_CLIENT_KEY
value: {{ printf "%s/tls.key" (include "litellm.billingMetrics.certDir" .) | quote }}
{{- if .Values.billingMetrics.caSecretName }}
- name: LITELLM_BILLING_METRICS_CA_CERT
value: {{ printf "%s/ca.crt" (include "litellm.billingMetrics.caDir" .) | quote }}
{{- end }}
{{- with .Values.billingMetrics.exportIntervalMs }}
- name: LITELLM_BILLING_METRICS_EXPORT_INTERVAL_MS
value: {{ . | quote }}
{{- end }}
{{- end -}}
{{- define "litellm.billingMetricsVolumes" -}}
- name: billing-metrics-mtls
secret:
secretName: {{ required "billingMetrics.secretName is required when billingMetrics.enabled is true (an existing Secret with tls.crt and tls.key)" .Values.billingMetrics.secretName }}
{{- if .Values.billingMetrics.caSecretName }}
- name: billing-metrics-mtls-ca
secret:
secretName: {{ .Values.billingMetrics.caSecretName }}
{{- end }}
{{- end -}}
{{- define "litellm.billingMetricsVolumeMounts" -}}
- name: billing-metrics-mtls
mountPath: {{ include "litellm.billingMetrics.certDir" . }}
readOnly: true
{{- if .Values.billingMetrics.caSecretName }}
- name: billing-metrics-mtls-ca
mountPath: {{ include "litellm.billingMetrics.caDir" . }}
readOnly: true
{{- end }}
{{- end -}}
{{/*
Create the name of the service account to use
*/}}

9
helm/litellm-helm/templates/deployment.yaml
View file

@ -142,6 +142,9 @@ spec:
{{- with .Values.extraEnvVars }}
{{- toYaml . | nindent 12 }}
{{- end }}
{{- if .Values.billingMetrics.enabled }}
{{- include "litellm.billingMetricsEnv" . | nindent 12 }}
{{- end }}
{{- if .Values.migrationJob.enabled }}
# Schema updates are owned by the dedicated migrations Job; skip
# the proxy's startup `prisma db push` so N replicas don't race
@ -220,6 +223,9 @@ spec:
- name: npm
mountPath: /.npm
{{- end }}
{{- if .Values.billingMetrics.enabled }}
{{- include "litellm.billingMetricsVolumeMounts" . | nindent 12 }}
{{- end }}
{{- with .Values.volumeMounts }}
{{- toYaml . | nindent 12 }}
{{- end }}
@ -252,6 +258,9 @@ spec:
items:
- key: {{ .Values.proxyConfigMap.key | default "config.yaml" }}
path: "config.yaml"
{{- if .Values.billingMetrics.enabled }}
{{- include "litellm.billingMetricsVolumes" . | nindent 8 }}
{{- end }}
{{- with .Values.volumes }}
{{- toYaml . | nindent 8 }}
{{- end }}

297
helm/litellm-helm/tests/billing_metrics_tests.yaml Normal file
View file

@ -0,0 +1,297 @@
suite: test billingMetrics wiring on the proxy deployment
templates:
- deployment.yaml
- configmap-litellm.yaml
- migrations-job.yaml
tests:
- it: is off by default, adding no env, volume, or mount
template: deployment.yaml
asserts:
- notContains:
path: spec.template.spec.volumes
content:
name: billing-metrics-mtls
secret:
secretName: litellm-billing-metrics-mtls
- notContains:
path: spec.template.spec.containers[0].volumeMounts
content:
name: billing-metrics-mtls
mountPath: /etc/litellm/billing-mtls
readOnly: true
- notContains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_BILLING_METRICS_ENDPOINT
value: https://telemetry.litellm.ai
- it: renders the endpoint and the mounted cert paths when enabled
template: deployment.yaml
set:
billingMetrics:
enabled: true
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_BILLING_METRICS_ENDPOINT
value: https://telemetry.litellm.ai
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_BILLING_METRICS_CLIENT_CERT
value: /etc/litellm/billing-mtls/tls.crt
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_BILLING_METRICS_CLIENT_KEY
value: /etc/litellm/billing-mtls/tls.key
# The conventional Secret name is the default, so enabling the block is enough.
- it: mounts the default cert secret read-only alongside the config volume
template: deployment.yaml
set:
billingMetrics:
enabled: true
asserts:
- contains:
path: spec.template.spec.volumes
content:
name: billing-metrics-mtls
secret:
secretName: litellm-billing-metrics-mtls
- contains:
path: spec.template.spec.containers[0].volumeMounts
content:
name: billing-metrics-mtls
mountPath: /etc/litellm/billing-mtls
readOnly: true
- it: honours a secretName override
template: deployment.yaml
set:
billingMetrics:
enabled: true
secretName: my-billing-mtls
asserts:
- contains:
path: spec.template.spec.volumes
content:
name: billing-metrics-mtls
secret:
secretName: my-billing-mtls
- notContains:
path: spec.template.spec.volumes
content:
name: billing-metrics-mtls
secret:
secretName: litellm-billing-metrics-mtls
- it: honours an endpoint override
template: deployment.yaml
set:
billingMetrics:
enabled: true
endpoint: https://collector.internal:4318
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_BILLING_METRICS_ENDPOINT
value: https://collector.internal:4318
# The production collector presents a public web-PKI certificate, so the CA
# override must stay absent unless a private collector is configured.
- it: omits the CA env, volume, and mount when no caSecretName is set
template: deployment.yaml
set:
billingMetrics:
enabled: true
asserts:
- notContains:
path: spec.template.spec.volumes
content:
name: billing-metrics-mtls-ca
secret:
secretName: billing-ca
- notContains:
path: spec.template.spec.containers[0].volumeMounts
content:
name: billing-metrics-mtls-ca
mountPath: /etc/litellm/billing-mtls-ca
readOnly: true
- notContains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_BILLING_METRICS_CA_CERT
value: /etc/litellm/billing-mtls-ca/ca.crt
- it: mounts the CA secret when caSecretName is set
template: deployment.yaml
set:
billingMetrics:
enabled: true
caSecretName: billing-ca
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_BILLING_METRICS_CA_CERT
value: /etc/litellm/billing-mtls-ca/ca.crt
- contains:
path: spec.template.spec.volumes
content:
name: billing-metrics-mtls-ca
secret:
secretName: billing-ca
- contains:
path: spec.template.spec.containers[0].volumeMounts
content:
name: billing-metrics-mtls-ca
mountPath: /etc/litellm/billing-mtls-ca
readOnly: true
- it: passes the export interval through only when set
template: deployment.yaml
set:
billingMetrics:
enabled: true
exportIntervalMs: 5000
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_BILLING_METRICS_EXPORT_INTERVAL_MS
value: "5000"
- it: omits the export interval when unset
template: deployment.yaml
set:
billingMetrics:
enabled: true
asserts:
- notContains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_BILLING_METRICS_EXPORT_INTERVAL_MS
value: "60000"
# Kubernetes resolves duplicate env names last-wins, so the chart-owned billing
# entries must render after .Values.envVars or a user could silently redirect
# the metering export. The three billing entries are the last ones emitted here
# (migrationJob, which appends DISABLE_SCHEMA_UPDATE, is off for this case).
- it: renders the billing endpoint after envVars so it cannot be shadowed
template: deployment.yaml
set:
migrationJob:
enabled: false
billingMetrics:
enabled: true
envVars:
LITELLM_BILLING_METRICS_ENDPOINT: https://shadowed.example
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_BILLING_METRICS_ENDPOINT
value: https://shadowed.example
- equal:
path: spec.template.spec.containers[0].env[-3]
value:
name: LITELLM_BILLING_METRICS_ENDPOINT
value: https://telemetry.litellm.ai
- equal:
path: spec.template.spec.containers[0].env[-2].name
value: LITELLM_BILLING_METRICS_CLIENT_CERT
- equal:
path: spec.template.spec.containers[0].env[-1].name
value: LITELLM_BILLING_METRICS_CLIENT_KEY
- it: keeps user-supplied volumes and mounts alongside the billing secret
template: deployment.yaml
set:
billingMetrics:
enabled: true
volumes:
- name: custom-callbacks
configMap:
name: my-callbacks
volumeMounts:
- name: custom-callbacks
mountPath: /app/callbacks
asserts:
- contains:
path: spec.template.spec.volumes
content:
name: custom-callbacks
configMap:
name: my-callbacks
- contains:
path: spec.template.spec.volumes
content:
name: billing-metrics-mtls
secret:
secretName: litellm-billing-metrics-mtls
- contains:
path: spec.template.spec.containers[0].volumeMounts
content:
name: custom-callbacks
mountPath: /app/callbacks
- contains:
path: spec.template.spec.containers[0].volumeMounts
content:
name: billing-metrics-mtls
mountPath: /etc/litellm/billing-mtls
readOnly: true
- it: still mounts the proxy config when enabled
template: deployment.yaml
set:
billingMetrics:
enabled: true
asserts:
- contains:
path: spec.template.spec.containers[0].volumeMounts
content:
name: litellm-config
mountPath: /etc/litellm/config.yaml
subPath: config.yaml
# Only the proxy serves billable traffic. The migrations Job must never mount
# the client certificate, and it renders its own env and volumes, so nothing
# stops a future edit from wiring the billing include into it by mistake.
- it: does not touch the migrations job when enabled
template: migrations-job.yaml
set:
billingMetrics:
enabled: true
asserts:
- notContains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_BILLING_METRICS_ENDPOINT
value: https://telemetry.litellm.ai
- notExists:
path: spec.template.spec.containers[0].volumeMounts
- notExists:
path: spec.template.spec.volumes
- it: fails loudly when enabled with an emptied secretName
template: deployment.yaml
set:
billingMetrics:
enabled: true
secretName: ""
asserts:
- failedTemplate:
errorMessage: billingMetrics.secretName is required when billingMetrics.enabled is true (an existing Secret with tls.crt and tls.key)
- it: fails loudly when enabled without an endpoint
template: deployment.yaml
set:
billingMetrics:
enabled: true
endpoint: ""
asserts:
- failedTemplate:
errorMessage: billingMetrics.endpoint is required when billingMetrics.enabled is true

14
helm/litellm-helm/values.yaml
View file

@ -139,6 +139,20 @@ masterkeySecretName: ""
# if set, use this secret key for the master key; otherwise, use the default key
masterkeySecretKey: ""
# Optional: enterprise billable-request metering. When enabled, the proxy counts
# successful requests to inference, MCP, and A2A endpoints and pushes them to
# LiteLLM's collector over mutual TLS. Requires an enterprise license.
# The client certificate identifies the deployment, so it is mounted read-only
# from an existing Secret and never passed through the environment.
billingMetrics:
enabled: false
endpoint: https://telemetry.litellm.ai # collector to push the counter to
secretName: litellm-billing-metrics-mtls # existing Secret holding tls.crt and tls.key
# Only for private or test collectors whose server certificate is not on the
# public web PKI. The production collector needs no CA override.
caSecretName: "" # existing Secret holding ca.crt
exportIntervalMs: "" # push cadence; the proxy defaults to 60000
proxyConfigMap:
# when true, creates a new configmap
create: true

5
helm/litellm/templates/NOTES.txt
View file

@ -46,4 +46,9 @@ Reminders:
- gateway.config.proxy_config (rendered into a ConfigMap and mounted at
/app/config/config.yaml; gateway reads it via
CONFIG_FILE_PATH)
- {component}.pdb.{enabled,minAvailable,maxUnavailable} (per-component PodDisruptionBudget; disabled by
default — with hpa.minReplicas of 1, minAvailable: 1
would block node drains)
- {component}.topologySpreadConstraints (standard k8s list, e.g. spread replicas across
topology.kubernetes.io/zone)
- Enable ingress.enabled=true to dispatch / → ui, gateway data-plane prefixes → gateway, and the catch-all → backend.

97
helm/litellm/templates/_helpers.tpl
View file

@ -34,6 +34,57 @@ app.kubernetes.io/managed-by: {{ .Release.Service }}
helm.sh/chart: {{ printf "%s-%s" .Chart.Name .Chart.Version | replace "+" "_" }}
{{- end -}}
{{/*
Enterprise billable-request metering. Wired into gateway and backend, not the
migrations job. The gateway serves nearly all billable traffic, but the backend
keeps the named-server MCP transport (/{mcp_server_name}/mcp), which writes a
SpendLogs row, so metering only the gateway would silently drop that traffic.
The client certificate identifies the deployment to LiteLLM's collector, so it is
mounted read-only from an existing Secret rather than passed through the
environment.
*/}}
{{- define "litellm.billingMetrics.certDir" -}}/etc/litellm/billing-mtls{{- end -}}
{{- define "litellm.billingMetrics.caDir" -}}/etc/litellm/billing-mtls-ca{{- end -}}
{{- define "litellm.billingMetricsEnv" -}}
- name: LITELLM_BILLING_METRICS_ENDPOINT
value: {{ required "billingMetrics.endpoint is required when billingMetrics.enabled is true" .Values.billingMetrics.endpoint | quote }}
- name: LITELLM_BILLING_METRICS_CLIENT_CERT
value: {{ printf "%s/tls.crt" (include "litellm.billingMetrics.certDir" .) | quote }}
- name: LITELLM_BILLING_METRICS_CLIENT_KEY
value: {{ printf "%s/tls.key" (include "litellm.billingMetrics.certDir" .) | quote }}
{{- if .Values.billingMetrics.caSecretName }}
- name: LITELLM_BILLING_METRICS_CA_CERT
value: {{ printf "%s/ca.crt" (include "litellm.billingMetrics.caDir" .) | quote }}
{{- end }}
{{- with .Values.billingMetrics.exportIntervalMs }}
- name: LITELLM_BILLING_METRICS_EXPORT_INTERVAL_MS
value: {{ . | quote }}
{{- end }}
{{- end -}}
{{- define "litellm.billingMetricsVolumes" -}}
- name: billing-metrics-mtls
secret:
secretName: {{ required "billingMetrics.secretName is required when billingMetrics.enabled is true (an existing Secret with tls.crt and tls.key)" .Values.billingMetrics.secretName }}
{{- if .Values.billingMetrics.caSecretName }}
- name: billing-metrics-mtls-ca
secret:
secretName: {{ .Values.billingMetrics.caSecretName }}
{{- end }}
{{- end -}}
{{- define "litellm.billingMetricsVolumeMounts" -}}
- name: billing-metrics-mtls
mountPath: {{ include "litellm.billingMetrics.certDir" . }}
readOnly: true
{{- if .Values.billingMetrics.caSecretName }}
- name: billing-metrics-mtls-ca
mountPath: {{ include "litellm.billingMetrics.caDir" . }}
readOnly: true
{{- end }}
{{- end -}}
{{/*
Per-component selector labels — used in both Service selectors and Deployment matchLabels.
*/}}
@ -244,6 +295,52 @@ harmless no-op for the Job and authoritative for the app pods.
{{- end }}
{{- end -}}
{{/*
PodDisruptionBudget shared by gateway, backend, and ui.
Invoke with a dict:
(dict "root" $ "component" .Values.gateway "componentName" "gateway"
"fullname" (include "litellm.gateway.fullname" .)
"selectorLabels" (include "litellm.gateway.selectorLabels" .))
Renders nothing unless both the component and its `pdb.enabled` are on.
Only one of minAvailable / maxUnavailable should be set; if both are,
minAvailable wins. If neither is set, falls back to `maxUnavailable: 1` so
an enabled-but-unconfigured PDB still permits node drains.
"Set" means non-nil and non-empty-string, so an explicit 0 (e.g.
`maxUnavailable: 0` to forbid all voluntary disruptions) is honored rather
than silently replaced by the fallback.
*/}}
{{- define "litellm.pdb" -}}
{{- $root := .root -}}
{{- $component := .component -}}
{{- $min := $component.pdb.minAvailable -}}
{{- $max := $component.pdb.maxUnavailable -}}
{{- $minSet := not (or (kindIs "invalid" $min) (eq (printf "%v" $min) "")) -}}
{{- $maxSet := not (or (kindIs "invalid" $max) (eq (printf "%v" $max) "")) -}}
{{- if and $component.enabled $component.pdb $component.pdb.enabled }}
apiVersion: policy/v1
kind: PodDisruptionBudget
metadata:
name: {{ .fullname }}
labels:
{{- include "litellm.commonLabels" $root | nindent 4 }}
app.kubernetes.io/component: {{ .componentName }}
spec:
selector:
matchLabels:
{{- .selectorLabels | nindent 6 }}
{{- if $minSet }}
minAvailable: {{ $min }}
{{- else if $maxSet }}
maxUnavailable: {{ $max }}
{{- else }}
maxUnavailable: 1
{{- end }}
{{- end }}
{{- end -}}
{{/*
Renders `envFrom:` block for a component's `envConfigMaps` / `envSecrets`
lists. Each entry is a resource name; the chart wires the whole ConfigMap /

17
helm/litellm/templates/backend/deployment.yaml
View file

@ -44,14 +44,20 @@ spec:
- name: CONFIG_FILE_PATH
value: /app/config/config.yaml
{{- end }}
{{- if .Values.billingMetrics.enabled }}
{{- include "litellm.billingMetricsEnv" . | nindent 12 }}
{{- end }}
{{- include "litellm.envFrom" .Values.backend | nindent 10 }}
{{- if or .Values.gateway.config.create .Values.backend.volumeMounts }}
{{- if or .Values.gateway.config.create .Values.backend.volumeMounts .Values.billingMetrics.enabled }}
volumeMounts:
{{- if .Values.gateway.config.create }}
- name: gateway-config
mountPath: /app/config/config.yaml
subPath: config.yaml
{{- end }}
{{- if .Values.billingMetrics.enabled }}
{{- include "litellm.billingMetricsVolumeMounts" . | nindent 12 }}
{{- end }}
{{- with .Values.backend.volumeMounts }}
{{- toYaml . | nindent 12 }}
{{- end }}
@ -66,13 +72,16 @@ spec:
{{- end }}
resources:
{{- toYaml .Values.backend.resources | nindent 12 }}
{{- if or .Values.gateway.config.create .Values.backend.volumes }}
{{- if or .Values.gateway.config.create .Values.backend.volumes .Values.billingMetrics.enabled }}
volumes:
{{- if .Values.gateway.config.create }}
- name: gateway-config
configMap:
name: {{ include "litellm.gateway.fullname" . }}-config
{{- end }}
{{- if .Values.billingMetrics.enabled }}
{{- include "litellm.billingMetricsVolumes" . | nindent 8 }}
{{- end }}
{{- with .Values.backend.volumes }}
{{- toYaml . | nindent 8 }}
{{- end }}
@ -89,4 +98,8 @@ spec:
tolerations:
{{- toYaml . | nindent 8 }}
{{- end }}
{{- with .Values.backend.topologySpreadConstraints }}
topologySpreadConstraints:
{{- toYaml . | nindent 8 }}
{{- end }}
{{- end }}

6
helm/litellm/templates/backend/poddisruptionbudget.yaml Normal file
View file

@ -0,0 +1,6 @@
{{- include "litellm.pdb" (dict
"root" $
"component" .Values.backend
"componentName" "backend"
"fullname" (include "litellm.backend.fullname" .)
"selectorLabels" (include "litellm.backend.selectorLabels" .)) }}

17
helm/litellm/templates/gateway/deployment.yaml
View file

@ -46,14 +46,20 @@ spec:
- name: NUM_WORKERS
value: {{ .Values.gateway.numWorkers | quote }}
{{- end }}
{{- if .Values.billingMetrics.enabled }}
{{- include "litellm.billingMetricsEnv" . | nindent 12 }}
{{- end }}
{{- include "litellm.envFrom" .Values.gateway | nindent 10 }}
{{- if or .Values.gateway.config.create .Values.gateway.volumeMounts }}
{{- if or .Values.gateway.config.create .Values.gateway.volumeMounts .Values.billingMetrics.enabled }}
volumeMounts:
{{- if .Values.gateway.config.create }}
- name: gateway-config
mountPath: /app/config/config.yaml
subPath: config.yaml
{{- end }}
{{- if .Values.billingMetrics.enabled }}
{{- include "litellm.billingMetricsVolumeMounts" . | nindent 12 }}
{{- end }}
{{- with .Values.gateway.volumeMounts }}
{{- toYaml . | nindent 12 }}
{{- end }}
@ -68,13 +74,16 @@ spec:
{{- end }}
resources:
{{- toYaml .Values.gateway.resources | nindent 12 }}
{{- if or .Values.gateway.config.create .Values.gateway.volumes }}
{{- if or .Values.gateway.config.create .Values.gateway.volumes .Values.billingMetrics.enabled }}
volumes:
{{- if .Values.gateway.config.create }}
- name: gateway-config
configMap:
name: {{ include "litellm.gateway.fullname" . }}-config
{{- end }}
{{- if .Values.billingMetrics.enabled }}
{{- include "litellm.billingMetricsVolumes" . | nindent 8 }}
{{- end }}
{{- with .Values.gateway.volumes }}
{{- toYaml . | nindent 8 }}
{{- end }}
@ -91,4 +100,8 @@ spec:
tolerations:
{{- toYaml . | nindent 8 }}
{{- end }}
{{- with .Values.gateway.topologySpreadConstraints }}
topologySpreadConstraints:
{{- toYaml . | nindent 8 }}
{{- end }}
{{- end }}

6
helm/litellm/templates/gateway/poddisruptionbudget.yaml Normal file
View file

@ -0,0 +1,6 @@
{{- include "litellm.pdb" (dict
"root" $
"component" .Values.gateway
"componentName" "gateway"
"fullname" (include "litellm.gateway.fullname" .)
"selectorLabels" (include "litellm.gateway.selectorLabels" .)) }}

4
helm/litellm/templates/ui/deployment.yaml
View file

@ -76,4 +76,8 @@ spec:
tolerations:
{{- toYaml . | nindent 8 }}
{{- end }}
{{- with .Values.ui.topologySpreadConstraints }}
topologySpreadConstraints:
{{- toYaml . | nindent 8 }}
{{- end }}
{{- end }}

6
helm/litellm/templates/ui/poddisruptionbudget.yaml Normal file
View file

@ -0,0 +1,6 @@
{{- include "litellm.pdb" (dict
"root" $
"component" .Values.ui
"componentName" "ui"
"fullname" (include "litellm.ui.fullname" .)
"selectorLabels" (include "litellm.ui.selectorLabels" .)) }}

249
helm/litellm/tests/billing_metrics_tests.yaml Normal file
View file

@ -0,0 +1,249 @@
suite: test billingMetrics wiring on gateway and backend
templates:
- gateway/deployment.yaml
- gateway/configmap.yaml
- backend/deployment.yaml
- migrations-job.yaml
values:
- ./values/required.yaml
tests:
- it: is off by default, adding no env, volume, or mount
template: gateway/deployment.yaml
asserts:
- notContains:
path: spec.template.spec.volumes
content:
name: billing-metrics-mtls
secret:
secretName: billing-mtls
- equal:
path: spec.template.spec.containers[0].volumeMounts
value:
- name: gateway-config
mountPath: /app/config/config.yaml
subPath: config.yaml
- notContains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_BILLING_METRICS_ENDPOINT
value: https://telemetry.litellm.ai
- it: renders the endpoint and the mounted cert paths when enabled
template: gateway/deployment.yaml
set:
billingMetrics:
enabled: true
secretName: billing-mtls
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_BILLING_METRICS_ENDPOINT
value: https://telemetry.litellm.ai
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_BILLING_METRICS_CLIENT_CERT
value: /etc/litellm/billing-mtls/tls.crt
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_BILLING_METRICS_CLIENT_KEY
value: /etc/litellm/billing-mtls/tls.key
- it: mounts the cert secret read-only alongside the config volume
template: gateway/deployment.yaml
set:
billingMetrics:
enabled: true
secretName: billing-mtls
asserts:
- contains:
path: spec.template.spec.volumes
content:
name: billing-metrics-mtls
secret:
secretName: billing-mtls
- contains:
path: spec.template.spec.containers[0].volumeMounts
content:
name: billing-metrics-mtls
mountPath: /etc/litellm/billing-mtls
readOnly: true
# The production collector presents a public web-PKI certificate, so the CA
# override must stay absent unless a private collector is configured.
- it: omits the CA env, volume, and mount when no caSecretName is set
template: gateway/deployment.yaml
set:
billingMetrics:
enabled: true
secretName: billing-mtls
asserts:
- notContains:
path: spec.template.spec.volumes
content:
name: billing-metrics-mtls-ca
secret:
secretName: billing-ca
- notContains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_BILLING_METRICS_CA_CERT
value: /etc/litellm/billing-mtls-ca/ca.crt
- it: mounts the CA secret when caSecretName is set
template: gateway/deployment.yaml
set:
billingMetrics:
enabled: true
secretName: billing-mtls
caSecretName: billing-ca
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_BILLING_METRICS_CA_CERT
value: /etc/litellm/billing-mtls-ca/ca.crt
- contains:
path: spec.template.spec.volumes
content:
name: billing-metrics-mtls-ca
secret:
secretName: billing-ca
- contains:
path: spec.template.spec.containers[0].volumeMounts
content:
name: billing-metrics-mtls-ca
mountPath: /etc/litellm/billing-mtls-ca
readOnly: true
- it: passes the export interval through only when set
template: gateway/deployment.yaml
set:
billingMetrics:
enabled: true
secretName: billing-mtls
exportIntervalMs: 5000
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_BILLING_METRICS_EXPORT_INTERVAL_MS
value: "5000"
- it: keeps user-supplied gateway volumes alongside the billing secret
template: gateway/deployment.yaml
set:
billingMetrics:
enabled: true
secretName: billing-mtls
gateway.volumes:
- name: custom-callbacks
configMap:
name: my-callbacks
gateway.volumeMounts:
- name: custom-callbacks
mountPath: /app/callbacks
asserts:
- contains:
path: spec.template.spec.volumes
content:
name: custom-callbacks
configMap:
name: my-callbacks
- contains:
path: spec.template.spec.volumes
content:
name: billing-metrics-mtls
secret:
secretName: billing-mtls
# The backend keeps the named-server MCP transport (/{mcp_server_name}/mcp),
# which writes a SpendLogs row, so it must meter too or that traffic is lost.
- it: meters the backend as well, since it serves the MCP transport
template: backend/deployment.yaml
set:
billingMetrics:
enabled: true
secretName: billing-mtls
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_BILLING_METRICS_ENDPOINT
value: https://telemetry.litellm.ai
- contains:
path: spec.template.spec.containers[0].volumeMounts
content:
name: billing-metrics-mtls
mountPath: /etc/litellm/billing-mtls
readOnly: true
- contains:
path: spec.template.spec.volumes
content:
name: billing-metrics-mtls
secret:
secretName: billing-mtls
- it: leaves the backend alone when metering is off
template: backend/deployment.yaml
asserts:
- notContains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_BILLING_METRICS_ENDPOINT
value: https://telemetry.litellm.ai
# The migrations job runs prisma and serves no traffic; it must never receive
# the client key.
- it: never mounts the billing cert on the migrations job
template: migrations-job.yaml
set:
billingMetrics:
enabled: true
secretName: billing-mtls
asserts:
- notContains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_BILLING_METRICS_ENDPOINT
value: https://telemetry.litellm.ai
- isNull:
path: spec.template.spec.volumes
# The conventional Secret name is the default, so enabling metering needs no
# secretName at all; the guard below only fires on an explicitly blanked one.
- it: uses the conventional secret name by default
template: gateway/deployment.yaml
set:
billingMetrics:
enabled: true
asserts:
- contains:
path: spec.template.spec.volumes
content:
name: billing-metrics-mtls
secret:
secretName: litellm-billing-metrics-mtls
- it: fails loudly when the secretName is explicitly blanked
template: gateway/deployment.yaml
set:
billingMetrics:
enabled: true
secretName: ""
asserts:
- failedTemplate:
errorMessage: billingMetrics.secretName is required when billingMetrics.enabled is true (an existing Secret with tls.crt and tls.key)
- it: fails loudly when enabled without an endpoint
template: gateway/deployment.yaml
set:
billingMetrics:
enabled: true
endpoint: ""
secretName: billing-mtls
asserts:
- failedTemplate:
errorMessage: billingMetrics.endpoint is required when billingMetrics.enabled is true

188
helm/litellm/tests/pdb_topology_spread_tests.yaml Normal file
View file

@ -0,0 +1,188 @@
suite: test pod disruption budgets and topology spread constraints
templates:
- gateway/poddisruptionbudget.yaml
- backend/poddisruptionbudget.yaml
- ui/poddisruptionbudget.yaml
- gateway/deployment.yaml
- gateway/configmap.yaml
- backend/deployment.yaml
- ui/deployment.yaml
values:
- ./values/required.yaml
tests:
- it: renders no PDB by default
templates:
- gateway/poddisruptionbudget.yaml
- backend/poddisruptionbudget.yaml
- ui/poddisruptionbudget.yaml
asserts:
- hasDocuments:
count: 0
- it: gateway PDB uses minAvailable and matches the gateway selector labels
template: gateway/poddisruptionbudget.yaml
set:
gateway.pdb.enabled: true
gateway.pdb.minAvailable: 1
asserts:
- isKind:
of: PodDisruptionBudget
- equal:
path: apiVersion
value: policy/v1
- equal:
path: metadata.name
value: RELEASE-NAME-litellm-gateway
- equal:
path: spec.minAvailable
value: 1
- notExists:
path: spec.maxUnavailable
- equal:
path: spec.selector.matchLabels
value:
app.kubernetes.io/name: litellm
app.kubernetes.io/instance: RELEASE-NAME
app.kubernetes.io/component: gateway
- it: backend PDB uses maxUnavailable when minAvailable is unset
template: backend/poddisruptionbudget.yaml
set:
backend.pdb.enabled: true
backend.pdb.maxUnavailable: 25%
asserts:
- equal:
path: spec.maxUnavailable
value: 25%
- notExists:
path: spec.minAvailable
- equal:
path: spec.selector.matchLabels
value:
app.kubernetes.io/name: litellm
app.kubernetes.io/instance: RELEASE-NAME
app.kubernetes.io/component: backend
- it: minAvailable wins when both minAvailable and maxUnavailable are set
template: gateway/poddisruptionbudget.yaml
set:
gateway.pdb.enabled: true
gateway.pdb.minAvailable: 2
gateway.pdb.maxUnavailable: 1
asserts:
- equal:
path: spec.minAvailable
value: 2
- notExists:
path: spec.maxUnavailable
- it: an explicit maxUnavailable 0 is honored instead of the fallback
template: backend/poddisruptionbudget.yaml
set:
backend.pdb.enabled: true
backend.pdb.maxUnavailable: 0
asserts:
- equal:
path: spec.maxUnavailable
value: 0
- notExists:
path: spec.minAvailable
- it: an explicit minAvailable 0 is honored and beats a set maxUnavailable
template: gateway/poddisruptionbudget.yaml
set:
gateway.pdb.enabled: true
gateway.pdb.minAvailable: 0
gateway.pdb.maxUnavailable: 1
asserts:
- equal:
path: spec.minAvailable
value: 0
- notExists:
path: spec.maxUnavailable
- it: enabled PDB with neither knob set falls back to maxUnavailable 1
template: ui/poddisruptionbudget.yaml
set:
ui.pdb.enabled: true
asserts:
- equal:
path: spec.maxUnavailable
value: 1
- notExists:
path: spec.minAvailable
- equal:
path: spec.selector.matchLabels
value:
app.kubernetes.io/name: litellm
app.kubernetes.io/instance: RELEASE-NAME
app.kubernetes.io/component: ui
- it: renders no PDB for a disabled component even when its pdb is enabled
template: gateway/poddisruptionbudget.yaml
set:
gateway.enabled: false
gateway.pdb.enabled: true
asserts:
- hasDocuments:
count: 0
- it: deployments omit topologySpreadConstraints by default
templates:
- gateway/deployment.yaml
- backend/deployment.yaml
- ui/deployment.yaml
asserts:
- notExists:
path: spec.template.spec.topologySpreadConstraints
- it: gateway deployment renders configured topologySpreadConstraints
template: gateway/deployment.yaml
set:
gateway.topologySpreadConstraints:
- maxSkew: 1
topologyKey: topology.kubernetes.io/zone
whenUnsatisfiable: ScheduleAnyway
labelSelector:
matchLabels:
app.kubernetes.io/component: gateway
asserts:
- equal:
path: spec.template.spec.topologySpreadConstraints
value:
- maxSkew: 1
topologyKey: topology.kubernetes.io/zone
whenUnsatisfiable: ScheduleAnyway
labelSelector:
matchLabels:
app.kubernetes.io/component: gateway
- it: backend deployment renders configured topologySpreadConstraints
template: backend/deployment.yaml
set:
backend.topologySpreadConstraints:
- maxSkew: 1
topologyKey: kubernetes.io/hostname
whenUnsatisfiable: DoNotSchedule
labelSelector:
matchLabels:
app.kubernetes.io/component: backend
asserts:
- equal:
path: spec.template.spec.topologySpreadConstraints[0].topologyKey
value: kubernetes.io/hostname
- equal:
path: spec.template.spec.topologySpreadConstraints[0].whenUnsatisfiable
value: DoNotSchedule
- it: ui deployment renders configured topologySpreadConstraints
template: ui/deployment.yaml
set:
ui.topologySpreadConstraints:
- maxSkew: 1
topologyKey: topology.kubernetes.io/zone
whenUnsatisfiable: ScheduleAnyway
asserts:
- equal:
path: spec.template.spec.topologySpreadConstraints[0].topologyKey
value: topology.kubernetes.io/zone

51
helm/litellm/values.yaml
View file

@ -73,6 +73,25 @@ masterKey:
secretName: litellm-master-key-secret # name of a Secret containing the master key
secretKey: master-key
# Optional: enterprise billable-request metering. When enabled, the gateway and
# backend count successful requests to inference, MCP, and A2A endpoints and push
# them to LiteLLM's collector over mutual TLS. Both components serve billable
# routes: the backend keeps the named-server MCP transport. Requires an
# enterprise license. The client certificate identifies the deployment, so it is
# mounted read-only from an existing Secret and never passed through the env.
billingMetrics:
enabled: false
endpoint: https://telemetry.litellm.ai # collector to push the counter to
# An existing Secret holding the client certificate under tls.crt and its key
# under tls.key, usually created from the onboarding artifact. The default is
# the conventional name, so the common path is to create that Secret and set
# enabled: true. Override only if yours is named differently.
secretName: litellm-billing-metrics-mtls
# Only for private or test collectors whose server certificate is not on the
# public web PKI. The production collector needs no CA override.
caSecretName: "" # existing Secret holding ca.crt
exportIntervalMs: "" # push cadence; the proxy defaults to 60000
# External Postgres connection.
database:
writer:
@ -171,10 +190,28 @@ gateway:
maxReplicas: 10
targetCPUUtilizationPercentage: 70
targetMemoryUtilizationPercentage: 80
# PodDisruptionBudget for the gateway pods. Set exactly one of
# `minAvailable` / `maxUnavailable` (minAvailable wins if both are set;
# enabling without either falls back to `maxUnavailable: 1`). Disabled by
# default: with the default hpa.minReplicas of 1, a `minAvailable: 1` PDB
# would block node drains entirely.
pdb:
enabled: false
minAvailable: ""
maxUnavailable: ""
podAnnotations: {}
nodeSelector: {}
tolerations: []
affinity: {}
# Standard k8s topologySpreadConstraints for the gateway pods, e.g. to
# spread replicas across zones:
# - maxSkew: 1
# topologyKey: topology.kubernetes.io/zone
# whenUnsatisfiable: ScheduleAnyway
# labelSelector:
# matchLabels:
# app.kubernetes.io/component: gateway
topologySpreadConstraints: []
# ---------- backend (UI / management API) ----------
backend:
@ -214,10 +251,17 @@ backend:
minReplicas: 1
maxReplicas: 4
targetCPUUtilizationPercentage: 70
# Same shape as gateway.pdb.
pdb:
enabled: false
minAvailable: ""
maxUnavailable: ""
podAnnotations: {}
nodeSelector: {}
tolerations: []
affinity: {}
# Same shape as gateway.topologySpreadConstraints.
topologySpreadConstraints: []
# ---------- ui (Next.js static dashboard) ----------
ui:
@ -260,7 +304,14 @@ ui:
minReplicas: 1
maxReplicas: 3
targetCPUUtilizationPercentage: 80
# Same shape as gateway.pdb.
pdb:
enabled: false
minAvailable: ""
maxUnavailable: ""
podAnnotations: {}
nodeSelector: {}
tolerations: []
affinity: {}
# Same shape as gateway.topologySpreadConstraints.
topologySpreadConstraints: []

2
litellm-proxy-extras/litellm_proxy_extras/migrations/20260715000000_add_issuer_to_mcp_server_table/migration.sql Normal file
View file

@ -0,0 +1,2 @@
-- AlterTable
ALTER TABLE "LiteLLM_MCPServerTable" ADD COLUMN IF NOT EXISTS "issuer" TEXT;

1
litellm-proxy-extras/litellm_proxy_extras/schema.prisma
View file

@ -325,6 +325,7 @@ model LiteLLM_MCPServerTable {
command String?
args String[] @default([])
env Json? @default("{}")
issuer String?
authorization_url String?
token_url String?
registration_url String?

4
litellm-proxy-extras/pyproject.toml
View file

@ -1,6 +1,6 @@
[project]
name = "litellm-proxy-extras"
version = "0.4.77"
version = "0.4.78"
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.77"
version = "0.4.78"
version_files = [
"pyproject.toml:^version",
"../pyproject.toml:litellm-proxy-extras==",

94
litellm-rust/Cargo.lock generated
View file

@ -19,12 +19,6 @@ version = "1.1.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "1505bd5d3d116872e7271a6d4e16d81d0c8570876c8de68093a09ac269d8aac0"
[[package]]
name = "autocfg"
version = "1.5.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "f2032f911046de80f0a198e0901378627c33f59ea0ac00e363d481118bd70a53"
[[package]]
name = "axum"
version = "0.7.9"
@ -233,21 +227,6 @@ dependencies = [
"percent-encoding",
]
[[package]]
name = "futures"
version = "0.3.32"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "8b147ee9d1f6d097cef9ce628cd2ee62288d963e16fb287bd9286455b241382d"
dependencies = [
"futures-channel",
"futures-core",
"futures-executor",
"futures-io",
"futures-sink",
"futures-task",
"futures-util",
]
[[package]]
name = "futures-channel"
version = "0.3.32"
@ -264,17 +243,6 @@ version = "0.3.32"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "7e3450815272ef58cec6d564423f6e755e25379b217b0bc688e295ba24df6b1d"
[[package]]
name = "futures-executor"
version = "0.3.32"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "baf29c38818342a3b26b5b923639e7b1f4a61fc5e76102d4b1981c6dc7a7579d"
dependencies = [
"futures-core",
"futures-task",
"futures-util",
]
[[package]]
name = "futures-io"
version = "0.3.32"
@ -310,7 +278,6 @@ version = "0.3.32"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "389ca41296e6190b48053de0321d02a77f32f8a5d2461dd38762c0593805c6d6"
dependencies = [
"futures-channel",
"futures-core",
"futures-io",
"futures-macro",
@ -608,15 +575,6 @@ dependencies = [
"hashbrown",
]
[[package]]
name = "indoc"
version = "2.0.7"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "79cf5c93f93228cf8efb3ba362535fb11199ac548a09ce117c9b1adc3030d706"
dependencies = [
"rustversion",
]
[[package]]
name = "ipnet"
version = "2.12.0"
@ -718,15 +676,6 @@ version = "2.8.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "88904434abc2901f197fe8cc55f0445e7ded921dba5911dad2e2b39b48e663c4"
[[package]]
name = "memoffset"
version = "0.9.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "488016bfae457b036d996092f6cb448677611ce4449e970ceaf42695203f218a"
dependencies = [
"autocfg",
]
[[package]]
name = "mime"
version = "0.3.17"
@ -803,29 +752,26 @@ dependencies = [
[[package]]
name = "pyo3"
version = "0.23.5"
version = "0.29.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "7778bffd85cf38175ac1f545509665d0b9b92a198ca7941f131f85f7a4f9a872"
checksum = "cd274650b21d4bfc26a0a47587962c1edb425f69287324355cd040c3ea66071c"
dependencies = [
"cfg-if",
"indoc",
"libc",
"memoffset",
"once_cell",
"portable-atomic",
"pyo3-build-config",
"pyo3-ffi",
"pyo3-macros",
"unindent",
]
[[package]]
name = "pyo3-async-runtimes"
version = "0.23.0"
version = "0.29.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "977dc837525cfd22919ba6a831413854beb7c99a256c03bf8624ad707e45810e"
checksum = "b3ef68daa7316a3fac65e5e18b2203f010346de1c1c53456811a2624673ab046"
dependencies = [
"futures",
"futures-channel",
"futures-util",
"once_cell",
"pin-project-lite",
"pyo3",
@ -834,19 +780,18 @@ dependencies = [
[[package]]
name = "pyo3-build-config"
version = "0.23.5"
version = "0.29.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "94f6cbe86ef3bf18998d9df6e0f3fc1050a8c5efa409bf712e661a4366e010fb"
checksum = "c5e2a7d2f0d013342f295c048ad19237add5154a55b1c5a254c0ec93d4109078"
dependencies = [
"once_cell",
"target-lexicon",
]
[[package]]
name = "pyo3-ffi"
version = "0.23.5"
version = "0.29.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "e9f1b4c431c0bb1c8fb0a338709859eed0d030ff6daa34368d3b152a63dfdd8d"
checksum = "ca85c467da1bbc8d866eea5deff9cf29ea5f7785054a17da36e65bda9c05845b"
dependencies = [
"libc",
"pyo3-build-config",
@ -854,9 +799,9 @@ dependencies = [
[[package]]
name = "pyo3-macros"
version = "0.23.5"
version = "0.29.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "fbc2201328f63c4710f68abdf653c89d8dbc2858b88c5d88b0ff38a75288a9da"
checksum = "9ac53762fd065daa3194dd09337a38bd793a188100fd1a9304c4ab312d901771"
dependencies = [
"proc-macro2",
"pyo3-macros-backend",
@ -866,13 +811,12 @@ dependencies = [
[[package]]
name = "pyo3-macros-backend"
version = "0.23.5"
version = "0.29.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "fca6726ad0f3da9c9de093d6f116a93c1a38e417ed73bf138472cf4064f72028"
checksum = "4ca3a1557399783172dc5bf39cfca835157732532cba56b71d2292161e53b362"
dependencies = [
"heck",
"proc-macro2",
"pyo3-build-config",
"quote",
"syn",
]
@ -1321,9 +1265,9 @@ dependencies = [
[[package]]
name = "target-lexicon"
version = "0.12.16"
version = "0.13.5"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "61c41af27dd6d1e27b1b16b489db798443478cef1f06a660c96db617ba5de3b1"
checksum = "adb6935a6f5c20170eeceb1a3835a49e12e19d792f6dd344ccc76a985ca5a6ca"
[[package]]
name = "thiserror"
@ -1559,12 +1503,6 @@ version = "1.0.24"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "e6e4313cd5fcd3dad5cafa179702e2b244f760991f45397d14d4ebf38247da75"
[[package]]
name = "unindent"
version = "0.2.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "7264e107f553ccae879d21fbea1d6724ac785e8c3bfc762137959b5802826ef3"
[[package]]
name = "untrusted"
version = "0.9.0"

4
litellm-rust/Cargo.toml
View file

@ -15,8 +15,8 @@ repository = "https://github.com/BerriAI/litellm"
litellm-core = { path = "crates/core" }
litellm-ai-gateway = { path = "crates/ai-gateway", default-features = false }
axum = "0.7"
pyo3 = "0.23.5"
pyo3-async-runtimes = { version = "0.23.0", features = ["tokio-runtime"] }
pyo3 = "0.29.0"
pyo3-async-runtimes = { version = "0.29.0", features = ["tokio-runtime"] }
rand = "0.8"
reqwest = { version = "0.12", default-features = false, features = ["blocking", "json", "rustls-tls", "http2", "stream"] }
serde = { version = "1.0", features = ["derive"] }

2
litellm-rust/crates/ai-gateway/src/python/config.rs
View file

@ -17,7 +17,7 @@ use crate::gil;
/// Load the router's `model_list` from `config_path` via the Python reader.
pub fn load_router_from_config(config_path: &str) -> CoreResult<Router> {
gil::record_acquisition();
Python::with_gil(|py| {
Python::attach(|py| {
let model_list = py
.import("litellm.proxy.read_model_list")
.and_then(|module| module.getattr("read_model_list"))

4
litellm-rust/crates/python-bridge/src/gil.rs
View file

@ -2,7 +2,7 @@
//!
//! A single chokepoint for releasing the GIL around blocking work. Every
//! blocking call in the bridge goes through [`release_gil`] instead of calling
//! `Python::allow_threads` directly, so the release count stays accurate and we
//! `Python::detach` directly, so the release count stays accurate and we
//! have one place to extend later (timing histograms, per-call labels, etc.).
use std::sync::atomic::{AtomicU64, Ordering};
@ -23,7 +23,7 @@ where
T: Send,
{
GIL_RELEASES.fetch_add(1, Ordering::Relaxed);
py.allow_threads(f)
py.detach(f)
}
/// Total GIL releases performed by the bridge so far.

2
litellm-rust/crates/python-bridge/src/lib.rs
View file

@ -167,7 +167,7 @@ fn aocr(
.await
.map_err(core_error_to_pyerr)?;
Python::with_gil(|py| json_to_py(py, value))
Python::attach(|py| json_to_py(py, value))
})
}

5
litellm/__init__.py
View file

@ -315,6 +315,11 @@ disable_token_counter: bool = False
disable_add_transform_inline_image_block: bool = False
disable_add_user_agent_to_request_tags: bool = False
disable_anthropic_gemini_context_caching_transform: bool = False
enable_anthropic_prompt_caching: bool = os.getenv("LITELLM_ENABLE_ANTHROPIC_PROMPT_CACHING", "false").lower() == "true"
_anthropic_prompt_caching_ttl_env: Optional[str] = os.getenv("LITELLM_ANTHROPIC_PROMPT_CACHING_TTL")
anthropic_prompt_caching_ttl: Optional[Literal["5m", "1h"]] = (
"1h" if _anthropic_prompt_caching_ttl_env == "1h" else "5m" if _anthropic_prompt_caching_ttl_env == "5m" else None
)
disable_vertex_batch_output_transformation: bool = False
extra_spend_tag_headers: Optional[List[str]] = None
in_memory_llm_clients_cache: "LLMClientCache"

6
litellm/_redis.py
View file

@ -688,10 +688,8 @@ def get_redis_connection_pool(
elif redis_connect_func and hasattr(redis_connect_func, "_gcp_service_account"):
redis_kwargs["credential_provider"] = GCPIAMCredentialProvider(redis_connect_func._gcp_service_account)
connection_class = async_redis.Connection
if redis_kwargs.pop("ssl", False):
connection_class = async_redis.SSLConnection
redis_kwargs["connection_class"] = connection_class
if redis_kwargs.pop("ssl", None):
redis_kwargs["connection_class"] = async_redis.SSLConnection
return async_redis.BlockingConnectionPool(timeout=REDIS_CONNECTION_POOL_TIMEOUT, **redis_kwargs)

22
litellm/caching/caching_handler.py
View file

@ -85,6 +85,22 @@ class CachingHandlerResponse(BaseModel):
in_memory_cache_obj = InMemoryCache()
def _drop_logging_obj_from_kwargs(request_kwargs: dict[str, object]) -> dict[str, object]:
"""
The caching handler is stored on the Logging object
(``logging_obj._llm_caching_handler``), so keeping ``litellm_logging_obj``
inside ``request_kwargs`` closes a reference cycle
(Logging -> LLMCachingHandler -> kwargs -> Logging) that keeps the full
request payload (messages included) alive until a generational GC pass
instead of being freed by refcount when the request ends. Nothing in the
caching layer reads the logging object from these kwargs; cache-key
generation ignores litellm-internal params.
"""
if "litellm_logging_obj" not in request_kwargs:
return request_kwargs
return {k: v for k, v in request_kwargs.items() if k != "litellm_logging_obj"}
def _is_chat_completion_cached_dict(cached_result: dict) -> bool:
cached_id = cached_result.get("id")
if isinstance(cached_id, str) and cached_id.startswith("chatcmpl"):
@ -118,7 +134,7 @@ class LLMCachingHandler:
self.async_streaming_chunks: List[ModelResponse] = []
self.sync_streaming_chunks: List[ModelResponse] = []
self.request_kwargs = request_kwargs
self.request_kwargs = _drop_logging_obj_from_kwargs(request_kwargs)
self.preset_cache_key: Optional[str] = None
self.original_function = original_function
self.start_time = start_time
@ -297,7 +313,7 @@ class LLMCachingHandler:
new_kwargs.pop("metadata", None)
if new_kwargs.get("stream") is True and "cache_key" not in new_kwargs:
new_kwargs["cache_key"] = litellm.cache.get_cache_key(**new_kwargs)
self.request_kwargs = new_kwargs
self.request_kwargs = _drop_logging_obj_from_kwargs(new_kwargs)
print_verbose("Checking Sync Cache")
cached_result = litellm.cache.get_cache(**new_kwargs)
if cached_result is not None:
@ -693,7 +709,7 @@ class LLMCachingHandler:
new_kwargs.pop("metadata", None)
if new_kwargs.get("stream") is True and "cache_key" not in new_kwargs:
new_kwargs["cache_key"] = litellm.cache.get_cache_key(**new_kwargs)
self.request_kwargs = new_kwargs
self.request_kwargs = _drop_logging_obj_from_kwargs(new_kwargs)
cached_result: Optional[Any] = None
if call_type == CallTypes.aembedding.value:
if isinstance(new_kwargs["input"], str):

18
litellm/caching/dual_cache.py
View file

@ -103,6 +103,18 @@ class DualCache(BaseCache):
if default_redis_ttl is not None:
self.default_redis_ttl = default_redis_ttl
def _backfill_kwargs(self, kwargs: "dict[str, object]") -> "dict[str, object]":
"""
Kwargs for writing a Redis read result into the in-memory tier.
Applies ``default_in_memory_ttl`` exactly like the write paths do;
without it, backfilled entries fall to ``InMemoryCache``'s own default
TTL and can outlive the TTL this cache was configured with.
"""
if "ttl" not in kwargs and self.default_in_memory_ttl is not None:
return {**kwargs, "ttl": self.default_in_memory_ttl}
return kwargs
def set_cache(self, key, value, local_only: bool = False, **kwargs):
# Update both Redis and in-memory cache
try:
@ -160,7 +172,7 @@ class DualCache(BaseCache):
if redis_result is not None:
# Update in-memory cache with the value from Redis
self.in_memory_cache.set_cache(key, redis_result, **kwargs)
self.in_memory_cache.set_cache(key, redis_result, **self._backfill_kwargs(kwargs))
result = redis_result
@ -226,7 +238,7 @@ class DualCache(BaseCache):
if redis_result is not None:
# Update in-memory cache with the value from Redis
await self.in_memory_cache.async_set_cache(key, redis_result, **kwargs)
await self.in_memory_cache.async_set_cache(key, redis_result, **self._backfill_kwargs(kwargs))
result = redis_result
@ -318,7 +330,7 @@ class DualCache(BaseCache):
result[key_to_index[key]] = value
if value is not None and self.in_memory_cache is not None:
await self.in_memory_cache.async_set_cache(key, value, **kwargs)
await self.in_memory_cache.async_set_cache(key, value, **self._backfill_kwargs(kwargs))
return result
except Exception:

18
litellm/constants.py
View file

@ -2,7 +2,7 @@ import os
import sys
from typing import List, Literal, Optional
from litellm.litellm_core_utils.env_utils import get_env_int
from litellm.litellm_core_utils.env_utils import get_env_int, get_env_int_or_none
DEFAULT_HEALTH_CHECK_PROMPT = str(os.getenv("DEFAULT_HEALTH_CHECK_PROMPT", "test from litellm"))
AZURE_DEFAULT_RESPONSES_API_VERSION = str(os.getenv("AZURE_DEFAULT_RESPONSES_API_VERSION", "preview"))
@ -269,9 +269,18 @@ TOOL_POLICY_CACHE_TTL_SECONDS = int(os.getenv("TOOL_POLICY_CACHE_TTL_SECONDS", 6
MAX_SIZE_IN_MEMORY_QUEUE = int(os.getenv("MAX_SIZE_IN_MEMORY_QUEUE", int(LITELLM_ASYNCIO_QUEUE_MAXSIZE * 0.8)))
MAX_IN_MEMORY_QUEUE_FLUSH_COUNT = int(os.getenv("MAX_IN_MEMORY_QUEUE_FLUSH_COUNT", 1000))
###############################################################################################
MINIMUM_PROMPT_CACHE_TOKEN_COUNT = int(
os.getenv("MINIMUM_PROMPT_CACHE_TOKEN_COUNT", 1024)
) # minimum number of tokens to cache a prompt by Anthropic
# Providers will not cache a prefix below a minimum size. That minimum is per-model, not global:
# Anthropic's ranges from 512 to 4096 depending on the model, and can differ per platform for the
# same model. The real minimum is resolved from `prompt_cache_min_tokens` in the model cost map;
# this value is only the fallback for models the cost map has no entry for, and doubles as a global
# escape hatch when `MINIMUM_PROMPT_CACHE_TOKEN_COUNT` is explicitly set.
MINIMUM_PROMPT_CACHE_TOKEN_COUNT_OVERRIDE: int | None = get_env_int_or_none("MINIMUM_PROMPT_CACHE_TOKEN_COUNT")
DEFAULT_MINIMUM_PROMPT_CACHE_TOKEN_COUNT = 1024
MINIMUM_PROMPT_CACHE_TOKEN_COUNT = (
MINIMUM_PROMPT_CACHE_TOKEN_COUNT_OVERRIDE
if MINIMUM_PROMPT_CACHE_TOKEN_COUNT_OVERRIDE is not None
else DEFAULT_MINIMUM_PROMPT_CACHE_TOKEN_COUNT
)
DEFAULT_TRIM_RATIO = float(
os.getenv("DEFAULT_TRIM_RATIO", 0.75)
) # default ratio of tokens to trim from the end of a prompt
@ -1496,6 +1505,7 @@ MAX_TEAM_LIST_LIMIT = int(os.getenv("MAX_TEAM_LIST_LIMIT", 20))
MAX_POLICY_ESTIMATE_IMPACT_ROWS = int(os.getenv("MAX_POLICY_ESTIMATE_IMPACT_ROWS", 1000))
DEFAULT_PROMPT_INJECTION_SIMILARITY_THRESHOLD = float(os.getenv("DEFAULT_PROMPT_INJECTION_SIMILARITY_THRESHOLD", 0.7))
LENGTH_OF_LITELLM_GENERATED_KEY = int(os.getenv("LENGTH_OF_LITELLM_GENERATED_KEY", 16))
MINIMUM_CUSTOM_KEY_LENGTH = int(os.getenv("MINIMUM_CUSTOM_KEY_LENGTH", 16))
SECRET_MANAGER_REFRESH_INTERVAL = int(os.getenv("SECRET_MANAGER_REFRESH_INTERVAL", 86400))
LITELLM_SETTINGS_SAFE_DB_OVERRIDES = [
"default_internal_user_params",

4
litellm/exceptions.py
View file

@ -966,11 +966,15 @@ class BudgetExceededError(Exception):
max_budget: float,
message: Optional[str] = None,
llm_provider: Optional[str] = None,
entity_type: Optional[str] = None,
entity_id: Optional[str] = None,
):
self.current_cost = current_cost
self.max_budget = max_budget
self.status_code = 429
self.llm_provider = llm_provider or ""
self.entity_type = entity_type
self.entity_id = entity_id
# Surface unified rate-limit fields without joining the RateLimitError
# hierarchy so existing `except BudgetExceededError:` handlers keep
# working; custom callbacks reading StandardLoggingPayload pick these

12
litellm/google_genai/main.py
View file

@ -17,6 +17,7 @@ from litellm.llms.base_llm.google_genai.transformation import (
)
from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler
from litellm.types.router import GenericLiteLLMParams
from litellm.types.utils import CallTypes
from litellm.utils import ProviderConfigManager, client
if TYPE_CHECKING:
@ -39,6 +40,11 @@ base_llm_http_handler = BaseLLMHTTPHandler()
#################################################
def _mark_async_entrypoint(logging_obj: LiteLLMLoggingObj | None, marker: str, is_async: bool) -> None:
if logging_obj is not None:
logging_obj.model_call_details.setdefault("litellm_params", {})[marker] = is_async
class GenerateContentSetupResult(BaseModel):
"""Internal Type - Result of setting up a generate content call"""
@ -315,6 +321,8 @@ def generate_content(
try:
_is_async = kwargs.pop("agenerate_content", False)
_mark_async_entrypoint(kwargs.get("litellm_logging_obj"), CallTypes.agenerate_content.value, _is_async)
# Handle generationConfig parameter from kwargs for backward compatibility
if "generationConfig" in kwargs and config is None:
config = kwargs.pop("generationConfig")
@ -403,6 +411,8 @@ async def agenerate_content_stream(
try:
kwargs["agenerate_content_stream"] = True
_mark_async_entrypoint(kwargs.get("litellm_logging_obj"), CallTypes.agenerate_content_stream.value, True)
# Handle generationConfig parameter from kwargs for backward compatibility
if "generationConfig" in kwargs and config is None:
config = kwargs.pop("generationConfig")
@ -497,6 +507,8 @@ def generate_content_stream(
# Remove any async-related flags since this is the sync function
_is_async = kwargs.pop("agenerate_content_stream", False)
_mark_async_entrypoint(kwargs.get("litellm_logging_obj"), CallTypes.agenerate_content_stream.value, _is_async)
# Handle generationConfig parameter from kwargs for backward compatibility
if "generationConfig" in kwargs and config is None:
config = kwargs.pop("generationConfig")

132
litellm/integrations/anthropic_cache_control_hook.py
View file

@ -296,18 +296,148 @@ class AnthropicCacheControlHook(CustomPromptManagement):
return processed_messages, processed_system, remaining_points
@staticmethod
def _default_control() -> ChatCompletionCachedContent:
"""Build the cache_control block for auto-injected breakpoints.
Defaults to Anthropic's 5-minute ephemeral cache; honors the optional
``litellm.anthropic_prompt_caching_ttl`` override ("5m" or "1h").
"""
import litellm
ttl = litellm.anthropic_prompt_caching_ttl
if ttl == "5m" or ttl == "1h":
return ChatCompletionCachedContent(type="ephemeral", ttl=ttl)
return ChatCompletionCachedContent(type="ephemeral")
@staticmethod
def _request_has_cache_control(
messages: list[AllMessageValues],
system: str | list | None,
tools: list | None = None,
) -> bool:
"""Return True if the request already carries any client-supplied cache_control.
When the client (e.g. Claude Code) already marks its own breakpoints we
stand down entirely rather than add more, per the auto-caching contract.
Tools count: they are a breakpoint the client can mark, they count toward
the provider's four-block limit, and caching only the tool definitions is
a common pattern, so injecting alongside them can exceed the cap.
"""
if any(AnthropicCacheControlHook._count_cache_control_blocks(msg) for msg in messages):
return True
if isinstance(system, list):
if any(isinstance(block, dict) and block.get("cache_control") is not None for block in system):
return True
if tools is not None:
return any(isinstance(tool, dict) and tool.get("cache_control") is not None for tool in tools)
return False
@staticmethod
def get_default_injection_points(
messages: list[AllMessageValues],
system: str | list | None,
model: str,
custom_llm_provider: str | None,
tools: list | None = None,
) -> list[CacheControlInjectionPoint]:
"""Default breakpoints when ``litellm.enable_anthropic_prompt_caching`` is on.
Caches the system prompt and the trailing turn, so the stable prefix
(system + tools + history) is reused while the breakpoint advances with
the conversation. Returns [] (stand down) when the flag is off, the
provider does not consume cache_control breakpoints (only anthropic /
bedrock do), the model lacks prompt-caching support, or the request
already carries client-supplied cache_control.
"""
import litellm
if litellm.enable_anthropic_prompt_caching is not True:
return []
provider = custom_llm_provider
if provider is None:
from litellm.litellm_core_utils.get_llm_provider_logic import (
get_llm_provider,
)
try:
_, provider, _, _ = get_llm_provider(model=model)
except Exception: # noqa: BLE001 # unroutable model must never block the call, just skip auto-caching
return []
if provider not in ("anthropic", "bedrock"):
return []
from litellm.utils import supports_prompt_caching
if not supports_prompt_caching(model=model, custom_llm_provider=provider):
return []
if AnthropicCacheControlHook._request_has_cache_control(messages, system, tools):
return []
control = AnthropicCacheControlHook._default_control()
points: list[CacheControlInjectionPoint] = [
CacheControlMessageInjectionPoint(location="message", role="system", index=None, control=control),
CacheControlMessageInjectionPoint(location="message", role=None, index=-1, control=control),
]
return points
@staticmethod
def maybe_seed_default_injection_points(
non_default_params: dict[str, Any],
messages: list[AllMessageValues],
model: str,
custom_llm_provider: str | None,
tools: list | None = None,
) -> None:
"""For /chat/completions: add default injection points to the request params.
No-op when injection points are already configured (explicit config wins).
Seeding the param lets the existing prompt-management gate and the
AnthropicCacheControlHook run unchanged.
"""
if non_default_params.get("cache_control_injection_points"):
return
points = AnthropicCacheControlHook.get_default_injection_points(
messages=messages,
system=None,
model=model,
custom_llm_provider=custom_llm_provider,
tools=tools,
)
if points:
non_default_params["cache_control_injection_points"] = points
@staticmethod
def maybe_inject_cache_control(
messages: List[Dict],
system: str | list | None,
kwargs: Dict[str, Any],
model: str | None = None,
custom_llm_provider: str | None = None,
tools: list[dict] | None = None,
) -> Tuple[List[Dict], str | list | None]:
"""Extract cache_control_injection_points from kwargs and apply if present.
When none are configured but ``litellm.enable_anthropic_prompt_caching``
is on, synthesize default breakpoints for the native /v1/messages path.
Pops the key from kwargs; if remaining (non-message) points exist they
are written back so downstream transforms can handle them.
"""
injection_points = kwargs.pop("cache_control_injection_points", None)
configured = cast( # cast-ok: kwargs is untyped; this key only holds the documented injection-point list
list[CacheControlInjectionPoint] | None, kwargs.pop("cache_control_injection_points", None)
)
injection_points: list[CacheControlInjectionPoint] = configured or []
if not injection_points and model is not None:
injection_points = AnthropicCacheControlHook.get_default_injection_points(
messages=cast(list[AllMessageValues], messages), # cast-ok: Anthropic-shaped dicts from v1/messages
system=system,
tools=tools,
model=model,
custom_llm_provider=custom_llm_provider,
)
if not injection_points:
return messages, system

79
litellm/integrations/langfuse/langfuse_otel.py
View file

@ -267,29 +267,10 @@ class LangfuseOtelLogger(OpenTelemetry):
# If no keys, return default from env (likely logging to console or something else)
return OpenTelemetryConfig.from_env()
# Determine endpoint - default to US cloud
langfuse_host = LangfuseOtelLogger._get_langfuse_otel_host()
if langfuse_host:
# If LANGFUSE_HOST is provided, construct OTEL endpoint from it
if not langfuse_host.startswith("http"):
langfuse_host = "https://" + langfuse_host
endpoint = f"{langfuse_host.rstrip('/')}/api/public/otel"
verbose_logger.debug(f"Using Langfuse OTEL endpoint from host: {endpoint}")
else:
# Default to US cloud endpoint
endpoint = LANGFUSE_CLOUD_US_ENDPOINT
verbose_logger.debug(f"Using Langfuse US cloud endpoint: {endpoint}")
auth_header = LangfuseOtelLogger._get_langfuse_authorization_header(
public_key=public_key, secret_key=secret_key
)
otlp_auth_headers = f"Authorization={auth_header}"
return OpenTelemetryConfig(
exporter="otlp_http",
endpoint=endpoint,
headers=otlp_auth_headers,
return LangfuseOtelLogger._build_langfuse_otel_config(
public_key=public_key,
secret_key=secret_key,
langfuse_host=LangfuseOtelLogger._get_langfuse_otel_host(),
)
@staticmethod
@ -316,33 +297,36 @@ class LangfuseOtelLogger(OpenTelemetry):
"LANGFUSE_PUBLIC_KEY and LANGFUSE_SECRET_KEY must be set for Langfuse OpenTelemetry integration."
)
# Determine endpoint - default to US cloud
langfuse_host = LangfuseOtelLogger._get_langfuse_otel_host()
return LangfuseOtelLogger._build_langfuse_otel_config(
public_key=public_key,
secret_key=secret_key,
langfuse_host=LangfuseOtelLogger._get_langfuse_otel_host(),
)
@staticmethod
def _build_langfuse_otel_config(
public_key: str, secret_key: str, langfuse_host: Optional[str]
) -> "OpenTelemetryConfig":
"""
Builds an OTLP HTTP config pointing at the Langfuse OTEL endpoint for the
given host (US cloud when no host is provided), authorized with the given keys.
"""
if langfuse_host:
# If LANGFUSE_HOST is provided, construct OTEL endpoint from it
if not langfuse_host.startswith("http"):
langfuse_host = "https://" + langfuse_host
endpoint = f"{langfuse_host.rstrip('/')}/api/public/otel"
normalized_host = langfuse_host if langfuse_host.startswith("http") else f"https://{langfuse_host}"
endpoint = f"{normalized_host.rstrip('/')}/api/public/otel"
verbose_logger.debug(f"Using Langfuse OTEL endpoint from host: {endpoint}")
else:
# Default to US cloud endpoint
endpoint = LANGFUSE_CLOUD_US_ENDPOINT
verbose_logger.debug(f"Using Langfuse US cloud endpoint: {endpoint}")
auth_header = LangfuseOtelLogger._get_langfuse_authorization_header(
public_key=public_key, secret_key=secret_key
)
otlp_auth_headers = f"Authorization={auth_header}"
# Prevent modification of global env vars which causes leakage
# os.environ["OTEL_EXPORTER_OTLP_ENDPOINT"] = endpoint
# os.environ["OTEL_EXPORTER_OTLP_HEADERS"] = otlp_auth_headers
return OpenTelemetryConfig(
exporter="otlp_http",
endpoint=endpoint,
headers=otlp_auth_headers,
headers=f"Authorization={auth_header}",
)
@staticmethod
@ -378,6 +362,29 @@ class LangfuseOtelLogger(OpenTelemetry):
return dynamic_headers
def construct_dynamic_otel_config(
self, standard_callback_dynamic_params: StandardCallbackDynamicParams
) -> Optional["OpenTelemetryConfig"]:
"""
Build a full per-request OTLP config from team/key dynamic Langfuse credentials.
Key-scoped credentials must define the export target, not just the auth
headers: without this, a proxy with no global LANGFUSE_* env vars keeps its
init-time fallback exporter (console), so key-level langfuse_otel silently
never reaches Langfuse.
"""
public_key = standard_callback_dynamic_params.get("langfuse_public_key")
secret_key = standard_callback_dynamic_params.get("langfuse_secret_key")
if not public_key or not secret_key:
return None
langfuse_host = standard_callback_dynamic_params.get("langfuse_host") or self._get_langfuse_otel_host()
return LangfuseOtelLogger._build_langfuse_otel_config(
public_key=public_key,
secret_key=secret_key,
langfuse_host=langfuse_host,
)
def create_litellm_proxy_request_started_span(
self,
start_time: datetime,

111
litellm/integrations/opentelemetry.py
View file

@ -28,6 +28,7 @@ from litellm.integrations.opentelemetry_utils.gen_ai_semconv import (
parse_semconv_opt_in,
)
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
from litellm.litellm_core_utils.secret_redaction import redact_string
from litellm.secret_managers.main import get_secret_bool, str_to_bool
from litellm.types.services import ServiceLoggerPayload
from litellm.types.utils import (
@ -948,12 +949,22 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger):
Returns:
Tracer: The tracer to use for this request
"""
dynamic_config = self._get_dynamic_otel_config_from_kwargs(kwargs)
if dynamic_config is not None:
verbose_logger.debug(
"[OTEL DEBUG] Using DYNAMIC config tracer with endpoint: %s",
dynamic_config.endpoint,
)
return self._get_tracer_with_dynamic_config(dynamic_config)
dynamic_headers = self._get_dynamic_otel_headers_from_kwargs(kwargs)
if dynamic_headers is not None:
# Create spans using a temporary tracer with dynamic headers
tracer_to_use = self._get_tracer_with_dynamic_headers(dynamic_headers)
verbose_logger.debug("[OTEL DEBUG] Using DYNAMIC tracer with headers: %s", dynamic_headers)
verbose_logger.debug(
"[OTEL DEBUG] Using DYNAMIC tracer with headers: %s", redact_string(str(dynamic_headers))
)
else:
# For langfuse_otel without dynamic headers, create a provider with env var credentials
if hasattr(self, "callback_name") and self.callback_name == "langfuse_otel":
@ -989,6 +1000,32 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger):
return dynamic_headers if dynamic_headers else None
def _get_dynamic_otel_config_from_kwargs(self, kwargs: dict) -> Optional[OpenTelemetryConfig]:
"""Extract a full dynamic exporter config from kwargs if available."""
standard_callback_dynamic_params: Optional[StandardCallbackDynamicParams] = kwargs.get(
"standard_callback_dynamic_params"
)
if not standard_callback_dynamic_params:
return None
return self.construct_dynamic_otel_config(standard_callback_dynamic_params=standard_callback_dynamic_params)
def _get_tracer_with_dynamic_config(self, dynamic_config: OpenTelemetryConfig):
"""Create (or reuse) a tracer whose exporter target comes from a per-request config."""
from opentelemetry.sdk.trace import TracerProvider
cache_key = f"dynamic_config:{dynamic_config.exporter}:{dynamic_config.endpoint}:{dynamic_config.headers}"
if cache_key in self._tracer_provider_cache:
return self._tracer_provider_cache[cache_key].get_tracer(LITELLM_TRACER_NAME)
temp_provider = TracerProvider(resource=self._get_litellm_resource(self.config))
temp_provider.add_span_processor(self._get_span_processor(config_override=dynamic_config))
self._tracer_provider_cache[cache_key] = temp_provider
return temp_provider.get_tracer(LITELLM_TRACER_NAME)
def _get_tracer_with_dynamic_headers(self, dynamic_headers: dict):
"""Create a temporary tracer with dynamic headers for this request only."""
from opentelemetry.sdk.trace import TracerProvider
@ -1020,6 +1057,19 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger):
"""
return None
def construct_dynamic_otel_config(
self, standard_callback_dynamic_params: StandardCallbackDynamicParams
) -> Optional[OpenTelemetryConfig]:
"""
Construct a full exporter config from standard callback dynamic params.
Override this when team/key dynamic params must control the export
target (exporter kind + endpoint), not just the request headers. When
this returns a config, it takes precedence over
construct_dynamic_otel_headers for the request.
"""
return None
#########################################################
# End of Team/Key Based Logging Control Flow
#########################################################
@ -2747,7 +2797,11 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger):
verbose_logger.debug("OpenTelemetry: No parent context found, creating root span")
return None, None
def _get_span_processor(self, dynamic_headers: Optional[dict] = None):
def _get_span_processor(
self,
dynamic_headers: Optional[dict] = None,
config_override: Optional[OpenTelemetryConfig] = None,
):
from opentelemetry.sdk.trace.export import (
BatchSpanProcessor,
ConsoleSpanExporter,
@ -2755,40 +2809,45 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger):
SpanExporter,
)
otel_exporter = config_override.exporter if config_override else self.OTEL_EXPORTER
otel_endpoint = config_override.endpoint if config_override else self.OTEL_ENDPOINT
otel_headers = config_override.headers if config_override else self.OTEL_HEADERS
verbose_logger.debug(
"OpenTelemetry Logger, initializing span processor \nself.OTEL_EXPORTER: %s\nself.OTEL_ENDPOINT: %s\nself.OTEL_HEADERS: %s",
self.OTEL_EXPORTER,
self.OTEL_ENDPOINT,
self.OTEL_HEADERS,
"OpenTelemetry Logger, initializing span processor \nexporter: %s\nendpoint: %s\nheaders: %s",
otel_exporter,
otel_endpoint,
redact_string(str(otel_headers)),
)
_split_otel_headers = OpenTelemetry._get_headers_dictionary(headers=dynamic_headers or self.OTEL_HEADERS)
_split_otel_headers = OpenTelemetry._get_headers_dictionary(headers=dynamic_headers or otel_headers)
if dynamic_headers:
verbose_logger.debug(
"[OTEL DEBUG] Creating span processor with DYNAMIC headers: %s",
{k: v[:20] + "..." if len(str(v)) > 20 else v for k, v in _split_otel_headers.items()},
redact_string(str(_split_otel_headers)),
)
elif config_override:
verbose_logger.debug(
"[OTEL DEBUG] Creating span processor with DYNAMIC config, endpoint: %s",
otel_endpoint,
)
else:
verbose_logger.debug("[OTEL DEBUG] Creating span processor with GLOBAL headers")
if hasattr(self.OTEL_EXPORTER, "export"): # Check if it has the export method that SpanExporter requires
if hasattr(otel_exporter, "export"): # Check if it has the export method that SpanExporter requires
verbose_logger.debug(
"OpenTelemetry: intiializing SpanExporter. Value of OTEL_EXPORTER: %s",
self.OTEL_EXPORTER,
otel_exporter,
)
return SimpleSpanProcessor(cast(SpanExporter, self.OTEL_EXPORTER))
return SimpleSpanProcessor(cast(SpanExporter, otel_exporter))
if self.OTEL_EXPORTER == "console":
if otel_exporter == "console":
verbose_logger.debug(
"OpenTelemetry: intiializing console exporter. Value of OTEL_EXPORTER: %s",
self.OTEL_EXPORTER,
otel_exporter,
)
return BatchSpanProcessor(ConsoleSpanExporter())
elif (
self.OTEL_EXPORTER == "otlp_http"
or self.OTEL_EXPORTER == "http/protobuf"
or self.OTEL_EXPORTER == "http/json"
):
elif otel_exporter == "otlp_http" or otel_exporter == "http/protobuf" or otel_exporter == "http/json":
try:
from opentelemetry.exporter.otlp.proto.http.trace_exporter import (
OTLPSpanExporter as OTLPSpanExporterHTTP,
@ -2801,13 +2860,13 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger):
verbose_logger.debug(
"OpenTelemetry: intiializing http exporter. Value of OTEL_EXPORTER: %s",
self.OTEL_EXPORTER,
otel_exporter,
)
normalized_endpoint = self._normalize_otel_endpoint(self.OTEL_ENDPOINT, "traces")
normalized_endpoint = self._normalize_otel_endpoint(otel_endpoint, "traces")
return BatchSpanProcessor(
OTLPSpanExporterHTTP(endpoint=normalized_endpoint, headers=_split_otel_headers),
)
elif self.OTEL_EXPORTER == "otlp_grpc" or self.OTEL_EXPORTER == "grpc":
elif otel_exporter == "otlp_grpc" or otel_exporter == "grpc":
try:
from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import (
OTLPSpanExporter as OTLPSpanExporterGRPC,
@ -2820,16 +2879,16 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger):
verbose_logger.debug(
"OpenTelemetry: intiializing grpc exporter. Value of OTEL_EXPORTER: %s",
self.OTEL_EXPORTER,
otel_exporter,
)
normalized_endpoint = self._normalize_otel_endpoint(self.OTEL_ENDPOINT, "traces")
normalized_endpoint = self._normalize_otel_endpoint(otel_endpoint, "traces")
return BatchSpanProcessor(
OTLPSpanExporterGRPC(endpoint=normalized_endpoint, headers=_split_otel_headers),
)
else:
verbose_logger.debug(
"OpenTelemetry: intiializing console exporter. Value of OTEL_EXPORTER: %s",
self.OTEL_EXPORTER,
otel_exporter,
)
return BatchSpanProcessor(ConsoleSpanExporter())
@ -2841,7 +2900,7 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger):
"OpenTelemetry Logger, initializing log exporter \nself.OTEL_EXPORTER: %s\nself.OTEL_ENDPOINT: %s\nself.OTEL_HEADERS: %s",
self.OTEL_EXPORTER,
self.OTEL_ENDPOINT,
self.OTEL_HEADERS,
redact_string(str(self.OTEL_HEADERS)),
)
_split_otel_headers = OpenTelemetry._get_headers_dictionary(self.OTEL_HEADERS)
@ -2928,7 +2987,7 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger):
"OpenTelemetry Logger, initializing metric reader\nself.OTEL_EXPORTER: %s\nself.OTEL_ENDPOINT: %s\nself.OTEL_HEADERS: %s",
self.OTEL_EXPORTER,
self.OTEL_ENDPOINT,
self.OTEL_HEADERS,
redact_string(str(self.OTEL_HEADERS)),
)
_split_otel_headers = OpenTelemetry._get_headers_dictionary(self.OTEL_HEADERS)

50
litellm/integrations/opik/utils.py
View file

@ -1,40 +1,38 @@
import configparser
import os
import time
import uuid
from typing import Any, Dict, Final, List, Optional, Tuple
CONFIG_FILE_PATH_DEFAULT: Final[str] = "~/.opik.config"
def create_uuid7():
ns = time.time_ns()
last = [0, 0, 0, 0]
def create_uuid7() -> str:
"""Generate an RFC 9562 conformant UUIDv7 string.
# Simple uuid7 implementation
sixteen_secs = 16_000_000_000
t1, rest1 = divmod(ns, sixteen_secs)
t2, rest2 = divmod(rest1 << 16, sixteen_secs)
t3, _ = divmod(rest2 << 12, sixteen_secs)
t3 |= 7 << 12 # Put uuid version in top 4 bits, which are 0 in t3
The top 48 bits encode the Unix timestamp in milliseconds. Opik's backend
validates this embedded timestamp on ingestion (it must fall within a window
around "now"), so the encoding has to be correct or trace/span batches are
rejected with HTTP 400. Implemented with the standard library only, so no
extra dependency is added to litellm. See ``opik.id_helpers`` for the
reference implementation.
"""
unix_ts_ms = int(time.time() * 1000)
# The next two bytes are an int (t4) with two bits for
# the variant 2 and a 14 bit sequence counter which increments
# if the time is unchanged.
if t1 == last[0] and t2 == last[1] and t3 == last[2]:
# Stop the seq counter wrapping past 0x3FFF.
# This won't happen in practice, but if it does,
# uuids after the 16383rd with that same timestamp
# will not longer be correctly ordered but
# are still unique due to the 6 random bytes.
if last[3] < 0x3FFF:
last[3] += 1
else:
last[:] = (t1, t2, t3, 0)
t4 = (2 << 14) | last[3] # Put variant 0b10 in top two bits
# Fill the 16-byte buffer with random data, then overwrite the structured
# parts (timestamp, version, variant) defined by the UUIDv7 layout.
uuid_bytes = bytearray(os.urandom(16))
# Six random bytes for the lower part of the uuid
rand = os.urandom(6)
return f"{t1:>08x}-{t2:>04x}-{t3:>04x}-{t4:>04x}-{rand.hex()}"
# First 48 bits (6 bytes): Unix timestamp in milliseconds.
uuid_bytes[0:6] = unix_ts_ms.to_bytes(6, byteorder="big")
# Version 7 in the top 4 bits of byte 6.
uuid_bytes[6] = 0x70 | (uuid_bytes[6] & 0x0F)
# Variant 0b10 in the top 2 bits of byte 8.
uuid_bytes[8] = 0x80 | (uuid_bytes[8] & 0x3F)
return str(uuid.UUID(bytes=bytes(uuid_bytes)))
def _read_opik_config_file() -> Dict[str, str]:

2
litellm/litellm_core_utils/dd_tracing.py
View file

@ -10,7 +10,7 @@ from typing import TYPE_CHECKING, Any, Optional, Union
from litellm.secret_managers.main import get_secret_bool
if TYPE_CHECKING:
from ddtrace.tracer import Tracer as DD_TRACER
from ddtrace.trace import Tracer as DD_TRACER
else:
DD_TRACER = Any

16
litellm/litellm_core_utils/env_utils.py
View file

@ -19,3 +19,19 @@ def get_env_int(env_var: str, default: int) -> int:
return int(raw)
except (ValueError, TypeError):
return default
def get_env_int_or_none(env_var: str) -> int | None:
"""Parse an environment variable as an integer, returning None when it is unset or unusable.
Use this instead of `get_env_int` when callers must distinguish "explicitly configured"
from "left at the default", for example when an override should take precedence over a
value resolved from somewhere else.
"""
raw = os.getenv(env_var)
if raw is None:
return None
try:
return int(raw.strip())
except (ValueError, TypeError):
return None

28
litellm/litellm_core_utils/litellm_logging.py
View file

@ -38,6 +38,7 @@ from litellm import (
)
from litellm._logging import _is_debugging_on, _redact_string, verbose_logger
from litellm.exceptions import (
BudgetExceededError,
validate_rate_limit_category,
validate_rate_limit_type,
)
@ -925,7 +926,6 @@ class Logging(LiteLLMLoggingBaseClass):
def pre_call(self, input, api_key, model=None, additional_args={}):
# Log the exact input to the LLM API
litellm.error_logs["PRE_CALL"] = locals()
try:
self._pre_call(
input=input,
@ -1135,7 +1135,6 @@ class Logging(LiteLLMLoggingBaseClass):
def post_call(self, original_response, input=None, api_key=None, additional_args={}):
# Log the exact result from the LLM API, for streaming - log the type of response received
litellm.error_logs["POST_CALL"] = locals()
if isinstance(original_response, dict):
original_response = json.dumps(original_response, default=str)
try:
@ -1454,6 +1453,9 @@ class Logging(LiteLLMLoggingBaseClass):
response_cost = litellm.response_cost_calculator(**response_cost_calculator_kwargs)
verbose_logger.debug(f"response_cost: {response_cost}")
additional_response_cost: object = self.model_call_details.get("additional_response_cost")
if isinstance(additional_response_cost, (int, float)) and additional_response_cost > 0:
return (response_cost or 0.0) + additional_response_cost
return response_cost
except Exception as e: # error calculating cost
debug_info = StandardLoggingModelCostFailureDebugInformation(
@ -1532,6 +1534,9 @@ class Logging(LiteLLMLoggingBaseClass):
and litellm_params.get(CallTypes.aimage_generation.value, False) is not True
and litellm_params.get(CallTypes.atranscription.value, False) is not True
and litellm_params.get(CallTypes.allm_passthrough_route.value, False) is not True
and litellm_params.get(CallTypes.aanthropic_messages.value, False) is not True
and litellm_params.get(CallTypes.agenerate_content.value, False) is not True
and litellm_params.get(CallTypes.agenerate_content_stream.value, False) is not True
)
def _is_assembled_stream_success(self, result=None) -> bool:
@ -3074,7 +3079,7 @@ class Logging(LiteLLMLoggingBaseClass):
def get_combined_callback_list(self, dynamic_success_callbacks: Optional[List], global_callbacks: List) -> List:
if dynamic_success_callbacks is None:
return list(global_callbacks)
return list(set(dynamic_success_callbacks + global_callbacks))
return list(dict.fromkeys(dynamic_success_callbacks + global_callbacks))
def _remove_internal_litellm_callbacks(self, callbacks: List) -> List:
"""
@ -4599,6 +4604,10 @@ class StandardLoggingPayloadSetup:
user_api_key_spend=None,
user_api_key_max_budget=None,
user_api_key_budget_reset_at=None,
user_api_key_user_spend=None,
user_api_key_user_max_budget=None,
user_api_key_team_spend=None,
user_api_key_team_max_budget=None,
user_api_key_team_id=None,
user_api_key_org_id=None,
user_api_key_org_alias=None,
@ -4945,6 +4954,7 @@ class StandardLoggingPayloadSetup:
rate_limit_category = validate_rate_limit_category(getattr(original_exception, "category", None))
rate_limit_type = validate_rate_limit_type(getattr(original_exception, "rate_limit_type", None))
budget_error = original_exception if isinstance(original_exception, BudgetExceededError) else None
return StandardLoggingPayloadErrorInformation(
error_code=error_status,
@ -4954,6 +4964,10 @@ class StandardLoggingPayloadSetup:
error_message=error_message,
error_rate_limit_category=rate_limit_category,
error_rate_limit_type=rate_limit_type,
error_budget_entity_type=budget_error.entity_type if budget_error else None,
error_budget_entity_id=budget_error.entity_id if budget_error else None,
error_budget_limit=budget_error.max_budget if budget_error else None,
error_budget_spend=budget_error.current_cost if budget_error else None,
)
@staticmethod
@ -5430,6 +5444,10 @@ def get_standard_logging_metadata(
user_api_key_spend=None,
user_api_key_max_budget=None,
user_api_key_budget_reset_at=None,
user_api_key_user_spend=None,
user_api_key_user_max_budget=None,
user_api_key_team_spend=None,
user_api_key_team_max_budget=None,
user_api_key_team_id=None,
user_api_key_org_id=None,
user_api_key_org_alias=None,
@ -5529,6 +5547,10 @@ def create_dummy_standard_logging_payload() -> StandardLoggingPayload:
user_api_key_team_id=str("test_team"),
user_api_key_user_id=str("test_user"),
user_api_key_team_alias=str("test_team_alias"),
user_api_key_user_spend=None,
user_api_key_user_max_budget=None,
user_api_key_team_spend=None,
user_api_key_team_max_budget=None,
user_api_key_org_id=None,
spend_logs_metadata=None,
requester_ip_address=str("127.0.0.1"),

41
litellm/litellm_core_utils/redact_messages.py
View file

@ -77,6 +77,22 @@ def _redact_streaming_response(streaming_response):
streaming_response.reasoning = None
def _redact_tool_calls(tool_calls) -> None:
"""Redact tool call arguments (assistant tool calls carry prompt-derived data)."""
if not tool_calls:
return
for tool_call in tool_calls:
function = getattr(tool_call, "function", None)
if function is not None and hasattr(function, "arguments"):
function.arguments = "redacted-by-litellm"
def _redact_function_call(function_call) -> None:
"""Redact legacy assistant function_call arguments."""
if function_call is not None and hasattr(function_call, "arguments"):
function_call.arguments = "redacted-by-litellm"
def _redact_choice_content(choice):
"""Helper to redact content in a choice (message or delta)."""
if isinstance(choice, litellm.Choices):
@ -85,12 +101,16 @@ def _redact_choice_content(choice):
choice.message.reasoning_content = "redacted-by-litellm"
if hasattr(choice.message, "thinking_blocks"):
choice.message.thinking_blocks = None
_redact_tool_calls(getattr(choice.message, "tool_calls", None))
_redact_function_call(getattr(choice.message, "function_call", None))
elif isinstance(choice, litellm.utils.StreamingChoices):
choice.delta.content = "redacted-by-litellm"
if hasattr(choice.delta, "reasoning_content"):
choice.delta.reasoning_content = "redacted-by-litellm"
if hasattr(choice.delta, "thinking_blocks"):
choice.delta.thinking_blocks = None
_redact_tool_calls(getattr(choice.delta, "tool_calls", None))
_redact_function_call(getattr(choice.delta, "function_call", None))
def _redact_responses_api_output(output_items):
@ -111,6 +131,9 @@ def _redact_responses_api_output(output_items):
if hasattr(summary_item, "text"):
summary_item.text = "redacted-by-litellm"
if hasattr(output_item, "type") and output_item.type == "function_call" and hasattr(output_item, "arguments"):
output_item.arguments = "redacted-by-litellm"
def _redact_responses_api_output_dict(output_items, redacted_str: str):
"""Helper to redact ResponsesAPIResponse output items in dict form."""
@ -131,6 +154,9 @@ def _redact_responses_api_output_dict(output_items, redacted_str: str):
if isinstance(summary_item, dict) and "text" in summary_item:
summary_item["text"] = redacted_str
if output_item.get("type") == "function_call" and "arguments" in output_item:
output_item["arguments"] = redacted_str
def _redact_standard_logging_object(model_call_details: dict):
"""Redact messages and response inside standard_logging_object if present."""
@ -162,6 +188,19 @@ def _redact_standard_logging_object(model_call_details: dict):
standard_logging_object["response"] = {"text": redacted_str}
def _redact_tool_calls_dict(message: dict, redacted_str: str) -> None:
"""Redact tool call / function_call arguments in a dict-form message or delta."""
tool_calls = message.get("tool_calls")
if isinstance(tool_calls, list):
for tool_call in tool_calls:
if isinstance(tool_call, dict) and isinstance(tool_call.get("function"), dict):
tool_call["function"]["arguments"] = redacted_str
function_call = message.get("function_call")
if isinstance(function_call, dict) and "arguments" in function_call:
function_call["arguments"] = redacted_str
def _redact_model_response_dict_choices(choices, redacted_str: str):
for choice in choices:
if isinstance(choice, dict):
@ -173,6 +212,7 @@ def _redact_model_response_dict_choices(choices, redacted_str: str):
choice["message"]["thinking_blocks"] = None
if "audio" in choice["message"]:
choice["message"]["audio"] = None
_redact_tool_calls_dict(choice["message"], redacted_str)
elif "delta" in choice and isinstance(choice["delta"], dict):
choice["delta"]["content"] = redacted_str
if "reasoning_content" in choice["delta"]:
@ -181,6 +221,7 @@ def _redact_model_response_dict_choices(choices, redacted_str: str):
choice["delta"]["thinking_blocks"] = None
if "audio" in choice["delta"]:
choice["delta"]["audio"] = None
_redact_tool_calls_dict(choice["delta"], redacted_str)
else:
_redact_choice_content(choice)

4
litellm/litellm_core_utils/secret_redaction.py
View file

@ -9,6 +9,8 @@ secrets from strings without depending on the logging-configuration module.
import re
from typing import List
from litellm.constants import MINIMUM_CUSTOM_KEY_LENGTH
_REDACTED = "REDACTED"
@ -30,7 +32,7 @@ def _build_secret_patterns() -> "re.Pattern[str]":
# Basic auth headers
r"Basic\s+[A-Za-z0-9+/]{10,}={0,2}",
# OpenAI / Anthropic sk- prefixed keys
r"sk-[A-Za-z0-9\-_]{20,}",
rf"sk-[A-Za-z0-9\-_]{{{MINIMUM_CUSTOM_KEY_LENGTH - len('sk-')},}}",
# Generic api_key / api-key / apikey (handles 'key': 'value' dict repr)
r"(?:api[_-]?key)['\"]?\s*[:=]\s*['\"]?[^\s,'\"})\]{}>]{8,}",
# x-api-key / api-key header values (handles 'key': 'value' dict repr)

40
litellm/litellm_core_utils/streaming_chunk_builder_utils.py
View file

@ -467,6 +467,7 @@ class ChunkProcessor:
cache_read_input_tokens: Optional[int] = None
completion_tokens_details: Optional[CompletionTokensDetails] = None
prompt_tokens_details: Optional[PromptTokensDetailsWrapper] = None
cost: Optional[float] = None
if "prompt_tokens" in usage_chunk:
prompt_tokens = usage_chunk.get("prompt_tokens", 0) or 0
@ -476,6 +477,8 @@ class ChunkProcessor:
cache_creation_input_tokens = usage_chunk.get("cache_creation_input_tokens")
if "cache_read_input_tokens" in usage_chunk:
cache_read_input_tokens = usage_chunk.get("cache_read_input_tokens")
if "cost" in usage_chunk:
cost = usage_chunk.get("cost")
if hasattr(usage_chunk, "completion_tokens_details"):
if isinstance(usage_chunk.completion_tokens_details, dict):
completion_tokens_details = CompletionTokensDetails(**usage_chunk.completion_tokens_details)
@ -494,6 +497,7 @@ class ChunkProcessor:
"cache_read_input_tokens": cache_read_input_tokens,
"completion_tokens_details": completion_tokens_details,
"prompt_tokens_details": prompt_tokens_details,
"cost": cost,
}
def count_reasoning_tokens(self, response: ModelResponse) -> Optional[int]:
@ -512,6 +516,22 @@ class ChunkProcessor:
return reasoning_tokens
@staticmethod
def _extract_usage_chunk(chunk: dict[str, Any] | ModelResponse | ModelResponseStream) -> Usage | None:
usage_chunk: Usage | dict[str, Any] | None = None
if hasattr(chunk, "usage") and chunk.usage is not None:
usage_chunk = chunk.usage
elif "usage" in chunk:
usage_chunk = chunk["usage"]
elif (isinstance(chunk, ModelResponse) or isinstance(chunk, ModelResponseStream)) and hasattr(
chunk, "_hidden_params"
):
usage_chunk = chunk._hidden_params.get("usage", None)
if isinstance(usage_chunk, dict):
return Usage(**usage_chunk)
return usage_chunk
def _calculate_usage_per_chunk(
self,
chunks: List[Union[Dict[str, Any], ModelResponse]],
@ -548,18 +568,12 @@ class ChunkProcessor:
# is last-wins, so without preserving this separately the 1h breakdown is
# lost and 1h cache writes get billed at the 5m rate.
cache_creation_token_details: Optional[CacheCreationTokenDetails] = None
cost: Optional[float] = None
for chunk in chunks:
usage_chunk: Optional[Usage] = None
if "usage" in chunk:
usage_chunk = chunk["usage"]
elif (isinstance(chunk, ModelResponse) or isinstance(chunk, ModelResponseStream)) and hasattr(
chunk, "_hidden_params"
):
usage_chunk = chunk._hidden_params.get("usage", None)
usage_chunk = self._extract_usage_chunk(chunk)
if usage_chunk is not None:
if isinstance(usage_chunk, dict):
usage_chunk = Usage(**usage_chunk)
usage_chunk_dict = self._usage_chunk_calculation_helper(usage_chunk)
if usage_chunk_dict["prompt_tokens"] is not None and usage_chunk_dict["prompt_tokens"] > 0:
prompt_tokens = usage_chunk_dict["prompt_tokens"]
@ -610,6 +624,9 @@ class ChunkProcessor:
prompt_tokens_details, cache_creation_token_details
)
if usage_chunk_dict["cost"] is not None:
cost = usage_chunk_dict["cost"]
prompt_tokens_details = self._attach_cache_creation_token_details(
prompt_tokens_details, cache_creation_token_details
)
@ -629,6 +646,7 @@ class ChunkProcessor:
web_search_requests=web_search_requests,
completion_tokens_details=completion_tokens_details,
prompt_tokens_details=prompt_tokens_details,
cost=cost,
)
@staticmethod
@ -727,6 +745,7 @@ class ChunkProcessor:
prompt_tokens_details: Optional[PromptTokensDetailsWrapper] = calculated_usage_per_chunk[
"prompt_tokens_details"
]
cost: Optional[float] = calculated_usage_per_chunk["cost"]
try:
returned_usage.prompt_tokens = prompt_tokens or token_counter(model=model, messages=messages)
@ -784,6 +803,9 @@ class ChunkProcessor:
else:
returned_usage.prompt_tokens_details.web_search_requests = web_search_requests
if cost is not None:
setattr(returned_usage, "cost", cost)
# Return a new usage object with the new values
returned_usage = Usage(**returned_usage.model_dump())

275
litellm/litellm_core_utils/streaming_handler.py
View file

@ -962,10 +962,11 @@ class CustomStreamWrapper:
if self.custom_llm_provider == "bedrock" and "trace" in model_response:
return model_response
# Default - return StopIteration
if hasattr(model_response, "usage"):
self.chunks.append(model_response)
raise StopIteration
# Don't raise StopIteration here - some providers (like OpenRouter)
# send usage/cost data in chunks after the finish_reason chunk
if hasattr(model_response, "usage") and model_response.usage is not None:
return model_response
return
# flush any remaining holding chunk
if len(self.holding_chunk) > 0:
if model_response.choices[0].delta.content is None:
@ -1474,12 +1475,16 @@ class CustomStreamWrapper:
self.tool_call = True
if hasattr(chunk, "usage") and chunk.usage is not None:
model_response.usage = chunk.usage
## RETURN ARG
return self.return_processed_chunk_logic(
result = self.return_processed_chunk_logic(
completion_obj=completion_obj,
model_response=model_response, # type: ignore
response_obj=response_obj,
)
return result
except StopIteration:
raise StopIteration
@ -1686,6 +1691,21 @@ class CustomStreamWrapper:
model_response.choices[0].finish_reason = "tool_calls"
return model_response
@staticmethod
def _propagate_usage_cost_to_hidden_params(
response: "ModelResponse",
) -> None:
"""
If the assembled response carries a provider-reported cost on
usage.cost, copy it into _hidden_params so litellm's cost
calculator uses it instead of a token-based estimate.
"""
_usage = getattr(response, "usage", None)
if _usage is not None and hasattr(_usage, "cost") and _usage.cost is not None:
if "additional_headers" not in response._hidden_params:
response._hidden_params["additional_headers"] = {}
response._hidden_params["additional_headers"]["llm_provider-x-litellm-response-cost"] = float(_usage.cost)
def __next__(self) -> "ModelResponseStream":
cache_hit = False
if self.custom_llm_provider is not None and self.custom_llm_provider == "cached_response":
@ -1741,6 +1761,10 @@ class CustomStreamWrapper:
# hasattr(response, "usage") is always True — must check
# `is not None` to avoid running this path on every chunk.
if getattr(response, "usage", None) is not None:
usage_to_preserve = response.usage
if usage_to_preserve:
response._hidden_params["usage"] = usage_to_preserve
obj_dict = response.model_dump()
if "usage" in obj_dict:
@ -1789,6 +1813,8 @@ class CustomStreamWrapper:
response = self.model_response_creator()
if complete_streaming_response is not None:
self._propagate_usage_cost_to_hidden_params(complete_streaming_response)
setattr(
response,
"usage",
@ -1974,97 +2000,7 @@ class CustomStreamWrapper:
self.chunks.append(processed_chunk)
return processed_chunk
except (StopAsyncIteration, StopIteration):
if self.sent_last_chunk is True:
# log the final chunk with accurate streaming values
try:
complete_streaming_response = litellm.stream_chunk_builder(
chunks=self.chunks,
messages=self.messages,
logging_obj=self.logging_obj,
)
except Exception as e:
# see sync __next__: a raise from stream_chunk_builder inside this
# except handler escapes __anext__ and drops the request from SpendLogs.
# Recover best-effort usage from the raw chunks so cost is still tracked
verbose_logger.warning(
"stream_chunk_builder raised at end-of-stream (%s); logging best-effort usage from chunks.",
str(e),
)
try:
complete_streaming_response = self.model_response_creator(
chunk={"usage": calculate_total_usage(chunks=self.chunks)}
)
except Exception:
complete_streaming_response = None
response = self.model_response_creator()
if complete_streaming_response is not None:
setattr(
response,
"usage",
getattr(complete_streaming_response, "usage"),
)
try:
_copy = complete_streaming_response.model_copy(deep=True)
except RuntimeError:
_copy = complete_streaming_response.model_copy()
asyncio.create_task(
self.async_cache_streaming_response(
processed_chunk=_copy,
cache_hit=cache_hit,
)
)
# Update hidden_params with final usage from
# stream_chunk_builder (see sync __next__ for full comment).
if (
self.stream_options is None
and complete_streaming_response is not None
and self._last_returned_hidden_params is not None
):
final_usage = getattr(complete_streaming_response, "usage", None)
if final_usage is not None:
self._last_returned_hidden_params["usage"] = final_usage
if self.sent_stream_usage is False and self.send_stream_usage is True:
self.sent_stream_usage = True
return response
_deferred_cb = getattr(
self.logging_obj,
"_on_deferred_stream_complete",
None,
)
if _deferred_cb is not None:
# Proxy has post-call guardrails. Store the assembled
# response so the outer streaming consumer
# (ProxyLogging.async_post_call_streaming_iterator_hook)
# can fire the deferred callback AFTER all guardrail
# end-of-stream blocks complete. Scheduling here via
# create_task would race with unified_guardrail's
# end-of-stream block for short-stream providers.
self.logging_obj._deferred_stream_complete_args = ( # type: ignore[attr-defined]
complete_streaming_response,
cache_hit,
)
else:
# prefer_async_handlers routes CustomLogger to async_success_handler
# when consumers use ``async for`` on sync-SDK streams. Legacy string
# callbacks still run via executor.submit inside dispatch_success_handlers.
asyncio.create_task(
self.logging_obj.dispatch_success_handlers(
complete_streaming_response,
cache_hit=cache_hit,
start_time=None,
end_time=None,
prefer_async_handlers=True,
)
)
raise StopAsyncIteration # Re-raise StopIteration
else:
self.sent_last_chunk = True
processed_chunk = self.finish_reason_handler()
return processed_chunk
return await self._finalize_completed_stream(cache_hit=cache_hit)
except httpx.TimeoutException as e: # if httpx read timeout error occues
traceback_exception = traceback.format_exc()
## ADD DEBUG INFORMATION - E.G. LITELLM REQUEST TIMEOUT
@ -2079,20 +2015,122 @@ class CustomStreamWrapper:
# Handle any exceptions that might occur during streaming
asyncio.create_task(self.logging_obj.async_failure_handler(e, traceback_exception))
self._handle_stream_fallback_error(e)
except (httpx.ReadError, httpx.RemoteProtocolError) as e:
if self.received_finish_reason is None:
self._log_stream_failure_and_raise(e)
return await self._finalize_completed_stream(cache_hit=cache_hit)
except Exception as e:
traceback_exception = traceback.format_exc()
if self.logging_obj is not None:
self._record_partial_usage_for_failure()
## LOGGING
threading.Thread(
target=self.logging_obj.failure_handler,
args=(e, traceback_exception),
).start() # log response
# Handle any exceptions that might occur during streaming
asyncio.create_task(
self.logging_obj.async_failure_handler(e, traceback_exception) # type: ignore
self._log_stream_failure_and_raise(e)
async def _finalize_completed_stream(self, cache_hit: bool) -> "ModelResponseStream":
if self.sent_last_chunk is True:
# log the final chunk with accurate streaming values
try:
complete_streaming_response = litellm.stream_chunk_builder(
chunks=self.chunks,
messages=self.messages,
logging_obj=self.logging_obj,
)
self._handle_stream_fallback_error(e)
except Exception as e:
# see sync __next__: a raise from stream_chunk_builder inside this
# except handler escapes __anext__ and drops the request from SpendLogs.
# Recover best-effort usage from the raw chunks so cost is still tracked
verbose_logger.warning(
"stream_chunk_builder raised at end-of-stream (%s); logging best-effort usage from chunks.",
str(e),
)
try:
complete_streaming_response = self.model_response_creator(
chunk={"usage": calculate_total_usage(chunks=self.chunks)}
)
except Exception:
complete_streaming_response = None
response = self.model_response_creator()
if complete_streaming_response is not None:
self._propagate_usage_cost_to_hidden_params(complete_streaming_response)
setattr(
response,
"usage",
getattr(complete_streaming_response, "usage"),
)
try:
_copy = complete_streaming_response.model_copy(deep=True)
except RuntimeError:
_copy = complete_streaming_response.model_copy()
asyncio.create_task(
self.async_cache_streaming_response(
processed_chunk=_copy,
cache_hit=cache_hit,
)
)
# Update hidden_params with final usage from
# stream_chunk_builder (see sync __next__ for full comment).
if (
self.stream_options is None
and complete_streaming_response is not None
and self._last_returned_hidden_params is not None
):
final_usage = getattr(complete_streaming_response, "usage", None)
if final_usage is not None:
self._last_returned_hidden_params["usage"] = final_usage
if self.sent_stream_usage is False and self.send_stream_usage is True:
self.sent_stream_usage = True
return response
_deferred_cb = getattr(
self.logging_obj,
"_on_deferred_stream_complete",
None,
)
if _deferred_cb is not None:
# Proxy has post-call guardrails. Store the assembled
# response so the outer streaming consumer
# (ProxyLogging.async_post_call_streaming_iterator_hook)
# can fire the deferred callback AFTER all guardrail
# end-of-stream blocks complete. Scheduling here via
# create_task would race with unified_guardrail's
# end-of-stream block for short-stream providers.
self.logging_obj._deferred_stream_complete_args = ( # type: ignore[attr-defined]
complete_streaming_response,
cache_hit,
)
else:
# prefer_async_handlers routes CustomLogger to async_success_handler
# when consumers use ``async for`` on sync-SDK streams. Legacy string
# callbacks still run via executor.submit inside dispatch_success_handlers.
asyncio.create_task(
self.logging_obj.dispatch_success_handlers(
complete_streaming_response,
cache_hit=cache_hit,
start_time=None,
end_time=None,
prefer_async_handlers=True,
)
)
raise StopAsyncIteration # Re-raise StopIteration
else:
self.sent_last_chunk = True
processed_chunk = self.finish_reason_handler()
return processed_chunk
def _log_stream_failure_and_raise(self, e: Exception) -> NoReturn:
traceback_exception = traceback.format_exc()
if self.logging_obj is not None:
self._record_partial_usage_for_failure()
## LOGGING
threading.Thread(
target=self.logging_obj.failure_handler,
args=(e, traceback_exception),
).start() # log response
# Handle any exceptions that might occur during streaming
asyncio.create_task(
self.logging_obj.async_failure_handler(e, traceback_exception) # type: ignore
)
self._handle_stream_fallback_error(e)
def _record_partial_usage_for_failure(self) -> None:
"""
@ -2228,12 +2266,16 @@ def calculate_total_usage(chunks: List[ModelResponse]) -> Usage:
"""Assume most recent usage chunk has total usage uptil then."""
prompt_tokens: int = 0
completion_tokens: int = 0
latest_usage_chunk = None
for chunk in chunks:
if "usage" in chunk and chunk["usage"] is not None:
if "prompt_tokens" in chunk["usage"]:
prompt_tokens = chunk["usage"].get("prompt_tokens", 0) or 0
if "completion_tokens" in chunk["usage"]:
completion_tokens = chunk["usage"].get("completion_tokens", 0) or 0
usage = chunk["usage"]
latest_usage_chunk = usage
if "prompt_tokens" in usage:
prompt_tokens = usage.get("prompt_tokens", 0) or 0
if "completion_tokens" in usage:
completion_tokens = usage.get("completion_tokens", 0) or 0
returned_usage_chunk = Usage(
prompt_tokens=prompt_tokens,
@ -2241,6 +2283,15 @@ def calculate_total_usage(chunks: List[ModelResponse]) -> Usage:
total_tokens=prompt_tokens + completion_tokens,
)
if latest_usage_chunk is not None:
latest_cost = (
latest_usage_chunk.get("cost")
if isinstance(latest_usage_chunk, dict)
else getattr(latest_usage_chunk, "cost", None)
)
if latest_cost is not None:
returned_usage_chunk.cost = latest_cost
return returned_usage_chunk

9
litellm/llms/anthropic/common_utils.py
View file

@ -906,16 +906,17 @@ def strip_advisor_blocks_from_messages(messages: List[Any], replace_with_text: b
def is_anthropic_invalid_thinking_signature_error(error_text: str) -> bool:
"""
Detect Anthropic 400 when encrypted thinking signatures in history do not match
the current deployment (e.g. user rotated API key or switched model endpoint).
Detect Anthropic 400 errors caused by missing or invalid thinking signatures.
Example API message:
Known error formats:
{"message":"messages.2.content.0.thinking.signature.str: Input should be a valid string"}
messages.N.content.M.thinking.signature.str: Input should be a valid string
messages.N.content.M: Invalid `signature` in `thinking` block
"""
if not error_text:
return False
lower = error_text.lower()
return "invalid" in lower and "signature" in lower and "thinking" in lower and "block" in lower
return "thinking" in lower and "signature" in lower and ("invalid" in lower or "valid string" in lower)
def strip_thinking_blocks_from_anthropic_messages(messages: List[Any]) -> List[Any]:

12
litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py
View file

@ -1403,11 +1403,6 @@ class LiteLLMAnthropicMessagesAdapter:
assert isinstance(thinking, str)
assert isinstance(signature, str)
if thinking and signature:
raise ValueError(
"Both `thinking` and `signature` in a single streaming chunk isn't supported."
)
return "thinking", ChatCompletionThinkingBlock(
type="thinking", thinking=thinking, signature=signature
)
@ -1463,17 +1458,14 @@ class LiteLLMAnthropicMessagesAdapter:
if choice.delta.reasoning_content is not None:
reasoning_content += choice.delta.reasoning_content
if reasoning_content and reasoning_signature:
raise ValueError("Both `reasoning` and `signature` in a single streaming chunk isn't supported.")
if partial_json is not None:
return "input_json_delta", ContentJsonBlockDelta(type="input_json_delta", partial_json=partial_json)
elif reasoning_content:
return "thinking_delta", ContentThinkingBlockDelta(type="thinking_delta", thinking=reasoning_content)
elif reasoning_signature:
return "signature_delta", ContentThinkingSignatureBlockDelta(
type="signature_delta", signature=reasoning_signature
)
elif reasoning_content:
return "thinking_delta", ContentThinkingBlockDelta(type="thinking_delta", thinking=reasoning_content)
else:
return "text_delta", ContentTextBlockDelta(type="text_delta", text=text)

12
litellm/llms/anthropic/experimental_pass_through/messages/handler.py
View file

@ -36,6 +36,7 @@ from litellm.types.llms.anthropic_messages.anthropic_response import (
AnthropicMessagesResponse,
)
from litellm.types.router import GenericLiteLLMParams
from litellm.types.utils import CallTypes
from litellm.utils import ProviderConfigManager, client
from ..utils import is_reasoning_auto_summary_enabled
@ -236,7 +237,9 @@ async def anthropic_messages(
AnthropicCacheControlHook,
)
messages, system = AnthropicCacheControlHook.maybe_inject_cache_control(messages, system, kwargs)
messages, system = AnthropicCacheControlHook.maybe_inject_cache_control(
messages, system, kwargs, model=model, custom_llm_provider=custom_llm_provider, tools=tools
)
original_stream = stream or kwargs.get("_websearch_interception_converted_stream", False)
@ -425,7 +428,9 @@ def anthropic_messages_handler(
AnthropicCacheControlHook,
)
messages, system = AnthropicCacheControlHook.maybe_inject_cache_control(messages, system, kwargs)
messages, system = AnthropicCacheControlHook.maybe_inject_cache_control(
messages, system, kwargs, model=model, custom_llm_provider=custom_llm_provider, tools=tools
)
metadata = validate_anthropic_api_metadata(metadata)
@ -463,6 +468,9 @@ def anthropic_messages_handler(
"model": original_model,
"custom_llm_provider": custom_llm_provider,
}
litellm_logging_obj.model_call_details.setdefault("litellm_params", {})[CallTypes.aanthropic_messages.value] = (
is_async
)
# Check if stream was converted for WebSearch interception
# This is set in the async wrapper above when stream=True is converted to stream=False

5
litellm/llms/anthropic/experimental_pass_through/responses_adapters/streaming_iterator.py
View file

@ -75,8 +75,9 @@ class AnthropicResponsesStreamWrapper:
# ---- message_start ----
if event_type == "response.created":
self._sent_message_start = True
self._chunk_queue.append(self._make_message_start())
if not self._sent_message_start:
self._sent_message_start = True
self._chunk_queue.append(self._make_message_start())
return
# ---- content_block_start for a new output message item ----

25
litellm/llms/custom_httpx/aiohttp_transport.py
View file

@ -85,29 +85,12 @@ class AiohttpResponseStream(httpx.AsyncByteStream):
try:
async for chunk in self._aiohttp_response.content.iter_chunked(self.CHUNK_SIZE):
yield chunk
except (
aiohttp.ClientPayloadError,
aiohttp.client_exceptions.ClientPayloadError,
) as e:
# Handle incomplete transfers more gracefully
# Log the error but don't re-raise if we've already yielded some data
verbose_logger.debug(f"Transfer incomplete, but continuing: {e}")
# If the error is due to incomplete transfer encoding, we can still
# return what we've received so far, similar to how httpx handles it
return
except RuntimeError as e:
# Some providers (e.g., SSE streams) may close the connection
# causing aiohttp StreamReader to raise a generic RuntimeError
# with message "Connection closed.". Treat this as a graceful
# end-of-stream so downstream consumers don't error.
if "Connection closed" in str(e):
verbose_logger.debug("Upstream closed streaming connection; ending iterator gracefully")
return
raise
if "Connection closed" not in str(e):
raise
raise httpx.ReadError(str(e)) from e
except aiohttp.http_exceptions.TransferEncodingError as e:
# Handle transfer encoding errors gracefully
verbose_logger.debug(f"Transfer encoding error, but continuing: {e}")
return
raise httpx.ReadError(str(e)) from e
except Exception:
# For other exceptions, use the normal mapping
with map_aiohttp_exceptions():

33
litellm/llms/custom_httpx/llm_http_handler.py
View file

@ -1879,6 +1879,7 @@ class BaseLLMHTTPHandler:
litellm_params: GenericLiteLLMParams,
api_key: Optional[str],
model: str,
timeout: Optional[Union[float, httpx.Timeout]] = None,
) -> httpx.Response:
max_attempts = max(provider_config.max_retry_on_anthropic_messages_http_error, 1)
litellm_params_dict = dict(litellm_params)
@ -1891,6 +1892,7 @@ class BaseLLMHTTPHandler:
data=signed_json_body or json.dumps(request_body),
stream=stream or False,
logging_obj=logging_obj,
timeout=timeout,
)
response.raise_for_status()
return response
@ -1925,6 +1927,32 @@ class BaseLLMHTTPHandler:
raise RuntimeError("unreachable: anthropic messages HTTP retry loop exited without return")
@staticmethod
def _resolve_anthropic_messages_timeout(
litellm_params: GenericLiteLLMParams,
stream: bool,
custom_llm_provider: str,
) -> Optional[Union[float, httpx.Timeout]]:
from litellm.litellm_core_utils.completion_timeout import CompletionTimeout
from litellm.litellm_core_utils.request_timeout_resolver import (
get_configured_request_timeout,
)
from litellm.utils import supports_httpx_timeout
stream_timeout = litellm_params.get("stream_timeout") if stream else None
model_timeout = stream_timeout if stream_timeout is not None else litellm_params.get("timeout")
request_timeout = litellm_params.get("request_timeout")
global_timeout = get_configured_request_timeout()
if model_timeout is None and request_timeout is None and global_timeout is None:
return None
return CompletionTimeout.resolve(
model_timeout,
{"request_timeout": request_timeout},
custom_llm_provider,
global_timeout=global_timeout,
supports_httpx_timeout=supports_httpx_timeout,
)
async def async_anthropic_messages_handler(
self,
model: str,
@ -2075,6 +2103,11 @@ class BaseLLMHTTPHandler:
litellm_params=litellm_params,
api_key=api_key,
model=model,
timeout=self._resolve_anthropic_messages_timeout(
litellm_params=litellm_params,
stream=stream or False,
custom_llm_provider=custom_llm_provider,
),
)
# used for logging + cost tracking

17
litellm/llms/fireworks_ai/cost_calculator.py
View file

@ -75,10 +75,23 @@ def cost_per_token(model: str, usage: Usage) -> Tuple[float, float]:
model_info = get_model_info(model=base_model, custom_llm_provider="fireworks_ai")
## CALCULATE INPUT COST
prompt_tokens_details = usage.prompt_tokens_details
cached_tokens: int = (
prompt_tokens_details.cached_tokens
if prompt_tokens_details is not None and prompt_tokens_details.cached_tokens is not None
else 0
)
input_cost_per_token: float = model_info["input_cost_per_token"] or 0.0
cache_read_input_token_cost = model_info.get("cache_read_input_token_cost")
cache_read_cost_per_token: float = (
cache_read_input_token_cost if cache_read_input_token_cost is not None else input_cost_per_token
)
non_cached_prompt_tokens: int = max(usage.prompt_tokens - cached_tokens, 0)
prompt_cost: float = usage["prompt_tokens"] * model_info["input_cost_per_token"]
prompt_cost: float = non_cached_prompt_tokens * input_cost_per_token + cached_tokens * cache_read_cost_per_token
## CALCULATE OUTPUT COST
completion_cost = usage["completion_tokens"] * model_info["output_cost_per_token"]
output_cost_per_token: float = model_info["output_cost_per_token"] or 0.0
completion_cost: float = usage.completion_tokens * output_cost_per_token
return prompt_cost, completion_cost

4
litellm/llms/openai/responses/guardrail_translation/handler.py
View file

@ -247,6 +247,8 @@ class OpenAIResponsesHandler(BaseTranslation):
"""
Merge remapped guardrailed tools with original tools that were not sent
to the guardrail (e.g. web_search, web_search_preview), preserving order.
Tools a guardrail appended (``remapped`` longer than ``original_tools``)
have no original slot and are kept so an injected tool is not dropped.
"""
if not original_tools:
return remapped
@ -262,6 +264,8 @@ class OpenAIResponsesHandler(BaseTranslation):
if j < len(remapped):
result.append(remapped[j])
j += 1
# Keep guardrail-appended tools that matched no original slot above.
result.extend(remapped[j:])
return result
def _apply_guardrailed_tools_to_data(

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

@ -1731,18 +1731,18 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
"""
Check if the candidate token count is inclusive of the thinking token count
if prompttokencount + candidatesTokenCount == totalTokenCount, then the candidate token count is inclusive of the thinking token count
if promptTokenCount + candidatesTokenCount + toolUsePromptTokenCount == totalTokenCount, then the candidate token count is inclusive of the thinking token count
else the candidate token count is exclusive of the thinking token count
Addresses - https://github.com/BerriAI/litellm/pull/10141#discussion_r2052272035
"""
if usage_metadata.get("promptTokenCount", 0) + usage_metadata.get(
"candidatesTokenCount", 0
) == usage_metadata.get("totalTokenCount", 0):
return True
else:
return False
non_thinking_tokens = (
usage_metadata.get("promptTokenCount", 0)
+ usage_metadata.get("candidatesTokenCount", 0)
+ usage_metadata.get("toolUsePromptTokenCount", 0)
)
return non_thinking_tokens == usage_metadata.get("totalTokenCount", 0)
@staticmethod
def _calculate_usage(
@ -1888,12 +1888,15 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
response_tokens_details = CompletionTokensDetailsWrapper()
response_tokens_details.reasoning_tokens = reasoning_tokens
tool_use_prompt_tokens = usage_metadata.get("toolUsePromptTokenCount") or None
prompt_tokens_details = PromptTokensDetailsWrapper(
cached_tokens=cached_tokens,
audio_tokens=prompt_audio_tokens,
text_tokens=prompt_text_tokens,
image_tokens=prompt_image_tokens,
video_tokens=prompt_video_tokens,
tool_use_tokens=tool_use_prompt_tokens,
)
completion_tokens = response_tokens or completion_response["usageMetadata"].get("candidatesTokenCount", 0)
@ -1901,7 +1904,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
completion_tokens = reasoning_tokens + completion_tokens
## GET USAGE ##
usage = Usage(
prompt_tokens=usage_metadata.get("promptTokenCount", 0),
prompt_tokens=usage_metadata.get("promptTokenCount", 0) + (tool_use_prompt_tokens or 0),
completion_tokens=completion_tokens,
total_tokens=usage_metadata.get("totalTokenCount", 0),
prompt_tokens_details=prompt_tokens_details,

27
litellm/main.py
View file

@ -510,6 +510,20 @@ async def acompletion(
#########################################################
#########################################################
litellm_logging_obj = kwargs.get("litellm_logging_obj", None)
from litellm.integrations.anthropic_cache_control_hook import (
AnthropicCacheControlHook,
)
from litellm.types.llms.openai import AllMessageValues
AnthropicCacheControlHook.maybe_seed_default_injection_points(
non_default_params=kwargs,
messages=cast(list[AllMessageValues], messages), # cast-ok: acompletion types messages as a bare List
model=model,
custom_llm_provider=cast(Optional[str], custom_llm_provider), # cast-ok: read from untyped kwargs
tools=tools,
)
if isinstance(litellm_logging_obj, LiteLLMLoggingObj) and (
litellm_logging_obj.should_run_prompt_management_hooks(
prompt_id=kwargs.get("prompt_id", None),
@ -5055,6 +5069,19 @@ def completion( # type: ignore
litellm_params = {} # used to prevent unbound var errors
## PROMPT MANAGEMENT HOOKS ##
from litellm.integrations.anthropic_cache_control_hook import (
AnthropicCacheControlHook,
)
from litellm.types.llms.openai import AllMessageValues
AnthropicCacheControlHook.maybe_seed_default_injection_points(
non_default_params=non_default_params,
messages=cast(list[AllMessageValues], messages), # cast-ok: completion types messages as a bare List
model=model,
custom_llm_provider=cast(Optional[str], kwargs.get("custom_llm_provider")), # cast-ok: untyped kwargs
tools=tools,
)
if isinstance(litellm_logging_obj, LiteLLMLoggingObj) and (
litellm_logging_obj.should_run_prompt_management_hooks(
prompt_id=prompt_id, non_default_params=non_default_params

483
litellm/model_prices_and_context_window_backup.json
View file

File diff suppressed because it is too large Load diff

1
litellm/models/mcp_server.py
View file

@ -79,6 +79,7 @@ class LiteLLM_MCPServerTable(LiteLLMPydanticObjectBase):
command: Optional[str] = None
args: List[str] = Field(default_factory=list)
env: Dict[str, str] = Field(default_factory=dict)
issuer: Optional[str] = None
authorization_url: Optional[str] = None
token_url: Optional[str] = None
registration_url: Optional[str] = None

32
litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py
View file

@ -1220,11 +1220,29 @@ class MCPRequestHandler:
global_mcp_server_manager,
)
key_tools = (
key_direct_tools = (
global_mcp_server_manager.expand_tool_permissions(key_obj_perm.mcp_tool_permissions).get(server_id)
if key_obj_perm
else None
)
# Tools granted through the key's toolsets restrict this server exactly
# as direct tool permissions do; union with any direct grants so the
# tool-level check sees the key's full effective tool scope
key_toolset_ids = (key_obj_perm.mcp_toolsets or []) if key_obj_perm else []
key_toolset_tools = (
(await global_mcp_server_manager.resolve_toolset_tool_permissions(toolset_ids=key_toolset_ids)).get(
server_id
)
if key_toolset_ids
else None
)
key_tools = (
list(set(key_direct_tools or []) | set(key_toolset_tools or []))
if key_direct_tools is not None or key_toolset_tools is not None
else None
)
team_tools = (
global_mcp_server_manager.expand_tool_permissions(team_obj_perm.mcp_tool_permissions).get(server_id)
if team_obj_perm
@ -1430,8 +1448,18 @@ class MCPRequestHandler:
global_mcp_server_manager.expand_tool_permissions(key_object_permission.mcp_tool_permissions).keys()
)
# servers referenced by the key's toolset grants are part of the key's
# scope on every path (list, call, REST), subject to the same team/org
# ceilings as any other key-level grant
toolset_ids = key_object_permission.mcp_toolsets or []
toolset_servers = (
list((await global_mcp_server_manager.resolve_toolset_tool_permissions(toolset_ids=toolset_ids)).keys())
if toolset_ids
else []
)
# Combine all lists
all_servers = direct_mcp_servers + access_group_servers + tool_perm_servers
all_servers = direct_mcp_servers + access_group_servers + tool_perm_servers + toolset_servers
return list(set(all_servers))
except Exception as e:
verbose_logger.warning(f"Failed to get allowed MCP servers for key: {str(e)}")

36
litellm/proxy/_experimental/mcp_server/db.py
View file

@ -48,6 +48,7 @@ if TYPE_CHECKING:
_AUTH_FLOW_SCOPED_FIELDS: frozenset = frozenset(
{
"issuer",
"authorization_url",
"token_url",
"registration_url",
@ -60,6 +61,13 @@ _AUTH_FLOW_SCOPED_FIELDS: frozenset = frozenset(
}
)
def _blank_to_none(value: Optional[str]) -> Optional[str]:
if not isinstance(value, str):
return None
return value.strip() or None
# Token-exchange settings with dedicated columns that also exist on
# ``MCPCredentials`` as a legacy shape (rows and REST callers that predate the
# columns). Every write lifts blob values into the columns and strips them from
@ -697,13 +705,15 @@ async def update_mcp_server(
# of being reset to a schema default (transport=sse, allow_all_keys=False...).
data_dict = _prepare_mcp_server_data(data, exclude_unset=True, fields_set=fields_set)
# Pre-fetch existing record once if we need it for auth_type or credential logic
# Pre-fetch existing record once if we need it for auth_type, url, or credential logic
existing = None
has_credentials = "credentials" in data_dict and data_dict["credentials"] is not None
# An explicit token-exchange column write (set or clear) also migrates the
# legacy blob copies below, so the existing row is needed for those updates.
explicit_te_write = bool(_TOKEN_EXCHANGE_COLUMN_FIELDS & data_dict.keys())
if data.auth_type or has_credentials or explicit_te_write:
url_provided = "url" in data_dict and data_dict["url"] is not None
issuer_provided = "issuer" in data_dict
if data.auth_type or has_credentials or explicit_te_write or url_provided or issuer_provided:
existing = await MCPServerRepository(prisma_client).table.find_unique(where={"server_id": data.server_id})
auth_type_changed = bool(
@ -711,13 +721,30 @@ async def update_mcp_server(
and existing
and _credential_auth_class(existing.auth_type) != _credential_auth_class(data.auth_type)
)
# A url change re-points the server at a potentially different upstream, so any discovered or
# trust-on-first-use OAuth endpoints/issuer belong to the old upstream and must re-discover.
url_changed = bool(url_provided and existing and existing.url != data_dict["url"])
old_issuer = _blank_to_none(getattr(existing, "issuer", None)) if existing else None
issuer_changed = bool(
issuer_provided and old_issuer is not None and _blank_to_none(data_dict.get("issuer")) != old_issuer
)
# Clear stale credentials when auth_type changes but no new credentials provided
if auth_type_changed and "credentials" not in data_dict:
data_dict["credentials"] = None
if auth_type_changed:
data_dict.update({field: None for field in _AUTH_FLOW_SCOPED_FIELDS if field not in data_dict})
if auth_type_changed or url_changed or issuer_changed:
# Clear each auth-flow-scoped field that the caller either omitted (partial update) or
# resubmitted unchanged. The edit form re-sends every field, so a stale issuer/endpoint
# belonging to the old upstream would otherwise survive a url/auth_type change and win in the
# resolution merge; only a genuinely new submitted value is kept.
data_dict.update(
{
field: None
for field in _AUTH_FLOW_SCOPED_FIELDS
if field not in data_dict or data_dict[field] == getattr(existing, field, None)
}
)
# An explicit column write that does not touch credentials must still migrate
# the row's legacy blob copies: lift values for columns the caller left
@ -1181,6 +1208,7 @@ def mcp_oauth_token_identity(server: object) -> tuple[object, ...]:
getattr(server, "spec_path", None),
getattr(server, "auth_type", None),
getattr(server, "oauth2_flow", None),
getattr(server, "issuer", None),
getattr(server, "authorization_url", None),
getattr(server, "token_url", None),
getattr(server, "registration_url", None),

16
litellm/proxy/_experimental/mcp_server/exceptions.py
View file

@ -88,3 +88,19 @@ class MCPToolResultError(Exception):
into two identities, breaking ``isinstance`` checks against instances
created before the reload.
"""
class MCPServerListError(Exception):
"""Carrier for a classified per-server listing fault (``faults.list_outcomes.ServerListFault``).
Raised where a server fetch used to silently return an empty tool list, so each boundary can
apply its own policy: the aggregate listing absorbs it into that server's outcome, while
single-server routes relay a truthful HTTP status instead of empty-success. The fault value is
typed as ``object`` here only to avoid a circular import with the faults package; construction
sites always pass a ``ServerListFault``.
"""
def __init__(self, fault: object, server_name: str) -> None:
self.fault = fault
self.server_name = server_name
super().__init__(f"Listing tools from MCP server {server_name!r} failed")

190
litellm/proxy/_experimental/mcp_server/faults/list_outcomes.py Normal file
View file

@ -0,0 +1,190 @@
"""Per-server outcomes for the aggregate MCP tools/list fan-out.
The aggregate listing deliberately keeps serving the healthy subset when one server fails, but a
failed server must contribute a classified outcome instead of silently shrinking the list: an empty
contribution with no signal makes a broken upstream indistinguishable from a healthy server with no
tools. Outcomes carry only machine fields (category and status code) so nothing from an upstream
body crosses the trust boundary; classification is total, so any exception out of a server fetch
becomes an outcome, never a second failure.
"""
from __future__ import annotations
from collections.abc import Iterator
from typing import Literal, NamedTuple, NoReturn, TypeAlias
import httpx
from mcp.types import Tool as MCPTool
from pydantic import BaseModel, ConfigDict
from typing_extensions import assert_never
from litellm.proxy._experimental.mcp_server.exceptions import (
MCPServerListError,
MCPUpstreamAuthError,
)
ListFaultCategory: TypeAlias = Literal[
"auth_required",
"forbidden",
"timeout",
"unreachable",
"upstream_error",
"internal",
]
class ServerListOk(BaseModel):
model_config = ConfigDict(frozen=True)
tag: Literal["ok"] = "ok"
tool_count: int
class ServerListFault(BaseModel):
"""Why a server contributed nothing to a listing: the caller must authenticate upstream
(``auth_required``/``forbidden``), the upstream did not answer (``timeout``/``unreachable``),
the upstream answered outside its contract (``upstream_error``), or the gateway itself failed
(``internal``). ``status_code`` is the upstream HTTP status when one exists."""
model_config = ConfigDict(frozen=True)
tag: ListFaultCategory
status_code: int | None = None
ServerOutcome: TypeAlias = ServerListOk | ServerListFault
SERVER_OUTCOMES_META_KEY = "litellm.ai/server_outcomes"
"""The tools/list result ``_meta`` key carrying per-server outcomes. Prefixed with the litellm.ai
domain per the MCP spec's ``_meta`` key format so it cannot collide with spec-reserved names."""
class AggregateToolListing(NamedTuple):
tools: list[MCPTool]
outcomes: dict[str, ServerOutcome]
def _iter_upstream_responses(exc: BaseException) -> Iterator[httpx.Response]:
"""Yield every ``httpx.Response`` in the exception tree (``__cause__``/``__context__``/
ExceptionGroup members) in deliberate order, mirroring how upstream failures surface through the
MCP SDK's task groups. Explicit links come first: each node's ``raise ... from`` cause, then
group members in raise order, then the incidental ``__context__`` chain, so a response raised
while handling the real failure can never shadow one on the explicit causal chain. Consumers
apply their own predicate over the stream: selecting the first response and THEN testing it
would miss a causal auth response sitting behind an unrelated earlier one."""
seen: set[int] = set()
stack = [exc]
while stack:
current = stack.pop()
if id(current) in seen:
continue
seen.add(id(current))
response = getattr(current, "response", None)
if isinstance(response, httpx.Response):
yield response
if current.__context__ is not None:
stack.append(current.__context__)
exceptions = getattr(current, "exceptions", None)
if isinstance(exceptions, tuple):
stack.extend(reversed(exceptions))
if current.__cause__ is not None:
stack.append(current.__cause__)
def _find_upstream_response(exc: BaseException) -> httpx.Response | None:
return next(_iter_upstream_responses(exc), None)
def upstream_auth_challenge(exc: BaseException) -> tuple[int, str | None] | None:
"""The first upstream 401/403 in deliberate order and its ``WWW-Authenticate`` challenge, both
read from the SAME response, so the status that picks the carrier channel and the challenge that
rides with it can never come from two different responses in the tree. Non-auth responses do not
end the scan: a causal 401 behind an unrelated 5xx must still be found, or the client never
receives the challenge it needs to re-authenticate."""
for response in _iter_upstream_responses(exc):
if response.status_code in (401, 403):
return response.status_code, response.headers.get("www-authenticate")
return None
def raise_classified_list_failure(
exc: BaseException,
server_name: str,
suppress_challenge: bool = False,
) -> NoReturn:
"""The one place a failed server fetch chooses its carrier: an upstream 401/403 travels as
``MCPUpstreamAuthError`` with the upstream's own challenge preserved (a challenge is only ever
fabricated at the HTTP edge, and only for a 401), everything else as ``MCPServerListError`` with
a classified fault. Every fetch site delegates here so the two channels cannot drift apart per
call site. ``suppress_challenge`` is for dcr_bridge servers, whose upstream challenge points
clients at the wrong protected-resource metadata and must never relay."""
auth = upstream_auth_challenge(exc)
if auth is not None:
status_code, challenge = auth
raise MCPUpstreamAuthError(
status_code=status_code,
www_authenticate=None if suppress_challenge else challenge,
server_name=server_name,
) from exc
raise MCPServerListError(classify_list_exception(exc), server_name) from exc
def classify_list_exception(exc: BaseException) -> ServerListFault:
"""Classify a per-server listing failure into exactly one outcome. Total: an exception this
function cannot recognize is the gateway's own fault (``internal``), never a re-raise."""
if isinstance(exc, MCPServerListError) and isinstance(exc.fault, ServerListFault):
return exc.fault
if isinstance(exc, MCPUpstreamAuthError):
tag = "forbidden" if exc.status_code == 403 else "auth_required"
return ServerListFault(tag=tag, status_code=exc.status_code)
if isinstance(exc, TimeoutError):
return ServerListFault(tag="timeout")
if isinstance(exc, ConnectionError):
return ServerListFault(tag="unreachable")
auth = upstream_auth_challenge(exc)
if auth is not None:
status_code, _ = auth
return ServerListFault(
tag="forbidden" if status_code == 403 else "auth_required",
status_code=status_code,
)
response = _find_upstream_response(exc)
if response is not None:
return ServerListFault(tag="upstream_error", status_code=response.status_code)
if isinstance(exc, (httpx.TimeoutException,)):
return ServerListFault(tag="timeout")
if isinstance(exc, httpx.TransportError):
return ServerListFault(tag="unreachable")
return ServerListFault(tag="internal")
def outcome_wire_value(outcome: ServerOutcome) -> dict[str, object]:
"""The client-visible form of one outcome, for the tools/list result ``_meta`` and the REST
response: category plus status code only, never upstream prose or URLs."""
match outcome.tag:
case "ok":
return {"status": "ok", "tool_count": outcome.tool_count}
case "auth_required" | "forbidden" | "timeout" | "unreachable" | "upstream_error" | "internal":
return {
"status": outcome.tag,
**({"http_status": outcome.status_code} if outcome.status_code is not None else {}),
}
case _:
assert_never(outcome.tag)
def list_fault_http_status(fault: ServerListFault) -> int:
"""The truthful HTTP status for a single-upstream listing fault per RFC 9110: the upstream's own
401/403 for auth, 504 for a timeout, 502 for an unreachable or misbehaving upstream, and 500 only
for the gateway's own failure."""
match fault.tag:
case "auth_required":
return fault.status_code or 401
case "forbidden":
return 403
case "timeout":
return 504
case "unreachable" | "upstream_error":
return 502
case "internal":
return 500
case _:
assert_never(fault.tag)

549
litellm/proxy/_experimental/mcp_server/mcp_server_manager.py
View file

@ -50,7 +50,15 @@ from litellm.llms.custom_httpx.http_handler import get_async_httpx_client
from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import (
MCPRequestHandler,
)
from litellm.proxy._experimental.mcp_server.exceptions import MCPUpstreamAuthError
from litellm.proxy._experimental.mcp_server.exceptions import (
MCPServerListError,
MCPUpstreamAuthError,
)
from litellm.proxy._experimental.mcp_server.faults.list_outcomes import (
ServerListFault,
raise_classified_list_failure,
upstream_auth_challenge,
)
from litellm.proxy._experimental.mcp_server.elicitation_handler import (
MCP_ELICITATION_AVAILABLE,
)
@ -186,6 +194,104 @@ _UPSTREAM_OAUTH_DISCOVERY_AUTH_TYPES: tuple[MCPAuth, ...] = (
)
def _blank_to_none(value: str | None) -> str | None:
"""Collapse an absent, empty, or whitespace-only string to ``None``.
OAuth endpoint fields are consumed by truthiness-based merges (``row or discovered``) and by the
corroboration gate. A whitespace-only value is truthy to ``or`` but is not a usable endpoint, so
without this the merge would keep the blank value for redirects while the gate treats it as
unpinned and backfills the other fields, yielding a broken half-discovered config. Normalizing
the pinned fields once, at each build entry point, gives every downstream consumer a single
notion of "blank" so those code paths cannot disagree.
"""
if not isinstance(value, str):
return None
return value.strip() or None
def _uses_issuer_anchor(manual_issuer: str | None, is_discovery_auth_type: bool) -> bool:
"""Whether the endpoints are authoritatively anchored to an admin-pinned issuer (RFC 8414 §3.3).
This is the trust/provenance property, distinct from whether the ``issuer`` field is merely
populated: a trust-on-first-use discovered issuer sets ``issuer`` for token identity but is NOT
anchored, so its endpoints stay resource-rooted. Anchoring holds only when the issuer was pinned
(present on the row/config) on a discovery auth type. Every consumer of "is this anchored" reads
this one definition, so the answer cannot diverge across build paths.
"""
return _blank_to_none(manual_issuer) is not None and is_discovery_auth_type
def _endpoints_yield_to_issuer(
issuer: str | None,
is_discovery_auth_type: bool,
authorization_url: str | None,
token_url: str | None,
registration_url: str | None,
) -> tuple[str | None, str | None, str | None]:
"""The single rule that makes an admin-configured ``issuer`` the sole authoritative endpoint
source (RFC 8414 §3.3): when it is set for a discovery auth type, the stored/manual
``authorization_url``/``token_url``/``registration_url`` do not apply. They neither anchor nor
short-circuit discovery, never override the issuer document in the merge, and never substitute for
it when the issuer fetch fails (fail-closed). Returns the endpoint values that remain in force,
i.e. all ``None`` when issuer-anchored, else the inputs unchanged. Called at every resolution site
so the invariant holds in one place instead of being re-derived per merge.
"""
if issuer is not None and is_discovery_auth_type:
return None, None, None
return authorization_url, token_url, registration_url
def _normalized_authorize_endpoint(url: str) -> str:
"""Compare authorize endpoints on scheme, host, and path only. The default port is elided and
the host is lowercased so ``https://IDP.example.com:443/authorize/`` and
``https://idp.example.com/authorize`` are the same identity; query and trailing slash are not."""
parsed = urlparse(url)
scheme = parsed.scheme.lower()
host = (parsed.hostname or "").lower()
default_port = {"https": 443, "http": 80}.get(scheme)
try:
port = parsed.port
except ValueError:
port = None
authority = host if port is None or port == default_port else f"{host}:{port}"
return f"{scheme}://{authority}{parsed.path.rstrip('/')}"
def _issuer_matches(claimed_issuer: object, configured_issuer: str) -> bool:
"""RFC 8414 §3.3 issuer equality between the metadata document's self-attested ``issuer`` and the
admin-configured issuer, tolerant only of URL-insignificant differences (scheme/host case, the
default port, a trailing slash). A non-string or empty claimed issuer never matches, so a
document that omits ``issuer`` fails closed under issuer-anchored discovery.
"""
if not isinstance(claimed_issuer, str) or not claimed_issuer:
return False
return _normalized_authorize_endpoint(claimed_issuer) == _normalized_authorize_endpoint(configured_issuer)
def _endpoints_corroborate_authorization_url(
source_authorization_url: str | None,
trusted_authorization_url: str | None,
) -> bool:
"""Whether a source's ``token_url``/``registration_url`` may be paired with a trusted authorize
endpoint. This is the single trust rule for adopting OAuth endpoints from any non-manual source.
Discovery is rooted at the MCP resource (RFC 9728), so a compromised upstream can advertise an
attacker-run authorization server. When ``authorization_url`` is admin-pinned, pairing it with a
``token_url`` from a different source is the RFC 9700 authorization-server mix-up: the user signs
in at the trusted authorize endpoint while the gateway redeems the code, with the stored client
secret and PKCE verifier, at the attacker's token endpoint. Endpoints are trustworthy together
only when they share an authorization server, so a source's endpoints are adopted only when the
same source advertised an ``authorization_endpoint`` matching the pinned value. With no pinned
value (``trusted_authorization_url is None``) there is nothing to protect: the authorize endpoint
comes from the same source as the token endpoint, so they corroborate each other by construction.
"""
if not (trusted_authorization_url and trusted_authorization_url.strip()):
return True
return bool(source_authorization_url) and _normalized_authorize_endpoint(
source_authorization_url
) == _normalized_authorize_endpoint(trusted_authorization_url)
def _carry_forward_resolved_oauth_endpoints(new_server: MCPServer, previous_server: MCPServer | None) -> None:
"""Keep the last known good OAuth endpoints when a rebuild's re-discovery comes back empty.
@ -193,26 +299,98 @@ def _carry_forward_resolved_oauth_endpoints(new_server: MCPServer, previous_serv
during re-discovery downgrades a working server (``authorization_url`` set) to a broken one
(``None``, /authorize 400s) with no configuration change. Mirrors the ``short_prefix``
carry-forward. Skipped when the server's ``url`` or ``auth_type`` changed, since the previous
endpoints may then belong to a different upstream. ``registration_url`` IS carried here even
though ``_persist_discovered_oauth_endpoints`` refuses to write it to the row: carrying only
restores the same in-memory value the previous build already ran with, while persisting it
would flip ``_dcr_bridge_relays_client_registration`` (which keys off the stored column) for
dcr_bridge servers that never had one configured.
endpoints may then belong to a different upstream. ``registration_url`` IS carried even though
``_persist_discovered_oauth_endpoints`` refuses to write it to the row: carrying only restores
the same in-memory value the previous build already ran with, while persisting it would flip
``_dcr_bridge_relays_client_registration`` (which keys off the stored column) for dcr_bridge
servers that never had one configured.
Carry-forward is a non-manual endpoint source, so the same trust rule as discovery applies: the
previous ``token_url``/``registration_url``/``scopes`` are carried only when the previous
``authorization_url`` corroborates the authorize endpoint this build will use, i.e. when the
incoming build has no pinned authorize endpoint (``None`` -> we adopt the previous one too, a
consistent group) or pins the same one. An admin re-pointing ``authorization_url`` to a different
server must not keep serving the old server's token endpoint or granted scopes.
When the server is issuer-anchored (``issuer_is_anchored`` -- a pinned issuer on a discovery auth
type), the endpoints come solely from the §3.3-validated issuer document, so carry-forward is
skipped entirely for its endpoints: a failed issuer fetch leaves them ``None`` and must stay
``None`` (fail-closed), never resurrected from the previous registry entry. A merely discovered
(trust-on-first-use) issuer is NOT anchored -- ``issuer`` is set for token identity but the
endpoints are resource-rooted, so they still carry forward as last-known-good, gated by the
corroboration check below like any other resource-rooted server. Scopes stay resource-driven and
can carry either way.
"""
if previous_server is None:
return
if previous_server.url != new_server.url or previous_server.auth_type != new_server.auth_type:
return
if new_server.issuer_is_anchored:
# Endpoints come solely from the §3.3-validated issuer document; a failed fetch stays
# fail-closed and must not be resurrected from the previous entry. Only the resource-driven
# scopes carry as last-known-good.
if not new_server.scopes and previous_server.scopes:
new_server.scopes = previous_server.scopes
return
may_carry = _endpoints_corroborate_authorization_url(
previous_server.authorization_url, new_server.authorization_url
)
if new_server.authorization_url is None and previous_server.authorization_url:
new_server.authorization_url = previous_server.authorization_url
if new_server.token_url is None and previous_server.token_url:
if may_carry and new_server.token_url is None and previous_server.token_url:
new_server.token_url = previous_server.token_url
if new_server.registration_url is None and previous_server.registration_url:
if may_carry and new_server.registration_url is None and previous_server.registration_url:
new_server.registration_url = previous_server.registration_url
if not new_server.scopes and previous_server.scopes:
if may_carry and not new_server.scopes and previous_server.scopes:
new_server.scopes = previous_server.scopes
def _restrict_discovery_to_corroborated_authorization_server(
metadata: MCPOAuthMetadata | None,
manual_authorization_url: str | None,
server_identifier: str,
is_dcr_bridge: bool,
) -> MCPOAuthMetadata | None:
"""Reject discovered token/registration endpoints a manually pinned authorize endpoint cannot
vouch for (the RFC 9700 authorization-server mix-up).
Discovery is rooted at the MCP resource, so a compromised upstream can advertise an attacker
``token_endpoint``: with ``authorization_url`` admin-pinned but ``token_url`` blank, the merge
would pair the trusted authorize endpoint with that attacker token endpoint, and the gateway would
post the authorization code and client secret there. So the discovered ``token_url`` and
``registration_url`` are kept only if the document corroborates the pin (its
``authorization_endpoint`` matches). ``scopes`` are deliberately NOT gated here: per the MCP
authorization spec Scope Selection Strategy and RFC 9700 §2.3, the scopes a client requests are
resource-driven (the WWW-Authenticate challenge or the RFC 9728 protected-resource
``scopes_supported``), and scope inflation by a compromised resource is bounded by the
authorization server and user consent (RFC 6749 §3.3), not by the client second-guessing the
request. With no pin there is no trust anchor to protect, so discovery is returned as-is.
"""
if metadata is None or not (manual_authorization_url and manual_authorization_url.strip()):
return metadata
if _endpoints_corroborate_authorization_url(metadata.authorization_url, manual_authorization_url):
return metadata
if not metadata.token_url and not metadata.registration_url:
return metadata
bridge_note = (
" The discovered registration_url is rejected with it, so this dcr_bridge server stays on the"
" short-circuit registration arm."
if is_dcr_bridge and metadata.registration_url
else ""
)
verbose_logger.warning(
"MCP OAuth discovery for server %s advertised authorization_endpoint %s, which does not match the "
"manually configured authorization_url %s; rejecting the discovered token_url/registration_url so "
"authorization codes and client credentials only follow the configured authorization server. "
"Configure Token URL manually if the mismatch is intentional.%s",
server_identifier,
_normalized_authorize_endpoint(metadata.authorization_url) if metadata.authorization_url else "<absent>",
_normalized_authorize_endpoint(manual_authorization_url),
bridge_note,
)
return metadata.model_copy(update={"token_url": None, "registration_url": None})
def invalidate_user_env_vars_cache(user_id: str, server_id: str) -> None:
"""Drop a cached entry after the user stores or clears their env var values
so the next request reads the fresh value instead of a stale one."""
@ -463,49 +641,14 @@ def _caller_authorization_fans_out(
def _extract_upstream_auth_failure(
exc: BaseException,
) -> Optional[tuple[int, Optional[str]]]:
"""Walk the exception tree looking for an HTTP 401/403 response from the
upstream MCP server.
"""The upstream 401/403 and its ``WWW-Authenticate`` header from the exception tree, or ``None``.
The MCP SDK wraps transport errors in anyio ``ExceptionGroup`` objects and
may chain through ``__cause__`` / ``__context__``. We inspect all of those
layers for an ``httpx.Response``-bearing exception (typically
``httpx.HTTPStatusError``) and extract the status code and any upstream
``WWW-Authenticate`` header.
Returns ``(status_code, www_authenticate)`` on match, else ``None``.
"""
seen: set[int] = set()
stack: list[BaseException] = [exc]
while stack:
current = stack.pop()
if id(current) in seen:
continue
seen.add(id(current))
response = getattr(current, "response", None)
if response is not None:
status_code = getattr(response, "status_code", None)
if isinstance(status_code, int) and status_code in (401, 403):
www_authenticate: Optional[str] = None
headers = getattr(response, "headers", None)
if headers is not None:
try:
www_authenticate = headers.get("www-authenticate")
except Exception:
www_authenticate = None
return status_code, www_authenticate
# anyio / PEP 654 ExceptionGroup
sub_exceptions = getattr(current, "exceptions", None)
if sub_exceptions:
stack.extend(sub_exceptions)
if current.__cause__ is not None:
stack.append(current.__cause__)
if current.__context__ is not None and current.__context__ is not current.__cause__:
stack.append(current.__context__)
return None
Delegates to the shared traversal in ``faults.list_outcomes`` so every consumer (tool listing,
tool calls, the connect-time probe) selects the same response with the same deliberate order:
explicit ``raise ... from`` causes first, ExceptionGroup members in raise order, the incidental
``__context__`` chain last. A response raised while handling the real failure can therefore never
shadow the causal one."""
return upstream_auth_challenge(exc)
def _warn_on_server_name_fields(
@ -1026,36 +1169,68 @@ class MCPServerManager:
)
auth_type = server_config.get("auth_type", None)
if server_url and (
auth_type in _UPSTREAM_OAUTH_DISCOVERY_AUTH_TYPES
manual_issuer = _blank_to_none(server_config.get("issuer"))
manual_authorization_url = _blank_to_none(server_config.get("authorization_url"))
manual_token_url = _blank_to_none(server_config.get("token_url"))
manual_registration_url = _blank_to_none(server_config.get("registration_url"))
is_discovery_auth_type = auth_type in _UPSTREAM_OAUTH_DISCOVERY_AUTH_TYPES
use_issuer_anchor = _uses_issuer_anchor(manual_issuer, is_discovery_auth_type)
manual_authorization_url, manual_token_url, manual_registration_url = _endpoints_yield_to_issuer(
manual_issuer,
is_discovery_auth_type,
manual_authorization_url,
manual_token_url,
manual_registration_url,
)
should_discover = bool(server_url) and (
is_discovery_auth_type
or self._obo_needs_endpoint_discovery(
auth_type,
server_config.get("token_exchange_endpoint"),
server_config.get("token_url"),
manual_token_url,
)
):
)
if not should_discover:
mcp_oauth_metadata = None
elif manual_issuer is not None and is_discovery_auth_type:
mcp_oauth_metadata = await self._fetch_issuer_anchored_oauth_metadata(manual_issuer, server_url)
else:
mcp_oauth_metadata = await self._descovery_metadata(
server_url=server_url,
allow_origin_fallback=auth_type in _UPSTREAM_OAUTH_DISCOVERY_AUTH_TYPES,
allow_origin_fallback=is_discovery_auth_type,
)
if use_issuer_anchor:
gated_oauth_metadata = mcp_oauth_metadata
elif is_discovery_auth_type:
gated_oauth_metadata = _restrict_discovery_to_corroborated_authorization_server(
mcp_oauth_metadata,
manual_authorization_url,
server_name or server_id,
bool(server_config.get("dcr_bridge")),
)
else:
mcp_oauth_metadata = None
gated_oauth_metadata = mcp_oauth_metadata
# Filter blank scopes (e.g. YAML ``scopes: [""]``) the same way the DB-build path does, so
# an all-blank list normalizes to None rather than a ``("",)`` tuple that skips the
# entra_obo fail-closed scope precondition and POSTs an empty scope to the IdP.
resolved_scopes = self._extract_scopes(server_config.get("scopes")) or (
mcp_oauth_metadata.scopes if mcp_oauth_metadata else None
gated_oauth_metadata.scopes if gated_oauth_metadata else None
)
resolved_authorization_url = server_config.get("authorization_url") or (
mcp_oauth_metadata.authorization_url if mcp_oauth_metadata else None
resolved_authorization_url = manual_authorization_url or (
gated_oauth_metadata.authorization_url if gated_oauth_metadata else None
)
resolved_token_url = server_config.get("token_url") or (
mcp_oauth_metadata.token_url if mcp_oauth_metadata else None
resolved_token_url = manual_token_url or (gated_oauth_metadata.token_url if gated_oauth_metadata else None)
resolved_registration_url = manual_registration_url or (
gated_oauth_metadata.registration_url if gated_oauth_metadata else None
)
resolved_registration_url = server_config.get("registration_url") or (
mcp_oauth_metadata.registration_url if mcp_oauth_metadata else None
discovered_issuer = (
gated_oauth_metadata.discovered_issuer
if gated_oauth_metadata and not gated_oauth_metadata.from_origin_fallback
else None
)
effective_issuer = manual_issuer or discovered_issuer
config_oauth2_flow = server_config.get("oauth2_flow", None)
if auth_type == MCPAuth.oauth2 and config_oauth2_flow not in (
@ -1104,6 +1279,8 @@ class MCPServerManager:
client_secret=server_config.get("client_secret", None),
oauth2_flow=self._explicit_oauth2_flow(config_oauth2_flow),
scopes=resolved_scopes,
issuer=effective_issuer,
issuer_is_anchored=use_issuer_anchor,
authorization_url=resolved_authorization_url,
token_url=resolved_token_url,
registration_url=resolved_registration_url,
@ -1364,6 +1541,52 @@ class MCPServerManager:
decrypt_global_env_var_values(env_vars_list)
return env_vars_list
async def _resolve_table_oauth_metadata(
self,
*,
mcp_server: LiteLLM_MCPServerTable,
auth_type: MCPAuthType,
server_url: Optional[str],
manual_issuer: Optional[str],
manual_authorization_url: Optional[str],
manual_token_url: Optional[str],
is_discovery_auth_type: bool,
use_issuer_anchor: bool,
scopes: Optional[list[str]],
token_exchange_endpoint: Optional[str],
) -> Optional[MCPOAuthMetadata]:
has_all_upstream_oauth_fields = bool(manual_authorization_url and manual_token_url and scopes)
needs_discovery = bool(server_url) and (
(is_discovery_auth_type and not has_all_upstream_oauth_fields)
or self._obo_needs_endpoint_discovery(auth_type, token_exchange_endpoint, manual_token_url)
)
if not needs_discovery:
mcp_oauth_metadata: Optional[MCPOAuthMetadata] = None
elif use_issuer_anchor and manual_issuer is not None:
mcp_oauth_metadata = await self._fetch_issuer_anchored_oauth_metadata(manual_issuer, server_url)
else:
mcp_oauth_metadata = await self._descovery_metadata(
server_url=server_url, # type: ignore[arg-type]
allow_origin_fallback=is_discovery_auth_type,
)
if needs_discovery and not use_issuer_anchor and mcp_oauth_metadata is None:
verbose_logger.warning(
"MCP OAuth discovery yielded no metadata for server %s (%s); "
"OAuth endpoints/scopes stay unresolved until a rebuild succeeds",
mcp_server.server_id,
server_url,
)
if use_issuer_anchor:
return mcp_oauth_metadata
if is_discovery_auth_type:
return _restrict_discovery_to_corroborated_authorization_server(
mcp_oauth_metadata,
manual_authorization_url,
mcp_server.server_id,
bool(getattr(mcp_server, "dcr_bridge", None)),
)
return mcp_oauth_metadata
async def build_mcp_server_from_table(
self,
mcp_server: LiteLLM_MCPServerTable,
@ -1447,32 +1670,38 @@ class MCPServerManager:
auth_type = cast(MCPAuthType, mcp_server.auth_type)
server_url = mcp_server.url
needs_discovery = bool(server_url) and (
(auth_type in _UPSTREAM_OAUTH_DISCOVERY_AUTH_TYPES and not mcp_server.authorization_url)
or self._obo_needs_endpoint_discovery(
auth_type,
mcp_server.token_exchange_endpoint
or (credentials_dict.get("token_exchange_endpoint") if credentials_dict else None),
mcp_server.token_url,
)
manual_issuer = _blank_to_none(mcp_server.issuer)
manual_authorization_url = _blank_to_none(mcp_server.authorization_url)
manual_token_url = _blank_to_none(mcp_server.token_url)
manual_registration_url = _blank_to_none(mcp_server.registration_url)
is_discovery_auth_type = auth_type in _UPSTREAM_OAUTH_DISCOVERY_AUTH_TYPES
use_issuer_anchor = _uses_issuer_anchor(manual_issuer, is_discovery_auth_type)
manual_authorization_url, manual_token_url, manual_registration_url = _endpoints_yield_to_issuer(
manual_issuer, is_discovery_auth_type, manual_authorization_url, manual_token_url, manual_registration_url
)
mcp_oauth_metadata = (
await self._descovery_metadata(
server_url=server_url, # type: ignore[arg-type]
allow_origin_fallback=auth_type in _UPSTREAM_OAUTH_DISCOVERY_AUTH_TYPES,
)
if needs_discovery
token_exchange_endpoint = mcp_server.token_exchange_endpoint or (
credentials_dict.get("token_exchange_endpoint") if credentials_dict else None
)
gated_oauth_metadata = await self._resolve_table_oauth_metadata(
mcp_server=mcp_server,
auth_type=auth_type,
server_url=server_url,
manual_issuer=manual_issuer,
manual_authorization_url=manual_authorization_url,
manual_token_url=manual_token_url,
is_discovery_auth_type=is_discovery_auth_type,
use_issuer_anchor=use_issuer_anchor,
scopes=scopes,
token_exchange_endpoint=token_exchange_endpoint,
)
resolved_scopes = scopes or (gated_oauth_metadata.scopes if gated_oauth_metadata else None)
discovered_issuer = (
gated_oauth_metadata.discovered_issuer
if gated_oauth_metadata and not gated_oauth_metadata.from_origin_fallback
else None
)
if needs_discovery and mcp_oauth_metadata is None:
verbose_logger.warning(
"MCP OAuth discovery yielded no metadata for server %s (%s); "
"OAuth endpoints stay unresolved until a rebuild succeeds",
mcp_server.server_id,
server_url,
)
resolved_scopes = scopes or (mcp_oauth_metadata.scopes if mcp_oauth_metadata else None)
effective_issuer = manual_issuer or discovered_issuer
new_server = MCPServer(
server_id=mcp_server.server_id,
@ -1492,9 +1721,11 @@ class MCPServerManager:
client_secret=client_secret_value or getattr(mcp_server, "client_secret", None),
oauth2_flow=self._explicit_oauth2_flow(getattr(mcp_server, "oauth2_flow", None)),
scopes=resolved_scopes,
authorization_url=mcp_server.authorization_url or getattr(mcp_oauth_metadata, "authorization_url", None),
token_url=mcp_server.token_url or getattr(mcp_oauth_metadata, "token_url", None),
registration_url=mcp_server.registration_url or getattr(mcp_oauth_metadata, "registration_url", None),
issuer=effective_issuer,
issuer_is_anchored=use_issuer_anchor,
authorization_url=manual_authorization_url or getattr(gated_oauth_metadata, "authorization_url", None),
token_url=manual_token_url or getattr(gated_oauth_metadata, "token_url", None),
registration_url=manual_registration_url or getattr(gated_oauth_metadata, "registration_url", None),
token_endpoint_auth_method=(
credentials_dict.get("token_endpoint_auth_method") if credentials_dict else None
),
@ -1545,16 +1776,18 @@ class MCPServerManager:
await self._persist_discovered_obo_token_url(
server_id=mcp_server.server_id,
auth_type=auth_type,
existing_token_url=mcp_server.token_url,
existing_token_url=manual_token_url,
discovered_token_url=new_server.token_url,
)
await self._persist_discovered_oauth_endpoints(
server_id=mcp_server.server_id,
auth_type=auth_type,
existing_authorization_url=mcp_server.authorization_url,
existing_token_url=mcp_server.token_url,
existing_issuer=manual_issuer,
existing_authorization_url=manual_authorization_url,
existing_token_url=manual_token_url,
existing_scopes=scopes,
metadata=mcp_oauth_metadata,
metadata=gated_oauth_metadata,
is_issuer_anchored=use_issuer_anchor,
)
return new_server
@ -1598,10 +1831,12 @@ class MCPServerManager:
*,
server_id: str,
auth_type: MCPAuthType | None,
existing_issuer: str | None,
existing_authorization_url: str | None,
existing_token_url: str | None,
existing_scopes: list[str] | None,
metadata: MCPOAuthMetadata | None,
is_issuer_anchored: bool = False,
) -> None:
"""Write freshly discovered OAuth endpoints back onto the DB row.
@ -1615,19 +1850,37 @@ class MCPServerManager:
because ``_dcr_bridge_relays_client_registration`` keys off that column. Best-effort: a
failed write re-discovers on the next build. Scopes go through ``update_mcp_server`` so
they merge into the credentials blob without touching the stored client credentials.
For an issuer-anchored server (``is_issuer_anchored``) the endpoints are re-derived from the
§3.3-validated issuer document on every build, so they are NOT persisted into the endpoint
columns: persisting them would make the next build see populated endpoints and treat them as
authoritative stored values, defeating the "endpoints come solely from the issuer" invariant.
Only the resource-driven scopes are persisted for such servers.
"""
if auth_type not in _UPSTREAM_OAUTH_DISCOVERY_AUTH_TYPES:
return
if metadata is None or metadata.from_origin_fallback:
return
issuer_update = (
{"issuer": metadata.discovered_issuer} if metadata.discovered_issuer and not existing_issuer else {}
)
authorization_url_update = (
{"authorization_url": metadata.authorization_url}
if metadata.authorization_url and not existing_authorization_url
if metadata.authorization_url and not existing_authorization_url and not is_issuer_anchored
else {}
)
token_url_update = (
{"token_url": metadata.token_url}
if metadata.token_url and not existing_token_url and not is_issuer_anchored
else {}
)
token_url_update = {"token_url": metadata.token_url} if metadata.token_url and not existing_token_url else {}
scopes_update = {"credentials": {"scopes": metadata.scopes}} if metadata.scopes and not existing_scopes else {}
updates: dict[str, object] = {**authorization_url_update, **token_url_update, **scopes_update}
updates: dict[str, object] = {
**issuer_update,
**authorization_url_update,
**token_url_update,
**scopes_update,
}
if not updates:
return
from litellm.proxy._experimental.mcp_server.db import ( # noqa: PLC0415 # db.py imports this module at load
@ -2767,10 +3020,12 @@ class MCPServerManager:
server_name=server.name,
) from e
verbose_logger.warning(f"Failed to get tools from server {server.name}: {str(e)}")
return []
raise MCPServerListError(ServerListFault(tag="internal", status_code=e.status_code), server.name) from e
except MCPServerListError:
raise
except Exception as e:
verbose_logger.warning(f"Failed to get tools from server {server.name}: {str(e)}")
return []
raise_classified_list_failure(e, server.name, suppress_challenge=server.is_dcr_bridge)
async def get_prompts_from_server(
self,
@ -3200,8 +3455,41 @@ class MCPServerManager:
return metadata
return None
async def _fetch_issuer_anchored_oauth_metadata(
self, issuer: str, server_url: Optional[str]
) -> Optional[MCPOAuthMetadata]:
"""RFC 8414 issuer-anchored discovery for the OAuth endpoints, with resource-driven scopes.
Fetch authorization-server metadata from the admin-configured issuer's own origin and adopt
its ``token_endpoint``/``registration_endpoint`` only when the document self-attests that same
issuer (RFC 8414 §3.3). Because the trust anchor is the pinned issuer rather than anything the
MCP resource advertises, the endpoints are authoritative for that issuer and cannot be
substituted by a compromised resource. Fails closed (returns None) on a §3.3 mismatch or a
fetch failure. The issuer is passed as its own ``server_url`` so the endpoint fetch is treated
as same-authority and is not subject to the resource-scoped SSRF shortcut.
Scopes are NOT taken from the issuer document. Per the MCP authorization spec Scope Selection
Strategy and RFC 9728, the scopes a client requests are resource-driven (the WWW-Authenticate
challenge or the protected-resource ``scopes_supported``), so the resource's advertised scopes
are fetched separately and used; the resource can influence only the requested scope, which
the authorization server and user consent bound (RFC 6749 §3.3), never the token endpoint.
"""
metadata = await self._fetch_single_authorization_server_metadata(issuer, issuer, require_issuer=issuer)
if metadata is None:
verbose_logger.warning(
"MCP OAuth issuer-anchored discovery for issuer %s yielded no metadata whose issuer "
"matched (RFC 8414 §3.3); OAuth endpoints stay unresolved until a rebuild succeeds",
issuer,
)
return None
resource_metadata = (
await self._descovery_metadata(server_url, allow_origin_fallback=False) if server_url else None
)
resource_scopes = resource_metadata.scopes if resource_metadata else None
return metadata.model_copy(update={"scopes": resource_scopes})
async def _fetch_single_authorization_server_metadata(
self, issuer_url: str, server_url: str
self, issuer_url: str, server_url: str, require_issuer: Optional[str] = None
) -> Optional[MCPOAuthMetadata]:
try:
parsed = urlparse(issuer_url)
@ -3245,20 +3533,33 @@ class MCPServerManager:
)
continue
scopes = self._extract_scopes(data.get("scopes_supported"))
claimed_issuer = data.get("issuer")
verbose_logger.debug(
"Authorization server metadata from %s: issuer=%s grant_types_supported=%s "
"token_endpoint_auth_methods_supported=%s",
url,
data.get("issuer"),
claimed_issuer,
data.get("grant_types_supported"),
data.get("token_endpoint_auth_methods_supported"),
)
if require_issuer is not None and not _issuer_matches(claimed_issuer, require_issuer):
verbose_logger.warning(
"MCP OAuth issuer-anchored discovery: metadata at %s self-attests issuer %r, which "
"does not match the configured issuer %r (RFC 8414 §3.3); rejecting so a compromised "
"resource cannot substitute an attacker authorization server",
url,
claimed_issuer,
require_issuer,
)
continue
scopes = self._extract_scopes(data.get("scopes_supported"))
metadata = MCPOAuthMetadata(
scopes=scopes,
authorization_url=data.get("authorization_endpoint"),
token_url=data.get("token_endpoint"),
registration_url=data.get("registration_endpoint"),
discovered_issuer=claimed_issuer if isinstance(claimed_issuer, str) and claimed_issuer else None,
)
if any(
@ -3356,16 +3657,17 @@ class MCPServerManager:
Uses anyio.fail_after() instead of asyncio.wait_for() to avoid conflicts
with the MCP SDK's anyio TaskGroup. See GitHub issue #20715 for details.
An upstream HTTP 401 is converted into :class:`MCPUpstreamAuthError`
instead of being swallowed to an empty tool list, regardless of the
server's auth_type. Callers route it by surface: the single-server HTTP
routes turn it into a 401 + ``WWW-Authenticate`` challenge so standards-
compliant MCP clients trigger the upstream OAuth flow, while the
multi-server ``/mcp`` aggregator absorbs it to an empty list so one
unauthenticated server doesn't fail the whole listing. Only a 401
(missing/invalid credential) drives the re-auth challenge; a 403
(authenticated but forbidden, e.g. insufficient scope) is not a re-auth
signal and, like other non-auth errors, returns an empty list.
Failures never return an empty tool list. An upstream 401 or 403 raises
:class:`MCPUpstreamAuthError` carrying the upstream's own
``WWW-Authenticate`` challenge when one was sent (a challenge is only
ever fabricated at the HTTP edge, and only for a 401: a 403 means the
caller is authenticated but not allowed, so prompting re-auth would be
wrong, while an upstream-sent 403 challenge is the RFC 6750
insufficient_scope step-up and relays verbatim). Every other failure
raises :class:`MCPServerListError` with a classified fault. Each
boundary then applies its own policy: single-server routes relay the
truthful status, the multi-server aggregator absorbs the failure into
that server's listing outcome.
Args:
client: MCP client instance
@ -3379,27 +3681,18 @@ class MCPServerManager:
tools = await client.list_tools(raise_on_error=True)
verbose_logger.debug(f"Tools from {server_name}: {tools}")
return tools
except TimeoutError:
except TimeoutError as e:
verbose_logger.warning(f"Timeout while listing tools from {server_name}")
return []
except asyncio.CancelledError:
raise MCPServerListError(ServerListFault(tag="timeout"), server_name) from e
except asyncio.CancelledError as e:
verbose_logger.warning(f"Task cancelled while listing tools from {server_name}")
return []
raise MCPServerListError(ServerListFault(tag="internal"), server_name) from e
except ConnectionError as e:
verbose_logger.warning(f"Connection error while listing tools from {server_name}: {str(e)}")
return []
raise MCPServerListError(ServerListFault(tag="unreachable"), server_name) from e
except Exception as e:
auth_info = _extract_upstream_auth_failure(e)
if auth_info is not None and auth_info[0] == 401:
_, www_authenticate = auth_info
verbose_logger.info(f"Upstream auth failure from MCP server {server_name}: HTTP 401")
raise MCPUpstreamAuthError(
status_code=401,
www_authenticate=www_authenticate,
server_name=server_name,
) from e
verbose_logger.warning(f"Error listing tools from {server_name}: {str(e)}")
return []
raise_classified_list_failure(e, server_name)
_SHORT_PREFIX_MAX_REHASH_ATTEMPTS = 1024
@ -4979,6 +5272,7 @@ class MCPServerManager:
command=getattr(server, "command", None),
args=getattr(server, "args", None) or [],
env=getattr(server, "env", None) or {},
issuer=server.issuer,
authorization_url=server.authorization_url,
token_url=server.token_url,
registration_url=server.registration_url,
@ -5088,6 +5382,7 @@ class MCPServerManager:
command=getattr(server, "command", None),
args=getattr(server, "args", None) or [],
env=getattr(server, "env", None) or {},
issuer=server.issuer,
authorization_url=server.authorization_url,
token_url=server.token_url,
registration_url=server.registration_url,

18
litellm/proxy/_experimental/mcp_server/oauth2_token_cache.py
View file

@ -175,16 +175,18 @@ mcp_oauth2_token_cache = MCPOAuth2TokenCache()
def _compute_per_user_token_ttl(server: "MCPServer", expires_in: Optional[int]) -> int:
"""Compute Redis TTL for a per-user token.
Uses server.token_storage_ttl_seconds when configured; otherwise derives
TTL from expires_in minus the expiry buffer; falls back to the default TTL.
Uses server.token_storage_ttl_seconds when configured, capped at the token's
remaining lifetime (expires_in minus the expiry buffer) so a cached entry never
outlives the token itself; otherwise derives TTL from expires_in minus the
expiry buffer; falls back to the default TTL.
"""
lifetime_bound = expires_in - MCP_PER_USER_TOKEN_EXPIRY_BUFFER_SECONDS if expires_in is not None else None
if server.token_storage_ttl_seconds is not None:
return max(server.token_storage_ttl_seconds, 1)
if expires_in is not None:
return max(
expires_in - MCP_PER_USER_TOKEN_EXPIRY_BUFFER_SECONDS,
1,
)
if lifetime_bound is None:
return max(server.token_storage_ttl_seconds, 1)
return max(min(server.token_storage_ttl_seconds, lifetime_bound), 1)
if lifetime_bound is not None:
return max(lifetime_bound, 1)
return MCP_PER_USER_TOKEN_DEFAULT_TTL

59
litellm/proxy/_experimental/mcp_server/rest_endpoints.py
View file

@ -19,12 +19,20 @@ import httpx
from fastapi import APIRouter, Depends, HTTPException, Query, Request, status
from litellm._logging import verbose_logger
from litellm.proxy._experimental.mcp_server.exceptions import MCPUpstreamAuthError
from litellm.proxy._experimental.mcp_server.exceptions import (
MCPServerListError,
MCPUpstreamAuthError,
)
from litellm.proxy._experimental.mcp_server.faults.list_outcomes import (
classify_list_exception,
list_fault_http_status,
)
from litellm.proxy._experimental.mcp_server.ui_session_utils import (
build_effective_auth_contexts,
)
from litellm.proxy._experimental.mcp_server.utils import (
MCPMissingUserEnvVarsError,
get_server_prefix,
merge_mcp_headers,
)
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
@ -515,20 +523,19 @@ if MCP_AVAILABLE:
# enforced even when no allowlist is set (matches the SSE/HTTP path).
tools = filter_tools_by_allowed_tools(tools, server)
# Filter tools based on user_api_key_auth.object_permission.mcp_tool_permissions
# This provides per-key/team/org control over which tools can be accessed
if (
user_api_key_auth
and user_api_key_auth.object_permission
and user_api_key_auth.object_permission.mcp_tool_permissions
):
# Dict keys may be server_ids OR names/aliases; normalize so lookup
# by concrete server_id resolves name-keyed restrictions too.
allowed_tools_for_server = global_mcp_server_manager.expand_tool_permissions(
user_api_key_auth.object_permission.mcp_tool_permissions
).get(server.server_id)
if allowed_tools_for_server is not None and len(allowed_tools_for_server) > 0:
# Filter tools to only include those in the allowed list
# Filter by the key's effective tool permissions through the same
# primitive the MCP protocol path uses (direct grants, toolset grants,
# and team/agent/org ceilings), so REST listing cannot drift from it
if user_api_key_auth:
from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import (
MCPRequestHandler,
)
allowed_tools_for_server = await MCPRequestHandler.get_allowed_tools_for_server(
server_id=server.server_id,
user_api_key_auth=user_api_key_auth,
)
if allowed_tools_for_server is not None:
tools = [tool for tool in tools if _tool_name_matches(tool.name, allowed_tools_for_server)]
return _create_tool_response_objects(tools, server)
@ -627,6 +634,16 @@ if MCP_AVAILABLE:
# matching status code and WWW-Authenticate challenge; that is what
# lets standards-compliant MCP clients run the upstream OAuth flow.
raise
except MCPServerListError as e:
fault = classify_list_exception(e)
verbose_logger.info(f"Listing tools from {server.name} failed with a {fault.tag} fault")
raise HTTPException(
status_code=list_fault_http_status(fault),
detail={
"error": fault.tag,
"message": f"Failed to list tools from server {get_server_prefix(server)}",
},
) from e
except Exception as e:
verbose_logger.exception(f"Error getting tools from {server.name}: {e}")
return {
@ -838,7 +855,11 @@ if MCP_AVAILABLE:
list_tools_result.extend(tools_result)
except Exception as e:
verbose_logger.exception(f"Error getting tools from {server.name}: {e}")
errors.append(f"{server.name}: {str(e)}")
errors.append(
f"{get_server_prefix(server)}: {classify_list_exception(e).tag}"
if isinstance(e, (MCPServerListError, MCPUpstreamAuthError))
else f"{get_server_prefix(server)}: {str(e)}"
)
continue
if errors and not list_tools_result:
@ -858,7 +879,10 @@ if MCP_AVAILABLE:
request_path=request.scope.get("_original_path") or request.url.path,
)
except HTTPException as http_exc:
if http_exc.status_code == status.HTTP_404_NOT_FOUND:
if http_exc.status_code == status.HTTP_404_NOT_FOUND or server_id:
# Single-server requests relay the truthful status (a 502/504 upstream fault must
# not masquerade as a 200 empty-success body); only the multi-server aggregate
# keeps the legacy error-dict response shape below.
raise
# Internal access/IP 403s keep the legacy error-dict response shape
# so the existing contract stays intact.
@ -1138,6 +1162,7 @@ if MCP_AVAILABLE:
static_headers=request.static_headers,
client_id=client_id,
client_secret=client_secret,
issuer=request.issuer,
token_url=request.token_url,
scopes=scopes,
authorization_url=request.authorization_url,

105
litellm/proxy/_experimental/mcp_server/semantic_tool_filter.py
View file

@ -4,6 +4,7 @@ Semantic MCP Tool Filtering using semantic-router
Filters MCP tools semantically for /chat/completions and /responses endpoints.
"""
import asyncio
from typing import TYPE_CHECKING, Any, Dict, List, Optional
from litellm._logging import verbose_logger
@ -76,6 +77,7 @@ class SemanticMCPToolFilter:
self.tool_router: Optional["SemanticRouter"] = None
self.context_window_error: Optional[str] = None
self._tool_map: Dict[str, Any] = {} # MCPTool objects or OpenAI function dicts
self._index_sync_lock = asyncio.Lock()
async def build_router_from_mcp_registry(self) -> None:
"""Build semantic router from all MCP tools in the registry (no auth checks)."""
@ -182,6 +184,81 @@ class SemanticMCPToolFilter:
return
raise
def _has_tools_missing_from_index(self, tools: list[Any]) -> bool:
"""Allocation-free check for any named tool not yet in the semantic index."""
return any(name and name not in self._tool_map for name in (self._extract_tool_info(t)[0] for t in tools))
def _tools_missing_from_index(self, tools: list[Any]) -> dict[str, Any]:
"""Map name -> tool for every named tool not yet in the semantic index."""
return {
name: tool
for name, tool in ((self._extract_tool_info(t)[0], t) for t in tools)
if name and name not in self._tool_map
}
async def _ensure_tools_indexed(self, available_tools: list[Any]) -> None:
"""
Index request-time tools the startup build never saw.
The startup index lists every registered MCP server WITHOUT per-user
credentials, so servers requiring per-user auth (interactive OAuth
tokens, user-scoped env vars) contribute zero routes. Tools reaching
the filter came through an authenticated expansion; without indexing
them here they can never be selected, so requests either bypass
filtering entirely (N->N) or lose every tool to unrelated matches.
Runs async-only (no synchronous embedding on the request path) and
never writes shared error state: an embedding failure here raises and
is scoped to the requesting call, so one request's oversized tool
description cannot poison the filter for other users on the worker.
"""
from semantic_router.routers import SemanticRouter
from semantic_router.routers.base import Route
from litellm.router_strategy.auto_router.litellm_encoder import (
LiteLLMRouterEncoder,
)
if not self._has_tools_missing_from_index(available_tools):
return
async with self._index_sync_lock:
missing = self._tools_missing_from_index(available_tools)
if not missing:
return
descriptions = {name: self._extract_tool_info(tool)[1] for name, tool in missing.items()}
routes = [
Route(
name=name,
description=description,
utterances=[description],
score_threshold=self.similarity_threshold,
)
for name, description in descriptions.items()
]
if self.tool_router is None:
router = SemanticRouter(
routes=[],
encoder=LiteLLMRouterEncoder(
litellm_router_instance=self.router_instance,
model_name=self.embedding_model,
score_threshold=self.similarity_threshold,
),
auto_sync="local",
top_k=self.top_k,
)
await router.aadd(routes)
self.tool_router = router
else:
await self.tool_router.aadd(routes)
self._tool_map.update(missing)
verbose_logger.info(
f"Semantic tool filter indexed {len(routes)} request-time tools missing from the startup index"
)
async def filter_tools(
self,
query: str,
@ -216,22 +293,34 @@ class SemanticMCPToolFilter:
if not query or not query.strip():
return available_tools
# Router should be built on startup - if not, something went wrong
if self.tool_router is None:
verbose_logger.warning("Router not initialized - was build_router_from_mcp_registry() called on startup?")
return available_tools
# Run semantic filtering
try:
await self._ensure_tools_indexed(available_tools)
if self.tool_router is None:
verbose_logger.warning("Semantic router could not be built from the request's tools")
return available_tools
available_names = [name for name in (self._extract_tool_info(t)[0] for t in available_tools) if name]
if not available_names:
return available_tools
limit = top_k or self.top_k
matches = self.tool_router(text=query, limit=limit)
if self.tool_router.top_k < limit:
self.tool_router.top_k = limit
matches = self.tool_router(text=query, limit=limit, route_filter=available_names)
matched_tool_names = self._extract_tool_names_from_matches(matches)
if not matched_tool_names:
return available_tools
return self._get_tools_by_names(matched_tool_names, available_tools)
filtered_tools = self._get_tools_by_names(matched_tool_names, available_tools)
if not filtered_tools:
return available_tools
return filtered_tools
except SemanticToolFilterContextWindowError:
raise
except Exception as e:
if _is_context_window_error(e):
verbose_logger.error(
@ -240,7 +329,7 @@ class SemanticMCPToolFilter:
)
raise SemanticToolFilterContextWindowError(
embedding_model=self.embedding_model,
stage="the user query",
stage="the user query or the MCP tool descriptions being indexed",
original_error=str(e),
) from e
verbose_logger.error(f"Semantic tool filter failed: {e}", exc_info=True)

243
litellm/proxy/_experimental/mcp_server/server.py
View file

@ -348,6 +348,7 @@ if MCP_AVAILABLE:
CallToolResult,
EmbeddedResource,
ImageContent,
ListToolsResult,
Prompt,
TextContent,
)
@ -356,6 +357,14 @@ if MCP_AVAILABLE:
from litellm.proxy._experimental.mcp_server.auth.litellm_auth_handler import (
MCPAuthenticatedUser,
)
from litellm.proxy._experimental.mcp_server.faults.list_outcomes import (
SERVER_OUTCOMES_META_KEY,
AggregateToolListing,
ServerListOk,
ServerOutcome,
classify_list_exception,
outcome_wire_value,
)
from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
MCPServerManager,
_caller_authorization_fans_out,
@ -664,9 +673,12 @@ if MCP_AVAILABLE:
########################################################
@server.list_tools()
async def handle_list_tools() -> List[Tool]:
async def handle_list_tools() -> "ListToolsResult | List[Tool]":
"""
List all available tools.
List all available tools, with each server's listing outcome attached to the result's
``_meta`` (SERVER_OUTCOMES_META_KEY) so a broken upstream is distinguishable from a healthy
server with no tools. Returning a ListToolsResult (rather than a bare list) makes the MCP SDK
pass the result through unwrapped, which is what lets the ``_meta`` survive to the client.
Also captures the active session for propagation to callbacks.
"""
from mcp.server.lowlevel.server import request_ctx
@ -709,7 +721,7 @@ if MCP_AVAILABLE:
# Get mcp_servers from context variable
verbose_logger.debug("MCP list_tools - Calling _list_mcp_tools")
tools = await _list_mcp_tools(
listing = await _list_mcp_tools(
user_api_key_auth=user_api_key_auth,
mcp_auth_header=mcp_auth_header,
mcp_servers=mcp_servers,
@ -719,8 +731,15 @@ if MCP_AVAILABLE:
log_list_tools_to_spendlogs=True,
list_tools_log_source="mcp_protocol",
)
verbose_logger.info(f"MCP list_tools - Successfully returned {len(tools)} tools")
return tools
verbose_logger.info(f"MCP list_tools - Successfully returned {len(listing.tools)} tools")
if not listing.outcomes:
return listing.tools
outcome_meta = {
SERVER_OUTCOMES_META_KEY: {
key: outcome_wire_value(outcome) for key, outcome in listing.outcomes.items()
}
}
return ListToolsResult.model_validate({"tools": listing.tools, "_meta": outcome_meta})
except Exception as e:
verbose_logger.exception(f"Error in list_tools endpoint: {str(e)}")
# Return empty list instead of failing completely
@ -1746,6 +1765,13 @@ if MCP_AVAILABLE:
_mcp_gateway_initialize_instructions.reset(instructions_token)
_mcp_gateway_server_name.reset(server_name_token)
def _aggregate_server_key(server: MCPServer) -> str:
"""The client-visible key for a server in listing outcomes and spend metadata: the same
display prefix (alias, or the short prefix when that mode is enabled) the caller already
sees on the tool names. Canonical internal server names never key a caller-readable
surface; when the display naming deliberately hides them, the outcome keys must too."""
return get_server_prefix(server) or "unknown"
async def _get_tools_from_mcp_servers(
user_api_key_auth: Optional[UserAPIKeyAuth],
mcp_auth_header: Optional[str],
@ -1758,7 +1784,7 @@ if MCP_AVAILABLE:
litellm_trace_id: Optional[str] = None,
request_tags: Optional[list[str]] = None,
client_ip: Optional[str] = None,
) -> List[MCPTool]:
) -> AggregateToolListing:
"""
Helper method to fetch tools from MCP servers based on server filtering criteria.
@ -1770,10 +1796,11 @@ if MCP_AVAILABLE:
oauth2_headers: Optional dict of oauth2 headers
Returns:
List[MCPTool]: Combined list of tools from filtered servers
AggregateToolListing: Combined tools from filtered servers plus each server's
classified listing outcome
"""
if not MCP_AVAILABLE:
return []
return AggregateToolListing(tools=[], outcomes={})
list_tools_start_time = datetime.now()
litellm_logging_obj: Optional[LiteLLMLoggingObj] = None
@ -1858,10 +1885,12 @@ if MCP_AVAILABLE:
async def _fetch_and_filter_server_tools(
server: MCPServer,
) -> List[MCPTool]:
"""Fetch and filter tools from a single server with error handling."""
) -> "tuple[List[MCPTool], ServerOutcome]":
"""Fetch and filter tools from a single server, classifying any failure into that
server's outcome so the aggregate can keep serving the healthy subset without a
broken server masquerading as an empty one."""
if server is None:
return []
return [], ServerListOk(tool_count=0)
server_auth_header, extra_headers = _prepare_mcp_server_headers(
server=server,
@ -1931,8 +1960,8 @@ if MCP_AVAILABLE:
verbose_logger.debug(
f"Successfully fetched {len(tools)} tools from server {server.name}, {len(filtered_tools)} after filtering"
)
return filtered_tools
except MCPUpstreamAuthError:
return filtered_tools, ServerListOk(tool_count=len(filtered_tools))
except MCPUpstreamAuthError as e:
# Absorb so one unauthenticated server does not empty every other server's
# tools. Surfacing the upstream 401 to the client as a re-auth challenge is
# intentionally not done here: raising from this list handler cannot produce a
@ -1940,31 +1969,30 @@ if MCP_AVAILABLE:
# error). Single-server routes surface it via the request-scope preemptive
# check in _raise_preemptive_401_for_unauthenticated_servers instead.
verbose_logger.debug(f"MCP list_tools: omitting {server.name}; it needs upstream auth")
return []
return [], classify_list_exception(e)
except Exception as e:
verbose_logger.exception(f"Error getting tools from server {server.name}: {str(e)}")
return []
return [], classify_list_exception(e)
# Fetch tools from all servers in parallel
tasks = [_fetch_and_filter_server_tools(server) for server in allowed_mcp_servers]
results = await asyncio.gather(*tasks)
# Flatten results into single list
all_tools: List[MCPTool] = [tool for tools in results for tool in tools]
all_tools: List[MCPTool] = [tool for tools, _ in results for tool in tools]
server_outcomes: Dict[str, ServerOutcome] = {
_aggregate_server_key(server): outcome
for server, (_, outcome) in zip(allowed_mcp_servers, results)
if server is not None
}
# If logging is enabled, enrich spend_logs_metadata with counts
if litellm_logging_obj:
per_server_tool_counts: Dict[str, int] = {}
for server, server_tools in zip(allowed_mcp_servers, results):
if server is None:
continue
server_key = (
getattr(server, "server_name", None)
or getattr(server, "alias", None)
or getattr(server, "name", None)
or "unknown"
)
per_server_tool_counts[str(server_key)] = len(server_tools)
per_server_tool_counts: Dict[str, int] = {
_aggregate_server_key(server): len(server_tools)
for server, (server_tools, _) in zip(allowed_mcp_servers, results)
if server is not None
}
metadata_dict = litellm_logging_obj.model_call_details.get("metadata")
if isinstance(metadata_dict, dict):
@ -1975,6 +2003,9 @@ if MCP_AVAILABLE:
spend_meta["allowed_server_count"] = len(allowed_mcp_servers)
spend_meta["tool_count_total"] = len(all_tools)
spend_meta["per_server_tool_counts"] = per_server_tool_counts
spend_meta["per_server_list_outcomes"] = {
key: outcome_wire_value(outcome) for key, outcome in server_outcomes.items()
}
end_time = datetime.now()
try:
@ -1995,7 +2026,7 @@ if MCP_AVAILABLE:
verbose_logger.info(f"Successfully fetched {len(all_tools)} tools total from all MCP servers")
return all_tools
return AggregateToolListing(tools=all_tools, outcomes=server_outcomes)
except Exception as e:
# Only fire failure hook if logging was requested for this list-tools execution
if log_list_tools_to_spendlogs and user_api_key_auth is not None:
@ -2218,43 +2249,6 @@ if MCP_AVAILABLE:
server = global_mcp_server_manager.get_mcp_server_by_id(server_id)
return [t for t in tools if strip_known_server_prefix(t.name, server) in allowed_tool_names]
async def _merge_toolset_permissions(
user_api_key_auth: Optional[UserAPIKeyAuth],
) -> Optional[UserAPIKeyAuth]:
"""
Resolve mcp_toolsets on the key's object_permission into tool-level permissions
and merge them (union) into object_permission.mcp_tool_permissions.
Returns the (possibly mutated copy of) user_api_key_auth.
"""
if user_api_key_auth is None:
return None
op = user_api_key_auth.object_permission
if op is None:
return user_api_key_auth
toolset_ids = getattr(op, "mcp_toolsets", None) or []
if not toolset_ids:
return user_api_key_auth
toolset_perms = await global_mcp_server_manager.resolve_toolset_tool_permissions(toolset_ids=toolset_ids)
if not toolset_perms:
return user_api_key_auth
# Merge toolset_perms into existing mcp_tool_permissions (union)
existing = dict(op.mcp_tool_permissions or {})
for server_id, tool_names in toolset_perms.items():
existing_tools = existing.get(server_id, [])
merged = list(set(existing_tools) | set(tool_names))
existing[server_id] = merged
# Build updated object_permission with merged tool permissions and server IDs.
# Union the toolset's server IDs into mcp_servers so downstream server-level
# filtering doesn't silently drop servers that the toolset references but that
# aren't already in the key's explicit mcp_servers list.
merged_servers = list(set(op.mcp_servers or []) | set(existing.keys()))
updated_op = op.model_copy(update={"mcp_servers": merged_servers, "mcp_tool_permissions": existing})
return user_api_key_auth.model_copy(update={"object_permission": updated_op})
async def _list_mcp_tools(
user_api_key_auth: Optional[UserAPIKeyAuth] = None,
mcp_auth_header: Optional[str] = None,
@ -2265,7 +2259,7 @@ if MCP_AVAILABLE:
log_list_tools_to_spendlogs: bool = False,
list_tools_log_source: Optional[str] = None,
client_ip: Optional[str] = None,
) -> List[MCPTool]:
) -> AggregateToolListing:
"""
List all available MCP tools.
@ -2277,19 +2271,14 @@ if MCP_AVAILABLE:
client_ip: Client IP for IP-based server access control
Returns:
List[MCPTool]: Combined list of tools from all accessible servers
AggregateToolListing: Combined tools from all accessible servers plus each server's
classified listing outcome
"""
if not MCP_AVAILABLE:
return []
return AggregateToolListing(tools=[], outcomes={})
# Resolve toolset permissions and merge into the key's object_permission
# so that the existing filter_tools_by_key_team_permissions logic picks them up.
user_api_key_auth = await _merge_toolset_permissions(user_api_key_auth)
# Get tools from managed MCP servers with error handling
managed_tools = []
try:
managed_tools = await _get_tools_from_mcp_servers(
listing = await _get_tools_from_mcp_servers(
user_api_key_auth=user_api_key_auth,
mcp_auth_header=mcp_auth_header,
mcp_servers=mcp_servers,
@ -2300,12 +2289,12 @@ if MCP_AVAILABLE:
list_tools_log_source=list_tools_log_source,
client_ip=client_ip,
)
verbose_logger.debug(f"Successfully fetched {len(managed_tools)} tools from managed MCP servers")
verbose_logger.debug(f"Successfully fetched {len(listing.tools)} tools from managed MCP servers")
return listing
except Exception as e:
verbose_logger.exception(f"Error getting tools from managed MCP servers: {str(e)}")
# Continue with empty managed tools list instead of failing completely
return managed_tools
# Continue with an empty listing instead of failing completely
return AggregateToolListing(tools=[], outcomes={})
async def _list_mcp_prompts(
user_api_key_auth: Optional[UserAPIKeyAuth] = None,
@ -3582,48 +3571,70 @@ if MCP_AVAILABLE:
# preemptive challenge and let downstream authorization
# return 403.
continue
if server and server.auth_type == MCPAuth.oauth2 and not oauth2_headers:
# For per-user OAuth servers, only skip the pre-emptive 401 when
# a stored token actually exists for this user+server pair.
# If no stored token exists, fail fast with 401 so clients can
# kick off PKCE/interactive OAuth flow immediately.
if server.needs_user_oauth_token:
if getattr(server, "delegate_auth_to_upstream", False) is True:
# Delegate-auth servers run upstream PKCE: challenge with
# the proxied resource_metadata (RFC 9728), not the
# gateway authorization_uri below which would authorize
# against the gateway instead of the upstream IdP.
www_authenticate = _get_passthrough_www_authenticate(
scope=scope,
server_name=server_name,
)
raise HTTPException(
status_code=401,
detail="Unauthorized",
headers={"www-authenticate": www_authenticate},
)
# The v2 resolver owns the existence check, so every authorization_code
# resolution (egress and this discovery challenge) runs through it.
if server and server.auth_type == MCPAuth.oauth2:
# The challenge decision is per oauth2 sub-mode, not per header:
# gateway-managed modes (M2M and interactive authorization_code)
# never receive a client-supplied upstream token, so a bearer in
# Authorization is a LiteLLM key (surfaced here as oauth2_headers)
# and must not suppress the challenge. Only the delegate mode
# treats a present bearer as the upstream token. The sub-mode is
# resolved the same way egress resolves it, via
# effective_oauth2_flow: an unstamped (null oauth2_flow) row with
# the M2M shape resolves to client_credentials, so the bare
# has_client_credentials column is never trusted here.
if MCPServerManager.effective_oauth2_flow(server) == "client_credentials":
# M2M: the gateway mints its own token at egress from the
# stored client credentials, so there is nothing to challenge.
continue
if getattr(server, "delegate_auth_to_upstream", False) is not True:
# Gateway-managed interactive (authorization_code): the only
# thing that authorizes egress is a stored per-user token, so
# challenge whenever one is absent, regardless of any bearer.
# The v2 resolver owns the existence check, so every
# authorization_code resolution (egress and this discovery
# challenge) runs through it.
if await global_mcp_server_manager.has_user_oauth_token(server, user_api_key_auth):
continue
request = StarletteRequest(scope)
base_url = get_request_base_url(request)
_path = scope.get("_original_path") or scope.get("path", "") or ""
request = StarletteRequest(scope)
base_url = get_request_base_url(request)
_path = scope.get("_original_path") or scope.get("path", "") or ""
# Pick the well-known AS-metadata form that matches the inbound route
# so strict RFC 9728 §3.2 clients can resolve it correctly.
if _path.startswith(f"/mcp/{server_name}"):
_as_url = f"{base_url}/.well-known/oauth-authorization-server/mcp/{server_name}"
else:
_as_url = f"{base_url}/.well-known/oauth-authorization-server/{server_name}"
authorization_uri = f'Bearer authorization_uri="{_as_url}"'
# Pick the well-known AS-metadata form that matches the inbound route
# so strict RFC 9728 §3.2 clients can resolve it correctly.
if _path.startswith(f"/mcp/{server_name}"):
_as_url = f"{base_url}/.well-known/oauth-authorization-server/mcp/{server_name}"
else:
_as_url = f"{base_url}/.well-known/oauth-authorization-server/{server_name}"
authorization_uri = f'Bearer authorization_uri="{_as_url}"'
raise HTTPException(
status_code=401,
detail="Unauthorized",
headers={"www-authenticate": authorization_uri},
)
raise HTTPException(
status_code=401,
detail="Unauthorized",
headers={"www-authenticate": authorization_uri},
)
if not oauth2_headers:
# Delegate-auth servers run upstream PKCE: a present bearer is
# the upstream token, so only challenge when it is absent, with
# the proxied resource_metadata (RFC 9728), not the gateway
# authorization_uri above which would authorize against the
# gateway instead of the upstream IdP.
www_authenticate = _get_passthrough_www_authenticate(
scope=scope,
server_name=server_name,
)
raise HTTPException(
status_code=401,
detail="Unauthorized",
headers={"www-authenticate": www_authenticate},
)
# Delegate server with a bearer present: it is the upstream token,
# so admit the session and move to the next target. Every oauth2
# sub-mode is terminal here (continue or raise) so no oauth2 server
# reaches the token_exchange / pass-through blocks below.
continue
# token_exchange (OBO): the caller supplied no subject token. Challenge at connect
# (transport level, where WWW-Authenticate survives) with the RFC 9728 resource_metadata

3
litellm/proxy/_experimental/mcp_server/tool_search.py
View file

@ -91,7 +91,7 @@ async def handle_mcp_tool_search(
from litellm.proxy._experimental.mcp_server.server import _list_mcp_tools
mcp_tools = await _list_mcp_tools(
mcp_listing = await _list_mcp_tools(
user_api_key_auth=user_api_key_dict,
mcp_servers=mcp_servers,
client_ip=client_ip,
@ -100,6 +100,7 @@ async def handle_mcp_tool_search(
oauth2_headers=oauth2_headers,
raw_headers=raw_headers,
)
mcp_tools = mcp_listing.tools
tools = [
{
"name": t.name,

12
litellm/proxy/_lazy_openapi_snapshot.json
View file

@ -7613,6 +7613,18 @@
],
"title": "Messages"
},
"metadata": {
"anyOf": [
{
"additionalProperties": true,
"type": "object"
},
{
"type": "null"
}
],
"title": "Metadata"
},
"text": {
"title": "Text",
"type": "string"

6
litellm/proxy/_types.py
View file

@ -1263,6 +1263,7 @@ class NewMCPServerRequest(LiteLLMPydanticObjectBase):
command: Optional[str] = None
args: List[str] = Field(default_factory=list)
env: Dict[str, str] = Field(default_factory=dict)
issuer: Optional[str] = None
authorization_url: Optional[str] = None
token_url: Optional[str] = None
registration_url: Optional[str] = None
@ -1368,6 +1369,7 @@ class UpdateMCPServerRequest(LiteLLMPydanticObjectBase):
command: Optional[str] = None
args: List[str] = Field(default_factory=list)
env: Dict[str, str] = Field(default_factory=dict)
issuer: Optional[str] = None
authorization_url: Optional[str] = None
token_url: Optional[str] = None
registration_url: Optional[str] = None
@ -2367,6 +2369,10 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase):
None,
description="If True, stores request messages and responses in spend logs. Default is False.",
)
disable_auto_add_proxy_admin_to_teams: bool | None = Field(
None,
description="By default, the user calling /team/new is automatically added to the new team as a team admin. If True, proxy admins are no longer auto-added; members explicitly listed in members_with_roles are unaffected. Default is False.",
)
maximum_spend_logs_retention_period: Optional[str] = Field(
None,
description="Maximum retention period for spend logs (e.g., '7d' for 7 days). Logs older than this will be deleted.",

27
litellm/proxy/auth/auth_checks.py
View file

@ -361,7 +361,11 @@ def _global_proxy_budget_check(global_proxy_spend: Optional[float], skip_budget_
and route != "/models"
):
if math.isfinite(litellm.max_budget) and global_proxy_spend > litellm.max_budget:
raise litellm.BudgetExceededError(current_cost=global_proxy_spend, max_budget=litellm.max_budget)
raise litellm.BudgetExceededError(
current_cost=global_proxy_spend,
max_budget=litellm.max_budget,
entity_type=Litellm_EntityType.PROXY.value,
)
_GUARDRAIL_MODIFICATION_KEYS: tuple = (
@ -523,6 +527,7 @@ async def common_checks(
request_headers=_safe_get_request_headers(request=request),
request_query_params=_safe_get_request_query_params(request=request),
llm_router=llm_router,
request=request,
)
if route in MODEL_DISCOVERY_ROUTES:
@ -648,6 +653,8 @@ async def common_checks(
current_cost=user_spend,
max_budget=user_budget,
message=f"ExceededBudget: User={user_object.user_id} over budget. Spend={user_spend}, Budget={user_budget}",
entity_type=Litellm_EntityType.USER.value,
entity_id=user_object.user_id,
)
# Each scope reads a distinct counter key with no cross-scope ordering
@ -1093,6 +1100,8 @@ async def _check_end_user_budget(
current_cost=end_user_spend,
max_budget=end_user_budget,
message=f"ExceededBudget: End User={end_user_obj.user_id} over budget. Spend={end_user_spend}, Budget={end_user_budget}",
entity_type=Litellm_EntityType.END_USER.value,
entity_id=end_user_obj.user_id,
)
@ -3552,6 +3561,8 @@ async def _virtual_key_max_budget_check(
current_cost=spend,
max_budget=valid_token.max_budget,
message=f"Budget has been exceeded! Key={key_descriptor} Current cost: {spend}, Max budget: {valid_token.max_budget}",
entity_type=Litellm_EntityType.KEY.value,
entity_id=valid_token.token,
)
@ -3593,6 +3604,8 @@ async def _virtual_key_multi_budget_check(
f"ExceededBudget: Key over {w['budget_duration']} budget. "
f"Spend=${window_spend:.4f}, Limit=${w['max_budget']:.2f}"
),
entity_type=Litellm_EntityType.KEY.value,
entity_id=valid_token.token,
)
@ -3824,6 +3837,8 @@ async def _check_team_member_budget(
current_cost=team_member_spend,
max_budget=team_member_budget,
message=f"Budget has been exceeded! User={valid_token.user_id} in Team={team_object.team_id} Current cost: {team_member_spend}, Max budget: {team_member_budget}",
entity_type=Litellm_EntityType.TEAM_MEMBER.value,
entity_id=f"{valid_token.user_id}:{team_object.team_id}",
)
@ -3923,6 +3938,8 @@ async def _team_max_budget_check(
current_cost=spend,
max_budget=team_object.max_budget,
message=f"Budget has been exceeded! Team={team_object.team_id} Current cost: {spend}, Max budget: {team_object.max_budget}",
entity_type=Litellm_EntityType.TEAM.value,
entity_id=team_object.team_id,
)
@ -3960,6 +3977,8 @@ async def _team_multi_budget_check(
f"ExceededBudget: Team={team_object.team_id} over {w['budget_duration']} budget. "
f"Spend=${window_spend:.4f}, Limit=${w['max_budget']:.2f}"
),
entity_type=Litellm_EntityType.TEAM.value,
entity_id=team_object.team_id,
)
@ -4081,6 +4100,8 @@ async def _project_max_budget_check(
current_cost=project_object.spend,
max_budget=max_budget,
message=f"Budget has been exceeded! Project={project_object.project_id} Current cost: {project_object.spend}, Max budget: {max_budget}",
entity_type=Litellm_EntityType.PROJECT.value,
entity_id=project_object.project_id,
)
@ -4269,6 +4290,8 @@ async def _organization_max_budget_check(
current_cost=org_spend,
max_budget=org_max_budget,
message=f"Budget has been exceeded! Organization={org_id} Current cost: {org_spend}, Max budget: {org_max_budget}",
entity_type=Litellm_EntityType.ORGANIZATION.value,
entity_id=org_id,
)
@ -4326,6 +4349,8 @@ async def _tag_max_budget_check(
current_cost=tag_spend,
max_budget=tag_object.litellm_budget_table.max_budget,
message=f"Budget has been exceeded! Tag={tag_name} Current cost: {tag_spend}, Max budget: {tag_object.litellm_budget_table.max_budget}",
entity_type=Litellm_EntityType.TAG.value,
entity_id=tag_name,
)

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

@ -10,10 +10,13 @@ from fastapi import HTTPException, Request, status
import litellm
from litellm import Router, provider_list
from litellm._logging import verbose_proxy_logger
from litellm.constants import STANDARD_CUSTOMER_ID_HEADERS
from litellm.constants import MINIMUM_CUSTOM_KEY_LENGTH, STANDARD_CUSTOMER_ID_HEADERS
from litellm.litellm_core_utils.safe_json_loads import safe_json_loads
from litellm.litellm_core_utils.url_utils import SSRFError, validate_url
from litellm.proxy._types import *
from litellm.types.passthrough_endpoints.pass_through_endpoints import (
LITELLM_PASS_THROUGH_ENDPOINT_MARKER,
)
from litellm.types.router import CONFIGURABLE_CLIENTSIDE_AUTH_PARAMS
from litellm.types.utils import CustomPricingLiteLLMParams
@ -1482,13 +1485,50 @@ def _format_model_candidates(
return candidates
def _request_dispatched_to_pass_through_endpoint(request: Request | None) -> bool:
"""Whether FastAPI resolved this request to a user-defined pass-through handler.
Reads the marker set by ``create_pass_through_route`` off the dispatched endpoint
(``request.scope["endpoint"]``). Because routing has already run by the time auth
dependencies execute, this reflects the handler that actually serves the request:
a custom path colliding with a built-in route resolves to the built-in handler,
which carries no marker, so model-access checks are never wrongly skipped.
"""
if request is None:
return False
scope = getattr(request, "scope", None)
if not isinstance(scope, dict):
return False
endpoint = scope.get("endpoint")
# Identity check against True (not truthiness): the marker is set to the literal
# True, and this keeps a spec'd Mock request (whose attribute access yields truthy
# child mocks) from being misread as a pass-through dispatch.
return getattr(endpoint, LITELLM_PASS_THROUGH_ENDPOINT_MARKER, False) is True
def get_model_from_request(
request_data: dict,
route: str,
request_headers: Optional[Mapping[str, Any]] = None,
request_query_params: Optional[Mapping[str, Any]] = None,
llm_router: Optional[Router] = None,
request: Request | None = None,
) -> Optional[Union[str, List[str]]]:
"""Resolve the model(s) a request targets, for model-access and budget checks.
Returns ``None`` when the request was dispatched to a user-defined pass-through
endpoint: its body is forwarded verbatim to the configured upstream, so a
``model`` field there names an upstream model, not a LiteLLM-managed one, and
enforcing key/team model allowlists against it would reject valid requests. The
check reads the FastAPI-resolved endpoint (``request.scope["endpoint"]``), not the
request path, so a custom path that collides with a built-in route never
suppresses model-access checks: on a collision the built-in handler is dispatched
and does not carry the marker. Built-in provider passthrough routes
(``/vertex_ai``, ``/gemini``, ...) are separate handlers and keep model enforcement.
"""
if _request_dispatched_to_pass_through_endpoint(request):
return None
candidates = _extract_model_candidates_from_request(
request_data=request_data,
route=route,
@ -1533,4 +1573,6 @@ def get_model_from_request(
def abbreviate_api_key(api_key: str) -> str:
if len(api_key) < MINIMUM_CUSTOM_KEY_LENGTH:
return "sk-..."
return f"sk-...{api_key[-4:]}"

13
litellm/proxy/auth/user_api_key_auth.py
View file

@ -162,6 +162,7 @@ def _get_model_from_request_context(
request_headers=_safe_get_request_headers(request=request),
request_query_params=_safe_get_request_query_params(request=request),
llm_router=llm_router,
request=request,
)
@ -1797,6 +1798,8 @@ async def _user_api_key_auth_builder(
raise litellm.BudgetExceededError(
current_cost=team_member_spend,
max_budget=team_member_budget,
entity_type=Litellm_EntityType.TEAM_MEMBER.value,
entity_id=f"{valid_token.user_id}:{valid_token.team_id}",
)
# Check 3. If token is expired
@ -1994,16 +1997,6 @@ async def _user_api_key_auth_builder(
raise HTTPException(401, detail="Invalid API key, no token associated")
api_key = valid_token.token
# Add hashed token to cache
asyncio.create_task(
_cache_key_object(
hashed_token=api_key,
user_api_key_obj=valid_token,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
)
)
valid_token_dict = valid_token.model_dump(exclude_none=True)
valid_token_dict.pop("token", None)
# budget_throttle_pct is excluded from model_dump (it must not leak

103
litellm/proxy/client/cli/README.md
View file

@ -471,6 +471,109 @@ The token minted by `lite login` is a short-lived, per-session agent credential,
The credential is short-lived by design (default 24h, configurable via `LITELLM_CLI_JWT_EXPIRATION_HOURS`); run `lite login` again to refresh it, which also re-reads your latest team and user settings. It does not appear in the Keys UI and cannot be rotated or revoked mid-session. `lite auth print-token` (usable as Claude Code's `apiKeyHelper`) prints it while it's still fresh and fails once it expires -- there is no silent renewal, so a long-running session needs a fresh `lite login` once a day. `lite claude`, `lite codex`, and `lite opencode` work with it on a default deployment; `EXPERIMENTAL_UI_LOGIN` is not required. If you need a long-lived, rotatable key that shows up in the Keys UI, create a dedicated virtual key in the dashboard and pass it via `--api-key` or `LITELLM_PROXY_API_KEY` instead.
### Route Every Claude Code Session Through the Proxy
`lite claude` wraps a single invocation, but `lite up` goes further: it patches `~/.claude/settings.json`, Claude Code's own config file, so that every Claude Code session started afterward -- from any terminal, launched normally with just `claude`, no wrapper needed -- routes through your LiteLLM proxy. It sets `env.ANTHROPIC_BASE_URL` to the proxy URL and `apiKeyHelper` to a `lite auth print-token` invocation, drops any stray static `ANTHROPIC_API_KEY` so the helper-issued token wins, and leaves every other setting in the file untouched. It backs up the original file before patching it.
Two things need to already be true: you've run `lite login`, since the apiKeyHelper depends on that stored token, and the proxy is already reachable, since `lite up` does not start one for you.
```bash
lite login
litellm --config litellm/proxy/dev_config.yaml &
lite up
```
`lite up` runs in the foreground and blocks. Press Ctrl-C to stop it, which restores the original settings file and exits. If the process is ever killed uncleanly instead -- `kill -9`, a crash -- the settings file is left patched, and `lite down` is the manual recovery path: run it at any later point to restore from the same backup.
This is a one-time file patch and restore, not a live traffic interceptor. A Claude Code session already running before `lite up` started keeps whatever `ANTHROPIC_BASE_URL` and token it loaded at its own startup, and a session still running when `lite up` stops keeps routing through the proxy until it exits; only sessions *started* while the patch is in effect are affected, and only *new* sessions after a restore go back to Anthropic directly.
Cursor is not supported: it has no equivalent file-based config to hot-patch this way, since its model routing lives in its own app storage and is configured through its GUI.
### QA Complexity-Based Auto-Routing Against Your Real Proxy
`lite autoroute` lets you try LiteLLM's complexity-based auto-routing -- picking a cheaper or more expensive model depending on how complex a prompt looks -- against models your key already has access to on your real, running proxy, without editing that proxy's `config.yaml` and without any real request ever bypassing it. It builds a second, throwaway proxy locally that forwards every request back to your real proxy, and points Claude Code at that local proxy for the duration of the session.
#### Install the CLI
`lite autoroute up` builds and runs a throwaway litellm proxy locally, so unlike the rest of this CLI it needs the proxy server runtime, not just the thin `litellm[cli]` client. Install `litellm[proxy]` (which ships the `lite` command too) with a single curl command -- no existing Python tooling required, `uv` is bootstrapped automatically if missing:
```bash
curl -fsSL https://raw.githubusercontent.com/BerriAI/litellm/main/scripts/install.sh | sh
```
To QA an unreleased branch or commit instead of the latest PyPI release, set `LITELLM_CLI_REF`:
```bash
curl -fsSL https://raw.githubusercontent.com/BerriAI/litellm/<branch-or-commit>/scripts/install.sh | \
LITELLM_CLI_REF=<branch-or-commit> sh
```
The thin `scripts/install-cli.sh` installs only `litellm[cli]`, which is enough for `lite login`, `lite claude`, and `lite up`, but not for `lite autoroute up`; running it against a `litellm[cli]` install fails fast with a message telling you to install the proxy runtime.
Point the CLI at your real proxy and key before running any `lite model-groups` or `lite autoroute` command -- like every other command in this CLI, they read `LITELLM_PROXY_URL`/`LITELLM_PROXY_API_KEY` (or `--base-url`/`--api-key`), no `lite login` required:
```bash
export LITELLM_PROXY_URL=http://localhost:4000
export LITELLM_PROXY_API_KEY=sk-...
```
#### List Your Accessible Model Groups
```bash
lite model-groups list [--format table|json]
```
Lists the model groups your key can reach on the proxy, via `/model_group/info`, along with each group's mode (`chat`, `embedding`, etc.) and per-token pricing. This is also what `lite autoroute configure` uses internally to discover what it can offer you.
#### Configure the Auto-Router
```bash
lite autoroute configure
```
An interactive wizard. It runs the same model-group discovery as above, splits the results into chat-capable and embedding-capable pools, and asks you to assign one or more models from the chat pool to each of the four complexity tiers -- SIMPLE, MEDIUM, COMPLEX, REASONING. Each tier's picker is a type-to-filter fuzzy search (fzf-style) rather than a scrollable numbered list, so it stays usable even with hundreds of model groups: type a substring to narrow the list, tab to toggle a model into the selection, enter to confirm (assigning more than one model to a tier is exactly when this matters -- complexity_router picks randomly among a tier's pool per request, and adaptive mode specifically depends on having more than one candidate to choose from). From there it optionally offers: classifying prompt complexity with an LLM (again picked from your discovered pool) instead of the free built-in heuristic scorer, semantic keyword matching for tier assignment (needs an embedding model from the pool), and adaptive (bandit-based) selection layered on top of tiering.
The wizard writes the result to `~/.litellm/autorouter/config.yaml` with `0600` permissions, since the file embeds your real proxy API key. Every model referenced anywhere in that config -- tier targets, the classifier model, the embedding model -- becomes its own `litellm_proxy/<model-name>` deployment whose `api_base` and `api_key` point back at your real proxy. That is the trick that keeps your real proxy's config untouched: every actual network call this generates, whether it is the routed completion, an LLM-classifier call, or an embedding call, forwards transparently through your real, already-running proxy with your real key.
You do not need to tell Claude Code to request `autorouter` by name yourself: `lite autoroute up` also sets `ANTHROPIC_DEFAULT_SONNET_MODEL`, `ANTHROPIC_DEFAULT_HAIKU_MODEL`, and `ANTHROPIC_DEFAULT_OPUS_MODEL` to `autorouter` in `~/.claude/settings.json`, so every one of Claude Code's own model tiers requests it directly regardless of `/model` or whatever it defaults to otherwise. (A bare `model_name: "*"` deployment looks like the obvious way to catch any request instead, but litellm's Router looks up auto-router deployments by the literal requested model string with no wildcard resolution, so a `"*"` entry would never actually match real traffic -- these env var overrides are what makes it work.)
You must run `configure` at least once before `up`; running `up` first fails with a clear error telling you to configure first.
#### Launch the Ephemeral Auto-Router Proxy
```bash
lite autoroute up
```
Starts a local, throwaway litellm proxy on a random free port, running the config `configure` generated, with a freshly-minted random API key baked in for this session only (your real proxy key never leaves the generated config -- it only appears there, forwarding to your real proxy). It waits for the ephemeral proxy to report healthy, then patches `~/.claude/settings.json` the same way `lite up` does, except with a static `ANTHROPIC_AUTH_TOKEN` env var instead of an `apiKeyHelper`, since this key is short-lived and self-issued rather than something needing SSO refresh. Any `claude` session started afterward, from any terminal, routes through the ephemeral proxy.
`lite autoroute up` runs in the foreground and streams the ephemeral proxy's own log file into your terminal, so you can watch its routing decisions -- which tier and model got picked for each request -- as you use Claude Code normally. Press Ctrl-C (or send SIGTERM) to stop it; this kills the child proxy process and restores your original Claude Code settings, in that order.
#### Recover From an Unclean Shutdown
```bash
lite autoroute down
```
If the `lite autoroute up` process dies uncleanly -- `kill -9`, a crash -- rather than being stopped with Ctrl-C, `down` is the manual recovery path: it kills any leftover ephemeral proxy process found via a recorded pid file and restores Claude Code's settings from whatever backup is on disk.
#### Example
```bash
lite autoroute configure
lite autoroute up
# use Claude Code as normal in another terminal; routing decisions stream live
lite autoroute down # only needed if `up` was killed uncleanly instead of Ctrl-C'd
```
#### Caveats
Adaptive mode's learned state does not persist across `lite autoroute up` sessions -- there is no local database, so every session starts adaptive selection cold. A Claude Code session already running before `up` started, or still running when it stops, keeps whatever settings it loaded at its own startup; like `lite up`, this is a one-time file patch and restore, not a live traffic interceptor. Only Claude Code is supported, for the same reason as `lite up`: no other supported agent (for example Cursor) has an equivalent hot-patchable config file.
A session that outlives `up` (or is still running the moment you stop it) keeps sending requests, master key included, to that now-freed loopback port until you restart it. Once the ephemeral proxy process exits, nothing stops another local account on the same machine from binding that same port and receiving those requests instead -- unlike `lite up`'s `apiKeyHelper`, which is re-resolved per request, `autoroute`'s master key is a static value, so whoever receives them gets a live-looking token along with the prompt content. Restart any Claude Code session before you consider the machine clean, run `lite autoroute down` promptly rather than leaving a stopped session's settings patched, and do not run `lite autoroute up` on a shared or multi-tenant host.
Do not run `lite up` and `lite autoroute up` at the same time. Each patches `~/.claude/settings.json` and keeps its own separate backup, with no coordination between them: whichever one you stop or crash out of last is the one whose backup gets restored, which can silently leave the *other* mode's settings (a static master key and a now-dead loopback URL, or a stale `apiKeyHelper`) active. Run `lite down` or `lite autoroute down` (whichever applies) before switching to the other mode.
## Environment Variables
The CLI respects the following environment variables:

5
litellm/proxy/client/cli/commands/agents.py
View file

@ -212,7 +212,7 @@ def _is_interactive() -> bool:
return sys.stdin.isatty()
def _resolve_api_key(ctx: click.Context) -> str:
def resolve_api_key(ctx: click.Context) -> str:
base_url = ctx.obj["base_url"]
api_key = ctx.obj.get("api_key")
if api_key:
@ -238,7 +238,7 @@ _SKIP_VERIFY_HELP = "Skip the pre-launch key check against the proxy."
def _launch(ctx: click.Context, binary: str, args: Sequence[str], *, skip_verify: bool) -> None:
base_url = ctx.obj["base_url"]
started_interactive = _is_interactive()
api_key = _resolve_api_key(ctx)
api_key = resolve_api_key(ctx)
display_name, _ = agent_profile(binary)
click.echo(f"litellm: routing {display_name} through proxy at {base_url.rstrip('/')}")
@ -288,5 +288,6 @@ __all__ = [
"agent_launch_args",
"verify_proxy_key",
"agent_profile",
"resolve_api_key",
"AgentRunError",
]

139
litellm/proxy/client/cli/commands/auth.py
View file

@ -72,7 +72,7 @@ def display_teams_table(teams: List[Dict[str, Any]]) -> None:
console = Console()
if not teams:
console.print("❌ No teams found for your user.")
console.print("No teams found for your user.")
return
table = Table(title="Available Teams")
@ -162,7 +162,7 @@ def display_interactive_team_selection(teams: List[Dict[str, Any]], selected_ind
# Clear the screen using Rich's method
console.clear()
console.print("🎯 Select a Team (Use ↑↓ arrows, Enter to select, 'q' to skip):\n")
console.print("Select a Team (Use up/down arrows, Enter to select, 'q' to skip):\n")
for i, team in enumerate(teams):
team_alias = team.get("team_alias") or "N/A"
@ -184,7 +184,7 @@ def display_interactive_team_selection(teams: List[Dict[str, Any]], selected_ind
# Highlight the selected item
if i == selected_index:
console.print(f"➤ [bold cyan]{team_alias}[/bold cyan] ({team_id})")
console.print(f"> [bold cyan]{team_alias}[/bold cyan] ({team_id})")
console.print(f" Models: [yellow]{models_str}[/yellow]")
console.print(f" Budget: [blue]{budget_str}[/blue]\n")
else:
@ -220,15 +220,13 @@ def prompt_team_selection(teams: List[Dict[str, Any]]) -> Optional[Dict[str, Any
# Clear screen and show selection
console = Console()
console.clear()
click.echo(
f"✅ Selected team: {selected_team.get('team_alias', 'N/A')} ({selected_team.get('team_id')})"
)
click.echo(f"Selected team: {selected_team.get('team_alias', 'N/A')} ({selected_team.get('team_id')})")
return selected_team
elif key == "quit" or key == "escape":
# Clear screen
console = Console()
console.clear()
click.echo("ℹ️ Team selection skipped.")
click.echo("Team selection skipped.")
return None
elif key is None:
# If we can't get key input, fall back to simple selection
@ -237,7 +235,7 @@ def prompt_team_selection(teams: List[Dict[str, Any]]) -> Optional[Dict[str, Any
except KeyboardInterrupt:
console = Console()
console.clear()
click.echo("\n❌ Team selection cancelled.")
click.echo("\nTeam selection cancelled.")
return None
except Exception:
# If interactive mode fails, fall back to simple selection
@ -265,18 +263,40 @@ def prompt_team_selection_fallback(
if 0 <= index < len(teams):
selected_team = teams[index]
click.echo(
f"\n✅ Selected team: {selected_team.get('team_alias', 'N/A')} ({selected_team.get('team_id')})"
f"\nSelected team: {selected_team.get('team_alias', 'N/A')} ({selected_team.get('team_id')})"
)
return selected_team
else:
click.echo(f"❌ Invalid selection. Please enter a number between 1 and {len(teams)}")
click.echo(f"Invalid selection. Please enter a number between 1 and {len(teams)}")
except ValueError:
click.echo("❌ Invalid input. Please enter a number or 'skip'")
click.echo("Invalid input. Please enter a number or 'skip'")
except KeyboardInterrupt:
click.echo("\n❌ Team selection cancelled.")
click.echo("\nTeam selection cancelled.")
return None
def _response_error_detail(response: requests.Response) -> str | None:
try:
body = response.json()
except ValueError:
return None
detail = body.get("detail") if isinstance(body, dict) else None
if isinstance(detail, str) and detail:
return detail
return None
def _polling_error_message(response: requests.Response) -> str:
detail = _response_error_detail(response)
if detail:
return f"Polling error: HTTP {response.status_code}: {detail}"
return f"Polling error: HTTP {response.status_code}"
def _is_permanent_polling_error(status_code: int) -> bool:
return 400 <= status_code < 500 and status_code != 429
# Polling-based authentication - no local server needed
def _poll_for_ready_data(
url: str,
@ -308,8 +328,14 @@ def _poll_for_ready_data(
click.echo(pending_message)
elif other_status_message and other_status_log_every > 0 and attempt % other_status_log_every == 0:
click.echo(other_status_message)
elif _is_permanent_polling_error(response.status_code):
detail = _response_error_detail(response)
raise ValueError(
f"The proxy rejected the login session with HTTP {response.status_code}"
+ (f": {detail}" if detail else f" and no error detail (from {url})")
)
elif http_error_log_every > 0 and attempt % http_error_log_every == 0:
click.echo(f"Polling error: HTTP {response.status_code}")
click.echo(_polling_error_message(response))
except requests.RequestException as e:
if connection_error_log_every > 0 and attempt % connection_error_log_every == 0:
click.echo(f"Connection error (will retry): {e}")
@ -342,12 +368,45 @@ def _normalize_teams(teams, team_details):
def _start_cli_sso_flow(base_url: str) -> Dict[str, Any]:
response = requests.post(f"{base_url}/sso/cli/start", timeout=10)
response.raise_for_status()
data = response.json()
start_url = f"{base_url}/sso/cli/start"
try:
response = requests.post(start_url, timeout=10)
except requests.RequestException as e:
raise ValueError(
f"Could not reach the proxy at {start_url}: {e}. "
"Check that the proxy is running and that --base-url points at it."
) from e
if response.status_code in (404, 405):
raise ValueError(
f"POST {start_url} returned HTTP {response.status_code}. "
"Either --base-url is wrong, or the proxy is older than this CLI and does not support "
"the CLI SSO login flow; upgrade the proxy or use a CLI version that matches it."
)
if response.status_code != 200:
detail = _response_error_detail(response)
raise ValueError(
f"Starting CLI login failed: HTTP {response.status_code} from {start_url}"
+ (f": {detail}" if detail else "")
)
try:
data = response.json()
except ValueError:
content_type = response.headers.get("content-type", "unknown")
raise ValueError(
f"The proxy returned a non-JSON response from {start_url} (content-type: {content_type}). "
"A proxy, load balancer, or auth gateway in front of LiteLLM may be intercepting the request. "
f"Response starts with: {response.text[:200]!r}"
)
required_fields = ("login_id", "poll_secret", "user_code")
if not all(isinstance(data.get(field), str) for field in required_fields):
raise ValueError("Invalid CLI SSO start response")
missing_fields = tuple(field for field in required_fields if not isinstance(data.get(field), str))
if missing_fields:
raise ValueError(
f"The response from {start_url} is missing required field(s): {', '.join(missing_fields)}. "
"The proxy version may not match this CLI; upgrade whichever is older."
)
return data
@ -376,7 +435,7 @@ def _poll_for_authentication(base_url: str, key_id: str, poll_secret: str) -> Op
user_id = data.get("user_id")
normalized_teams: List[Dict[str, Any]] = _normalize_teams(teams, team_details)
if not normalized_teams:
click.echo("⚠️ No teams available for selection.")
click.echo("Warning: No teams available for selection.")
return None
# User has multiple teams - let them select
@ -396,7 +455,7 @@ def _poll_for_authentication(base_url: str, key_id: str, poll_secret: str) -> Op
"team_id": None, # Set by server in JWT
}
click.echo("❌ Team selection cancelled or JWT generation failed.")
click.echo("Team selection cancelled or JWT generation failed.")
return None
# JWT is ready (single team or team already selected)
@ -407,7 +466,7 @@ def _poll_for_authentication(base_url: str, key_id: str, poll_secret: str) -> Op
# Show which team was assigned
if team_id and len(teams) == 1:
click.echo(f"\n✅ Automatically assigned to team: {team_id}")
click.echo(f"\nAutomatically assigned to team: {team_id}")
if api_key:
return {
@ -433,19 +492,19 @@ def _handle_team_selection_during_polling(
The JWT token with the selected team, or None if selection was skipped
"""
if not teams:
click.echo("ℹ️ No teams found. You can create or join teams using the web interface.")
click.echo("No teams found. You can create or join teams using the web interface.")
return None
click.echo("\n" + "=" * 60)
click.echo("📋 Select a team for your CLI session...")
click.echo("Select a team for your CLI session...")
team_id = _render_and_prompt_for_team_selection(teams)
if not team_id:
click.echo("ℹ️ No team selected.")
click.echo("No team selected.")
return None
click.echo(f"\n🔄 Generating JWT for team: {team_id}")
click.echo(f"\nGenerating JWT for team: {team_id}")
poll_url = f"{base_url}/sso/cli/poll/{key_id}?team_id={team_id}"
data = _poll_for_ready_data(
@ -459,7 +518,7 @@ def _handle_team_selection_during_polling(
return None
jwt_token = data.get("key")
if jwt_token:
click.echo(f"✅ Successfully generated JWT for team: {team_id}")
click.echo(f"Successfully generated JWT for team: {team_id}")
return jwt_token
return None
@ -507,14 +566,14 @@ def _render_and_prompt_for_team_selection(teams: List[Dict[str, Any]]) -> Option
selected_team = teams[index]
team_id = str(selected_team.get("team_id"))
team_alias = selected_team.get("team_alias") or team_id
click.echo(f"\n✅ Selected team: {team_alias} ({team_id})")
click.echo(f"\nSelected team: {team_alias} ({team_id})")
return team_id
click.echo(f"❌ Invalid selection. Please enter a number between 1 and {len(teams)}")
click.echo(f"Invalid selection. Please enter a number between 1 and {len(teams)}")
except ValueError:
click.echo("❌ Invalid input. Please enter a number or 'skip'")
click.echo("Invalid input. Please enter a number or 'skip'")
except KeyboardInterrupt:
click.echo("\n❌ Team selection cancelled.")
click.echo("\nTeam selection cancelled.")
return None
@ -567,7 +626,7 @@ def login(ctx: click.Context):
}
)
click.echo("\n✅ Login successful!")
click.echo("\nLogin successful!")
click.echo(f"JWT Token: {api_key[:20]}...")
click.echo("You can now use the CLI without specifying --api-key")
@ -576,14 +635,18 @@ def login(ctx: click.Context):
show_commands()
return
else:
click.echo("❌ Authentication timed out. Please try again.")
click.echo("Authentication timed out. Please try again.")
click.echo(
"The proxy never reported the browser sign-in as finished. If you did complete it, "
"check the proxy logs for /sso/callback errors and confirm SSO is configured on the proxy."
)
return
except KeyboardInterrupt:
click.echo("\n❌ Authentication cancelled by user.")
click.echo("\nAuthentication cancelled by user.")
return
except Exception as e:
click.echo(f"❌ Authentication failed: {e}")
click.echo(f"Authentication failed: {e}")
return
@ -591,7 +654,7 @@ def login(ctx: click.Context):
def logout():
"""Logout and clear stored authentication"""
clear_token()
click.echo("✅ Logged out successfully. Authentication token cleared.")
click.echo("Logged out successfully. Authentication token cleared.")
@click.command(name="print-token")
@ -638,10 +701,10 @@ def whoami():
token_data = load_token()
if not token_data:
click.echo("❌ Not authenticated. Run 'lite login' to authenticate.")
click.echo("Not authenticated. Run 'lite login' to authenticate.")
return
click.echo("✅ Authenticated")
click.echo("Authenticated")
click.echo(f"User Email: {token_data.get('user_email', 'Unknown')}")
click.echo(f"User ID: {token_data.get('user_id', 'Unknown')}")
click.echo(f"User Role: {token_data.get('user_role', 'Unknown')}")
@ -652,7 +715,7 @@ def whoami():
click.echo(f"Token age: {age_hours:.1f} hours")
if age_hours > CLI_JWT_EXPIRATION_HOURS:
click.echo(f"⚠️ Warning: Token is more than {CLI_JWT_EXPIRATION_HOURS} hours old and may have expired.")
click.echo(f"Warning: Token is more than {CLI_JWT_EXPIRATION_HOURS} hours old and may have expired.")
@click.group(name="auth")

0
tests/e2e/claude_code/_builder_unit_tests/__init__.py → litellm/proxy/client/cli/commands/autoroute/__init__.py
View file

207
litellm/proxy/client/cli/commands/autoroute/commands.py Normal file
View file

@ -0,0 +1,207 @@
import atexit
import json
import secrets
import signal
import threading
from types import FrameType
import click
import yaml
from pydantic import JsonValue, TypeAdapter, ValidationError
from ..up import CLAUDE_SETTINGS_PATH, UpError, load_json_or_empty, restore_claude_settings, write_backup
from ..up import BackupRecord as ClaudeBackupRecord
from .process import (
AUTOROUTE_DIR,
CONFIG_PATH,
LOG_PATH,
PidRecord,
ProcessLaunchError,
allocate_free_port,
clear_pid_record,
is_running,
launch_proxy,
missing_proxy_runtime_modules,
poll_liveliness,
read_pid_record,
secure_create,
stream_log,
terminate,
write_pid_record,
)
from .settings import merge_claude_settings_static_token
from .wizard import run_configure_wizard
AUTOROUTE_BACKUP_PATH = AUTOROUTE_DIR / "claude_settings_backup.json"
_GENERATED_CONFIG_ADAPTER = TypeAdapter(dict[str, JsonValue])
def _mint_and_embed_master_key() -> str:
"""Generate a fresh key for this session and write it into the generated config.yaml.
Must go under general_settings, not litellm_settings -- the proxy server only ever
reads general_settings.master_key (proxy_server.py:4530) to authenticate requests. A
key placed under litellm_settings is silently ignored, leaving the ephemeral proxy with
no real auth: any request reaches it regardless of the token Claude Code sends.
"""
master_key = secrets.token_urlsafe(32)
with open(CONFIG_PATH, "r") as f:
try:
generated = _GENERATED_CONFIG_ADAPTER.validate_python(yaml.safe_load(f))
except (yaml.YAMLError, ValidationError):
raise click.ClickException(
f"{CONFIG_PATH} is empty or corrupt. Run `lite autoroute configure` again to regenerate it."
)
general_settings = generated.get("general_settings")
updated_settings: dict[str, JsonValue] = {
**(general_settings if isinstance(general_settings, dict) else {}),
"master_key": master_key,
}
updated: dict[str, JsonValue] = {**generated, "general_settings": updated_settings}
with secure_create(CONFIG_PATH) as f:
yaml.safe_dump(updated, f, sort_keys=False)
return master_key
@click.group(name="autoroute")
def autoroute_group() -> None:
"""QA complexity-based auto-routing against models your key can already use"""
@autoroute_group.command("configure")
@click.pass_context
def configure(ctx: click.Context) -> None:
"""Discover accessible models and generate an ephemeral auto-router config"""
run_configure_wizard(ctx)
@autoroute_group.command("up")
def up() -> None:
"""Launch the ephemeral auto-router proxy and route Claude Code through it"""
if not CONFIG_PATH.exists():
raise click.ClickException("No config found. Run `lite autoroute configure` first.")
missing = missing_proxy_runtime_modules()
if missing:
raise click.ClickException(
"lite autoroute up launches a local litellm proxy, which needs the proxy runtime that the "
f"thin `litellm[cli]` install does not include (missing: {', '.join(missing)}). Install the "
"proxy runtime with `uv tool install --force 'litellm[proxy]'`, or to QA a branch, "
"`curl -fsSL https://raw.githubusercontent.com/BerriAI/litellm/<branch>/scripts/install.sh | "
"LITELLM_CLI_REF=<branch> sh`."
)
try:
existing_pid = read_pid_record()
except UpError as e:
raise click.ClickException(str(e))
if existing_pid is not None and is_running(existing_pid.pid):
raise click.ClickException(
"An ephemeral proxy is already running (lite autoroute up looks already active). "
"Run `lite autoroute down` first."
)
if AUTOROUTE_BACKUP_PATH.exists():
raise click.ClickException(
f"{AUTOROUTE_BACKUP_PATH} already exists -- `lite autoroute up` looks like it's already "
"running (or crashed without cleanup). Run `lite autoroute down` first."
)
master_key = _mint_and_embed_master_key()
port = allocate_free_port()
base_url = f"http://127.0.0.1:{port}"
process = launch_proxy(CONFIG_PATH, port, LOG_PATH)
write_pid_record(PidRecord(pid=process.pid, port=port, config_path=str(CONFIG_PATH), log_path=str(LOG_PATH)))
try:
poll_liveliness(base_url, LOG_PATH, process)
except ProcessLaunchError as e:
terminate(process.pid)
clear_pid_record()
raise click.ClickException(str(e))
try:
original_existed = CLAUDE_SETTINGS_PATH.exists()
original_settings = load_json_or_empty(CLAUDE_SETTINGS_PATH)
write_backup(
ClaudeBackupRecord(existed=original_existed, content=original_settings if original_existed else None),
AUTOROUTE_BACKUP_PATH,
)
merged = merge_claude_settings_static_token(original_settings, base_url, master_key)
CLAUDE_SETTINGS_PATH.parent.mkdir(parents=True, exist_ok=True)
with secure_create(CLAUDE_SETTINGS_PATH) as f:
json.dump(merged, f, indent=2)
except UpError as e:
terminate(process.pid)
clear_pid_record()
raise click.ClickException(str(e))
click.echo(f"litellm: ephemeral auto-router proxy up at {base_url} (pid {process.pid})")
click.echo("Claude Code sessions started now will route through it. Press Ctrl-C to stop and restore.")
stop_event = threading.Event()
restored = threading.Lock()
def _teardown() -> None:
if not restored.acquire(blocking=False):
return
terminate(process.pid)
clear_pid_record()
try:
restore_claude_settings(CLAUDE_SETTINGS_PATH, AUTOROUTE_BACKUP_PATH)
except UpError as e:
# Runs from atexit/a signal handler too, outside Click's own exception
# handling -- raising here would only produce an unhandled-exception
# warning on stderr, not a clean message.
click.echo(str(e), err=True)
return
click.echo("\nStopped ephemeral proxy and restored Claude Code settings.")
click.echo(
f"Restart any Claude Code session still open from this session, or another local account could "
f"bind the now-free port {port} and receive its requests. Do not use `lite autoroute up` on a "
f"shared or multi-tenant host."
)
def _handle_signal(_signum: int, _frame: FrameType | None) -> None:
stop_event.set()
signal.signal(signal.SIGINT, _handle_signal)
signal.signal(signal.SIGTERM, _handle_signal)
atexit.register(_teardown)
log_thread = threading.Thread(target=stream_log, args=(LOG_PATH, stop_event), daemon=True)
log_thread.start()
stop_event.wait()
_teardown()
@autoroute_group.command("down")
def down() -> None:
"""Restore Claude Code settings and stop a leftover ephemeral proxy, if any"""
try:
record: PidRecord | None = read_pid_record()
except UpError as e:
# down is the crash-recovery path -- a corrupt pid record must not block it; clear the
# unusable record and keep going rather than leaving the user with no way to clean up.
click.echo(f"{e} Clearing it and continuing cleanup.", err=True)
record = None
if record is not None and is_running(record.pid):
terminate(record.pid)
click.echo(f"Stopped leftover ephemeral proxy (pid {record.pid}).")
clear_pid_record()
try:
restored = restore_claude_settings(CLAUDE_SETTINGS_PATH, AUTOROUTE_BACKUP_PATH)
except UpError as e:
raise click.ClickException(str(e))
if restored is None:
click.echo("Nothing to restore.")
elif restored.existed:
click.echo(f"Restored {CLAUDE_SETTINGS_PATH} to its original contents.")
else:
click.echo(f"Removed {CLAUDE_SETTINGS_PATH} (it did not exist before `lite autoroute up`).")
__all__ = ["autoroute_group"]

249
litellm/proxy/client/cli/commands/autoroute/config.py Normal file
View file

@ -0,0 +1,249 @@
from typing import Literal, Union
from pydantic import BaseModel, ConfigDict, Field, JsonValue, TypeAdapter
TIER_NAMES: tuple[str, ...] = ("SIMPLE", "MEDIUM", "COMPLEX", "REASONING")
AUTOROUTER_MODEL_NAME = "autorouter"
class ConfigGenerationError(Exception):
"""Raised when an AutorouteConfig references a model the discovery step didn't find."""
class DiscoveredModel(BaseModel):
model_config = ConfigDict(frozen=True)
name: str
mode: str = "chat"
input_cost_per_token: float | None = None
output_cost_per_token: float | None = None
class _RawModelGroup(BaseModel):
model_config = ConfigDict(extra="ignore")
model_group: str
# Optional: some real deployments return an explicit `"mode": null` for models that
# were registered without a mode (seen for embedding models like voyage-4-large).
# ModelGroupInfo's own "chat" default (litellm/types/router.py) only applies when the
# key is missing entirely, not when it's present as null, so this must tolerate None.
mode: str | None = "chat"
input_cost_per_token: float | None = None
output_cost_per_token: float | None = None
_RAW_MODEL_GROUPS_ADAPTER = TypeAdapter(list[_RawModelGroup])
def parse_discovered_models(raw: list[JsonValue]) -> tuple[DiscoveredModel, ...]:
"""Validate a raw `/model_group/info` response into typed models."""
parsed = _RAW_MODEL_GROUPS_ADAPTER.validate_python(raw)
return tuple(
DiscoveredModel(
name=group.model_group,
# A null mode means the server genuinely doesn't know what this model does;
# "unknown" (rather than guessing "chat") keeps it out of both chat_models()
# and embedding_models() instead of risking a wrong-mode deployment.
mode=group.mode or "unknown",
input_cost_per_token=group.input_cost_per_token,
output_cost_per_token=group.output_cost_per_token,
)
for group in parsed
)
def chat_models(models: tuple[DiscoveredModel, ...]) -> tuple[DiscoveredModel, ...]:
return tuple(m for m in models if m.mode == "chat")
def embedding_models(models: tuple[DiscoveredModel, ...]) -> tuple[DiscoveredModel, ...]:
return tuple(m for m in models if m.mode == "embedding")
class HeuristicClassifier(BaseModel):
model_config = ConfigDict(frozen=True)
kind: Literal["heuristic"] = "heuristic"
class LLMClassifier(BaseModel):
model_config = ConfigDict(frozen=True)
kind: Literal["llm"] = "llm"
model: str
timeout_ms: int = 3000
ClassifierChoice = Union[HeuristicClassifier, LLMClassifier]
class NoSemanticMatching(BaseModel):
model_config = ConfigDict(frozen=True)
kind: Literal["none"] = "none"
class KeywordTierRule(BaseModel):
model_config = ConfigDict(frozen=True)
keywords: tuple[str, ...]
tier: str
# Satisfies complexity_router's "semantic matching requires non-empty keyword_tier_rules"
# invariant with a sane starting point; the wizard lets the user override these per tier.
DEFAULT_KEYWORD_TIER_RULES: tuple[KeywordTierRule, ...] = (
KeywordTierRule(keywords=("hi", "hello", "thanks"), tier="SIMPLE"),
KeywordTierRule(keywords=("explain", "how does"), tier="MEDIUM"),
KeywordTierRule(keywords=("refactor", "implement", "debug"), tier="COMPLEX"),
KeywordTierRule(keywords=("step by step", "think through", "prove"), tier="REASONING"),
)
class SemanticMatching(BaseModel):
model_config = ConfigDict(frozen=True)
kind: Literal["semantic"] = "semantic"
embedding_model: str
match_threshold: float = 0.5
keyword_tier_rules: tuple[KeywordTierRule, ...] = DEFAULT_KEYWORD_TIER_RULES
SemanticMatchingChoice = Union[NoSemanticMatching, SemanticMatching]
class AutorouteConfig(BaseModel):
model_config = ConfigDict(frozen=True)
base_url: str
api_key: str
# Each tier maps to a pool of one or more models; complexity_router picks randomly among
# them per request (or, in adaptive mode, learns which to prefer within the pool).
tiers: dict[str, tuple[str, ...]]
default_model: str
classifier: ClassifierChoice = Field(default_factory=HeuristicClassifier)
semantic_matching: SemanticMatchingChoice = Field(default_factory=NoSemanticMatching)
adaptive: bool = False
def validate_config(config: AutorouteConfig, discovered: tuple[DiscoveredModel, ...]) -> None:
"""Raise ConfigGenerationError if config references a model discovery didn't return."""
chat_names: frozenset[str] = frozenset(m.name for m in chat_models(discovered))
embedding_names: frozenset[str] = frozenset(m.name for m in embedding_models(discovered))
for tier, models in config.tiers.items():
for model in models:
if model not in chat_names:
raise ConfigGenerationError(f"Tier {tier} references unknown chat model '{model}'")
if config.default_model not in chat_names:
raise ConfigGenerationError(f"default_model '{config.default_model}' is not a known chat model")
if isinstance(config.classifier, LLMClassifier) and config.classifier.model not in chat_names:
raise ConfigGenerationError(f"classifier model '{config.classifier.model}' is not a known chat model")
if (
isinstance(config.semantic_matching, SemanticMatching)
and config.semantic_matching.embedding_model not in embedding_names
):
raise ConfigGenerationError(
f"embedding model '{config.semantic_matching.embedding_model}' is not a known embedding model"
)
def _litellm_proxy_deployment(name: str, base_url: str, api_key: str) -> dict[str, JsonValue]:
return {
"model_name": name,
"litellm_params": {
"model": f"litellm_proxy/{name}",
"api_base": base_url,
"api_key": api_key,
},
}
def build_generated_model_list(config: AutorouteConfig) -> list[JsonValue]:
"""Build the model_list for the ephemeral proxy's config.yaml.
Every real model referenced anywhere (tier targets, classifier, embedding) is deduplicated
to exactly one `litellm_proxy/<name>` deployment forwarding to the customer's real proxy,
plus one `auto_router/complexity_router` deployment tying the tiers together.
"""
referenced_names = {model for models in config.tiers.values() for model in models}
referenced_names.add(config.default_model)
if isinstance(config.classifier, LLMClassifier):
referenced_names.add(config.classifier.model)
if isinstance(config.semantic_matching, SemanticMatching):
referenced_names.add(config.semantic_matching.embedding_model)
proxy_deployments = [
_litellm_proxy_deployment(name, config.base_url, config.api_key) for name in sorted(referenced_names)
]
complexity_router_config: dict[str, JsonValue] = {
"tiers": {tier: list(models) for tier, models in config.tiers.items()},
"default_model": config.default_model,
}
if isinstance(config.classifier, LLMClassifier):
complexity_router_config["classifier_type"] = "llm"
complexity_router_config["classifier_llm_config"] = {
"model": config.classifier.model,
"timeout_ms": config.classifier.timeout_ms,
}
if isinstance(config.semantic_matching, SemanticMatching):
complexity_router_config["semantic_keyword_matching"] = True
complexity_router_config["embedding_model"] = config.semantic_matching.embedding_model
complexity_router_config["match_threshold"] = config.semantic_matching.match_threshold
complexity_router_config["keyword_tier_rules"] = [
{"keywords": list(rule.keywords), "tier": rule.tier} for rule in config.semantic_matching.keyword_tier_rules
]
if config.adaptive:
complexity_router_config["adaptive"] = True
auto_router_litellm_params: dict[str, JsonValue] = {
"model": "auto_router/complexity_router",
"complexity_router_config": complexity_router_config,
}
# A bare "*" model_name looks like the obvious way to catch every request Claude Code
# might send regardless of which model it thinks it's using, but Router's auto-router
# registry is keyed by the literal requested model string (router.py:10711-10717), not
# resolved through pattern/wildcard matching first -- so a "*" entry here would only ever
# match a client that literally sends model="*", never an actual wildcard catch-all. Callers
# instead need to make Claude Code request this "autorouter" name directly (see
# ANTHROPIC_DEFAULT_*_MODEL in settings.py's merge_claude_settings_static_token).
return [
*proxy_deployments,
{"model_name": AUTOROUTER_MODEL_NAME, "litellm_params": auto_router_litellm_params},
]
def build_generated_proxy_config(config: AutorouteConfig, master_key: str) -> dict[str, JsonValue]:
"""Full config.yaml content for the ephemeral proxy, including its own auth key.
master_key must live under general_settings, not litellm_settings -- the proxy server
only ever reads general_settings.master_key (proxy_server.py:4530) to authenticate
requests; a key placed under litellm_settings is silently ignored, leaving the proxy
with no real auth at all.
"""
return {
"model_list": build_generated_model_list(config),
"general_settings": {"master_key": master_key},
}
__all__ = [
"AUTOROUTER_MODEL_NAME",
"TIER_NAMES",
"AutorouteConfig",
"ClassifierChoice",
"ConfigGenerationError",
"DEFAULT_KEYWORD_TIER_RULES",
"DiscoveredModel",
"HeuristicClassifier",
"KeywordTierRule",
"LLMClassifier",
"NoSemanticMatching",
"SemanticMatching",
"SemanticMatchingChoice",
"build_generated_model_list",
"build_generated_proxy_config",
"chat_models",
"embedding_models",
"parse_discovered_models",
"validate_config",
]

190
litellm/proxy/client/cli/commands/autoroute/process.py Normal file
View file

@ -0,0 +1,190 @@
import contextlib
import importlib.util
import json
import os
import signal
import socket
import subprocess
import sys
import threading
import time
from dataclasses import dataclass
from pathlib import Path
import click
import requests
from pydantic import TypeAdapter, ValidationError
from ..up import UpError, secure_create
AUTOROUTE_DIR = Path.home() / ".litellm" / "autorouter"
CONFIG_PATH = AUTOROUTE_DIR / "config.yaml"
LOG_PATH = AUTOROUTE_DIR / "proxy.log"
PID_RECORD_PATH = AUTOROUTE_DIR / "proxy.pid.json"
class ProcessLaunchError(Exception):
"""Raised when the ephemeral proxy subprocess fails to come up healthy."""
@dataclass(frozen=True, slots=True)
class PidRecord:
pid: int
port: int
config_path: str
log_path: str
_PID_RECORD_ADAPTER = TypeAdapter(PidRecord)
_PROXY_RUNTIME_MODULES: tuple[str, ...] = ("fastapi", "uvicorn", "backoff", "orjson", "websockets", "apscheduler")
def missing_proxy_runtime_modules() -> tuple[str, ...]:
"""Proxy-server modules that ``lite autoroute up`` needs but the thin CLI install lacks.
``launch_proxy`` runs the full ``litellm.proxy.proxy_cli`` server, whose dependencies live in
the ``proxy`` extra, not the ``cli`` extra that installs the ``lite`` command. On a thin
``litellm[cli]`` install the subprocess dies with a bare ``ModuleNotFoundError``; detecting the
gap here lets ``up`` fail with an actionable message instead.
"""
return tuple(name for name in _PROXY_RUNTIME_MODULES if importlib.util.find_spec(name) is None)
def allocate_free_port() -> int:
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
sock.bind(("127.0.0.1", 0))
return int(sock.getsockname()[1])
def launch_proxy(config_path: Path, port: int, log_path: Path) -> "subprocess.Popen[bytes]":
log_path.parent.mkdir(parents=True, exist_ok=True)
with open(log_path, "w") as log_file:
return subprocess.Popen(
[
sys.executable,
"-m",
"litellm.proxy.proxy_cli",
"--config",
str(config_path),
"--port",
str(port),
"--host",
"127.0.0.1",
],
stdout=log_file,
stderr=subprocess.STDOUT,
)
def _tail(log_path: Path, lines: int = 40) -> str:
if not log_path.exists():
return "(no log output captured)"
return "\n".join(log_path.read_text(errors="replace").splitlines()[-lines:])
def poll_liveliness(base_url: str, log_path: Path, process: "subprocess.Popen[bytes]", timeout: float = 30.0) -> None:
"""Poll /health/liveliness until it responds, the process dies, or timeout elapses."""
deadline = time.monotonic() + timeout
url = base_url.rstrip("/") + "/health/liveliness"
while time.monotonic() < deadline:
if process.poll() is not None:
raise ProcessLaunchError(
f"Ephemeral proxy exited early (code {process.returncode}). Last log lines:\n{_tail(log_path)}"
)
with contextlib.suppress(requests.RequestException):
if requests.get(url, timeout=2).status_code == 200:
return
time.sleep(0.5)
raise ProcessLaunchError(
f"Ephemeral proxy never became healthy within {timeout}s. Last log lines:\n{_tail(log_path)}"
)
def write_pid_record(record: PidRecord, path: Path | None = None) -> None:
resolved_path = path if path is not None else PID_RECORD_PATH
resolved_path.parent.mkdir(parents=True, exist_ok=True)
with open(resolved_path, "w") as f:
json.dump(
{"pid": record.pid, "port": record.port, "config_path": record.config_path, "log_path": record.log_path},
f,
indent=2,
)
def read_pid_record(path: Path | None = None) -> PidRecord | None:
resolved_path = path if path is not None else PID_RECORD_PATH
if not resolved_path.exists():
return None
with open(resolved_path, "r") as f:
content = f.read()
try:
return _PID_RECORD_ADAPTER.validate_json(content)
except ValidationError:
raise UpError(f"{resolved_path} contains invalid or unexpected JSON; cannot proceed safely.")
def clear_pid_record(path: Path | None = None) -> None:
resolved_path = path if path is not None else PID_RECORD_PATH
resolved_path.unlink(missing_ok=True)
def is_running(pid: int) -> bool:
try:
os.kill(pid, 0)
except ProcessLookupError:
return False
except PermissionError:
return True
return True
def terminate(pid: int, grace_period: float = 5.0) -> None:
"""Terminate a process by pid, escalating from SIGTERM to SIGKILL if needed."""
if not is_running(pid):
return
with contextlib.suppress(ProcessLookupError):
os.kill(pid, signal.SIGTERM)
deadline = time.monotonic() + grace_period
while time.monotonic() < deadline and is_running(pid):
time.sleep(0.2)
if is_running(pid):
with contextlib.suppress(ProcessLookupError):
os.kill(pid, signal.SIGKILL)
def stream_log(log_path: Path, stop_event: threading.Event) -> None:
"""Print new lines appended to log_path until stop_event is set. Blocks the calling thread."""
while not log_path.exists() and not stop_event.is_set():
time.sleep(0.1)
if stop_event.is_set() or not log_path.exists():
return
with open(log_path, "r") as f:
while not stop_event.is_set():
line = f.readline()
if line:
click.echo(line, nl=False)
else:
time.sleep(0.2)
__all__ = [
"AUTOROUTE_DIR",
"CONFIG_PATH",
"LOG_PATH",
"PID_RECORD_PATH",
"PidRecord",
"ProcessLaunchError",
"allocate_free_port",
"clear_pid_record",
"is_running",
"launch_proxy",
"missing_proxy_runtime_modules",
"poll_liveliness",
"read_pid_record",
"secure_create",
"stream_log",
"terminate",
"write_pid_record",
]

46
litellm/proxy/client/cli/commands/autoroute/settings.py Normal file
View file

@ -0,0 +1,46 @@
from pydantic import JsonValue
from .config import AUTOROUTER_MODEL_NAME
ENV_KEY = "env"
API_KEY_HELPER_KEY = "apiKeyHelper"
ANTHROPIC_API_KEY_KEY = "ANTHROPIC_API_KEY"
ANTHROPIC_AUTH_TOKEN_KEY = "ANTHROPIC_AUTH_TOKEN"
ANTHROPIC_BASE_URL_KEY = "ANTHROPIC_BASE_URL"
# Force every one of Claude Code's own model tiers to request the auto-router by name.
# Router's auto-router registry is keyed by the literal requested model string
# (litellm/router.py:10711-10717) with no wildcard/pattern resolution, so a bare "*"
# model_name can never work as a catch-all -- these overrides are what actually makes
# Claude Code send "autorouter" regardless of /model or its own version-specific defaults.
ANTHROPIC_DEFAULT_MODEL_ENV_KEYS = (
"ANTHROPIC_DEFAULT_SONNET_MODEL",
"ANTHROPIC_DEFAULT_HAIKU_MODEL",
"ANTHROPIC_DEFAULT_OPUS_MODEL",
)
def merge_claude_settings_static_token(
settings: dict[str, JsonValue], base_url: str, auth_token: str
) -> dict[str, JsonValue]:
"""Return a new settings dict wired to a local ephemeral proxy with a static token.
Unlike up.py's merge_claude_settings (which sets apiKeyHelper for a long-lived, real
remote proxy needing refreshable SSO tokens), this proxy is ephemeral and its key was just
minted for this session, so a plain env var is simpler and correct. Any existing
apiKeyHelper is cleared so it can't fight with the static token.
"""
raw_env = settings.get(ENV_KEY, {})
base_env = raw_env if isinstance(raw_env, dict) else {}
env: dict[str, JsonValue] = {
**base_env,
ANTHROPIC_BASE_URL_KEY: base_url.rstrip("/"),
ANTHROPIC_AUTH_TOKEN_KEY: auth_token,
**{key: AUTOROUTER_MODEL_NAME for key in ANTHROPIC_DEFAULT_MODEL_ENV_KEYS},
}
env.pop(ANTHROPIC_API_KEY_KEY, None)
merged: dict[str, JsonValue] = {**settings, ENV_KEY: env}
merged.pop(API_KEY_HELPER_KEY, None)
return merged
__all__ = ["merge_claude_settings_static_token"]

150
litellm/proxy/client/cli/commands/autoroute/wizard.py Normal file
View file

@ -0,0 +1,150 @@
import sys
from pathlib import Path
import click
import yaml
from InquirerPy import inquirer
from InquirerPy.base.control import Choice
from .... import Client
from .config import (
DEFAULT_KEYWORD_TIER_RULES,
TIER_NAMES,
AutorouteConfig,
ConfigGenerationError,
DiscoveredModel,
HeuristicClassifier,
KeywordTierRule,
LLMClassifier,
NoSemanticMatching,
SemanticMatching,
build_generated_model_list,
chat_models,
embedding_models,
parse_discovered_models,
validate_config,
)
from .process import CONFIG_PATH, secure_create
def _is_interactive() -> bool:
return sys.stdin.isatty()
def _fuzzy_pick(models: tuple[DiscoveredModel, ...], prompt_label: str, multiselect: bool) -> list[str]:
"""Type-to-filter picker over a (possibly huge) model pool, using InquirerPy's fzf-style fuzzy prompt.
A plain numbered table + typed index does not scale past a handful of models -- proxies with
hundreds of model groups made that interaction unusable. This lets the user narrow the pool by
typing a substring instead of scrolling/counting.
Assumes the caller already checked interactivity (run_configure_wizard does, once, up front) --
checking here too would check the wrong thing under test, where InquirerPy is driven through its
own injected input/output rather than the real process stdin.
"""
choices = [Choice(value=model.name, name=model.name) for model in models]
toggle_hint = "tab to toggle, " if multiselect else ""
while True:
result = inquirer.fuzzy(
message=f"{prompt_label}: type to filter, {toggle_hint}enter to confirm",
choices=choices,
multiselect=multiselect,
max_height="70%",
).execute()
selected = result if multiselect else [result]
if selected:
return selected
click.echo("Select at least one model.")
def _render_and_prompt_for_model(models: tuple[DiscoveredModel, ...], prompt_label: str) -> str:
return _fuzzy_pick(models, prompt_label, multiselect=False)[0]
def _render_and_prompt_for_models(models: tuple[DiscoveredModel, ...], prompt_label: str) -> tuple[str, ...]:
return tuple(_fuzzy_pick(models, prompt_label, multiselect=True))
def _parse_keywords(raw: str) -> tuple[str, ...]:
return tuple(keyword.strip() for keyword in raw.split(",") if keyword.strip())
def _prompt_for_keyword_tier_rules() -> tuple[KeywordTierRule, ...]:
"""Let the user supply the semantic-matching keywords per tier, since matching those
keywords against the request is the whole point of enabling it. Each prompt is prefilled
with the built-in default, so pressing enter keeps it."""
click.echo("\nEnter example keywords/phrases per tier (comma-separated); press enter to keep the default:")
defaults = {rule.tier: rule.keywords for rule in DEFAULT_KEYWORD_TIER_RULES}
def _rule_for(tier: str) -> KeywordTierRule:
default_keywords = defaults.get(tier, ())
raw = click.prompt(f" {tier} keywords", default=", ".join(default_keywords), show_default=True)
return KeywordTierRule(keywords=_parse_keywords(raw) or default_keywords, tier=tier)
return tuple(_rule_for(tier) for tier in TIER_NAMES)
def run_configure_wizard(ctx: click.Context) -> Path:
"""Discover the caller's accessible models, walk them through tier assignment, write config."""
base_url = ctx.obj["base_url"]
api_key = ctx.obj["api_key"]
client = Client(base_url=base_url, api_key=api_key)
raw_groups = client.model_groups.info()
if not isinstance(raw_groups, list):
raise click.ClickException(
f"Unexpected response from /model_group/info: expected a list, got {type(raw_groups).__name__}"
)
discovered = parse_discovered_models(raw_groups)
chat_pool = chat_models(discovered)
embedding_pool = embedding_models(discovered)
if not chat_pool:
raise click.ClickException("Your key has no chat-capable models available on this proxy.")
if not _is_interactive():
raise click.ClickException("`lite autoroute configure` requires an interactive terminal.")
click.echo("Assign model(s) to each complexity tier (from what your key can access):")
tiers = {tier: _render_and_prompt_for_models(chat_pool, tier) for tier in TIER_NAMES}
default_model = tiers["MEDIUM"][0]
classifier = HeuristicClassifier()
if click.confirm("\nUse an LLM classifier instead of the free heuristic scorer?", default=False):
classifier_model = _render_and_prompt_for_model(chat_pool, "LLM classifier")
classifier = LLMClassifier(model=classifier_model)
semantic_matching = NoSemanticMatching()
if embedding_pool and click.confirm("\nEnable semantic keyword matching?", default=False):
embedding_model = _render_and_prompt_for_model(embedding_pool, "semantic embeddings")
keyword_tier_rules = _prompt_for_keyword_tier_rules()
semantic_matching = SemanticMatching(embedding_model=embedding_model, keyword_tier_rules=keyword_tier_rules)
adaptive = click.confirm("\nEnable adaptive (bandit) selection on top of tiering?", default=False)
config = AutorouteConfig(
base_url=base_url,
api_key=api_key,
tiers=tiers,
default_model=default_model,
classifier=classifier,
semantic_matching=semantic_matching,
adaptive=adaptive,
)
try:
validate_config(config, discovered)
except ConfigGenerationError as e:
raise click.ClickException(str(e))
model_list = build_generated_model_list(config)
CONFIG_PATH.parent.mkdir(parents=True, exist_ok=True)
with secure_create(CONFIG_PATH) as f:
yaml.safe_dump({"model_list": model_list}, f, sort_keys=False)
click.echo(f"\nWrote {CONFIG_PATH}")
for tier, models in tiers.items():
click.echo(f" {tier}: {', '.join(models)}")
return CONFIG_PATH
__all__ = ["run_configure_wizard"]

2
litellm/proxy/client/cli/commands/chat.py
View file

@ -150,7 +150,7 @@ def chat(
f"Max Tokens: [yellow]{max_tokens or 'unlimited'}[/yellow]\n\n"
f"Type your messages and press Enter. Type '/quit' or '/exit' to end the session.\n"
f"Type '/help' for more commands.",
title="🤖 Chat Session",
title="Chat Session",
)
)

2
litellm/proxy/client/cli/commands/encryption.py
View file

@ -32,7 +32,7 @@ def migrate(ctx: click.Context, check_only: bool, dry_run: bool):
Requires the proxy to be started with
``general_settings.encryption_algorithm: aes-256-gcm``. Idempotent and
resumable — safe to re-run after an interruption.
resumable; safe to re-run after an interruption.
Examples:
litellm-proxy encryption migrate --check # attestation scan, no writes

4
litellm/proxy/client/cli/commands/keys.py
View file

@ -309,12 +309,12 @@ def _import_keys_to_destination(
imported_count += 1
key_alias = key.get("key_alias", "N/A")
click.echo(f"✓ Imported key: {key_alias}")
click.echo(f"Imported key: {key_alias}")
except Exception as e:
failed_count += 1
key_alias = key.get("key_alias", "N/A")
click.echo(f"✗ Failed to import key {key_alias}: {str(e)}", err=True)
click.echo(f"Failed to import key {key_alias}: {str(e)}", err=True)
return imported_count, failed_count

57
litellm/proxy/client/cli/commands/model_groups.py Normal file
View file

@ -0,0 +1,57 @@
from typing import Literal
import click
import rich
import rich.table
from ... import Client
def create_client(ctx: click.Context) -> Client:
return Client(base_url=ctx.obj["base_url"], api_key=ctx.obj["api_key"])
@click.group(name="model-groups")
def model_groups() -> None:
"""Inspect model groups your key can access on the proxy"""
@model_groups.command("list")
@click.option(
"--format",
"output_format",
type=click.Choice(["table", "json"]),
default="table",
help="Output format (table or json)",
)
@click.pass_context
def list_model_groups(ctx: click.Context, output_format: Literal["table", "json"]) -> None:
"""List model groups accessible to your key, with mode and pricing"""
client = create_client(ctx)
groups = client.model_groups.info()
if not isinstance(groups, list):
raise click.ClickException(
f"Unexpected response from /model_group/info: expected a list, got {type(groups).__name__}"
)
if output_format == "json":
rich.print_json(data=groups)
return
table = rich.table.Table(title="Accessible Model Groups")
table.add_column("Model", style="cyan")
table.add_column("Mode", style="green")
table.add_column("Input $/token", style="yellow")
table.add_column("Output $/token", style="yellow")
for group in groups:
table.add_row(
str(group.get("model_group", "")),
str(group.get("mode", "chat")),
str(group.get("input_cost_per_token", "")),
str(group.get("output_cost_per_token", "")),
)
rich.print(table)
__all__ = ["model_groups"]

20
litellm/proxy/client/cli/commands/teams.py
View file

@ -21,7 +21,7 @@ def display_teams_table(teams: List[Dict[str, Any]]) -> None:
console = Console()
if not teams:
console.print("❌ No teams found for your user.")
console.print("No teams found for your user.")
return
table = Table(title="Available Teams")
@ -91,10 +91,10 @@ def available(ctx: click.Context):
teams = client.teams.get_available()
if teams:
console = Console()
console.print("\n🎯 Available Teams to Join:")
console.print("\nAvailable Teams to Join:")
display_teams_table(teams)
else:
click.echo("ℹ️ No available teams to join.")
click.echo("No available teams to join.")
except requests.exceptions.HTTPError as e:
click.echo(f"Error: HTTP {e.response.status_code}", err=True)
error_body = e.response.json()
@ -113,7 +113,7 @@ def assign_key(ctx: click.Context, team_id: Optional[str]):
api_key = ctx.obj["api_key"]
if not api_key:
click.echo("❌ No API key found. Please login first using 'litellm login'")
click.echo("No API key found. Please login first using 'litellm login'")
raise click.Abort()
try:
@ -122,7 +122,7 @@ def assign_key(ctx: click.Context, team_id: Optional[str]):
teams = client.teams.list()
if not teams:
click.echo("❌ No teams found for your user.")
click.echo("No teams found for your user.")
return
# Use interactive selection from auth module
@ -133,14 +133,14 @@ def assign_key(ctx: click.Context, team_id: Optional[str]):
if selected_team:
team_id = selected_team.get("team_id")
else:
click.echo("❌ Operation cancelled.")
click.echo("Operation cancelled.")
return
# Update the key with the selected team
if team_id:
click.echo(f"\n🔄 Assigning your key to team: {team_id}")
click.echo(f"\nAssigning your key to team: {team_id}")
client.keys.update(key=api_key, team_id=team_id)
click.echo(f"✅ Successfully assigned key to team: {team_id}")
click.echo(f"Successfully assigned key to team: {team_id}")
# Show team details if available
teams = client.teams.list()
@ -148,9 +148,9 @@ def assign_key(ctx: click.Context, team_id: Optional[str]):
if team.get("team_id") == team_id:
models = team.get("models", [])
if models:
click.echo(f"🎯 You can now access models: {', '.join(models)}")
click.echo(f"You can now access models: {', '.join(models)}")
else:
click.echo("🎯 You can now access all available models")
click.echo("You can now access all available models")
break
except requests.exceptions.HTTPError as e:

283
litellm/proxy/client/cli/commands/up.py Normal file
View file

@ -0,0 +1,283 @@
import atexit
import contextlib
import json
import os
import shlex
import shutil
import signal
import sys
import threading
from dataclasses import dataclass
from pathlib import Path
from types import FrameType
from typing import IO, Iterator, Mapping
import click
from pydantic import JsonValue, TypeAdapter, ValidationError
from litellm.litellm_core_utils.cli_token_utils import is_cli_token_fresh
from .agents import AgentRunError, resolve_api_key, verify_proxy_key
from .auth import load_token, login
ENV_KEY = "env"
API_KEY_HELPER_KEY = "apiKeyHelper"
ANTHROPIC_BASE_URL_KEY = "ANTHROPIC_BASE_URL"
ANTHROPIC_API_KEY_KEY = "ANTHROPIC_API_KEY"
CLAUDE_SETTINGS_PATH = Path.home() / ".claude" / "settings.json"
BACKUP_PATH = Path.home() / ".litellm" / "claude_settings_backup.json"
class UpError(Exception):
"""Raised for any user-actionable failure while starting/stopping interception."""
@dataclass(frozen=True, slots=True)
class BackupRecord:
"""Snapshot of ~/.claude/settings.json taken right before `lite up` patches it."""
existed: bool
content: dict[str, JsonValue] | None
_SETTINGS_ADAPTER = TypeAdapter(dict[str, JsonValue])
_BACKUP_RECORD_ADAPTER = TypeAdapter(BackupRecord)
def load_json_or_empty(path: Path) -> dict[str, JsonValue]:
if not path.exists():
return {}
with open(path, "r") as f:
content = f.read()
if not content.strip():
return {}
try:
return _SETTINGS_ADAPTER.validate_json(content)
except ValidationError:
raise UpError(f"{path} contains invalid JSON (or its root is not an object); cannot proceed safely.")
def merge_claude_settings(
settings: Mapping[str, JsonValue], base_url: str, api_key_helper: str
) -> dict[str, JsonValue]:
"""Return a new settings dict wired to route Claude Code through the proxy.
Only env.ANTHROPIC_BASE_URL and the top-level apiKeyHelper are overridden; a
stray env.ANTHROPIC_API_KEY is dropped so it cannot outrank the helper-issued
token (same reasoning as build_agent_env in agents.py). Every other key is
preserved untouched.
"""
raw_env = settings.get(ENV_KEY, {})
base_env = raw_env if isinstance(raw_env, dict) else {}
env = {**base_env, ANTHROPIC_BASE_URL_KEY: base_url.rstrip("/")}
env.pop(ANTHROPIC_API_KEY_KEY, None)
return {**settings, ENV_KEY: env, API_KEY_HELPER_KEY: api_key_helper}
@contextlib.contextmanager
def secure_create(path: Path) -> Iterator[IO[str]]:
"""Open path for writing with mode 0600 fixed up before any content is written.
A plain `open(path, "w")` creates a *new* file at the umask-derived default (commonly 0644)
and leaves it world- or group-readable until a later `chmod` call catches up -- a real window
in which a file holding a credential is readable by another local account. Passing the mode to
`os.open` closes that window for a brand-new file, but `O_CREAT`'s mode argument is only
applied on creation: if the file already exists its old, broader permissions carry over
untouched. `os.fchmod` right after opening -- before a single byte of the new content is
written -- covers both cases.
"""
fd = os.open(path, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600)
os.fchmod(fd, 0o600)
f: IO[str] = os.fdopen(fd, "w")
try:
yield f
finally:
f.close()
def write_backup(record: BackupRecord, backup_path: Path | None = None) -> None:
path = backup_path if backup_path is not None else BACKUP_PATH
path.parent.mkdir(exist_ok=True)
with secure_create(path) as f:
json.dump({"existed": record.existed, "content": record.content}, f, indent=2)
def read_backup(backup_path: Path | None = None) -> BackupRecord | None:
path = backup_path if backup_path is not None else BACKUP_PATH
if not path.exists():
return None
with open(path, "r") as f:
content = f.read()
try:
return _BACKUP_RECORD_ADAPTER.validate_json(content)
except ValidationError:
raise UpError(f"{path} contains invalid or unexpected JSON; cannot restore from it safely.")
def restore_claude_settings(settings_path: Path | None = None, backup_path: Path | None = None) -> BackupRecord | None:
"""Restore settings_path from the backup at backup_path, then delete the backup.
Returns the restored record, or None if there was nothing to restore.
"""
resolved_settings_path = settings_path if settings_path is not None else CLAUDE_SETTINGS_PATH
resolved_backup_path = backup_path if backup_path is not None else BACKUP_PATH
record = read_backup(resolved_backup_path)
if record is None:
return None
if record.existed and record.content is not None:
resolved_settings_path.parent.mkdir(parents=True, exist_ok=True)
with open(resolved_settings_path, "w") as f:
json.dump(record.content, f, indent=2)
elif resolved_settings_path.exists():
resolved_settings_path.unlink()
resolved_backup_path.unlink()
return record
def resolve_api_key_helper(base_url: str) -> str:
"""Build the shell command Claude Code should run for its apiKeyHelper.
Resolves `lite` to an absolute path so the helper works regardless of the
PATH visible to whatever subprocess Claude Code spawns it from. Passing
--base-url explicitly (rather than relying on the bare invocation Claude
Code would otherwise use) makes `print-token` enforce that the cached
token was actually issued for this proxy -- without it, a token minted
for a different, previously-logged-into proxy would be handed to
whichever server `up` currently points at.
"""
lite_path = shutil.which("lite")
if lite_path is None:
raise UpError(
"Could not find `lite` on your PATH. Claude Code's apiKeyHelper needs "
"an absolute path to it, so `lite up` cannot continue."
)
return f"{shlex.quote(lite_path)} auth print-token --base-url {shlex.quote(base_url)}"
def _ensure_fresh_login(ctx: click.Context) -> None:
base_url = ctx.obj["base_url"].rstrip("/")
token_data = load_token()
if token_data and token_data.get("base_url") == base_url and is_cli_token_fresh(token_data):
return
if not sys.stdin.isatty():
raise UpError(
"No fresh LiteLLM login found for this proxy. Run `lite login` first (apiKeyHelper "
"reads this token on every Claude Code request)."
)
click.echo("No fresh LiteLLM login found for this proxy; starting login...")
ctx.invoke(login)
token_data = load_token()
if not token_data or token_data.get("base_url") != base_url or not is_cli_token_fresh(token_data):
raise UpError("Login did not produce a usable token; cannot start `lite up`.")
def _restore_and_report() -> None:
record = restore_claude_settings()
if record is None:
click.echo("Nothing to restore.")
return
if record.existed:
click.echo(f"Restored {CLAUDE_SETTINGS_PATH} to its original contents.")
else:
click.echo(f"Removed {CLAUDE_SETTINGS_PATH} (it did not exist before `lite up`).")
@click.command(name="up")
@click.pass_context
def up(ctx: click.Context) -> None:
"""Route every Claude Code session through your LiteLLM proxy until stopped.
Patches ~/.claude/settings.json so Claude Code picks up the proxy on its own
next startup, from any terminal -- no need to launch it through `lite`.
Press Ctrl-C to stop and restore your original settings. Assumes the proxy
is already running (this does not start one for you). Cursor is not
supported: it has no equivalent file-based config to patch.
"""
base_url = ctx.obj["base_url"]
try:
_ensure_fresh_login(ctx)
api_key = resolve_api_key(ctx)
verify_proxy_key(base_url, api_key)
if BACKUP_PATH.exists():
raise UpError(
f"{BACKUP_PATH} already exists -- `lite up` looks like it's already "
"running (or crashed without cleanup). Run `lite down` first."
)
api_key_helper = resolve_api_key_helper(base_url)
original_existed = CLAUDE_SETTINGS_PATH.exists()
original_settings = load_json_or_empty(CLAUDE_SETTINGS_PATH)
write_backup(
BackupRecord(
existed=original_existed,
content=original_settings if original_existed else None,
)
)
CLAUDE_SETTINGS_PATH.parent.mkdir(exist_ok=True)
merged = merge_claude_settings(original_settings, base_url, api_key_helper)
with open(CLAUDE_SETTINGS_PATH, "w") as f:
json.dump(merged, f, indent=2)
except (AgentRunError, UpError) as e:
raise click.ClickException(str(e))
click.echo(f"litellm: routing Claude Code through proxy at {base_url.rstrip('/')}")
click.echo("Press Ctrl-C to stop and restore your original settings.")
stop_event = threading.Event()
restored = threading.Lock()
def _handle_signal(_signum: int, _frame: FrameType | None) -> None:
stop_event.set()
def _restore_once() -> None:
if not restored.acquire(blocking=False):
return
try:
_restore_and_report()
except UpError as e:
# Runs from atexit/a signal handler, outside Click's own exception
# handling -- raising here would only produce an unhandled-exception
# warning on stderr, not a clean message.
click.echo(str(e), err=True)
signal.signal(signal.SIGINT, _handle_signal)
signal.signal(signal.SIGTERM, _handle_signal)
atexit.register(_restore_once)
stop_event.wait()
_restore_once()
@click.command(name="down")
def down() -> None:
"""Restore ~/.claude/settings.json if a `lite up` session left it patched.
Use this after a `lite up` process was killed uncleanly (e.g. `kill -9`)
instead of stopped with Ctrl-C.
"""
try:
_restore_and_report()
except UpError as e:
raise click.ClickException(str(e))
__all__ = [
"BACKUP_PATH",
"CLAUDE_SETTINGS_PATH",
"BackupRecord",
"UpError",
"down",
"load_json_or_empty",
"merge_claude_settings",
"read_backup",
"resolve_api_key_helper",
"restore_claude_settings",
"up",
"write_backup",
]

14
litellm/proxy/client/cli/interface.py
View file

@ -27,13 +27,13 @@ def styled_prompt():
verbose_logger.debug(f"Error getting terminal size: {e}")
click.echo("\n" * 3)
# Unicode box drawing characters
top_left = "┌"
top_right = "┐"
bottom_left = "└"
bottom_right = "┘"
horizontal = "─"
vertical = "│"
# ASCII box drawing characters
top_left = "+"
top_right = "+"
bottom_left = "+"
bottom_right = "+"
horizontal = "-"
vertical = "|"
# Create the box with increased width
width = 80

10
litellm/proxy/client/cli/main.py
View file

@ -9,15 +9,18 @@ from litellm.proxy.client.health import HealthManagementClient
from .commands.agents import agent_commands
from .commands.auth import auth_group, get_stored_api_key, login, logout, whoami
from .commands.autoroute.commands import autoroute_group
from .commands.chat import chat
from .commands.credentials import credentials
from .commands.encryption import encryption
from .commands.http import http
from .commands.keys import keys
from .commands.model_groups import model_groups
# local imports
from .commands.models import models
from .commands.teams import teams
from .commands.up import down, up
from .commands.users import users
from .interface import interactive_shell
@ -131,6 +134,13 @@ cli.add_command(users)
# Add a top-level command per coding agent (claude, codex, opencode, ...)
for agent_command in agent_commands():
cli.add_command(agent_command)
# Add the up/down commands (route Claude Code through the local LiteLLM proxy)
cli.add_command(up)
cli.add_command(down)
# Add the model-groups command group (discover models your key can access)
cli.add_command(model_groups)
# Add the autoroute command group (QA auto-routing against your real proxy)
cli.add_command(autoroute_group, name="autoroute")
if __name__ == "__main__":

134
litellm/proxy/common_request_processing.py
View file

@ -151,6 +151,80 @@ async def _record_streaming_client_disconnect_if_needed(
return True
def _deferred_stream_logging_is_armed(request_data: dict) -> bool:
logging_obj = request_data.get("litellm_logging_obj")
if logging_obj is None:
return False
return (
getattr(logging_obj, "_on_deferred_stream_complete", None) is not None
and getattr(logging_obj, "_deferred_stream_complete_args", None) is not None
)
async def _bill_partial_streamed_spend_on_disconnect(request_data: dict, response: object) -> bool:
"""
A client disconnect throws GeneratorExit/CancelledError into the streaming
generator, so neither the success nor the failure logging callback fires
and the chunks already streamed (plus any sub-call cost folded into the
logging object) would never reach spend tracking. Assemble the partial
response from the wrapper's collected chunks and dispatch success logging
for it; dispatch_success_handlers dedups against a natural end-of-stream
dispatch via has_dispatched_final_stream_success.
Awaited directly by the shielded cleanup rather than scheduled with
create_task: the client is already gone so the extra latency is harmless,
and an unrooted task could be garbage-collected before it bills.
Returns True when a disconnect-time success event owns the request's
max_parallel_requests slot release (one was dispatched here, or one had
already been dispatched for this stream), so the caller can skip the
explicit slot release and avoid a double release. Returns False when no
success event fired (logging disabled, nothing streamed, or assembly
failed) and the caller must release the slot itself.
"""
if litellm.disable_streaming_logging is True:
return False
logging_obj = request_data.get("litellm_logging_obj")
if not isinstance(logging_obj, LiteLLMLoggingObj):
return False
if logging_obj.model_call_details.get("has_dispatched_final_stream_success"):
# A natural end-of-stream success event already fired and released the
# slot; do not bill again, and let the caller skip the slot release.
return True
chunks: object = getattr(response, "chunks", None)
if not isinstance(chunks, list) or not chunks:
return False
verbose_proxy_logger.debug(
"Billing partial streamed spend for %s chunks after client disconnect, litellm_call_id=%s",
len(chunks),
request_data.get("litellm_call_id"),
)
messages: object = getattr(response, "messages", None)
try:
partial_response = litellm.stream_chunk_builder(
chunks=chunks,
messages=messages if isinstance(messages, list) else None,
logging_obj=logging_obj,
)
except Exception as e: # noqa: BLE001 # partial billing is best-effort; never break stream teardown
verbose_proxy_logger.debug("Failed to assemble partial streamed response for disconnect billing: %s", e)
return False
if partial_response is None:
return False
try:
await logging_obj.dispatch_success_handlers(
partial_response,
cache_hit=False,
start_time=None,
end_time=None,
prefer_async_handlers=True,
)
except Exception as e: # noqa: BLE001 # partial billing is best-effort; never break stream teardown
verbose_proxy_logger.debug("Failed to dispatch disconnect billing event: %s", e)
return False
return True
async def _cancel_pending_gather_tasks(tasks: list["asyncio.Task[Any]"]) -> None:
pending_tasks = [task for task in tasks if not task.done()]
for task in pending_tasks:
@ -851,9 +925,12 @@ class ProxyBaseLLMRequestProcessing:
# If conversion fails, use original spend
pass
model_name = ProxyBaseLLMRequestProcessing._get_deployment_model_name(litellm_logging_obj)
headers = {
"x-litellm-call-id": call_id,
"x-litellm-model-id": model_id,
"x-litellm-model-name": model_name,
"x-litellm-cache-key": cache_key,
"x-litellm-model-api-base": (
api_base.split("?")[0] if api_base else None
@ -1322,6 +1399,27 @@ class ProxyBaseLLMRequestProcessing:
model_id = model_info.get("id", "") or ""
return model_id
@staticmethod
def _get_deployment_model_name(
litellm_logging_obj: LiteLLMLoggingObj | None,
) -> str | None:
"""Extract the underlying deployment model string (e.g. ``azure/gpt-4o``).
The router rewrites the response ``model`` field to the model-group alias
the client requested, so neither the response body nor the existing
headers expose the concrete deployment model. The router records it under
``litellm_params`` metadata as ``deployment``, so read it back from there.
"""
litellm_params = getattr(litellm_logging_obj, "litellm_params", None)
if not isinstance(litellm_params, dict):
return None
for key in ("litellm_metadata", "metadata"):
metadata = litellm_params.get(key, {}) or {}
deployment = metadata.get("deployment")
if deployment:
return deployment
return None
@staticmethod
def _response_cost_from_logging_obj(
*,
@ -2575,6 +2673,8 @@ class ProxyBaseLLMRequestProcessing:
response: Any,
stream_completed: bool = False,
client_disconnected: bool = False,
user_api_key_dict: UserAPIKeyAuth | None = None,
proxy_logging_obj: ProxyLogging | None = None,
) -> None:
with anyio.CancelScope(shield=True):
should_record_client_disconnect = client_disconnected or (not stream_completed)
@ -2586,7 +2686,28 @@ class ProxyBaseLLMRequestProcessing:
client_disconnected,
)
if recorded_client_disconnect:
deferred_stream_logging_armed = _deferred_stream_logging_is_armed(request_data)
ProxyLogging._fire_deferred_stream_logging(request_data)
# A disconnect-time success event (the deferred-guardrail flush
# above, or the partial-spend billing below) releases the
# request's max_parallel_requests slot through the limiter's
# own success callback. Release the slot explicitly only when
# no such event fires, so exactly one release happens; two
# concurrent releases would race and double-decrement under the
# limiter's in-memory fallback.
success_event_owns_slot_release = deferred_stream_logging_armed
if not deferred_stream_logging_armed:
success_event_owns_slot_release = await _bill_partial_streamed_spend_on_disconnect(
request_data, response
)
if (
not success_event_owns_slot_release
and proxy_logging_obj is not None
and user_api_key_dict is not None
):
await proxy_logging_obj._arelease_max_parallel_requests_on_disconnect(
user_api_key_dict, request_data
)
if hasattr(response, "aclose"):
try:
@ -2675,12 +2796,13 @@ class ProxyBaseLLMRequestProcessing:
except (asyncio.CancelledError, GeneratorExit):
# Client disconnected mid-stream. CancelledError / GeneratorExit
# are BaseException and bypass the success/failure logging
# callbacks that release the pre-call max_parallel_requests +1;
# release it here. This is the outermost generator Starlette closes
# on disconnect, so the nested iterator hook (which only sees
# GeneratorExit on GC) cannot own the refund.
# callbacks that release the pre-call max_parallel_requests +1.
# Flag the disconnect; the shielded cleanup in `finally` owns the
# slot release so it can coordinate with disconnect-time success
# billing and release exactly once. This is the outermost generator
# Starlette closes on disconnect, so the nested iterator hook (which
# only sees GeneratorExit on GC) cannot own the refund.
if not stream_completed:
proxy_logging_obj._release_max_parallel_requests_on_disconnect(user_api_key_dict)
client_disconnected = True
if not delivered_chunk:
from litellm.proxy.spend_tracking.budget_reservation import (
@ -2723,6 +2845,8 @@ class ProxyBaseLLMRequestProcessing:
response=response,
stream_completed=stream_completed,
client_disconnected=client_disconnected,
user_api_key_dict=user_api_key_dict,
proxy_logging_obj=proxy_logging_obj,
)
@staticmethod

12
litellm/proxy/db/db_spend_update_writer.py
View file

@ -27,7 +27,7 @@ from typing import (
import litellm
from litellm._logging import verbose_proxy_logger
from litellm.caching import DualCache, RedisCache
from litellm.caching import RedisCache
from litellm.constants import (
DB_SPEND_UPDATE_JOB_NAME,
DB_DAILY_TAG_SPEND_UPDATE_JOB_NAME,
@ -44,7 +44,6 @@ from litellm.proxy._types import (
DailyUserSpendTransaction,
DBSpendUpdateTransactions,
Litellm_EntityType,
LiteLLM_UserTable,
SpendLogsMetadata,
SpendLogsPayload,
SpendUpdateQueueItem,
@ -137,7 +136,6 @@ class DBSpendUpdateWriter:
disable_spend_logs,
litellm_proxy_budget_name,
prisma_client,
user_api_key_cache,
)
from litellm.proxy.utils import ProxyUpdateSpend, hash_token
@ -195,7 +193,6 @@ class DBSpendUpdateWriter:
org_id=org_id,
end_user_id=end_user_id,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
litellm_proxy_budget_name=litellm_proxy_budget_name,
payload=payload,
)
@ -326,7 +323,6 @@ class DBSpendUpdateWriter:
org_id: Optional[str],
end_user_id: Optional[str],
prisma_client: Optional[PrismaClient],
user_api_key_cache: DualCache,
litellm_proxy_budget_name: Optional[str],
payload: SpendLogsPayload,
):
@ -345,7 +341,6 @@ class DBSpendUpdateWriter:
response_cost=response_cost,
user_id=user_id,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
litellm_proxy_budget_name=litellm_proxy_budget_name,
end_user_id=end_user_id,
)
@ -510,7 +505,6 @@ class DBSpendUpdateWriter:
response_cost: Optional[float],
user_id: Optional[str],
prisma_client: Optional[PrismaClient],
user_api_key_cache: DualCache,
litellm_proxy_budget_name: Optional[str],
end_user_id: Optional[str] = None,
):
@ -518,10 +512,6 @@ class DBSpendUpdateWriter:
- Update that user's row
- Update litellm-proxy-budget row (global proxy spend)
"""
## if an end-user is passed in, do an upsert - we can't guarantee they already exist in db
existing_user_obj = await user_api_key_cache.async_get_cache(key=user_id)
if existing_user_obj is not None and isinstance(existing_user_obj, dict):
existing_user_obj = LiteLLM_UserTable(**existing_user_obj)
try:
if prisma_client is not None: # update
user_ids = [user_id]

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