From 6e3670ddcac22d6c52ec9af3cb9db9ae332bf167 Mon Sep 17 00:00:00 2001 From: mateo Date: Wed, 22 Jul 2026 18:27:33 +0000 Subject: [PATCH 001/144] feat(proxy): serve the Claude Code gateway protocol under /claude_code_gateway Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/_lazy_features.py | 5 + litellm/proxy/_types.py | 8 + .../anthropic_endpoints/gateway_endpoints.py | 289 ++++++++++++++++++ .../test_gateway_endpoints.py | 219 +++++++++++++ 4 files changed, 521 insertions(+) create mode 100644 litellm/proxy/anthropic_endpoints/gateway_endpoints.py create mode 100644 tests/test_litellm/proxy/anthropic_endpoints/test_gateway_endpoints.py diff --git a/litellm/proxy/_lazy_features.py b/litellm/proxy/_lazy_features.py index 1dda1f29fb9..1426d100713 100644 --- a/litellm/proxy/_lazy_features.py +++ b/litellm/proxy/_lazy_features.py @@ -196,6 +196,11 @@ LAZY_FEATURES: Tuple[LazyFeature, ...] = ( module_path="litellm.proxy.anthropic_endpoints.skills_endpoints", path_prefixes=("/v1/skills", "/skills"), ), + LazyFeature( + name="claude_code_gateway", + module_path="litellm.proxy.anthropic_endpoints.gateway_endpoints", + path_prefixes=("/claude_code_gateway",), + ), LazyFeature( name="langfuse_passthrough", module_path="litellm.proxy.vertex_ai_endpoints.langfuse_endpoints", diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 7df725bf965..5577c5caff4 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -2223,6 +2223,14 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase): None, description="opt-in to RFC 8628 verification_uri_complete for the CLI SSO device flow, pre-filling the user_code in the browser. Off by default; intended for same-host clients where the device that starts the flow and the browser run on the same machine", ) + enable_claude_code_gateway: bool | None = Field( + None, + description="serve the Claude Code gateway protocol (https://code.claude.com/docs/en/claude-apps-gateway) under /claude_code_gateway: OAuth device-flow sign-in reusing proxy SSO, plus managed settings and OTLP telemetry ingestion. Off by default", + ) + claude_code_gateway_managed_settings: Dict[str, Any] | None = Field( + None, + description="Claude Code managed-settings.json served verbatim at the gateway's /claude_code_gateway/managed/settings endpoint. When unset the endpoint returns 404 (no managed policy)", + ) database_url: Optional[str] = Field( None, description="connect to a postgres db - needed for generating temporary keys + tracking spend / key", diff --git a/litellm/proxy/anthropic_endpoints/gateway_endpoints.py b/litellm/proxy/anthropic_endpoints/gateway_endpoints.py new file mode 100644 index 00000000000..1ec4cd488c5 --- /dev/null +++ b/litellm/proxy/anthropic_endpoints/gateway_endpoints.py @@ -0,0 +1,289 @@ +""" +Claude Code gateway protocol. + +Implements the wire contract the Claude Code CLI uses to talk to a gateway: +OAuth 2.0 device-authorization sign-in (RFC 8414 / RFC 8628), inference via the +Anthropic Messages API, managed settings, and OTLP telemetry ingestion. See +https://code.claude.com/docs/en/claude-apps-gateway. + +Everything lives under the ``/claude_code_gateway`` base so operators point +Claude Code at ``https:///claude_code_gateway`` via ``/login``. The +device flow reuses the proxy's existing SSO login machinery: the browser leg is +served by ``/sso/key/generate`` and the shared ``cli_sso_session_cache`` flow, +so the bearer token minted here is the same session JWT the LiteLLM CLI uses and +is accepted by every bearer-authenticated proxy route. +""" + +import hashlib +import json +import secrets +from typing import Any + +from fastapi import APIRouter, Depends, Request, Response +from fastapi.responses import JSONResponse + +from litellm.constants import ( + CLI_JWT_EXPIRATION_HOURS, + CLI_SSO_SESSION_TTL_SECONDS, + LITELLM_CLI_SOURCE_IDENTIFIER, +) +from litellm.proxy.anthropic_endpoints.endpoints import anthropic_response, count_tokens +from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + +GATEWAY_PREFIX = "/claude_code_gateway" +_DEVICE_CODE_GRANT = "urn:ietf:params:oauth:grant-type:device_code" +_REFRESH_TOKEN_GRANT = "refresh_token" +_DEVICE_POLL_INTERVAL_SECONDS = 5 + + +def _is_gateway_enabled() -> bool: + from litellm.proxy.proxy_server import general_settings + + return bool((general_settings or {}).get("enable_claude_code_gateway", False)) + + +def ensure_gateway_enabled() -> None: + from fastapi import HTTPException + + if not _is_gateway_enabled(): + raise HTTPException(status_code=404, detail="Claude Code gateway is not enabled") + + +def _managed_settings() -> dict[str, Any] | None: + from litellm.proxy.proxy_server import general_settings + + settings = (general_settings or {}).get("claude_code_gateway_managed_settings") + return settings if isinstance(settings, dict) else None + + +def _oauth_error(*, status_code: int, error: str, description: str | None = None) -> "_OAuthError": + return _OAuthError(status_code=status_code, error=error, description=description) + + +class _OAuthError(Exception): + def __init__(self, *, status_code: int, error: str, description: str | None) -> None: + self.status_code = status_code + self.error = error + self.description = description + + +def _oauth_error_response(err: _OAuthError) -> JSONResponse: + body: dict[str, str] = {"error": err.error} + if err.description is not None: + body["error_description"] = err.description + return JSONResponse(status_code=err.status_code, content=body) + + +router = APIRouter(prefix=GATEWAY_PREFIX, tags=["Claude Code gateway"]) + +router.add_api_route( + "/v1/messages", + anthropic_response, + methods=["POST"], + dependencies=[Depends(ensure_gateway_enabled)], + include_in_schema=False, +) +router.add_api_route( + "/v1/messages/count_tokens", + count_tokens, + methods=["POST"], + dependencies=[Depends(ensure_gateway_enabled)], + include_in_schema=False, +) + + +@router.get("/.well-known/oauth-authorization-server", include_in_schema=False) +async def oauth_authorization_server(request: Request) -> JSONResponse: + if not _is_gateway_enabled(): + return _oauth_error_response(_oauth_error(status_code=404, error="not_found")) + + from litellm.proxy.utils import get_custom_url + + request_base_url = str(request.base_url) + issuer = get_custom_url(request_base_url=request_base_url, route="claude_code_gateway") + return JSONResponse( + content={ + "issuer": issuer, + "device_authorization_endpoint": get_custom_url( + request_base_url=request_base_url, route="claude_code_gateway/oauth/device_authorization" + ), + "token_endpoint": get_custom_url( + request_base_url=request_base_url, route="claude_code_gateway/oauth/token" + ), + "grant_types_supported": [_DEVICE_CODE_GRANT, _REFRESH_TOKEN_GRANT], + } + ) + + +@router.post("/oauth/device_authorization", include_in_schema=False) +async def device_authorization(request: Request) -> JSONResponse: + from urllib.parse import urlencode + + from litellm.proxy.management_endpoints.ui_sso import ( + _check_cli_sso_start_rate_limit, + _generate_cli_sso_user_code, + _hash_cli_sso_secret, + _normalize_cli_sso_user_code, + _set_cli_sso_flow, + ) + from litellm.proxy.proxy_server import cli_sso_session_cache, general_settings + from litellm.proxy.utils import get_custom_url + + if not _is_gateway_enabled(): + return _oauth_error_response(_oauth_error(status_code=404, error="not_found")) + + _check_cli_sso_start_rate_limit( + request=request, + cache=cli_sso_session_cache, + use_x_forwarded_for=bool((general_settings or {}).get("use_x_forwarded_for", False)), + ) + + device_code = f"cli-{secrets.token_urlsafe(24)}" + user_code = _generate_cli_sso_user_code() + flow = { + "poll_secret_hash": _hash_cli_sso_secret(device_code), + "user_code_hash": _hash_cli_sso_secret(_normalize_cli_sso_user_code(user_code)), + "sso_complete": False, + "user_code_verified": False, + "session_data": None, + } + _set_cli_sso_flow(login_id=device_code, cache=cli_sso_session_cache, flow=flow) + + request_base_url = str(request.base_url) + verification_uri = get_custom_url(request_base_url=request_base_url, route="sso/key/generate") + verification_uri_complete = ( + verification_uri + + "?" + + urlencode({"source": LITELLM_CLI_SOURCE_IDENTIFIER, "key": device_code, "user_code": user_code}) + ) + verification_uri_no_code = ( + verification_uri + "?" + urlencode({"source": LITELLM_CLI_SOURCE_IDENTIFIER, "key": device_code}) + ) + return JSONResponse( + content={ + "device_code": device_code, + "user_code": user_code, + "verification_uri": verification_uri_no_code, + "verification_uri_complete": verification_uri_complete, + "expires_in": CLI_SSO_SESSION_TTL_SECONDS, + "interval": _DEVICE_POLL_INTERVAL_SECONDS, + } + ) + + +def _mint_access_token_from_flow(flow: dict[str, Any]) -> str: + from litellm.proxy._types import LiteLLM_UserTable + from litellm.proxy.auth.auth_checks import ExperimentalUIJWTToken + + session_data = flow.get("session_data") + if not isinstance(session_data, dict): + raise _oauth_error(status_code=400, error="authorization_pending") + + teams = session_data.get("teams") or [] + team_id = teams[0] if isinstance(teams, list) and teams else None + user_info = LiteLLM_UserTable( + user_id=session_data["user_id"], + user_role=session_data["user_role"], + models=session_data.get("models", []), + ) + return ExperimentalUIJWTToken.get_cli_jwt_auth_token(user_info=user_info, team_id=team_id) + + +async def _handle_device_code_grant(device_code: str | None) -> JSONResponse: + from fastapi import HTTPException + + from litellm.proxy.management_endpoints.ui_sso import ( + _get_cli_sso_flow_cache_key, + _get_cli_sso_flow_or_raise, + ) + from litellm.proxy.proxy_server import cli_sso_session_cache + + if not device_code: + return _oauth_error_response( + _oauth_error(status_code=400, error="invalid_request", description="device_code is required") + ) + + try: + flow = _get_cli_sso_flow_or_raise(login_id=device_code, cache=cli_sso_session_cache) + except HTTPException: + return _oauth_error_response(_oauth_error(status_code=400, error="expired_token")) + + if not flow.get("sso_complete") or not flow.get("user_code_verified"): + return _oauth_error_response(_oauth_error(status_code=400, error="authorization_pending")) + + try: + access_token = _mint_access_token_from_flow(flow) + except _OAuthError as err: + return _oauth_error_response(err) + + cli_sso_session_cache.delete_cache(key=_get_cli_sso_flow_cache_key(device_code)) + return JSONResponse( + content={ + "access_token": access_token, + "token_type": "Bearer", + "expires_in": CLI_JWT_EXPIRATION_HOURS * 3600, + } + ) + + +@router.post("/oauth/token", include_in_schema=False) +async def oauth_token(request: Request) -> JSONResponse: + if not _is_gateway_enabled(): + return _oauth_error_response(_oauth_error(status_code=404, error="not_found")) + + form = await request.form() + grant_type = form.get("grant_type") + + if grant_type == _DEVICE_CODE_GRANT: + device_code = form.get("device_code") + return await _handle_device_code_grant(device_code if isinstance(device_code, str) else None) + + if grant_type == _REFRESH_TOKEN_GRANT: + return _oauth_error_response( + _oauth_error( + status_code=401, + error="invalid_grant", + description="This gateway does not issue refresh tokens; sign in again", + ) + ) + + return _oauth_error_response( + _oauth_error(status_code=400, error="unsupported_grant_type", description=f"Unsupported grant_type: {grant_type}") + ) + + +@router.get("/managed/settings", include_in_schema=False, dependencies=[Depends(user_api_key_auth)]) +async def managed_settings(request: Request) -> Response: + ensure_gateway_enabled() + + settings = _managed_settings() + if settings is None: + return Response(status_code=404) + + body = json.dumps(settings, sort_keys=True, separators=(",", ":")) + etag = '"' + hashlib.sha256(body.encode("utf-8")).hexdigest() + '"' + if_none_match = request.headers.get("If-None-Match") + if if_none_match is not None and if_none_match == etag: + return Response(status_code=304, headers={"ETag": etag}) + return Response(content=body, media_type="application/json", headers={"ETag": etag}) + + +async def _accept_otlp(request: Request) -> Response: + ensure_gateway_enabled() + await request.body() + return Response(status_code=200) + + +@router.post("/v1/metrics", include_in_schema=False, dependencies=[Depends(user_api_key_auth)]) +async def otlp_metrics(request: Request) -> Response: + return await _accept_otlp(request) + + +@router.post("/v1/logs", include_in_schema=False, dependencies=[Depends(user_api_key_auth)]) +async def otlp_logs(request: Request) -> Response: + return await _accept_otlp(request) + + +@router.post("/v1/traces", include_in_schema=False, dependencies=[Depends(user_api_key_auth)]) +async def otlp_traces(request: Request) -> Response: + return await _accept_otlp(request) diff --git a/tests/test_litellm/proxy/anthropic_endpoints/test_gateway_endpoints.py b/tests/test_litellm/proxy/anthropic_endpoints/test_gateway_endpoints.py new file mode 100644 index 00000000000..8645f4a8680 --- /dev/null +++ b/tests/test_litellm/proxy/anthropic_endpoints/test_gateway_endpoints.py @@ -0,0 +1,219 @@ +""" +Tests for the Claude Code gateway protocol (anthropic_endpoints/gateway_endpoints.py). + +Covers the OAuth device-flow surface (RFC 8414 discovery, RFC 8628 device +authorization + token), managed settings, OTLP ingestion, and the enable flag. +""" + +from contextlib import contextmanager +from typing import Any, Iterator, Optional +from unittest.mock import patch + +import pytest +from fastapi import FastAPI +from fastapi.testclient import TestClient + +from litellm.caching.dual_cache import DualCache +from litellm.proxy.anthropic_endpoints import gateway_endpoints +from litellm.proxy.management_endpoints.ui_sso import _get_cli_sso_flow_cache_key + + +@contextmanager +def _gateway_env( + *, + enabled: bool = True, + managed_settings: Optional[dict[str, Any]] = None, +) -> Iterator[tuple[TestClient, DualCache]]: + general_settings: dict[str, Any] = {"enable_claude_code_gateway": enabled} + if managed_settings is not None: + general_settings["claude_code_gateway_managed_settings"] = managed_settings + cache = DualCache(default_in_memory_ttl=600) + + app = FastAPI() + app.include_router(gateway_endpoints.router) + + async def _fake_auth() -> Any: + return object() + + app.dependency_overrides[gateway_endpoints.user_api_key_auth] = _fake_auth + + with patch("litellm.proxy.proxy_server.general_settings", general_settings), patch( + "litellm.proxy.proxy_server.cli_sso_session_cache", cache + ): + with TestClient(app) as client: + yield client, cache + + +def _complete_flow(cache: DualCache, device_code: str) -> None: + key = _get_cli_sso_flow_cache_key(device_code) + flow = cache.get_cache(key=key) + assert isinstance(flow, dict) + flow["sso_complete"] = True + flow["user_code_verified"] = True + flow["session_data"] = { + "user_id": "user-123", + "user_role": "internal_user", + "models": ["claude-sonnet-4-5"], + "teams": ["team-a"], + } + cache.set_cache(key=key, value=flow, ttl=600) + + +def test_discovery_shape(): + with _gateway_env() as (client, _): + resp = client.get("/claude_code_gateway/.well-known/oauth-authorization-server") + assert resp.status_code == 200 + body = resp.json() + assert body["device_authorization_endpoint"].endswith("/claude_code_gateway/oauth/device_authorization") + assert body["token_endpoint"].endswith("/claude_code_gateway/oauth/token") + assert body["grant_types_supported"] == [ + "urn:ietf:params:oauth:grant-type:device_code", + "refresh_token", + ] + # authorization_endpoint is intentionally absent (device flow only). + assert "authorization_endpoint" not in body + # Both endpoints must be same-origin with the issuer. + assert body["device_authorization_endpoint"].startswith(body["issuer"]) + assert body["token_endpoint"].startswith(body["issuer"]) + + +def test_discovery_404_when_disabled(): + with _gateway_env(enabled=False) as (client, _): + resp = client.get("/claude_code_gateway/.well-known/oauth-authorization-server") + assert resp.status_code == 404 + + +def test_device_authorization_returns_rfc8628_shape_and_persists_flow(): + with _gateway_env() as (client, cache): + resp = client.post("/claude_code_gateway/oauth/device_authorization") + assert resp.status_code == 200 + body = resp.json() + device_code = body["device_code"] + assert device_code.startswith("cli-") + assert body["user_code"] + assert body["expires_in"] == 600 + assert body["interval"] == 5 + # verification_uri_complete carries the user_code; the short uri does not. + assert f"user_code={body['user_code']}" in body["verification_uri_complete"] + assert "user_code=" not in body["verification_uri"] + assert f"key={device_code}" in body["verification_uri"] + # The device flow is stored under the device_code so the browser SSO leg can complete it. + stored = cache.get_cache(key=_get_cli_sso_flow_cache_key(device_code)) + assert isinstance(stored, dict) + assert stored["sso_complete"] is False + + +def test_token_authorization_pending_before_browser_completes(): + with _gateway_env() as (client, _): + device_code = client.post("/claude_code_gateway/oauth/device_authorization").json()["device_code"] + resp = client.post( + "/claude_code_gateway/oauth/token", + data={"grant_type": "urn:ietf:params:oauth:grant-type:device_code", "device_code": device_code}, + ) + assert resp.status_code == 400 + assert resp.json()["error"] == "authorization_pending" + + +def test_token_success_mints_bearer_and_is_single_use(): + with _gateway_env() as (client, cache): + device_code = client.post("/claude_code_gateway/oauth/device_authorization").json()["device_code"] + _complete_flow(cache, device_code) + + with patch( + "litellm.proxy.auth.auth_checks.ExperimentalUIJWTToken.get_cli_jwt_auth_token", + return_value="sk-litellm-session-token", + ) as mint: + resp = client.post( + "/claude_code_gateway/oauth/token", + data={"grant_type": "urn:ietf:params:oauth:grant-type:device_code", "device_code": device_code}, + ) + assert resp.status_code == 200 + body = resp.json() + assert body["access_token"] == "sk-litellm-session-token" + assert body["token_type"] == "Bearer" + assert body["expires_in"] > 0 + + called_user = mint.call_args.kwargs["user_info"] + assert called_user.user_id == "user-123" + assert mint.call_args.kwargs["team_id"] == "team-a" + + # Single-use: the flow is deleted, so a replay returns expired_token. + replay = client.post( + "/claude_code_gateway/oauth/token", + data={"grant_type": "urn:ietf:params:oauth:grant-type:device_code", "device_code": device_code}, + ) + assert replay.status_code == 400 + assert replay.json()["error"] == "expired_token" + + +def test_token_unknown_device_code_is_expired_token(): + with _gateway_env() as (client, _): + resp = client.post( + "/claude_code_gateway/oauth/token", + data={"grant_type": "urn:ietf:params:oauth:grant-type:device_code", "device_code": "cli-does-not-exist"}, + ) + assert resp.status_code == 400 + assert resp.json()["error"] == "expired_token" + + +def test_refresh_grant_forces_relogin(): + with _gateway_env() as (client, _): + resp = client.post( + "/claude_code_gateway/oauth/token", + data={"grant_type": "refresh_token", "refresh_token": "whatever"}, + ) + assert resp.status_code == 401 + assert resp.json()["error"] == "invalid_grant" + + +def test_unsupported_grant_type(): + with _gateway_env() as (client, _): + resp = client.post("/claude_code_gateway/oauth/token", data={"grant_type": "password"}) + assert resp.status_code == 400 + assert resp.json()["error"] == "unsupported_grant_type" + + +def test_managed_settings_404_when_unset(): + with _gateway_env() as (client, _): + resp = client.get("/claude_code_gateway/managed/settings") + assert resp.status_code == 404 + + +def test_managed_settings_returns_json_with_etag_and_304(): + settings = {"permissions": {"defaultMode": "acceptEdits"}, "env": {"FOO": "bar"}} + with _gateway_env(managed_settings=settings) as (client, _): + resp = client.get("/claude_code_gateway/managed/settings") + assert resp.status_code == 200 + assert resp.json() == settings + etag = resp.headers["ETag"] + assert etag + + not_modified = client.get("/claude_code_gateway/managed/settings", headers={"If-None-Match": etag}) + assert not_modified.status_code == 304 + assert not_modified.headers["ETag"] == etag + + +def test_managed_settings_404_when_gateway_disabled(): + with _gateway_env(enabled=False, managed_settings={"env": {}}) as (client, _): + resp = client.get("/claude_code_gateway/managed/settings") + assert resp.status_code == 404 + + +@pytest.mark.parametrize("signal", ["metrics", "logs", "traces"]) +def test_otlp_endpoints_accept_and_return_200(signal: str): + with _gateway_env() as (client, _): + resp = client.post(f"/claude_code_gateway/v1/{signal}", content=b"\x00\x01binary-otlp") + assert resp.status_code == 200 + + +@pytest.mark.parametrize("signal", ["metrics", "logs", "traces"]) +def test_otlp_endpoints_404_when_disabled(signal: str): + with _gateway_env(enabled=False) as (client, _): + resp = client.post(f"/claude_code_gateway/v1/{signal}", content=b"payload") + assert resp.status_code == 404 + + +def test_messages_gated_by_enable_flag(): + with _gateway_env(enabled=False) as (client, _): + resp = client.post("/claude_code_gateway/v1/messages", json={"model": "claude-sonnet-4-5", "messages": []}) + assert resp.status_code == 404 From 575da405f231b43d55eaabd55057140f60f84629 Mon Sep 17 00:00:00 2001 From: mateo Date: Wed, 22 Jul 2026 18:31:12 +0000 Subject: [PATCH 002/144] fix(proxy): do not re-read request body in Claude Code gateway OTLP handlers Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../anthropic_endpoints/gateway_endpoints.py | 15 +++++++-------- 1 file changed, 7 insertions(+), 8 deletions(-) diff --git a/litellm/proxy/anthropic_endpoints/gateway_endpoints.py b/litellm/proxy/anthropic_endpoints/gateway_endpoints.py index 1ec4cd488c5..d0861cccf9a 100644 --- a/litellm/proxy/anthropic_endpoints/gateway_endpoints.py +++ b/litellm/proxy/anthropic_endpoints/gateway_endpoints.py @@ -268,22 +268,21 @@ async def managed_settings(request: Request) -> Response: return Response(content=body, media_type="application/json", headers={"ETag": etag}) -async def _accept_otlp(request: Request) -> Response: +def _accept_otlp() -> Response: ensure_gateway_enabled() - await request.body() return Response(status_code=200) @router.post("/v1/metrics", include_in_schema=False, dependencies=[Depends(user_api_key_auth)]) -async def otlp_metrics(request: Request) -> Response: - return await _accept_otlp(request) +async def otlp_metrics() -> Response: + return _accept_otlp() @router.post("/v1/logs", include_in_schema=False, dependencies=[Depends(user_api_key_auth)]) -async def otlp_logs(request: Request) -> Response: - return await _accept_otlp(request) +async def otlp_logs() -> Response: + return _accept_otlp() @router.post("/v1/traces", include_in_schema=False, dependencies=[Depends(user_api_key_auth)]) -async def otlp_traces(request: Request) -> Response: - return await _accept_otlp(request) +async def otlp_traces() -> Response: + return _accept_otlp() From 1c445f36166e56f14b3d716a7700e554ec2f0c0d Mon Sep 17 00:00:00 2001 From: mateo Date: Wed, 22 Jul 2026 18:37:11 +0000 Subject: [PATCH 003/144] style(proxy): apply ruff format to gateway endpoints Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/anthropic_endpoints/gateway_endpoints.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/anthropic_endpoints/gateway_endpoints.py b/litellm/proxy/anthropic_endpoints/gateway_endpoints.py index d0861cccf9a..e0c35117bf6 100644 --- a/litellm/proxy/anthropic_endpoints/gateway_endpoints.py +++ b/litellm/proxy/anthropic_endpoints/gateway_endpoints.py @@ -248,7 +248,9 @@ async def oauth_token(request: Request) -> JSONResponse: ) return _oauth_error_response( - _oauth_error(status_code=400, error="unsupported_grant_type", description=f"Unsupported grant_type: {grant_type}") + _oauth_error( + status_code=400, error="unsupported_grant_type", description=f"Unsupported grant_type: {grant_type}" + ) ) From b9f63c7bd9f69e5f18f91ab1302ddb7a45e0d809 Mon Sep 17 00:00:00 2001 From: mateo Date: Wed, 22 Jul 2026 18:43:19 +0000 Subject: [PATCH 004/144] chore(ui): regenerate schema.d.ts for Claude Code gateway config fields Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ui/litellm-dashboard/src/lib/http/schema.d.ts | 46 +++++++++++++++++++ 1 file changed, 46 insertions(+) diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 203057f23f1..e061f4277cc 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -1569,6 +1569,40 @@ export interface paths { patch?: never; trace?: never; }; + "/claude_code_gateway": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** claude_code_gateway */ + get: { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description OK */ + 200: { + headers: { + [name: string]: unknown; + }; + content?: never; + }; + }; + }; + put?: never; + post?: never; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; "/cloudzero/delete": { parameters: { query?: never; @@ -22451,6 +22485,13 @@ export interface components { * @description cancel the in-flight upstream LLM request (non-streaming) when the client disconnects, freeing backend capacity (e.g. a vLLM GPU slot); the request is logged as a 499 failure */ cancel_on_disconnect?: boolean | null; + /** + * Claude Code Gateway Managed Settings + * @description Claude Code managed-settings.json served verbatim at the gateway's /claude_code_gateway/managed/settings endpoint. When unset the endpoint returns 404 (no managed policy) + */ + claude_code_gateway_managed_settings?: { + [key: string]: unknown; + } | null; /** * Completion Model * @description proxy level default model for all chat completion calls @@ -22519,6 +22560,11 @@ export interface components { * @description If True, disables the optimistic per-request budget reservation introduced in v1.84.0. WARNING: This weakens hard budget enforcement. Without the reservation, a burst of concurrent requests from a single key can each pass the read-time spend check before any of them is charged, allowing a configured budget to be exceeded under high concurrency. Budgets are still evaluated on every request at read time, so an already-exhausted budget is still rejected. Enable only if your deployment is experiencing phantom BudgetExceededError responses caused by leaked reservations (see GitHub issue #27639). A proxy-level WARNING is logged on every request while this flag is active as a reminder that hard enforcement is relaxed. */ disable_budget_reservation?: boolean | null; + /** + * Enable Claude Code Gateway + * @description serve the Claude Code gateway protocol (https://code.claude.com/docs/en/claude-apps-gateway) under /claude_code_gateway: OAuth device-flow sign-in reusing proxy SSO, plus managed settings and OTLP telemetry ingestion. Off by default + */ + enable_claude_code_gateway?: boolean | null; /** * Enable Public Model Hub * @description Public model hub for users to see what models they have access to, supported openai params, etc. From 6230a379b91470a2690d9ec54ccd40a909947fd3 Mon Sep 17 00:00:00 2001 From: mateo Date: Sat, 15 Aug 2026 17:58:43 +0000 Subject: [PATCH 005/144] fix(proxy): use builtin dict annotation for gateway managed settings config Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/_types.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 3e26d6230fd..a2d0a2cdcd0 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -2340,7 +2340,7 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase): None, description="serve the Claude Code gateway protocol (https://code.claude.com/docs/en/claude-apps-gateway) under /claude_code_gateway: OAuth device-flow sign-in reusing proxy SSO, plus managed settings and OTLP telemetry ingestion. Off by default", ) - claude_code_gateway_managed_settings: Dict[str, Any] | None = Field( + claude_code_gateway_managed_settings: dict[str, Any] | None = Field( None, description="Claude Code managed-settings.json served verbatim at the gateway's /claude_code_gateway/managed/settings endpoint. When unset the endpoint returns 404 (no managed policy)", ) From b5bcdbccd4a268f33b9ee10e4ef41d325dc27af8 Mon Sep 17 00:00:00 2001 From: mateo Date: Sat, 15 Aug 2026 18:33:09 +0000 Subject: [PATCH 006/144] refactor(proxy): type gateway protocol payloads with pydantic models Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../anthropic_endpoints/gateway_endpoints.py | 229 ++++++++++-------- 1 file changed, 131 insertions(+), 98 deletions(-) diff --git a/litellm/proxy/anthropic_endpoints/gateway_endpoints.py b/litellm/proxy/anthropic_endpoints/gateway_endpoints.py index e0c35117bf6..5a4a4d0eb78 100644 --- a/litellm/proxy/anthropic_endpoints/gateway_endpoints.py +++ b/litellm/proxy/anthropic_endpoints/gateway_endpoints.py @@ -17,10 +17,13 @@ is accepted by every bearer-authenticated proxy route. import hashlib import json import secrets -from typing import Any +from collections.abc import Mapping +from types import MappingProxyType +from typing import Final from fastapi import APIRouter, Depends, Request, Response from fastapi.responses import JSONResponse +from pydantic import BaseModel, Field, TypeAdapter from litellm.constants import ( CLI_JWT_EXPIRATION_HOURS, @@ -30,16 +33,58 @@ from litellm.constants import ( from litellm.proxy.anthropic_endpoints.endpoints import anthropic_response, count_tokens from litellm.proxy.auth.user_api_key_auth import user_api_key_auth -GATEWAY_PREFIX = "/claude_code_gateway" -_DEVICE_CODE_GRANT = "urn:ietf:params:oauth:grant-type:device_code" -_REFRESH_TOKEN_GRANT = "refresh_token" -_DEVICE_POLL_INTERVAL_SECONDS = 5 +GATEWAY_PREFIX: Final = "/claude_code_gateway" +_DEVICE_CODE_GRANT: Final = "urn:ietf:params:oauth:grant-type:device_code" +_REFRESH_TOKEN_GRANT: Final = "refresh_token" +_DEVICE_POLL_INTERVAL_SECONDS: Final = 5 +_SECONDS_PER_HOUR: Final = 3600 +_MANAGED_SETTINGS_ADAPTER: Final = TypeAdapter(dict[str, object]) +_NO_SETTINGS: Final = MappingProxyType({}) +_POST_ONLY: Final = ["POST"] # mutable-ok: FastAPI's add_api_route only accepts a list of methods + + +class _GatewaySessionData(BaseModel): + user_id: str + user_role: str | None = None + models: list[str] = Field(default_factory=list) + teams: tuple[str, ...] = () + + +class _OAuthErrorBody(BaseModel): + error: str + error_description: str | None = None + + +class _AuthorizationServerMetadata(BaseModel): + issuer: str + device_authorization_endpoint: str + token_endpoint: str + grant_types_supported: tuple[str, ...] + + +class _DeviceAuthorizationBody(BaseModel): + device_code: str + user_code: str + verification_uri: str + verification_uri_complete: str + expires_in: int + interval: int + + +class _AccessTokenBody(BaseModel): + access_token: str + expires_in: int + token_type: str = "Bearer" + + +def _general_settings() -> Mapping[str, object]: + from litellm.proxy.proxy_server import general_settings + + return general_settings or _NO_SETTINGS def _is_gateway_enabled() -> bool: - from litellm.proxy.proxy_server import general_settings - - return bool((general_settings or {}).get("enable_claude_code_gateway", False)) + return bool(_general_settings().get("enable_claude_code_gateway", False)) def ensure_gateway_enabled() -> None: @@ -49,11 +94,11 @@ def ensure_gateway_enabled() -> None: raise HTTPException(status_code=404, detail="Claude Code gateway is not enabled") -def _managed_settings() -> dict[str, Any] | None: - from litellm.proxy.proxy_server import general_settings - - settings = (general_settings or {}).get("claude_code_gateway_managed_settings") - return settings if isinstance(settings, dict) else None +def _managed_settings() -> dict[str, object] | None: + settings: Final[object] = _general_settings().get("claude_code_gateway_managed_settings") + if not isinstance(settings, dict): + return None + return _MANAGED_SETTINGS_ADAPTER.validate_python(settings) def _oauth_error(*, status_code: int, error: str, description: str | None = None) -> "_OAuthError": @@ -68,26 +113,29 @@ class _OAuthError(Exception): def _oauth_error_response(err: _OAuthError) -> JSONResponse: - body: dict[str, str] = {"error": err.error} - if err.description is not None: - body["error_description"] = err.description - return JSONResponse(status_code=err.status_code, content=body) + body: Final = _OAuthErrorBody(error=err.error, error_description=err.description) + return JSONResponse(status_code=err.status_code, content=body.model_dump(exclude_none=True)) -router = APIRouter(prefix=GATEWAY_PREFIX, tags=["Claude Code gateway"]) +router: Final = APIRouter( + prefix=GATEWAY_PREFIX, + tags=["Claude Code gateway"], # mutable-ok: FastAPI's APIRouter only accepts a list of tags +) +_GATEWAY_ENABLED: Final = (Depends(ensure_gateway_enabled),) +_AUTHENTICATED: Final = (Depends(user_api_key_auth),) router.add_api_route( "/v1/messages", anthropic_response, - methods=["POST"], - dependencies=[Depends(ensure_gateway_enabled)], + methods=_POST_ONLY, + dependencies=_GATEWAY_ENABLED, include_in_schema=False, ) router.add_api_route( "/v1/messages/count_tokens", count_tokens, - methods=["POST"], - dependencies=[Depends(ensure_gateway_enabled)], + methods=_POST_ONLY, + dependencies=_GATEWAY_ENABLED, include_in_schema=False, ) @@ -99,20 +147,16 @@ async def oauth_authorization_server(request: Request) -> JSONResponse: from litellm.proxy.utils import get_custom_url - request_base_url = str(request.base_url) - issuer = get_custom_url(request_base_url=request_base_url, route="claude_code_gateway") - return JSONResponse( - content={ - "issuer": issuer, - "device_authorization_endpoint": get_custom_url( - request_base_url=request_base_url, route="claude_code_gateway/oauth/device_authorization" - ), - "token_endpoint": get_custom_url( - request_base_url=request_base_url, route="claude_code_gateway/oauth/token" - ), - "grant_types_supported": [_DEVICE_CODE_GRANT, _REFRESH_TOKEN_GRANT], - } + request_base_url: Final = str(request.base_url) + metadata: Final = _AuthorizationServerMetadata( + issuer=get_custom_url(request_base_url=request_base_url, route="claude_code_gateway"), + device_authorization_endpoint=get_custom_url( + request_base_url=request_base_url, route="claude_code_gateway/oauth/device_authorization" + ), + token_endpoint=get_custom_url(request_base_url=request_base_url, route="claude_code_gateway/oauth/token"), + grant_types_supported=(_DEVICE_CODE_GRANT, _REFRESH_TOKEN_GRANT), ) + return JSONResponse(content=metadata.model_dump()) @router.post("/oauth/device_authorization", include_in_schema=False) @@ -120,13 +164,13 @@ async def device_authorization(request: Request) -> JSONResponse: from urllib.parse import urlencode from litellm.proxy.management_endpoints.ui_sso import ( - _check_cli_sso_start_rate_limit, - _generate_cli_sso_user_code, - _hash_cli_sso_secret, - _normalize_cli_sso_user_code, - _set_cli_sso_flow, + _check_cli_sso_start_rate_limit, # pyright: ignore[reportPrivateUsage] # shared device-flow helper + _generate_cli_sso_user_code, # pyright: ignore[reportPrivateUsage] # shared device-flow helper + _hash_cli_sso_secret, # pyright: ignore[reportPrivateUsage] # shared device-flow helper + _normalize_cli_sso_user_code, # pyright: ignore[reportPrivateUsage] # shared device-flow helper + _set_cli_sso_flow, # pyright: ignore[reportPrivateUsage] # shared device-flow helper ) - from litellm.proxy.proxy_server import cli_sso_session_cache, general_settings + from litellm.proxy.proxy_server import cli_sso_session_cache from litellm.proxy.utils import get_custom_url if not _is_gateway_enabled(): @@ -135,12 +179,12 @@ async def device_authorization(request: Request) -> JSONResponse: _check_cli_sso_start_rate_limit( request=request, cache=cli_sso_session_cache, - use_x_forwarded_for=bool((general_settings or {}).get("use_x_forwarded_for", False)), + use_x_forwarded_for=bool(_general_settings().get("use_x_forwarded_for", False)), ) - device_code = f"cli-{secrets.token_urlsafe(24)}" - user_code = _generate_cli_sso_user_code() - flow = { + device_code: Final = f"cli-{secrets.token_urlsafe(24)}" + user_code: Final = _generate_cli_sso_user_code() + flow: Final = { # mutable-ok: the shared CLI SSO cache entry is a dict the browser leg mutates "poll_secret_hash": _hash_cli_sso_secret(device_code), "user_code_hash": _hash_cli_sso_secret(_normalize_cli_sso_user_code(user_code)), "sso_complete": False, @@ -149,42 +193,36 @@ async def device_authorization(request: Request) -> JSONResponse: } _set_cli_sso_flow(login_id=device_code, cache=cli_sso_session_cache, flow=flow) - request_base_url = str(request.base_url) - verification_uri = get_custom_url(request_base_url=request_base_url, route="sso/key/generate") - verification_uri_complete = ( - verification_uri - + "?" - + urlencode({"source": LITELLM_CLI_SOURCE_IDENTIFIER, "key": device_code, "user_code": user_code}) - ) - verification_uri_no_code = ( - verification_uri + "?" + urlencode({"source": LITELLM_CLI_SOURCE_IDENTIFIER, "key": device_code}) - ) - return JSONResponse( - content={ - "device_code": device_code, - "user_code": user_code, - "verification_uri": verification_uri_no_code, - "verification_uri_complete": verification_uri_complete, - "expires_in": CLI_SSO_SESSION_TTL_SECONDS, - "interval": _DEVICE_POLL_INTERVAL_SECONDS, - } + request_base_url: Final = str(request.base_url) + verification_uri: Final = get_custom_url(request_base_url=request_base_url, route="sso/key/generate") + query: Final = MappingProxyType({"source": LITELLM_CLI_SOURCE_IDENTIFIER, "key": device_code}) + body: Final = _DeviceAuthorizationBody( + device_code=device_code, + user_code=user_code, + verification_uri=f"{verification_uri}?{urlencode(query)}", + verification_uri_complete=( + f"{verification_uri}?{urlencode(MappingProxyType({**query, 'user_code': user_code}))}" + ), + expires_in=CLI_SSO_SESSION_TTL_SECONDS, + interval=_DEVICE_POLL_INTERVAL_SECONDS, ) + return JSONResponse(content=body.model_dump()) -def _mint_access_token_from_flow(flow: dict[str, Any]) -> str: +def _mint_access_token_from_flow(flow: Mapping[str, object]) -> str: from litellm.proxy._types import LiteLLM_UserTable from litellm.proxy.auth.auth_checks import ExperimentalUIJWTToken - session_data = flow.get("session_data") - if not isinstance(session_data, dict): + raw_session_data: Final = flow.get("session_data") + if not isinstance(raw_session_data, dict): raise _oauth_error(status_code=400, error="authorization_pending") - teams = session_data.get("teams") or [] - team_id = teams[0] if isinstance(teams, list) and teams else None - user_info = LiteLLM_UserTable( - user_id=session_data["user_id"], - user_role=session_data["user_role"], - models=session_data.get("models", []), + session_data: Final = _GatewaySessionData.model_validate(raw_session_data) + team_id: Final = session_data.teams[0] if session_data.teams else None + user_info: Final = LiteLLM_UserTable( + user_id=session_data.user_id, + user_role=session_data.user_role, + models=session_data.models, ) return ExperimentalUIJWTToken.get_cli_jwt_auth_token(user_info=user_info, team_id=team_id) @@ -193,8 +231,8 @@ async def _handle_device_code_grant(device_code: str | None) -> JSONResponse: from fastapi import HTTPException from litellm.proxy.management_endpoints.ui_sso import ( - _get_cli_sso_flow_cache_key, - _get_cli_sso_flow_or_raise, + _get_cli_sso_flow_cache_key, # pyright: ignore[reportPrivateUsage] # shared device-flow helper + _get_cli_sso_flow_or_raise, # pyright: ignore[reportPrivateUsage] # shared device-flow helper ) from litellm.proxy.proxy_server import cli_sso_session_cache @@ -204,7 +242,7 @@ async def _handle_device_code_grant(device_code: str | None) -> JSONResponse: ) try: - flow = _get_cli_sso_flow_or_raise(login_id=device_code, cache=cli_sso_session_cache) + flow: Final = _get_cli_sso_flow_or_raise(login_id=device_code, cache=cli_sso_session_cache) except HTTPException: return _oauth_error_response(_oauth_error(status_code=400, error="expired_token")) @@ -212,18 +250,13 @@ async def _handle_device_code_grant(device_code: str | None) -> JSONResponse: return _oauth_error_response(_oauth_error(status_code=400, error="authorization_pending")) try: - access_token = _mint_access_token_from_flow(flow) + access_token: Final = _mint_access_token_from_flow(flow) except _OAuthError as err: return _oauth_error_response(err) cli_sso_session_cache.delete_cache(key=_get_cli_sso_flow_cache_key(device_code)) - return JSONResponse( - content={ - "access_token": access_token, - "token_type": "Bearer", - "expires_in": CLI_JWT_EXPIRATION_HOURS * 3600, - } - ) + body: Final = _AccessTokenBody(access_token=access_token, expires_in=CLI_JWT_EXPIRATION_HOURS * _SECONDS_PER_HOUR) + return JSONResponse(content=body.model_dump()) @router.post("/oauth/token", include_in_schema=False) @@ -231,11 +264,11 @@ async def oauth_token(request: Request) -> JSONResponse: if not _is_gateway_enabled(): return _oauth_error_response(_oauth_error(status_code=404, error="not_found")) - form = await request.form() - grant_type = form.get("grant_type") + form: Final = await request.form() + grant_type: Final = form.get("grant_type") if grant_type == _DEVICE_CODE_GRANT: - device_code = form.get("device_code") + device_code: Final = form.get("device_code") return await _handle_device_code_grant(device_code if isinstance(device_code, str) else None) if grant_type == _REFRESH_TOKEN_GRANT: @@ -254,20 +287,20 @@ async def oauth_token(request: Request) -> JSONResponse: ) -@router.get("/managed/settings", include_in_schema=False, dependencies=[Depends(user_api_key_auth)]) +@router.get("/managed/settings", include_in_schema=False, dependencies=_AUTHENTICATED) async def managed_settings(request: Request) -> Response: ensure_gateway_enabled() - settings = _managed_settings() + settings: Final = _managed_settings() if settings is None: return Response(status_code=404) - body = json.dumps(settings, sort_keys=True, separators=(",", ":")) - etag = '"' + hashlib.sha256(body.encode("utf-8")).hexdigest() + '"' - if_none_match = request.headers.get("If-None-Match") - if if_none_match is not None and if_none_match == etag: - return Response(status_code=304, headers={"ETag": etag}) - return Response(content=body, media_type="application/json", headers={"ETag": etag}) + body: Final = json.dumps(settings, sort_keys=True, separators=(",", ":")) + etag: Final = '"' + hashlib.sha256(body.encode("utf-8")).hexdigest() + '"' + headers: Final = MappingProxyType({"ETag": etag}) + if request.headers.get("If-None-Match") == etag: + return Response(status_code=304, headers=headers) + return Response(content=body, media_type="application/json", headers=headers) def _accept_otlp() -> Response: @@ -275,16 +308,16 @@ def _accept_otlp() -> Response: return Response(status_code=200) -@router.post("/v1/metrics", include_in_schema=False, dependencies=[Depends(user_api_key_auth)]) +@router.post("/v1/metrics", include_in_schema=False, dependencies=_AUTHENTICATED) async def otlp_metrics() -> Response: return _accept_otlp() -@router.post("/v1/logs", include_in_schema=False, dependencies=[Depends(user_api_key_auth)]) +@router.post("/v1/logs", include_in_schema=False, dependencies=_AUTHENTICATED) async def otlp_logs() -> Response: return _accept_otlp() -@router.post("/v1/traces", include_in_schema=False, dependencies=[Depends(user_api_key_auth)]) +@router.post("/v1/traces", include_in_schema=False, dependencies=_AUTHENTICATED) async def otlp_traces() -> Response: return _accept_otlp() From 9d4bab3b700c1170b7451bedfbda67cf830f58f3 Mon Sep 17 00:00:00 2001 From: zoroyihan7 Date: Tue, 8 Sep 2026 10:46:16 +0000 Subject: [PATCH 007/144] fix(responses): emit typed streaming failure events --- litellm/exceptions.py | 3 +- .../common_utils/responses_stream_errors.py | 119 ++++++++++++++++++ litellm/proxy/proxy_server.py | 21 +++- .../proxy/response_api_endpoints/endpoints.py | 6 +- litellm/responses/streaming_iterator.py | 5 + .../proxy_server/test_streaming_helpers.py | 85 ++++++++++++- .../response_api_endpoints/test_endpoints.py | 95 +++++++++++++- 7 files changed, 328 insertions(+), 6 deletions(-) create mode 100644 litellm/proxy/common_utils/responses_stream_errors.py diff --git a/litellm/exceptions.py b/litellm/exceptions.py index f9215267bf3..318b227ffa8 100644 --- a/litellm/exceptions.py +++ b/litellm/exceptions.py @@ -789,6 +789,7 @@ class APIError(openai.APIError): litellm_debug_info: str | None = None, max_retries: int | None = None, num_retries: int | None = None, + body: object | None = None, ): self.status_code = status_code self.message = f"litellm.APIError: {message}" @@ -799,7 +800,7 @@ class APIError(openai.APIError): self.num_retries = num_retries if request is None: request = httpx.Request(method="POST", url="https://api.openai.com/v1") - super().__init__(self.message, request=request, body=None) + super().__init__(self.message, request=request, body=body) def __str__(self): _message = self.message diff --git a/litellm/proxy/common_utils/responses_stream_errors.py b/litellm/proxy/common_utils/responses_stream_errors.py new file mode 100644 index 00000000000..a3bee06e912 --- /dev/null +++ b/litellm/proxy/common_utils/responses_stream_errors.py @@ -0,0 +1,119 @@ +import time +from collections.abc import Mapping +from types import MappingProxyType +from typing import Final + +from pydantic import BaseModel, ConfigDict + +from litellm._logging import redact_internal_details_from_client_message +from litellm._uuid import uuid +from litellm.exceptions import MidStreamFallbackError +from litellm.types.llms.openai import ResponseFailedEvent, ResponsesAPIResponse, ResponsesAPIStreamEvents + + +class _ResponseIdentity(BaseModel): + model_config = ConfigDict(frozen=True, from_attributes=True) + + id: str | None = None + model: str | None = None + created_at: int | None = None + + +class _StreamEvent(BaseModel): + model_config = ConfigDict(frozen=True, from_attributes=True) + + type: str | None = None + sequence_number: int | None = None + response: _ResponseIdentity | None = None + + +class _FailureDetails(BaseModel): + model_config = ConfigDict(frozen=True, from_attributes=True) + + message: str | None = None + code: str | int | None = None + type: str | None = None + status_code: int | None = None + + +def _original_failure(exception: Exception) -> Exception: + if isinstance(exception, MidStreamFallbackError) and exception.original_exception is not None: + return _original_failure(exception.original_exception) + return exception + + +def _response_error_code(details: _FailureDetails) -> str: + for value in (details.code, details.type): + if value == "insufficient_quota": + return "insufficient_quota" + if value in (429, "429") or isinstance(value, str) and value.startswith("rate_limit"): + return "rate_limit_exceeded" + if isinstance(details.code, str) and details.code and not details.code.isdecimal(): + return details.code + if details.status_code == 429: + return "rate_limit_exceeded" + return "server_error" + + +class ResponsesStreamErrorState: + def __init__(self) -> None: + self.response_id: str | None = None + self.model: str | None = None + self.created_at: int | None = None + self.sequence_number = -1 + self.terminal_emitted = False + + @staticmethod + def observe_chunk(chunk: object) -> _StreamEvent | None: + if not isinstance(chunk, (BaseModel, Mapping)): + return None + return _StreamEvent.model_validate(chunk) + + def mark_emitted(self, event: _StreamEvent | None) -> None: + if event is None: + return + if event.sequence_number is not None: + self.sequence_number = max(self.sequence_number, event.sequence_number) + if event.response is not None: + self.response_id = event.response.id or self.response_id + self.model = event.response.model or self.model + if event.response.created_at is not None: + self.created_at = event.response.created_at + if event.type in ("response.completed", "response.failed", "response.incomplete"): + self.terminal_emitted = True + + def format_failure(self, exception: Exception) -> str | None: + if self.terminal_emitted: + return None + original: Final = _original_failure(exception) + details: Final = _FailureDetails.model_validate(original) + response: Final = ResponsesAPIResponse.model_validate( + MappingProxyType( + { + "id": self.response_id or f"resp_{uuid.uuid4().hex}", + "object": "response", + "created_at": self.created_at if self.created_at is not None else int(time.time()), + "model": self.model, + "status": "failed", + "output": (), + "error": MappingProxyType( + { + "code": _response_error_code(details), + "message": redact_internal_details_from_client_message(details.message or str(original)), + } + ), + } + ) + ) + event: Final = ResponseFailedEvent.model_validate( + MappingProxyType( + { + "type": ResponsesAPIStreamEvents.RESPONSE_FAILED, + "response": response, + "sequence_number": self.sequence_number + 1, + } + ) + ) + payload: Final = event.model_dump_json(exclude_none=True) + self.terminal_emitted = True + return f"event: response.failed\ndata: {payload}\n\n" diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 32b6b841af7..9f1d69acefa 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -385,6 +385,7 @@ from litellm.proxy.common_utils.periodic_reload_schedule import ( ) from litellm.proxy.common_utils.proxy_state import ProxyState from litellm.proxy.common_utils.reset_budget_job import ResetBudgetJob +from litellm.proxy.common_utils.responses_stream_errors import ResponsesStreamErrorState from litellm.proxy.common_utils.scheduled_job_stagger import ( apply_scheduled_job_stagger, attach_job_timing_logger, @@ -8723,10 +8724,13 @@ async def async_data_generator( user_api_key_dict: UserAPIKeyAuth, request_data: dict, request: Request | None = None, + *, + responses_stream_errors: bool = False, ): verbose_proxy_logger.debug("inside generator") stream_completed = False client_disconnected = False + error_state: Final = ResponsesStreamErrorState() if responses_stream_errors else None try: error_message: str | None = None requested_model_from_client: Final = _get_client_requested_model_for_streaming(request_data=request_data) @@ -8837,6 +8841,7 @@ async def async_data_generator( fallback_metadata_event_sent = True continue + responses_event: Final = error_state.observe_chunk(chunk) if error_state is not None else None raw_passthrough = False if isinstance(chunk, BaseModel): chunk = _serialize_streaming_chunk(chunk) @@ -8871,8 +8876,13 @@ async def async_data_generator( if not raw_passthrough: try: - yield _format_streaming_sse_chunk(chunk=chunk) + formatted_chunk: Final = _format_streaming_sse_chunk(chunk=chunk) + if error_state is not None: + error_state.mark_emitted(responses_event) + yield formatted_chunk except Exception as e: + if error_state is not None: + raise yield f"data: {e}\n\n" if pending_fallback_event: @@ -8922,6 +8932,12 @@ async def async_data_generator( e, ) + if error_state is not None: + stream_completed = True + error_frame: Final = error_state.format_failure(e) + if error_frame is not None: + yield error_frame + return if isinstance(e, HTTPException): raise e elif isinstance(e, StreamingCallbackError): @@ -8958,12 +8974,15 @@ def select_data_generator( user_api_key_dict: UserAPIKeyAuth, request_data: dict, request: Request | None = None, + *, + responses_stream_errors: bool = False, ): return async_data_generator( response=response, user_api_key_dict=user_api_key_dict, request_data=request_data, request=request, + responses_stream_errors=responses_stream_errors, ) diff --git a/litellm/proxy/response_api_endpoints/endpoints.py b/litellm/proxy/response_api_endpoints/endpoints.py index 5907ffc64eb..3202fabe74e 100644 --- a/litellm/proxy/response_api_endpoints/endpoints.py +++ b/litellm/proxy/response_api_endpoints/endpoints.py @@ -3,6 +3,7 @@ import json import time from collections.abc import AsyncIterator, Awaitable, Mapping from enum import Enum +from functools import partial from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, NamedTuple, Protocol, cast, get_args from uuid import uuid4 @@ -243,6 +244,7 @@ async def responses_api( version, ) + native_data_generator: Final = partial(select_data_generator, responses_stream_errors=True) data = await _read_request_body(request=request) # Check if polling via cache should be used for this request @@ -329,7 +331,7 @@ async def responses_api( llm_router=llm_router, proxy_config=proxy_config, proxy_logging_obj=proxy_logging_obj, - select_data_generator=select_data_generator, + select_data_generator=native_data_generator, user_model=user_model, user_temperature=user_temperature, user_request_timeout=user_request_timeout, @@ -355,7 +357,7 @@ async def responses_api( llm_router=llm_router, general_settings=general_settings, proxy_config=proxy_config, - select_data_generator=select_data_generator, + select_data_generator=native_data_generator, model=None, user_model=user_model, user_temperature=user_temperature, diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 9f9016c5a7f..1134f6b07e3 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -579,6 +579,11 @@ class BaseResponsesAPIStreamingIterator: message=error_message, llm_provider=self.custom_llm_provider or "", model=self.model or "", + body={ # mutable-ok: OpenAI APIError reads code/type only from a dict body + "code": error_code, + "type": error_type, + "message": error_message, + }, ) if 400 <= status_code < 500 and status_code != 429: raise mapped_exception diff --git a/tests/test_litellm/proxy/proxy_server/test_streaming_helpers.py b/tests/test_litellm/proxy/proxy_server/test_streaming_helpers.py index 87e10ce7e8d..6f3a47a4b79 100644 --- a/tests/test_litellm/proxy/proxy_server/test_streaming_helpers.py +++ b/tests/test_litellm/proxy/proxy_server/test_streaming_helpers.py @@ -17,15 +17,18 @@ from __future__ import annotations import asyncio import json +from collections.abc import AsyncIterator +from typing import Final, Literal from unittest.mock import AsyncMock, MagicMock import pytest from fastapi import Response from fastapi.responses import StreamingResponse +from pydantic import BaseModel import litellm -from litellm.constants import RETURN_RAW_MODEL_NAME_METADATA_KEY import litellm.proxy.proxy_server as ps +from litellm.constants import RETURN_RAW_MODEL_NAME_METADATA_KEY from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.proxy_server import ( _apply_streaming_chunk_hooks, @@ -42,6 +45,12 @@ from litellm.proxy.proxy_server import ( data_generator, select_data_generator, ) +from litellm.types.llms.openai import ( + ResponseCompletedEvent, + ResponseCreatedEvent, + ResponseFailedEvent, + ResponsesAPIResponse, +) from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices, Usage from .conftest import normalize @@ -872,6 +881,80 @@ async def test_async_data_generator_mid_stream_exception_yields_error_payload( assert any(isinstance(item, str) and item.startswith('data: {"error":') for item in out) +@pytest.mark.asyncio +@pytest.mark.parametrize("terminal", ["completed", "serialization_failure", "failure_after_completed"]) +async def test_responses_stream_keeps_tool_deltas_and_only_emits_a_valid_terminal( + terminal: Literal["completed", "serialization_failure", "failure_after_completed"], +) -> None: + class ToolDelta(BaseModel): + type: Literal["response.function_call_arguments.delta"] + sequence_number: int + item_id: str + output_index: int + delta: str + + class UnserializableTerminal(BaseModel): + type: Literal["response.completed"] + sequence_number: int + response: ResponsesAPIResponse + invalid: object + + response: Final = ResponsesAPIResponse(id="resp_visible", created_at=1, model="gpt-6-astra", output=[]) + created: Final = ResponseCreatedEvent.model_validate( + {"type": "response.created", "sequence_number": 0, "response": response} + ) + completed: Final = ResponseCompletedEvent.model_validate( + {"type": "response.completed", "sequence_number": 2, "response": response} + ) + tool_delta: Final = ToolDelta( + type="response.function_call_arguments.delta", sequence_number=1, item_id="fc_stream_error", + output_index=0, delta='{"path":"partial', + ) + + async def upstream() -> AsyncIterator[BaseModel]: + yield created + yield tool_delta + yield ( + UnserializableTerminal(type="response.completed", sequence_number=2, response=response, invalid=object()) + if terminal == "serialization_failure" else completed + ) + if terminal == "failure_after_completed": + raise litellm.APIError(status_code=500, message="Stream close failed", llm_provider="openai", model="gpt-6-astra") + + frames: Final = [ + frame + async for frame in select_data_generator( + response=upstream(), + user_api_key_dict=_user_auth(), + request_data={}, + responses_stream_errors=True, + ) + ] + decoded: Final = tuple(frame.decode() if isinstance(frame, bytes) else frame for frame in frames) + event_frames: Final = tuple(frame for frame in decoded if frame != "data: [DONE]\n\n") + payloads: Final = tuple( + json.loads(next(line[6:] for line in frame.splitlines() if line.startswith("data: "))) + for frame in event_frames + ) + + assert payloads[0]["response"]["id"] == "resp_visible" + assert payloads[1] == tool_delta.model_dump() + assert len(payloads) == 3 + if terminal == "serialization_failure": + failure: Final = ResponseFailedEvent.model_validate(payloads[-1]) + assert event_frames[-1].startswith("event: response.failed\n") + assert failure.response.id == "resp_visible" + assert failure.response.status == "failed" + assert failure.response.error is not None + assert failure.response.error["code"] == "server_error" + assert "serialize" in failure.response.error["message"].lower() + assert payloads[-1]["sequence_number"] > payloads[1]["sequence_number"] + else: + assert payloads[-1]["type"] == "response.completed" + assert payloads[-1]["sequence_number"] == 2 + assert "error" not in payloads[-1] + + # --------------------------------------------------------------------------- # select_data_generator # --------------------------------------------------------------------------- diff --git a/tests/test_litellm/proxy/response_api_endpoints/test_endpoints.py b/tests/test_litellm/proxy/response_api_endpoints/test_endpoints.py index d7010de6405..9e2f70a3d64 100644 --- a/tests/test_litellm/proxy/response_api_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/response_api_endpoints/test_endpoints.py @@ -3,10 +3,12 @@ Test for response_api_endpoints/endpoints.py """ import unittest -from typing import Any +from typing import Any, Final, Literal from unittest.mock import AsyncMock, MagicMock, patch +import httpx import pytest +import respx from fastapi.testclient import TestClient from httpx import Response @@ -14,6 +16,97 @@ import litellm from litellm.proxy.proxy_server import app +@pytest.mark.asyncio +@pytest.mark.parametrize( + "path,error_kind", + [ + ("/v1/responses", "rate_limit"), + ("/v1/responses", "numeric_rate_limit"), + ("/v1/responses", "server_error"), + ("/v1/responses", "response_failed"), + ("/cursor/chat/completions", "server_error"), + ("/v1/chat/completions", "server_error"), + ], +) +async def test_streaming_upstream_errors_keep_the_client_protocol( + monkeypatch: pytest.MonkeyPatch, + path: str, + error_kind: Literal["rate_limit", "numeric_rate_limit", "server_error", "response_failed"], +) -> None: + import litellm.proxy.proxy_server as ps + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + + model: Final = "gpt-6-astra" + message: Final = "Upstream cannot complete this response" + code: Final = { + "rate_limit": "rate_limit_exceeded", "numeric_rate_limit": "429", + "server_error": "server_error", "response_failed": "server_error", + }[error_kind] + error: Final = {"message": message, "code": code, "type": None, "param": "input"} + response: Final = {"id": "resp_upstream", "object": "response", "created_at": 1, + "status": "in_progress", "model": model, "output": [], + "parallel_tool_calls": True, "tool_choice": "auto", "tools": []} + created: Final = {"type": "response.created", "sequence_number": 0, "response": response} + tool_added: Final = {"type": "response.output_item.added", "sequence_number": 1, "output_index": 0, + "item": {"type": "function_call", "id": "fc_partial", "call_id": "call_partial", + "name": "read_file", "arguments": "", "status": "in_progress"}} + tool_delta: Final = {"type": "response.function_call_arguments.delta", "sequence_number": 2, + "item_id": "fc_partial", "output_index": 0, "delta": '{"path":"partial'} + failed: Final = ( + {"type": "response.failed", "sequence_number": 9, + "response": {**response, "status": "failed", "error": error}} + if error_kind == "response_failed" else {"type": "error", "error": error} + ) + chat: Final = {"id": "chatcmpl_partial", "object": "chat.completion.chunk", "created": 1, + "model": model, "choices": [{"index": 0, "delta": {"content": "partial"}, + "finish_reason": None}]} + is_chat: Final = path == "/v1/chat/completions" + upstream_events: Final = (chat, {"error": error}) if is_chat else (created, tool_added, tool_delta, failed) + wire: Final = "".join("data: " + json.dumps(event) + "\n\n" for event in upstream_events) + upstream_url: Final = "https://streaming.example/v1" + router: Final = litellm.Router( + model_list=[{"model_name": model, "litellm_params": { + "model": "openai/" + model, "api_base": upstream_url, "api_key": "fixture-key"}}], + num_retries=0, + ) + monkeypatch.setattr(ps, "llm_router", router) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setitem(app.dependency_overrides, user_api_key_auth, _auth_override) + with respx.mock as transport: + transport.post(upstream_url + ("/chat/completions" if is_chat else "/responses")).respond( + 200, content=wire, headers={"Content-Type": "text/event-stream"} + ) + async with httpx.AsyncClient(transport=httpx.ASGITransport(app), base_url="http://testserver") as client: + result: Final = await client.post( + path, json={"model": model, "stream": True, + **({"messages": [{"role": "user", "content": "hello"}]} if is_chat else {"input": "hello"})}, + ) + frames: Final = tuple(frame for frame in result.text.split("\n\n") if "data: " in frame) + events: Final = tuple( + json.loads(next(line[6:] for line in frame.splitlines() if line.startswith("data: "))) + for frame in frames if "data: [DONE]" not in frame + ) + + assert result.status_code == 200, result.text + assert message in result.text + if path == "/v1/responses": + assert frames[-1].startswith("event: response.failed\n"), result.text + assert [event["type"] for event in events] == [ + "response.created", "response.output_item.added", "response.function_call_arguments.delta", "response.failed" + ] + assert events[2]["delta"] == tool_delta["delta"] + assert events[-1]["sequence_number"] == events[-2]["sequence_number"] + 1 + assert events[-1]["response"]["id"] == events[0]["response"]["id"] + assert events[-1]["response"]["status"] == "failed" + assert events[-1]["response"]["error"]["code"] == ( + "rate_limit_exceeded" if error_kind in ("rate_limit", "numeric_rate_limit") else "server_error" + ) + else: + assert events[0]["object"] == "chat.completion.chunk", result.text + assert "response.failed" not in result.text + assert "error" in events[-1] + + class TestResponsesAPIEndpoints(unittest.TestCase): @pytest.mark.asyncio @patch("litellm.proxy.proxy_server.llm_router") From 089703d10b2eebbec1a5b780494686451ab72161 Mon Sep 17 00:00:00 2001 From: zoroyihan7 Date: Tue, 8 Sep 2026 11:02:38 +0000 Subject: [PATCH 008/144] fix(responses): preserve failure metadata at streaming boundaries --- .../common_utils/responses_stream_errors.py | 45 +++++++++---- litellm/proxy/proxy_server.py | 9 +-- .../proxy_server/test_streaming_helpers.py | 63 ++++++++++++++++--- .../response_api_endpoints/test_endpoints.py | 28 ++++++--- 4 files changed, 114 insertions(+), 31 deletions(-) diff --git a/litellm/proxy/common_utils/responses_stream_errors.py b/litellm/proxy/common_utils/responses_stream_errors.py index a3bee06e912..356e948a2df 100644 --- a/litellm/proxy/common_utils/responses_stream_errors.py +++ b/litellm/proxy/common_utils/responses_stream_errors.py @@ -1,9 +1,10 @@ import time from collections.abc import Mapping +from http import HTTPStatus from types import MappingProxyType from typing import Final -from pydantic import BaseModel, ConfigDict +from pydantic import BaseModel, ConfigDict, field_validator from litellm._logging import redact_internal_details_from_client_message from litellm._uuid import uuid @@ -35,6 +36,16 @@ class _FailureDetails(BaseModel): type: str | None = None status_code: int | None = None + @field_validator("code", mode="before") + @classmethod + def normalize_code(cls, value: object) -> str | int | None: + return value if isinstance(value, (str, int)) and not isinstance(value, bool) else None + + @field_validator("type", mode="before") + @classmethod + def normalize_type(cls, value: object) -> str | None: + return value if isinstance(value, str) else None + def _original_failure(exception: Exception) -> Exception: if isinstance(exception, MidStreamFallbackError) and exception.original_exception is not None: @@ -50,9 +61,21 @@ def _response_error_code(details: _FailureDetails) -> str: return "rate_limit_exceeded" if isinstance(details.code, str) and details.code and not details.code.isdecimal(): return details.code - if details.status_code == 429: - return "rate_limit_exceeded" - return "server_error" + match details.status_code: + case HTTPStatus.UNAUTHORIZED: + return "authentication_error" + case HTTPStatus.FORBIDDEN: + return "permission_error" + case HTTPStatus.NOT_FOUND: + return "not_found_error" + case HTTPStatus.REQUEST_TIMEOUT: + return "request_timeout" + case HTTPStatus.TOO_MANY_REQUESTS: + return "rate_limit_exceeded" + case int(status) if HTTPStatus.BAD_REQUEST <= status < HTTPStatus.INTERNAL_SERVER_ERROR: + return "invalid_request_error" + case _: + return "server_error" class ResponsesStreamErrorState: @@ -62,16 +85,15 @@ class ResponsesStreamErrorState: self.created_at: int | None = None self.sequence_number = -1 self.terminal_emitted = False + self._pending_event: _StreamEvent | None = None - @staticmethod - def observe_chunk(chunk: object) -> _StreamEvent | None: - if not isinstance(chunk, (BaseModel, Mapping)): - return None - return _StreamEvent.model_validate(chunk) + def observe_chunk(self, chunk: object) -> None: + self._pending_event = _StreamEvent.model_validate(chunk) if isinstance(chunk, (BaseModel, Mapping)) else None - def mark_emitted(self, event: _StreamEvent | None) -> None: + def mark_emitted(self, frame: str | bytes) -> str | bytes: + event: Final = self._pending_event if event is None: - return + return frame if event.sequence_number is not None: self.sequence_number = max(self.sequence_number, event.sequence_number) if event.response is not None: @@ -81,6 +103,7 @@ class ResponsesStreamErrorState: self.created_at = event.response.created_at if event.type in ("response.completed", "response.failed", "response.incomplete"): self.terminal_emitted = True + return frame def format_failure(self, exception: Exception) -> str | None: if self.terminal_emitted: diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 9f1d69acefa..870cd78aa85 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -8841,7 +8841,8 @@ async def async_data_generator( fallback_metadata_event_sent = True continue - responses_event: Final = error_state.observe_chunk(chunk) if error_state is not None else None + if error_state is not None: + error_state.observe_chunk(cast(object, chunk)) # cast-ok: the helper validates legacy untyped chunks raw_passthrough = False if isinstance(chunk, BaseModel): chunk = _serialize_streaming_chunk(chunk) @@ -8876,10 +8877,10 @@ async def async_data_generator( if not raw_passthrough: try: - formatted_chunk: Final = _format_streaming_sse_chunk(chunk=chunk) if error_state is not None: - error_state.mark_emitted(responses_event) - yield formatted_chunk + yield error_state.mark_emitted(_format_streaming_sse_chunk(chunk=chunk)) + else: + yield _format_streaming_sse_chunk(chunk=chunk) except Exception as e: if error_state is not None: raise diff --git a/tests/test_litellm/proxy/proxy_server/test_streaming_helpers.py b/tests/test_litellm/proxy/proxy_server/test_streaming_helpers.py index 6f3a47a4b79..53e055882e8 100644 --- a/tests/test_litellm/proxy/proxy_server/test_streaming_helpers.py +++ b/tests/test_litellm/proxy/proxy_server/test_streaming_helpers.py @@ -21,9 +21,11 @@ from collections.abc import AsyncIterator from typing import Final, Literal from unittest.mock import AsyncMock, MagicMock +import httpx import pytest -from fastapi import Response +from fastapi import HTTPException, Response from fastapi.responses import StreamingResponse +from openai import APIError as OpenAIAPIError from pydantic import BaseModel import litellm @@ -882,9 +884,44 @@ async def test_async_data_generator_mid_stream_exception_yields_error_payload( @pytest.mark.asyncio -@pytest.mark.parametrize("terminal", ["completed", "serialization_failure", "failure_after_completed"]) +@pytest.mark.parametrize( + "terminal,upstream_error,expected_code", + [ + ("completed", None, None), + ("serialization_failure", None, "server_error"), + ("failure_after_completed", None, None), + pytest.param( + "upstream_failure", + litellm.AuthenticationError( + message="Upstream rejected request", llm_provider="openai", model="gpt-6-astra" + ), + "authentication_error", id="authentication_error", + ), + pytest.param( + "upstream_failure", + OpenAIAPIError( + message="Upstream rejected request", + request=httpx.Request("POST", "https://streaming.example/v1/responses"), + body={"code": {"reason": "overloaded"}, "type": {"unexpected": "object"}}, + ), + "server_error", id="structured_provider_error_fields", + ), + *( + pytest.param( + "upstream_failure", HTTPException(status_code=status, detail="Upstream rejected request"), + code, id=f"http_{status}", + ) + for status, code in ( + (400, "invalid_request_error"), (403, "permission_error"), (404, "not_found_error"), + (408, "request_timeout"), (422, "invalid_request_error"), (500, "server_error"), (503, "server_error"), + ) + ), + ], +) async def test_responses_stream_keeps_tool_deltas_and_only_emits_a_valid_terminal( - terminal: Literal["completed", "serialization_failure", "failure_after_completed"], + terminal: Literal["completed", "serialization_failure", "failure_after_completed", "upstream_failure"], + upstream_error: HTTPException | OpenAIAPIError | None, + expected_code: str | None, ) -> None: class ToolDelta(BaseModel): type: Literal["response.function_call_arguments.delta"] @@ -910,16 +947,23 @@ async def test_responses_stream_keeps_tool_deltas_and_only_emits_a_valid_termina type="response.function_call_arguments.delta", sequence_number=1, item_id="fc_stream_error", output_index=0, delta='{"path":"partial', ) + original_status: Final = ( + upstream_error.status_code if isinstance(upstream_error, (HTTPException, litellm.AuthenticationError)) else None + ) async def upstream() -> AsyncIterator[BaseModel]: yield created yield tool_delta + if upstream_error is not None: + raise upstream_error yield ( UnserializableTerminal(type="response.completed", sequence_number=2, response=response, invalid=object()) if terminal == "serialization_failure" else completed ) if terminal == "failure_after_completed": - raise litellm.APIError(status_code=500, message="Stream close failed", llm_provider="openai", model="gpt-6-astra") + raise litellm.APIError( + status_code=500, message="Stream close failed", llm_provider="openai", model="gpt-6-astra" + ) frames: Final = [ frame @@ -940,14 +984,19 @@ async def test_responses_stream_keeps_tool_deltas_and_only_emits_a_valid_termina assert payloads[0]["response"]["id"] == "resp_visible" assert payloads[1] == tool_delta.model_dump() assert len(payloads) == 3 - if terminal == "serialization_failure": + if terminal in ("serialization_failure", "upstream_failure"): failure: Final = ResponseFailedEvent.model_validate(payloads[-1]) assert event_frames[-1].startswith("event: response.failed\n") assert failure.response.id == "resp_visible" assert failure.response.status == "failed" assert failure.response.error is not None - assert failure.response.error["code"] == "server_error" - assert "serialize" in failure.response.error["message"].lower() + assert failure.response.error["code"] == expected_code + if upstream_error is None: + assert "serialize" in failure.response.error["message"].lower() + else: + assert "Upstream rejected request" in failure.response.error["message"] + if isinstance(upstream_error, (HTTPException, litellm.AuthenticationError)): + assert upstream_error.status_code == original_status assert payloads[-1]["sequence_number"] > payloads[1]["sequence_number"] else: assert payloads[-1]["type"] == "response.completed" diff --git a/tests/test_litellm/proxy/response_api_endpoints/test_endpoints.py b/tests/test_litellm/proxy/response_api_endpoints/test_endpoints.py index 9e2f70a3d64..7b79e1b8613 100644 --- a/tests/test_litellm/proxy/response_api_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/response_api_endpoints/test_endpoints.py @@ -61,7 +61,9 @@ async def test_streaming_upstream_errors_keep_the_client_protocol( "model": model, "choices": [{"index": 0, "delta": {"content": "partial"}, "finish_reason": None}]} is_chat: Final = path == "/v1/chat/completions" - upstream_events: Final = (chat, {"error": error}) if is_chat else (created, tool_added, tool_delta, failed) + partial: Final = path != "/v1/responses" or error_kind in ("numeric_rate_limit", "response_failed") + response_events: Final = (created, tool_added, tool_delta, failed) if partial else (failed,) + upstream_events: Final = (chat, {"error": error}) if is_chat else response_events wire: Final = "".join("data: " + json.dumps(event) + "\n\n" for event in upstream_events) upstream_url: Final = "https://streaming.example/v1" router: Final = litellm.Router( @@ -78,8 +80,10 @@ async def test_streaming_upstream_errors_keep_the_client_protocol( ) async with httpx.AsyncClient(transport=httpx.ASGITransport(app), base_url="http://testserver") as client: result: Final = await client.post( - path, json={"model": model, "stream": True, - **({"messages": [{"role": "user", "content": "hello"}]} if is_chat else {"input": "hello"})}, + path, json={ + "model": model, "stream": True, + **({"messages": [{"role": "user", "content": "hello"}]} if is_chat else {"input": "hello"}), + }, ) frames: Final = tuple(frame for frame in result.text.split("\n\n") if "data: " in frame) events: Final = tuple( @@ -91,12 +95,18 @@ async def test_streaming_upstream_errors_keep_the_client_protocol( assert message in result.text if path == "/v1/responses": assert frames[-1].startswith("event: response.failed\n"), result.text - assert [event["type"] for event in events] == [ - "response.created", "response.output_item.added", "response.function_call_arguments.delta", "response.failed" - ] - assert events[2]["delta"] == tool_delta["delta"] - assert events[-1]["sequence_number"] == events[-2]["sequence_number"] + 1 - assert events[-1]["response"]["id"] == events[0]["response"]["id"] + if partial: + assert [event["type"] for event in events] == [ + "response.created", "response.output_item.added", + "response.function_call_arguments.delta", "response.failed", + ] + assert events[2]["delta"] == tool_delta["delta"] + assert events[-1]["sequence_number"] == events[-2]["sequence_number"] + 1 + assert events[-1]["response"]["id"] == events[0]["response"]["id"] + else: + assert [event["type"] for event in events] == ["response.failed"] + assert events[0]["sequence_number"] == 0 + assert events[0]["response"]["id"].startswith("resp_") assert events[-1]["response"]["status"] == "failed" assert events[-1]["response"]["error"]["code"] == ( "rate_limit_exceeded" if error_kind in ("rate_limit", "numeric_rate_limit") else "server_error" From dc7895c1eacf55733953c3f802f46ea9103a76b4 Mon Sep 17 00:00:00 2001 From: zoroyihan7 Date: Tue, 8 Sep 2026 11:44:41 +0000 Subject: [PATCH 009/144] fix(responses): satisfy streaming regression checks --- litellm/proxy/common_utils/responses_stream_errors.py | 7 ++++--- .../test_response_polling_pre_call_checks.py | 1 - 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/litellm/proxy/common_utils/responses_stream_errors.py b/litellm/proxy/common_utils/responses_stream_errors.py index 356e948a2df..706b4c298d7 100644 --- a/litellm/proxy/common_utils/responses_stream_errors.py +++ b/litellm/proxy/common_utils/responses_stream_errors.py @@ -48,9 +48,10 @@ class _FailureDetails(BaseModel): def _original_failure(exception: Exception) -> Exception: - if isinstance(exception, MidStreamFallbackError) and exception.original_exception is not None: - return _original_failure(exception.original_exception) - return exception + current = exception # rebind-ok: the recursion gate requires iterative wrapper traversal + while isinstance(current, MidStreamFallbackError) and current.original_exception is not None: + current = current.original_exception + return current def _response_error_code(details: _FailureDetails) -> str: diff --git a/tests/proxy_unit_tests/test_response_polling_pre_call_checks.py b/tests/proxy_unit_tests/test_response_polling_pre_call_checks.py index 459834d0fd2..3fabdcefe5a 100644 --- a/tests/proxy_unit_tests/test_response_polling_pre_call_checks.py +++ b/tests/proxy_unit_tests/test_response_polling_pre_call_checks.py @@ -130,7 +130,6 @@ class TestPollingEndpointPreCallGuard: "litellm.proxy.proxy_server.proxy_config": MagicMock(), "litellm.proxy.proxy_server.proxy_logging_obj": AsyncMock(), "litellm.proxy.proxy_server.redis_usage_cache": AsyncMock(), - "litellm.proxy.proxy_server.select_data_generator": None, "litellm.proxy.proxy_server.user_api_base": None, "litellm.proxy.proxy_server.user_max_tokens": None, "litellm.proxy.proxy_server.user_model": None, From 1a749d84bdd66706bb41cafdba28e7a8b6a20fa9 Mon Sep 17 00:00:00 2001 From: ryan Date: Wed, 16 Sep 2026 01:56:00 +0000 Subject: [PATCH 010/144] fix(proxy): track project spend and enforce project budgets additively Project-scoped keys never wrote spend to LiteLLM_ProjectTable, so /project/info stayed at 0 and project budgets could not block. Wire the PROJECT entity through the spend queue, redis buffer, and db writer, reserve and increment a spend:project counter, reseed it from the project row, reset project spend in the budget cascade, and read the live counter in the project max budget check. Team member budgets keep gating project-scoped keys alongside the project budget Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/_types.py | 1 + litellm/proxy/auth/auth_checks.py | 34 ++-- .../proxy/common_utils/reset_budget_job.py | 24 +++ .../proxy/common_utils/user_api_key_cache.py | 10 ++ litellm/proxy/db/db_spend_update_writer.py | 79 ++++++++- .../redis_update_buffer.py | 7 + .../spend_update_queue.py | 4 + litellm/proxy/db/spend_counter_reseed.py | 5 + .../proxy/hooks/proxy_track_cost_callback.py | 6 + litellm/proxy/proxy_server.py | 30 ++++ .../spend_tracking/budget_reservation.py | 41 +++++ .../spend_tracking/spend_counter_batch.py | 11 +- litellm/repositories/prisma_protocols.py | 3 + litellm/repositories/unit_of_work.py | 2 + .../proxy/auth/test_auth_checks.py | 61 +++++++ .../common_utils/test_reset_budget_job.py | 30 +++- .../proxy/db/test_db_spend_update_writer.py | 76 +++++++++ .../proxy/db/test_spend_counter_reseed.py | 19 +++ .../hooks/test_proxy_track_cost_callback.py | 1 + .../test_spend_tracking_utils.py | 1 + .../proxy/test_budget_reservation.py | 156 ++++++++++++++++++ .../repositories/test_unit_of_work.py | 3 + 22 files changed, 583 insertions(+), 21 deletions(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 321f8190f13..228a91ad446 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -5261,6 +5261,7 @@ class DBSpendUpdateTransactions(TypedDict): team_member_list_transactions: dict[str, float] | None org_list_transactions: dict[str, float] | None org_member_list_transactions: ReadOnly[dict[str, float] | None] + project_list_transactions: ReadOnly[dict[str, float] | None] tag_list_transactions: dict[str, float] | None agent_list_transactions: dict[str, float] | None model_access_group_list_transactions: ReadOnly[dict[str, float] | None] diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index e3783c94dc7..ef17913d9ec 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -92,6 +92,8 @@ from litellm.proxy.common_utils.user_api_key_cache import ( model_access_group_registry_cache_key, model_access_group_spend_counter_key, object_permission_cache_key, + project_cache_key, + project_spend_counter_key, tag_cache_key, tag_registry_cache_key, team_membership_auth_cache_key, @@ -5586,16 +5588,22 @@ async def _project_max_budget_check( if project_object.litellm_budget_table is not None: max_budget = project_object.litellm_budget_table.max_budget - if ( - max_budget is not None - and project_object.spend is not None - and math.isfinite(max_budget) - and project_object.spend > max_budget - ): + if max_budget is None or not math.isfinite(max_budget): + return + + from litellm.proxy.proxy_server import get_current_spend + + project_spend: Final = await get_current_spend( + counter_key=project_spend_counter_key(project_object.project_id), + fallback_spend=project_object.spend or 0.0, + max_budget=max_budget, + ) + + if project_spend >= max_budget: if valid_token: call_info: Final = CallInfo( token=valid_token.token, - spend=project_object.spend, + spend=project_spend, max_budget=max_budget, user_id=valid_token.user_id, team_id=valid_token.team_id, @@ -5611,9 +5619,9 @@ async def _project_max_budget_check( ) raise litellm.BudgetExceededError( - current_cost=project_object.spend, + current_cost=project_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}", + message=f"Budget has been exceeded! Project={project_object.project_id} Current cost: {project_spend}, Max budget: {max_budget}", entity_type=Litellm_EntityType.PROJECT.value, entity_id=project_object.project_id, ) @@ -5663,10 +5671,6 @@ async def _project_soft_budget_check( ) -def _project_cache_key(project_id: str) -> str: - return f"project_id:{project_id}" - - async def get_project_object( project_id: str, prisma_client: PrismaClient | None, @@ -5684,7 +5688,7 @@ async def get_project_object( return None # Check cache first - cache_key: Final = _project_cache_key(project_id) + cache_key: Final = project_cache_key(project_id) deserialized_project: Final = await user_api_key_cache.async_get_cache( key=cache_key, model_type=LiteLLM_ProjectTableCachedObj, @@ -5726,7 +5730,7 @@ async def delete_cached_project_object( from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import evict_and_broadcast await evict_and_broadcast( - cache_keys=(_project_cache_key(project_id),), + cache_keys=(project_cache_key(project_id),), user_api_key_cache=user_api_key_cache, ) diff --git a/litellm/proxy/common_utils/reset_budget_job.py b/litellm/proxy/common_utils/reset_budget_job.py index acb51e73daf..2baefa89943 100644 --- a/litellm/proxy/common_utils/reset_budget_job.py +++ b/litellm/proxy/common_utils/reset_budget_job.py @@ -41,6 +41,8 @@ from litellm.proxy.common_utils.user_api_key_cache import ( end_user_cache_key, model_access_group_cache_key, model_access_group_spend_counter_key, + project_cache_key, + project_spend_counter_key, tag_cache_key, ) from litellm.proxy.db.budget_window_spend_writer import roll_window_spend_row @@ -49,6 +51,7 @@ from litellm.proxy.db.exception_handler import call_with_db_reconnect_retry from litellm.proxy.utils import PrismaClient, ProxyLogging from litellm.repositories.organization_repository import OrganizationRepository from litellm.repositories.prisma_protocols import SpendLinkedTable +from litellm.repositories.project_repository import ProjectRepository from litellm.repositories.table_repositories import ( EndUserRepository, ModelAccessGroupBudgetRepository, @@ -115,6 +118,11 @@ class _ModelAccessGroupRow(_BudgetLinkedRow, Protocol): def access_group_name(self) -> str: ... +class _ProjectRow(_BudgetLinkedRow, Protocol): + @property + def project_id(self) -> str: ... + + class _EndUserRow(_BudgetLinkedRow, Protocol): @property def user_id(self) -> str: ... @@ -185,6 +193,14 @@ def _model_access_group_cache_keys(row: _ModelAccessGroupRow) -> tuple[str, ...] return (model_access_group_cache_key(row.access_group_name),) +def _project_counter_key(row: _ProjectRow) -> str: + return project_spend_counter_key(row.project_id) + + +def _project_cache_keys(row: _ProjectRow) -> tuple[str, ...]: + return (project_cache_key(row.project_id),) + + def _enduser_counter_key(row: _EndUserRow) -> str: return f"spend:end_user:{row.user_id}" @@ -661,6 +677,11 @@ class ResetBudgetJob: where=_budget_link_where(budget_ids, _SPENT_ROWS_WHERE), log_subject="model access groups", ) + projects: Final[tuple[_ProjectRow, ...]] = await self._fetch_linked_rows( + table=ProjectRepository(self.prisma_client).table, + where=_budget_link_where(budget_ids, _SPENT_ROWS_WHERE), + log_subject="projects", + ) rollover_caps: Final[Mapping[str, float]] = MappingProxyType( { # mutable-ok: MappingProxyType wraps a one-shot dict comprehension b.budget_id: cap @@ -695,6 +716,7 @@ class ResetBudgetJob: (_model_access_group_counter_key(row), _row_carried_spend(row, rollover_caps)) for row in model_access_groups ), + *((_project_counter_key(row), _row_carried_spend(row, rollover_caps)) for row in projects), *((_enduser_counter_key(row), _enduser_carried_spend(row, rollover_caps)) for row in endusers), ), rollover_caps=rollover_caps, @@ -704,6 +726,7 @@ class ResetBudgetJob: *(key for row in orgs for key in _org_cache_keys(row)), *(key for row in tags for key in _tag_cache_keys(row)), *(key for row in model_access_groups for key in _model_access_group_cache_keys(row)), + *(key for row in projects for key in _project_cache_keys(row)), *(key for row in endusers for key in _enduser_cache_keys(row)), ), ) @@ -731,6 +754,7 @@ class ResetBudgetJob: _queue_budget_linked_resets(uow.organizations, cascade, extra=_SPENT_ROWS_WHERE) _queue_budget_linked_resets(uow.tags, cascade, extra=_SPENT_ROWS_WHERE) _queue_budget_linked_resets(uow.model_access_groups, cascade, extra=_SPENT_ROWS_WHERE) + _queue_budget_linked_resets(uow.projects, cascade, extra=_SPENT_ROWS_WHERE) _queue_enduser_resets(uow.endusers, cascade) for budget_id, budget_reset_at in cascade.budget_resets: uow.budgets.queue_window_advance(budget_id=budget_id, budget_reset_at=budget_reset_at) diff --git a/litellm/proxy/common_utils/user_api_key_cache.py b/litellm/proxy/common_utils/user_api_key_cache.py index 1c7a379897f..2187ed63ea5 100644 --- a/litellm/proxy/common_utils/user_api_key_cache.py +++ b/litellm/proxy/common_utils/user_api_key_cache.py @@ -306,6 +306,16 @@ def model_access_group_spend_counter_key(access_group_name: str) -> str: return f"spend:model_access_group:{access_group_name}" +def project_cache_key(project_id: str) -> str: + """Cache key one project row is stored under; shared by auth, spend tracking and the spend writer.""" + return f"project_id:{project_id}" + + +def project_spend_counter_key(project_id: str) -> str: + """Spend counter key for one project; the reservation, cost callback, auth and reseed paths all read it.""" + return f"spend:project:{project_id}" + + #: Cached under ``end_user_restricted_registry_cache_key`` when the restricted set exceeds #: ``END_USER_RESTRICTED_REGISTRY_MAX_SIZE``: registry unusable, fall back to the per-id fetch. END_USER_RESTRICTED_REGISTRY_OVERFLOW_SENTINEL: Final = "__end_user_restricted_registry_overflow__" diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index a90d1351fd7..51d00b9789c 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -44,6 +44,7 @@ from litellm.proxy._types import ( SpendUpdateQueueItem, ToolDiscoveryQueueItem, ) +from litellm.proxy.common_utils.user_api_key_cache import project_cache_key from litellm.proxy.db.daily_spend_bulk_upsert import ( DAILY_SPEND_TABLES, build_bulk_upsert, @@ -116,6 +117,7 @@ class _SpendBatch(Protocol): litellm_teammembership: BatchTable litellm_organizationtable: BatchTable litellm_organizationmembership: BatchTable + litellm_projecttable: BatchTable litellm_tagtable: BatchTable litellm_agentstable: BatchTable litellm_modelaccessgroupbudgettable: BatchTable @@ -254,6 +256,7 @@ class DBSpendUpdateWriter: start_time: datetime | None, end_time: datetime | None, response_cost: float | None, + project_id: str | None = None, ) -> bool: """Record the request's spend, answering whether its cost still needs charging. @@ -335,6 +338,7 @@ class DBSpendUpdateWriter: hashed_token=hashed_token, team_id=team_id, org_id=org_id, + project_id=project_id, end_user_id=end_user_id, prisma_client=prisma_client, litellm_proxy_budget_name=litellm_proxy_budget_name, @@ -631,6 +635,7 @@ class DBSpendUpdateWriter: litellm_proxy_budget_name: str | None, payload: SpendLogsPayload, request_model_access_groups: Sequence[str] = (), + project_id: str | None = None, ): """ Runs all 13 spend-update helpers sequentially inside a single asyncio task. @@ -694,6 +699,18 @@ class DBSpendUpdateWriter: traceback.format_exc(), ) + try: + await self._update_project_db( + response_cost=response_cost, + project_id=project_id, + prisma_client=prisma_client, + ) + except Exception: + verbose_proxy_logger.debug( + "_batch_database_updates: _update_project_db failed: %s", + traceback.format_exc(), + ) + try: await self._update_tag_db( response_cost=response_cost, @@ -956,6 +973,33 @@ class DBSpendUpdateWriter: ) raise e + async def _update_project_db( + self, + response_cost: float | None, + project_id: str | None, + prisma_client: PrismaClient | None, + ): + try: + if project_id is None or prisma_client is None: + return + + await self.spend_update_queue.add_update( + update=SpendUpdateQueueItem( + entity_type=Litellm_EntityType.PROJECT, + entity_id=project_id, + response_cost=response_cost, + ) + ) + except Exception as e: + spend_log_error( + "Spend tracking - failed to enqueue project spend update. project_id=%s, response_cost=%s - %s", + project_id, + response_cost, + str(e), + exc=e, + ) + raise e + async def _update_agent_db( self, response_cost: float | None, @@ -1193,8 +1237,8 @@ class DBSpendUpdateWriter: if db_spend_update_transactions is not None: verbose_proxy_logger.info( "Spend tracking - committing spend updates from Redis to DB: " - "keys=%d, users=%d, teams=%d, orgs=%d, end_users=%d, team_members=%d, org_members=%d, tags=%d, " - "agents=%d, model_access_groups=%d", + "keys=%d, users=%d, teams=%d, orgs=%d, end_users=%d, team_members=%d, org_members=%d, " + "projects=%d, tags=%d, agents=%d, model_access_groups=%d", len(db_spend_update_transactions.get("key_list_transactions") or {}), len(db_spend_update_transactions.get("user_list_transactions") or {}), len(db_spend_update_transactions.get("team_list_transactions") or {}), @@ -1202,6 +1246,7 @@ class DBSpendUpdateWriter: len(db_spend_update_transactions.get("end_user_list_transactions") or {}), len(db_spend_update_transactions.get("team_member_list_transactions") or {}), len(db_spend_update_transactions.get("org_member_list_transactions") or {}), + len(db_spend_update_transactions.get("project_list_transactions") or {}), len(db_spend_update_transactions.get("tag_list_transactions") or {}), len(db_spend_update_transactions.get("agent_list_transactions") or {}), len(db_spend_update_transactions.get("model_access_group_list_transactions") or {}), @@ -1762,6 +1807,22 @@ class DBSpendUpdateWriter: proxy_logging_obj=proxy_logging_obj, ) + ### UPDATE PROJECT TABLE ### + project_list_transactions: Final = db_spend_update_transactions.get("project_list_transactions") + await DBSpendUpdateWriter._update_entity_spend_in_db( + entity_name="Project", + transactions=project_list_transactions, + table_accessor="litellm_projecttable", + where_field="project_id", + n_retry_times=n_retry_times, + prisma_client=prisma_client, + proxy_logging_obj=proxy_logging_obj, + ) + await DBSpendUpdateWriter._invalidate_project_caches( + project_ids=tuple(project_list_transactions or ()), + proxy_logging_obj=proxy_logging_obj, + ) + ### UPDATE TAG TABLE ### tag_list_transactions: Final = db_spend_update_transactions["tag_list_transactions"] await DBSpendUpdateWriter._update_entity_spend_in_db( @@ -1800,11 +1861,23 @@ class DBSpendUpdateWriter: proxy_logging_obj=proxy_logging_obj, ) + @staticmethod + async def _invalidate_project_caches(project_ids: Sequence[str], proxy_logging_obj: ProxyLogging | None) -> None: + if not project_ids or proxy_logging_obj is None: + return + user_api_key_cache: Final = proxy_logging_obj.call_details.get("user_api_key_cache") + if user_api_key_cache is None: + return + for project_id in project_ids: + await user_api_key_cache.async_delete_cache(key=project_cache_key(project_id)) + @staticmethod async def _update_entity_spend_in_db( entity_name: str, transactions: dict[str, float] | None, - table_accessor: Literal["litellm_tagtable", "litellm_agentstable", "litellm_modelaccessgroupbudgettable"], + table_accessor: Literal[ + "litellm_tagtable", "litellm_agentstable", "litellm_modelaccessgroupbudgettable", "litellm_projecttable" + ], where_field: str, n_retry_times: int, prisma_client: PrismaClient, diff --git a/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py b/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py index 6f49a00b763..cead63795a2 100644 --- a/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py +++ b/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py @@ -70,6 +70,7 @@ _SpendTransactionField: TypeAlias = Literal[ "team_member_list_transactions", "org_list_transactions", "org_member_list_transactions", + "project_list_transactions", "tag_list_transactions", "agent_list_transactions", "model_access_group_list_transactions", @@ -83,6 +84,7 @@ _SPEND_TRANSACTION_FIELDS: Final[tuple[_SpendTransactionField, ...]] = ( "team_member_list_transactions", "org_list_transactions", "org_member_list_transactions", + "project_list_transactions", "tag_list_transactions", "agent_list_transactions", "model_access_group_list_transactions", @@ -418,6 +420,10 @@ class RedisUpdateBuffer: Litellm_EntityType.ORGANIZATION_MEMBER, db_spend_update_transactions.get("org_member_list_transactions"), ), + ( + Litellm_EntityType.PROJECT, + db_spend_update_transactions.get("project_list_transactions"), + ), ( Litellm_EntityType.TAG, db_spend_update_transactions.get("tag_list_transactions"), @@ -885,6 +891,7 @@ class RedisUpdateBuffer: org_member_list_transactions=_merged_entity_transactions( list_of_transactions, "org_member_list_transactions" ), + project_list_transactions=_merged_entity_transactions(list_of_transactions, "project_list_transactions"), tag_list_transactions=_merged_entity_transactions(list_of_transactions, "tag_list_transactions"), agent_list_transactions=_merged_entity_transactions(list_of_transactions, "agent_list_transactions"), model_access_group_list_transactions=_merged_entity_transactions( diff --git a/litellm/proxy/db/db_transaction_queue/spend_update_queue.py b/litellm/proxy/db/db_transaction_queue/spend_update_queue.py index bc068d10daf..2b8535cb113 100644 --- a/litellm/proxy/db/db_transaction_queue/spend_update_queue.py +++ b/litellm/proxy/db/db_transaction_queue/spend_update_queue.py @@ -138,6 +138,7 @@ class SpendUpdateQueue(BaseUpdateQueue): team_member_list_transactions={}, org_list_transactions={}, org_member_list_transactions={}, + project_list_transactions={}, tag_list_transactions={}, agent_list_transactions={}, model_access_group_list_transactions={}, @@ -152,6 +153,7 @@ class SpendUpdateQueue(BaseUpdateQueue): Litellm_EntityType.TEAM_MEMBER: "team_member_list_transactions", Litellm_EntityType.ORGANIZATION: "org_list_transactions", Litellm_EntityType.ORGANIZATION_MEMBER: "org_member_list_transactions", + Litellm_EntityType.PROJECT: "project_list_transactions", Litellm_EntityType.TAG: "tag_list_transactions", Litellm_EntityType.AGENT: "agent_list_transactions", Litellm_EntityType.MODEL_ACCESS_GROUP: "model_access_group_list_transactions", @@ -192,6 +194,8 @@ class SpendUpdateQueue(BaseUpdateQueue): transactions_dict = db_spend_update_transactions["org_list_transactions"] elif dict_key == "org_member_list_transactions": transactions_dict = db_spend_update_transactions["org_member_list_transactions"] + elif dict_key == "project_list_transactions": + transactions_dict = db_spend_update_transactions["project_list_transactions"] elif dict_key == "tag_list_transactions": transactions_dict = db_spend_update_transactions["tag_list_transactions"] elif dict_key == "agent_list_transactions": diff --git a/litellm/proxy/db/spend_counter_reseed.py b/litellm/proxy/db/spend_counter_reseed.py index 89a07234c6c..2dd028454d6 100644 --- a/litellm/proxy/db/spend_counter_reseed.py +++ b/litellm/proxy/db/spend_counter_reseed.py @@ -26,6 +26,7 @@ from litellm.proxy._types import Litellm_EntityType from litellm.proxy.db.db_lookup_gate import db_lookup_gate from litellm.proxy.spend_tracking.spend_counter_batch import read_batched_spend_counter, record_spend_counter_value from litellm.repositories.organization_repository import OrganizationRepository +from litellm.repositories.project_repository import ProjectRepository from litellm.repositories.table_repositories import ( BudgetWindowSpendRepository, EndUserRepository, @@ -77,6 +78,7 @@ class SpendCounterReseed: spend:team_member:{uid}:{tid} -> LiteLLM_TeamMembership.spend spend:user:{user_id} -> LiteLLM_UserTable.spend spend:org:{org_id} -> LiteLLM_OrganizationTable.spend + spend:project:{project_id} -> LiteLLM_ProjectTable.spend End-user and tag spend counters intentionally do not reseed here. Their auth paths already load the corresponding objects via get_end_user_object() @@ -157,6 +159,9 @@ class SpendCounterReseed: row = await OrganizationRepository(prisma_client).table.find_unique( where={"organization_id": org_id} ) + elif counter_key.startswith("spend:project:"): + project_id: Final = counter_key[len("spend:project:") :] + row = await ProjectRepository(prisma_client).table.find_unique(where={"project_id": project_id}) else: return None except Exception: diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index 1ae106be390..0c562cf37ef 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -267,6 +267,7 @@ class _ProxyDBLogger(CustomLogger): start_time=actual_start_time, end_time=datetime.now(), org_id=user_api_key_dict.org_id, + project_id=user_api_key_dict.project_id, ) @log_db_metrics @@ -318,6 +319,7 @@ class _ProxyDBLogger(CustomLogger): user_id: Final = cast(str | None, metadata.get("user_api_key_user_id", None)) team_id: Final = cast(str | None, metadata.get("user_api_key_team_id", None)) org_id: Final = cast(str | None, metadata.get("user_api_key_org_id", None)) + project_id: Final = cast(str | None, metadata.get("user_api_key_project_id", None)) key_alias: Final = cast(str | None, metadata.get("user_api_key_alias", None)) end_user_max_budget: Final = metadata.get("user_api_end_user_max_budget", None) sl_object: Final[StandardLoggingPayload | None] = kwargs.get("standard_logging_object", None) @@ -368,6 +370,7 @@ class _ProxyDBLogger(CustomLogger): budget_reservation=budget_reservation, request_tags=tags, model_access_groups=model_access_groups, + project_id=project_id, ) if not charged: return @@ -651,6 +654,7 @@ async def _update_database_and_spend_counters( budget_reservation: dict | None, request_tags: list[str] | None = None, model_access_groups: Sequence[str] | None = None, + project_id: str | None = None, ) -> bool: if budget_reservation is not None: await _reconcile_budget_reservation_before_db_update( @@ -668,6 +672,7 @@ async def _update_database_and_spend_counters( start_time=start_time, end_time=end_time, org_id=org_id, + project_id=project_id, ) except Exception: if budget_reservation is not None: @@ -698,6 +703,7 @@ async def _update_database_and_spend_counters( tags=request_tags, request_started_at=start_time, model_access_groups=model_access_groups, + project_id=project_id, ) except Exception: if budget_reservation is not None: diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index d7964556531..1e84e5f56e2 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -417,6 +417,8 @@ from litellm.proxy.common_utils.user_api_key_cache import ( get_management_object_ttl, model_access_group_cache_key, model_access_group_spend_counter_key, + project_cache_key, + project_spend_counter_key, tag_cache_key, ) from litellm.proxy.config_resolvers import resolve_fields @@ -2780,6 +2782,7 @@ async def increment_spend_counters( tags: list[str] | None = None, request_started_at: datetime | None = None, model_access_groups: Sequence[str] | None = None, + project_id: str | None = None, ): """ Atomically increment spend counters for budget enforcement. @@ -2801,6 +2804,7 @@ async def increment_spend_counters( end_user_id=end_user_id, tags=tags, model_access_groups=model_access_groups, + project_id=project_id, ), ): await _increment_spend_counters_batched( @@ -2814,6 +2818,7 @@ async def increment_spend_counters( tags=tags, request_started_at=request_started_at, model_access_groups=model_access_groups, + project_id=project_id, ) @@ -2828,6 +2833,7 @@ async def _increment_spend_counters_batched( tags: list[str] | None, request_started_at: datetime | None, model_access_groups: Sequence[str] | None, + project_id: str | None = None, ): """Runs inside one spend counter batch: the reservation reconcile and the warm checks share a single MGET.""" reserved_counter_keys: Final = await _reconcile_budget_reservation_for_counter_update( @@ -3028,6 +3034,13 @@ async def _increment_spend_counters_batched( ) if org_id is not None else None, + _prepare_project_spend_increment( + project_id=project_id, + response_cost=cost, + reserved_counter_keys=reserved_counter_keys, + ) + if project_id is not None + else None, ) if coro is not None ) @@ -3180,6 +3193,23 @@ async def _prepare_org_spend_increment( return (pending,) if pending is not None else () +async def _prepare_project_spend_increment( + project_id: str | None, + response_cost: float, + reserved_counter_keys: set[str], +) -> tuple[PendingSpendIncrement, ...]: + if project_id is None: + return () + + pending: Final = await _prepare_unreserved_spend_counter_increment( + counter_key=project_spend_counter_key(project_id), + source_cache_key=project_cache_key(project_id), + increment=response_cost, + reserved_counter_keys=reserved_counter_keys, + ) + return (pending,) if pending is not None else () + + async def _prepare_unreserved_spend_counter_increment( counter_key: str, source_cache_key: str | list[str], diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index 373f2d0fe36..24b8470145c 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -2,6 +2,7 @@ from __future__ import annotations import asyncio import json +import math from collections.abc import Mapping, Sequence from dataclasses import dataclass from datetime import datetime, timedelta, timezone @@ -29,6 +30,8 @@ from litellm.proxy.common_utils.user_api_key_cache import ( end_user_cache_key, model_access_group_cache_key, model_access_group_spend_counter_key, + project_cache_key, + project_spend_counter_key, tag_cache_key, team_membership_reservation_cache_key, ) @@ -62,6 +65,7 @@ _COUNTER_ENTITY_TYPES: Final[Mapping[str, str]] = { "Tag": Litellm_EntityType.TAG.value, "Model access group": Litellm_EntityType.MODEL_ACCESS_GROUP.value, "Organization": Litellm_EntityType.ORGANIZATION.value, + "Project": Litellm_EntityType.PROJECT.value, } @@ -542,6 +546,13 @@ async def _get_budget_counters( if org_counter is not None: counters.append(org_counter) + project_counter: Final = await _get_project_budget_counter( + valid_token=valid_token, + user_api_key_cache=user_api_key_cache, + ) + if project_counter is not None: + counters.append(project_counter) + return counters @@ -751,6 +762,36 @@ async def _get_org_budget_counter( ) +async def _get_project_budget_counter( + valid_token: UserAPIKeyAuth, + user_api_key_cache: UserApiKeyCache, +) -> _BudgetCounter | None: + if valid_token.project_id is None: + return None + + source_cache_key: Final = project_cache_key(valid_token.project_id) + project_object: Final = await user_api_key_cache.async_get_cache(key=source_cache_key) + if project_object is None: + return None + + project_budget_table: Final = _get_value(project_object, "litellm_budget_table") + if project_budget_table is None: + return None + + project_max_budget: Final = _to_float(_get_value(project_budget_table, "max_budget")) + if project_max_budget is None or project_max_budget <= 0 or not math.isfinite(project_max_budget): + return None + + return _BudgetCounter( + counter_key=project_spend_counter_key(valid_token.project_id), + source_cache_key=source_cache_key, + max_budget=project_max_budget, + fallback_spend=_to_float(_get_value(project_object, "spend")) or 0.0, + entity_type="Project", + entity_id=valid_token.project_id, + ) + + def _get_budget_limit_counters( entity_prefix: str, entity_type: str, diff --git a/litellm/proxy/spend_tracking/spend_counter_batch.py b/litellm/proxy/spend_tracking/spend_counter_batch.py index 7106d88c655..ddb074ae023 100644 --- a/litellm/proxy/spend_tracking/spend_counter_batch.py +++ b/litellm/proxy/spend_tracking/spend_counter_batch.py @@ -12,7 +12,10 @@ from pydantic import TypeAdapter from litellm._logging import verbose_proxy_logger from litellm.caching.redis_cache import RedisCache from litellm.proxy._types import UserAPIKeyAuth -from litellm.proxy.common_utils.user_api_key_cache import model_access_group_spend_counter_key +from litellm.proxy.common_utils.user_api_key_cache import ( + model_access_group_spend_counter_key, + project_spend_counter_key, +) _CounterValues: Final = TypeAdapter(dict[str, float | None]) _NO_VALUES: Final[Mapping[str, float | None]] = MappingProxyType({}) @@ -154,6 +157,8 @@ def _iter_admission_counter_keys(token: UserAPIKeyAuth, end_user_id: str | None) yield f"spend:end_user:{end_user_id}" if token.org_id is not None: yield f"spend:org:{token.org_id}" + if token.project_id is not None: + yield project_spend_counter_key(token.project_id) def admission_counter_keys(token: UserAPIKeyAuth, end_user_id: str | None) -> frozenset[str]: @@ -168,10 +173,12 @@ def post_call_counter_keys( end_user_id: str | None, tags: Sequence[object] | None, model_access_groups: Sequence[object] | None, + project_id: str | None = None, ) -> frozenset[str]: """Every counter ``increment_spend_counters`` warm-checks, except budget windows which bind on read.""" entity_keys: Final = admission_counter_keys( - UserAPIKeyAuth(token=token, team_id=team_id, user_id=user_id, org_id=org_id), end_user_id + UserAPIKeyAuth(token=token, team_id=team_id, user_id=user_id, org_id=org_id, project_id=project_id), + end_user_id, ) tag_keys: Final = frozenset(f"spend:tag:{tag}" for tag in tags or () if tag and isinstance(tag, str)) group_keys: Final = frozenset( diff --git a/litellm/repositories/prisma_protocols.py b/litellm/repositories/prisma_protocols.py index 93b8c5c7cd7..60c16fbd746 100644 --- a/litellm/repositories/prisma_protocols.py +++ b/litellm/repositories/prisma_protocols.py @@ -152,4 +152,7 @@ class PrismaBatch(Protocol): @property def litellm_modelaccessgroupbudgettable(self) -> BatchTable: ... + @property + def litellm_projecttable(self) -> BatchTable: ... + async def commit(self) -> None: ... diff --git a/litellm/repositories/unit_of_work.py b/litellm/repositories/unit_of_work.py index 0cdce307f9b..c09e5eb75d4 100644 --- a/litellm/repositories/unit_of_work.py +++ b/litellm/repositories/unit_of_work.py @@ -109,6 +109,7 @@ class BudgetCascadeUnitOfWork: organizations: LinkedSpendResetWrites tags: LinkedSpendResetWrites model_access_groups: LinkedSpendResetWrites + projects: LinkedSpendResetWrites endusers: LinkedSpendResetWrites budgets: BudgetWindowWrites @@ -135,6 +136,7 @@ async def budget_cascade_unit_of_work( organizations=LinkedSpendResetWrites(table=batch.litellm_organizationtable), tags=LinkedSpendResetWrites(table=batch.litellm_tagtable), model_access_groups=LinkedSpendResetWrites(table=batch.litellm_modelaccessgroupbudgettable), + projects=LinkedSpendResetWrites(table=batch.litellm_projecttable), endusers=LinkedSpendResetWrites(table=batch.litellm_endusertable), budgets=BudgetWindowWrites(table=batch.litellm_budgettable), ) diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index 8c8b755195f..0d5c0dd5d72 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -7138,6 +7138,67 @@ async def test_project_allowlist_enforced_when_key_models_empty(): assert exc_info.value.code == "403" +def _project_with_budget(spend: float, max_budget: float): + from litellm.proxy._types import LiteLLM_BudgetTable, LiteLLM_ProjectTableCachedObj + + return LiteLLM_ProjectTableCachedObj( + project_id="p-budget", + team_id="t-1", + budget_id="b-1", + spend=spend, + litellm_budget_table=LiteLLM_BudgetTable(budget_id="b-1", max_budget=max_budget), + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "counter_spend, db_spend, blocks", + [ + pytest.param(5.0, 0.0, True, id="counter-at-budget-blocks-despite-stale-db-row"), + pytest.param(4.99, 0.0, False, id="counter-under-budget-admits"), + pytest.param(None, 5.0, True, id="no-counter-falls-back-to-persisted-spend"), + pytest.param(None, 0.0, False, id="no-counter-and-no-persisted-spend-admits"), + ], +) +async def test_project_max_budget_check_reads_live_spend_counter(counter_spend, db_spend, blocks): + """LIT-3269: project budget enforcement must read the cross-pod + ``spend:project:{id}`` counter first and only fall back to the cached row's + spend, matching key/team/org checks. The boundary is inclusive (>=).""" + from litellm.caching.dual_cache import DualCache + from litellm.proxy.auth.auth_checks import _project_max_budget_check + + real_spend_counter_cache = DualCache() + if counter_spend is not None: + real_spend_counter_cache.in_memory_cache.set_cache(key="spend:project:p-budget", value=counter_spend) + valid_token = UserAPIKeyAuth(api_key="hashed-key", project_id="p-budget", team_id="t-1", user_id="u-1") + proxy_logging_obj = MagicMock() + proxy_logging_obj.budget_alerts = AsyncMock() + + with patch( # test-quality-ok: injects a real DualCache for the module global, not a behavior mock + "litellm.proxy.proxy_server.spend_counter_cache", real_spend_counter_cache + ): + if not blocks: + await _project_max_budget_check( + project_object=_project_with_budget(spend=db_spend, max_budget=5.0), + valid_token=valid_token, + proxy_logging_obj=proxy_logging_obj, + ) + return + with pytest.raises(litellm.BudgetExceededError) as exc_info: + await _project_max_budget_check( + project_object=_project_with_budget(spend=db_spend, max_budget=5.0), + valid_token=valid_token, + proxy_logging_obj=proxy_logging_obj, + ) + await asyncio.sleep(0) + + assert exc_info.value.entity_type == Litellm_EntityType.PROJECT.value + assert exc_info.value.entity_id == "p-budget" + assert exc_info.value.current_cost == 5.0 + proxy_logging_obj.budget_alerts.assert_awaited_once() + assert proxy_logging_obj.budget_alerts.await_args.kwargs["type"] == "project_budget" + + def test_is_user_proxy_admin_rejects_view_only_admin(): """This predicate skips `non_proxy_admin_allowed_routes_check` entirely, so an Admin Viewer answering True here would gain every write route. Read parity for diff --git a/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py b/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py index 943a6c905c0..8c254686385 100644 --- a/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py +++ b/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py @@ -78,6 +78,7 @@ class MockBatcher: self.litellm_organizationtable = _Table("org", self) self.litellm_tagtable = _Table("tag", self) self.litellm_modelaccessgroupbudgettable = _Table("model_access_group", self) + self.litellm_projecttable = _Table("project", self) self.litellm_endusertable = _Table("enduser", self) async def commit(self): @@ -93,6 +94,7 @@ class MockDB: self.litellm_organizationtable = MockTable() self.litellm_tagtable = MockTable() self.litellm_modelaccessgroupbudgettable = MockTable() + self.litellm_projecttable = MockTable() self.batch_calls: List[Dict[str, Any]] = [] self.batchers: List[MockBatcher] = [] @@ -1507,13 +1509,19 @@ _INVALIDATION_CASES = [ "spend:model_access_group:gpt-4-group", {"model_access_group:gpt-4-group"}, ), + ( + "litellm_projecttable", + type("Project", (), {"project_id": "proj-1"}), + "spend:project:proj-1", + {"project_id:proj-1"}, + ), ] @pytest.mark.parametrize( "table_attr, linked_row, counter_key, cache_keys", _INVALIDATION_CASES, - ids=["team_membership", "key", "org", "tag", "model_access_group"], + ids=["team_membership", "key", "org", "tag", "model_access_group", "project"], ) def test_budget_table_reset_invalidates_counters_and_management_cache( reset_budget_job, mock_prisma_client, monkeypatch, table_attr, linked_row, counter_key, cache_keys @@ -1657,6 +1665,25 @@ def test_budget_table_reset_invalidates_every_access_group_not_just_the_first( counter_cache.in_memory_cache.delete_cache.assert_any_call(key=f"spend:model_access_group:{name}") +def test_project_reset_zeroes_spend_on_due_tiers(reset_budget_job, mock_prisma_client, monkeypatch): + """A project linked to an expiring budget tier has its spend zeroed in the same cascade transaction.""" + _make_counter_invalidation_job(monkeypatch) + mock_prisma_client.data["budget"] = [_budget_row(budget_id="budget-due", budget_duration="7d")] + mock_prisma_client.db.litellm_projecttable.set_find_many_results( + [type("Project", (), {"project_id": "proj-1", "spend": 12.0, "budget_id": "budget-due"})] + ) + + asyncio.run(reset_budget_job.reset_budget_for_litellm_budget_table()) + + expected_where = {"budget_id": {"in": ["budget-due"]}, "spend": {"gt": 0}} + assert mock_prisma_client.db.litellm_projecttable.find_many_calls == [{"where": expected_where}] + writes = _batch_writes(mock_prisma_client, "project", op="update_many") + assert len(writes) == 1 + assert writes[0]["where"] == expected_where + assert writes[0]["data"] == {"spend": 0} + assert mock_prisma_client.db.batchers[0].committed is True + + def test_budget_cascade_carries_access_group_overage_when_rollover_enabled( rollover_enabled, reset_budget_job, mock_prisma_client, monkeypatch ): @@ -1802,6 +1829,7 @@ def test_budget_cascade_writes_land_in_a_single_transaction(reset_budget_job, mo ("org", "update_many"), ("tag", "update_many"), ("model_access_group", "update_many"), + ("project", "update_many"), ("enduser", "update_many"), ("budget", "update_many"), } diff --git a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py index c547d06904b..da5879a375a 100644 --- a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py +++ b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py @@ -1060,6 +1060,82 @@ async def test_batch_database_updates_queues_org_member_spend_for_the_request_us assert transactions["org_member_list_transactions"] == {"organization_id::org1::user_id::u1": 0.1} +@pytest.mark.asyncio +async def test_project_spend_is_persisted_to_project_table_and_project_cache_is_evicted(): + """Regression for LIT-3269: a request made with a project-scoped key must + increment LiteLLM_ProjectTable.spend, otherwise /project/info stays at 0 + and the project budget never blocks. The cached project row is evicted so + the next auth check reads the fresh spend.""" + db_writer: Final = DBSpendUpdateWriter() + await db_writer._batch_database_updates( + response_cost=0.25, + user_id="u1", + hashed_token="t1", + team_id="team-1", + org_id=None, + end_user_id=None, + prisma_client=MagicMock(), + litellm_proxy_budget_name=None, + payload={"request_id": "req-1", "model": "gpt-4o-mini", "spend": 0.25}, + project_id="proj-1", + ) + await db_writer._batch_database_updates( + response_cost=0.5, + user_id="u1", + hashed_token="t1", + team_id="team-1", + org_id=None, + end_user_id=None, + prisma_client=MagicMock(), + litellm_proxy_budget_name=None, + payload={"request_id": "req-2", "model": "gpt-4o-mini", "spend": 0.5}, + project_id="proj-1", + ) + transactions: Final = await db_writer.spend_update_queue.flush_and_get_aggregated_db_spend_update_transactions() + assert transactions["project_list_transactions"] == {"proj-1": 0.75} + assert transactions["team_member_list_transactions"] == {"team_id::team-1::user_id::u1": 0.75} + + mock_batcher: Final = MagicMock() + mock_prisma_client: Final = MagicMock() + mock_prisma_client.db.tx = MagicMock(return_value=_good_tx(mock_batcher)) + user_api_key_cache: Final = MagicMock() + user_api_key_cache.async_delete_cache = AsyncMock() + proxy_logging: Final = MagicMock() + proxy_logging.call_details = {"user_api_key_cache": user_api_key_cache} + + await db_writer._commit_spend_updates_to_db( + prisma_client=mock_prisma_client, + n_retry_times=0, + proxy_logging_obj=proxy_logging, + db_spend_update_transactions=transactions, + ) + + mock_batcher.litellm_projecttable.update_many.assert_called_once_with( + where={"project_id": "proj-1"}, + data={"spend": {"increment": 0.75}}, + ) + user_api_key_cache.async_delete_cache.assert_any_await(key="project_id:proj-1") + + +@pytest.mark.asyncio +async def test_batch_database_updates_without_project_id_touches_no_project_row(): + db_writer: Final = DBSpendUpdateWriter() + await db_writer._batch_database_updates( + response_cost=0.1, + user_id="u1", + hashed_token="t1", + team_id=None, + org_id=None, + end_user_id=None, + prisma_client=MagicMock(), + litellm_proxy_budget_name=None, + payload={"request_id": "req-1", "model": "gpt-4o-mini", "spend": 0.1}, + ) + transactions: Final = await db_writer.spend_update_queue.flush_and_get_aggregated_db_spend_update_transactions() + + assert transactions["project_list_transactions"] == {} + + @pytest.mark.asyncio async def test_add_spend_log_transaction_to_daily_tag_transaction_with_request_id(): """ diff --git a/tests/test_litellm/proxy/db/test_spend_counter_reseed.py b/tests/test_litellm/proxy/db/test_spend_counter_reseed.py index bca6344b3f7..53e91b8792e 100644 --- a/tests/test_litellm/proxy/db/test_spend_counter_reseed.py +++ b/tests/test_litellm/proxy/db/test_spend_counter_reseed.py @@ -67,12 +67,14 @@ class _FakePrismaClient: error: Exception | None = None, end_user_row: SimpleNamespace | None = None, end_user_error: Exception | None = None, + project_row: SimpleNamespace | None = None, ) -> None: self.db = SimpleNamespace( litellm_budgetwindowspend=_FakeFindUniqueTable(row=row, error=error), litellm_spendlogs=_FakeSpendLogsTable(total=spend_logs_total), litellm_endusertable=_FakeFindUniqueTable(row=end_user_row, error=end_user_error), litellm_verificationtoken=_InFlightCountingTable(), + litellm_projecttable=_FakeFindUniqueTable(row=project_row), ) @@ -428,6 +430,23 @@ async def test_from_db_bounds_in_flight_prisma_requests_across_counter_keys(): assert prisma.db.litellm_verificationtoken.max_in_flight == PROXY_DB_LOOKUP_MAX_CONCURRENCY +@pytest.mark.asyncio +async def test_from_db_reseeds_project_counter_from_the_project_row(): + """LIT-3269: a cold ``spend:project:{id}`` counter seeds from LiteLLM_ProjectTable.spend, + so a fresh pod enforces the project budget against persisted spend rather than 0.""" + prisma: Final = _FakePrismaClient(project_row=SimpleNamespace(project_id="proj-1", spend=7.25)) + + assert await SpendCounterReseed.from_db(prisma_client=prisma, counter_key="spend:project:proj-1") == 7.25 + assert prisma.db.litellm_projecttable.where_clauses == [{"project_id": "proj-1"}] + + +@pytest.mark.asyncio +async def test_from_db_returns_none_for_a_missing_project_row(): + prisma: Final = _FakePrismaClient(project_row=None) + + assert await SpendCounterReseed.from_db(prisma_client=prisma, counter_key="spend:project:proj-1") is None + + @pytest.mark.asyncio async def test_from_db_still_never_reads_the_end_user_row(): """A cold end-user counter keeps seeding from the cached end-user object the auth diff --git a/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py b/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py index dfc95db3e14..965e134772d 100644 --- a/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py +++ b/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py @@ -586,6 +586,7 @@ async def test_update_database_and_spend_counters_updates_counters_after_db_upda tags=["tag-a"], request_started_at=start_time, model_access_groups=("premium",), + project_id=None, ) diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py b/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py index 8b105e94d19..834cf8d100d 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py @@ -1187,6 +1187,7 @@ async def test_api_key_preserved_through_failure_hook_to_database(): start_time, end_time, org_id, + project_id=None, ): """Mock update_database and capture the payload it creates""" from litellm.proxy.spend_tracking.spend_tracking_utils import ( diff --git a/tests/test_litellm/proxy/test_budget_reservation.py b/tests/test_litellm/proxy/test_budget_reservation.py index 032722d3259..014f240d9cc 100644 --- a/tests/test_litellm/proxy/test_budget_reservation.py +++ b/tests/test_litellm/proxy/test_budget_reservation.py @@ -22,6 +22,7 @@ from litellm.proxy._types import ( LiteLLM_EndUserTable, Litellm_EntityType, LiteLLM_OrganizationTable, + LiteLLM_ProjectTableCachedObj, LiteLLM_TagTable, LiteLLM_TeamMembership, LiteLLM_TeamTable, @@ -631,6 +632,161 @@ async def test_should_reserve_team_member_and_org_budget_counters(spend_counter_ await release_budget_reservation(reservation) +def _project_scoped_token() -> UserAPIKeyAuth: + return UserAPIKeyAuth( + token="key-project-scoped", + spend=0.0, + user_id="user-proj", + team_id="team-proj", + project_id="proj-1", + ) + + +async def _seed_project_scoped_budgets( + key_cache: DualCache, + team_member_spend: float, + team_member_max_budget: float, + project_spend: float, + project_max_budget: float, +) -> None: + await key_cache.async_set_cache( + key="team_membership:user-proj:team-proj", + value=LiteLLM_TeamMembership( + user_id="user-proj", + team_id="team-proj", + spend=team_member_spend, + litellm_budget_table=LiteLLM_BudgetTable(max_budget=team_member_max_budget), + ).model_dump(), + ) + await key_cache.async_set_cache( + key="project_id:proj-1", + value=LiteLLM_ProjectTableCachedObj( + project_id="proj-1", + team_id="team-proj", + budget_id="project-budget-id", + spend=project_spend, + litellm_budget_table=LiteLLM_BudgetTable(max_budget=project_max_budget), + ).model_dump(), + ) + + +@pytest.mark.asyncio +async def test_should_reserve_project_and_team_member_counters_for_project_scoped_key(spend_counter_state): + """LIT-3269: a key carrying user_id, team_id and project_id reserves against + both the team member counter and the project counter; neither replaces the + other. After the call the project counter reflects the real cost once, not + the reservation plus the post-call increment.""" + counter_cache, key_cache = spend_counter_state + proxy_logging_obj = ProxyLogging(user_api_key_cache=key_cache) + await _seed_project_scoped_budgets( + key_cache, + team_member_spend=0.1, + team_member_max_budget=1.0, + project_spend=0.2, + project_max_budget=1.0, + ) + + estimated = estimate_request_max_cost(request_body=_request_body(), route="/chat/completions", llm_router=None) + assert estimated is not None and estimated > 0 + + reservation = await reserve_budget_for_request( + request_body=_request_body(), + route="/chat/completions", + llm_router=None, + valid_token=_project_scoped_token(), + team_object=LiteLLM_TeamTable(team_id="team-proj", spend=0.0, max_budget=None), + user_object=LiteLLM_UserTable(user_id="user-proj", spend=0.0), + prisma_client=None, + user_api_key_cache=key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + + assert reservation is not None + assert counter_cache.in_memory_cache.get_cache(key="spend:team_member:user-proj:team-proj") == pytest.approx( + 0.1 + estimated + ) + assert counter_cache.in_memory_cache.get_cache(key="spend:project:proj-1") == pytest.approx(0.2 + estimated) + + from litellm.proxy.proxy_server import increment_spend_counters + + await increment_spend_counters( + token="key-project-scoped", + team_id="team-proj", + user_id="user-proj", + response_cost=0.05, + budget_reservation=reservation, + project_id="proj-1", + ) + + assert counter_cache.in_memory_cache.get_cache(key="spend:project:proj-1") == pytest.approx(0.25) + assert counter_cache.in_memory_cache.get_cache(key="spend:team_member:user-proj:team-proj") == pytest.approx(0.15) + + +@pytest.mark.asyncio +async def test_exhausted_team_member_budget_still_blocks_project_scoped_key(spend_counter_state): + """LIT-3269: the project budget is additive. A project with plenty of + headroom must not let a key through once its team member budget is spent.""" + counter_cache, key_cache = spend_counter_state + proxy_logging_obj = ProxyLogging(user_api_key_cache=key_cache) + await _seed_project_scoped_budgets( + key_cache, + team_member_spend=1.0, + team_member_max_budget=1.0, + project_spend=0.0, + project_max_budget=100.0, + ) + + with pytest.raises(litellm.BudgetExceededError) as exc_info: + await reserve_budget_for_request( + request_body=_request_body(), + route="/chat/completions", + llm_router=None, + valid_token=_project_scoped_token(), + team_object=LiteLLM_TeamTable(team_id="team-proj", spend=0.0, max_budget=None), + user_object=LiteLLM_UserTable(user_id="user-proj", spend=0.0), + prisma_client=None, + user_api_key_cache=key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + + assert "TeamMember=user-proj:team-proj" in str(exc_info.value) + assert counter_cache.in_memory_cache.get_cache(key="spend:project:proj-1") in (None, pytest.approx(0.0)) + + +@pytest.mark.asyncio +async def test_exhausted_project_budget_blocks_project_scoped_key(spend_counter_state): + """LIT-3269: with team member headroom left, the project budget alone blocks the key.""" + counter_cache, key_cache = spend_counter_state + proxy_logging_obj = ProxyLogging(user_api_key_cache=key_cache) + await _seed_project_scoped_budgets( + key_cache, + team_member_spend=0.0, + team_member_max_budget=100.0, + project_spend=5.0, + project_max_budget=5.0, + ) + + with pytest.raises(litellm.BudgetExceededError) as exc_info: + await reserve_budget_for_request( + request_body=_request_body(), + route="/chat/completions", + llm_router=None, + valid_token=_project_scoped_token(), + team_object=LiteLLM_TeamTable(team_id="team-proj", spend=0.0, max_budget=None), + user_object=LiteLLM_UserTable(user_id="user-proj", spend=0.0), + prisma_client=None, + user_api_key_cache=key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + + assert "Project=proj-1" in str(exc_info.value) + assert exc_info.value.entity_type == Litellm_EntityType.PROJECT.value + assert counter_cache.in_memory_cache.get_cache(key="spend:team_member:user-proj:team-proj") in ( + None, + pytest.approx(0.0), + ) + + @pytest.mark.asyncio async def test_should_not_reserve_user_budget_counter_for_team_key(spend_counter_state): """The reservation path mirrors the read path: no personal user counter for a team key. diff --git a/tests/test_litellm/repositories/test_unit_of_work.py b/tests/test_litellm/repositories/test_unit_of_work.py index b52b8ced31e..9eff248b917 100644 --- a/tests/test_litellm/repositories/test_unit_of_work.py +++ b/tests/test_litellm/repositories/test_unit_of_work.py @@ -35,6 +35,7 @@ class FakeBatch: self.litellm_organizationtable = FakeBatchTable("litellm_organizationtable", self.calls) self.litellm_tagtable = FakeBatchTable("litellm_tagtable", self.calls) self.litellm_modelaccessgroupbudgettable = FakeBatchTable("litellm_modelaccessgroupbudgettable", self.calls) + self.litellm_projecttable = FakeBatchTable("litellm_projecttable", self.calls) self.litellm_endusertable = FakeBatchTable("litellm_endusertable", self.calls) async def commit(self) -> None: @@ -94,6 +95,7 @@ async def test_budget_cascade_dependents_and_window_advance_share_one_batch(): uow.organizations.queue_spend_zero(where=linked) uow.tags.queue_spend_zero(where=linked) uow.model_access_groups.queue_spend_zero(where=linked) + uow.projects.queue_spend_zero(where=linked) uow.endusers.queue_spend_zero(where={"user_id": {"in": ["enduser-1"]}}) uow.budgets.queue_window_advance(budget_id="budget-1", budget_reset_at=reset_at) assert batch.commit_count == 0 @@ -105,6 +107,7 @@ async def test_budget_cascade_dependents_and_window_advance_share_one_batch(): ("litellm_organizationtable.update_many", linked, {"spend": 0}), ("litellm_tagtable.update_many", linked, {"spend": 0}), ("litellm_modelaccessgroupbudgettable.update_many", linked, {"spend": 0}), + ("litellm_projecttable.update_many", linked, {"spend": 0}), ("litellm_endusertable.update_many", {"user_id": {"in": ["enduser-1"]}}, {"spend": 0}), ("litellm_budgettable.update_many", {"budget_id": "budget-1"}, {"budget_reset_at": reset_at}), ] From b00cd15bd73a380c77aad0d04672d27ee4bc13cf Mon Sep 17 00:00:00 2001 From: ryan Date: Wed, 16 Sep 2026 02:07:11 +0000 Subject: [PATCH 011/144] test(proxy): import project_cache_key from user_api_key_cache Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../management_endpoints/test_key_management_endpoints.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index 4e70063015d..b57f8b7857d 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -38,11 +38,10 @@ from litellm.proxy._types import ( from litellm.models.object_permission import LiteLLM_ObjectPermissionTable from litellm.proxy.auth.auth_checks import ( _delete_cache_key_object, - _project_cache_key, jwt_key_mapping_cache_key, ) from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth -from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache, project_cache_key from litellm.litellm_core_utils.duration_parser import duration_in_seconds from litellm.proxy.management_endpoints.key_management_endpoints import ( _check_org_key_limits, @@ -18951,7 +18950,7 @@ async def test_regenerate_key_repoints_live_membership_not_the_key_row_it_read( async def _cache_with_project(project_id: str, project_models: list[str]) -> UserApiKeyCache: user_api_key_cache = UserApiKeyCache() await user_api_key_cache.async_set_cache( - key=_project_cache_key(project_id), + key=project_cache_key(project_id), value=LiteLLM_ProjectTableCachedObj(project_id=project_id, team_id="team-lit-5823", models=project_models), model_type=LiteLLM_ProjectTableCachedObj, ) From b20f1422eb204c2cb2fba26912a548454cd18b02 Mon Sep 17 00:00:00 2001 From: ryan Date: Wed, 16 Sep 2026 02:28:38 +0000 Subject: [PATCH 012/144] fix(proxy): carry project_id through key metadata enrichment and drop docstrings Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/common_utils/user_api_key_cache.py | 2 -- litellm/proxy/hooks/proxy_track_cost_callback.py | 2 ++ tests/test_litellm/proxy/auth/test_auth_checks.py | 3 --- .../proxy/common_utils/test_reset_budget_job.py | 1 - tests/test_litellm/proxy/db/test_db_spend_update_writer.py | 4 ---- tests/test_litellm/proxy/db/test_spend_counter_reseed.py | 2 -- .../proxy/hooks/test_proxy_track_cost_callback.py | 3 +++ tests/test_litellm/proxy/test_budget_reservation.py | 7 ------- 8 files changed, 5 insertions(+), 19 deletions(-) diff --git a/litellm/proxy/common_utils/user_api_key_cache.py b/litellm/proxy/common_utils/user_api_key_cache.py index 2187ed63ea5..0386b58070d 100644 --- a/litellm/proxy/common_utils/user_api_key_cache.py +++ b/litellm/proxy/common_utils/user_api_key_cache.py @@ -307,12 +307,10 @@ def model_access_group_spend_counter_key(access_group_name: str) -> str: def project_cache_key(project_id: str) -> str: - """Cache key one project row is stored under; shared by auth, spend tracking and the spend writer.""" return f"project_id:{project_id}" def project_spend_counter_key(project_id: str) -> str: - """Spend counter key for one project; the reservation, cost callback, auth and reseed paths all read it.""" return f"spend:project:{project_id}" diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index 0c562cf37ef..5e525108ade 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -504,6 +504,8 @@ class _ProxyDBLogger(CustomLogger): metadata["user_api_key_team_id"] = key_obj.team_id if metadata.get("user_api_key_org_id") is None: metadata["user_api_key_org_id"] = key_obj.org_id + if metadata.get("user_api_key_project_id") is None: + metadata["user_api_key_project_id"] = key_obj.project_id except Exception: verbose_proxy_logger.debug( "Failed to enrich failure metadata with key info for api_key=%s", diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index 0d5c0dd5d72..fe469fc574e 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -7161,9 +7161,6 @@ def _project_with_budget(spend: float, max_budget: float): ], ) async def test_project_max_budget_check_reads_live_spend_counter(counter_spend, db_spend, blocks): - """LIT-3269: project budget enforcement must read the cross-pod - ``spend:project:{id}`` counter first and only fall back to the cached row's - spend, matching key/team/org checks. The boundary is inclusive (>=).""" from litellm.caching.dual_cache import DualCache from litellm.proxy.auth.auth_checks import _project_max_budget_check diff --git a/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py b/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py index 8c254686385..48b06649237 100644 --- a/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py +++ b/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py @@ -1666,7 +1666,6 @@ def test_budget_table_reset_invalidates_every_access_group_not_just_the_first( def test_project_reset_zeroes_spend_on_due_tiers(reset_budget_job, mock_prisma_client, monkeypatch): - """A project linked to an expiring budget tier has its spend zeroed in the same cascade transaction.""" _make_counter_invalidation_job(monkeypatch) mock_prisma_client.data["budget"] = [_budget_row(budget_id="budget-due", budget_duration="7d")] mock_prisma_client.db.litellm_projecttable.set_find_many_results( diff --git a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py index da5879a375a..60cc742577b 100644 --- a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py +++ b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py @@ -1062,10 +1062,6 @@ async def test_batch_database_updates_queues_org_member_spend_for_the_request_us @pytest.mark.asyncio async def test_project_spend_is_persisted_to_project_table_and_project_cache_is_evicted(): - """Regression for LIT-3269: a request made with a project-scoped key must - increment LiteLLM_ProjectTable.spend, otherwise /project/info stays at 0 - and the project budget never blocks. The cached project row is evicted so - the next auth check reads the fresh spend.""" db_writer: Final = DBSpendUpdateWriter() await db_writer._batch_database_updates( response_cost=0.25, diff --git a/tests/test_litellm/proxy/db/test_spend_counter_reseed.py b/tests/test_litellm/proxy/db/test_spend_counter_reseed.py index 53e91b8792e..ff0b67d426b 100644 --- a/tests/test_litellm/proxy/db/test_spend_counter_reseed.py +++ b/tests/test_litellm/proxy/db/test_spend_counter_reseed.py @@ -432,8 +432,6 @@ async def test_from_db_bounds_in_flight_prisma_requests_across_counter_keys(): @pytest.mark.asyncio async def test_from_db_reseeds_project_counter_from_the_project_row(): - """LIT-3269: a cold ``spend:project:{id}`` counter seeds from LiteLLM_ProjectTable.spend, - so a fresh pod enforces the project budget against persisted spend rather than 0.""" prisma: Final = _FakePrismaClient(project_row=SimpleNamespace(project_id="proj-1", spend=7.25)) assert await SpendCounterReseed.from_db(prisma_client=prisma, counter_key="spend:project:proj-1") == 7.25 diff --git a/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py b/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py index 965e134772d..202495517ad 100644 --- a/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py +++ b/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py @@ -1372,6 +1372,7 @@ async def test_enrich_failure_metadata_with_full_key_lookup(): mock_key_obj.user_id = "fetched-user-id" mock_key_obj.team_id = "fetched-team-id" mock_key_obj.org_id = "fetched-org-id" + mock_key_obj.project_id = "fetched-project-id" mock_team_obj = MagicMock() mock_team_obj.team_alias = "fetched-team-alias" @@ -1395,12 +1396,14 @@ async def test_enrich_failure_metadata_with_full_key_lookup(): "user_api_key_team_id": None, "user_api_key_team_alias": None, "user_api_key_org_id": None, + "user_api_key_project_id": None, } result = await _ProxyDBLogger._enrich_failure_metadata_with_key_info(metadata) assert result["user_api_key_alias"] == "fetched-key-alias" assert result["user_api_key_user_id"] == "fetched-user-id" assert result["user_api_key_team_id"] == "fetched-team-id" assert result["user_api_key_org_id"] == "fetched-org-id" + assert result["user_api_key_project_id"] == "fetched-project-id" assert result["user_api_key_team_alias"] == "fetched-team-alias" diff --git a/tests/test_litellm/proxy/test_budget_reservation.py b/tests/test_litellm/proxy/test_budget_reservation.py index 014f240d9cc..c834ac05f0a 100644 --- a/tests/test_litellm/proxy/test_budget_reservation.py +++ b/tests/test_litellm/proxy/test_budget_reservation.py @@ -672,10 +672,6 @@ async def _seed_project_scoped_budgets( @pytest.mark.asyncio async def test_should_reserve_project_and_team_member_counters_for_project_scoped_key(spend_counter_state): - """LIT-3269: a key carrying user_id, team_id and project_id reserves against - both the team member counter and the project counter; neither replaces the - other. After the call the project counter reflects the real cost once, not - the reservation plus the post-call increment.""" counter_cache, key_cache = spend_counter_state proxy_logging_obj = ProxyLogging(user_api_key_cache=key_cache) await _seed_project_scoped_budgets( @@ -724,8 +720,6 @@ async def test_should_reserve_project_and_team_member_counters_for_project_scope @pytest.mark.asyncio async def test_exhausted_team_member_budget_still_blocks_project_scoped_key(spend_counter_state): - """LIT-3269: the project budget is additive. A project with plenty of - headroom must not let a key through once its team member budget is spent.""" counter_cache, key_cache = spend_counter_state proxy_logging_obj = ProxyLogging(user_api_key_cache=key_cache) await _seed_project_scoped_budgets( @@ -755,7 +749,6 @@ async def test_exhausted_team_member_budget_still_blocks_project_scoped_key(spen @pytest.mark.asyncio async def test_exhausted_project_budget_blocks_project_scoped_key(spend_counter_state): - """LIT-3269: with team member headroom left, the project budget alone blocks the key.""" counter_cache, key_cache = spend_counter_state proxy_logging_obj = ProxyLogging(user_api_key_cache=key_cache) await _seed_project_scoped_budgets( From 489a3ecf95bd354da16b4d50433552fd3c22f100 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 16 Sep 2026 13:19:45 -0700 Subject: [PATCH 013/144] feat(proxy): add RFC 8693 token exchange for IdP JWTs on the gateway token endpoint A registered gateway DCR client can now POST /token with grant_type=urn:ietf:params:oauth:grant-type:token-exchange and an IdP JWT as subject_token. The gateway proves the JWT the way its JWT auth does, resolves the user and team, and answers with the proxy-API credential and a refresh token, so a fresh laptop with only an IdP login gets a gateway key without a browser round trip. "/token" joins mcp_inference_routes so the default JWT team allowlist reaches the exchange, and the JWT auth builder accepts any header mapping so the request headers pass through unchanged. --- .../mcp_server/discoverable_endpoints.py | 13 +- .../mcp_server/gateway_dcr_flow.py | 148 ++++++++++++-- .../mcp_server/idp_token_exchange.py | 105 ++++++++++ litellm/proxy/_lazy_openapi_snapshot.json | 68 ++++++- litellm/proxy/_types.py | 1 + litellm/proxy/auth/handle_jwt.py | 8 +- .../mcp_server/test_gateway_dcr_flow.py | 186 +++++++++++++++++- .../mcp_server/test_idp_token_exchange.py | 115 +++++++++++ .../proxy/auth/test_auth_checks.py | 18 ++ .../proxy/auth/test_route_checks.py | 1 + ui/litellm-dashboard/src/lib/http/schema.d.ts | 12 ++ 11 files changed, 654 insertions(+), 21 deletions(-) create mode 100644 litellm/proxy/_experimental/mcp_server/idp_token_exchange.py create mode 100644 tests/test_litellm/proxy/_experimental/mcp_server/test_idp_token_exchange.py diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index ffb27d5f92e..ebd5f43bf02 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -46,6 +46,7 @@ from litellm.proxy._experimental.mcp_server.faults import ( render_token_fault, ) from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import ( + TOKEN_EXCHANGE_GRANT_TYPE, VendorCredentialState, aggregate_authorize, aggregate_token, @@ -60,6 +61,9 @@ from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import ( relative_request_url, revoke_refresh_token, ) +from litellm.proxy._experimental.mcp_server.idp_token_exchange import ( + exchange_idp_subject_token, +) from litellm.proxy._experimental.mcp_server.oauth_identity_binding import ( RefreshOwnershipProven, RefreshTokenPresented, @@ -1980,6 +1984,9 @@ async def token_endpoint( refresh_token: str | None = Form(None), scope: str | None = Form(None), resource: str | None = Form(None), + subject_token: str | None = Form(None), + subject_token_type: str | None = Form(None), + requested_token_type: str | None = Form(None), mcp_server_name: str | None = None, ): """ @@ -2010,6 +2017,10 @@ async def token_endpoint( cache=user_api_key_cache, resource=resource, mint_proxy_credential=mint_proxy_credential, + subject_token=subject_token, + subject_token_type=subject_token_type, + requested_token_type=requested_token_type, + exchange_subject_token=exchange_idp_subject_token, ) lookup_name: Final = mcp_server_name or client_id @@ -2638,7 +2649,7 @@ def _build_aggregate_authorization_server_response(request: Request) -> dict: "registration_endpoint": f"{request_base_url}/register", "response_types_supported": ["code"], "scopes_supported": [], - "grant_types_supported": ["authorization_code", "refresh_token"], + "grant_types_supported": ("authorization_code", "refresh_token", TOKEN_EXCHANGE_GRANT_TYPE), "code_challenge_methods_supported": ["S256"], "token_endpoint_auth_methods_supported": ["none", "client_secret_post"], } diff --git a/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py b/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py index f3fdd54b39d..9bdde3c5edc 100644 --- a/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py +++ b/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py @@ -51,7 +51,7 @@ from urllib.parse import parse_qsl, urlencode, urlparse, urlunparse from fastapi import HTTPException, Request from fastapi.responses import HTMLResponse, JSONResponse, RedirectResponse, Response from pydantic import BaseModel, ConfigDict, Field, ValidationError -from typing_extensions import ReadOnly, TypedDict, assert_never +from typing_extensions import NotRequired, ReadOnly, TypedDict, assert_never from litellm._logging import verbose_logger from litellm.caching.caching import DualCache @@ -187,6 +187,41 @@ class MintProxyCredential(Protocol): ) -> Awaitable[MintedProxyCredential | ProxyCredentialMintFailure]: ... +TOKEN_EXCHANGE_GRANT_TYPE: Final = "urn:ietf:params:oauth:grant-type:token-exchange" +"""RFC 8693: a native client that already holds a token from the customer's identity +provider trades it for the proxy-API credential without a browser round trip.""" + +_IssuedTokenType = Literal["urn:ietf:params:oauth:token-type:access_token"] +ACCESS_TOKEN_TOKEN_TYPE: Final[_IssuedTokenType] = "urn:ietf:params:oauth:token-type:access_token" +SUBJECT_TOKEN_TYPES: Final = frozenset( + { + "urn:ietf:params:oauth:token-type:jwt", + "urn:ietf:params:oauth:token-type:id_token", + ACCESS_TOKEN_TOKEN_TYPE, + } +) + + +class SubjectIdentity(BaseModel): + model_config = ConfigDict(frozen=True) + user_id: str = Field(min_length=1) + team_id: str | None = None + + +class SubjectTokenRefusal(BaseModel): + model_config = ConfigDict(frozen=True) + error: Literal["unsupported_grant_type", "invalid_request"] + description: str = Field(min_length=1) + + +class ExchangeSubjectToken(Protocol): + """Injected RFC 8693 subject-token verifier ``(subject_token, request)``: proves the + IdP token the way the proxy's own JWT auth does and names the litellm user and team it + stands for, or says why this gateway will not take it.""" + + def __call__(self, subject_token: str, request: Request, /) -> Awaitable[SubjectIdentity | SubjectTokenRefusal]: ... + + class ConsentTeam(BaseModel): model_config = ConfigDict(frozen=True) team_id: str = Field(min_length=1) @@ -213,6 +248,12 @@ async def _refuse_proxy_credential(user_id: str, team_id: str | None) -> ProxyCr return "unresolvable" +async def _refuse_subject_token(subject_token: str, request: Request) -> SubjectTokenRefusal: + return SubjectTokenRefusal( + error="unsupported_grant_type", description="this gateway is not configured to exchange IdP tokens" + ) + + async def _unavailable_vendor_credential(user_id: str, server_id: str) -> VendorCredentialState: return "unavailable" @@ -382,7 +423,7 @@ async def register_aggregate_client(request: Request, request_body: Mapping[str, "client_id_issued_at": int(now.timestamp()), "redirect_uris": list(raw_uris), "token_endpoint_auth_method": "none", - "grant_types": ["authorization_code", "refresh_token"], + "grant_types": ["authorization_code", "refresh_token", TOKEN_EXCHANGE_GRANT_TYPE], "response_types": ["code"], }, ) @@ -595,7 +636,7 @@ def native_client_auth_contract(request: Request) -> NativeClientAuthContract: "revocation_endpoint": f"{base_url}/revoke", "resource": base_url, "response_types_supported": ("code",), - "grant_types_supported": ("authorization_code", "refresh_token"), + "grant_types_supported": ("authorization_code", "refresh_token", TOKEN_EXCHANGE_GRANT_TYPE), "code_challenge_methods_supported": ("S256",), "token_endpoint_auth_methods_supported": ("none",), "revocation_endpoint_auth_methods_supported": ("none",), @@ -1033,20 +1074,26 @@ class _ProxyCredentialTokenResponse(TypedDict): refresh_token: ReadOnly[str] user_id: ReadOnly[str] team_id: ReadOnly[str | None] + issued_token_type: NotRequired[ReadOnly[_IssuedTokenType]] def _proxy_credential_response( - minted: MintedProxyCredential, principal: SessionPrincipal, keys: SessionSigningKeys, now: datetime + minted: MintedProxyCredential, + principal: SessionPrincipal, + keys: SessionSigningKeys, + now: datetime, + issued_token_type: _IssuedTokenType | None = None, ) -> Response: """The proxy-API token response: the access token is the very credential ``lite login`` stores (accepted on every proxy route with user and team attribution), and the refresh token is a gateway-sealed rotating token bound to the team the credential - was minted for, so a renewal keeps the team the user consented to.""" + was minted for, so a renewal keeps the team the user consented to. A token exchange + also states ``issued_token_type``, which RFC 8693 section 2.2.1 requires.""" bound_principal: Final = principal.model_copy(update=MappingProxyType({"team_id": minted.team_id})) refresh: Final = mint_session_refresh_token(bound_principal, keys, now) if not isinstance(refresh, MintedSessionToken): return _oauth_error(500, "server_error", "failed to mint the session credential") - body: Final[_ProxyCredentialTokenResponse] = { + credential: Final[_ProxyCredentialTokenResponse] = { "access_token": minted.key, "token_type": "Bearer", "expires_in": minted.expires_in, @@ -1054,7 +1101,10 @@ def _proxy_credential_response( "user_id": minted.user_id, "team_id": minted.team_id, } - return JSONResponse(status_code=200, content=body, headers=TOKEN_NO_CACHE_HEADERS) + if issued_token_type is None: + return JSONResponse(status_code=200, content=credential, headers=TOKEN_NO_CACHE_HEADERS) + exchanged: Final[_ProxyCredentialTokenResponse] = {**credential, "issued_token_type": issued_token_type} + return JSONResponse(status_code=200, content=exchanged, headers=TOKEN_NO_CACHE_HEADERS) def _reload_failure_response(failure: ReloadUserFailure) -> Response: @@ -1116,11 +1166,16 @@ async def aggregate_token( cache: DualCache, resource: str | None = None, mint_proxy_credential: MintProxyCredential = _refuse_proxy_credential, + subject_token: str | None = None, + subject_token_type: str | None = None, + requested_token_type: str | None = None, + exchange_subject_token: ExchangeSubjectToken = _refuse_subject_token, ) -> Response: """The aggregate token verb: authorization_code and refresh_token grants for the identity-only session pair, or for the proxy-API credential when the grant was issued - with that audience. Every path re-validates the litellm user live before minting, so a - deactivated user cannot obtain or renew a session.""" + with that audience, and the RFC 8693 token exchange that turns an IdP token straight + into the proxy-API credential. Every path re-validates the litellm user live before + minting, so a deactivated user cannot obtain or renew a session.""" if master_key is None: verbose_logger.error("mcp_gateway_dcr token grant rejected: no master_key configured") return _oauth_error(500, "server_error", "the gateway has no master key configured") @@ -1159,7 +1214,20 @@ async def aggregate_token( now=now, issue=issue, ) - return _oauth_error(400, "unsupported_grant_type", "grant_type must be authorization_code or refresh_token") + if grant_type == TOKEN_EXCHANGE_GRANT_TYPE: + return await _token_exchange_grant( + subject_token=subject_token, + subject_token_type=subject_token_type, + requested_token_type=requested_token_type, + client_id=client_id, + exchange_subject_token=exchange_subject_token, + issue=issue, + ) + return _oauth_error( + 400, + "unsupported_grant_type", + f"grant_type must be authorization_code, refresh_token, or {TOKEN_EXCHANGE_GRANT_TYPE}", + ) class _GrantIssuer: @@ -1211,10 +1279,9 @@ class _GrantIssuer: async def _issue_proxy_credential( self, principal: SessionPrincipal, claim_key: str, claim_ttl_seconds: int, replayed: str ) -> Response: - if self._resource is not None and not is_proxy_api_resource(self._request, self._resource): - return _oauth_error( - 400, "invalid_target", "resource does not match the proxy API this grant was issued for" - ) + target_refusal: Final = self._proxy_api_target_refusal() + if target_refusal is not None: + return target_refusal minted: Final = await self._mint_proxy_credential(principal.user_id, principal.team_id) if not isinstance(minted, MintedProxyCredential): return _mint_failure_response(minted) @@ -1223,6 +1290,33 @@ class _GrantIssuer: return refusal return _proxy_credential_response(minted, principal, self._keys, self._now) + async def exchange( + self, subject_token: str, client_id: str, exchange_subject_token: ExchangeSubjectToken + ) -> Response: + """The RFC 8693 tail: prove the IdP token, then mint. No single-use marker, because + the subject token stays a valid proof for as long as the IdP says it is and every + exchange mints a fresh credential and refresh token of its own.""" + target_refusal: Final = self._proxy_api_target_refusal() + if target_refusal is not None: + return target_refusal + identity: Final = await exchange_subject_token(subject_token, self._request) + if isinstance(identity, SubjectTokenRefusal): + return _oauth_error(400, identity.error, identity.description) + principal: Final = SessionPrincipal( + user_id=identity.user_id, client_id=client_id, audience=PROXY_API_AUDIENCE, team_id=identity.team_id + ) + minted: Final = await self._mint_proxy_credential(principal.user_id, principal.team_id) + if not isinstance(minted, MintedProxyCredential): + return _mint_failure_response(minted) + return _proxy_credential_response( + minted, principal, self._keys, self._now, issued_token_type=ACCESS_TOKEN_TOKEN_TYPE + ) + + def _proxy_api_target_refusal(self) -> Response | None: + if self._resource is None or is_proxy_api_resource(self._request, self._resource): + return None + return _oauth_error(400, "invalid_target", "resource does not match the proxy API this grant was issued for") + async def _claim_refusal(self, claim_key: str, claim_ttl_seconds: int, replayed: str) -> Response | None: return _claim_refusal( await self._guard.claim(claim_key, claim_ttl_seconds), replayed=_oauth_error(400, "invalid_grant", replayed) @@ -1297,6 +1391,32 @@ async def _refresh_token_grant( ) +async def _token_exchange_grant( + subject_token: str | None, + subject_token_type: str | None, + requested_token_type: str | None, + client_id: str, + exchange_subject_token: ExchangeSubjectToken, + issue: _GrantIssuer, +) -> Response: + """RFC 8693 token exchange for a registered native client that already holds an IdP + token: the gateway proves the token the way its JWT auth does and answers with the + proxy-API credential, so a fresh laptop with only an IdP login gets a gateway key + without a browser round trip. The client must be registered because the refresh token + in the answer is bound to it.""" + if not is_gateway_dcr_client_id(client_id) or open_gateway_dcr_client(client_id) is None: + return _oauth_error(401, "invalid_client", "unknown or malformed client_id") + if not subject_token or not subject_token_type: + return _oauth_error(400, "invalid_request", "subject_token and subject_token_type are required") + if subject_token_type not in SUBJECT_TOKEN_TYPES: + return _oauth_error( + 400, "invalid_request", f"subject_token_type must be one of {', '.join(sorted(SUBJECT_TOKEN_TYPES))}" + ) + if requested_token_type is not None and requested_token_type != ACCESS_TOKEN_TOKEN_TYPE: + return _oauth_error(400, "invalid_request", f"requested_token_type must be {ACCESS_TOKEN_TOKEN_TYPE}") + return await issue.exchange(subject_token, client_id, exchange_subject_token) + + async def revoke_refresh_token(token: str, client_id: str, master_key: str | None, cache: DualCache) -> Response: """RFC 7009 revocation for the gateway's refresh tokens: burn the presented token's ``jti`` so neither the holder nor a thief can rotate it again. Access tokens are diff --git a/litellm/proxy/_experimental/mcp_server/idp_token_exchange.py b/litellm/proxy/_experimental/mcp_server/idp_token_exchange.py new file mode 100644 index 00000000000..69b81713810 --- /dev/null +++ b/litellm/proxy/_experimental/mcp_server/idp_token_exchange.py @@ -0,0 +1,105 @@ +"""The identity-provider side of the RFC 8693 token exchange on ``POST /token``: a native +client that already holds a JWT from the customer's IdP trades it for the same proxy-API +credential ``lite login`` stores, proven by the proxy's own JWT auth (signature, claims, +and the user and team sync it performs), so no browser round trip is needed.""" + +from __future__ import annotations + +from collections.abc import Awaitable, Callable, Mapping +from typing import Final, Protocol + +from fastapi import HTTPException, Request + +from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import SubjectIdentity, SubjectTokenRefusal +from litellm.proxy._types import JWTAuthBuilderResult, ProxyException +from litellm.proxy.auth.handle_jwt import JWTAuthManager + +EXCHANGE_ROUTE: Final = "/token" + + +class AuthorizeSubjectToken(Protocol): + """Injected JWT authorization ``(subject_token, request_headers)``: the proxy's + ``JWTAuthManager.auth_builder`` in production, which raises when the token is not + acceptable and otherwise names the user and team it resolved.""" + + def __call__( + self, subject_token: str, request_headers: Mapping[str, str], / + ) -> Awaitable[JWTAuthBuilderResult]: ... + + +async def exchange_idp_subject_token(subject_token: str, request: Request) -> SubjectIdentity | SubjectTokenRefusal: + from litellm.proxy.proxy_server import ( # noqa: PLC0415 # rebound after startup, so read them per call + general_settings, + jwt_handler, + premium_user, + prisma_client, + proxy_logging_obj, + user_api_key_cache, + ) + + async def authorize(token: str, request_headers: Mapping[str, str]) -> JWTAuthBuilderResult: + return await JWTAuthManager.auth_builder( + api_key=token, + jwt_handler=jwt_handler, + request_data={}, + general_settings=general_settings, + route=EXCHANGE_ROUTE, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=None, + proxy_logging_obj=proxy_logging_obj, + request_headers=request_headers, + request_method="POST", + ) + + return await identity_from_subject_token( + subject_token, + request_headers=request.headers, + jwt_auth_enabled=general_settings.get("enable_jwt_auth", False) is True, + has_database=prisma_client is not None, + licensed=premium_user is True, + is_jwt=jwt_handler.is_jwt, + authorize=authorize, + ) + + +async def identity_from_subject_token( + subject_token: str, + request_headers: Mapping[str, str], + jwt_auth_enabled: bool, + has_database: bool, + licensed: bool, + is_jwt: Callable[[str], bool], + authorize: AuthorizeSubjectToken, +) -> SubjectIdentity | SubjectTokenRefusal: + """Apply the same gates ``user_api_key_auth`` applies to a JWT bearer, then let the + proxy's JWT auth prove the token. A rejection comes back as ``invalid_request``, which + RFC 8693 section 2.2.2 prescribes for an invalid or unacceptable subject token.""" + if not jwt_auth_enabled: + return SubjectTokenRefusal( + error="unsupported_grant_type", + description="JWT auth is not enabled on this gateway, so it cannot exchange IdP tokens", + ) + if not has_database: + return SubjectTokenRefusal( + error="unsupported_grant_type", + description="this gateway has no database, so it cannot exchange IdP tokens", + ) + if not is_jwt(subject_token): + return SubjectTokenRefusal(error="invalid_request", description="subject_token is not a JWT") + if not licensed: + return SubjectTokenRefusal( + error="unsupported_grant_type", description="JWT auth is an enterprise only feature; no license is set" + ) + try: + result: Final = await authorize(subject_token, request_headers) + except HTTPException as denied: + return SubjectTokenRefusal(error="invalid_request", description=f"subject_token was rejected: {denied.detail}") + except ProxyException as denied: + return SubjectTokenRefusal(error="invalid_request", description=f"subject_token was rejected: {denied.message}") + except Exception as denied: # noqa: BLE001 # auth_jwt raises a plain Exception on signature and claim failures + return SubjectTokenRefusal(error="invalid_request", description=f"subject_token was rejected: {denied}") + user_id: Final = result["user_id"] + if user_id is None: + return SubjectTokenRefusal(error="invalid_request", description="subject_token names no user the gateway knows") + return SubjectIdentity(user_id=user_id, team_id=result["team_id"]) diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index 74f38b3ca6d..6893c06d2d7 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -19346,7 +19346,7 @@ } } }, - "description": "\n Unified rate-limit error.\n\n Every rate-limit condition surfaced by litellm \u2014 whether it originated from\n an upstream LLM provider, a vendor batch endpoint, or one of litellm's own\n proxy-side limiters (parallel-requests, dynamic-rate, batch-rate, budget,\n max-iterations, etc.) \u2014 is raised as an instance of this class.\n\n The :attr:`category` attribute lets callers distinguish the source. See\n :class:`RateLimitErrorCategory` for the available values.\n " + "description": "\nUnified rate-limit error.\n\nEvery rate-limit condition surfaced by litellm \u2014 whether it originated from\nan upstream LLM provider, a vendor batch endpoint, or one of litellm's own\nproxy-side limiters (parallel-requests, dynamic-rate, batch-rate, budget,\nmax-iterations, etc.) \u2014 is raised as an instance of this class.\n\nThe :attr:`category` attribute lets callers distinguish the source. See\n:class:`RateLimitErrorCategory` for the available values.\n" }, "500": { "content": { @@ -23115,6 +23115,17 @@ ], "title": "Refresh Token" }, + "requested_token_type": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Requested Token Type" + }, "resource": { "anyOf": [ { @@ -23136,6 +23147,28 @@ } ], "title": "Scope" + }, + "subject_token": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Subject Token" + }, + "subject_token_type": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Subject Token Type" } }, "required": [ @@ -23189,6 +23222,17 @@ ], "title": "Refresh Token" }, + "requested_token_type": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Requested Token Type" + }, "resource": { "anyOf": [ { @@ -23210,6 +23254,28 @@ } ], "title": "Scope" + }, + "subject_token": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Subject Token" + }, + "subject_token_type": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Subject Token Type" } }, "required": [ diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 63d0bfcc5b8..5f7365cedd5 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -525,6 +525,7 @@ class LiteLLMRoutes(enum.Enum): "/mcp-rest/tools/call", "/v1/mcp/tools", "/introspect", + "/token", ] # MCP server CRUD routes — control-plane. Gated by DISABLE_ADMIN_ENDPOINTS. diff --git a/litellm/proxy/auth/handle_jwt.py b/litellm/proxy/auth/handle_jwt.py index 6a28cd7ff99..803093ff93a 100644 --- a/litellm/proxy/auth/handle_jwt.py +++ b/litellm/proxy/auth/handle_jwt.py @@ -1867,7 +1867,7 @@ class JWTAuthManager: @staticmethod def get_team_id_from_header( - request_headers: dict | None, + request_headers: Mapping[str, str] | None, allowed_team_ids: set[str], fallback_to_db_teams: bool = False, ) -> str | None: @@ -2037,7 +2037,7 @@ class JWTAuthManager: async def _attach_team_from_header_for_admin( admin_result: JWTAuthBuilderResult, route: str, - request_headers: dict | None, + request_headers: Mapping[str, str] | None, jwt_handler: JWTHandler, prisma_client: PrismaClient | None, user_api_key_cache: UserApiKeyCache, @@ -2293,7 +2293,7 @@ class JWTAuthManager: user_api_key_cache: UserApiKeyCache, parent_otel_span: Span | None, proxy_logging_obj: ProxyLogging, - request_headers: dict | None = None, + request_headers: Mapping[str, str] | None = None, request_method: str | None = None, ) -> JWTAuthBuilderResult: return await JWTAuthManager.authorize_jwt( @@ -2390,7 +2390,7 @@ class JWTAuthManager: user_api_key_cache: UserApiKeyCache, parent_otel_span: Span | None, proxy_logging_obj: ProxyLogging, - request_headers: dict[str, str] | None = None, + request_headers: Mapping[str, str] | None = None, request_method: str | None = None, provisioning: _JWTProvisioning | None = None, ) -> JWTAuthBuilderResult: diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py index 7c80ee77cd7..1be6ebb6e22 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py @@ -15,13 +15,18 @@ from starlette.requests import Request from litellm.caching.caching import DualCache from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import ( _AUTH_CODE_DEBUG_KEY, + ACCESS_TOKEN_TOKEN_TYPE, CONNECT_FLOW_COOKIE_PREFIX, GATEWAY_AUTH_CODE_PREFIX, GATEWAY_AUTH_CODE_TTL_SECONDS, MANUAL_DELIVERY_AUTH_CODE_TTL_SECONDS, MAX_CLIENT_ID_LENGTH, + SUBJECT_TOKEN_TYPES, + TOKEN_EXCHANGE_GRANT_TYPE, ConsentTeam, MintedProxyCredential, + SubjectIdentity, + SubjectTokenRefusal, _GatewayAuthCode, _open_sealed, _seal, @@ -105,6 +110,7 @@ async def _reload_user_active(user_id: str): async def test_register_mints_stateless_public_client(): body = await _register([REDIRECT_URI]) assert body["token_endpoint_auth_method"] == "none" + assert body["grant_types"] == ["authorization_code", "refresh_token", TOKEN_EXCHANGE_GRANT_TYPE] assert "client_secret" not in body assert body["redirect_uris"] == [REDIRECT_URI] assert is_gateway_dcr_client_id(body["client_id"]) @@ -1957,7 +1963,11 @@ def test_native_client_auth_contract_points_every_endpoint_at_this_proxy(): "revocation_endpoint": "https://llm.example.com/revoke", "resource": "https://llm.example.com", "response_types_supported": ["code"], - "grant_types_supported": ["authorization_code", "refresh_token"], + "grant_types_supported": [ + "authorization_code", + "refresh_token", + "urn:ietf:params:oauth:grant-type:token-exchange", + ], "code_challenge_methods_supported": ["S256"], "token_endpoint_auth_methods_supported": ["none"], "revocation_endpoint_auth_methods_supported": ["none"], @@ -2148,3 +2158,177 @@ async def test_gateway_owned_resource_stays_scoped_through_consent_and_refresh(a ) assert renewed.status_code == 200 assert _opened_principal(json.loads(renewed.body)).resource_server_id == "github-id" + + +JWT_SUBJECT_TOKEN_TYPE = "urn:ietf:params:oauth:token-type:jwt" +IDP_TOKEN = "eyJhbGciOiJSUzI1NiJ9.eyJzdWIiOiJ1MSJ9.idp-signature" + + +class _Exchanger: + def __init__(self, result=None): + self.calls = [] + self.result = result + + async def __call__(self, subject_token, request): + self.calls.append((subject_token, request.url.path)) + if self.result is not None: + return self.result + return SubjectIdentity(user_id="u1", team_id="team-b") + + +async def _exchange_native(client_id, minter, exchanger, cache=None, **overrides): + arguments = { + "grant_type": TOKEN_EXCHANGE_GRANT_TYPE, + "subject_token": IDP_TOKEN, + "subject_token_type": JWT_SUBJECT_TOKEN_TYPE, + "exchange_subject_token": exchanger, + } + return await _redeem_native(None, client_id, minter, cache=cache, **{**arguments, **overrides}) + + +@pytest.mark.asyncio +async def test_token_exchange_mints_the_proxy_credential_for_the_idp_subject(): + """RFC 8693: a registered native client trades the IdP token it already holds for the + same credential the consent flow mints, attributed to the user and team the gateway's + JWT auth resolved, with a rotating refresh token bound to that team and the client. + The exchange can be repeated while the IdP token lives; nothing is burned.""" + client_id = (await _register([LOOPBACK_REDIRECT_URI]))["client_id"] + minter, exchanger, cache = _Minter(), _Exchanger(), DualCache() + response = await _exchange_native(client_id, minter, exchanger, cache=cache) + assert response.status_code == 200 + assert response.headers["cache-control"] == "no-store" + body = json.loads(response.body) + assert exchanger.calls == [(IDP_TOKEN, "/token")] + assert minter.calls == [("u1", "team-b")] + assert body["issued_token_type"] == ACCESS_TOKEN_TOKEN_TYPE + assert body["access_token"] == "sk-cli-u1" + assert body["token_type"] == "Bearer" + assert body["expires_in"] == 3600 + assert (body["user_id"], body["team_id"]) == ("u1", "team-b") + principal = _opened_refresh(body["refresh_token"], client_id) + assert (principal.user_id, principal.client_id, principal.audience, principal.team_id) == ( + "u1", + client_id, + "proxy_api", + "team-b", + ) + again = await _exchange_native(client_id, minter, exchanger, cache=cache) + assert again.status_code == 200 + assert json.loads(again.body)["refresh_token"] != body["refresh_token"] + assert minter.calls == [("u1", "team-b"), ("u1", "team-b")] + + +@pytest.mark.asyncio +async def test_exchanged_credential_refreshes_and_rotates_like_a_consented_one(): + client_id = (await _register([LOOPBACK_REDIRECT_URI]))["client_id"] + minter, cache = _Minter(), DualCache() + exchanged = json.loads((await _exchange_native(client_id, minter, _Exchanger(), cache=cache)).body) + refreshed = await _refresh_native(exchanged["refresh_token"], client_id, minter, cache) + assert refreshed.status_code == 200 + body = json.loads(refreshed.body) + assert "issued_token_type" not in body + assert (body["access_token"], body["user_id"], body["team_id"]) == ("sk-cli-u1", "u1", "team-b") + assert body["refresh_token"] != exchanged["refresh_token"] + assert minter.calls == [("u1", "team-b"), ("u1", "team-b")] + replay = await _refresh_native(exchanged["refresh_token"], client_id, minter, cache) + assert replay.status_code == 400 + assert json.loads(replay.body)["error"] == "invalid_grant" + + +@pytest.mark.asyncio +async def test_token_exchange_for_a_teamless_subject_mints_a_teamless_credential(): + client_id = (await _register([LOOPBACK_REDIRECT_URI]))["client_id"] + minter = _Minter() + response = await _exchange_native(client_id, minter, _Exchanger(SubjectIdentity(user_id="u2"))) + assert response.status_code == 200 + body = json.loads(response.body) + assert minter.calls == [("u2", None)] + assert (body["user_id"], body["team_id"]) == ("u2", None) + assert _opened_refresh(body["refresh_token"], client_id).team_id is None + + +@pytest.mark.asyncio +@pytest.mark.parametrize("subject_token_type", sorted(SUBJECT_TOKEN_TYPES)) +async def test_token_exchange_accepts_every_advertised_subject_token_type(subject_token_type): + client_id = (await _register([LOOPBACK_REDIRECT_URI]))["client_id"] + response = await _exchange_native(client_id, _Minter(), _Exchanger(), subject_token_type=subject_token_type) + assert response.status_code == 200 + + +@pytest.mark.asyncio +async def test_token_exchange_without_an_idp_exchanger_is_unsupported(): + """A gateway that wires no IdP verifier into the endpoint answers the way it always + answered an unknown grant, and never reaches the minter.""" + client_id = (await _register([LOOPBACK_REDIRECT_URI]))["client_id"] + minter = _Minter() + response = await _redeem_native( + None, + client_id, + minter, + grant_type=TOKEN_EXCHANGE_GRANT_TYPE, + subject_token=IDP_TOKEN, + subject_token_type=JWT_SUBJECT_TOKEN_TYPE, + ) + assert response.status_code == 400 + assert json.loads(response.body)["error"] == "unsupported_grant_type" + assert minter.calls == [] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "overrides, status, error", + [ + ({"subject_token": None}, 400, "invalid_request"), + ({"subject_token": ""}, 400, "invalid_request"), + ({"subject_token_type": None}, 400, "invalid_request"), + ({"subject_token_type": "urn:ietf:params:oauth:token-type:saml2"}, 400, "invalid_request"), + ({"requested_token_type": "urn:ietf:params:oauth:token-type:refresh_token"}, 400, "invalid_request"), + ({"resource": "https://other.example.com"}, 400, "invalid_target"), + ({"resource": "https://llm.example.com/mcp"}, 400, "invalid_target"), + ({"client_id": "llm_dcrc_forged"}, 401, "invalid_client"), + ({"client_id": "not-a-gateway-client"}, 401, "invalid_client"), + ], +) +async def test_token_exchange_refuses_a_malformed_request_before_touching_the_idp_token(overrides, status, error): + registered = (await _register([LOOPBACK_REDIRECT_URI]))["client_id"] + minter, exchanger = _Minter(), _Exchanger() + response = await _exchange_native( + overrides.get("client_id", registered), + minter, + exchanger, + **{name: value for name, value in overrides.items() if name != "client_id"}, + ) + assert response.status_code == status + assert json.loads(response.body)["error"] == error + assert exchanger.calls == [] + assert minter.calls == [] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("error", ["unsupported_grant_type", "invalid_request"]) +async def test_token_exchange_relays_the_idp_refusal_and_never_mints(error): + client_id = (await _register([LOOPBACK_REDIRECT_URI]))["client_id"] + minter = _Minter() + exchanger = _Exchanger(SubjectTokenRefusal(error=error, description="subject_token was rejected: bad signature")) + response = await _exchange_native(client_id, minter, exchanger) + assert response.status_code == 400 + body = json.loads(response.body) + assert (body["error"], body["error_description"]) == (error, "subject_token was rejected: bad signature") + assert minter.calls == [] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "failure, status, error", + [ + ("not_a_member", 400, "invalid_grant"), + ("team_required", 400, "invalid_grant"), + ("no_active_key", 400, "invalid_grant"), + ("unavailable", 503, "temporarily_unavailable"), + ], +) +async def test_token_exchange_relays_a_mint_refusal(failure, status, error): + client_id = (await _register([LOOPBACK_REDIRECT_URI]))["client_id"] + response = await _exchange_native(client_id, _Minter(failure), _Exchanger()) + assert response.status_code == status + assert json.loads(response.body)["error"] == error diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_idp_token_exchange.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_idp_token_exchange.py new file mode 100644 index 00000000000..00440614120 --- /dev/null +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_idp_token_exchange.py @@ -0,0 +1,115 @@ +import pytest +from fastapi import HTTPException + +from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import SubjectIdentity, SubjectTokenRefusal +from litellm.proxy._experimental.mcp_server.idp_token_exchange import identity_from_subject_token +from litellm.proxy._types import ProxyException +from litellm.proxy.auth.handle_jwt import JWTHandler + +IDP_JWT = "eyJhbGciOiJSUzI1NiJ9.eyJzdWIiOiJ1MSJ9.idp-signature" +REQUEST_HEADERS = {"x-litellm-team-id": "team-b", "user-agent": "lite/0.1"} + + +def _authorized(user_id="u1", team_id="team-b"): + return { + "is_proxy_admin": False, + "team_object": None, + "user_object": None, + "end_user_object": None, + "org_object": None, + "token": IDP_JWT, + "team_id": team_id, + "user_id": user_id, + "user_email": None, + "end_user_id": None, + "org_id": None, + "team_membership": None, + "jwt_claims": {"sub": user_id}, + "agent_id": None, + } + + +class _Authorizer: + def __init__(self, result=None, raises=None): + self.calls = [] + self.result = result if result is not None else _authorized() + self.raises = raises + + async def __call__(self, subject_token, request_headers): + self.calls.append((subject_token, dict(request_headers))) + if self.raises is not None: + raise self.raises + return self.result + + +async def _identity(authorizer, subject_token=IDP_JWT, **overrides): + arguments = { + "request_headers": REQUEST_HEADERS, + "jwt_auth_enabled": True, + "has_database": True, + "licensed": True, + "is_jwt": JWTHandler.is_jwt, + "authorize": authorizer, + } + return await identity_from_subject_token(subject_token, **{**arguments, **overrides}) + + +@pytest.mark.asyncio +async def test_a_jwt_the_proxy_accepts_names_its_user_and_team(): + """The subject token goes to the proxy's own JWT auth with the caller's headers (that is + where the team header is read), and the identity it resolved is what gets minted.""" + authorizer = _Authorizer() + assert await _identity(authorizer) == SubjectIdentity(user_id="u1", team_id="team-b") + assert authorizer.calls == [(IDP_JWT, REQUEST_HEADERS)] + + +@pytest.mark.asyncio +async def test_a_jwt_that_resolves_no_team_names_a_teamless_identity(): + assert await _identity(_Authorizer(_authorized(team_id=None))) == SubjectIdentity(user_id="u1", team_id=None) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "overrides, subject_token, error, mentions", + [ + ({"jwt_auth_enabled": False}, IDP_JWT, "unsupported_grant_type", "JWT auth is not enabled"), + ({"has_database": False}, IDP_JWT, "unsupported_grant_type", "no database"), + ({"licensed": False}, IDP_JWT, "unsupported_grant_type", "enterprise"), + ({}, "sk-litellm-virtual-key", "invalid_request", "not a JWT"), + ], +) +async def test_the_gates_user_api_key_auth_applies_refuse_before_any_verification( + overrides, subject_token, error, mentions +): + authorizer = _Authorizer() + refusal = await _identity(authorizer, subject_token=subject_token, **overrides) + assert isinstance(refusal, SubjectTokenRefusal) + assert refusal.error == error + assert mentions in refusal.description + assert authorizer.calls == [] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "raised, mentions", + [ + (HTTPException(status_code=403, detail="User not allowed to access this route"), "not allowed"), + (ProxyException(message="Token expired", type="auth_error", param="token", code=401), "Token expired"), + (Exception("Validation fails: signature verification failed"), "signature verification failed"), + (Exception("Invalid JWT Submitted"), "Invalid JWT"), + ], +) +async def test_a_jwt_the_proxy_rejects_is_an_invalid_subject_token(raised, mentions): + refusal = await _identity(_Authorizer(raises=raised)) + assert isinstance(refusal, SubjectTokenRefusal) + assert refusal.error == "invalid_request" + assert refusal.description.startswith("subject_token was rejected: ") + assert mentions in refusal.description + + +@pytest.mark.asyncio +async def test_a_jwt_that_resolves_no_user_cannot_be_exchanged(): + refusal = await _identity(_Authorizer(_authorized(user_id=None))) + assert refusal == SubjectTokenRefusal( + error="invalid_request", description="subject_token names no user the gateway knows" + ) diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index 26ae28a57d2..e4c03399b30 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -8388,3 +8388,21 @@ async def test_access_group_model_fallback_uses_the_injected_database(channel: s llm_router=None, prisma_client=client, ) is True reader.assert_awaited_once_with(where={"access_group_id": "group-a"}) + + +def test_jwt_team_role_reaches_the_gateway_token_endpoint_by_default(): + """The RFC 8693 token exchange authorizes the IdP JWT against ``POST /token`` itself, and JWT + auth only binds a team from a multi-team claim when that team may call the route, so the + default team allowlist has to cover the gateway's token endpoint or the exchange would mint + teamless credentials for every ``team_ids_jwt_field`` deployment.""" + from litellm.proxy._types import LiteLLM_JWTAuth + from litellm.proxy.auth.auth_checks import allowed_routes_check + + assert allowed_routes_check( + user_role=LitellmUserRoles.TEAM, user_route="/token", litellm_proxy_roles=LiteLLM_JWTAuth() + ) + assert not allowed_routes_check( + user_role=LitellmUserRoles.TEAM, + user_route="/token", + litellm_proxy_roles=LiteLLM_JWTAuth(team_allowed_routes=[]), + ) diff --git a/tests/test_litellm/proxy/auth/test_route_checks.py b/tests/test_litellm/proxy/auth/test_route_checks.py index 806c55d51ce..630f69bf0f1 100644 --- a/tests/test_litellm/proxy/auth/test_route_checks.py +++ b/tests/test_litellm/proxy/auth/test_route_checks.py @@ -627,6 +627,7 @@ def test_virtual_key_llm_api_routes_denies_spend_logs_v2(): "/mcp/tools/call", "/mcp-rest/tools/call", "/mcp/tools/list", + "/token", ], ) def test_mcp_inference_routes_classified_as_llm_api(route): diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 17ec8367324..b896f6901cb 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -24505,10 +24505,16 @@ export interface components { redirect_uri?: string; /** Refresh Token */ refresh_token?: string | null; + /** Requested Token Type */ + requested_token_type?: string | null; /** Resource */ resource?: string | null; /** Scope */ scope?: string | null; + /** Subject Token */ + subject_token?: string | null; + /** Subject Token Type */ + subject_token_type?: string | null; }; /** Body_token_endpoint_token_post */ Body_token_endpoint_token_post: { @@ -24526,10 +24532,16 @@ export interface components { redirect_uri?: string; /** Refresh Token */ refresh_token?: string | null; + /** Requested Token Type */ + requested_token_type?: string | null; /** Resource */ resource?: string | null; /** Scope */ scope?: string | null; + /** Subject Token */ + subject_token?: string | null; + /** Subject Token Type */ + subject_token_type?: string | null; }; /** Body_upload_logo_upload_logo_post */ Body_upload_logo_upload_logo_post: { From 0bacf26b1e147f709ea8e5c638ef81f6b7288599 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 16 Sep 2026 13:32:51 -0700 Subject: [PATCH 014/144] chore(proxy): keep the lazy OpenAPI snapshot as the CI Python renders it --- litellm/proxy/_lazy_openapi_snapshot.json | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index 6893c06d2d7..0ff155eee8e 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -19346,7 +19346,7 @@ } } }, - "description": "\nUnified rate-limit error.\n\nEvery rate-limit condition surfaced by litellm \u2014 whether it originated from\nan upstream LLM provider, a vendor batch endpoint, or one of litellm's own\nproxy-side limiters (parallel-requests, dynamic-rate, batch-rate, budget,\nmax-iterations, etc.) \u2014 is raised as an instance of this class.\n\nThe :attr:`category` attribute lets callers distinguish the source. See\n:class:`RateLimitErrorCategory` for the available values.\n" + "description": "\n Unified rate-limit error.\n\n Every rate-limit condition surfaced by litellm \u2014 whether it originated from\n an upstream LLM provider, a vendor batch endpoint, or one of litellm's own\n proxy-side limiters (parallel-requests, dynamic-rate, batch-rate, budget,\n max-iterations, etc.) \u2014 is raised as an instance of this class.\n\n The :attr:`category` attribute lets callers distinguish the source. See\n :class:`RateLimitErrorCategory` for the available values.\n " }, "500": { "content": { From 7cb01cf47f67e91137b4457e5c107104b10055c2 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 16 Sep 2026 13:53:10 -0700 Subject: [PATCH 015/144] fix(proxy): exchange role default, logged rejections, gated grant listing The exchange refused the very user JWT auth upserts (its row has no user_role) as "no longer active". The credential now carries the role the proxy already enforces for that user on every request, internal_user when the row has none, the same rule _get_user_role applies on the data plane. A rejected subject_token no longer echoes JWT auth's wording on the public /token endpoint: the response is a fixed invalid_request and the reason goes to the proxy log, since that wording can name the JWKS URL or relay the IdP's reply. The exchange grant is listed on /register, /.well-known/litellm-cli-auth, and the aggregate authorization-server metadata only when JWT auth is on, backed by a database, and licensed, so a client never selects a grant the gateway would then refuse. --- .../mcp_server/discoverable_endpoints.py | 17 ++-- .../mcp_server/gateway_dcr_flow.py | 21 +++- .../mcp_server/idp_token_exchange.py | 95 ++++++++++++++----- .../mcp_server/proxy_api_credentials.py | 12 ++- litellm/proxy/auth/auth_checks.py | 18 ++-- .../mcp_server/test_discoverable_endpoints.py | 19 ++++ .../mcp_server/test_gateway_dcr_flow.py | 27 +++++- .../mcp_server/test_idp_token_exchange.py | 74 +++++++++++---- .../mcp_server/test_proxy_api_credentials.py | 20 +++- 9 files changed, 225 insertions(+), 78 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index ebd5f43bf02..873c3d1baec 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -46,7 +46,6 @@ from litellm.proxy._experimental.mcp_server.faults import ( render_token_fault, ) from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import ( - TOKEN_EXCHANGE_GRANT_TYPE, VendorCredentialState, aggregate_authorize, aggregate_token, @@ -60,9 +59,11 @@ from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import ( register_aggregate_client, relative_request_url, revoke_refresh_token, + supported_grant_types, ) from litellm.proxy._experimental.mcp_server.idp_token_exchange import ( exchange_idp_subject_token, + token_exchange_available, ) from litellm.proxy._experimental.mcp_server.oauth_identity_binding import ( RefreshOwnershipProven, @@ -2142,7 +2143,9 @@ async def introspect_endpoint(token: str = Form(...)) -> Response: async def native_client_auth_discovery(request: Request) -> JSONResponse: """The versioned contract a native client (``lite login --pkce``, or a CLI in any other language) reads to sign a user in through the browser and obtain a proxy credential.""" - return JSONResponse(native_client_auth_contract(request), headers=TOKEN_NO_CACHE_HEADERS) + return JSONResponse( + native_client_auth_contract(request, token_exchange_available()), headers=TOKEN_NO_CACHE_HEADERS + ) # Per RFC 6749 §4.1.2.1, an IdP that rejects an OAuth authorization request @@ -2630,7 +2633,7 @@ def _build_aggregate_protected_resource_response(request: Request) -> dict: } -def _build_aggregate_authorization_server_response(request: Request) -> dict: +def _build_aggregate_authorization_server_response(request: Request, token_exchange_available: bool) -> dict: """RFC 8414 metadata for the gateway as the aggregate authorization server. The issuer is ``{base}/mcp`` and must stay equal to the value the @@ -2649,7 +2652,7 @@ def _build_aggregate_authorization_server_response(request: Request) -> dict: "registration_endpoint": f"{request_base_url}/register", "response_types_supported": ["code"], "scopes_supported": [], - "grant_types_supported": ("authorization_code", "refresh_token", TOKEN_EXCHANGE_GRANT_TYPE), + "grant_types_supported": supported_grant_types(token_exchange_available), "code_challenge_methods_supported": ["S256"], "token_endpoint_auth_methods_supported": ["none", "client_secret_post"], } @@ -2687,7 +2690,7 @@ async def oauth_authorization_server_aggregate(request: Request): per-server row win here instead would serve an issuer of {base} against a resource that advertised {base}/mcp, which fails the RFC 8414 issuer check and breaks the front door. """ - return _build_aggregate_authorization_server_response(request) + return _build_aggregate_authorization_server_response(request, token_exchange_available()) # Standard MCP pattern: /.well-known/oauth-protected-resource/mcp/{server_name} @@ -2913,7 +2916,9 @@ async def register_client(request: Request, mcp_server_name: str | None = None): # advertises that), so this does not affect it. A request without redirect_uris is not # a DCR request, so the legacy single-server-or-dummy fallback is kept for it. if data.get("redirect_uris"): - return await register_aggregate_client(request=request, request_body=data) + return await register_aggregate_client( + request=request, request_body=data, token_exchange_available=token_exchange_available() + ) resolved: Final = _resolve_oauth2_server_for_root_endpoints(client_ip=client_ip) if resolved: return await register_client_with_server( diff --git a/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py b/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py index 9bdde3c5edc..ba24e861d6e 100644 --- a/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py +++ b/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py @@ -188,6 +188,17 @@ class MintProxyCredential(Protocol): TOKEN_EXCHANGE_GRANT_TYPE: Final = "urn:ietf:params:oauth:grant-type:token-exchange" + + +def supported_grant_types(token_exchange_available: bool) -> tuple[str, ...]: + """The grants ``/token`` can serve on this deployment. The RFC 8693 exchange is listed + only where the JWT auth that proves a subject token is on, backed by a database, and + licensed, so a client never selects a grant the gateway would then refuse.""" + if token_exchange_available: + return ("authorization_code", "refresh_token", TOKEN_EXCHANGE_GRANT_TYPE) + return ("authorization_code", "refresh_token") + + """RFC 8693: a native client that already holds a token from the customer's identity provider trades it for the proxy-API credential without a browser round trip.""" @@ -359,7 +370,9 @@ def open_gateway_dcr_client(client_id: str) -> GatewayDcrClient | None: return _open_sealed(client_id, GATEWAY_DCR_CLIENT_ID_PREFIX, GatewayDcrClient, _CLIENT_RECORD_DEBUG_KEY) -async def register_aggregate_client(request: Request, request_body: Mapping[str, object]) -> Response: +async def register_aggregate_client( + request: Request, request_body: Mapping[str, object], token_exchange_available: bool +) -> Response: """RFC 7591 dynamic registration against the gateway itself, statelessly. Only ``redirect_uris`` is authoritative; every client is registered as a public @@ -423,7 +436,7 @@ async def register_aggregate_client(request: Request, request_body: Mapping[str, "client_id_issued_at": int(now.timestamp()), "redirect_uris": list(raw_uris), "token_endpoint_auth_method": "none", - "grant_types": ["authorization_code", "refresh_token", TOKEN_EXCHANGE_GRANT_TYPE], + "grant_types": list(supported_grant_types(token_exchange_available)), "response_types": ["code"], }, ) @@ -621,7 +634,7 @@ class NativeClientAuthContract(TypedDict): revocation_endpoint_auth_methods_supported: ReadOnly[tuple[str, ...]] -def native_client_auth_contract(request: Request) -> NativeClientAuthContract: +def native_client_auth_contract(request: Request, token_exchange_available: bool) -> NativeClientAuthContract: """The versioned discovery document at ``/.well-known/litellm-cli-auth``: everything a native client (in any language) needs to run the sign-in without reading LiteLLM source. ``resource`` is the exact value to send as the RFC 8707 ``resource`` parameter @@ -636,7 +649,7 @@ def native_client_auth_contract(request: Request) -> NativeClientAuthContract: "revocation_endpoint": f"{base_url}/revoke", "resource": base_url, "response_types_supported": ("code",), - "grant_types_supported": ("authorization_code", "refresh_token", TOKEN_EXCHANGE_GRANT_TYPE), + "grant_types_supported": supported_grant_types(token_exchange_available), "code_challenge_methods_supported": ("S256",), "token_endpoint_auth_methods_supported": ("none",), "revocation_endpoint_auth_methods_supported": ("none",), diff --git a/litellm/proxy/_experimental/mcp_server/idp_token_exchange.py b/litellm/proxy/_experimental/mcp_server/idp_token_exchange.py index 69b81713810..cdefaf76d49 100644 --- a/litellm/proxy/_experimental/mcp_server/idp_token_exchange.py +++ b/litellm/proxy/_experimental/mcp_server/idp_token_exchange.py @@ -6,15 +6,69 @@ and the user and team sync it performs), so no browser round trip is needed.""" from __future__ import annotations from collections.abc import Awaitable, Callable, Mapping +from dataclasses import dataclass from typing import Final, Protocol from fastapi import HTTPException, Request +from litellm._logging import verbose_proxy_logger from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import SubjectIdentity, SubjectTokenRefusal from litellm.proxy._types import JWTAuthBuilderResult, ProxyException from litellm.proxy.auth.handle_jwt import JWTAuthManager EXCHANGE_ROUTE: Final = "/token" +REJECTED_SUBJECT_TOKEN: Final = "subject_token was rejected by the gateway's JWT auth" + + +@dataclass(frozen=True, slots=True) +class TokenExchangePrerequisites: + """The deployment-level gates ``user_api_key_auth`` applies before it verifies any JWT + bearer. Discovery and registration advertise the exchange grant only when every one of + them holds, and an exchange attempt is refused naming the first one that does not.""" + + jwt_auth_enabled: bool + has_database: bool + licensed: bool + + @property + def available(self) -> bool: + return self.jwt_auth_enabled and self.has_database and self.licensed + + def refusal(self) -> SubjectTokenRefusal | None: + if not self.jwt_auth_enabled: + return SubjectTokenRefusal( + error="unsupported_grant_type", + description="JWT auth is not enabled on this gateway, so it cannot exchange IdP tokens", + ) + if not self.has_database: + return SubjectTokenRefusal( + error="unsupported_grant_type", + description="this gateway has no database, so it cannot exchange IdP tokens", + ) + if not self.licensed: + return SubjectTokenRefusal( + error="unsupported_grant_type", + description="JWT auth is an enterprise only feature; no license is set", + ) + return None + + +def read_token_exchange_prerequisites() -> TokenExchangePrerequisites: + from litellm.proxy.proxy_server import ( # noqa: PLC0415 # rebound after startup, so read them per call + general_settings, + premium_user, + prisma_client, + ) + + return TokenExchangePrerequisites( + jwt_auth_enabled=general_settings.get("enable_jwt_auth", False) is True, + has_database=prisma_client is not None, + licensed=premium_user is True, + ) + + +def token_exchange_available() -> bool: + return read_token_exchange_prerequisites().available class AuthorizeSubjectToken(Protocol): @@ -31,7 +85,6 @@ async def exchange_idp_subject_token(subject_token: str, request: Request) -> Su from litellm.proxy.proxy_server import ( # noqa: PLC0415 # rebound after startup, so read them per call general_settings, jwt_handler, - premium_user, prisma_client, proxy_logging_obj, user_api_key_cache, @@ -55,9 +108,7 @@ async def exchange_idp_subject_token(subject_token: str, request: Request) -> Su return await identity_from_subject_token( subject_token, request_headers=request.headers, - jwt_auth_enabled=general_settings.get("enable_jwt_auth", False) is True, - has_database=prisma_client is not None, - licensed=premium_user is True, + prerequisites=read_token_exchange_prerequisites(), is_jwt=jwt_handler.is_jwt, authorize=authorize, ) @@ -66,40 +117,34 @@ async def exchange_idp_subject_token(subject_token: str, request: Request) -> Su async def identity_from_subject_token( subject_token: str, request_headers: Mapping[str, str], - jwt_auth_enabled: bool, - has_database: bool, - licensed: bool, + prerequisites: TokenExchangePrerequisites, is_jwt: Callable[[str], bool], authorize: AuthorizeSubjectToken, ) -> SubjectIdentity | SubjectTokenRefusal: """Apply the same gates ``user_api_key_auth`` applies to a JWT bearer, then let the proxy's JWT auth prove the token. A rejection comes back as ``invalid_request``, which - RFC 8693 section 2.2.2 prescribes for an invalid or unacceptable subject token.""" - if not jwt_auth_enabled: - return SubjectTokenRefusal( - error="unsupported_grant_type", - description="JWT auth is not enabled on this gateway, so it cannot exchange IdP tokens", - ) - if not has_database: - return SubjectTokenRefusal( - error="unsupported_grant_type", - description="this gateway has no database, so it cannot exchange IdP tokens", - ) + RFC 8693 section 2.2.2 prescribes for an invalid or unacceptable subject token. The + reason stays in the proxy log: this endpoint is public and JWT auth's own wording can + name the JWKS URL it fetched or quote the IdP's response.""" + unmet: Final = prerequisites.refusal() + if unmet is not None: + return unmet if not is_jwt(subject_token): return SubjectTokenRefusal(error="invalid_request", description="subject_token is not a JWT") - if not licensed: - return SubjectTokenRefusal( - error="unsupported_grant_type", description="JWT auth is an enterprise only feature; no license is set" - ) try: result: Final = await authorize(subject_token, request_headers) except HTTPException as denied: - return SubjectTokenRefusal(error="invalid_request", description=f"subject_token was rejected: {denied.detail}") + return _rejected_by_jwt_auth(denied.detail) except ProxyException as denied: - return SubjectTokenRefusal(error="invalid_request", description=f"subject_token was rejected: {denied.message}") + return _rejected_by_jwt_auth(denied.message) except Exception as denied: # noqa: BLE001 # auth_jwt raises a plain Exception on signature and claim failures - return SubjectTokenRefusal(error="invalid_request", description=f"subject_token was rejected: {denied}") + return _rejected_by_jwt_auth(denied) user_id: Final = result["user_id"] if user_id is None: return SubjectTokenRefusal(error="invalid_request", description="subject_token names no user the gateway knows") return SubjectIdentity(user_id=user_id, team_id=result["team_id"]) + + +def _rejected_by_jwt_auth(reason: object) -> SubjectTokenRefusal: + verbose_proxy_logger.warning("token exchange refused a subject_token: %s", reason) + return SubjectTokenRefusal(error="invalid_request", description=REJECTED_SUBJECT_TOKEN) diff --git a/litellm/proxy/_experimental/mcp_server/proxy_api_credentials.py b/litellm/proxy/_experimental/mcp_server/proxy_api_credentials.py index 27d0ebbd5e6..a34119edf10 100644 --- a/litellm/proxy/_experimental/mcp_server/proxy_api_credentials.py +++ b/litellm/proxy/_experimental/mcp_server/proxy_api_credentials.py @@ -16,7 +16,7 @@ from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import ( ReloadUserFailure, ) from litellm.proxy._types import LiteLLM_UserTable -from litellm.proxy.auth.auth_checks import ExperimentalUIJWTToken +from litellm.proxy.auth.auth_checks import ExperimentalUIJWTToken, effective_user_role from litellm.proxy.management_endpoints.ui_sso import ( CliSsoTeamDetail, fetch_cli_sso_team_details, @@ -51,12 +51,12 @@ async def mint_proxy_credential( posting the consent form without one. Memberships whose team rows are gone count as no team at all, the way ``lite login`` treats them, so they can never lock a user out. The user row handed to the minter carries no team list, exactly like ``lite login``'s, so - the minter's own first-team fallback stays inert.""" + the minter's own first-team fallback stays inert. The credential carries the role the + proxy already enforces for the user on every request, so a row with no role (JWT auth's + upsert writes none) mints as an internal user instead of being refused.""" user: Final = await load_active_user_by_id(user_id) if isinstance(user, str): return user - if user.user_role is None: - return "no_active_key" if team_id is not None and team_id not in user.teams: return "not_a_member" details: Final = await _team_details(user.teams) if user.teams else () @@ -68,7 +68,9 @@ async def mint_proxy_credential( if selected is None: return "not_a_member" key: Final = ExperimentalUIJWTToken.get_cli_jwt_auth_token( - user_info=LiteLLM_UserTable(user_id=user.user_id, user_role=user.user_role, models=user.models), + user_info=LiteLLM_UserTable( + user_id=user.user_id, user_role=effective_user_role(user.user_role).value, models=user.models + ), team_id=team_id, team_alias=selected.team_alias, team_models=selected.team_models, diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index ba68dc8a17f..d90553db72b 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -1207,21 +1207,19 @@ async def common_checks( return True +def effective_user_role(user_role: str | None) -> LitellmUserRoles: + try: + return LitellmUserRoles(user_role) + except ValueError: + return LitellmUserRoles.INTERNAL_USER + + def _get_user_role( user_obj: LiteLLM_UserTable | None, ) -> LitellmUserRoles | None: if user_obj is None: return None - - _user: Final = user_obj - - _user_role: Final = _user.user_role - try: - role: Final = LitellmUserRoles(_user_role) - except ValueError: - return LitellmUserRoles.INTERNAL_USER - - return role + return effective_user_role(user_obj.user_role) def _is_api_route_allowed( diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py index aa45b2f6793..e4edb5fb4dd 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py @@ -11048,6 +11048,25 @@ def test_native_client_login_walks_discovery_consent_token_refresh_and_revoke(mo assert stranger.json()["error"] == "invalid_client" +@pytest.mark.parametrize("exchange_servable", [True, False]) +def test_discovery_advertises_the_exchange_grant_only_where_the_gateway_can_serve_it(monkeypatch, exchange_servable): + """Every document a native client reads before it picks a grant (the versioned contract, the + aggregate authorization-server metadata, and the registration response) lists the RFC 8693 + exchange exactly when the running proxy can serve it: JWT auth on, a database, and a license.""" + client, _session_cookie, _minted = _native_client_app(monkeypatch) + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {"enable_jwt_auth": exchange_servable}) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", object()) + monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True) + exchange_grant = ["urn:ietf:params:oauth:grant-type:token-exchange"] if exchange_servable else [] + expected = ["authorization_code", "refresh_token", *exchange_grant] + + assert client.get("/.well-known/litellm-cli-auth").json()["grant_types_supported"] == expected + assert client.get("/.well-known/oauth-authorization-server/mcp").json()["grant_types_supported"] == expected + registered = client.post("/register", json={"redirect_uris": ["http://127.0.0.1:51234/callback"]}) + assert registered.status_code == 201 + assert registered.json()["grant_types"] == expected + + def test_native_client_authorize_without_the_proxy_resource_keeps_the_mcp_flow(monkeypatch): """A registered client asking for the MCP resource (or no resource) never sees the consent page, so existing MCP clients are untouched by the native-client arm.""" diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py index 1be6ebb6e22..c42f8763e74 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py @@ -95,9 +95,11 @@ def _request(path="/authorize", query="", cookies=None, method="GET"): ) -async def _register(redirect_uris) -> dict: +async def _register(redirect_uris, token_exchange_available=True) -> dict: response = await register_aggregate_client( - request=_request(path="/register", method="POST"), request_body={"redirect_uris": redirect_uris} + request=_request(path="/register", method="POST"), + request_body={"redirect_uris": redirect_uris}, + token_exchange_available=token_exchange_available, ) return json.loads(response.body) @@ -119,11 +121,18 @@ async def test_register_mints_stateless_public_client(): assert record.redirect_uris == (REDIRECT_URI,) +@pytest.mark.asyncio +async def test_register_omits_the_exchange_grant_where_the_gateway_cannot_serve_it(): + body = await _register([REDIRECT_URI], token_exchange_available=False) + assert body["grant_types"] == ["authorization_code", "refresh_token"] + + @pytest.mark.asyncio @pytest.mark.parametrize("redirect_uris", [VSCODE_REDIRECT_URIS, MAX_LENGTH_REDIRECT_URIS]) async def test_register_four_callbacks_preserves_metadata(redirect_uris: tuple[str, ...]) -> None: response: Final = await register_aggregate_client( request=_request(path="/register", method="POST"), + token_exchange_available=True, request_body={ "client_name": "Visual Studio Code", "client_uri": "https://code.visualstudio.com", @@ -149,6 +158,7 @@ async def test_register_four_callbacks_preserves_metadata(redirect_uris: tuple[s async def test_register_rejects_five_valid_callbacks() -> None: response: Final = await register_aggregate_client( request=_request(path="/register", method="POST"), + token_exchange_available=True, request_body={"redirect_uris": [*VSCODE_REDIRECT_URIS, "http://127.0.0.1:33419/"]}, ) assert response.status_code == 400 @@ -162,6 +172,7 @@ async def test_register_rejects_five_valid_callbacks() -> None: async def test_register_four_callbacks_preserves_encoded_size_guard() -> None: response: Final = await register_aggregate_client( request=_request(path="/register", method="POST"), + token_exchange_available=True, request_body={"redirect_uris": [f"https://client.example/{index}/".ljust(256, "é") for index in range(4)]}, ) assert response.status_code == 400 @@ -214,6 +225,7 @@ async def test_register_rejects_userinfo_spoofed_origin(): response = await register_aggregate_client( request=_request(path="/register", method="POST"), request_body={"redirect_uris": ["https://claude.ai@attacker.example/callback"]}, + token_exchange_available=True, ) assert response.status_code == 400 assert json.loads(response.body)["error"] == "invalid_redirect_uri" @@ -234,7 +246,9 @@ async def test_register_rejects_userinfo_spoofed_origin(): ) async def test_register_rejects_bad_redirect_uris(redirect_uris): response = await register_aggregate_client( - request=_request(path="/register", method="POST"), request_body={"redirect_uris": redirect_uris} + request=_request(path="/register", method="POST"), + request_body={"redirect_uris": redirect_uris}, + token_exchange_available=True, ) assert response.status_code == 400 assert json.loads(response.body)["error"] in ("invalid_redirect_uri", "invalid_client_metadata") @@ -1954,7 +1968,7 @@ async def test_revoke_refuses_unknown_clients_and_a_missing_master_key(): def test_native_client_auth_contract_points_every_endpoint_at_this_proxy(): - assert json.loads(json.dumps(native_client_auth_contract(_request("/.well-known/litellm-cli-auth")))) == { + assert json.loads(json.dumps(native_client_auth_contract(_request("/.well-known/litellm-cli-auth"), True))) == { "contract_version": 1, "issuer": "https://llm.example.com", "authorization_endpoint": "https://llm.example.com/authorize", @@ -1974,6 +1988,11 @@ def test_native_client_auth_contract_points_every_endpoint_at_this_proxy(): } +def test_native_client_auth_contract_omits_the_exchange_grant_where_the_gateway_cannot_serve_it(): + contract = native_client_auth_contract(_request("/.well-known/litellm-cli-auth"), False) + assert list(contract["grant_types_supported"]) == ["authorization_code", "refresh_token"] + + @pytest.mark.parametrize( "resource, expected", [ diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_idp_token_exchange.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_idp_token_exchange.py index 00440614120..d1b049dddd5 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_idp_token_exchange.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_idp_token_exchange.py @@ -1,13 +1,22 @@ +import logging + import pytest from fastapi import HTTPException from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import SubjectIdentity, SubjectTokenRefusal -from litellm.proxy._experimental.mcp_server.idp_token_exchange import identity_from_subject_token +from litellm.proxy._experimental.mcp_server.idp_token_exchange import ( + REJECTED_SUBJECT_TOKEN, + TokenExchangePrerequisites, + identity_from_subject_token, + token_exchange_available, +) from litellm.proxy._types import ProxyException from litellm.proxy.auth.handle_jwt import JWTHandler IDP_JWT = "eyJhbGciOiJSUzI1NiJ9.eyJzdWIiOiJ1MSJ9.idp-signature" REQUEST_HEADERS = {"x-litellm-team-id": "team-b", "user-agent": "lite/0.1"} +EVERY_GATE_HOLDS = {"jwt_auth_enabled": True, "has_database": True, "licensed": True} +JWKS_URL = "https://idp.example.com/.well-known/jwks.json" def _authorized(user_id="u1", team_id="team-b"): @@ -42,16 +51,14 @@ class _Authorizer: return self.result -async def _identity(authorizer, subject_token=IDP_JWT, **overrides): - arguments = { - "request_headers": REQUEST_HEADERS, - "jwt_auth_enabled": True, - "has_database": True, - "licensed": True, - "is_jwt": JWTHandler.is_jwt, - "authorize": authorizer, - } - return await identity_from_subject_token(subject_token, **{**arguments, **overrides}) +async def _identity(authorizer, subject_token=IDP_JWT, **unmet): + return await identity_from_subject_token( + subject_token, + request_headers=REQUEST_HEADERS, + prerequisites=TokenExchangePrerequisites(**{**EVERY_GATE_HOLDS, **unmet}), + is_jwt=JWTHandler.is_jwt, + authorize=authorizer, + ) @pytest.mark.asyncio @@ -70,7 +77,7 @@ async def test_a_jwt_that_resolves_no_team_names_a_teamless_identity(): @pytest.mark.asyncio @pytest.mark.parametrize( - "overrides, subject_token, error, mentions", + "unmet, subject_token, error, mentions", [ ({"jwt_auth_enabled": False}, IDP_JWT, "unsupported_grant_type", "JWT auth is not enabled"), ({"has_database": False}, IDP_JWT, "unsupported_grant_type", "no database"), @@ -79,32 +86,59 @@ async def test_a_jwt_that_resolves_no_team_names_a_teamless_identity(): ], ) async def test_the_gates_user_api_key_auth_applies_refuse_before_any_verification( - overrides, subject_token, error, mentions + unmet, subject_token, error, mentions ): authorizer = _Authorizer() - refusal = await _identity(authorizer, subject_token=subject_token, **overrides) + refusal = await _identity(authorizer, subject_token=subject_token, **unmet) assert isinstance(refusal, SubjectTokenRefusal) assert refusal.error == error assert mentions in refusal.description assert authorizer.calls == [] +@pytest.mark.parametrize("unmet", [{}, {"jwt_auth_enabled": False}, {"has_database": False}, {"licensed": False}]) +def test_the_grant_is_available_exactly_when_every_gate_holds(unmet): + prerequisites = TokenExchangePrerequisites(**{**EVERY_GATE_HOLDS, **unmet}) + assert prerequisites.available is (unmet == {}) + assert (prerequisites.refusal() is None) is prerequisites.available + + +@pytest.mark.parametrize( + "general_settings, prisma_client, premium_user, expected", + [ + ({"enable_jwt_auth": True}, object(), True, True), + ({}, object(), True, False), + ({"enable_jwt_auth": True}, None, True, False), + ({"enable_jwt_auth": True}, object(), False, False), + ], +) +def test_availability_is_read_from_the_running_proxy( + monkeypatch, general_settings, prisma_client, premium_user, expected +): + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", general_settings) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", prisma_client) + monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", premium_user) + assert token_exchange_available() is expected + + @pytest.mark.asyncio @pytest.mark.parametrize( - "raised, mentions", + "raised, reason", [ (HTTPException(status_code=403, detail="User not allowed to access this route"), "not allowed"), (ProxyException(message="Token expired", type="auth_error", param="token", code=401), "Token expired"), (Exception("Validation fails: signature verification failed"), "signature verification failed"), (Exception("Invalid JWT Submitted"), "Invalid JWT"), + (Exception(f"Failed to fetch keys from {JWKS_URL}: 502 Bad Gateway from the IdP"), JWKS_URL), ], ) -async def test_a_jwt_the_proxy_rejects_is_an_invalid_subject_token(raised, mentions): +async def test_a_jwt_the_proxy_rejects_is_refused_with_the_reason_kept_in_the_log(raised, reason, caplog): + """The endpoint is public, so the response never quotes JWT auth's wording (it can name + the JWKS URL or relay the IdP's reply); the operator reads the reason in the proxy log.""" + caplog.set_level(logging.WARNING, logger="LiteLLM Proxy") refusal = await _identity(_Authorizer(raises=raised)) - assert isinstance(refusal, SubjectTokenRefusal) - assert refusal.error == "invalid_request" - assert refusal.description.startswith("subject_token was rejected: ") - assert mentions in refusal.description + assert refusal == SubjectTokenRefusal(error="invalid_request", description=REJECTED_SUBJECT_TOKEN) + assert reason in caplog.text @pytest.mark.asyncio diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_proxy_api_credentials.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_proxy_api_credentials.py index 8bb8bdada7d..85650c6a05a 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_proxy_api_credentials.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_proxy_api_credentials.py @@ -8,6 +8,7 @@ from litellm.constants import CLI_JWT_EXPIRATION_HOURS from litellm.models.user import LiteLLM_UserTable from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import ConsentTeam, MintedProxyCredential from litellm.proxy._experimental.mcp_server.proxy_api_credentials import lookup_consent_teams, mint_proxy_credential +from litellm.proxy._types import LitellmUserRoles from litellm.proxy.auth.auth_checks import ExperimentalUIJWTToken from litellm.proxy.management_endpoints.ui_sso import CliSsoTeamDetail @@ -67,10 +68,21 @@ async def test_mint_passes_user_lookup_failures_through(failure, load_user, fetc @pytest.mark.asyncio -async def test_mint_refuses_a_user_without_a_role(load_user, fetch_teams): - load_user.return_value = _user(user_role=None) - assert await mint_proxy_credential("u1", None) == "no_active_key" - fetch_teams.assert_not_awaited() +@pytest.mark.parametrize( + "stored_role, minted_role", + [ + (None, LitellmUserRoles.INTERNAL_USER), + ("made_up_role", LitellmUserRoles.INTERNAL_USER), + ("proxy_admin", LitellmUserRoles.PROXY_ADMIN), + ], +) +async def test_mint_carries_the_role_the_proxy_enforces_for_the_user(load_user, fetch_teams, stored_role, minted_role): + """A user JWT auth upserted has no role in the database, and the proxy already treats + such a user as an internal user on every request, so the credential says the same.""" + load_user.return_value = _user(user_role=stored_role) + minted = await mint_proxy_credential("u1", "team-a") + assert isinstance(minted, MintedProxyCredential) + assert _decoded(minted).user_role == minted_role @pytest.mark.asyncio From 3a8ac47e99b1383a8f3231e2ef88494e669a069c Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 16 Sep 2026 14:12:15 -0700 Subject: [PATCH 016/144] fix(proxy): mint the exchange credential off the database user row, not the cache JWT auth caches the user it creates before it adds that user to the JWT's team, and adding a team member never evicts the cached user row, so the mint read a row with no teams and refused the very first token exchange for a never-seen user as not a member. The loader now reads the row from the database and leaves the fresh row in the cache for the requests the credential makes next --- .../mcp_server/bridge_token_flow.py | 7 +++- .../mcp_server/test_discoverable_endpoints.py | 42 +++++++++++++++++-- 2 files changed, 44 insertions(+), 5 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/bridge_token_flow.py b/litellm/proxy/_experimental/mcp_server/bridge_token_flow.py index 2b13baa624b..4235471f2d9 100644 --- a/litellm/proxy/_experimental/mcp_server/bridge_token_flow.py +++ b/litellm/proxy/_experimental/mcp_server/bridge_token_flow.py @@ -273,7 +273,11 @@ async def load_active_user_by_id(user_id: str) -> "LiteLLM_UserTable | _KeyResol ``HTTPException``, a SCIM-deactivated user, and, unlike the key path, a missing user. ``get_user_object`` catches every DB failure and re-raises a bare ``ValueError`` (a deleted user and a real outage look identical, the original error surviving only as ``__context__``), so the outage check walks the cause - chain, and a missing user falls through to ``no_active_key`` rather than an opaque gateway fault.""" + chain, and a missing user falls through to ``no_active_key`` rather than an opaque gateway fault. The + row is read from the database, never the cache: JWT auth caches the user it creates before it adds + that user to the JWT's team and adding a member never evicts the cached row, so a credential minted + off the cache would refuse the very first exchange as not a member. The fresh row replaces the cached + one.""" from litellm.proxy._types import ( ProxyException, # noqa: PLC0415 # inline import avoids a module-load circular import ) @@ -296,6 +300,7 @@ async def load_active_user_by_id(user_id: str) -> "LiteLLM_UserTable | _KeyResol prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, user_id_upsert=False, + check_db_only=True, ) except (ProxyException, HTTPException): return "no_active_key" diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py index e4edb5fb4dd..30b179f1a26 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py @@ -7568,6 +7568,36 @@ async def test_reload_active_user_by_id_permanent_engine_fault_is_faulted(proxy_ assert await _reload_active_user_by_id("sso-user-7") == "faulted" +@pytest.mark.asyncio +async def test_load_active_user_by_id_reads_the_row_from_the_database_not_the_cache(proxy_globals): + """JWT auth caches the user it creates before it adds that user to the JWT's team, and adding a + member never evicts the cached row, so a credential minted off the cached row refused the very first + token exchange as not a member. The loader has to read the row from the database and leave the fresh + row in the cache for the requests the credential makes next.""" + from litellm.proxy._experimental.mcp_server.bridge_token_flow import load_active_user_by_id + from litellm.proxy._types import LiteLLM_UserTable + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + cache = UserApiKeyCache() + await cache.async_set_cache( + key="fresh-jwt-user", value=LiteLLM_UserTable(user_id="fresh-jwt-user", teams=[]), model_type=LiteLLM_UserTable + ) + prisma = MagicMock() + prisma.db.litellm_usertable.find_unique = AsyncMock( + return_value=LiteLLM_UserTable(user_id="fresh-jwt-user", teams=["team-a"]) + ) + proxy_globals.user_api_key_cache = cache + proxy_globals.prisma_client = prisma + + loaded = await load_active_user_by_id("fresh-jwt-user") + + assert not isinstance(loaded, str) + assert loaded.teams == ["team-a"] + cached = await cache.async_get_cache(key="fresh-jwt-user", model_type=LiteLLM_UserTable) + assert cached is not None + assert cached.teams == ["team-a"] + + @pytest.mark.asyncio async def test_token_endpoint_uses_client_secret_basic_when_configured(): """LIT-4091: a server with token_endpoint_auth_method=client_secret_basic must send the @@ -11866,13 +11896,17 @@ async def test_oauth_refresh_revalidates_the_same_active_user_rule( from litellm.proxy._experimental.mcp_server.bridge_token_flow import _reload_active_user_by_id handler, _ = jwt_oauth_identity - handler.user_api_key_cache.set_cache( - "jwt-owner", LiteLLM_UserTable(user_id="jwt-owner", metadata={"scim_active": state != "inactive"}) - ) + user_id: Final = f"jwt-owner-{state}" + row: Final = LiteLLM_UserTable(user_id=user_id, metadata={"scim_active": state != "inactive"}) + proxy_server.prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=row) if state == "missing_database": monkeypatch.setattr(proxy_server, "prisma_client", None) expected: Final = None if state == "active" else "no_active_key" if state == "inactive" else "unresolvable" - assert await _reload_active_user_by_id("jwt-owner") == expected + assert await _reload_active_user_by_id(user_id) == expected + if state != "missing_database": + cached: Final = handler.user_api_key_cache.get_cache(user_id, model_type=LiteLLM_UserTable) + assert cached is not None + assert cached.metadata == row.metadata @pytest.mark.asyncio From 167edf2769b28f76870b654014f99312843cc327 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 16 Sep 2026 15:10:03 -0700 Subject: [PATCH 017/144] fix(proxy): read the database user row only in the credential mint The token exchange mint keeps reading the user row from the database, since JWT auth caches the user it creates before adding it to the JWT's team and a mint off that cached row refused the first exchange for a new user. Introspection and the refresh revalidation go back to the cache read, so a resource server calling /introspect per request pays no database read. --- .../mcp_server/bridge_token_flow.py | 18 ++++++--- .../mcp_server/proxy_api_credentials.py | 4 +- .../mcp_server/test_discoverable_endpoints.py | 39 +++++++++++++++++-- .../mcp_server/test_proxy_api_credentials.py | 30 +++++++++++++- 4 files changed, 78 insertions(+), 13 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/bridge_token_flow.py b/litellm/proxy/_experimental/mcp_server/bridge_token_flow.py index 4235471f2d9..f19cb87ae18 100644 --- a/litellm/proxy/_experimental/mcp_server/bridge_token_flow.py +++ b/litellm/proxy/_experimental/mcp_server/bridge_token_flow.py @@ -262,7 +262,12 @@ async def _reload_active_user_by_id(user_id: str) -> "_KeyResolutionFailure | No return loaded if isinstance(loaded, str) else None -async def load_active_user_by_id(user_id: str) -> "LiteLLM_UserTable | _KeyResolutionFailure": +UserRowSource = Literal["cache", "database"] + + +async def load_active_user_by_id( + user_id: str, source: UserRowSource = "cache" +) -> "LiteLLM_UserTable | _KeyResolutionFailure": """Load a live litellm user by id, returning the record when the user is active or a precise failure otherwise. The interactive DCR client authenticates via SSO, so its refresh envelope seals a user subject; renewing it must re-check the user is still live (present and not SCIM-deactivated) so a @@ -273,11 +278,12 @@ async def load_active_user_by_id(user_id: str) -> "LiteLLM_UserTable | _KeyResol ``HTTPException``, a SCIM-deactivated user, and, unlike the key path, a missing user. ``get_user_object`` catches every DB failure and re-raises a bare ``ValueError`` (a deleted user and a real outage look identical, the original error surviving only as ``__context__``), so the outage check walks the cause - chain, and a missing user falls through to ``no_active_key`` rather than an opaque gateway fault. The - row is read from the database, never the cache: JWT auth caches the user it creates before it adds + chain, and a missing user falls through to ``no_active_key`` rather than an opaque gateway fault. + ``source="database"`` reads the row from the database, never the cache, and leaves the fresh row in the + cache for the requests the credential makes next: JWT auth caches the user it creates before it adds that user to the JWT's team and adding a member never evicts the cached row, so a credential minted - off the cache would refuse the very first exchange as not a member. The fresh row replaces the cached - one.""" + off the cache would refuse the very first exchange as not a member. Every other caller keeps the cache + read, so introspection, which a resource server may call per request, stays off the database.""" from litellm.proxy._types import ( ProxyException, # noqa: PLC0415 # inline import avoids a module-load circular import ) @@ -300,7 +306,7 @@ async def load_active_user_by_id(user_id: str) -> "LiteLLM_UserTable | _KeyResol prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, user_id_upsert=False, - check_db_only=True, + check_db_only=source == "database", ) except (ProxyException, HTTPException): return "no_active_key" diff --git a/litellm/proxy/_experimental/mcp_server/proxy_api_credentials.py b/litellm/proxy/_experimental/mcp_server/proxy_api_credentials.py index a34119edf10..2f7fcaef645 100644 --- a/litellm/proxy/_experimental/mcp_server/proxy_api_credentials.py +++ b/litellm/proxy/_experimental/mcp_server/proxy_api_credentials.py @@ -42,7 +42,7 @@ async def mint_proxy_credential( user_id: str, team_id: str | None ) -> MintedProxyCredential | ProxyCredentialMintFailure: """Mint the ``lite login`` credential for a consented grant. Membership is checked - live, so a team the user left between consent and redemption (or between refreshes) + live against the database row, so a team the user left between consent and redemption (or between refreshes) refuses the grant instead of minting a credential attributed to a team they are no longer on. The team is exactly the one the consent page sealed into the grant; nothing is picked on the user's behalf here, so a refresh can never move the credential, and a @@ -54,7 +54,7 @@ async def mint_proxy_credential( the minter's own first-team fallback stays inert. The credential carries the role the proxy already enforces for the user on every request, so a row with no role (JWT auth's upsert writes none) mints as an internal user instead of being refused.""" - user: Final = await load_active_user_by_id(user_id) + user: Final = await load_active_user_by_id(user_id, source="database") if isinstance(user, str): return user if team_id is not None and team_id not in user.teams: diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py index 30b179f1a26..6965b3b4ebe 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py @@ -7572,8 +7572,8 @@ async def test_reload_active_user_by_id_permanent_engine_fault_is_faulted(proxy_ async def test_load_active_user_by_id_reads_the_row_from_the_database_not_the_cache(proxy_globals): """JWT auth caches the user it creates before it adds that user to the JWT's team, and adding a member never evicts the cached row, so a credential minted off the cached row refused the very first - token exchange as not a member. The loader has to read the row from the database and leave the fresh - row in the cache for the requests the credential makes next.""" + token exchange as not a member. The database source has to read the row from the database and leave + the fresh row in the cache for the requests the credential makes next.""" from litellm.proxy._experimental.mcp_server.bridge_token_flow import load_active_user_by_id from litellm.proxy._types import LiteLLM_UserTable from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache @@ -7589,7 +7589,7 @@ async def test_load_active_user_by_id_reads_the_row_from_the_database_not_the_ca proxy_globals.user_api_key_cache = cache proxy_globals.prisma_client = prisma - loaded = await load_active_user_by_id("fresh-jwt-user") + loaded = await load_active_user_by_id("fresh-jwt-user", source="database") assert not isinstance(loaded, str) assert loaded.teams == ["team-a"] @@ -7598,6 +7598,39 @@ async def test_load_active_user_by_id_reads_the_row_from_the_database_not_the_ca assert cached.teams == ["team-a"] +@pytest.mark.asyncio +async def test_load_active_user_by_id_serves_a_cached_row_without_a_database_read(proxy_globals): + """Introspection and refresh revalidation run per call, so the loader's default source is the cache: a + cached row answers without a database read, and only a caller that asks for the database row pays for + one.""" + from litellm.proxy._experimental.mcp_server.bridge_token_flow import ( + _reload_active_user_by_id, + load_active_user_by_id, + ) + from litellm.proxy._types import LiteLLM_UserTable + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + cache = UserApiKeyCache() + await cache.async_set_cache( + key="cached-jwt-user", + value=LiteLLM_UserTable(user_id="cached-jwt-user", teams=["team-a"]), + model_type=LiteLLM_UserTable, + ) + prisma = MagicMock() + prisma.db.litellm_usertable.find_unique = AsyncMock( + return_value=LiteLLM_UserTable(user_id="cached-jwt-user", teams=[]) + ) + proxy_globals.user_api_key_cache = cache + proxy_globals.prisma_client = prisma + + loaded = await load_active_user_by_id("cached-jwt-user") + + assert not isinstance(loaded, str) + assert loaded.teams == ["team-a"] + assert await _reload_active_user_by_id("cached-jwt-user") is None + prisma.db.litellm_usertable.find_unique.assert_not_awaited() + + @pytest.mark.asyncio async def test_token_endpoint_uses_client_secret_basic_when_configured(): """LIT-4091: a server with token_endpoint_auth_method=client_secret_basic must send the diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_proxy_api_credentials.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_proxy_api_credentials.py index 85650c6a05a..04fbbe4a6ce 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_proxy_api_credentials.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_proxy_api_credentials.py @@ -1,6 +1,6 @@ """Tests for minting the ``lite login`` credential from a consented native-client grant.""" -from unittest.mock import ANY, AsyncMock +from unittest.mock import ANY, AsyncMock, MagicMock import pytest @@ -10,6 +10,7 @@ from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import ConsentTeam, from litellm.proxy._experimental.mcp_server.proxy_api_credentials import lookup_consent_teams, mint_proxy_credential from litellm.proxy._types import LitellmUserRoles from litellm.proxy.auth.auth_checks import ExperimentalUIJWTToken +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.proxy.management_endpoints.ui_sso import CliSsoTeamDetail _LOAD_USER = "litellm.proxy._experimental.mcp_server.proxy_api_credentials.load_active_user_by_id" @@ -91,7 +92,7 @@ async def test_mint_refuses_a_teamless_grant_for_a_team_member(load_user, fetch_ is refused for a user with teams instead of minting an unscoped credential or drifting onto the first team, on redemption and on every refresh alike.""" assert await mint_proxy_credential("u1", None) == "team_required" - load_user.assert_awaited_once_with("u1") + load_user.assert_awaited_once_with("u1", source="database") fetch_teams.assert_awaited_once_with(ANY, ["team-a", "team-b"]) @@ -126,6 +127,31 @@ async def test_mint_honors_the_consented_team(load_user, fetch_teams): assert decoded.team_model_aliases == {"fast": "gpt-5.4-mini"} +@pytest.mark.asyncio +async def test_mint_reads_the_users_teams_from_the_database_not_a_stale_cached_row(fetch_teams, monkeypatch): + """JWT auth caches the user it creates before it adds that user to the JWT's team, and adding a member + never evicts the cached row, so a mint off the cached row refused the very first token exchange as not + a member. The mint has to read the database row, whatever the cache holds.""" + from litellm.proxy import proxy_server + + cache = UserApiKeyCache() + await cache.async_set_cache( + key="stale-cache-user", value=_user(user_id="stale-cache-user", teams=[]), model_type=LiteLLM_UserTable + ) + prisma = MagicMock() + prisma.db.litellm_usertable.find_unique = AsyncMock( + return_value=_user(user_id="stale-cache-user", teams=["team-a"]) + ) + monkeypatch.setattr(proxy_server, "user_api_key_cache", cache) + monkeypatch.setattr(proxy_server, "prisma_client", prisma) + + minted = await mint_proxy_credential("stale-cache-user", "team-a") + + assert isinstance(minted, MintedProxyCredential) + assert minted.team_id == "team-a" + assert _decoded(minted).team_id == "team-a" + + @pytest.mark.asyncio async def test_mint_refuses_a_team_the_user_is_not_on(load_user, fetch_teams): assert await mint_proxy_credential("u1", "team-c") == "not_a_member" From ce722ab1b30d4b1331504364eea7adb362887c38 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 16 Sep 2026 16:31:28 -0700 Subject: [PATCH 018/144] fix(proxy): evict the member's cached user row on team member add JWT auth caches the user row before it adds the user to the JWT's team, and admission checks the credential's team against the cached row on whichever worker takes the next request. On a two-worker gateway the credential minted for a newly joined team answered 403 "not in your team memberships" until the management-object TTL ran out, because /team/member_add only evicted the membership spend sentinel. The add now evicts the added members' cached user rows and broadcasts the eviction to the other workers, the way /team/member_delete already did The mint test now also covers a user SCIM deactivated after the cache last saw them active: the database read refuses the mint while the cached row still says active --- .../mcp_server/bridge_token_flow.py | 9 ++- .../management_endpoints/team_endpoints.py | 5 ++ .../mcp_server/test_proxy_api_credentials.py | 22 ++++++ .../test_team_endpoints.py | 71 +++++++++++++++++++ 4 files changed, 102 insertions(+), 5 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/bridge_token_flow.py b/litellm/proxy/_experimental/mcp_server/bridge_token_flow.py index f19cb87ae18..37a893973e3 100644 --- a/litellm/proxy/_experimental/mcp_server/bridge_token_flow.py +++ b/litellm/proxy/_experimental/mcp_server/bridge_token_flow.py @@ -279,11 +279,10 @@ async def load_active_user_by_id( catches every DB failure and re-raises a bare ``ValueError`` (a deleted user and a real outage look identical, the original error surviving only as ``__context__``), so the outage check walks the cause chain, and a missing user falls through to ``no_active_key`` rather than an opaque gateway fault. - ``source="database"`` reads the row from the database, never the cache, and leaves the fresh row in the - cache for the requests the credential makes next: JWT auth caches the user it creates before it adds - that user to the JWT's team and adding a member never evicts the cached row, so a credential minted - off the cache would refuse the very first exchange as not a member. Every other caller keeps the cache - read, so introspection, which a resource server may call per request, stays off the database.""" + ``source="database"`` reads the row from the database, never the cache, so the credential mint refuses + a user that a writer deactivated or deleted without evicting the cached row, and it leaves the fresh + row in the cache for the requests the credential makes next. Every other caller keeps the cache read, + so introspection, which a resource server may call per request, stays off the database.""" from litellm.proxy._types import ( ProxyException, # noqa: PLC0415 # inline import avoids a module-load circular import ) diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index e719d6d761a..84e3e9f87e6 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -3154,6 +3154,7 @@ async def team_member_add( ``` """ + from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import evict_and_broadcast from litellm.proxy.proxy_server import ( litellm_proxy_admin_name, premium_user, @@ -3248,6 +3249,10 @@ async def team_member_add( litellm_proxy_admin_name=litellm_proxy_admin_name, ) + await evict_and_broadcast( + cache_keys=tuple(sorted(user.user_id for user in updated_users)), + user_api_key_cache=user_api_key_cache, + ) await _evict_created_membership_caches( user_ids=(tm.user_id for tm in updated_team_memberships), team_id=data.team_id, diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_proxy_api_credentials.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_proxy_api_credentials.py index 04fbbe4a6ce..ed3e5f48516 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_proxy_api_credentials.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_proxy_api_credentials.py @@ -152,6 +152,28 @@ async def test_mint_reads_the_users_teams_from_the_database_not_a_stale_cached_r assert _decoded(minted).team_id == "team-a" +@pytest.mark.asyncio +async def test_mint_refuses_a_user_scim_deactivated_after_the_cache_last_saw_them_active(fetch_teams, monkeypatch): + """SCIM deactivation writes the user row without evicting the cached copy, so a mint off the cache would + keep issuing credentials for the management-object TTL. The mint reads the database row, so the + deactivated user is refused on the first refresh after the deactivation.""" + from litellm.proxy import proxy_server + + cache = UserApiKeyCache() + await cache.async_set_cache( + key="deactivated-user", value=_user(user_id="deactivated-user", teams=["team-a"]), model_type=LiteLLM_UserTable + ) + prisma = MagicMock() + prisma.db.litellm_usertable.find_unique = AsyncMock( + return_value=_user(user_id="deactivated-user", teams=["team-a"], metadata={"scim_active": False}) + ) + monkeypatch.setattr(proxy_server, "user_api_key_cache", cache) + monkeypatch.setattr(proxy_server, "prisma_client", prisma) + + assert await mint_proxy_credential("deactivated-user", "team-a") == "no_active_key" + fetch_teams.assert_not_awaited() + + @pytest.mark.asyncio async def test_mint_refuses_a_team_the_user_is_not_on(load_user, fetch_teams): assert await mint_proxy_credential("u1", "team-c") == "not_a_member" diff --git a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py index ebbedc6541e..748cdbbc175 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -13090,6 +13090,77 @@ async def test_team_member_add_audits_a_user_created_from_a_list_payload(monkeyp assert created_user_id not in mock_audit.call_args.kwargs["existing_user_ids"] +@pytest.mark.asyncio +async def test_team_member_add_evicts_the_new_members_cached_user_row_on_every_worker(monkeypatch): + """Auth admits a team-bound credential off the teams list of the cached user row. The add wrote the + new team to the database row only, so a worker still holding the old row refused the member's + credential with 403 until the management-object TTL expired. The add now evicts the row here and + broadcasts the eviction to the other workers, the way /team/member_delete already does.""" + from litellm.proxy._types import TeamMemberAddRequest + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + from litellm.proxy.management_endpoints.team_endpoints import team_member_add + + team_id = "team-b" + user_id = "dev-1" + cache = UserApiKeyCache() + await cache.async_set_cache( + key=user_id, value=LiteLLM_UserTable(user_id=user_id, teams=["team-a"]), model_type=LiteLLM_UserTable + ) + broadcast = AsyncMock() + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", AsyncMock()) + monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True) + monkeypatch.setattr("litellm.proxy.proxy_server.litellm_proxy_admin_name", "default_user_id") + monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", cache) + monkeypatch.setattr( + "litellm.proxy.common_utils.auth_cache_invalidation_pubsub.publish_auth_cache_invalidation", broadcast + ) + + updated_team = MagicMock() + updated_team.model_dump.return_value = { + "team_id": team_id, + "members_with_roles": [{"user_id": user_id, "role": "user"}], + } + + async def fake_add_team_members_to_team(**kwargs): + return updated_team, [LiteLLM_UserTable(user_id=user_id, teams=["team-a", team_id])], [] + + with ( + patch( # test-quality-ok: team_member_add has no injection seam for its prisma-backed helpers + "litellm.proxy.management_endpoints.team_endpoints.get_team_object", + new_callable=AsyncMock, + return_value=LiteLLM_TeamTable(team_id=team_id, members_with_roles=[]), + ), + patch( # test-quality-ok: team_member_add has no injection seam for its prisma-backed helpers + "litellm.proxy.management_endpoints.team_endpoints._validate_team_member_add_permissions", + new_callable=AsyncMock, + ), + patch( # test-quality-ok: team_member_add has no injection seam for its prisma-backed helpers + "litellm.proxy.management_endpoints.team_endpoints._validate_and_populate_member_user_info", + new_callable=AsyncMock, + ), + patch( # test-quality-ok: team_member_add has no injection seam for its prisma-backed helpers + "litellm.proxy.management_endpoints.team_endpoints._resolve_existing_member_user_ids", + new_callable=AsyncMock, + return_value=frozenset({user_id}), + ), + patch( # test-quality-ok: team_member_add has no injection seam for its prisma-backed helpers + "litellm.proxy.management_endpoints.team_endpoints._add_team_members_to_team", + side_effect=fake_add_team_members_to_team, + ), + patch( # test-quality-ok: team_member_add has no injection seam for its prisma-backed helpers + "litellm.proxy.management_endpoints.team_endpoints._create_team_member_add_audit_logs", + new_callable=AsyncMock, + ), + ): + await team_member_add( + data=TeamMemberAddRequest(team_id=team_id, member=Member(user_id=user_id, role="user")), + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin-1"), + ) + + assert await cache.async_get_cache(key=user_id, model_type=LiteLLM_UserTable) is None + broadcast.assert_awaited_once_with(cache_key=user_id) + + def test_validate_member_user_id_provisioning_caps_the_ids_it_echoes_back(): """A large member list must not echo every id back in the error body.""" from litellm.proxy.management_endpoints.team_endpoints import ( From 769b47457ee2aa35ac20e6b2815bc7800ef95bf2 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 16 Sep 2026 17:22:27 -0700 Subject: [PATCH 019/144] fix(proxy): keep the token exchange off gateways that map JWTs to virtual keys --- .../mcp_server/idp_token_exchange.py | 26 ++++++++-- .../mcp_server/test_discoverable_endpoints.py | 26 ++++++++-- .../mcp_server/test_idp_token_exchange.py | 50 +++++++++++++++---- 3 files changed, 85 insertions(+), 17 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/idp_token_exchange.py b/litellm/proxy/_experimental/mcp_server/idp_token_exchange.py index cdefaf76d49..7a453e85cce 100644 --- a/litellm/proxy/_experimental/mcp_server/idp_token_exchange.py +++ b/litellm/proxy/_experimental/mcp_server/idp_token_exchange.py @@ -14,7 +14,7 @@ from fastapi import HTTPException, Request from litellm._logging import verbose_proxy_logger from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import SubjectIdentity, SubjectTokenRefusal from litellm.proxy._types import JWTAuthBuilderResult, ProxyException -from litellm.proxy.auth.handle_jwt import JWTAuthManager +from litellm.proxy.auth.handle_jwt import JWTAuthManager, JWTHandler EXCHANGE_ROUTE: Final = "/token" REJECTED_SUBJECT_TOKEN: Final = "subject_token was rejected by the gateway's JWT auth" @@ -23,16 +23,21 @@ REJECTED_SUBJECT_TOKEN: Final = "subject_token was rejected by the gateway's JWT @dataclass(frozen=True, slots=True) class TokenExchangePrerequisites: """The deployment-level gates ``user_api_key_auth`` applies before it verifies any JWT - bearer. Discovery and registration advertise the exchange grant only when every one of - them holds, and an exchange attempt is refused naming the first one that does not.""" + bearer, plus the JWT-to-virtual-key mapping it consults first: a gateway that maps + tokens authenticates a JWT as its mapped key, with that key's models and budget, or + refuses an unmapped one, and the exchange proves the token through ``auth_builder`` + alone, so it would mint the user's own credential past that policy. Discovery and + registration advertise the exchange grant only when every gate holds, and an exchange + attempt is refused naming the first one that does not.""" jwt_auth_enabled: bool has_database: bool licensed: bool + maps_jwts_to_virtual_keys: bool @property def available(self) -> bool: - return self.jwt_auth_enabled and self.has_database and self.licensed + return self.jwt_auth_enabled and self.has_database and self.licensed and not self.maps_jwts_to_virtual_keys def refusal(self) -> SubjectTokenRefusal | None: if not self.jwt_auth_enabled: @@ -50,12 +55,18 @@ class TokenExchangePrerequisites: error="unsupported_grant_type", description="JWT auth is an enterprise only feature; no license is set", ) + if self.maps_jwts_to_virtual_keys: + return SubjectTokenRefusal( + error="unsupported_grant_type", + description="this gateway maps IdP tokens to virtual keys, which the exchange does not serve", + ) return None def read_token_exchange_prerequisites() -> TokenExchangePrerequisites: from litellm.proxy.proxy_server import ( # noqa: PLC0415 # rebound after startup, so read them per call general_settings, + jwt_handler, premium_user, prisma_client, ) @@ -64,9 +75,16 @@ def read_token_exchange_prerequisites() -> TokenExchangePrerequisites: jwt_auth_enabled=general_settings.get("enable_jwt_auth", False) is True, has_database=prisma_client is not None, licensed=premium_user is True, + maps_jwts_to_virtual_keys=_maps_jwts_to_virtual_keys(jwt_handler), ) +def _maps_jwts_to_virtual_keys(jwt_handler: JWTHandler) -> bool: + if not hasattr(jwt_handler, "litellm_jwtauth"): + return False + return jwt_handler.litellm_jwtauth.is_virtual_key_mapping_configured() + + def token_exchange_available() -> bool: return read_token_exchange_prerequisites().available diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py index 6965b3b4ebe..d7666f5e694 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py @@ -11111,13 +11111,31 @@ def test_native_client_login_walks_discovery_consent_token_refresh_and_revoke(mo assert stranger.json()["error"] == "invalid_client" -@pytest.mark.parametrize("exchange_servable", [True, False]) -def test_discovery_advertises_the_exchange_grant_only_where_the_gateway_can_serve_it(monkeypatch, exchange_servable): +@pytest.mark.parametrize( + "jwt_auth_enabled, virtual_key_claim_field, exchange_servable", + [(True, None, True), (False, None, False), (True, "client_id", False)], + ids=["jwt auth on", "jwt auth off", "jwts mapped to virtual keys"], +) +def test_discovery_advertises_the_exchange_grant_only_where_the_gateway_can_serve_it( + monkeypatch, jwt_auth_enabled, virtual_key_claim_field, exchange_servable +): """Every document a native client reads before it picks a grant (the versioned contract, the aggregate authorization-server metadata, and the registration response) lists the RFC 8693 - exchange exactly when the running proxy can serve it: JWT auth on, a database, and a license.""" + exchange exactly when the running proxy can serve it: JWT auth on, a database, a license, and + no JWT-to-virtual-key mapping, since the exchange would mint past the mapped key's policy.""" + from litellm.caching.caching import DualCache + from litellm.proxy._types import LiteLLM_JWTAuth + from litellm.proxy.auth.handle_jwt import JWTHandler + client, _session_cookie, _minted = _native_client_app(monkeypatch) - monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {"enable_jwt_auth": exchange_servable}) + handler: Final = JWTHandler() + handler.update_environment( + prisma_client=None, + user_api_key_cache=DualCache(), + litellm_jwtauth=LiteLLM_JWTAuth(virtual_key_claim_field=virtual_key_claim_field), + ) + monkeypatch.setattr("litellm.proxy.proxy_server.jwt_handler", handler) + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {"enable_jwt_auth": jwt_auth_enabled}) monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", object()) monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True) exchange_grant = ["urn:ietf:params:oauth:grant-type:token-exchange"] if exchange_servable else [] diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_idp_token_exchange.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_idp_token_exchange.py index d1b049dddd5..e12c8823f99 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_idp_token_exchange.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_idp_token_exchange.py @@ -3,6 +3,7 @@ import logging import pytest from fastapi import HTTPException +from litellm.caching.caching import DualCache from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import SubjectIdentity, SubjectTokenRefusal from litellm.proxy._experimental.mcp_server.idp_token_exchange import ( REJECTED_SUBJECT_TOKEN, @@ -10,12 +11,17 @@ from litellm.proxy._experimental.mcp_server.idp_token_exchange import ( identity_from_subject_token, token_exchange_available, ) -from litellm.proxy._types import ProxyException +from litellm.proxy._types import JWTIssuerConfig, LiteLLM_JWTAuth, ProxyException from litellm.proxy.auth.handle_jwt import JWTHandler IDP_JWT = "eyJhbGciOiJSUzI1NiJ9.eyJzdWIiOiJ1MSJ9.idp-signature" REQUEST_HEADERS = {"x-litellm-team-id": "team-b", "user-agent": "lite/0.1"} -EVERY_GATE_HOLDS = {"jwt_auth_enabled": True, "has_database": True, "licensed": True} +EVERY_GATE_HOLDS = { + "jwt_auth_enabled": True, + "has_database": True, + "licensed": True, + "maps_jwts_to_virtual_keys": False, +} JWKS_URL = "https://idp.example.com/.well-known/jwks.json" @@ -82,6 +88,7 @@ async def test_a_jwt_that_resolves_no_team_names_a_teamless_identity(): ({"jwt_auth_enabled": False}, IDP_JWT, "unsupported_grant_type", "JWT auth is not enabled"), ({"has_database": False}, IDP_JWT, "unsupported_grant_type", "no database"), ({"licensed": False}, IDP_JWT, "unsupported_grant_type", "enterprise"), + ({"maps_jwts_to_virtual_keys": True}, IDP_JWT, "unsupported_grant_type", "virtual keys"), ({}, "sk-litellm-virtual-key", "invalid_request", "not a JWT"), ], ) @@ -96,28 +103,53 @@ async def test_the_gates_user_api_key_auth_applies_refuse_before_any_verificatio assert authorizer.calls == [] -@pytest.mark.parametrize("unmet", [{}, {"jwt_auth_enabled": False}, {"has_database": False}, {"licensed": False}]) +@pytest.mark.parametrize( + "unmet", + [ + {}, + {"jwt_auth_enabled": False}, + {"has_database": False}, + {"licensed": False}, + {"maps_jwts_to_virtual_keys": True}, + ], +) def test_the_grant_is_available_exactly_when_every_gate_holds(unmet): prerequisites = TokenExchangePrerequisites(**{**EVERY_GATE_HOLDS, **unmet}) assert prerequisites.available is (unmet == {}) assert (prerequisites.refusal() is None) is prerequisites.available +MAPPED_ISSUER = JWTIssuerConfig( + issuer="https://idp.example.test", audience="litellm-gateway", virtual_key_claim_field="client_id" +) + + +def _running_jwt_handler(litellm_jwtauth): + handler = JWTHandler() + if litellm_jwtauth is not None: + handler.update_environment(prisma_client=None, user_api_key_cache=DualCache(), litellm_jwtauth=litellm_jwtauth) + return handler + + @pytest.mark.parametrize( - "general_settings, prisma_client, premium_user, expected", + "general_settings, prisma_client, premium_user, litellm_jwtauth, expected", [ - ({"enable_jwt_auth": True}, object(), True, True), - ({}, object(), True, False), - ({"enable_jwt_auth": True}, None, True, False), - ({"enable_jwt_auth": True}, object(), False, False), + ({"enable_jwt_auth": True}, object(), True, LiteLLM_JWTAuth(), True), + ({"enable_jwt_auth": True}, object(), True, None, True), + ({}, object(), True, LiteLLM_JWTAuth(), False), + ({"enable_jwt_auth": True}, None, True, LiteLLM_JWTAuth(), False), + ({"enable_jwt_auth": True}, object(), False, LiteLLM_JWTAuth(), False), + ({"enable_jwt_auth": True}, object(), True, LiteLLM_JWTAuth(virtual_key_claim_field="client_id"), False), + ({"enable_jwt_auth": True}, object(), True, LiteLLM_JWTAuth(issuers=[MAPPED_ISSUER]), False), ], ) def test_availability_is_read_from_the_running_proxy( - monkeypatch, general_settings, prisma_client, premium_user, expected + monkeypatch, general_settings, prisma_client, premium_user, litellm_jwtauth, expected ): monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", general_settings) monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", prisma_client) monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", premium_user) + monkeypatch.setattr("litellm.proxy.proxy_server.jwt_handler", _running_jwt_handler(litellm_jwtauth)) assert token_exchange_available() is expected From 441021fc96eb24680ec41f78f16000402deff8b9 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 17 Sep 2026 04:46:26 +0000 Subject: [PATCH 020/144] fix(responses): announce message item before text events in the chat completions bridge Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../streaming_iterator.py | 67 ++++++---- .../test_streaming_iterator_transformation.py | 118 ++++++++++++++++++ 2 files changed, 160 insertions(+), 25 deletions(-) diff --git a/litellm/responses/litellm_completion_transformation/streaming_iterator.py b/litellm/responses/litellm_completion_transformation/streaming_iterator.py index 1b9f39449cf..7af62e9bfef 100644 --- a/litellm/responses/litellm_completion_transformation/streaming_iterator.py +++ b/litellm/responses/litellm_completion_transformation/streaming_iterator.py @@ -102,6 +102,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): self.sent_response_created_event: bool = False self.sent_response_in_progress_event: bool = False self.sent_output_item_added_event: bool = False + self.sent_message_item_added_event: bool = False self.sent_content_part_added_event: bool = False self.sent_output_text_done_event: bool = False self.sent_output_content_part_done_event: bool = False @@ -592,6 +593,29 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): event.__dict__["sequence_number"] = self._sequence_number return event + def _queue_message_item_added_events(self) -> None: + if self._cached_item_id is None: + self._cached_item_id = f"msg_{uuid.uuid4()}" + self.sent_message_item_added_event = True + self.sent_content_part_added_event = True + self._sequence_number += 1 + event: Final = OutputItemAddedEvent( + type=ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED, + output_index=0, + item=BaseLiteLLMOpenAIResponseObject( + **{ + "id": self._cached_item_id, + "type": "message", + "role": "assistant", + "status": "in_progress", + "content": [], + } + ), + ) + event.__dict__["sequence_number"] = self._sequence_number + self._pending_response_events.append(event) + self._pending_response_events.append(self.create_content_part_added_event()) + def _merge_provider_specific_fields(self, src: dict) -> None: """Merge provider_specific_fields using last-value-wins for lists. @@ -832,6 +856,15 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): def return_default_done_events( self, litellm_complete_object: ModelResponse ) -> BaseLiteLLMOpenAIResponseObject | None: + if self.sent_message_item_added_event is False: + final_content: Final = litellm_complete_object.choices[0].message.content or "" + if not final_content: + self.sent_output_text_done_event = True + self.sent_output_content_part_done_event = True + self.sent_output_item_done_event = True + return None + self._queue_message_item_added_events() + return self._pending_response_events.pop(0) if self.sent_output_text_done_event is False: self.sent_output_text_done_event = True return self.create_output_text_done_event(litellm_complete_object) @@ -898,6 +931,12 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): def _ensure_output_item_for_chunk(self, chunk: ModelResponseStream) -> None: # Change: Never return a value, just enqueue output item events if self.sent_output_item_added_event: + if ( + not self.sent_message_item_added_event + and chunk.choices + and self._get_delta_string_from_streaming_choices(chunk.choices) + ): + self._queue_message_item_added_events() return if not chunk.choices: return @@ -936,31 +975,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): return # Default: message - self._cached_item_id = self._cached_item_id or f"msg_{uuid.uuid4()}" - event = OutputItemAddedEvent( - type=ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED, - output_index=0, - item=BaseLiteLLMOpenAIResponseObject( - **{ - "id": self._cached_item_id, - "type": "message", - "role": "assistant", - "status": "in_progress", - "content": [], - } - ), - ) - event.__dict__["sequence_number"] = self._sequence_number - self._pending_response_events.append(event) - - # Emit content_part.added immediately after output_item.added for message - # items. The OpenAI Responses spec requires this event before any - # output_text.delta events so downstream parsers can initialize the - # text part structure. - if not self.sent_content_part_added_event: - self.sent_content_part_added_event = True - content_part_event: Final = self.create_content_part_added_event() - self._pending_response_events.append(content_part_event) + self._queue_message_item_added_events() return async def __anext__( @@ -1189,6 +1204,8 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): # Priority 2: Handle text deltas delta_content: Final = self._get_delta_string_from_streaming_choices(chunk.choices) if delta_content: + if not self.sent_message_item_added_event: + self._queue_message_item_added_events() self._sequence_number += 1 text_delta_event: Final = OutputTextDeltaEvent( type=ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA, diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py b/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py index 343fc873fa4..7a7500f666b 100644 --- a/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py +++ b/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py @@ -957,3 +957,121 @@ def test_streamed_unrecognized_tool_choice_is_echoed_as_auto() -> None: ] assert [event.response.tool_choice for event in response_events] == ["auto", "auto", "auto"] + + +def _reasoning_chunk(reasoning: str, finish_reason: str | None = None) -> ModelResponseStream: + return ModelResponseStream( + id=CHAT_COMPLETION_ID, + created=1748575031, + model="claude-haiku-4-5", + object="chat.completion.chunk", + choices=[ + StreamingChoices( + index=0, + delta=Delta(role="assistant", reasoning_content=reasoning), + finish_reason=finish_reason, + ) + ], + ) + + +async def _collect_events(iterator: LiteLLMCompletionStreamingIterator, sync_mode: bool) -> list: + if sync_mode: + return list(iterator) + return [event async for event in iterator] + + +def _is_message_item(event) -> bool: + return getattr(getattr(event, "item", None), "type", None) == "message" + + +@pytest.mark.parametrize("sync_mode", [True, False]) +@pytest.mark.asyncio +async def test_tool_only_stream_emits_no_message_item_events(sync_mode): + """ + A turn that only calls tools must not announce or close a message output item: + Vercel AI SDK clients reject text/item events that reference a message id they + never saw in response.output_item.added. + """ + iterator: Final = _build_iterator([_tool_call_chunk(), _chunk("", finish_reason="tool_calls")]) + + events: Final = await _collect_events(iterator, sync_mode) + + message_item_events = [ + event + for event in events + if getattr(event, "type", None) + in (ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED, ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE) + and _is_message_item(event) + ] + assert message_item_events == [] + assert [ + event + for event in events + if str(getattr(event, "type", "")).startswith("response.output_text") + or getattr(event, "type", None) + in (ResponsesAPIStreamEvents.CONTENT_PART_ADDED, ResponsesAPIStreamEvents.CONTENT_PART_DONE) + ] == [] + assert any(getattr(event, "type", None) == ResponsesAPIStreamEvents.RESPONSE_COMPLETED for event in events) + + +@pytest.mark.parametrize("sync_mode", [True, False]) +@pytest.mark.asyncio +async def test_reasoning_then_text_announces_message_item_before_text_events(sync_mode): + """ + When reasoning is announced first, a later text delta still has to be preceded by + the message output_item.added/content_part.added, and every text-scoped event must + reference that announced message item id. + """ + iterator: Final = _build_iterator( + [ + _reasoning_chunk("let me think"), + _chunk("Hello"), + _chunk("!", finish_reason="stop"), + ] + ) + + events: Final = await _collect_events(iterator, sync_mode) + + announced_message_ids: set[str] = set() + content_part_added_seen = False + saw_text_delta = False + for event in events: + event_type = getattr(event, "type", None) + if event_type == ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED and _is_message_item(event): + announced_message_ids.add(event.item.id) + elif event_type == ResponsesAPIStreamEvents.CONTENT_PART_ADDED: + content_part_added_seen = True + elif event_type in ( + ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA, + ResponsesAPIStreamEvents.OUTPUT_TEXT_DONE, + ResponsesAPIStreamEvents.CONTENT_PART_DONE, + ): + assert event.item_id in announced_message_ids + if event_type == ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA: + assert content_part_added_seen + saw_text_delta = True + elif event_type == ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE and _is_message_item(event): + assert event.item.id in announced_message_ids + assert saw_text_delta + + +@pytest.mark.parametrize("sync_mode", [True, False]) +@pytest.mark.asyncio +async def test_plain_text_stream_announces_exactly_one_message_item(sync_mode): + iterator: Final = _build_iterator([_chunk("Hel"), _chunk("lo", finish_reason="stop")]) + + events: Final = await _collect_events(iterator, sync_mode) + + message_item_adds = [ + event + for event in events + if getattr(event, "type", None) == ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED and _is_message_item(event) + ] + assert len(message_item_adds) == 1 + for event in events: + if getattr(event, "type", None) in ( + ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA, + ResponsesAPIStreamEvents.OUTPUT_TEXT_DONE, + ): + assert event.item_id == message_item_adds[0].item.id From d4e54a0f345aedcfca85f120441ff2a56f21d5f7 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 17 Sep 2026 05:05:54 +0000 Subject: [PATCH 021/144] fix(responses): keep sync text deltas and give the message item its own output index Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../streaming_iterator.py | 26 ++++++++++--------- .../test_streaming_iterator_transformation.py | 11 ++++++-- 2 files changed, 23 insertions(+), 14 deletions(-) diff --git a/litellm/responses/litellm_completion_transformation/streaming_iterator.py b/litellm/responses/litellm_completion_transformation/streaming_iterator.py index 7af62e9bfef..6b13c9d4297 100644 --- a/litellm/responses/litellm_completion_transformation/streaming_iterator.py +++ b/litellm/responses/litellm_completion_transformation/streaming_iterator.py @@ -112,6 +112,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): self.completed_response = None self.final_text: str = "" self._cached_item_id: str | None = None + self._message_output_index: int = 0 self._cached_response_id: str | None = None self._buffered_chunk: ModelResponseStream | None = None self._upstream_exhausted: bool = False @@ -564,7 +565,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): self._sequence_number += 1 event: Final = OutputItemAddedEvent( type=ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED, - output_index=0, + output_index=self._message_output_index, item=BaseLiteLLMOpenAIResponseObject( **{ "id": self._cached_item_id, @@ -586,7 +587,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): event: Final = ContentPartAddedEvent( type=ResponsesAPIStreamEvents.CONTENT_PART_ADDED, item_id=self._cached_item_id, - output_index=0, + output_index=self._message_output_index, content_index=0, part=BaseLiteLLMOpenAIResponseObject(**{"type": "output_text", "text": "", "annotations": []}), ) @@ -598,10 +599,11 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): self._cached_item_id = f"msg_{uuid.uuid4()}" self.sent_message_item_added_event = True self.sent_content_part_added_event = True + self._message_output_index = 1 if self._cached_reasoning_item_id is not None else 0 self._sequence_number += 1 event: Final = OutputItemAddedEvent( type=ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED, - output_index=0, + output_index=self._message_output_index, item=BaseLiteLLMOpenAIResponseObject( **{ "id": self._cached_item_id, @@ -735,7 +737,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): return OutputTextDoneEvent( type=ResponsesAPIStreamEvents.OUTPUT_TEXT_DONE, item_id=self._cached_item_id, - output_index=0, + output_index=self._message_output_index, content_index=0, text=getattr(litellm_complete_object.choices[0].message, "content", "") or "", ) @@ -771,7 +773,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): return ContentPartDoneEvent( type=ResponsesAPIStreamEvents.CONTENT_PART_DONE, item_id=self._cached_item_id, - output_index=0, + output_index=self._message_output_index, content_index=0, part=part, ) @@ -790,7 +792,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): ) return OutputItemDoneEvent( type=ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE, - output_index=0, + output_index=self._message_output_index, sequence_number=1, item=BaseLiteLLMOpenAIResponseObject( **{ @@ -951,6 +953,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): if self._cached_reasoning_item_id is None: self._cached_reasoning_item_id = f"rs_{uuid.uuid4()}" self._reasoning_item_id = self._cached_reasoning_item_id + self._next_tool_output_index = max(self._next_tool_output_index, 2) event = OutputItemAddedEvent( type=ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED, @@ -1130,12 +1133,11 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): self.collected_chat_completion_chunks.append( self._snapshot_chunk_for_stream_chunk_builder(cast(ModelResponseStream, chunk)) ) - # Emit any just-queued output_item event - if self._pending_response_events: - return self._pending_response_events.pop(0) response_api_chunk = self._transform_chat_completion_chunk_to_response_api_chunk(chunk) if response_api_chunk: - return response_api_chunk + self._pending_response_events.append(response_api_chunk) + if self._pending_response_events: + return self._pending_response_events.pop(0) # Otherwise, loop to next chunk except StopIteration: return self.common_done_event_logic(sync_mode=True) @@ -1177,7 +1179,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): event = OutputTextAnnotationAddedEvent( type=ResponsesAPIStreamEvents.OUTPUT_TEXT_ANNOTATION_ADDED, item_id=item_id, - output_index=0, + output_index=self._message_output_index, content_index=0, annotation_index=idx, annotation=annotation_dict, @@ -1210,7 +1212,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): text_delta_event: Final = OutputTextDeltaEvent( type=ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA, item_id=item_id, - output_index=0, + output_index=self._message_output_index, content_index=0, delta=delta_content, ) diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py b/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py index 7a7500f666b..895f59632c7 100644 --- a/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py +++ b/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py @@ -1034,12 +1034,15 @@ async def test_reasoning_then_text_announces_message_item_before_text_events(syn events: Final = await _collect_events(iterator, sync_mode) announced_message_ids: set[str] = set() + announced_indexes_by_item_type: dict[str, int] = {} content_part_added_seen = False saw_text_delta = False for event in events: event_type = getattr(event, "type", None) - if event_type == ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED and _is_message_item(event): - announced_message_ids.add(event.item.id) + if event_type == ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED: + announced_indexes_by_item_type[event.item.type] = event.output_index + if _is_message_item(event): + announced_message_ids.add(event.item.id) elif event_type == ResponsesAPIStreamEvents.CONTENT_PART_ADDED: content_part_added_seen = True elif event_type in ( @@ -1054,6 +1057,10 @@ async def test_reasoning_then_text_announces_message_item_before_text_events(syn elif event_type == ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE and _is_message_item(event): assert event.item.id in announced_message_ids assert saw_text_delta + assert "".join( + event.delta for event in events if getattr(event, "type", None) == ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA + ) == "Hello!" + assert announced_indexes_by_item_type["message"] != announced_indexes_by_item_type["reasoning"] @pytest.mark.parametrize("sync_mode", [True, False]) From 42c4c8163333328fa053f9ee8967ddf2d319c186 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 17 Sep 2026 05:15:07 +0000 Subject: [PATCH 022/144] fix(responses): allocate the message output index from the shared item allocator Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../streaming_iterator.py | 5 ++-- .../test_streaming_iterator_transformation.py | 28 +++++++++++++++++++ 2 files changed, 31 insertions(+), 2 deletions(-) diff --git a/litellm/responses/litellm_completion_transformation/streaming_iterator.py b/litellm/responses/litellm_completion_transformation/streaming_iterator.py index 6b13c9d4297..beffd12a349 100644 --- a/litellm/responses/litellm_completion_transformation/streaming_iterator.py +++ b/litellm/responses/litellm_completion_transformation/streaming_iterator.py @@ -599,7 +599,9 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): self._cached_item_id = f"msg_{uuid.uuid4()}" self.sent_message_item_added_event = True self.sent_content_part_added_event = True - self._message_output_index = 1 if self._cached_reasoning_item_id is not None else 0 + if self._cached_reasoning_item_id is not None: + self._message_output_index = self._next_tool_output_index + self._next_tool_output_index += 1 self._sequence_number += 1 event: Final = OutputItemAddedEvent( type=ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED, @@ -953,7 +955,6 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): if self._cached_reasoning_item_id is None: self._cached_reasoning_item_id = f"rs_{uuid.uuid4()}" self._reasoning_item_id = self._cached_reasoning_item_id - self._next_tool_output_index = max(self._next_tool_output_index, 2) event = OutputItemAddedEvent( type=ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED, diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py b/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py index 895f59632c7..c727a5be4bc 100644 --- a/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py +++ b/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py @@ -1063,6 +1063,34 @@ async def test_reasoning_then_text_announces_message_item_before_text_events(syn assert announced_indexes_by_item_type["message"] != announced_indexes_by_item_type["reasoning"] +@pytest.mark.parametrize("sync_mode", [True, False]) +@pytest.mark.asyncio +async def test_tool_then_reasoning_then_text_gives_message_its_own_output_index(sync_mode): + iterator: Final = _build_iterator( + [ + _tool_call_chunk(), + _reasoning_chunk("thinking"), + _chunk("Hello"), + _chunk("!", finish_reason="stop"), + ] + ) + + events: Final = await _collect_events(iterator, sync_mode) + output_item_added_events: Final = [ + event for event in events if getattr(event, "type", None) == ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED + ] + message_item_adds: Final = [event for event in output_item_added_events if _is_message_item(event)] + function_call_adds: Final = [ + event for event in output_item_added_events if getattr(event.item, "type", None) == "function_call" + ] + + assert len(message_item_adds) == 1 + assert all(message_item_adds[0].output_index != event.output_index for event in function_call_adds) + + output_indexes_by_item_id: Final = {event.item.id: event.output_index for event in output_item_added_events} + assert len(output_indexes_by_item_id) == len(set(output_indexes_by_item_id.values())) + + @pytest.mark.parametrize("sync_mode", [True, False]) @pytest.mark.asyncio async def test_plain_text_stream_announces_exactly_one_message_item(sync_mode): From e9625ad069920a224be093cede8de0bb1f379c0a Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 17 Sep 2026 05:22:09 +0000 Subject: [PATCH 023/144] fix(responses): default the message output index when no reasoning item exists Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../litellm_completion_transformation/streaming_iterator.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/litellm/responses/litellm_completion_transformation/streaming_iterator.py b/litellm/responses/litellm_completion_transformation/streaming_iterator.py index beffd12a349..71502e33d5c 100644 --- a/litellm/responses/litellm_completion_transformation/streaming_iterator.py +++ b/litellm/responses/litellm_completion_transformation/streaming_iterator.py @@ -602,6 +602,8 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): if self._cached_reasoning_item_id is not None: self._message_output_index = self._next_tool_output_index self._next_tool_output_index += 1 + else: + self._message_output_index = 0 self._sequence_number += 1 event: Final = OutputItemAddedEvent( type=ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED, From 586124094667c30d9a81a810ec4faee1d9c68216 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 17 Sep 2026 05:43:14 +0000 Subject: [PATCH 024/144] test(responses): type new streaming bridge test parameters Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../test_streaming_iterator_transformation.py | 29 ++++++++----------- 1 file changed, 12 insertions(+), 17 deletions(-) diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py b/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py index c727a5be4bc..3581771bc63 100644 --- a/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py +++ b/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py @@ -20,7 +20,10 @@ from litellm.responses.litellm_completion_transformation.streaming_iterator impo LiteLLMCompletionStreamingIterator, ) from litellm.responses.utils import ResponsesAPIRequestUtils -from litellm.types.llms.openai import ResponsesAPIStreamEvents +from litellm.types.llms.openai import ( + BaseLiteLLMOpenAIResponseObject, + ResponsesAPIStreamEvents, +) from litellm.types.responses.main import build_web_search_call from litellm.types.utils import ( Delta, @@ -975,24 +978,21 @@ def _reasoning_chunk(reasoning: str, finish_reason: str | None = None) -> ModelR ) -async def _collect_events(iterator: LiteLLMCompletionStreamingIterator, sync_mode: bool) -> list: +async def _collect_events( + iterator: LiteLLMCompletionStreamingIterator, sync_mode: bool +) -> list[BaseLiteLLMOpenAIResponseObject]: if sync_mode: return list(iterator) return [event async for event in iterator] -def _is_message_item(event) -> bool: +def _is_message_item(event: BaseLiteLLMOpenAIResponseObject) -> bool: return getattr(getattr(event, "item", None), "type", None) == "message" @pytest.mark.parametrize("sync_mode", [True, False]) @pytest.mark.asyncio -async def test_tool_only_stream_emits_no_message_item_events(sync_mode): - """ - A turn that only calls tools must not announce or close a message output item: - Vercel AI SDK clients reject text/item events that reference a message id they - never saw in response.output_item.added. - """ +async def test_tool_only_stream_emits_no_message_item_events(sync_mode: bool): iterator: Final = _build_iterator([_tool_call_chunk(), _chunk("", finish_reason="tool_calls")]) events: Final = await _collect_events(iterator, sync_mode) @@ -1017,12 +1017,7 @@ async def test_tool_only_stream_emits_no_message_item_events(sync_mode): @pytest.mark.parametrize("sync_mode", [True, False]) @pytest.mark.asyncio -async def test_reasoning_then_text_announces_message_item_before_text_events(sync_mode): - """ - When reasoning is announced first, a later text delta still has to be preceded by - the message output_item.added/content_part.added, and every text-scoped event must - reference that announced message item id. - """ +async def test_reasoning_then_text_announces_message_item_before_text_events(sync_mode: bool): iterator: Final = _build_iterator( [ _reasoning_chunk("let me think"), @@ -1065,7 +1060,7 @@ async def test_reasoning_then_text_announces_message_item_before_text_events(syn @pytest.mark.parametrize("sync_mode", [True, False]) @pytest.mark.asyncio -async def test_tool_then_reasoning_then_text_gives_message_its_own_output_index(sync_mode): +async def test_tool_then_reasoning_then_text_gives_message_its_own_output_index(sync_mode: bool): iterator: Final = _build_iterator( [ _tool_call_chunk(), @@ -1093,7 +1088,7 @@ async def test_tool_then_reasoning_then_text_gives_message_its_own_output_index( @pytest.mark.parametrize("sync_mode", [True, False]) @pytest.mark.asyncio -async def test_plain_text_stream_announces_exactly_one_message_item(sync_mode): +async def test_plain_text_stream_announces_exactly_one_message_item(sync_mode: bool): iterator: Final = _build_iterator([_chunk("Hel"), _chunk("lo", finish_reason="stop")]) events: Final = await _collect_events(iterator, sync_mode) From 8ce2887888648fbea603ae91deffdc6e794926e9 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 17 Sep 2026 05:56:28 +0000 Subject: [PATCH 025/144] fix(responses): close the message content part as output_text on reasoning turns Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../streaming_iterator.py | 30 +++++++------------ 1 file changed, 10 insertions(+), 20 deletions(-) diff --git a/litellm/responses/litellm_completion_transformation/streaming_iterator.py b/litellm/responses/litellm_completion_transformation/streaming_iterator.py index 71502e33d5c..0cda83d979d 100644 --- a/litellm/responses/litellm_completion_transformation/streaming_iterator.py +++ b/litellm/responses/litellm_completion_transformation/streaming_iterator.py @@ -22,7 +22,6 @@ from litellm.types.llms.openai import ( ContentPartAddedEvent, ContentPartDoneEvent, ContentPartDonePartOutputText, - ContentPartDonePartReasoningText, FunctionCallArgumentsDeltaEvent, FunctionCallArgumentsDoneEvent, OutputItemAddedEvent, @@ -751,28 +750,19 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): self._cached_item_id = f"msg_{uuid.uuid4()}" text: Final = getattr(litellm_complete_object.choices[0].message, "content", "") or "" - reasoning_content = getattr(litellm_complete_object.choices[0].message, "reasoning_content", "") or "" annotations: Final = getattr(litellm_complete_object.choices[0].message, "annotations", None) - part: PART_UNION_TYPES | None = None - if reasoning_content: - part = ContentPartDonePartReasoningText( - type="reasoning_text", - reasoning=reasoning_content, - ) - - else: - response_annotations: Final = ( - LiteLLMCompletionResponsesConfig._transform_chat_completion_annotations_to_response_output_annotations( - annotations=annotations - ) - ) - part = ContentPartDonePartOutputText( - type="output_text", - text=text, - annotations=response_annotations, - logprobs=None, + response_annotations: Final = ( + LiteLLMCompletionResponsesConfig._transform_chat_completion_annotations_to_response_output_annotations( + annotations=annotations ) + ) + part: Final[PART_UNION_TYPES] = ContentPartDonePartOutputText( + type="output_text", + text=text, + annotations=response_annotations, + logprobs=None, + ) return ContentPartDoneEvent( type=ResponsesAPIStreamEvents.CONTENT_PART_DONE, From 47be6c8aeb578c71c13f56eea8d1db6a5b81ba3f Mon Sep 17 00:00:00 2001 From: yassin Date: Thu, 17 Sep 2026 18:34:36 +0000 Subject: [PATCH 026/144] feat(mcp): allowlist client applications for MCP gateway access Adds the mcp_allowed_clients general setting, enforced against the clientInfo.name each MCP client sends in its initialize request. A client not on the list, or one that does not identify itself, is rejected with 403 before any stateful session is created. The setting is configurable from config.yaml and from the Admin UI MCP network settings page Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../mcp_server/client_allowlist.py | 67 +++++ .../proxy/_experimental/mcp_server/server.py | 66 ++++- litellm/proxy/_types.py | 4 + litellm/proxy/proxy_server.py | 4 + .../mcp_server/test_client_allowlist.py | 105 ++++++++ .../mcp_server/test_mcp_server.py | 244 ++++++++++++++++++ tests/test_litellm/proxy/test_proxy_server.py | 165 +++++++++--- .../_components/MCPNetworkSettings.test.tsx | 66 ++++- .../_components/MCPNetworkSettings.tsx | 77 +++++- 9 files changed, 745 insertions(+), 53 deletions(-) create mode 100644 litellm/proxy/_experimental/mcp_server/client_allowlist.py create mode 100644 tests/test_litellm/proxy/_experimental/mcp_server/test_client_allowlist.py diff --git a/litellm/proxy/_experimental/mcp_server/client_allowlist.py b/litellm/proxy/_experimental/mcp_server/client_allowlist.py new file mode 100644 index 00000000000..57343c39565 --- /dev/null +++ b/litellm/proxy/_experimental/mcp_server/client_allowlist.py @@ -0,0 +1,67 @@ +""" +Gateway-level allowlist of MCP client applications, matched against the +``clientInfo.name`` a client sends in its JSON-RPC ``initialize`` request. The +name is client-supplied, so this is a policy control and not a security boundary. +""" + +import json +from dataclasses import dataclass +from typing import Final + +from pydantic import TypeAdapter, ValidationError + +from litellm._logging import verbose_logger + +MCP_ALLOWED_CLIENTS_SETTING: Final = "mcp_allowed_clients" + +_ALLOWED_CLIENTS_ADAPTER: Final = TypeAdapter(list[str]) + + +@dataclass(frozen=True, slots=True) +class MCPClientRejection: + client_name: str | None + + @property + def details(self) -> str: + if self.client_name is None: + return ( + "MCP initialize request did not identify the client application (clientInfo.name). " + f"This gateway only admits clients listed in {MCP_ALLOWED_CLIENTS_SETTING}." + ) + return f"MCP client '{self.client_name}' is not listed in this gateway's {MCP_ALLOWED_CLIENTS_SETTING}." + + +def parse_allowed_mcp_clients(raw_setting: object) -> frozenset[str] | None: + """None when the setting is absent (not enforced). A malformed setting admits nobody.""" + if raw_setting is None: + return None + try: + return frozenset(_ALLOWED_CLIENTS_ADAPTER.validate_python(raw_setting)) + except ValidationError: + verbose_logger.warning( + "%s is not a list of client names (%r); rejecting every MCP client until it is fixed", + MCP_ALLOWED_CLIENTS_SETTING, + raw_setting, + ) + return frozenset() + + +def extract_mcp_client_name(body: bytes) -> str | None: + try: + data: Final = json.loads(body) + except (json.JSONDecodeError, UnicodeDecodeError): + return None + params: Final = data.get("params") if isinstance(data, dict) else None + client_info: Final = params.get("clientInfo") if isinstance(params, dict) else None + name: Final = client_info.get("name") if isinstance(client_info, dict) else None + return name if isinstance(name, str) and name else None + + +def check_mcp_client_allowed(body: bytes, allowed_clients: frozenset[str] | None) -> MCPClientRejection | None: + """None when the initialize is admitted, otherwise the rejection to send back as a 403.""" + if allowed_clients is None: + return None + client_name: Final = extract_mcp_client_name(body) + if client_name is not None and client_name in allowed_clients: + return None + return MCPClientRejection(client_name=client_name) diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 7feb1fd468d..cca867d4d2a 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -463,6 +463,11 @@ if MCP_AVAILABLE: from litellm.proxy._experimental.mcp_server.auth.litellm_auth_handler import ( MCPAuthenticatedUser, ) + from litellm.proxy._experimental.mcp_server.client_allowlist import ( + MCP_ALLOWED_CLIENTS_SETTING, + check_mcp_client_allowed, + parse_allowed_mcp_clients, + ) from litellm.proxy._experimental.mcp_server.faults.list_outcomes import ( SERVER_OUTCOMES_META_KEY, AggregateToolListing, @@ -3810,6 +3815,43 @@ if MCP_AVAILABLE: except (json.JSONDecodeError, TypeError): return False + def _load_allowed_mcp_clients() -> frozenset[str] | None: + from litellm.proxy.proxy_server import general_settings + + return parse_allowed_mcp_clients(general_settings.get(MCP_ALLOWED_CLIENTS_SETTING)) + + async def _reject_initialize_from_disallowed_client( + scope: Scope, + receive: Receive, + send: Send, + body: bytes, + client_ip: str | None, + ) -> bool: + """Send a 403 and return True when the initialize body names a client the gateway does not admit.""" + rejection: Final = check_mcp_client_allowed(body, _load_allowed_mcp_clients()) + if rejection is None: + return False + verbose_logger.warning( + "Rejecting MCP initialize from client %r (ip=%s): not listed in %s", + rejection.client_name, + client_ip, + MCP_ALLOWED_CLIENTS_SETTING, + ) + forbidden: Final = JSONResponse( + status_code=403, + content={"error": "Forbidden", "details": rejection.details}, + ) + await forbidden(scope, receive, send) + return True + + def _replay_consumed_messages(consumed_messages: list[Message], receive: Receive) -> Receive: + async def wrapped_receive() -> Message: + if consumed_messages: + return consumed_messages.pop(0) + return await receive() + + return wrapped_receive + async def _read_request_body_for_routing( receive: Receive, ) -> tuple[list[Message], bytes]: @@ -4510,6 +4552,10 @@ if MCP_AVAILABLE: if scope.get("method") == "POST": consumed_messages, body = await _read_request_body_for_routing(receive) is_initialize = _is_initialize_request(body) + if is_initialize and await _reject_initialize_from_disallowed_client( + scope, receive, send, body, _client_ip + ): + return use_stateful: Final = bool(session_id or is_initialize) target_manager: Final = session_manager_stateful if use_stateful else session_manager_stateless @@ -4540,15 +4586,8 @@ if MCP_AVAILABLE: return # Replay body messages if we consumed them for peeking - original_receive: Final = receive if consumed_messages: - - async def wrapped_receive(): - if consumed_messages: - return consumed_messages.pop(0) - return await original_receive() - - receive = wrapped_receive + receive = _replay_consumed_messages(consumed_messages, receive) # Serialize requests on the same stateful session so concurrent # callers don't clobber each other's auth context mid-flight. @@ -4785,6 +4824,15 @@ if MCP_AVAILABLE: await initialize_session_managers() await asyncio.sleep(0.1) + sse_consumed_messages, sse_body = ( + await _read_request_body_for_routing(receive) if scope.get("method") == "POST" else ([], b"") + ) + if _is_initialize_request(sse_body) and await _reject_initialize_from_disallowed_client( + scope, receive, send, sse_body, _sse_client_ip + ): + return + sse_receive: Final = _replay_consumed_messages(sse_consumed_messages, receive) + async with _gateway_initialize_instructions_request_scope( user_api_key_auth, mcp_servers, @@ -4792,7 +4840,7 @@ if MCP_AVAILABLE: scoped_server_endpoint=scoped_server_endpoint, is_initialize=scope.get("method") == "GET", ): - await sse_session_manager.handle_request(scope, receive, send) + await sse_session_manager.handle_request(scope, sse_receive, send) except MCPUpstreamAuthError as e: # Upstream delegated auth returned 401; surface it to the client so # standards-compliant MCP clients trigger the upstream OAuth flow. diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 680c63393e8..bbebb99055f 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -2853,6 +2853,10 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase): None, description="Custom CIDR ranges that define internal/private networks for MCP access control. When set, only these ranges are treated as internal. Defaults to RFC 1918 private ranges (10.0.0.0/8, 172.16.0.0/12, 192.168.0.0/16, 127.0.0.0/8).", ) + mcp_allowed_clients: list[str] | None = Field( + None, + description="MCP client applications admitted by the gateway, matched exactly against the clientInfo.name the client sends in its initialize request (for example 'claude-code'). When set, an initialize from any other client, or one that does not identify itself, is rejected with 403. Unset means every client is admitted. The name is client-supplied, so this is a policy control rather than a security boundary.", + ) mcp_trusted_proxy_ranges: list[str] | None = Field( None, description="CIDR ranges of trusted reverse proxies. When set, X-Forwarded-For and X-Forwarded-* origin headers are only trusted from these IPs.", diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index d7d8413d2ce..aa98491cf3b 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -7179,6 +7179,9 @@ class ProxyConfig: "enable_openai_websocket_passthrough" ) + if "mcp_allowed_clients" not in self._yaml_general_settings_keys: + general_settings["mcp_allowed_clients"] = _general_settings.get("mcp_allowed_clients") + if "user_api_key_cache_max_size" not in self._yaml_general_settings_keys: db_cache_max_size: Final = _general_settings.get("user_api_key_cache_max_size") try: @@ -17137,6 +17140,7 @@ _GENERAL_SETTINGS_CONFIG_LIST_FIELD_TYPES: Final[Mapping[str, str]] = MappingPro "maximum_spend_logs_cleanup_run_budget": "String", "maximum_spend_logs_cleanup_batch_timeout": "String", "mcp_internal_ip_ranges": "List", + "mcp_allowed_clients": "List", "mcp_trusted_proxy_ranges": "List", "mcp_xff_num_trusted_hops": "Integer", "always_include_stream_usage": "Boolean", diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_client_allowlist.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_client_allowlist.py new file mode 100644 index 00000000000..fce8a0a6c3b --- /dev/null +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_client_allowlist.py @@ -0,0 +1,105 @@ +import json +from typing import Final + +import pytest + +from litellm.proxy._experimental.mcp_server.client_allowlist import ( + MCP_ALLOWED_CLIENTS_SETTING, + MCPClientRejection, + check_mcp_client_allowed, + extract_mcp_client_name, + parse_allowed_mcp_clients, +) + + +def _initialize_body(client_info: object) -> bytes: + return json.dumps( + { + "jsonrpc": "2.0", + "id": 0, + "method": "initialize", + "params": {"protocolVersion": "2025-06-18", "capabilities": {}, "clientInfo": client_info}, + } + ).encode() + + +CLAUDE_CODE: Final = _initialize_body({"name": "claude-code", "version": "2.1.274"}) +ANTIGRAVITY: Final = _initialize_body({"name": "antigravity-cli", "version": "1.0.0"}) + + +@pytest.mark.parametrize( + ("raw_setting", "expected"), + ( + (None, None), + ([], frozenset()), + (["antigravity-cli"], frozenset({"antigravity-cli"})), + (["antigravity-cli", "codex-mcp-client"], frozenset({"antigravity-cli", "codex-mcp-client"})), + ("antigravity-cli", frozenset()), + ([1, "antigravity-cli"], frozenset()), + ({"name": "antigravity-cli"}, frozenset()), + ), +) +def test_parse_allowed_mcp_clients(raw_setting: object, expected: frozenset[str] | None) -> None: + assert parse_allowed_mcp_clients(raw_setting) == expected + + +@pytest.mark.parametrize( + ("body", "expected"), + ( + (CLAUDE_CODE, "claude-code"), + (_initialize_body({"name": "", "version": "1"}), None), + (_initialize_body({"version": "1"}), None), + (_initialize_body({"name": 7}), None), + (_initialize_body("claude-code"), None), + (b'{"jsonrpc":"2.0","id":1,"method":"initialize","params":{}}', None), + (b'{"jsonrpc":"2.0","id":1,"method":"initialize","params":[]}', None), + (b'["not", "an", "object"]', None), + (b'{"jsonrpc":"2.0","id":1,"method":"initialize","params":{"clientInfo":{"name":"clau', None), + (b"\xff\xfe", None), + (b"", None), + ), +) +def test_extract_mcp_client_name(body: bytes, expected: str | None) -> None: + assert extract_mcp_client_name(body) == expected + + +def test_unconfigured_allowlist_admits_every_client_including_unidentified_ones() -> None: + assert check_mcp_client_allowed(CLAUDE_CODE, None) is None + assert check_mcp_client_allowed(b'{"method":"initialize","params":{}}', None) is None + assert check_mcp_client_allowed(b"garbage", None) is None + + +def test_listed_client_is_admitted_and_unlisted_client_is_rejected_by_name() -> None: + allowed: Final = frozenset({"antigravity-cli"}) + assert check_mcp_client_allowed(ANTIGRAVITY, allowed) is None + assert check_mcp_client_allowed(CLAUDE_CODE, allowed) == MCPClientRejection(client_name="claude-code") + + +def test_matching_is_exact_not_prefix_or_case_insensitive() -> None: + allowed: Final = frozenset({"claude-code"}) + assert check_mcp_client_allowed(_initialize_body({"name": "Claude-Code"}), allowed) is not None + assert check_mcp_client_allowed(_initialize_body({"name": "claude-code-sdk"}), allowed) is not None + assert check_mcp_client_allowed(_initialize_body({"name": " claude-code"}), allowed) is not None + + +def test_empty_allowlist_rejects_every_client() -> None: + assert check_mcp_client_allowed(ANTIGRAVITY, frozenset()) == MCPClientRejection(client_name="antigravity-cli") + assert check_mcp_client_allowed(CLAUDE_CODE, frozenset()) == MCPClientRejection(client_name="claude-code") + + +def test_missing_or_malformed_client_metadata_is_rejected_when_allowlist_is_set() -> None: + allowed: Final = frozenset({"antigravity-cli"}) + assert check_mcp_client_allowed(_initialize_body({"version": "1"}), allowed) == MCPClientRejection(None) + assert check_mcp_client_allowed(b'{"method":"initialize","params":{}}', allowed) == MCPClientRejection(None) + assert check_mcp_client_allowed(b"{not json", allowed) == MCPClientRejection(None) + + +def test_rejection_details_name_the_setting_and_the_offending_client() -> None: + named: Final = MCPClientRejection(client_name="claude-code").details + assert "claude-code" in named + assert MCP_ALLOWED_CLIENTS_SETTING in named + + anonymous: Final = MCPClientRejection(client_name=None).details + assert "clientInfo.name" in anonymous + assert MCP_ALLOWED_CLIENTS_SETTING in anonymous + assert "None" not in anonymous diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py index f5e4a420496..8550226e19d 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py @@ -1,4 +1,5 @@ import asyncio +import contextlib import contextvars import os from datetime import datetime, timedelta @@ -2035,6 +2036,249 @@ async def test_mcp_routing_initialize_to_stateful_no_session_to_stateless( assert not any(name.startswith(b"x-mcp-debug") for name in headers) +_CLAUDE_CODE_INITIALIZE: Final = ( + b'{"jsonrpc":"2.0","id":0,"method":"initialize","params":{"protocolVersion":"2025-06-18",' + b'"capabilities":{},"clientInfo":{"name":"claude-code","version":"2.1.274"}}}' +) +_ANTIGRAVITY_INITIALIZE: Final = ( + b'{"jsonrpc":"2.0","id":0,"method":"initialize","params":{"protocolVersion":"2025-06-18",' + b'"capabilities":{},"clientInfo":{"name":"antigravity-cli","version":"1.0.0"}}}' +) +_ANONYMOUS_INITIALIZE: Final = b'{"jsonrpc":"2.0","id":0,"method":"initialize","params":{"capabilities":{}}}' +_TOOLS_LIST: Final = b'{"jsonrpc":"2.0","id":2,"method":"tools/list","params":{}}' + + +async def _drain_body(receive) -> bytes: + chunks: list[bytes] = [] + while True: + message = await receive() + chunks.append(message.get("body", b"")) + if not message.get("more_body", False): + return b"".join(chunks) + + +def _forbidden_client_response(send: AsyncMock) -> tuple[int, dict[str, str]]: + import json as _json + + start: Final = send.call_args_list[0].args[0] + body: Final = b"".join(call.args[0].get("body", b"") for call in send.call_args_list[1:]) + return start["status"], _json.loads(body) + + +@contextlib.contextmanager +def _client_allowlist_patches(allowed_clients: object): + settings: Final = {} if allowed_clients is None else {"mcp_allowed_clients": allowed_clients} + with ( + patch( + "litellm.proxy._experimental.mcp_server.server.extract_mcp_auth_context", + new_callable=AsyncMock, + return_value=(UserAPIKeyAuth(user_id="allowlist-user"), None, None, None, None, {}), + ), + patch("litellm.proxy._experimental.mcp_server.server._SESSION_MANAGERS_INITIALIZED", True), + patch("litellm.proxy.proxy_server.general_settings", settings), + ): + yield + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("request_body", "expected_details"), + ( + ( + _CLAUDE_CODE_INITIALIZE, + "MCP client 'claude-code' is not listed in this gateway's mcp_allowed_clients.", + ), + ( + _ANONYMOUS_INITIALIZE, + "MCP initialize request did not identify the client application (clientInfo.name). " + "This gateway only admits clients listed in mcp_allowed_clients.", + ), + ), +) +async def test_streamable_http_rejects_initialize_from_unlisted_client_before_session_creation( + request_body: bytes, expected_details: str +) -> None: + from starlette.types import Scope + + from litellm.proxy._experimental.mcp_server import server as mcp_module + + scope: Final[Scope] = {"type": "http", "method": "POST", "path": "/mcp", "headers": []} + receive: Final = AsyncMock(return_value={"type": "http.request", "body": request_body, "more_body": False}) + send: Final = AsyncMock() + stateful_handle: Final = AsyncMock() + stateless_handle: Final = AsyncMock() + session_cap: Final = AsyncMock(return_value=True) + + with ( + _client_allowlist_patches(["antigravity-cli"]), + patch( + "litellm.proxy._experimental.mcp_server.server.session_manager_stateful", + SimpleNamespace(handle_request=stateful_handle), + ), + patch( + "litellm.proxy._experimental.mcp_server.server.session_manager_stateless", + SimpleNamespace(handle_request=stateless_handle), + ), + patch("litellm.proxy._experimental.mcp_server.server._enforce_stateful_session_cap_for_owner", session_cap), + ): + await mcp_module.handle_streamable_http_mcp(scope, receive, send) + + assert _forbidden_client_response(send) == (403, {"error": "Forbidden", "details": expected_details}) + stateful_handle.assert_not_awaited() + stateless_handle.assert_not_awaited() + session_cap.assert_not_awaited() + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("allowed_clients", "request_body"), + ( + (["antigravity-cli"], _ANTIGRAVITY_INITIALIZE), + (["claude-code", "antigravity-cli"], _CLAUDE_CODE_INITIALIZE), + (None, _CLAUDE_CODE_INITIALIZE), + (None, _ANONYMOUS_INITIALIZE), + ), +) +async def test_streamable_http_admits_listed_or_unrestricted_initialize_and_replays_body( + allowed_clients: list[str] | None, request_body: bytes +) -> None: + from starlette.types import Receive, Scope, Send + + from litellm.proxy._experimental.mcp_server import server as mcp_module + + scope: Final[Scope] = {"type": "http", "method": "POST", "path": "/mcp", "headers": []} + receive: Final = AsyncMock( + side_effect=[ + {"type": "http.request", "body": request_body[:20], "more_body": True}, + {"type": "http.request", "body": request_body[20:], "more_body": False}, + ] + ) + send: Final = AsyncMock() + downstream_bodies: Final[list[bytes]] = [] + + async def handle_request(_: Scope, downstream_receive: Receive, __: Send) -> None: + downstream_bodies.append(await _drain_body(downstream_receive)) + + stateful_handle: Final = AsyncMock(side_effect=handle_request) + stateless_handle: Final = AsyncMock() + + with ( + _client_allowlist_patches(allowed_clients), + patch( + "litellm.proxy._experimental.mcp_server.server.session_manager_stateful", + SimpleNamespace(handle_request=stateful_handle), + ), + patch( + "litellm.proxy._experimental.mcp_server.server.session_manager_stateless", + SimpleNamespace(handle_request=stateless_handle), + ), + ): + await mcp_module.handle_streamable_http_mcp(scope, receive, send) + + assert downstream_bodies == [request_body] + stateless_handle.assert_not_awaited() + send.assert_not_awaited() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("allowed_clients", ([], "claude-code", [{"name": "claude-code"}])) +async def test_streamable_http_empty_or_malformed_allowlist_admits_nobody(allowed_clients: object) -> None: + from starlette.types import Scope + + from litellm.proxy._experimental.mcp_server import server as mcp_module + + scope: Final[Scope] = {"type": "http", "method": "POST", "path": "/mcp", "headers": []} + receive: Final = AsyncMock( + return_value={"type": "http.request", "body": _CLAUDE_CODE_INITIALIZE, "more_body": False} + ) + send: Final = AsyncMock() + stateful_handle: Final = AsyncMock() + + with ( + _client_allowlist_patches(allowed_clients), + patch( + "litellm.proxy._experimental.mcp_server.server.session_manager_stateful", + SimpleNamespace(handle_request=stateful_handle), + ), + ): + await mcp_module.handle_streamable_http_mcp(scope, receive, send) + + status, body = _forbidden_client_response(send) + assert status == 403 + assert body["details"] == "MCP client 'claude-code' is not listed in this gateway's mcp_allowed_clients." + stateful_handle.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_streamable_http_allowlist_only_inspects_initialize_requests() -> None: + from starlette.types import Receive, Scope, Send + + from litellm.proxy._experimental.mcp_server import server as mcp_module + + scope: Final[Scope] = {"type": "http", "method": "POST", "path": "/mcp", "headers": []} + receive: Final = AsyncMock(side_effect=[{"type": "http.request", "body": _TOOLS_LIST, "more_body": False}]) + send: Final = AsyncMock() + downstream_bodies: Final[list[bytes]] = [] + + async def handle_request(_: Scope, downstream_receive: Receive, __: Send) -> None: + downstream_bodies.append(await _drain_body(downstream_receive)) + + with ( + _client_allowlist_patches(["antigravity-cli"]), + patch( + "litellm.proxy._experimental.mcp_server.server.session_manager_stateless", + SimpleNamespace(handle_request=AsyncMock(side_effect=handle_request)), + ), + ): + await mcp_module.handle_streamable_http_mcp(scope, receive, send) + + assert downstream_bodies == [_TOOLS_LIST] + send.assert_not_awaited() + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("request_body", "admitted"), + ((_ANTIGRAVITY_INITIALIZE, True), (_CLAUDE_CODE_INITIALIZE, False), (_ANONYMOUS_INITIALIZE, False)), +) +async def test_sse_endpoint_applies_the_same_client_allowlist(request_body: bytes, admitted: bool) -> None: + from starlette.types import Receive, Scope, Send + + from litellm.proxy._experimental.mcp_server import server as mcp_module + + scope: Final[Scope] = {"type": "http", "method": "POST", "path": "/mcp/sse", "headers": []} + receive: Final = AsyncMock(side_effect=[{"type": "http.request", "body": request_body, "more_body": False}]) + send: Final = AsyncMock() + downstream_bodies: Final[list[bytes]] = [] + + async def handle_request(_: Scope, downstream_receive: Receive, __: Send) -> None: + downstream_bodies.append(await _drain_body(downstream_receive)) + + with ( + _client_allowlist_patches(["antigravity-cli"]), + patch( + "litellm.proxy._experimental.mcp_server.server._raise_preemptive_401_for_unauthenticated_servers", + new_callable=AsyncMock, + ), + patch( + "litellm.proxy._experimental.mcp_server.server._check_passthrough_upstream_auth", + new_callable=AsyncMock, + ), + patch.object(mcp_module.sse_session_manager, "handle_request", side_effect=handle_request), + ): + await mcp_module.handle_sse_mcp(scope, receive, send) + + if admitted: + assert downstream_bodies == [request_body] + send.assert_not_awaited() + return + assert downstream_bodies == [] + status, body = _forbidden_client_response(send) + assert status == 403 + assert body["error"] == "Forbidden" + assert "mcp_allowed_clients" in body["details"] + + @pytest.mark.asyncio async def test_mcp_routing_chunked_initialize_to_stateful(): """ diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 41c4956dba6..4a0f543dc38 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -7428,10 +7428,18 @@ async def test_update_general_settings_keeps_yaml_pass_through_endpoints_next_to request.query_params = {} return request - settings: Final = patch("litellm.proxy.proxy_server.general_settings", {"pass_through_endpoints": [yaml_endpoint]}) # test-quality-ok: the method reads this module global; no injection seam - yaml_endpoints: Final = patch("litellm.proxy.proxy_server.config_passthrough_endpoints", [yaml_endpoint]) # test-quality-ok: module global holding the YAML endpoints the fix merges in - initialize: Final = patch("litellm.proxy.proxy_server.initialize_pass_through_endpoints", AsyncMock()) # test-quality-ok: route registration needs the FastAPI app; auth is the observable here - master_key: Final = patch("litellm.proxy.proxy_server.master_key", "sk-master") # test-quality-ok: a set master key is what makes a missing Authorization header a 401 + settings: Final = patch( + "litellm.proxy.proxy_server.general_settings", {"pass_through_endpoints": [yaml_endpoint]} + ) # test-quality-ok: the method reads this module global; no injection seam + yaml_endpoints: Final = patch( + "litellm.proxy.proxy_server.config_passthrough_endpoints", [yaml_endpoint] + ) # test-quality-ok: module global holding the YAML endpoints the fix merges in + initialize: Final = patch( + "litellm.proxy.proxy_server.initialize_pass_through_endpoints", AsyncMock() + ) # test-quality-ok: route registration needs the FastAPI app; auth is the observable here + master_key: Final = patch( + "litellm.proxy.proxy_server.master_key", "sk-master" + ) # test-quality-ok: a set master key is what makes a missing Authorization header a 401 with settings, yaml_endpoints, initialize, master_key: await ProxyConfig()._update_general_settings(db_general_settings={"pass_through_endpoints": [db_endpoint]}) @@ -7479,10 +7487,18 @@ async def test_update_general_settings_db_pass_through_endpoint_overrides_yaml_e request.headers = {} request.query_params = {} - settings: Final = patch("litellm.proxy.proxy_server.general_settings", {"pass_through_endpoints": [yaml_endpoint]}) # test-quality-ok: the method reads this module global; no injection seam - yaml_endpoints: Final = patch("litellm.proxy.proxy_server.config_passthrough_endpoints", [yaml_endpoint]) # test-quality-ok: module global holding the YAML endpoints the fix merges in - initialize: Final = patch("litellm.proxy.proxy_server.initialize_pass_through_endpoints", AsyncMock()) # test-quality-ok: route registration needs the FastAPI app; auth is the observable here - master_key: Final = patch("litellm.proxy.proxy_server.master_key", "sk-master") # test-quality-ok: a set master key is what makes a missing Authorization header a 401 + settings: Final = patch( + "litellm.proxy.proxy_server.general_settings", {"pass_through_endpoints": [yaml_endpoint]} + ) # test-quality-ok: the method reads this module global; no injection seam + yaml_endpoints: Final = patch( + "litellm.proxy.proxy_server.config_passthrough_endpoints", [yaml_endpoint] + ) # test-quality-ok: module global holding the YAML endpoints the fix merges in + initialize: Final = patch( + "litellm.proxy.proxy_server.initialize_pass_through_endpoints", AsyncMock() + ) # test-quality-ok: route registration needs the FastAPI app; auth is the observable here + master_key: Final = patch( + "litellm.proxy.proxy_server.master_key", "sk-master" + ) # test-quality-ok: a set master key is what makes a missing Authorization header a 401 with settings, yaml_endpoints, initialize, master_key: await ProxyConfig()._update_general_settings(db_general_settings={"pass_through_endpoints": [db_endpoint]}) @@ -8597,9 +8613,7 @@ async def test_increment_spend_counters_finalizes_after_unreserved_increments(): async def assert_reservation_not_finalized_yet(**kwargs): assert budget_reservation["finalized"] is False incremented_counters.append(kwargs["counter_key"]) - return ps.PendingSpendIncrement( - counter_key=kwargs["counter_key"], increment=kwargs["increment"] - ) + return ps.PendingSpendIncrement(counter_key=kwargs["counter_key"], increment=kwargs["increment"]) import litellm.proxy.proxy_server as ps @@ -10144,9 +10158,15 @@ async def _lit6973_drive_realtime_session( side_effect=pre_call_error, return_value=({"model": "vertex_ai/gemini-live-2.5-flash"}, logging_obj) ) ws: Final = websocket if websocket is not None else _lit6973_fake_realtime_ws() - can_call = patch.object(ps, "can_key_call_resolved_model", new=AsyncMock(side_effect=model_access_error)) # test-quality-ok: no HTTP boundary; fakes in-process auth to reach the exit under test - pre = patch.object(ps.ProxyBaseLLMRequestProcessing, "common_processing_pre_call_logic", new=pre_call) # test-quality-ok: fakes phase-1 wiring; assertion checks observable reservation state - route = patch.object(ps, "route_request", new=AsyncMock(return_value=fake_llm_call())) # test-quality-ok: fakes the relay whose success/refusal outcome the endpoint reads off the logging object + can_call = patch.object( + ps, "can_key_call_resolved_model", new=AsyncMock(side_effect=model_access_error) + ) # test-quality-ok: no HTTP boundary; fakes in-process auth to reach the exit under test + pre = patch.object( + ps.ProxyBaseLLMRequestProcessing, "common_processing_pre_call_logic", new=pre_call + ) # test-quality-ok: fakes phase-1 wiring; assertion checks observable reservation state + route = patch.object( + ps, "route_request", new=AsyncMock(return_value=fake_llm_call()) + ) # test-quality-ok: fakes the relay whose success/refusal outcome the endpoint reads off the logging object with can_call, pre, route: await ps.realtime_websocket_endpoint( websocket=ws, @@ -10278,13 +10298,9 @@ async def _lit6463_drive_realtime_session_holding_a_max_parallel_slot( from litellm.proxy.utils import InternalUsageCache dual_cache: Final = DualCache() - await dual_cache.async_set_cache( - key=_LIT6463_COUNTER_KEY, value={"slot-1": 1.0, "slot-2": 2.0}, local_only=True - ) + await dual_cache.async_set_cache(key=_LIT6463_COUNTER_KEY, value={"slot-1": 1.0, "slot-2": 2.0}, local_only=True) limiter: Final = _PROXY_MaxParallelRequestsHandler_v3(internal_usage_cache=InternalUsageCache(dual_cache)) - stash: Final = RequestRateLimiterStash( - parallel_slot={"slot_id": "slot-1", "counter_keys": [_LIT6463_COUNTER_KEY]} - ) + stash: Final = RequestRateLimiterStash(parallel_slot={"slot_id": "slot-1", "counter_keys": [_LIT6463_COUNTER_KEY]}) reservation: Final = {"reserved_cost": 0.55, "input_cost": 0.0, "finalized": False, "entries": []} stash_token: Final = _request_stash.set(stash) @@ -10336,9 +10352,7 @@ async def test_successful_realtime_session_leaves_the_max_parallel_slot_for_the_ limiter's integer in-memory fallback, double-decrement the counter so the key admits more sessions than max_parallel_requests allows. With the success stamp present the route leaves the slot and the stash alone.""" - dual_cache, stash = await _lit6463_drive_realtime_session_holding_a_max_parallel_slot( - backend_logged_success=True - ) + dual_cache, stash = await _lit6463_drive_realtime_session_holding_a_max_parallel_slot(backend_logged_success=True) assert await dual_cache.async_get_cache(key=_LIT6463_COUNTER_KEY, local_only=True) == { "slot-1": 1.0, @@ -10384,8 +10398,12 @@ async def test_release_or_invalidate_falls_back_to_invalidating_the_counters(): async def _record(counter_key: str) -> None: invalidated.append(counter_key) - failing_release = patch.object(br, "release_budget_reservation", new=AsyncMock(side_effect=RuntimeError("counter store down"))) # test-quality-ok: forces the failure branch; assertion observes which counter key got invalidated - sink = patch.object(ps, "_invalidate_spend_counter", new=_record) # test-quality-ok: fakes the counter-store sink so the invalidated key is observable + failing_release = patch.object( + br, "release_budget_reservation", new=AsyncMock(side_effect=RuntimeError("counter store down")) + ) # test-quality-ok: forces the failure branch; assertion observes which counter key got invalidated + sink = patch.object( + ps, "_invalidate_spend_counter", new=_record + ) # test-quality-ok: fakes the counter-store sink so the invalidated key is observable with failing_release, sink: await br.release_or_invalidate_budget_reservation(budget_reservation=reservation) @@ -10401,8 +10419,12 @@ async def test_release_or_invalidate_finalizes_even_when_the_invalidate_fallback from litellm.proxy.spend_tracking import budget_reservation as br reservation: Final = {"reserved_cost": 0.55, "input_cost": 0.0, "finalized": False, "entries": []} - failing_release = patch.object(br, "release_budget_reservation", new=AsyncMock(side_effect=RuntimeError("counter store down"))) # test-quality-ok: forces the fallback branch - failing_invalidate = patch.object(br, "invalidate_budget_reservation_counters", new=AsyncMock(side_effect=RuntimeError("still down"))) # test-quality-ok: forces the fallback itself to fail + failing_release = patch.object( + br, "release_budget_reservation", new=AsyncMock(side_effect=RuntimeError("counter store down")) + ) # test-quality-ok: forces the fallback branch + failing_invalidate = patch.object( + br, "invalidate_budget_reservation_counters", new=AsyncMock(side_effect=RuntimeError("still down")) + ) # test-quality-ok: forces the fallback itself to fail with failing_release, failing_invalidate: await br.release_or_invalidate_budget_reservation(budget_reservation=reservation) @@ -12987,9 +13009,15 @@ async def test_moderations_response_carries_litellm_call_id_header(): user_api_key_dict = UserAPIKeyAuth(api_key="sk-test", spend=0.0) with ( - patch.object(proxy_server_module, "add_litellm_data_to_request", new=passthrough_add_litellm_data), # test-quality-ok: the route reads this module global, no injection point - patch.object(proxy_server_module, "route_request", new=AsyncMock(return_value=fake_llm_call())), # test-quality-ok: fakes the provider call so the response headers assembled by the real route are observable - patch.object(proxy_server_module, "proxy_logging_obj") as mock_logging, # test-quality-ok: module global, no injection point + patch.object( + proxy_server_module, "add_litellm_data_to_request", new=passthrough_add_litellm_data + ), # test-quality-ok: the route reads this module global, no injection point + patch.object( + proxy_server_module, "route_request", new=AsyncMock(return_value=fake_llm_call()) + ), # test-quality-ok: fakes the provider call so the response headers assembled by the real route are observable + patch.object( + proxy_server_module, "proxy_logging_obj" + ) as mock_logging, # test-quality-ok: module global, no injection point ): mock_logging.pre_call_hook = AsyncMock(side_effect=lambda user_api_key_dict, data, call_type: data) mock_logging.update_request_status = AsyncMock() @@ -13026,9 +13054,15 @@ async def test_moderations_failure_log_carries_the_callers_litellm_call_id(caplo verbose_proxy_logger.propagate = True try: with ( - patch.object(proxy_server_module, "add_litellm_data_to_request", new=passthrough_add_litellm_data), # test-quality-ok: the route reads this module global, no injection point - patch.object(proxy_server_module, "route_request", new=AsyncMock(side_effect=Exception("bad key"))), # test-quality-ok: fakes the provider failure so the real route's error log is observable - patch.object(proxy_server_module, "proxy_logging_obj", new=fake_logging), # test-quality-ok: module global, no injection point + patch.object( + proxy_server_module, "add_litellm_data_to_request", new=passthrough_add_litellm_data + ), # test-quality-ok: the route reads this module global, no injection point + patch.object( + proxy_server_module, "route_request", new=AsyncMock(side_effect=Exception("bad key")) + ), # test-quality-ok: fakes the provider failure so the real route's error log is observable + patch.object( + proxy_server_module, "proxy_logging_obj", new=fake_logging + ), # test-quality-ok: module global, no injection point caplog.at_level(logging.ERROR, logger="LiteLLM Proxy"), pytest.raises(ProxyException) as raised, ): @@ -13061,7 +13095,9 @@ async def test_moderations_unparseable_body_bills_the_callers_litellm_call_id(): fake_logging.post_call_failure_hook = AsyncMock() with ( - patch.object(proxy_server_module, "proxy_logging_obj", new=fake_logging), # test-quality-ok: module global, no injection point + patch.object( + proxy_server_module, "proxy_logging_obj", new=fake_logging + ), # test-quality-ok: module global, no injection point pytest.raises(ProxyException) as raised, ): await proxy_server_module.moderations( @@ -13089,8 +13125,12 @@ async def test_moderations_already_shaped_failure_answers_with_the_callers_litel fake_logging.post_call_failure_hook = AsyncMock() with ( - patch.object(proxy_server_module, "add_litellm_data_to_request", new=AsyncMock(side_effect=exc)), # test-quality-ok: the route reads this module global, no injection point - patch.object(proxy_server_module, "proxy_logging_obj", new=fake_logging), # test-quality-ok: module global, no injection point + patch.object( + proxy_server_module, "add_litellm_data_to_request", new=AsyncMock(side_effect=exc) + ), # test-quality-ok: the route reads this module global, no injection point + patch.object( + proxy_server_module, "proxy_logging_obj", new=fake_logging + ), # test-quality-ok: module global, no injection point pytest.raises(ProxyException) as raised, ): await proxy_server_module.moderations( @@ -13125,8 +13165,12 @@ async def test_audio_speech_already_shaped_failure_answers_with_the_callers_lite fake_logging.post_call_failure_hook = AsyncMock() with ( - patch.object(proxy_server_module, "add_litellm_data_to_request", new=AsyncMock(side_effect=exc)), # test-quality-ok: the route reads this module global, no injection point - patch.object(proxy_server_module, "proxy_logging_obj", new=fake_logging), # test-quality-ok: module global, no injection point + patch.object( + proxy_server_module, "add_litellm_data_to_request", new=AsyncMock(side_effect=exc) + ), # test-quality-ok: the route reads this module global, no injection point + patch.object( + proxy_server_module, "proxy_logging_obj", new=fake_logging + ), # test-quality-ok: module global, no injection point pytest.raises(type(exc)) as raised, ): await proxy_server_module.audio_speech( @@ -13814,6 +13858,43 @@ async def test_update_general_settings_keeps_yaml_openai_websocket_passthrough() assert ps.general_settings["enable_openai_websocket_passthrough"] is False +@pytest.mark.asyncio +@pytest.mark.parametrize( + "db_general_settings, expected", + [ + ({"mcp_allowed_clients": ["antigravity-cli"]}, ["antigravity-cli"]), + ({"mcp_allowed_clients": []}, []), + ({}, None), + ], +) +async def test_update_general_settings_propagates_mcp_allowed_clients(db_general_settings, expected): + from litellm.proxy.proxy_server import ProxyConfig + + proxy_config = ProxyConfig() + + with patch("litellm.proxy.proxy_server.general_settings", {"mcp_allowed_clients": ["claude-code"]}): + await proxy_config._update_general_settings(db_general_settings=db_general_settings) + + import litellm.proxy.proxy_server as ps + + assert ps.general_settings["mcp_allowed_clients"] == expected + + +@pytest.mark.asyncio +async def test_update_general_settings_keeps_yaml_mcp_allowed_clients(): + from litellm.proxy.proxy_server import ProxyConfig + + proxy_config = ProxyConfig() + proxy_config._yaml_general_settings_keys = {"mcp_allowed_clients"} + + with patch("litellm.proxy.proxy_server.general_settings", {"mcp_allowed_clients": ["claude-code"]}): + await proxy_config._update_general_settings(db_general_settings={"mcp_allowed_clients": ["codex-mcp-client"]}) + + import litellm.proxy.proxy_server as ps + + assert ps.general_settings["mcp_allowed_clients"] == ["claude-code"] + + async def test_token_counter_keeps_the_event_loop_free_during_a_huggingface_count(monkeypatch): from tests.large_text import text from tests.test_litellm.litellm_core_utils.event_loop_lag import ( @@ -13855,14 +13936,18 @@ async def test_token_counter_loads_a_custom_tokenizer_off_the_event_loop(monkeyp { "model_name": "self-hosted", "litellm_params": {"model": "openai/self-hosted-model", "api_base": "http://localhost:8080/v1"}, - "model_info": {"custom_tokenizer": {"identifier": "my-org/tokenizer", "revision": "main", "auth_token": None}}, + "model_info": { + "custom_tokenizer": {"identifier": "my-org/tokenizer", "revision": "main", "auth_token": None} + }, } ] ), ) response, took, lags = await timed_with_loop_lags( - lambda: proxy_server_module.token_counter(TokenCountRequest(model="self-hosted", prompt="count me off the loop")) + lambda: proxy_server_module.token_counter( + TokenCountRequest(model="self-hosted", prompt="count me off the loop") + ) ) assert response.tokenizer_type == "huggingface_tokenizer" diff --git a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPNetworkSettings.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPNetworkSettings.test.tsx index 9526f5de074..b6521acd7e2 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPNetworkSettings.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPNetworkSettings.test.tsx @@ -93,7 +93,7 @@ describe("MCPNetworkSettings", () => { await waitFor(() => expect(updateConfigFieldSetting).toHaveBeenCalledWith("tok", "mcp_internal_ip_ranges", ["10.0.0.0/8"]), ); - expect(deleteConfigFieldSetting).not.toHaveBeenCalled(); + expect(deleteConfigFieldSetting).not.toHaveBeenCalledWith("tok", "mcp_internal_ip_ranges"); }); it("clears the setting instead of saving an empty list", async () => { @@ -103,4 +103,68 @@ describe("MCPNetworkSettings", () => { await waitFor(() => expect(deleteConfigFieldSetting).toHaveBeenCalledWith("tok", "mcp_internal_ip_ranges")); expect(updateConfigFieldSetting).not.toHaveBeenCalled(); }); + + it("renders the stored allowed client names once settings load", async () => { + vi.mocked(getGeneralSettingsCall).mockResolvedValue([ + { field_name: "mcp_allowed_clients", field_value: ["antigravity-cli", "codex-mcp-client"] }, + ]); + + renderSettings(); + + expect(await screen.findByText("antigravity-cli")).toBeInTheDocument(); + expect(screen.getByText("codex-mcp-client")).toBeInTheDocument(); + }); + + it("adds typed client names on Enter and saves them under mcp_allowed_clients", async () => { + renderSettings(); + const input = await screen.findByRole("textbox", { name: "Allowed client names" }); + + await userEvent.type(input, "antigravity-cli, codex-mcp-client{Enter}"); + + expect(screen.getByText("antigravity-cli")).toBeInTheDocument(); + expect(screen.getByText("codex-mcp-client")).toBeInTheDocument(); + expect(input).toHaveValue(""); + + await userEvent.click(screen.getByRole("button", { name: /Save/ })); + + await waitFor(() => + expect(updateConfigFieldSetting).toHaveBeenCalledWith("tok", "mcp_allowed_clients", [ + "antigravity-cli", + "codex-mcp-client", + ]), + ); + expect(deleteConfigFieldSetting).not.toHaveBeenCalledWith("tok", "mcp_allowed_clients"); + }); + + it("removes a client name and clears the setting when the list becomes empty", async () => { + vi.mocked(getGeneralSettingsCall).mockResolvedValue([ + { field_name: "mcp_allowed_clients", field_value: ["claude-code"] }, + ]); + + renderSettings(); + await userEvent.click(await screen.findByRole("button", { name: "Remove claude-code" })); + + expect(screen.queryByText("claude-code")).not.toBeInTheDocument(); + + await userEvent.click(screen.getByRole("button", { name: /Save/ })); + + await waitFor(() => expect(deleteConfigFieldSetting).toHaveBeenCalledWith("tok", "mcp_allowed_clients")); + expect(updateConfigFieldSetting).not.toHaveBeenCalledWith("tok", "mcp_allowed_clients", expect.anything()); + }); + + it("keeps the private ranges and the allowed clients as independent settings on save", async () => { + vi.mocked(getGeneralSettingsCall).mockResolvedValue([ + { field_name: "mcp_internal_ip_ranges", field_value: ["10.0.0.0/8"] }, + { field_name: "mcp_allowed_clients", field_value: ["antigravity-cli"] }, + ]); + + renderSettings(); + await userEvent.click(await screen.findByRole("button", { name: /Save/ })); + + await waitFor(() => + expect(updateConfigFieldSetting).toHaveBeenCalledWith("tok", "mcp_allowed_clients", ["antigravity-cli"]), + ); + expect(updateConfigFieldSetting).toHaveBeenCalledWith("tok", "mcp_internal_ip_ranges", ["10.0.0.0/8"]); + expect(deleteConfigFieldSetting).not.toHaveBeenCalled(); + }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPNetworkSettings.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPNetworkSettings.tsx index 8b4d2a58652..18377a4bb82 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPNetworkSettings.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPNetworkSettings.tsx @@ -30,8 +30,10 @@ const MCPNetworkSettings: React.FC = ({ accessToken }) const [loading, setLoading] = useState(true); const [saving, setSaving] = useState(false); const [privateRanges, setPrivateRanges] = useState([]); + const [allowedClients, setAllowedClients] = useState([]); const [currentIp, setCurrentIp] = useState(null); const [rangeDraft, setRangeDraft] = useState(""); + const [clientDraft, setClientDraft] = useState(""); useEffect(() => { loadSettings(); @@ -47,6 +49,9 @@ const MCPNetworkSettings: React.FC = ({ accessToken }) if (field.field_name === "mcp_internal_ip_ranges" && field.field_value) { setPrivateRanges(field.field_value); } + if (field.field_name === "mcp_allowed_clients" && field.field_value) { + setAllowedClients(field.field_value); + } } } catch (error) { console.error("Failed to load MCP network settings:", error); @@ -72,6 +77,11 @@ const MCPNetworkSettings: React.FC = ({ accessToken }) } else { await deleteConfigFieldSetting(accessToken, "mcp_internal_ip_ranges"); } + if (allowedClients.length > 0) { + await updateConfigFieldSetting(accessToken, "mcp_allowed_clients", allowedClients); + } else { + await deleteConfigFieldSetting(accessToken, "mcp_allowed_clients"); + } } catch (error) { console.error("Failed to save MCP network settings:", error); } finally { @@ -86,17 +96,28 @@ const MCPNetworkSettings: React.FC = ({ accessToken }) }; // Commas separate entries, matching the old tokenised input. - const commitDraft = () => { - const added = rangeDraft + const splitDraft = (draft: string, existing: string[]) => + draft .split(",") .map((r) => r.trim()) - .filter((r) => r !== "" && !privateRanges.includes(r)); + .filter((r) => r !== "" && !existing.includes(r)); + + const commitDraft = () => { + const added = splitDraft(rangeDraft, privateRanges); if (added.length > 0) { setPrivateRanges([...privateRanges, ...added]); } setRangeDraft(""); }; + const commitClientDraft = () => { + const added = splitDraft(clientDraft, allowedClients); + if (added.length > 0) { + setAllowedClients([...allowedClients, ...added]); + } + setClientDraft(""); + }; + if (loading) { return (
@@ -178,6 +199,56 @@ const MCPNetworkSettings: React.FC = ({ accessToken })

+
+

Allowed Client Applications

+

+ Only the MCP client applications listed here can connect to the gateway. Names are matched exactly against + the clientInfo.name each client sends in its MCP initialize request (for example claude-code or + codex-mcp-client). Leave empty to allow every client. Clients choose the name they send, so treat this as a + policy control rather than a security boundary. +

+
+ + +
+

Allowed Client Names

+
+ {allowedClients.length > 0 && ( +
+ {allowedClients.map((client) => ( + + {client} + + + ))} +
+ )} + setClientDraft(e.target.value)} + onBlur={commitClientDraft} + onKeyDown={(e) => { + if (e.key === "Enter" || e.key === ",") { + e.preventDefault(); + commitClientDraft(); + } + }} + /> +

+ Enter the clientInfo.name values to admit. Any other client, or one that does not identify itself, gets a + 403 on its MCP initialize request. +

+ +
- + + + ))}
)} - setClientDraft(e.target.value)} - onBlur={commitClientDraft} - onKeyDown={(e) => { - if (e.key === "Enter" || e.key === ",") { - e.preventDefault(); - commitClientDraft(); - } - }} - /> +

- Enter the exact JWT claim or header values to admit. Every MCP request from any other client, or from one with - no resolvable identity, gets a 403. + The alias is the name shown here and in gateway logs. The value is the exact JWT claim or header value that + identifies the client, such as the OAuth client ID your identity provider issues. Leave the list empty to + allow every client. Every MCP request from an unlisted client, or from one with no resolvable identity, gets a + 403.

diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 19c0e29f17c..45a11aff0bc 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -26849,9 +26849,9 @@ export interface components { maximum_spend_logs_retention_period?: string | null; /** * Mcp Allowed Clients - * @description MCP client applications admitted by the gateway. When set, every MCP request must carry a client identity that matches one of these values exactly: a JWT caller is identified by the claim named in litellm_jwtauth.mcp_client_id_jwt_field, any other caller by the header named in mcp_client_id_header. A request with no resolvable identity, or an unlisted one, is rejected with 403. Unset means every client is admitted. + * @description MCP client applications admitted by the gateway, each an {alias, value} pair where alias is the name shown in the dashboard and logs and value is the identity that must match exactly. When set, every MCP request must carry a client identity equal to one of the values: a JWT caller is identified by the claim named in litellm_jwtauth.mcp_client_id_jwt_field, any other caller by the header named in mcp_client_id_header. A request with no resolvable identity, or an unlisted one, is rejected with 403. Unset means every client is admitted. */ - mcp_allowed_clients?: string[] | null; + mcp_allowed_clients?: components["schemas"]["MCPAllowedClient"][] | null; /** * Mcp Client Id Header * @description Request header whose value names the calling MCP client application (for example 'x-mcp-client') for callers that did not authenticate with a JWT, used only while mcp_allowed_clients is set. The client picks this value itself, so it is a policy control rather than a security boundary; prefer litellm_jwtauth.mcp_client_id_jwt_field where callers use JWTs. @@ -32491,6 +32491,22 @@ export interface components { */ status?: "healthy" | "unhealthy"; }; + /** + * MCPAllowedClient + * @description One entry of `general_settings.mcp_allowed_clients`. + */ + MCPAllowedClient: { + /** + * Alias + * @description Human-readable name for this client application, shown in the dashboard and in gateway logs. + */ + alias: string; + /** + * Value + * @description Exact value of the JWT claim named in litellm_jwtauth.mcp_client_id_jwt_field, or of the mcp_client_id_header header, that identifies this client application. Matched case-sensitively. + */ + value: string; + }; /** MCPConnectorEntry */ MCPConnectorEntry: { /** Args */ From 7dede188f8dbed8b1815791d3d06a29de031b5de Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 18 Sep 2026 22:30:12 +0000 Subject: [PATCH 068/144] test(timing): type the logging object test helper Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com> --- .../llm_response_utils/test_response_metadata.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/tests/test_litellm/litellm_core_utils/llm_response_utils/test_response_metadata.py b/tests/test_litellm/litellm_core_utils/llm_response_utils/test_response_metadata.py index 40f964cd3fc..eeccbc719d3 100644 --- a/tests/test_litellm/litellm_core_utils/llm_response_utils/test_response_metadata.py +++ b/tests/test_litellm/litellm_core_utils/llm_response_utils/test_response_metadata.py @@ -233,7 +233,12 @@ class TestResponseTimingMetrics: START = datetime.datetime(2025, 1, 1, 0, 0, 0) END = datetime.datetime(2025, 1, 1, 0, 0, 1) - def _make_logging_obj(self, llm_api_duration_ms=None, caching_details=None, received_at=None): + def _make_logging_obj( + self, + llm_api_duration_ms: float | None = None, + caching_details: dict[str, object] | None = None, + received_at: datetime.datetime | str | None = None, + ) -> MagicMock: logging_obj = MagicMock() logging_obj.model_call_details = {} if llm_api_duration_ms is not None: From 7f9db61528bf3fd8a83b93851a3d3d5070868329 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 18 Sep 2026 15:35:59 -0700 Subject: [PATCH 069/144] fix(vertex_ai): carry turns across stream rotation and route by model info Rotating the Speech-to-Text stream at 240 s no longer ends the active turn: the turn and its billed seconds continue on the new stream, forced at 280 s. Bound the request and event queues (64 and 256) so a slow peer applies backpressure instead of growing memory. Route a model to the Chirp realtime path from its cost-map entry (mode audio_transcription plus /v1/realtime) instead of a hardcoded name. Return on every branch of the recv and transform helpers (CodeQL mixed returns), have the shared protocol helper take the provider's error class so the Meta tests assert MuseProtocolError again, and pin google-cloud-speech in the ci group so unit shards import it. --- .../realtime/transcription_protocol.py | 84 ++++--- .../llms/base_llm/realtime/transformation.py | 3 +- litellm/llms/meta/realtime/transformation.py | 8 +- .../audio_transcription/realtime_backend.py | 209 +++++++++++------- .../realtime_transformation.py | 22 +- pyproject.toml | 1 + .../test_meta_realtime_transformation.py | 5 +- .../test_vertex_ai_realtime_backend.py | 162 ++++++++++++-- .../test_vertex_ai_realtime_transformation.py | 34 ++- uv.lock | 4 +- 10 files changed, 374 insertions(+), 158 deletions(-) diff --git a/litellm/llms/base_llm/realtime/transcription_protocol.py b/litellm/llms/base_llm/realtime/transcription_protocol.py index 05151bd4a78..c3264911d46 100644 --- a/litellm/llms/base_llm/realtime/transcription_protocol.py +++ b/litellm/llms/base_llm/realtime/transcription_protocol.py @@ -59,37 +59,46 @@ class TranscriptionSessionUpdate: return None if self.turn_detection is None else self.turn_detection.get("type") -def json_object(payload: str) -> Mapping[str, JsonValue]: +ProtocolErrorType = type[RealtimeTranscriptionProtocolError] + + +def json_object(payload: str, error: ProtocolErrorType = RealtimeTranscriptionProtocolError) -> Mapping[str, JsonValue]: try: value: Final = _JSON_ADAPTER.validate_json(payload) except ValidationError: - raise RealtimeTranscriptionProtocolError("invalid JSON object") from None + raise error("invalid JSON object") from None if not isinstance(value, dict): - raise RealtimeTranscriptionProtocolError("message must be a JSON object") + raise error("message must be a JSON object") return value -def json_mapping(value: JsonValue | None, name: str) -> Mapping[str, JsonValue]: +def json_mapping( + value: JsonValue | None, name: str, error: ProtocolErrorType = RealtimeTranscriptionProtocolError +) -> Mapping[str, JsonValue]: if value is None: return EMPTY_JSON_OBJECT if not isinstance(value, dict): - raise RealtimeTranscriptionProtocolError(f"{name} must be an object") + raise error(f"{name} must be an object") return value -def json_string(value: JsonValue | None, name: str) -> str | None: +def json_string( + value: JsonValue | None, name: str, error: ProtocolErrorType = RealtimeTranscriptionProtocolError +) -> str | None: if value is None: return None if not isinstance(value, str): - raise RealtimeTranscriptionProtocolError(f"{name} must be a string") + raise error(f"{name} must be a string") return value -def json_integer(value: JsonValue | None, name: str) -> int | None: +def json_integer( + value: JsonValue | None, name: str, error: ProtocolErrorType = RealtimeTranscriptionProtocolError +) -> int | None: if value is None: return None if isinstance(value, bool) or not isinstance(value, int): - raise RealtimeTranscriptionProtocolError(f"{name} must be an integer") + raise error(f"{name} must be an integer") return value @@ -97,49 +106,52 @@ def new_event_id() -> str: return f"event_{uuid.uuid4().hex}" -def parse_transcription_session_update(payload: str) -> TranscriptionSessionUpdate: - message: Final = json_object(payload) +def parse_transcription_session_update( + payload: str, error: ProtocolErrorType = RealtimeTranscriptionProtocolError +) -> TranscriptionSessionUpdate: + message: Final = json_object(payload, error) if message.get("type") not in SESSION_UPDATE_EVENT_TYPES: - raise RealtimeTranscriptionProtocolError("expected session.update") - session: Final = json_mapping(message.get("session"), "session") + raise error("expected session.update") + session: Final = json_mapping(message.get("session"), "session", error) if not session: - raise RealtimeTranscriptionProtocolError("session.update requires a session object") - audio: Final = json_mapping(session.get("audio"), "session.audio") - audio_input: Final = json_mapping(audio.get("input"), "session.audio.input") + raise error("session.update requires a session object") + audio: Final = json_mapping(session.get("audio"), "session.audio", error) + audio_input: Final = json_mapping(audio.get("input"), "session.audio.input", error) beta_transcription: Final = session.get("input_audio_transcription") ga_transcription: Final = audio_input.get("transcription") if beta_transcription is not None and ga_transcription is not None: - raise RealtimeTranscriptionProtocolError("input transcription must use either beta or GA layout") + raise error("input transcription must use either beta or GA layout") transcription: Final = json_mapping( beta_transcription if beta_transcription is not None else ga_transcription, "input audio transcription", + error, ) turn_detection_present: Final = "turn_detection" in session or "turn_detection" in audio_input turn_detection: Final = session.get("turn_detection", audio_input.get("turn_detection")) return TranscriptionSessionUpdate( - session_type=json_string(session.get("type"), "session.type"), - audio_format=_parse_audio_format(session, audio_input), - model=json_string(transcription.get("model"), "transcription model"), - language=json_string(transcription.get("language"), "language"), + session_type=json_string(session.get("type"), "session.type", error), + audio_format=_parse_audio_format(session, audio_input, error), + model=json_string(transcription.get("model"), "transcription model", error), + language=json_string(transcription.get("language"), "language", error), unsupported_transcription_keys=tuple( sorted(key for key in transcription if key not in _SUPPORTED_TRANSCRIPTION_KEYS) ), - turn_detection=None if turn_detection is None else json_mapping(turn_detection, "turn_detection"), + turn_detection=None if turn_detection is None else json_mapping(turn_detection, "turn_detection", error), turn_detection_disabled=turn_detection_present and turn_detection is None, ) def _parse_audio_format( - session: Mapping[str, JsonValue], audio_input: Mapping[str, JsonValue] + session: Mapping[str, JsonValue], audio_input: Mapping[str, JsonValue], error: ProtocolErrorType ) -> TranscriptionAudioFormat | None: beta_format: Final = session.get("input_audio_format") ga_format: Final = audio_input.get("format") if beta_format is not None and ga_format is not None: - raise RealtimeTranscriptionProtocolError("input audio format must use either beta or GA layout") + raise error("input audio format must use either beta or GA layout") if beta_format is not None: return TranscriptionAudioFormat( layout="beta", - encoding=json_string(beta_format, "session.input_audio_format"), + encoding=json_string(beta_format, "session.input_audio_format", error), rate=None, channels=None, ) @@ -147,26 +159,30 @@ def _parse_audio_format( return None if isinstance(ga_format, str): return TranscriptionAudioFormat(layout="ga", encoding=ga_format, rate=None, channels=None) - format_mapping: Final = json_mapping(ga_format, "session.audio.input.format") + format_mapping: Final = json_mapping(ga_format, "session.audio.input.format", error) return TranscriptionAudioFormat( layout="ga", - encoding=json_string(format_mapping.get("type"), "session.audio.input.format.type"), - rate=json_integer(format_mapping.get("rate"), "session.audio.input.format.rate"), - channels=json_integer(format_mapping.get("channels"), "session.audio.input.format.channels"), + encoding=json_string(format_mapping.get("type"), "session.audio.input.format.type", error), + rate=json_integer(format_mapping.get("rate"), "session.audio.input.format.rate", error), + channels=json_integer(format_mapping.get("channels"), "session.audio.input.format.channels", error), ) -def decode_pcm16_append(audio: JsonValue | None, max_encoded_bytes: int | None = None) -> bytes: +def decode_pcm16_append( + audio: JsonValue | None, + max_encoded_bytes: int | None = None, + error: ProtocolErrorType = RealtimeTranscriptionProtocolError, +) -> bytes: if not isinstance(audio, str): - raise RealtimeTranscriptionProtocolError("Audio must be a base64 string") + raise error("Audio must be a base64 string") if max_encoded_bytes is not None and len(audio) > max_encoded_bytes: - raise RealtimeTranscriptionProtocolError("Audio append exceeds the four-second backlog limit") + raise error("Audio append exceeds the four-second backlog limit") try: decoded: Final = base64.b64decode(audio, validate=True) except (binascii.Error, ValueError): - raise RealtimeTranscriptionProtocolError("Audio must be valid base64") from None + raise error("Audio must be valid base64") from None if len(decoded) % 2: - raise RealtimeTranscriptionProtocolError("PCM16 audio must contain complete samples") + raise error("PCM16 audio must contain complete samples") return decoded diff --git a/litellm/llms/base_llm/realtime/transformation.py b/litellm/llms/base_llm/realtime/transformation.py index e1b16a5985d..1f4ad29fa74 100644 --- a/litellm/llms/base_llm/realtime/transformation.py +++ b/litellm/llms/base_llm/realtime/transformation.py @@ -1,9 +1,10 @@ from abc import ABC, abstractmethod from collections.abc import Mapping, Sequence from types import TracebackType -from typing import TYPE_CHECKING, Any, Protocol, Self +from typing import TYPE_CHECKING, Any, Protocol import httpx +from typing_extensions import Self from litellm.types.llms.openai import OpenAIRealtimeStreamSessionEvents from litellm.types.realtime import ( diff --git a/litellm/llms/meta/realtime/transformation.py b/litellm/llms/meta/realtime/transformation.py index f0ca448bea0..442c79255af 100644 --- a/litellm/llms/meta/realtime/transformation.py +++ b/litellm/llms/meta/realtime/transformation.py @@ -239,7 +239,7 @@ def _parse_mode(update: TranscriptionSessionUpdate) -> MuseMode: def parse_session_update(payload: str, expected_model: str) -> MuseSessionConfig: - update: Final = parse_transcription_session_update(payload) + update: Final = parse_transcription_session_update(payload, MuseProtocolError) if update.session_type not in (None, "transcription", "realtime"): raise MuseProtocolError("Muse Voice supports transcription sessions only") if update.unsupported_transcription_keys: @@ -486,7 +486,7 @@ class MetaRealtimeConfig(BaseRealtimeConfig): model: str, session_configuration_request: str | None = None, ) -> tuple[str | bytes, ...]: - request: Final = json_object(message) + request: Final = json_object(message, MuseProtocolError) event_type: Final = request.get("type") if event_type in ("session.update", "transcription_session.update"): return self._configure(message, model) @@ -538,7 +538,7 @@ class MetaRealtimeConfig(BaseRealtimeConfig): return result def _backend_events(self, payload: str) -> tuple[OpenAIRealtimeEvents, ...]: - frame: Final = json_object(payload) + frame: Final = json_object(payload, MuseProtocolError) session_id: Final = frame.get("sessionId") if session_id is None: return self._transformer.transform(frame) @@ -560,7 +560,7 @@ class MetaRealtimeConfig(BaseRealtimeConfig): def _append_audio(self, request: Mapping[str, JsonValue]) -> tuple[bytes, ...]: config: Final = self._require_config() - audio: Final = decode_pcm16_append(request.get("audio"), config.max_encoded_append_bytes) + audio: Final = decode_pcm16_append(request.get("audio"), config.max_encoded_append_bytes, MuseProtocolError) buffered: Final = self._pending_audio + audio packet_end: Final = len(buffered) - len(buffered) % config.packet_bytes self._pending_audio = buffered[packet_end:] diff --git a/litellm/llms/vertex_ai/audio_transcription/realtime_backend.py b/litellm/llms/vertex_ai/audio_transcription/realtime_backend.py index 0ecbe7209f6..87c72c193bc 100644 --- a/litellm/llms/vertex_ai/audio_transcription/realtime_backend.py +++ b/litellm/llms/vertex_ai/audio_transcription/realtime_backend.py @@ -4,9 +4,10 @@ from collections.abc import AsyncIterable, AsyncIterator, Awaitable, Callable from dataclasses import dataclass from datetime import timedelta from types import MappingProxyType, TracebackType -from typing import TYPE_CHECKING, Final, Literal, Protocol, Self +from typing import TYPE_CHECKING, Final, Literal, Protocol from pydantic import TypeAdapter +from typing_extensions import Self, assert_never from websockets.exceptions import ConnectionClosedError, ConnectionClosedOK from websockets.frames import Close @@ -37,6 +38,10 @@ SPEECH_SDK_INSTALL_HINT: Final = ( ) STREAM_FAILURE_CLOSE_CODE: Final = 1011 STREAM_ROTATION_SECONDS: Final = 240.0 +STREAM_ROTATION_DEADLINE_SECONDS: Final = 280.0 +REQUEST_QUEUE_SIZE: Final = 64 +OUTBOX_SIZE: Final = 256 +_LINK_QUEUE_SIZE: Final = 64 _CLOSE_REASON_MAX_CHARS: Final = 120 _CONFIGURED_EVENT: Final = VertexSpeechStreamingConfigured().model_dump_json() _TURN_FINISHED_EVENT: Final = VertexSpeechStreamingTurnFinished().model_dump_json() @@ -129,6 +134,10 @@ def _billed_seconds(response: "StreamingRecognizeResponse") -> float: return _TIMEDELTA_ADAPTER.validate_python(response.metadata.total_billed_duration).total_seconds() +def _normal_closure() -> ConnectionClosedOK: + return ConnectionClosedOK(rcvd=Close(1000, ""), sent=None) + + class _RecognizeStream: def __init__( self, @@ -136,63 +145,59 @@ class _RecognizeStream: client: SpeechStreamingClient, request_type: "type[StreamingRecognizeRequest]", first_request: "StreamingRecognizeRequest", - outbox: "asyncio.Queue[str | _StreamFailure | _Closed]", - previous: "_RecognizeStream | None", opened_at: float, ) -> None: self._client: Final = client self._request_type: Final = request_type - self._outbox: Final = outbox - self._previous: _RecognizeStream | None = previous self.opened_at: Final = opened_at - self._requests: Final[asyncio.Queue[StreamingRecognizeRequest | None]] = asyncio.Queue() + self._requests: Final[asyncio.Queue[StreamingRecognizeRequest | None]] = asyncio.Queue( + maxsize=REQUEST_QUEUE_SIZE + ) self._requests.put_nowait(first_request) - self._billed_seconds: float = 0.0 - self._base_billed_seconds: float = 0.0 - self._task: Final = asyncio.create_task(self._run()) + self.speech_active: bool = False + self.billed_seconds: float = 0.0 + self._cancelled: bool = False + self._task: asyncio.Task[None] | None = None - @property - def billed_seconds(self) -> float: - return self._base_billed_seconds + self._billed_seconds + async def send_audio(self, audio: bytes) -> None: + await self._requests.put(self._request_type(audio=audio)) - def send_audio(self, audio: bytes) -> None: - self._requests.put_nowait(self._request_type(audio=audio)) + async def half_close(self) -> None: + await self._requests.put(None) - def half_close(self) -> None: - self._requests.put_nowait(None) + def cancel(self) -> None: + self._cancelled = True + if self._task is not None: + self._task.cancel() - async def wait(self) -> None: - await asyncio.gather(self._task, return_exceptions=True) + async def relay(self, outbox: "asyncio.Queue[str | _StreamFailure | _Closed]", billed_before: float) -> float: + if self._cancelled: + return 0.0 + task: Final = asyncio.create_task(self._forward(outbox, billed_before)) + self._task = task + try: + await asyncio.wait((task,)) + except asyncio.CancelledError: + task.cancel() + await asyncio.wait((task,)) + raise + return self.billed_seconds - async def cancel(self) -> None: - self._task.cancel() - await self.wait() - - async def cancel_chain(self) -> None: - previous: Final = self._previous - self._task.cancel() - if previous is not None: - await previous.cancel_chain() - await self.wait() - - async def _run(self) -> None: - previous: Final = self._previous - if previous is not None: - await previous.wait() - self._base_billed_seconds = previous.billed_seconds - self._previous = None + async def _forward(self, outbox: "asyncio.Queue[str | _StreamFailure | _Closed]", billed_before: float) -> None: try: responses: Final = await self._client.streaming_recognize(self._drain()) async for response in responses: - self._billed_seconds = max(self._billed_seconds, _billed_seconds(response)) - await self._outbox.put(_response_event(response, self.billed_seconds)) - await self._outbox.put(_TURN_FINISHED_EVENT) - except asyncio.CancelledError: - self._outbox.put_nowait(_TURN_DISCARDED_EVENT) - raise + self._note(response) + await outbox.put(_response_event(response, billed_before + self.billed_seconds)) except Exception as e: # noqa: BLE001 # task boundary: a swallowed failure would hang the client session verbose_logger.warning("Google Speech-to-Text streaming failed: %s", e) - await self._outbox.put(_StreamFailure(reason=f"Google Speech-to-Text streaming failed: {e}")) + await outbox.put(_StreamFailure(reason=f"Google Speech-to-Text streaming failed: {e}")) + + def _note(self, response: "StreamingRecognizeResponse") -> None: + activity: Final = _SPEECH_EVENTS.get(response.speech_event_type.name) + if activity is not None: + self.speech_active = activity == "begin" + self.billed_seconds = max(self.billed_seconds, _billed_seconds(response)) async def _drain(self) -> "AsyncIterator[StreamingRecognizeRequest]": while (request := await self._requests.get()) is not None: @@ -207,16 +212,21 @@ class SpeechStreamingBackend: client_factory: Callable[[SpeechStreamingTarget], SpeechStreamingClient] = open_speech_client, clock: Callable[[], float] = time.monotonic, rotation_seconds: float = STREAM_ROTATION_SECONDS, + rotation_deadline_seconds: float = STREAM_ROTATION_DEADLINE_SECONDS, ) -> None: self._target: Final = target self._client_factory: Final = client_factory self._clock: Final = clock self._rotation_seconds: Final = rotation_seconds - self._outbox: Final[asyncio.Queue[str | _StreamFailure | _Closed]] = asyncio.Queue() + self._rotation_deadline_seconds: Final = rotation_deadline_seconds + self._outbox: Final[asyncio.Queue[str | _StreamFailure | _Closed]] = asyncio.Queue(maxsize=OUTBOX_SIZE) + self._links: Final[asyncio.Queue[_RecognizeStream | str]] = asyncio.Queue(maxsize=_LINK_QUEUE_SIZE) + self._pump: asyncio.Task[None] | None = None self._client: SpeechStreamingClient | None = None self._config: StreamingRecognitionConfig | None = None - self._stream: _RecognizeStream | None = None - self._last_stream: _RecognizeStream | None = None + self._turn: tuple[_RecognizeStream, ...] = () + self._billed_before: float = 0.0 + self._closed: bool = False async def __aenter__(self) -> Self: return self @@ -230,20 +240,26 @@ class SpeechStreamingBackend: await self.close() async def send(self, message: str | bytes) -> None: + if self._closed: + raise _normal_closure() if isinstance(message, bytes): - self._send_audio(message) + await self._send_audio(message) return command: Final = _COMMAND_ADAPTER.validate_json(message) match command: case VertexSpeechStreamingConfigure(): self._config = _streaming_config(command) - await self._outbox.put(_CONFIGURED_EVENT) + await self._link(_CONFIGURED_EVENT) case VertexSpeechStreamingFinishTurn(): - self._finish_turn() + await self._finish_turn() case VertexSpeechStreamingDiscardTurn(): await self._discard_turn() + case _: + assert_never(command) async def recv(self, decode: bool | None = None) -> str | bytes: + if self._closed and self._outbox.empty(): + raise _normal_closure() item: Final = await self._outbox.get() match item: case _StreamFailure(): @@ -251,35 +267,67 @@ class SpeechStreamingBackend: rcvd=Close(STREAM_FAILURE_CLOSE_CODE, item.reason[:_CLOSE_REASON_MAX_CHARS]), sent=None ) case _Closed(): - raise ConnectionClosedOK(rcvd=Close(1000, ""), sent=None) + raise _normal_closure() case str(): return item + case _: + assert_never(item) async def close(self) -> None: - self._stream = None - last_stream: Final = self._last_stream - self._last_stream = None - if last_stream is not None: - await last_stream.cancel_chain() + if self._closed: + return + self._closed = True + self._turn = () + pump: Final = self._pump + if pump is not None: + pump.cancel() + await asyncio.wait((pump,)) client: Final = self._client self._client = None if client is not None: await client.transport.close() - self._outbox.put_nowait(_Closed()) + if not self._outbox.full(): + self._outbox.put_nowait(_Closed()) - def _send_audio(self, audio: bytes) -> None: - self._rotate_expiring_stream() - stream: Final = self._stream if self._stream is not None else self._open_stream() - stream.send_audio(audio) + async def _link(self, item: _RecognizeStream | str) -> None: + if self._pump is None: + self._pump = asyncio.create_task(self._pump_links()) + await self._links.put(item) - def _rotate_expiring_stream(self) -> None: - stream: Final = self._stream - if stream is None or self._clock() - stream.opened_at < self._rotation_seconds: - return - self._stream = None - stream.half_close() + async def _pump_links(self) -> None: + while True: + await self._relay(await self._links.get()) - def _open_stream(self) -> _RecognizeStream: + async def _relay(self, link: _RecognizeStream | str) -> None: + match link: + case str(): + await self._outbox.put(link) + case _RecognizeStream(): + self._billed_before += await link.relay(self._outbox, self._billed_before) + case _: + assert_never(link) + + async def _send_audio(self, audio: bytes) -> None: + stream: Final = await self._turn_stream() + await stream.send_audio(audio) + + async def _turn_stream(self) -> _RecognizeStream: + current: Final = self._turn[-1] if self._turn else None + if current is not None and not self._expired(current): + return current + if current is not None: + await current.half_close() + stream: Final = await self._open_stream() + self._turn = (*self._turn, stream) + return stream + + def _expired(self, stream: _RecognizeStream) -> bool: + elapsed: Final = self._clock() - stream.opened_at + if elapsed >= self._rotation_deadline_seconds: + return True + return elapsed >= self._rotation_seconds and not stream.speech_active + + async def _open_stream(self) -> _RecognizeStream: from google.cloud.speech_v2.types import StreamingRecognizeRequest config: Final = self._config @@ -291,26 +339,21 @@ class SpeechStreamingBackend: client=self._client, request_type=StreamingRecognizeRequest, first_request=StreamingRecognizeRequest(recognizer=self._target.recognizer, streaming_config=config), - outbox=self._outbox, - previous=self._last_stream, opened_at=self._clock(), ) - self._stream = stream - self._last_stream = stream + await self._link(stream) return stream - def _finish_turn(self) -> None: - stream: Final = self._stream - self._stream = None - if stream is None: - self._outbox.put_nowait(_TURN_FINISHED_EVENT) - return - stream.half_close() + async def _finish_turn(self) -> None: + turn: Final = self._turn + self._turn = () + if turn: + await turn[-1].half_close() + await self._link(_TURN_FINISHED_EVENT) async def _discard_turn(self) -> None: - stream: Final = self._stream - self._stream = None - if stream is None: - await self._outbox.put(_TURN_DISCARDED_EVENT) - return - await stream.cancel() + turn: Final = self._turn + self._turn = () + for stream in turn: + stream.cancel() + await self._link(_TURN_DISCARDED_EVENT) diff --git a/litellm/llms/vertex_ai/audio_transcription/realtime_transformation.py b/litellm/llms/vertex_ai/audio_transcription/realtime_transformation.py index 5e2857f24c3..6ec7a21a134 100644 --- a/litellm/llms/vertex_ai/audio_transcription/realtime_transformation.py +++ b/litellm/llms/vertex_ai/audio_transcription/realtime_transformation.py @@ -3,7 +3,9 @@ from dataclasses import dataclass, replace from typing import Final from pydantic import JsonValue, TypeAdapter +from typing_extensions import assert_never +import litellm from litellm import verbose_logger from litellm._uuid import uuid from litellm.litellm_core_utils.audio_utils.utils import normalize_transcription_language_to_bcp47 @@ -56,7 +58,7 @@ DEFAULT_SAMPLE_RATE_HERTZ: Final = 24_000 MIN_SAMPLE_RATE_HERTZ: Final = 8_000 MAX_SAMPLE_RATE_HERTZ: Final = 48_000 MAX_AUDIO_MESSAGE_BYTES: Final = 25_000 -SPEECH_TO_TEXT_MODEL_PREFIX: Final = "chirp" +_SPEECH_TO_TEXT_ENDPOINTS: Final = frozenset({"/v1/audio/transcriptions", "/v1/realtime"}) _VERTEX_MODEL_PREFIX: Final = "vertex_ai/" _STREAMING_EVENT_ADAPTER: Final = TypeAdapter[VertexSpeechStreamingEventUnion](VertexSpeechStreamingEvent) _FINISH_TURN_COMMAND: Final = VertexSpeechStreamingFinishTurn().model_dump_json() @@ -99,7 +101,15 @@ class ChirpSessionConfig: def is_vertex_speech_to_text_model(model: str) -> bool: - return normalize_speech_to_text_model(model).startswith(SPEECH_TO_TEXT_MODEL_PREFIX) + try: + info: Final = litellm.get_model_info( + model=normalize_speech_to_text_model(model), custom_llm_provider="vertex_ai" + ) + except Exception: # noqa: BLE001 # get_model_info raises for unmapped models, which are not Speech-to-Text models + return False + if info.get("mode") != "audio_transcription": + return False + return _SPEECH_TO_TEXT_ENDPOINTS <= frozenset(info.get("supported_endpoints") or ()) def normalize_speech_to_text_model(model: str) -> str: @@ -116,7 +126,7 @@ def default_session_config(model: str) -> ChirpSessionConfig: def parse_chirp_session_update(payload: str, expected_model: str) -> ChirpSessionConfig: - update: Final = parse_transcription_session_update(payload) + update: Final = parse_transcription_session_update(payload, ChirpProtocolError) if update.session_type not in (None, "transcription", "realtime"): raise ChirpProtocolError("Speech-to-Text streaming supports transcription sessions only") if update.unsupported_transcription_keys: @@ -228,6 +238,8 @@ class ChirpEventTransformer: case VertexSpeechStreamingTurnDiscarded(): self._turn = None return () + case _: + assert_never(frame) def _response(self, frame: VertexSpeechStreamingResponse) -> tuple[OpenAIRealtimeEvents, ...]: self._billed_seconds = max(self._billed_seconds, frame.billed_seconds) @@ -367,7 +379,7 @@ class VertexChirpRealtimeConfig(BaseRealtimeConfig): model: str, session_configuration_request: str | None = None, ) -> tuple[str | bytes, ...]: - request: Final = json_object(message) + request: Final = json_object(message, ChirpProtocolError) event_type: Final = request.get("type") if event_type in ("session.update", "transcription_session.update"): return self._configure(message, model) @@ -417,7 +429,7 @@ class VertexChirpRealtimeConfig(BaseRealtimeConfig): def _append_audio(self, request: Mapping[str, JsonValue]) -> tuple[bytes, ...]: self._require_config() - audio: Final = decode_pcm16_append(request.get("audio")) + audio: Final = decode_pcm16_append(request.get("audio"), error=ChirpProtocolError) return tuple( audio[start : start + MAX_AUDIO_MESSAGE_BYTES] for start in range(0, len(audio), MAX_AUDIO_MESSAGE_BYTES) ) diff --git a/pyproject.toml b/pyproject.toml index c58fd1d2158..361a5309ef3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -277,6 +277,7 @@ ci = [ "langgraph>=1.2.4,<1.3.0", "langgraph-prebuilt>=1.1.0,<1.3.0", "claude-agent-sdk==0.1.44", + "google-cloud-speech==2.40.0", ] healthcheck = [ "httpx==0.28.1", diff --git a/tests/test_litellm/llms/meta/realtime/test_meta_realtime_transformation.py b/tests/test_litellm/llms/meta/realtime/test_meta_realtime_transformation.py index 1262d124d89..a5d7e47fb65 100644 --- a/tests/test_litellm/llms/meta/realtime/test_meta_realtime_transformation.py +++ b/tests/test_litellm/llms/meta/realtime/test_meta_realtime_transformation.py @@ -6,7 +6,6 @@ from unittest.mock import MagicMock import pytest -from litellm.llms.base_llm.realtime.transcription_protocol import RealtimeTranscriptionProtocolError from litellm.llms.meta.realtime.transformation import ( DEFAULT_MUSE_REALTIME_URL, MUSE_MODEL, @@ -161,7 +160,7 @@ def test_language_normalization_uses_official_muse_names(source: str, expected: ], ) def test_session_rejects_unsupported_audio_model_and_hints(session: dict[str, object], message: str): - with pytest.raises(RealtimeTranscriptionProtocolError, match=message): + with pytest.raises(MuseProtocolError, match=message): parse_session_update(_event("session.update", session={"type": "transcription", **session}), MUSE_MODEL) @@ -585,7 +584,7 @@ def test_pcm_is_packetized_into_raw_binary_frames(rate: int, packet_bytes: int): def test_invalid_audio_appends_are_rejected(audio: object, message: str): config = _configured() - with pytest.raises(RealtimeTranscriptionProtocolError, match=message): + with pytest.raises(MuseProtocolError, match=message): config.transform_realtime_request(_event("input_audio_buffer.append", audio=audio), MUSE_MODEL) diff --git a/tests/test_litellm/llms/vertex_ai/audio_transcription/test_vertex_ai_realtime_backend.py b/tests/test_litellm/llms/vertex_ai/audio_transcription/test_vertex_ai_realtime_backend.py index 87ab3da37a5..49758e62415 100644 --- a/tests/test_litellm/llms/vertex_ai/audio_transcription/test_vertex_ai_realtime_backend.py +++ b/tests/test_litellm/llms/vertex_ai/audio_transcription/test_vertex_ai_realtime_backend.py @@ -1,3 +1,4 @@ +import asyncio import json from collections.abc import AsyncIterator, Sequence from datetime import timedelta @@ -13,7 +14,7 @@ from google.cloud.speech_v2.types import ( ) from websockets.exceptions import ConnectionClosedError, ConnectionClosedOK -from litellm.llms.vertex_ai.audio_transcription.realtime_backend import SpeechStreamingBackend +from litellm.llms.vertex_ai.audio_transcription.realtime_backend import REQUEST_QUEUE_SIZE, SpeechStreamingBackend from litellm.llms.vertex_ai.audio_transcription.realtime_transformation import SpeechStreamingTarget TARGET: Final = SpeechStreamingTarget( @@ -26,7 +27,7 @@ CONFIGURE: Final = json.dumps( ) FINISH_TURN: Final = json.dumps({"kind": "finish_turn"}) DISCARD_TURN: Final = json.dumps({"kind": "discard_turn"}) -ScriptItem = StreamingRecognizeResponse | Exception +ScriptItem = StreamingRecognizeResponse | Exception | asyncio.Event def _response( @@ -84,13 +85,16 @@ class _FakeSpeechClient: async for request in requests: received.append(request) if request.audio and script: - yield self._next(script) + yield await self._next(script) while script: - yield self._next(script) + yield await self._next(script) @staticmethod - def _next(script: list[ScriptItem]) -> StreamingRecognizeResponse: + async def _next(script: list[ScriptItem]) -> StreamingRecognizeResponse: item: Final = script.pop(0) + if isinstance(item, asyncio.Event): + await item.wait() + return await _FakeSpeechClient._next(script) if isinstance(item, Exception): raise item return item @@ -101,11 +105,18 @@ def _backend(client: _FakeSpeechClient, **kwargs: object) -> SpeechStreamingBack async def _recv(backend: SpeechStreamingBackend) -> dict[str, object]: - message: Final = await backend.recv() + message: Final = await asyncio.wait_for(backend.recv(), timeout=2) assert isinstance(message, str) return json.loads(message) +async def _transcript(backend: SpeechStreamingBackend) -> str: + event: Final = await _recv(backend) + assert event["kind"] == "response", event + (result,) = event["results"] + return result["transcript"] + + async def _configure(backend: SpeechStreamingBackend) -> None: await backend.send(CONFIGURE) assert await _recv(backend) == {"kind": "configured"} @@ -149,7 +160,9 @@ async def test_audio_streams_through_one_recognize_call_with_the_config_first(): @pytest.mark.asyncio async def test_voice_activity_events_are_relayed(): - client = _FakeSpeechClient([_response(None, event="SPEECH_ACTIVITY_BEGIN"), _response(None, event="SPEECH_ACTIVITY_END")]) + client = _FakeSpeechClient( + [_response(None, event="SPEECH_ACTIVITY_BEGIN"), _response(None, event="SPEECH_ACTIVITY_END")] + ) async with _backend(client) as backend: await _configure(backend) await backend.send(b"\x00\x00") @@ -181,16 +194,17 @@ async def test_stream_failure_closes_the_session_with_1011_and_the_reason(): @pytest.mark.asyncio -async def test_close_discards_the_open_turn_then_reports_a_normal_closure(): +async def test_close_reports_a_normal_closure_to_both_directions(): client = _FakeSpeechClient([_response("hi")]) backend = _backend(client) await _configure(backend) await backend.send(b"\x00\x00") - assert (await _recv(backend))["results"][0]["transcript"] == "hi" + assert await _transcript(backend) == "hi" await backend.close() - assert await _recv(backend) == {"kind": "turn_discarded"} with pytest.raises(ConnectionClosedOK): await backend.recv() + with pytest.raises(ConnectionClosedOK): + await backend.send(b"\x00\x00") assert client.transport.closed @@ -210,17 +224,19 @@ async def test_discard_turn_cancels_the_open_stream_and_the_next_turn_starts_fre async with _backend(client) as backend: await _configure(backend) await backend.send(b"\x01\x01") - assert (await _recv(backend))["results"][0]["transcript"] == "draft" + assert await _transcript(backend) == "draft" await backend.send(DISCARD_TURN) assert await _recv(backend) == {"kind": "turn_discarded"} await backend.send(b"\x02\x02") - assert (await _recv(backend))["results"][0]["transcript"] == "again" + assert await _transcript(backend) == "again" assert [_audio(stream) for stream in client.streams] == [[b"\x01\x01"], [b"\x02\x02"]] @pytest.mark.asyncio async def test_billed_seconds_accumulate_across_turns(): - client = _FakeSpeechClient([_response("one", is_final=True, billed=2.0)], [_response("two", is_final=True, billed=3.0)]) + client = _FakeSpeechClient( + [_response("one", is_final=True, billed=2.0)], [_response("two", is_final=True, billed=3.0)] + ) async with _backend(client) as backend: await _configure(backend) await backend.send(b"\x00\x00") @@ -236,26 +252,132 @@ async def test_billed_seconds_accumulate_across_turns(): @pytest.mark.asyncio -async def test_streams_rotate_before_the_five_minute_limit_without_losing_audio(): +async def test_streams_rotate_before_the_five_minute_limit_without_ending_the_turn(): now = [0.0] client = _FakeSpeechClient( [_response("first"), _response("first half", is_final=True, billed=239.0)], - [_response("second")], + [_response("second", billed=1.0)], ) async with _backend(client, clock=lambda: now[0], rotation_seconds=240.0) as backend: await _configure(backend) await backend.send(b"\x01\x01") - assert (await _recv(backend))["results"][0]["transcript"] == "first" + assert await _transcript(backend) == "first" now[0] = 239.0 await backend.send(b"\x02\x02") - assert (await _recv(backend))["results"][0]["transcript"] == "first half" + assert await _transcript(backend) == "first half" now[0] = 240.0 await backend.send(b"\x03\x03") - assert await _recv(backend) == {"kind": "turn_finished"} second = await _recv(backend) - assert second["results"][0]["transcript"] == "second" - assert second["billed_seconds"] == 239.0 + assert second["results"] == [{"transcript": "second", "is_final": False}] + assert second["billed_seconds"] == 240.0 await backend.send(FINISH_TURN) assert await _recv(backend) == {"kind": "turn_finished"} assert [_audio(stream) for stream in client.streams] == [[b"\x01\x01", b"\x02\x02"], [b"\x03\x03"]] assert client.streams[1][0].streaming_config.config.model == "chirp_3" + + +@pytest.mark.asyncio +async def test_turn_finished_follows_results_that_arrive_after_a_rotation(): + now = [0.0] + client = _FakeSpeechClient( + [_response("one"), _response("one two", is_final=True)], + [_response("three")], + ) + async with _backend(client, clock=lambda: now[0], rotation_seconds=240.0) as backend: + await _configure(backend) + await backend.send(b"\x01\x01") + assert await _transcript(backend) == "one" + now[0] = 240.0 + await backend.send(b"\x02\x02") + await backend.send(FINISH_TURN) + assert await _transcript(backend) == "one two" + assert await _transcript(backend) == "three" + assert await _recv(backend) == {"kind": "turn_finished"} + + +@pytest.mark.asyncio +async def test_rotation_waits_for_a_pause_in_speech(): + now = [0.0] + client = _FakeSpeechClient( + [ + _response(None, event="SPEECH_ACTIVITY_BEGIN"), + _response("still talking"), + _response("still talking", is_final=True, event="SPEECH_ACTIVITY_END"), + ], + [_response("next")], + ) + async with _backend( + client, clock=lambda: now[0], rotation_seconds=240.0, rotation_deadline_seconds=280.0 + ) as backend: + await _configure(backend) + await backend.send(b"\x01\x01") + assert (await _recv(backend))["speech_event"] == "begin" + now[0] = 250.0 + await backend.send(b"\x02\x02") + assert await _transcript(backend) == "still talking" + now[0] = 260.0 + await backend.send(b"\x03\x03") + assert (await _recv(backend))["speech_event"] == "end" + now[0] = 261.0 + await backend.send(b"\x04\x04") + assert await _transcript(backend) == "next" + assert [_audio(stream) for stream in client.streams] == [[b"\x01\x01", b"\x02\x02", b"\x03\x03"], [b"\x04\x04"]] + + +@pytest.mark.asyncio +async def test_rotation_is_forced_at_the_deadline_during_continuous_speech(): + now = [0.0] + client = _FakeSpeechClient( + [_response(None, event="SPEECH_ACTIVITY_BEGIN"), _response("still talking")], + [_response("cut off")], + ) + async with _backend( + client, clock=lambda: now[0], rotation_seconds=240.0, rotation_deadline_seconds=280.0 + ) as backend: + await _configure(backend) + await backend.send(b"\x01\x01") + assert (await _recv(backend))["speech_event"] == "begin" + now[0] = 279.0 + await backend.send(b"\x02\x02") + assert await _transcript(backend) == "still talking" + now[0] = 280.0 + await backend.send(b"\x03\x03") + assert await _transcript(backend) == "cut off" + assert [_audio(stream) for stream in client.streams] == [[b"\x01\x01", b"\x02\x02"], [b"\x03\x03"]] + + +@pytest.mark.asyncio +async def test_discard_turn_cancels_every_stream_of_the_turn(): + now = [0.0] + hold = asyncio.Event() + client = _FakeSpeechClient( + [_response("draft"), hold, _response("never delivered")], + [_response("fresh", is_final=True)], + ) + async with _backend(client, clock=lambda: now[0], rotation_seconds=240.0) as backend: + await _configure(backend) + await backend.send(b"\x01\x01") + assert await _transcript(backend) == "draft" + now[0] = 240.0 + await backend.send(b"\x02\x02") + await backend.send(DISCARD_TURN) + assert await _recv(backend) == {"kind": "turn_discarded"} + await backend.send(b"\x03\x03") + assert await _transcript(backend) == "fresh" + assert [_audio(stream) for stream in client.streams] == [[b"\x01\x01"], [b"\x03\x03"]] + + +@pytest.mark.asyncio +async def test_audio_sends_block_once_the_request_queue_is_full(): + hold = asyncio.Event() + client = _FakeSpeechClient([hold, _response("late", is_final=True)]) + async with _backend(client) as backend: + await _configure(backend) + for _ in range(REQUEST_QUEUE_SIZE + 1): + await backend.send(b"\x00\x00") + blocked = asyncio.create_task(backend.send(b"\x00\x00")) + await asyncio.sleep(0) + assert not blocked.done() + hold.set() + await asyncio.wait_for(blocked, timeout=2) + assert await _transcript(backend) == "late" diff --git a/tests/test_litellm/llms/vertex_ai/audio_transcription/test_vertex_ai_realtime_transformation.py b/tests/test_litellm/llms/vertex_ai/audio_transcription/test_vertex_ai_realtime_transformation.py index e294c7cad36..fcbfad8bf12 100644 --- a/tests/test_litellm/llms/vertex_ai/audio_transcription/test_vertex_ai_realtime_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/audio_transcription/test_vertex_ai_realtime_transformation.py @@ -100,7 +100,10 @@ def _types(events: list[dict[str, object]]) -> list[object]: def _commands(config: VertexChirpRealtimeConfig, payload: str) -> list[object]: - return [json.loads(command) if isinstance(command, str) else command for command in config.transform_realtime_request(payload, MODEL)] + return [ + json.loads(command) if isinstance(command, str) else command + for command in config.transform_realtime_request(payload, MODEL) + ] @pytest.mark.parametrize( @@ -108,9 +111,11 @@ def _commands(config: VertexChirpRealtimeConfig, payload: str) -> list[object]: [ ("vertex_ai/chirp_3", True), ("chirp_3", True), - ("chirp_2", True), + ("chirp_2", False), ("gemini-live-2.5-flash", False), ("vertex_ai/gemini-2.0-flash-live-preview-04-09", False), + ("vertex_ai/gemini-3.5-transcribe-live-preview", False), + ("gemini-3.5-transcribe-preview", False), ], ) def test_is_vertex_speech_to_text_model(model: str, expected: bool): @@ -132,7 +137,11 @@ def test_beta_session_update_defaults_the_rate_and_auto_detects_the_language(): config = parse_chirp_session_update( _event( "transcription_session.update", - session={"input_audio_format": "pcm16", "input_audio_transcription": {"model": MODEL}, "turn_detection": None}, + session={ + "input_audio_format": "pcm16", + "input_audio_transcription": {"model": MODEL}, + "turn_detection": None, + }, ), MODEL, ) @@ -146,7 +155,10 @@ def test_beta_session_update_defaults_the_rate_and_auto_detects_the_language(): (_event("session.update", session={"type": "realtime_voice"}), "transcription sessions only"), (_ga_session_update(model="gemini-live-2.5-flash"), "cannot be changed"), (_event("session.update", session={"audio": {"input": {"format": {"type": "audio/pcmu"}}}}), "pcm16"), - (_event("session.update", session={"audio": {"input": {"format": {"type": "audio/pcm", "channels": 2}}}}), "mono"), + ( + _event("session.update", session={"audio": {"input": {"format": {"type": "audio/pcm", "channels": 2}}}}), + "mono", + ), (_ga_session_update(rate=4_000), "sample rates"), (_ga_session_update(rate=96_000), "sample rates"), (_ga_session_update(turn_detection="semantic_vad"), "server_vad"), @@ -169,7 +181,9 @@ def test_session_update_configures_once_and_later_updates_are_ignored(): def test_audio_and_commits_before_session_update_are_rejected(): config = _config() with pytest.raises(ChirpProtocolError, match=r"session\.update must configure"): - config.transform_realtime_request(_event("input_audio_buffer.append", audio=base64.b64encode(b"\x00\x00").decode()), MODEL) + config.transform_realtime_request( + _event("input_audio_buffer.append", audio=base64.b64encode(b"\x00\x00").decode()), MODEL + ) with pytest.raises(ChirpProtocolError, match=r"session\.update must configure"): config.transform_realtime_request(_event("input_audio_buffer.commit"), MODEL) @@ -177,8 +191,14 @@ def test_audio_and_commits_before_session_update_are_rejected(): def test_append_is_split_into_google_sized_chunks(): config = _configured() audio = bytes(range(256)) * 250 - chunks = config.transform_realtime_request(_event("input_audio_buffer.append", audio=base64.b64encode(audio).decode()), MODEL) - assert [len(chunk) for chunk in chunks] == [MAX_AUDIO_MESSAGE_BYTES, MAX_AUDIO_MESSAGE_BYTES, 64_000 - 2 * MAX_AUDIO_MESSAGE_BYTES] + chunks = config.transform_realtime_request( + _event("input_audio_buffer.append", audio=base64.b64encode(audio).decode()), MODEL + ) + assert [len(chunk) for chunk in chunks] == [ + MAX_AUDIO_MESSAGE_BYTES, + MAX_AUDIO_MESSAGE_BYTES, + 64_000 - 2 * MAX_AUDIO_MESSAGE_BYTES, + ] assert b"".join(chunk for chunk in chunks if isinstance(chunk, bytes)) == audio diff --git a/uv.lock b/uv.lock index de7ccb302e9..9d26d22c980 100644 --- a/uv.lock +++ b/uv.lock @@ -10,7 +10,7 @@ resolution-markers = [ ] [options] -exclude-newer = "2026-09-15T01:04:49.417319Z" +exclude-newer = "2026-09-15T21:32:43.124695Z" exclude-newer-span = "P3D" [manifest] @@ -4628,6 +4628,7 @@ ci = [ { name = "blockbuster" }, { name = "claude-agent-sdk" }, { name = "detect-secrets" }, + { name = "google-cloud-speech" }, { name = "google-generativeai" }, { name = "jsonlines" }, { name = "langchain" }, @@ -4822,6 +4823,7 @@ ci = [ { name = "blockbuster", specifier = "==1.5.26" }, { name = "claude-agent-sdk", specifier = "==0.1.44" }, { name = "detect-secrets", specifier = "==1.5.0" }, + { name = "google-cloud-speech", specifier = "==2.40.0" }, { name = "google-generativeai", specifier = "==0.8.6" }, { name = "jsonlines", specifier = "==4.0.0" }, { name = "langchain", specifier = "==1.3.9" }, From da603c629ba465b8e81709847506580623450cb7 Mon Sep 17 00:00:00 2001 From: yassin Date: Fri, 18 Sep 2026 22:37:30 +0000 Subject: [PATCH 070/144] fix(ui): surface a malformed stored MCP allowlist as deny-all and let Save replace or remove it Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../_components/MCPNetworkSettings.test.tsx | 29 +++++++++- .../_components/MCPNetworkSettings.tsx | 55 +++++++++++++------ 2 files changed, 65 insertions(+), 19 deletions(-) diff --git a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPNetworkSettings.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPNetworkSettings.test.tsx index 92f6f7554d4..d27c18c5ae3 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPNetworkSettings.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPNetworkSettings.test.tsx @@ -154,16 +154,39 @@ describe("MCPNetworkSettings", () => { expect(screen.getByRole("textbox", { name: "Client 2 value" })).toHaveValue("codex-mcp-client"); }); - it("ignores a stored allowlist in the old plain-string shape instead of rendering it", async () => { + it("warns that a stored allowlist in the old plain-string shape denies every client and lets Save remove it", async () => { vi.mocked(getGeneralSettingsCall).mockResolvedValue([ { field_name: "mcp_allowed_clients", field_value: ["antigravity-cli"] }, ]); renderSettings(); - await screen.findByText("Allowed Clients"); + expect(await screen.findByText(/stored allowlist is not a list of alias and value pairs/)).toBeVisible(); expect(screen.queryByRole("textbox", { name: "Client 1 value" })).not.toBeInTheDocument(); - expect(screen.queryByText(/every client is denied/)).not.toBeInTheDocument(); + + await userEvent.click(screen.getByRole("button", { name: /Save/ })); + + await waitFor(() => expect(deleteConfigFieldSetting).toHaveBeenCalledWith("tok", "mcp_allowed_clients")); + expect(updateConfigFieldSetting).not.toHaveBeenCalled(); + await waitFor(() => expect(screen.queryByText(/stored allowlist is not a list/)).not.toBeInTheDocument()); + }); + + it("replaces a stored allowlist in the old plain-string shape with the clients the admin adds", async () => { + vi.mocked(getGeneralSettingsCall).mockResolvedValue([ + { field_name: "mcp_allowed_clients", field_value: ["antigravity-cli"] }, + ]); + + renderSettings(); + + await screen.findByText(/stored allowlist is not a list of alias and value pairs/); + await addClient(ANTIGRAVITY.alias, ANTIGRAVITY.value); + await userEvent.click(screen.getByRole("button", { name: /Save/ })); + + await waitFor(() => + expect(updateConfigFieldSetting).toHaveBeenCalledWith("tok", "mcp_allowed_clients", [ANTIGRAVITY]), + ); + expect(deleteConfigFieldSetting).not.toHaveBeenCalledWith("tok", "mcp_allowed_clients"); + await waitFor(() => expect(screen.queryByText(/stored allowlist is not a list/)).not.toBeInTheDocument()); }); it("adds clients as alias and value pairs and saves them under mcp_allowed_clients", async () => { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPNetworkSettings.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPNetworkSettings.tsx index 2ef3ee8707d..ae1fad36599 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPNetworkSettings.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPNetworkSettings.tsx @@ -42,10 +42,20 @@ const isAllowedClient = (entry: unknown): entry is AllowedClient => { return typeof alias === "string" && typeof value === "string"; }; -const parseStoredClients = (fieldValue: unknown): AllowedClient[] | null => - Array.isArray(fieldValue) && fieldValue.every(isAllowedClient) - ? fieldValue.map(({ alias, value }) => ({ alias, value })) - : null; +type StoredAllowlist = + | { readonly kind: "absent" } + | { readonly kind: "clients"; readonly clients: AllowedClient[] } + | { readonly kind: "malformed" }; + +const ABSENT: StoredAllowlist = { kind: "absent" }; + +const parseStoredClients = (fieldValue: unknown): StoredAllowlist => { + if (fieldValue === null || fieldValue === undefined) return ABSENT; + if (Array.isArray(fieldValue) && fieldValue.every(isAllowedClient)) { + return { kind: "clients", clients: fieldValue.map(({ alias, value }) => ({ alias, value })) }; + } + return { kind: "malformed" }; +}; let nextRowKey = 0; const newRow = (client: AllowedClient = { alias: "", value: "" }): AllowedClientRow => ({ @@ -66,8 +76,16 @@ const sameClients = (a: AllowedClient[], b: AllowedClient[]) => const unchangedSinceLoad = (value: string[], stored: string[] | null) => stored === null ? value.length === 0 : value.length > 0 && sameList(value, stored); -const clientsUnchangedSinceLoad = (value: AllowedClient[], stored: AllowedClient[] | null) => - stored === null ? value.length === 0 : value.length > 0 && sameClients(value, stored); +const clientsUnchangedSinceLoad = (value: AllowedClient[], stored: StoredAllowlist) => { + switch (stored.kind) { + case "absent": + return value.length === 0; + case "clients": + return value.length > 0 && sameClients(value, stored.clients); + case "malformed": + return false; + } +}; const headerUnchangedSinceLoad = (value: string, stored: string | null) => stored === null ? value === "" : value !== "" && value === stored; @@ -79,7 +97,7 @@ const MCPNetworkSettings: React.FC = ({ accessToken }) const [allowedClients, setAllowedClients] = useState([]); const [clientIdHeader, setClientIdHeader] = useState(""); const [storedRanges, setStoredRanges] = useState(null); - const [storedClients, setStoredClients] = useState(null); + const [storedClients, setStoredClients] = useState(ABSENT); const [storedClientIdHeader, setStoredClientIdHeader] = useState(null); const [currentIp, setCurrentIp] = useState(null); const [rangeDraft, setRangeDraft] = useState(""); @@ -100,11 +118,9 @@ const MCPNetworkSettings: React.FC = ({ accessToken }) setStoredRanges(field.field_value); } if (field.field_name === "mcp_allowed_clients") { - const clients = parseStoredClients(field.field_value); - if (clients !== null) { - setAllowedClients(clients.map(newRow)); - setStoredClients(clients); - } + const stored = parseStoredClients(field.field_value); + setAllowedClients(stored.kind === "clients" ? stored.clients.map(newRow) : []); + setStoredClients(stored); } if (field.field_name === "mcp_client_id_header" && typeof field.field_value === "string") { setClientIdHeader(field.field_value); @@ -145,11 +161,11 @@ const MCPNetworkSettings: React.FC = ({ accessToken }) if (clientsUnchangedSinceLoad(clients, storedClients)) return; if (clients.length > 0) { await updateConfigFieldSetting(token, "mcp_allowed_clients", clients); - setStoredClients(clients); + setStoredClients({ kind: "clients", clients }); return; } await deleteConfigFieldSetting(token, "mcp_allowed_clients"); - setStoredClients(null); + setStoredClients(ABSENT); }; const persistClientIdHeader = async (token: string) => { @@ -216,7 +232,8 @@ const MCPNetworkSettings: React.FC = ({ accessToken }) } const suggestedRange = currentIp ? ipToSlash24(currentIp) : null; - const storedAllowlistDeniesEveryone = storedClients !== null && storedClients.length === 0; + const storedAllowlistIsMalformed = storedClients.kind === "malformed"; + const storedAllowlistIsEmpty = storedClients.kind === "clients" && storedClients.clients.length === 0; return (
@@ -303,7 +320,13 @@ const MCPNetworkSettings: React.FC = ({ accessToken })

Allowed Clients

- {storedAllowlistDeniesEveryone && ( + {storedAllowlistIsMalformed && ( +

+ The stored allowlist is not a list of alias and value pairs, so every client is denied. Add the clients you + want and save to replace it, or save with the list empty to remove it and allow every client again. +

+ )} + {storedAllowlistIsEmpty && (

An empty allowlist is currently stored, so every client is denied. Save with the list empty to remove it and allow every client again. From 1c15d9f291d6631c8af430e128231746f899c7f2 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 18 Sep 2026 15:44:06 -0700 Subject: [PATCH 071/144] fix(responses): restore encrypted_content and apply affinity on the native WebSocket relay --- litellm/llms/custom_httpx/llm_http_handler.py | 1 + litellm/proxy/_lazy_openapi_snapshot.json | 2 +- .../proxy/response_api_endpoints/endpoints.py | 54 +++- litellm/responses/main.py | 2 + litellm/responses/streaming_iterator.py | 167 +++++++++-- .../response_api_endpoints/test_endpoints.py | 95 +++++++ .../test_responses_api_request_body.py | 24 ++ .../test_responses_websocket_all_providers.py | 266 ++++++++++++++++++ 8 files changed, 570 insertions(+), 41 deletions(-) diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 98fe0014386..ab327299243 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -6742,6 +6742,7 @@ class BaseLLMHTTPHandler: output_guardrail_callbacks=_ws_output_guardrail_callbacks, quota_callbacks=_ws_quota_callbacks, authorized_model=model, + custom_llm_provider=custom_llm_provider, ) await streaming.bidirectional_forward() diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index 213cd88b6ce..40b64160b71 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -19616,7 +19616,7 @@ } } }, - "description": "\n Unified rate-limit error.\n\n Every rate-limit condition surfaced by litellm \u2014 whether it originated from\n an upstream LLM provider, a vendor batch endpoint, or one of litellm's own\n proxy-side limiters (parallel-requests, dynamic-rate, batch-rate, budget,\n max-iterations, etc.) \u2014 is raised as an instance of this class.\n\n The :attr:`category` attribute lets callers distinguish the source. See\n :class:`RateLimitErrorCategory` for the available values.\n " + "description": "\nUnified rate-limit error.\n\nEvery rate-limit condition surfaced by litellm \u2014 whether it originated from\nan upstream LLM provider, a vendor batch endpoint, or one of litellm's own\nproxy-side limiters (parallel-requests, dynamic-rate, batch-rate, budget,\nmax-iterations, etc.) \u2014 is raised as an instance of this class.\n\nThe :attr:`category` attribute lets callers distinguish the source. See\n:class:`RateLimitErrorCategory` for the available values.\n" }, "500": { "content": { diff --git a/litellm/proxy/response_api_endpoints/endpoints.py b/litellm/proxy/response_api_endpoints/endpoints.py index 5907ffc64eb..ea6b67fa026 100644 --- a/litellm/proxy/response_api_endpoints/endpoints.py +++ b/litellm/proxy/response_api_endpoints/endpoints.py @@ -11,10 +11,12 @@ import fastapi from fastapi import APIRouter, Depends, HTTPException, Request, Response from fastapi.responses import JSONResponse from openai.types.responses.response_create_params import ResponseInputParam +from pydantic import BaseModel, ConfigDict, ValidationError from starlette.websockets import WebSocket, WebSocketDisconnect from typing_extensions import ReadOnly, TypedDict from litellm._logging import verbose_proxy_logger +from litellm.constants import EMPTY_MAPPING from litellm.integrations.custom_guardrail import ModifyResponseException from litellm.llms.base_llm.guardrail_translation.utils import ( blocked_responses_api_usage as _blocked_responses_api_usage, @@ -1289,7 +1291,8 @@ async def cancel_response( async def _read_ws_model_from_first_frame( websocket: WebSocket, -) -> tuple | None: + query_model: str | None = None, +) -> tuple[str, str] | None: """Read the first WS frame and return (model, raw_message), or None on error. Sends an appropriate error frame and closes the socket before returning None. @@ -1338,7 +1341,7 @@ async def _read_ws_model_from_first_frame( await websocket.close(code=1008, reason="Invalid first message") return None - model: Final = _extract_model_from_first_ws_event(first_event) + model: Final = query_model or _extract_model_from_first_ws_event(first_event) if not model: await websocket.send_text( json.dumps( @@ -1369,6 +1372,29 @@ def _extract_model_from_first_ws_event(first_event: Any) -> str | None: return (nested.get("model") if isinstance(nested, dict) else None) or first_event.get("model") +class _ResponseCreateRoutingHints(BaseModel): + model_config = ConfigDict(extra="ignore", frozen=True) + + input: str | list[object] | None = None + previous_response_id: str | None = None + response: "_ResponseCreateRoutingHints | None" = None + + +def _routing_hints_from_first_ws_frame(first_message: str) -> Mapping[str, object]: + try: + frame: Final = _ResponseCreateRoutingHints.model_validate_json(first_message) + except ValidationError: + return EMPTY_MAPPING + nested: Final = frame.response or frame + hints: Final = { + "input": frame.input if nested.input is None else nested.input, + "previous_response_id": ( + frame.previous_response_id if nested.previous_response_id is None else nested.previous_response_id + ), + } + return MappingProxyType({key: value for key, value in hints.items() if value is not None}) + + async def _enforce_responses_ws_first_frame_model_auth( request: Request, model: str, @@ -1455,19 +1481,16 @@ async def responses_websocket_endpoint( accept_kwargs["subprotocol"] = requested_protocols[0] await websocket.accept(**accept_kwargs) - first_message: str | None = None - if not model: - result: Final = await _read_ws_model_from_first_frame(websocket) - if result is None: - return - model, first_message = result + result: Final = await _read_ws_model_from_first_frame(websocket, query_model=model) + if result is None: + return + resolved_model, first_message = result data: dict[str, object] = { - "model": model, + "model": resolved_model, "websocket": websocket, + "first_message": first_message, } - if first_message is not None: - data["first_message"] = first_message # Construct a synthetic Request for pre-call processing headers_list: Final = list(websocket.scope.get("headers") or []) @@ -1480,7 +1503,7 @@ async def responses_websocket_endpoint( request: Final = Request(scope=scope) request._url = websocket.url - _body_bytes: Final = json.dumps({"model": model}).encode() + _body_bytes: Final = json.dumps({"model": resolved_model}).encode() async def return_body(): return _body_bytes @@ -1490,10 +1513,10 @@ async def responses_websocket_endpoint( # Phase 1: pre-call processing (auth, guardrails, rate limits) base_llm_response_processor: Final = ProxyBaseLLMRequestProcessing(data=data) try: - if first_message is not None: + if not model: await _enforce_responses_ws_first_frame_model_auth( request=request, - model=model, + model=resolved_model, user_api_key_dict=user_api_key_dict, llm_router=llm_router, ) @@ -1512,7 +1535,7 @@ async def responses_websocket_endpoint( user_request_timeout=user_request_timeout, user_max_tokens=user_max_tokens, user_api_base=user_api_base, - model=model, + model=resolved_model, route_type="_aresponses_websocket", ) except Exception as e: @@ -1537,6 +1560,7 @@ async def responses_websocket_endpoint( # Phase 2: route to upstream provider try: data["user_api_key_dict"] = user_api_key_dict + data.update(_routing_hints_from_first_ws_frame(first_message)) llm_call: Final = await route_request( data=data, route_type="_aresponses_websocket", diff --git a/litellm/responses/main.py b/litellm/responses/main.py index 6dc34bb93ef..9705794d01d 100644 --- a/litellm/responses/main.py +++ b/litellm/responses/main.py @@ -2338,6 +2338,8 @@ async def _aresponses_websocket( "api_base", "api_key", "timeout", + "input", + "previous_response_id", } remaining_kwargs: Final = {k: v for k, v in kwargs.items() if k not in _explicit_keys} diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 8d766cf1cd0..c99a481db7d 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -225,6 +225,29 @@ def _status_code_for_error_fields(error_type: str | None, error_code: str | None ) +def _map_stream_error_to_exception(error_obj: object, model: str, custom_llm_provider: str) -> Exception: + from litellm.llms.base_llm.chat.transformation import BaseLLMException + + error_message, error_type, error_code = _error_event_fields(error_obj) + status_code: Final = _status_code_for_error_fields(error_type, error_code) + error_body: Final = {"message": error_message, "type": error_type, "code": error_code} + provider_exception: Final = BaseLLMException( + status_code=status_code, + message=f"Error code: {status_code} - {{'error': {error_body}}}", + body=error_body, + ) + try: + return litellm.exception_type( + model=model, + custom_llm_provider=custom_llm_provider, + original_exception=provider_exception, + completion_kwargs={}, + extra_kwargs={}, + ) + except Exception as mapped_exception: + return mapped_exception + + def _mid_stream_fallback_eligible(mapped_exception: Exception) -> bool: if isinstance(mapped_exception, litellm.ContentPolicyViolationError): return True @@ -588,26 +611,7 @@ class BaseResponsesAPIStreamingIterator: ) def _map_error_event_exception(self, error_obj: object) -> Exception: - from litellm.llms.base_llm.chat.transformation import BaseLLMException - - error_message, error_type, error_code = _error_event_fields(error_obj) - status_code: Final = _status_code_for_error_fields(error_type, error_code) - error_body: Final = {"message": error_message, "type": error_type, "code": error_code} - provider_exception: Final = BaseLLMException( - status_code=status_code, - message=f"Error code: {status_code} - {{'error': {error_body}}}", - body=error_body, - ) - try: - return litellm.exception_type( - model=self.model or "", - custom_llm_provider=self.custom_llm_provider or "", - original_exception=provider_exception, - completion_kwargs={}, - extra_kwargs={}, - ) - except Exception as mapped_exception: - return mapped_exception + return _map_stream_error_to_exception(error_obj, self.model or "", self.custom_llm_provider or "") def _maybe_raise_for_error_event(self, result: object) -> None: chunk_type: Final = getattr(result, "type", None) @@ -1691,6 +1695,65 @@ RESPONSES_WS_LOGGED_EVENT_TYPES: Final = [ RESPONSES_WS_MASKABLE_TEXT_BLOCK_TYPES: Final = frozenset({"input_text", "output_text", "text"}) +_RESPONSES_WS_FAILURE_EVENT_TYPES: Final = frozenset({"error", "response.failed"}) + +_RESPONSES_WS_OUTPUT_ITEM_EVENT_TYPES: Final = frozenset({"response.output_item.added", "response.output_item.done"}) + + +def _ws_event_error(event: _MutableJsonObject) -> object: + if event.get("type") == "error": + return event.get("error") + response: Final = event.get("response") + return response.get("error") if _is_json_object(response) else None + + +def _item_id_fields(item: object) -> tuple[object, object]: + return (item.get("id"), item.get("encrypted_content")) if _is_json_object(item) else (None, None) + + +def _restore_input_item_ids(items: list[object]) -> bool: + before: Final = tuple(_item_id_fields(item) for item in items) + ResponsesAPIRequestUtils._restore_encrypted_content_item_ids_in_input(items) # pyright: ignore[reportPrivateUsage] # same restore the HTTP responses path runs + return before != tuple(_item_id_fields(item) for item in items) + + +def _restore_wrapped_ids_in_container(container: _MutableJsonObject) -> bool: + input_items: Final = container.get("input") + input_restored: Final = _is_json_array(input_items) and _restore_input_item_ids(input_items) + previous_response_id: Final = container.get("previous_response_id") + if not isinstance(previous_response_id, str): + return input_restored + original_previous_response_id: Final = ( + ResponsesAPIRequestUtils.decode_previous_response_id_to_original_previous_response_id(previous_response_id) + ) + if original_previous_response_id == previous_response_id: + return input_restored + container["previous_response_id"] = original_previous_response_id + return True + + +def _restore_wrapped_ids_in_response_create(msg_obj: _MutableJsonObject) -> bool: + nested: Final = msg_obj.get("response") + containers: Final = (msg_obj, nested) if _is_json_object(nested) else (msg_obj,) + restored: Final = tuple(_restore_wrapped_ids_in_container(container) for container in containers) + return any(restored) + + +def _wrap_output_item_encrypted_content(event_obj: _MutableJsonObject, litellm_metadata: dict[str, object]) -> bool: + if not litellm_metadata.get("encrypted_content_affinity_enabled"): + return False + model_id: Final = _model_id_from_metadata(litellm_metadata) + item: Final = event_obj.get("item") + if model_id is None or not _is_json_object(item): + return False + encrypted_content: Final = item.get("encrypted_content") + if not isinstance(encrypted_content, str) or not encrypted_content: + return False + item["encrypted_content"] = ResponsesAPIRequestUtils._wrap_encrypted_content_with_model_id( # pyright: ignore[reportPrivateUsage] # same wrap the HTTP streaming path applies + encrypted_content=encrypted_content, model_id=model_id + ) + return True + class ResponsesWebSocketStreaming: """ @@ -1717,12 +1780,16 @@ class ResponsesWebSocketStreaming: output_guardrail_callbacks: list[PresidioGuardrailCallback] | None = None, quota_callbacks: Sequence[ProjectQuotaCallback] | None = None, authorized_model: str | None = None, + custom_llm_provider: str | None = None, ): self.websocket = websocket self.backend_ws = backend_ws self.logging_obj = logging_obj self.user_api_key_dict = user_api_key_dict self.request_data: dict[str, object] = request_data or {} + litellm_metadata: Final = self.request_data.get("litellm_metadata") + self.litellm_metadata: dict[str, object] = litellm_metadata if _is_json_object(litellm_metadata) else {} + self.custom_llm_provider: str | None = custom_llm_provider self.messages: list[_MutableJsonObject] = [] self.input_messages: list[dict[str, object]] = [] self.first_message = first_message @@ -1795,8 +1862,55 @@ class ResponsesWebSocketStreaming: return if self.input_messages: self.logging_obj.model_call_details["messages"] = self.input_messages - if self.messages: + if not self.messages: + return + failed_event: Final = next( + (event for event in self.messages if event.get("type") in _RESPONSES_WS_FAILURE_EVENT_TYPES), None + ) + if failed_event is None: asyncio.create_task(self.logging_obj.dispatch_success_handlers(self.messages, prefer_async_handlers=True)) + return + self._record_usage_for_failure() + exception: Final = _map_stream_error_to_exception( + _ws_event_error(failed_event), self.authorized_model or "", self.custom_llm_provider or "" + ) + traceback_exception: Final = "".join(traceback.format_exception(exception)) + asyncio.create_task( + self.logging_obj.dispatch_failure_handlers(exception, traceback_exception, prefer_async_handlers=True) + ) + + def _record_usage_for_failure(self) -> None: + from litellm.cost_calculator import ResponsesWebSocketTokenUsageProcessor + from litellm.types.utils import LiteLLMRealtimeStreamLoggingObject + + usage: Final = ResponsesWebSocketTokenUsageProcessor.collect_and_combine_usage_from_responses_ws_results( + self.messages + ) + tier_partition: Final = ResponsesWebSocketTokenUsageProcessor.partition_results_by_service_tier(self.messages) + service_tier: Final = next(iter(tier_partition)) if len(tier_partition) == 1 else None + logging_result: Final = LiteLLMRealtimeStreamLoggingObject( + usage=usage, results=self.messages, service_tier=service_tier + ) + response_cost: Final = self.logging_obj._response_cost_calculator(result=logging_result) or 0.0 # pyright: ignore[reportPrivateUsage] # as the HTTP streaming iterator does + self.logging_obj.record_partial_usage_for_failure(usage, response_cost) + + def _wrap_response_event(self, response_str: str) -> str: + try: + event_obj: Final = _load_json_object(response_str) + except (json.JSONDecodeError, TypeError): + return response_str + response: Final = event_obj.get("response") + if _is_json_object(response): + event_obj["response"] = ResponsesAPIRequestUtils._update_responses_api_response_id_with_model_id( # pyright: ignore[reportPrivateUsage] # same wrap the HTTP streaming path applies + responses_api_response=response, + custom_llm_provider=self.custom_llm_provider, + litellm_metadata=self.litellm_metadata, + ) + return json.dumps(event_obj) + if event_obj.get("type") not in _RESPONSES_WS_OUTPUT_ITEM_EVENT_TYPES: + return response_str + item_wrapped: Final = _wrap_output_item_encrypted_content(event_obj, self.litellm_metadata) + return json.dumps(event_obj) if item_wrapped else response_str async def backend_to_client(self) -> None: """Forward events from backend WebSocket to the client.""" @@ -1833,12 +1947,13 @@ class ResponsesWebSocketStreaming: unmasked_str = self._unmask_response_event(response_str) output_masked_str = await self._mask_response_completed(unmasked_str) + wrapped_str = self._wrap_response_event(output_masked_str) # Log the output-masked form so PII redacted by apply_to_output # guardrails does not appear in success logs. - self._store_event(output_masked_str) + self._store_event(wrapped_str) - await self.websocket.send_text(output_masked_str) + await self.websocket.send_text(wrapped_str) except websockets.exceptions.ConnectionClosed as e: verbose_logger.debug("Responses WS backend connection closed: %s", e) @@ -1898,14 +2013,16 @@ class ResponsesWebSocketStreaming: # Always enforce the authorized model, even when PII masking is off. model_modified: Final = self._enforce_authorized_model(msg_obj) + ids_restored: Final = _restore_wrapped_ids_in_response_create(msg_obj) + frame_modified: Final = model_modified or ids_restored if not self.guardrail_callbacks: - return json.dumps(msg_obj) if model_modified else message + return json.dumps(msg_obj) if frame_modified else message if "metadata" not in self.request_data: self.request_data["metadata"] = {} - modified = model_modified + modified = frame_modified guardrail_cbs: Final[tuple[PresidioGuardrailCallback, ...]] = tuple(self.guardrail_callbacks) for cb in guardrail_cbs: presidio_config = cb.get_presidio_settings_from_request_data(self.request_data) diff --git a/tests/test_litellm/proxy/response_api_endpoints/test_endpoints.py b/tests/test_litellm/proxy/response_api_endpoints/test_endpoints.py index d7010de6405..91d688fbacf 100644 --- a/tests/test_litellm/proxy/response_api_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/response_api_endpoints/test_endpoints.py @@ -510,6 +510,66 @@ class TestResponsesWSFirstFrameModelAuth: mock_model_auth.assert_awaited_once() + @pytest.mark.asyncio + @pytest.mark.parametrize("nested", [False, True]) + @pytest.mark.parametrize("query_model", [None, "gpt-4o-mini"]) + async def test_endpoint_routes_on_first_frame_input_and_previous_response_id(self, nested, query_model): + from litellm.proxy.response_api_endpoints.endpoints import ( + responses_websocket_endpoint, + ) + + replayed_input = [{"type": "reasoning", "id": "encitem_abc", "encrypted_content": "litellm_enc:abc;blob"}] + payload = {"model": "gpt-4o-mini", "input": replayed_input, "previous_response_id": "resp_prev"} + first_frame = {"type": "response.create", "response": payload} if nested else {"type": "response.create", **payload} + raw_first_frame = json.dumps(first_frame) + + ws = MagicMock() + ws.headers = {} + ws.query_params = {} + ws.scope = {"headers": []} + ws.url = "ws://testserver/v1/responses" + ws.accept = AsyncMock() + ws.receive_text = AsyncMock(return_value=raw_first_frame) + ws.close = AsyncMock() + + processor = MagicMock() + processor.common_processing_pre_call_logic = AsyncMock( + return_value=({"model": "gpt-4o-mini", "litellm_metadata": {}}, MagicMock()) + ) + + async def fake_llm_call(): + return None + + with ( + patch( # test-quality-ok: first-frame model auth needs a live router and key table and has its own tests below + "litellm.proxy.response_api_endpoints.endpoints._enforce_responses_ws_first_frame_model_auth", + new_callable=AsyncMock, + ), + patch( # test-quality-ok: the pre-call processor needs a live proxy; the payload it hands to routing is what is under test + "litellm.proxy.response_api_endpoints.endpoints.ProxyBaseLLMRequestProcessing", + return_value=processor, + ), + patch( # test-quality-ok: routing is the seam where the first frame's input and previous_response_id become observable + "litellm.proxy.route_llm_request.route_request", + new_callable=AsyncMock, + return_value=fake_llm_call(), + ) as mock_route_request, + ): + await responses_websocket_endpoint( + websocket=ws, + model=query_model, + user_api_key_dict=MagicMock(), + ) + + ws.receive_text.assert_awaited_once() + routed = mock_route_request.await_args.kwargs["data"] + assert routed["model"] == "gpt-4o-mini" + assert routed["input"] == replayed_input + assert routed["previous_response_id"] == "resp_prev" + assert processor.common_processing_pre_call_logic.await_args.kwargs["model"] == "gpt-4o-mini" + assert mock_route_request.await_args.kwargs["route_type"] == "_aresponses_websocket" + ws.close.assert_not_awaited() + @pytest.mark.asyncio async def test_reruns_model_auth_for_first_frame_model(self): from starlette.requests import Request @@ -636,6 +696,41 @@ class TestReadWSModelFromFirstFrameErrors: ws.send_text.assert_not_awaited() ws.close.assert_not_awaited() + @pytest.mark.asyncio + async def test_query_model_wins_over_first_frame_model(self): + from litellm.proxy.response_api_endpoints.endpoints import ( + _read_ws_model_from_first_frame, + ) + + raw = json.dumps({"type": "response.create", "model": "gpt-4o", "input": []}) + ws = MagicMock() + ws.receive_text = AsyncMock(return_value=raw) + ws.send_text = AsyncMock() + ws.close = AsyncMock() + + result = await _read_ws_model_from_first_frame(ws, query_model="reasoning-group") + + assert result == ("reasoning-group", raw) + ws.close.assert_not_awaited() + + @pytest.mark.asyncio + async def test_query_model_satisfies_a_first_frame_without_model(self): + from litellm.proxy.response_api_endpoints.endpoints import ( + _read_ws_model_from_first_frame, + ) + + raw = json.dumps({"type": "response.create", "input": []}) + ws = MagicMock() + ws.receive_text = AsyncMock(return_value=raw) + ws.send_text = AsyncMock() + ws.close = AsyncMock() + + result = await _read_ws_model_from_first_frame(ws, query_model="reasoning-group") + + assert result == ("reasoning-group", raw) + ws.send_text.assert_not_awaited() + ws.close.assert_not_awaited() + class TestManagedResponsesSameProvider: def _handler(self, model, custom_llm_provider=None): diff --git a/tests/test_litellm/responses/test_responses_api_request_body.py b/tests/test_litellm/responses/test_responses_api_request_body.py index 5fced458208..743ad237e45 100644 --- a/tests/test_litellm/responses/test_responses_api_request_body.py +++ b/tests/test_litellm/responses/test_responses_api_request_body.py @@ -424,6 +424,30 @@ async def test_aresponses_websocket_strips_responses_routing_prefix_from_openai_ assert mock_ws.call_args.kwargs["custom_llm_provider"] == "openai" +@pytest.mark.asyncio +async def test_aresponses_websocket_keeps_routing_hints_out_of_the_relay_kwargs(): # test-quality-ok: the relay kwargs are the only place a dropped key is observable; the provider socket behind them is the boundary + from unittest.mock import MagicMock + + from litellm.responses.main import _aresponses_websocket + + with patch.object( + import_module("litellm.responses.main").base_llm_http_handler, "async_responses_websocket", + new_callable=AsyncMock, + ) as mock_ws: + await _aresponses_websocket( + model="openai/gpt-5.6", + websocket=MagicMock(), + api_key="sk-test", + litellm_logging_obj=MagicMock(), + input=[{"type": "message", "role": "user", "content": "hi"}], + previous_response_id="resp_prev", + ) + + mock_ws.assert_awaited_once() + assert "input" not in mock_ws.call_args.kwargs + assert "previous_response_id" not in mock_ws.call_args.kwargs + + _INJECTION_POINT_INPUT = [{"role": "system", "content": "You are terse."}, {"role": "user", "content": "hi"}] _SYSTEM_POINT = {"location": "message", "role": "system"} _USER_POINT = {"location": "message", "role": "user"} diff --git a/tests/test_litellm/responses/test_responses_websocket_all_providers.py b/tests/test_litellm/responses/test_responses_websocket_all_providers.py index fe3c4a0640d..b671e60438e 100644 --- a/tests/test_litellm/responses/test_responses_websocket_all_providers.py +++ b/tests/test_litellm/responses/test_responses_websocket_all_providers.py @@ -2628,3 +2628,269 @@ class TestNativeWebSocketUrlConstruction: mock_config.get_websocket_url.assert_called_once() _, call_kwargs = mock_config.get_websocket_url.call_args assert call_kwargs["litellm_params"]["api_version"] == "2025-04-01-preview" + + +_AFFINITY_METADATA = { + "model_info": {"id": "dep-1"}, + "encrypted_content_affinity_enabled": True, +} + + +def _wrapped_reasoning_item(): + from litellm.responses.utils import ResponsesAPIRequestUtils + + return { + "type": "reasoning", + "id": ResponsesAPIRequestUtils._build_encrypted_item_id("dep-1", "rs_orig"), + "encrypted_content": ResponsesAPIRequestUtils._wrap_encrypted_content_with_model_id("gAAAA-blob", "dep-1"), + "summary": [], + } + + +class TestNativeWebSocketEncryptedContentAffinity: + """The native relay must restore and wrap ids the same way the HTTP /v1/responses path does.""" + + @pytest.mark.asyncio + @pytest.mark.parametrize("nested", [False, True]) + async def test_client_to_backend_restores_wrapped_ids(self, nested): + from unittest.mock import AsyncMock + + from litellm.responses.utils import ResponsesAPIRequestUtils + + wrapped_previous = ResponsesAPIRequestUtils._build_responses_api_response_id( + custom_llm_provider="openai", model_id="dep-1", response_id="resp_orig" + ) + payload = { + "input": [_wrapped_reasoning_item(), {"type": "message", "role": "user", "content": "hi"}], + "previous_response_id": wrapped_previous, + } + frame = {"type": "response.create", "response": payload} if nested else {"type": "response.create", **payload} + backend_ws = MagicMock() + backend_ws.send = AsyncMock() + websocket = MagicMock() + websocket.receive_text = AsyncMock(side_effect=[json.dumps(frame), Exception("stop")]) + handler = _make_streaming(websocket=websocket, backend_ws=backend_ws, request_data={}) + + await handler.client_to_backend() + + sent = json.loads(backend_ws.send.await_args_list[0][0][0]) + body = sent["response"] if nested else sent + assert body["input"][0]["id"] == "rs_orig" + assert body["input"][0]["encrypted_content"] == "gAAAA-blob" + assert body["input"][1] == {"type": "message", "role": "user", "content": "hi"} + assert body["previous_response_id"] == "resp_orig" + + @pytest.mark.asyncio + async def test_client_to_backend_leaves_unwrapped_frames_untouched(self): + from unittest.mock import AsyncMock + + frame = json.dumps({"type": "response.create", "input": "hello", "previous_response_id": "resp_raw"}) + backend_ws = MagicMock() + backend_ws.send = AsyncMock() + websocket = MagicMock() + websocket.receive_text = AsyncMock(side_effect=[frame, Exception("stop")]) + handler = _make_streaming(websocket=websocket, backend_ws=backend_ws, request_data={}) + + await handler.client_to_backend() + + assert backend_ws.send.await_args_list[0][0][0] == frame + + @pytest.mark.asyncio + async def test_backend_to_client_wraps_ids_when_affinity_is_enabled(self): + import asyncio + from unittest.mock import AsyncMock + + import websockets.exceptions # noqa: F401 (lazy submodule must be importable) + + from litellm.responses.utils import ResponsesAPIRequestUtils + + reasoning_item = {"type": "reasoning", "id": "rs_1", "encrypted_content": "gAAAA-blob", "summary": []} + websocket = MagicMock() + websocket.send_text = AsyncMock() + backend_ws = MagicMock() + backend_ws.recv = AsyncMock( + side_effect=[ + json.dumps({"type": "response.output_item.done", "output_index": 0, "item": dict(reasoning_item)}), + json.dumps( + { + "type": "response.completed", + "response": {"id": "resp_1", "output": [dict(reasoning_item)], "usage": {"total_tokens": 3}}, + } + ), + Exception("stop"), + ] + ) + logging_obj = MagicMock() + logging_obj.dispatch_success_handlers = AsyncMock() + handler = _make_streaming( + websocket=websocket, + backend_ws=backend_ws, + logging_obj=logging_obj, + request_data={"litellm_metadata": dict(_AFFINITY_METADATA)}, + custom_llm_provider="openai", + ) + + await handler.backend_to_client() + + wrapped_content = ResponsesAPIRequestUtils._wrap_encrypted_content_with_model_id("gAAAA-blob", "dep-1") + item_done = json.loads(websocket.send_text.await_args_list[0][0][0]) + assert item_done["item"]["encrypted_content"] == wrapped_content + completed = json.loads(websocket.send_text.await_args_list[1][0][0]) + assert completed["response"]["id"] == ResponsesAPIRequestUtils._build_responses_api_response_id( + custom_llm_provider="openai", model_id="dep-1", response_id="resp_1" + ) + assert completed["response"]["output"][0]["id"] == ResponsesAPIRequestUtils._build_encrypted_item_id( + "dep-1", "rs_1" + ) + assert completed["response"]["output"][0]["encrypted_content"] == wrapped_content + await asyncio.sleep(0) + logged = logging_obj.dispatch_success_handlers.await_args[0][0] + assert logged[0]["response"]["id"] == completed["response"]["id"] + + @pytest.mark.asyncio + async def test_backend_to_client_wraps_only_response_id_without_affinity(self): + from unittest.mock import AsyncMock + + import websockets.exceptions # noqa: F401 (lazy submodule must be importable) + + from litellm.responses.utils import ResponsesAPIRequestUtils + + reasoning_item = {"type": "reasoning", "id": "rs_1", "encrypted_content": "gAAAA-blob", "summary": []} + websocket = MagicMock() + websocket.send_text = AsyncMock() + backend_ws = MagicMock() + backend_ws.recv = AsyncMock( + side_effect=[ + json.dumps({"type": "response.output_item.done", "output_index": 0, "item": dict(reasoning_item)}), + json.dumps({"type": "response.completed", "response": {"id": "resp_1", "output": [dict(reasoning_item)]}}), + Exception("stop"), + ] + ) + logging_obj = MagicMock() + logging_obj.dispatch_success_handlers = AsyncMock() + handler = _make_streaming( + websocket=websocket, + backend_ws=backend_ws, + logging_obj=logging_obj, + request_data={"litellm_metadata": {"model_info": {"id": "dep-1"}}}, + custom_llm_provider="openai", + ) + + await handler.backend_to_client() + + item_done = json.loads(websocket.send_text.await_args_list[0][0][0]) + assert item_done["item"] == reasoning_item + completed = json.loads(websocket.send_text.await_args_list[1][0][0]) + assert completed["response"]["id"] == ResponsesAPIRequestUtils._build_responses_api_response_id( + custom_llm_provider="openai", model_id="dep-1", response_id="resp_1" + ) + assert completed["response"]["output"][0] == reasoning_item + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "failure_frame, expected_status", + [ + ( + { + "type": "error", + "error": { + "type": "invalid_request_error", + "code": "invalid_encrypted_content", + "message": "The encrypted content for item rs_1 could not be verified.", + }, + }, + 400, + ), + ( + { + "type": "response.failed", + "response": { + "id": "resp_1", + "status": "failed", + "error": {"code": "server_error", "message": "upstream blew up"}, + }, + }, + 500, + ), + ], + ) + async def test_backend_to_client_books_failure_frames_as_failures(self, failure_frame, expected_status): + import asyncio + from unittest.mock import AsyncMock + + import websockets.exceptions # noqa: F401 (lazy submodule must be importable) + + websocket = MagicMock() + websocket.send_text = AsyncMock() + backend_ws = MagicMock() + backend_ws.recv = AsyncMock( + side_effect=[ + json.dumps({"type": "response.created", "response": {"id": "resp_1", "status": "in_progress"}}), + json.dumps(failure_frame), + Exception("stop"), + ] + ) + logging_obj = MagicMock() + logging_obj.dispatch_success_handlers = AsyncMock() + logging_obj.dispatch_failure_handlers = AsyncMock() + logging_obj._response_cost_calculator = MagicMock(return_value=0.0) + handler = _make_streaming( + websocket=websocket, + backend_ws=backend_ws, + logging_obj=logging_obj, + request_data={}, + authorized_model="gpt-5.6", + custom_llm_provider="openai", + ) + + await handler.backend_to_client() + await asyncio.sleep(0) + + logging_obj.dispatch_success_handlers.assert_not_awaited() + logging_obj.dispatch_failure_handlers.assert_awaited_once() + exception = logging_obj.dispatch_failure_handlers.await_args[0][0] + assert exception.status_code == expected_status + assert failure_frame.get("error", failure_frame.get("response", {}).get("error"))["message"] in str(exception) + + @pytest.mark.asyncio + async def test_backend_to_client_bills_completed_turns_before_a_failure(self): + import asyncio + from unittest.mock import AsyncMock + + import websockets.exceptions # noqa: F401 (lazy submodule must be importable) + + websocket = MagicMock() + websocket.send_text = AsyncMock() + backend_ws = MagicMock() + backend_ws.recv = AsyncMock( + side_effect=[ + json.dumps( + { + "type": "response.completed", + "response": { + "id": "resp_1", + "status": "completed", + "output": [], + "usage": {"input_tokens": 10, "output_tokens": 5, "total_tokens": 15}, + }, + } + ), + json.dumps({"type": "error", "error": {"type": "invalid_request_error", "message": "bad turn"}}), + Exception("stop"), + ] + ) + logging_obj = MagicMock() + logging_obj.dispatch_success_handlers = AsyncMock() + logging_obj.dispatch_failure_handlers = AsyncMock() + logging_obj._response_cost_calculator = MagicMock(return_value=0.01) + handler = _make_streaming(websocket=websocket, backend_ws=backend_ws, logging_obj=logging_obj, request_data={}) + + await handler.backend_to_client() + await asyncio.sleep(0) + + logging_obj.record_partial_usage_for_failure.assert_called_once() + usage, response_cost = logging_obj.record_partial_usage_for_failure.call_args[0] + assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == (10, 5, 15) + assert response_cost == 0.01 + logging_obj.dispatch_success_handlers.assert_not_awaited() + logging_obj.dispatch_failure_handlers.assert_awaited_once() From c181c927d0b7a4ad214a4dd160c2f16f1373385e Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 18 Sep 2026 15:58:15 -0700 Subject: [PATCH 072/144] fix(proxy): record response.failed frames in background polling --- .../response_polling/background_streaming.py | 7 +++- .../test_response_polling_handler.py | 38 +++++++++++++++++++ 2 files changed, 43 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/response_polling/background_streaming.py b/litellm/proxy/response_polling/background_streaming.py index fac45d4391c..b13042dfb6c 100644 --- a/litellm/proxy/response_polling/background_streaming.py +++ b/litellm/proxy/response_polling/background_streaming.py @@ -75,6 +75,10 @@ class _StreamEventParser: parse: Callable[[str], _StreamEvent] = staticmethod(json.loads) +def _sse_frame_data(frame: str) -> str | None: + return next((line[6:].strip() for line in frame.splitlines() if line.startswith("data: ")), None) + + async def _never_receive() -> Message: await asyncio.Event().wait() raise AssertionError("unreachable") @@ -224,8 +228,7 @@ async def background_streaming_task( if isinstance(chunk, bytes): chunk = chunk.decode("utf-8") - if isinstance(chunk, str) and chunk.startswith("data: "): - chunk_data = chunk[6:].strip() + if isinstance(chunk, str) and (chunk_data := _sse_frame_data(chunk)) is not None: if chunk_data == "[DONE]": break diff --git a/tests/proxy_unit_tests/test_response_polling_handler.py b/tests/proxy_unit_tests/test_response_polling_handler.py index 467c1332325..81ca0114a8d 100644 --- a/tests/proxy_unit_tests/test_response_polling_handler.py +++ b/tests/proxy_unit_tests/test_response_polling_handler.py @@ -1482,6 +1482,44 @@ class TestBackgroundStreamingTerminalEvents: assert final_call.kwargs["status"] == "failed" assert final_call.kwargs["error"] == error_payload + @pytest.mark.asyncio + async def test_named_event_failed_frame_sets_failed_status_and_error(self): + from litellm.proxy.response_polling.background_streaming import ( + background_streaming_task, + ) + + error_payload = { + "code": "cyber_policy", + "message": "Your request was flagged for possible cybersecurity risk and was not completed", + } + failed_event = { + "type": "response.failed", + "sequence_number": 5, + "response": {"id": "resp_123", "status": "failed", "error": error_payload, "output": []}, + } + + async def _body_iterator(): + yield b'data: {"type": "response.in_progress"}\n\n' + yield f"event: response.failed\ndata: {json.dumps(failed_event)}\n\n".encode() + yield b"data: [DONE]\n\n" + + mock_response = Mock() + mock_response.body_iterator = _body_iterator() + handler = AsyncMock(spec=ResponsePollingHandler) + kwargs = _make_background_streaming_kwargs("poll_named_event", handler) + + with patch( # test-quality-ok: the processor is built inside the task, same idiom as the sibling tests + "litellm.proxy.response_polling.background_streaming.ProxyBaseLLMRequestProcessing" + ) as MockProcessor: + MockProcessor.return_value.base_process_llm_request = AsyncMock( + return_value=mock_response + ) + await background_streaming_task(**kwargs) + + final_call = handler.update_state.call_args_list[-1] + assert final_call.kwargs["status"] == "failed" + assert final_call.kwargs["error"] == error_payload + @pytest.mark.asyncio async def test_response_incomplete_sets_incomplete_status_and_details(self): """Test that a response.incomplete stream event results in incomplete status""" From b3cf45e9f232c094e2f2b2a6bf59f609464bf742 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 18 Sep 2026 15:58:24 -0700 Subject: [PATCH 073/144] fix(proxy): drop daily spend batches that cannot be re-sent safely instead of requeueing them --- litellm/proxy/db/db_spend_update_writer.py | 24 ++++++++- litellm/proxy/db/exception_handler.py | 18 +++++++ .../proxy/db/test_db_spend_update_writer.py | 51 ++++++++++++++++++- .../proxy/db/test_exception_handler.py | 22 ++++++++ 4 files changed, 112 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index b2e6f9dc54d..e9967fb0d67 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -31,6 +31,7 @@ from litellm.constants import ( from litellm.litellm_core_utils.litellm_logging import coerce_model_access_groups from litellm.litellm_core_utils.safe_json_loads import safe_json_loads from litellm.proxy._types import ( + DB_CONNECTION_ERROR_TYPES, DB_RETRY_SAFE_ERROR_TYPES, BaseDailySpendTransaction, DailyAgentSpendTransaction, @@ -64,6 +65,7 @@ from litellm.proxy.db.db_transaction_queue.window_spend_update_queue import ( WindowSpendTransaction, WindowSpendUpdateQueue, ) +from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler from litellm.proxy.route_llm_request import ROUTE_ENDPOINT_MAPPING from litellm.proxy.spend_tracking.compression_savings import ( extract_compression_saved_tokens, @@ -157,6 +159,16 @@ class _DailySpendCommit(Protocol[_DailySpendTransactionT]): ) -> None: ... +_DATA_REJECTED_SQLSTATE_CLASSES: Final = frozenset({"22", "23"}) + + +def _daily_spend_commit_failure_is_requeue_safe(e: Exception) -> bool: + if isinstance(e, DB_CONNECTION_ERROR_TYPES): + return isinstance(e, DB_RETRY_SAFE_ERROR_TYPES) + sqlstate: Final = PrismaDBExceptionHandler.postgres_sqlstate(e) + return sqlstate is None or sqlstate[:2] not in _DATA_REJECTED_SQLSTATE_CLASSES + + def _timed_request_duration_ms( payload: dict | SpendLogsPayload, request_status: Literal["success", "failure"], @@ -1319,7 +1331,17 @@ class DBSpendUpdateWriter: proxy_logging_obj=proxy_logging_obj, daily_spend_transactions=cast(dict[str, _DailySpendTransactionT], transactions), ) - except Exception as e: # noqa: BLE001 # the uncommitted rows go back on the queue; the other tables must still flush + except Exception as e: # noqa: BLE001 # whatever failed here, the other tables must still flush + if not _daily_spend_commit_failure_is_requeue_safe(e): + spend_log_error( + "Spend tracking - dropped %d daily %s spend rows: the failed commit may have applied " + "or the database refused the data, so re-sending it is not safe. Error: %s", + len(transactions), + entity_type, + str(e), + exc=e, + ) + return spend_log_error( "Spend tracking - failed to commit daily %s spend updates. " "Re-queued %d rows for retry on next tick. Error: %s", diff --git a/litellm/proxy/db/exception_handler.py b/litellm/proxy/db/exception_handler.py index 2cee5128c66..460bf5db3b1 100644 --- a/litellm/proxy/db/exception_handler.py +++ b/litellm/proxy/db/exception_handler.py @@ -1,6 +1,8 @@ from collections.abc import Awaitable, Callable, Iterator from typing import Any, Final, TypeVar +from pydantic import TypeAdapter, ValidationError + from litellm._logging import verbose_proxy_logger from litellm.proxy._types import ( DB_CONNECTION_ERROR_TYPES, @@ -17,6 +19,8 @@ _TRANSIENT_DB_UNAVAILABLE_MESSAGE: Final = ( "Service Unavailable, the authentication database is temporarily unreachable. Please retry shortly." ) +_DATABASE_ERROR_META: Final = TypeAdapter(dict[str, object]) + def _exception_chain(e: BaseException) -> Iterator[BaseException]: current = e # rebind-ok: advances one link per iteration of the bounded walk @@ -221,6 +225,20 @@ class PrismaDBExceptionHandler: or "write conflict or a deadlock" in error_message ) + @staticmethod + def postgres_sqlstate(e: Exception) -> str | None: + """The SQLSTATE Postgres attached to a failed statement, as prisma surfaces it, or None.""" + import prisma + + if not isinstance(e, _exception_types(prisma.errors.DataError)): + return None + try: + meta: Final = _DATABASE_ERROR_META.validate_python(getattr(e, "meta", None)) + except ValidationError: + return None + code: Final = meta.get("code") + return code if isinstance(code, str) else None + @staticmethod def is_read_only_transaction_error(e: Exception) -> bool: """True iff ``e`` is Postgres SQLSTATE 25006 surfaced through prisma: the diff --git a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py index b27b838133b..155bca656d5 100644 --- a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py +++ b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py @@ -11,7 +11,9 @@ from types import SimpleNamespace from typing import Final from unittest.mock import AsyncMock, MagicMock, call, patch +import httpx import pytest +from prisma.errors import RawQueryError from redis.exceptions import DataError import litellm @@ -2812,14 +2814,15 @@ async def test_failed_window_spend_commit_requeues_the_increments_and_continues_ class _DailySpendFakeDB(_WindowSpendFakeDB): """Records the daily rollup upserts it is handed and fails the ones aimed at one table.""" - def __init__(self, failing_table: str | None) -> None: + def __init__(self, failing_table: str | None, failure: Exception | None = None) -> None: super().__init__() self.failing_table = failing_table + self.failure = failure self.execute_raw_calls: list[Statement] = [] async def execute_raw(self, query: str, *args: object) -> int: if self.failing_table is not None and self.failing_table in query: - raise Exception("connection reset") + raise self.failure if self.failure is not None else Exception("connection reset") self.execute_raw_calls.append((query, args)) return len(args) @@ -2828,6 +2831,50 @@ def _daily_upserts(db: _DailySpendFakeDB, table: str) -> list[Statement]: return [statement for statement in db.execute_raw_calls if table in statement[0]] +def _postgres_rejection(sqlstate: str) -> RawQueryError: + return RawQueryError( + data={"user_facing_error": {"error_code": "P2010", "meta": {"code": sqlstate, "message": "db error"}}} + ) + + +@pytest.mark.parametrize( + ("failure", "lands_on_the_next_tick"), + [ + pytest.param(httpx.ReadTimeout("no reply"), False, id="reply lost after the statement was sent"), + pytest.param(httpx.ConnectError("refused"), True, id="statement never reached the database"), + pytest.param(_postgres_rejection("22021"), False, id="postgres refused the data itself"), + pytest.param(_postgres_rejection("23502"), False, id="postgres refused a constraint violation"), + pytest.param(_postgres_rejection("42P01"), True, id="table missing"), + pytest.param(_postgres_rejection("57014"), True, id="statement cancelled"), + ], +) +@pytest.mark.asyncio +async def test_failed_daily_spend_commit_is_requeued_only_when_the_rows_are_provably_uncommitted( + failure: Exception, lands_on_the_next_tick: bool +): + """A lost reply means the statement may already have applied, and re-sending it stacks a + second increment into the same transaction (LIT-4823); a row Postgres refuses would fail + every tick forever. Both are dropped loudly. Every other failure left nothing committed, + so its rows go back on the queue and land on the next tick.""" + db_writer = DBSpendUpdateWriter() + await db_writer.daily_spend_update_queue.add_update({"user-key": _daily_txn(user_id="user-1")}) + db = _DailySpendFakeDB(failing_table="LiteLLM_DailyUserSpend", failure=failure) + db_writer._flush_tool_discovery_queue = AsyncMock() + proxy_logging_obj = MagicMock() + proxy_logging_obj.failure_handler = AsyncMock() + + await db_writer._commit_spend_updates_to_db_without_redis_buffer( + prisma_client=_WindowSpendFakePrisma(db), n_retry_times=0, proxy_logging_obj=proxy_logging_obj + ) + db.failing_table = None + await db_writer._commit_spend_updates_to_db_without_redis_buffer( + prisma_client=_WindowSpendFakePrisma(db), n_retry_times=0, proxy_logging_obj=proxy_logging_obj + ) + + assert len(_daily_upserts(db, "LiteLLM_DailyUserSpend")) == (1 if lands_on_the_next_tick else 0) + assert db_writer.daily_spend_update_queue.update_queue.empty() + + @pytest.mark.asyncio async def test_failed_daily_spend_commit_requeues_the_rows_and_flushes_the_other_tables(): """With the Redis buffer off, a daily batch that failed to commit was discarded along diff --git a/tests/test_litellm/proxy/db/test_exception_handler.py b/tests/test_litellm/proxy/db/test_exception_handler.py index 3f009137a1c..26ac1ea65ad 100644 --- a/tests/test_litellm/proxy/db/test_exception_handler.py +++ b/tests/test_litellm/proxy/db/test_exception_handler.py @@ -665,6 +665,28 @@ def test_is_deadlock_error_excludes_non_deadlocks(error): assert PrismaDBExceptionHandler.is_deadlock_error(error) is False +@pytest.mark.parametrize( + ("error", "sqlstate"), + [ + ( + RawQueryError( + data={"user_facing_error": {"error_code": "P2010", "meta": {"code": "22021", "message": "m"}}} + ), + "22021", + ), + (RawQueryError(data={"user_facing_error": {"error_code": "P2010", "meta": {"message": "m"}}}), None), + (RawQueryError(data={"user_facing_error": {"error_code": "P2010", "meta": {"code": 42, "message": "m"}}}), None), + (prisma_errors.DataError(data={"user_facing_error": {"meta": None}}), None), + (PrismaError("db error"), None), + (httpx.ReadTimeout("no reply"), None), + ], +) +def test_postgres_sqlstate_reads_the_code_prisma_attached_to_the_failed_statement(error, sqlstate): + """Only a prisma data error carrying Postgres's own error code yields a SQLSTATE; a + codeless or malformed payload, an engine-level error, and a transport error yield None.""" + assert PrismaDBExceptionHandler.postgres_sqlstate(error) == sqlstate + + READ_ONLY_CONNECTOR_ERROR: Final = ( "Error occurred during query execution:\nConnectorError(ConnectorError { user_facing_error: None, " 'kind: QueryError(PostgresError { code: "25006", message: "cannot execute UPDATE in a read-only transaction", ' From 8005856411ae43091c86ba2632d389916b8fa0ec Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 18 Sep 2026 16:03:27 -0700 Subject: [PATCH 074/144] ci(build_and_test): seed the routing strategy through /config/update --- .circleci/config.yml | 6 ++++++ proxy_server_config.yaml | 1 - 2 files changed, 6 insertions(+), 1 deletion(-) diff --git a/.circleci/config.yml b/.circleci/config.yml index df17a9e4402..87f1ee604cf 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -1785,6 +1785,12 @@ jobs: - wait_for_service: url: http://localhost:4000 timeout: "300" + - run: + name: Seed the routing strategy through /config/update + command: | + curl --noproxy '*' -sSf -X POST http://localhost:4000/config/update \ + -H 'Authorization: Bearer sk-1234' -H 'Content-Type: application/json' \ + -d '{"router_settings": {"routing_strategy": "usage-based-routing-v2"}}' - run: name: Run tests command: | diff --git a/proxy_server_config.yaml b/proxy_server_config.yaml index 73990153227..703d56bc0cd 100644 --- a/proxy_server_config.yaml +++ b/proxy_server_config.yaml @@ -213,7 +213,6 @@ files_settings: api_key: os.environ/OPENAI_API_KEY router_settings: - routing_strategy: usage-based-routing-v2 redis_host: os.environ/REDIS_HOST redis_password: os.environ/REDIS_PASSWORD redis_port: os.environ/REDIS_PORT From 87ac68709c5b09b1f08f1e868913ab2a55bd3c65 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 18 Sep 2026 16:05:16 -0700 Subject: [PATCH 075/144] fix(claude_code_gateway): single-use device codes across replicas, protobuf telemetry, CLI user route access --- litellm/proxy/_lazy_openapi_snapshot.json | 8 +- litellm/proxy/_types.py | 7 + .../anthropic_endpoints/gateway_endpoints.py | 54 +++- .../proxy/common_utils/http_parsing_utils.py | 10 +- .../test_gateway_endpoints.py | 270 ++++++++++++++---- .../proxy/auth/test_route_checks.py | 30 ++ .../common_utils/test_http_parsing_utils.py | 7 + ui/litellm-dashboard/src/lib/http/schema.d.ts | 12 + 8 files changed, 337 insertions(+), 61 deletions(-) diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index 8a8d08c6887..80527f50d10 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -5235,6 +5235,12 @@ } } }, + "claude_code_gateway": { + "components": { + "schemas": {} + }, + "paths": {} + }, "claude_code_marketplace": { "components": { "schemas": { @@ -19394,7 +19400,7 @@ } } }, - "description": "\n Unified rate-limit error.\n\n Every rate-limit condition surfaced by litellm \u2014 whether it originated from\n an upstream LLM provider, a vendor batch endpoint, or one of litellm's own\n proxy-side limiters (parallel-requests, dynamic-rate, batch-rate, budget,\n max-iterations, etc.) \u2014 is raised as an instance of this class.\n\n The :attr:`category` attribute lets callers distinguish the source. See\n :class:`RateLimitErrorCategory` for the available values.\n " + "description": "\nUnified rate-limit error.\n\nEvery rate-limit condition surfaced by litellm \u2014 whether it originated from\nan upstream LLM provider, a vendor batch endpoint, or one of litellm's own\nproxy-side limiters (parallel-requests, dynamic-rate, batch-rate, budget,\nmax-iterations, etc.) \u2014 is raised as an instance of this class.\n\nThe :attr:`category` attribute lets callers distinguish the source. See\n:class:`RateLimitErrorCategory` for the available values.\n" }, "500": { "content": { diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 545b555f63f..e5042430568 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -508,6 +508,8 @@ class LiteLLMRoutes(enum.Enum): anthropic_routes = [ "/v1/messages", "/v1/messages/count_tokens", + "/claude_code_gateway/v1/messages", + "/claude_code_gateway/v1/messages/count_tokens", "/v1/skills", "/v1/skills/{skill_id}", "/claude-code/marketplace.json", @@ -885,6 +887,11 @@ class LiteLLMRoutes(enum.Enum): # of; a caller who administers none gets an empty result set. "/organization/daily/activity", "/user/available_roles", # read-only role metadata; any authenticated user may read + # Claude Code gateway: the signed-in CLI fetches its managed settings and posts its own telemetry + "/claude_code_gateway/managed/settings", + "/claude_code_gateway/v1/metrics", + "/claude_code_gateway/v1/logs", + "/claude_code_gateway/v1/traces", "/user/list", # org admins checked in endpoint; non-admins get 403 "/management/v1/users/bulk_delete", # proxy admins delete anyone, org admins only their orgs' users; others 403 "/model/{model_id}/update", diff --git a/litellm/proxy/anthropic_endpoints/gateway_endpoints.py b/litellm/proxy/anthropic_endpoints/gateway_endpoints.py index 5a4a4d0eb78..991dc67ab82 100644 --- a/litellm/proxy/anthropic_endpoints/gateway_endpoints.py +++ b/litellm/proxy/anthropic_endpoints/gateway_endpoints.py @@ -23,8 +23,10 @@ from typing import Final from fastapi import APIRouter, Depends, Request, Response from fastapi.responses import JSONResponse -from pydantic import BaseModel, Field, TypeAdapter +from pydantic import BaseModel, Field, TypeAdapter, ValidationError +from litellm._logging import verbose_proxy_logger +from litellm.caching.dual_cache import DualCache from litellm.constants import ( CLI_JWT_EXPIRATION_HOURS, CLI_SSO_SESSION_TTL_SECONDS, @@ -45,9 +47,10 @@ _POST_ONLY: Final = ["POST"] # mutable-ok: FastAPI's add_api_route only accepts class _GatewaySessionData(BaseModel): user_id: str - user_role: str | None = None + user_role: str | None models: list[str] = Field(default_factory=list) teams: tuple[str, ...] = () + team_details: object | None = None class _OAuthErrorBody(BaseModel): @@ -212,19 +215,51 @@ async def device_authorization(request: Request) -> JSONResponse: def _mint_access_token_from_flow(flow: Mapping[str, object]) -> str: from litellm.proxy._types import LiteLLM_UserTable from litellm.proxy.auth.auth_checks import ExperimentalUIJWTToken + from litellm.proxy.management_endpoints.ui_sso import selected_cli_sso_team_detail - raw_session_data: Final = flow.get("session_data") - if not isinstance(raw_session_data, dict): - raise _oauth_error(status_code=400, error="authorization_pending") + try: + session_data: Final = _GatewaySessionData.model_validate(flow.get("session_data")) + except ValidationError as err: + verbose_proxy_logger.warning("Claude Code gateway login session is malformed: %s", err) + raise _oauth_error( + status_code=400, error="invalid_grant", description="The login session is malformed; sign in again" + ) from err - session_data: Final = _GatewaySessionData.model_validate(raw_session_data) team_id: Final = session_data.teams[0] if session_data.teams else None + selected_team: Final = selected_cli_sso_team_detail(team_details=session_data.team_details, team_id=team_id) + if selected_team is None: + raise _oauth_error( + status_code=400, + error="invalid_grant", + description=f"Could not resolve the model grants for team {team_id}; sign in again", + ) + user_info: Final = LiteLLM_UserTable( user_id=session_data.user_id, user_role=session_data.user_role, models=session_data.models, ) - return ExperimentalUIJWTToken.get_cli_jwt_auth_token(user_info=user_info, team_id=team_id) + return ExperimentalUIJWTToken.get_cli_jwt_auth_token( + user_info=user_info, + team_id=team_id, + team_alias=selected_team.team_alias, + team_models=selected_team.team_models, + team_model_aliases=selected_team.team_model_aliases, + max_budget=None, + ) + + +async def _claim_device_code(device_code: str, cache: DualCache) -> bool: + from litellm.proxy.management_endpoints.ui_sso import ( + _get_cli_sso_flow_cache_key, # pyright: ignore[reportPrivateUsage] # shared device-flow helper + ) + + claims: Final = await cache.async_increment_cache( + key=f"{_get_cli_sso_flow_cache_key(device_code)}:claimed", + value=1, + ttl=CLI_SSO_SESSION_TTL_SECONDS, + ) + return claims == 1 async def _handle_device_code_grant(device_code: str | None) -> JSONResponse: @@ -249,12 +284,15 @@ async def _handle_device_code_grant(device_code: str | None) -> JSONResponse: if not flow.get("sso_complete") or not flow.get("user_code_verified"): return _oauth_error_response(_oauth_error(status_code=400, error="authorization_pending")) + if not await _claim_device_code(device_code, cli_sso_session_cache): + return _oauth_error_response(_oauth_error(status_code=400, error="expired_token")) + + await cli_sso_session_cache.async_delete_cache(key=_get_cli_sso_flow_cache_key(device_code)) try: access_token: Final = _mint_access_token_from_flow(flow) except _OAuthError as err: return _oauth_error_response(err) - cli_sso_session_cache.delete_cache(key=_get_cli_sso_flow_cache_key(device_code)) body: Final = _AccessTokenBody(access_token=access_token, expires_in=CLI_JWT_EXPIRATION_HOURS * _SECONDS_PER_HOUR) return JSONResponse(content=body.model_dump()) diff --git a/litellm/proxy/common_utils/http_parsing_utils.py b/litellm/proxy/common_utils/http_parsing_utils.py index f5b6a0a766d..592060e84ee 100644 --- a/litellm/proxy/common_utils/http_parsing_utils.py +++ b/litellm/proxy/common_utils/http_parsing_utils.py @@ -18,6 +18,8 @@ from litellm.types.router import Deployment _FORM_CONTENT_TYPES: Final[frozenset[str]] = frozenset({"application/x-www-form-urlencoded", "multipart/form-data"}) +_PROTOBUF_CONTENT_TYPES: Final[frozenset[str]] = frozenset({"application/x-protobuf", "application/protobuf"}) + _ANNOTATION_QUALIFIERS: Final[frozenset[object]] = frozenset({Annotated, NotRequired, ReadOnly, Required}) @@ -44,6 +46,10 @@ def is_json_content_type(content_type: str) -> bool: return _normalize_media_type(content_type) == "application/json" +def _is_protobuf_content_type(content_type: str) -> bool: + return _normalize_media_type(content_type) in _PROTOBUF_CONTENT_TYPES + + def _unqualified(annotation: object) -> object: """Which qualifiers ``get_type_hints`` already stripped varies by interpreter version, so peel them all.""" if get_origin(annotation) not in _ANNOTATION_QUALIFIERS: @@ -133,7 +139,9 @@ async def _read_request_body(request: Request | None) -> dict: _request_headers: Final[dict] = _safe_get_request_headers(request=request) content_type: Final = _request_headers.get("content-type", "") - if _is_form_content_type(content_type): + if _is_protobuf_content_type(content_type): + parsed_body = {} + elif _is_form_content_type(content_type): try: form_data: Final = await request.form() except Exception as e: diff --git a/tests/test_litellm/proxy/anthropic_endpoints/test_gateway_endpoints.py b/tests/test_litellm/proxy/anthropic_endpoints/test_gateway_endpoints.py index 8645f4a8680..c0a39b95c40 100644 --- a/tests/test_litellm/proxy/anthropic_endpoints/test_gateway_endpoints.py +++ b/tests/test_litellm/proxy/anthropic_endpoints/test_gateway_endpoints.py @@ -5,58 +5,161 @@ Covers the OAuth device-flow surface (RFC 8414 discovery, RFC 8628 device authorization + token), managed settings, OTLP ingestion, and the enable flag. """ -from contextlib import contextmanager -from typing import Any, Iterator, Optional -from unittest.mock import patch +import asyncio +from collections.abc import Iterator, Mapping +from contextlib import ExitStack, contextmanager +from types import MappingProxyType +from typing import Final +from unittest.mock import AsyncMock, MagicMock, patch +import httpx import pytest from fastapi import FastAPI from fastapi.testclient import TestClient from litellm.caching.dual_cache import DualCache +from litellm.proxy._types import ProxyException from litellm.proxy.anthropic_endpoints import gateway_endpoints -from litellm.proxy.management_endpoints.ui_sso import _get_cli_sso_flow_cache_key +from litellm.proxy.management_endpoints.ui_sso import _get_cli_sso_flow_cache_key, _set_cli_sso_flow + +_DEVICE_CODE_GRANT: Final = "urn:ietf:params:oauth:grant-type:device_code" +_MASTER_KEY: Final = "sk-master-key" +_MINT: Final = "litellm.proxy.auth.auth_checks.ExperimentalUIJWTToken.get_cli_jwt_auth_token" +_PROTOBUF_BODY: Final = b"\x0a\x05hello\x12\x03{{{" +_COMPLETED_SESSION: Final = MappingProxyType( + { + "user_id": "user-123", + "user_role": "internal_user", + "models": ["claude-sonnet-4-5"], + "teams": ["team-a"], + "team_details": [ + { + "team_id": "team-a", + "team_alias": "Team A", + "team_models": ["claude-sonnet-4-5"], + "team_model_aliases": None, + } + ], + } +) + + +class _SharedRedisFake: + def __init__(self) -> None: + self.values: Mapping[str, object] = MappingProxyType({}) + self.counters: Mapping[str, float] = MappingProxyType({}) + + def set_cache(self, key: str, value: object, **kwargs: object) -> None: + self.values = MappingProxyType({**self.values, key: value}) + + def get_cache(self, key: str, **kwargs: object) -> object: + return self.values.get(key) + + def delete_cache(self, key: str) -> None: + self.values = MappingProxyType({name: value for name, value in self.values.items() if name != key}) + + async def async_delete_cache(self, key: str) -> None: + self.delete_cache(key) + + async def async_increment(self, key: str, value: float, **kwargs: object) -> float: + incremented: Final = self.counters.get(key, 0) + value + self.counters = MappingProxyType({**self.counters, key: incremented}) + return incremented + + +def _replica(redis: _SharedRedisFake) -> DualCache: + return DualCache(redis_cache=redis, default_in_memory_ttl=600) # pyright: ignore[reportArgumentType] # duck-typed Redis double + + +def _real_auth_proxy_attrs() -> Mapping[str, object]: + proxy_logging_obj: Final = MagicMock() + proxy_logging_obj.internal_usage_cache.dual_cache = AsyncMock() + proxy_logging_obj.post_call_failure_hook = AsyncMock(return_value=None) + return MappingProxyType( + { + "master_key": _MASTER_KEY, + "prisma_client": None, + "user_api_key_cache": DualCache(), + "proxy_logging_obj": proxy_logging_obj, + "llm_router": None, + "llm_model_list": [], + "user_custom_auth": None, + "litellm_proxy_admin_name": "admin", + "jwt_handler": None, + "open_telemetry_logger": None, + "model_max_budget_limiter": MagicMock(), + } + ) @contextmanager def _gateway_env( *, enabled: bool = True, - managed_settings: Optional[dict[str, Any]] = None, + managed_settings: Mapping[str, object] | None = None, + cache: DualCache | None = None, + real_auth: bool = False, ) -> Iterator[tuple[TestClient, DualCache]]: - general_settings: dict[str, Any] = {"enable_claude_code_gateway": enabled} - if managed_settings is not None: - general_settings["claude_code_gateway_managed_settings"] = managed_settings - cache = DualCache(default_in_memory_ttl=600) + general_settings: Final = { + "enable_claude_code_gateway": enabled, + **({} if managed_settings is None else {"claude_code_gateway_managed_settings": dict(managed_settings)}), + } + session_cache: Final = cache or DualCache(default_in_memory_ttl=600) - app = FastAPI() + app: Final = FastAPI() app.include_router(gateway_endpoints.router) - async def _fake_auth() -> Any: + async def _fake_auth() -> object: return object() - app.dependency_overrides[gateway_endpoints.user_api_key_auth] = _fake_auth - - with patch("litellm.proxy.proxy_server.general_settings", general_settings), patch( - "litellm.proxy.proxy_server.cli_sso_session_cache", cache - ): + with ExitStack() as stack: + stack.enter_context( + patch( # test-quality-ok: the gateway reads this proxy_server module global and has no injection seam + "litellm.proxy.proxy_server.general_settings", general_settings + ) + ) + stack.enter_context( + patch( # test-quality-ok: the CLI SSO flow cache is this proxy_server module global shared with ui_sso + "litellm.proxy.proxy_server.cli_sso_session_cache", session_cache + ) + ) + if real_auth: + for name, value in _real_auth_proxy_attrs().items(): + stack.enter_context(patch(f"litellm.proxy.proxy_server.{name}", value)) + else: + app.dependency_overrides[gateway_endpoints.user_api_key_auth] = _fake_auth with TestClient(app) as client: - yield client, cache + yield client, session_cache -def _complete_flow(cache: DualCache, device_code: str) -> None: - key = _get_cli_sso_flow_cache_key(device_code) - flow = cache.get_cache(key=key) - assert isinstance(flow, dict) - flow["sso_complete"] = True - flow["user_code_verified"] = True - flow["session_data"] = { - "user_id": "user-123", - "user_role": "internal_user", - "models": ["claude-sonnet-4-5"], - "teams": ["team-a"], +def _start_device_flow(client: TestClient) -> str: + return client.post("/claude_code_gateway/oauth/device_authorization").json()["device_code"] + + +def _request_token(client: TestClient, device_code: str) -> httpx.Response: + return client.post( + "/claude_code_gateway/oauth/token", + data={"grant_type": _DEVICE_CODE_GRANT, "device_code": device_code}, + ) + + +def _completed_flow(session_data: Mapping[str, object] = _COMPLETED_SESSION) -> dict[str, object]: + return { + "poll_secret_hash": "unused", + "user_code_hash": "unused", + "sso_complete": True, + "user_code_verified": True, + "session_data": dict(session_data), } - cache.set_cache(key=key, value=flow, ttl=600) + + +def _complete_flow( + cache: DualCache, device_code: str, session_data: Mapping[str, object] = _COMPLETED_SESSION +) -> None: + key: Final = _get_cli_sso_flow_cache_key(device_code) + flow: Final = cache.get_cache(key=key) + assert isinstance(flow, dict) + cache.set_cache(key=key, value={**flow, **_completed_flow(session_data)}, ttl=600) def test_discovery_shape(): @@ -105,28 +208,18 @@ def test_device_authorization_returns_rfc8628_shape_and_persists_flow(): def test_token_authorization_pending_before_browser_completes(): with _gateway_env() as (client, _): - device_code = client.post("/claude_code_gateway/oauth/device_authorization").json()["device_code"] - resp = client.post( - "/claude_code_gateway/oauth/token", - data={"grant_type": "urn:ietf:params:oauth:grant-type:device_code", "device_code": device_code}, - ) + resp = _request_token(client, _start_device_flow(client)) assert resp.status_code == 400 assert resp.json()["error"] == "authorization_pending" def test_token_success_mints_bearer_and_is_single_use(): with _gateway_env() as (client, cache): - device_code = client.post("/claude_code_gateway/oauth/device_authorization").json()["device_code"] + device_code = _start_device_flow(client) _complete_flow(cache, device_code) - with patch( - "litellm.proxy.auth.auth_checks.ExperimentalUIJWTToken.get_cli_jwt_auth_token", - return_value="sk-litellm-session-token", - ) as mint: - resp = client.post( - "/claude_code_gateway/oauth/token", - data={"grant_type": "urn:ietf:params:oauth:grant-type:device_code", "device_code": device_code}, - ) + with patch(_MINT, return_value="sk-litellm-session-token") as mint: + resp = _request_token(client, device_code) assert resp.status_code == 200 body = resp.json() assert body["access_token"] == "sk-litellm-session-token" @@ -136,22 +229,77 @@ def test_token_success_mints_bearer_and_is_single_use(): called_user = mint.call_args.kwargs["user_info"] assert called_user.user_id == "user-123" assert mint.call_args.kwargs["team_id"] == "team-a" + assert mint.call_args.kwargs["team_alias"] == "Team A" + assert mint.call_args.kwargs["team_models"] == ("claude-sonnet-4-5",) # Single-use: the flow is deleted, so a replay returns expired_token. - replay = client.post( - "/claude_code_gateway/oauth/token", - data={"grant_type": "urn:ietf:params:oauth:grant-type:device_code", "device_code": device_code}, - ) + replay = _request_token(client, device_code) assert replay.status_code == 400 assert replay.json()["error"] == "expired_token" +def test_token_teamless_user_mints_without_a_team(): + with _gateway_env() as (client, cache): + device_code = _start_device_flow(client) + _complete_flow(cache, device_code, session_data={**_COMPLETED_SESSION, "teams": [], "team_details": []}) + with patch(_MINT, return_value="sk-litellm-session-token") as mint: + resp = _request_token(client, device_code) + assert resp.status_code == 200 + assert mint.call_args.kwargs["team_id"] is None + assert mint.call_args.kwargs["team_models"] == () + + +def test_token_malformed_session_is_invalid_grant(): + with _gateway_env() as (client, cache): + device_code = _start_device_flow(client) + _complete_flow(cache, device_code, session_data={"user_role": "internal_user"}) + with patch(_MINT) as mint: + resp = _request_token(client, device_code) + assert resp.status_code == 400 + assert resp.json()["error"] == "invalid_grant" + mint.assert_not_called() + + +def test_token_unknown_team_grants_is_invalid_grant(): + with _gateway_env() as (client, cache): + device_code = _start_device_flow(client) + _complete_flow(cache, device_code, session_data={**_COMPLETED_SESSION, "team_details": []}) + with patch(_MINT) as mint: + resp = _request_token(client, device_code) + assert resp.status_code == 400 + assert resp.json()["error"] == "invalid_grant" + mint.assert_not_called() + + +def test_token_mints_on_a_replica_that_did_not_start_the_login(): + redis: Final = _SharedRedisFake() + device_code: Final = "cli-shared-login-code" + _set_cli_sso_flow(login_id=device_code, cache=_replica(redis), flow=_completed_flow()) + + with _gateway_env(cache=_replica(redis)) as (client, _), patch(_MINT, return_value="sk-session") as mint: + resp = _request_token(client, device_code) + assert resp.status_code == 200 + assert resp.json()["access_token"] == "sk-session" + assert mint.call_args.kwargs["team_id"] == "team-a" + + +def test_token_refuses_a_device_code_another_replica_already_claimed(): + redis: Final = _SharedRedisFake() + replica_a: Final = _replica(redis) + device_code: Final = "cli-shared-login-code" + _set_cli_sso_flow(login_id=device_code, cache=replica_a, flow=_completed_flow()) + assert asyncio.run(gateway_endpoints._claim_device_code(device_code, replica_a)) is True + + with _gateway_env(cache=_replica(redis)) as (client, _), patch(_MINT) as mint: + resp = _request_token(client, device_code) + assert resp.status_code == 400 + assert resp.json()["error"] == "expired_token" + mint.assert_not_called() + + def test_token_unknown_device_code_is_expired_token(): with _gateway_env() as (client, _): - resp = client.post( - "/claude_code_gateway/oauth/token", - data={"grant_type": "urn:ietf:params:oauth:grant-type:device_code", "device_code": "cli-does-not-exist"}, - ) + resp = _request_token(client, "cli-does-not-exist") assert resp.status_code == 400 assert resp.json()["error"] == "expired_token" @@ -213,6 +361,26 @@ def test_otlp_endpoints_404_when_disabled(signal: str): assert resp.status_code == 404 +def test_otlp_protobuf_body_is_accepted_through_real_auth(): + with _gateway_env(real_auth=True) as (client, _): + resp = client.post( + "/claude_code_gateway/v1/metrics", + content=_PROTOBUF_BODY, + headers={"Authorization": f"Bearer {_MASTER_KEY}", "Content-Type": "application/x-protobuf"}, + ) + assert resp.status_code == 200 + + +def test_otlp_without_a_bearer_is_rejected_by_real_auth(): + with _gateway_env(real_auth=True) as (client, _), pytest.raises(ProxyException) as exc_info: + client.post( + "/claude_code_gateway/v1/metrics", + content=_PROTOBUF_BODY, + headers={"Content-Type": "application/x-protobuf"}, + ) + assert exc_info.value.code == "401" + + def test_messages_gated_by_enable_flag(): with _gateway_env(enabled=False) as (client, _): resp = client.post("/claude_code_gateway/v1/messages", json={"model": "claude-sonnet-4-5", "messages": []}) diff --git a/tests/test_litellm/proxy/auth/test_route_checks.py b/tests/test_litellm/proxy/auth/test_route_checks.py index 72c59223549..3ec4d2e63ad 100644 --- a/tests/test_litellm/proxy/auth/test_route_checks.py +++ b/tests/test_litellm/proxy/auth/test_route_checks.py @@ -910,6 +910,36 @@ def test_anthropic_count_tokens_route_accessible_to_internal_users(): assert RouteChecks.is_llm_api_route("/v1/messages") is True +_CLAUDE_CODE_GATEWAY_ROUTES: Final = ( + "/claude_code_gateway/v1/messages", + "/claude_code_gateway/v1/messages/count_tokens", + "/claude_code_gateway/managed/settings", + "/claude_code_gateway/v1/metrics", + "/claude_code_gateway/v1/logs", + "/claude_code_gateway/v1/traces", +) + + +@pytest.mark.parametrize("route", _CLAUDE_CODE_GATEWAY_ROUTES) +@pytest.mark.parametrize( + "role", [LitellmUserRoles.INTERNAL_USER.value, LitellmUserRoles.INTERNAL_USER_VIEW_ONLY.value] +) +def test_claude_code_gateway_routes_open_to_signed_in_cli_users(role: str, route: str): + user_obj: Final = LiteLLM_UserTable(user_id="test_user", user_email="test@example.com", user_role=role) + valid_token: Final = UserAPIKeyAuth(user_id="test_user", user_role=role) + request: Final = MagicMock(spec=Request) + request.query_params = {} + + RouteChecks.non_proxy_admin_allowed_routes_check( + user_obj=user_obj, + _user_role=role, + route=route, + request=request, + valid_token=valid_token, + request_data={}, + ) + + def test_virtual_key_llm_api_routes_allows_registered_pass_through_endpoints(): """ Virtual keys with llm_api_routes can access auth=true pass-through endpoints only when diff --git a/tests/test_litellm/proxy/common_utils/test_http_parsing_utils.py b/tests/test_litellm/proxy/common_utils/test_http_parsing_utils.py index 72cd7a218d3..bd9912a96ac 100644 --- a/tests/test_litellm/proxy/common_utils/test_http_parsing_utils.py +++ b/tests/test_litellm/proxy/common_utils/test_http_parsing_utils.py @@ -573,6 +573,13 @@ async def test_lone_surrogate_escape_is_rejected_with_400(content: bytes): assert parsed["messages"][0]["content"] == "say ok \U0001F600" +@pytest.mark.asyncio +@pytest.mark.parametrize("media_type", ["application/x-protobuf", "application/protobuf"]) +async def test_protobuf_body_is_left_unparsed(media_type: str): + request = _starlette_request(b"\x0a\x05hello\x12\x03{{{", media_type) + assert await _read_request_body(request) == {} + + @pytest.mark.asyncio async def test_get_form_data(): """ diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 645d6ec5ac4..bc972f913a4 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -26606,6 +26606,13 @@ export interface components { * @description cancel the in-flight upstream LLM request (non-streaming) when the client disconnects, freeing backend capacity (e.g. a vLLM GPU slot); the request is logged as a 499 failure */ cancel_on_disconnect?: boolean | null; + /** + * Claude Code Gateway Managed Settings + * @description Claude Code managed-settings.json served verbatim at the gateway's /claude_code_gateway/managed/settings endpoint. When unset the endpoint returns 404 (no managed policy) + */ + claude_code_gateway_managed_settings?: { + [key: string]: unknown; + } | null; /** * Completion Model * @description proxy level default model for all chat completion calls @@ -26700,6 +26707,11 @@ export interface components { * @description If True, disables ownership enforcement on Responses API ids. Keys may then retrieve, cancel, delete, and chain from any response id, including ids belonging to another user or team and ids this proxy never issued. WARNING: this removes tenant isolation on /v1/responses */ disable_responses_id_security?: boolean | null; + /** + * Enable Claude Code Gateway + * @description serve the Claude Code gateway protocol (https://code.claude.com/docs/en/claude-apps-gateway) under /claude_code_gateway: OAuth device-flow sign-in reusing proxy SSO, plus managed settings and OTLP telemetry ingestion. Off by default + */ + enable_claude_code_gateway?: boolean | null; /** * Enable Openai Websocket Passthrough * @description Serve the OpenAI pass-through WebSocket route, which relays frames to OpenAI under the proxy's own provider credential without reading them. Off by default. From 5db2a0c8850385b9e56b17bcd8053bab4fac0b59 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 18 Sep 2026 16:05:42 -0700 Subject: [PATCH 076/144] test(proxy): type the sqlstate test parameters --- tests/test_litellm/proxy/db/test_exception_handler.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_litellm/proxy/db/test_exception_handler.py b/tests/test_litellm/proxy/db/test_exception_handler.py index 26ac1ea65ad..f7cc5e3ed83 100644 --- a/tests/test_litellm/proxy/db/test_exception_handler.py +++ b/tests/test_litellm/proxy/db/test_exception_handler.py @@ -681,7 +681,7 @@ def test_is_deadlock_error_excludes_non_deadlocks(error): (httpx.ReadTimeout("no reply"), None), ], ) -def test_postgres_sqlstate_reads_the_code_prisma_attached_to_the_failed_statement(error, sqlstate): +def test_postgres_sqlstate_reads_the_code_prisma_attached_to_the_failed_statement(error: Exception, sqlstate: str | None): """Only a prisma data error carrying Postgres's own error code yields a SQLSTATE; a codeless or malformed payload, an engine-level error, and a transport error yield None.""" assert PrismaDBExceptionHandler.postgres_sqlstate(error) == sqlstate From 4356fc58d82a30e470ce25caa0949b002c51f3b4 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 18 Sep 2026 16:05:57 -0700 Subject: [PATCH 077/144] chore(proxy): restore the CI-generated lazy OpenAPI snapshot --- litellm/proxy/_lazy_openapi_snapshot.json | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index 40b64160b71..213cd88b6ce 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -19616,7 +19616,7 @@ } } }, - "description": "\nUnified rate-limit error.\n\nEvery rate-limit condition surfaced by litellm \u2014 whether it originated from\nan upstream LLM provider, a vendor batch endpoint, or one of litellm's own\nproxy-side limiters (parallel-requests, dynamic-rate, batch-rate, budget,\nmax-iterations, etc.) \u2014 is raised as an instance of this class.\n\nThe :attr:`category` attribute lets callers distinguish the source. See\n:class:`RateLimitErrorCategory` for the available values.\n" + "description": "\n Unified rate-limit error.\n\n Every rate-limit condition surfaced by litellm \u2014 whether it originated from\n an upstream LLM provider, a vendor batch endpoint, or one of litellm's own\n proxy-side limiters (parallel-requests, dynamic-rate, batch-rate, budget,\n max-iterations, etc.) \u2014 is raised as an instance of this class.\n\n The :attr:`category` attribute lets callers distinguish the source. See\n :class:`RateLimitErrorCategory` for the available values.\n " }, "500": { "content": { From 01d8d3c21807431c93d76cb3c13fe1516f1191fe Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 18 Sep 2026 16:18:15 -0700 Subject: [PATCH 078/144] fix(claude_code_gateway): wrap managed settings in the uuid, checksum, settings envelope the client requires --- .../anthropic_endpoints/gateway_endpoints.py | 14 +++++++-- .../test_gateway_endpoints.py | 31 ++++++++++++++----- 2 files changed, 35 insertions(+), 10 deletions(-) diff --git a/litellm/proxy/anthropic_endpoints/gateway_endpoints.py b/litellm/proxy/anthropic_endpoints/gateway_endpoints.py index 991dc67ab82..cc3106fce53 100644 --- a/litellm/proxy/anthropic_endpoints/gateway_endpoints.py +++ b/litellm/proxy/anthropic_endpoints/gateway_endpoints.py @@ -80,6 +80,12 @@ class _AccessTokenBody(BaseModel): token_type: str = "Bearer" +class _ManagedSettingsBody(BaseModel): + uuid: str + checksum: str + settings: dict[str, object] + + def _general_settings() -> Mapping[str, object]: from litellm.proxy.proxy_server import general_settings @@ -333,12 +339,14 @@ async def managed_settings(request: Request) -> Response: if settings is None: return Response(status_code=404) - body: Final = json.dumps(settings, sort_keys=True, separators=(",", ":")) - etag: Final = '"' + hashlib.sha256(body.encode("utf-8")).hexdigest() + '"' + canonical: Final = json.dumps(settings, sort_keys=True, separators=(",", ":")) + checksum: Final = "sha256:" + hashlib.sha256(canonical.encode("utf-8")).hexdigest() + etag: Final = f'"{checksum}"' headers: Final = MappingProxyType({"ETag": etag}) if request.headers.get("If-None-Match") == etag: return Response(status_code=304, headers=headers) - return Response(content=body, media_type="application/json", headers=headers) + body: Final = _ManagedSettingsBody(uuid=checksum, checksum=checksum, settings=settings) + return Response(content=body.model_dump_json(), media_type="application/json", headers=headers) def _accept_otlp() -> Response: diff --git a/tests/test_litellm/proxy/anthropic_endpoints/test_gateway_endpoints.py b/tests/test_litellm/proxy/anthropic_endpoints/test_gateway_endpoints.py index c0a39b95c40..158fe253796 100644 --- a/tests/test_litellm/proxy/anthropic_endpoints/test_gateway_endpoints.py +++ b/tests/test_litellm/proxy/anthropic_endpoints/test_gateway_endpoints.py @@ -327,18 +327,35 @@ def test_managed_settings_404_when_unset(): assert resp.status_code == 404 -def test_managed_settings_returns_json_with_etag_and_304(): +def test_managed_settings_returns_client_envelope_and_304_on_cached_checksum(): settings = {"permissions": {"defaultMode": "acceptEdits"}, "env": {"FOO": "bar"}} with _gateway_env(managed_settings=settings) as (client, _): resp = client.get("/claude_code_gateway/managed/settings") assert resp.status_code == 200 - assert resp.json() == settings - etag = resp.headers["ETag"] - assert etag + body = resp.json() + assert body["settings"] == settings + checksum = body["checksum"] + assert checksum.startswith("sha256:") + assert body["uuid"] == checksum + assert resp.headers["ETag"] == f'"{checksum}"' - not_modified = client.get("/claude_code_gateway/managed/settings", headers={"If-None-Match": etag}) - assert not_modified.status_code == 304 - assert not_modified.headers["ETag"] == etag + not_modified = client.get( + "/claude_code_gateway/managed/settings", headers={"If-None-Match": f'"{checksum}"'} + ) + assert not_modified.status_code == 304 + assert not_modified.headers["ETag"] == f'"{checksum}"' + + stale = client.get("/claude_code_gateway/managed/settings", headers={"If-None-Match": '"sha256:stale"'}) + assert stale.status_code == 200 + assert stale.json()["checksum"] == checksum + + +def test_managed_settings_checksum_tracks_policy_content(): + with _gateway_env(managed_settings={"env": {"FOO": "bar"}}) as (client, _): + first = client.get("/claude_code_gateway/managed/settings").json()["checksum"] + with _gateway_env(managed_settings={"env": {"FOO": "baz"}}) as (client, _): + second = client.get("/claude_code_gateway/managed/settings").json()["checksum"] + assert first != second def test_managed_settings_404_when_gateway_disabled(): From c6c8aed3f8594f89a86b91df553b84f6aab2fb20 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 18 Sep 2026 16:22:04 -0700 Subject: [PATCH 079/144] fix(proxy): drop only the daily spend batch whose failure cannot be re-sent, requeue the unsent ones --- litellm/proxy/db/db_spend_update_writer.py | 30 ++++++----- .../proxy/db/test_db_spend_update_writer.py | 54 +++++++++++++++++++ 2 files changed, 71 insertions(+), 13 deletions(-) diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index e9967fb0d67..37eac8604bd 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -1332,15 +1332,7 @@ class DBSpendUpdateWriter: daily_spend_transactions=cast(dict[str, _DailySpendTransactionT], transactions), ) except Exception as e: # noqa: BLE001 # whatever failed here, the other tables must still flush - if not _daily_spend_commit_failure_is_requeue_safe(e): - spend_log_error( - "Spend tracking - dropped %d daily %s spend rows: the failed commit may have applied " - "or the database refused the data, so re-sending it is not safe. Error: %s", - len(transactions), - entity_type, - str(e), - exc=e, - ) + if not transactions: return spend_log_error( "Spend tracking - failed to commit daily %s spend updates. " @@ -2050,13 +2042,25 @@ class DBSpendUpdateWriter: sql, params = build_bulk_upsert(table=table, batch=merged_batch) await prisma_client.db.execute_raw(sql, *params) except Exception as batch_error: - # Log detailed error information for debugging batch upsert failures - # This helps diagnose issues like unique constraint violations + if _daily_spend_commit_failure_is_requeue_safe(batch_error): + spend_log_error( + "Daily %s spend batch upsert failed. Table: %s, Rows: %d, Error: %s", + entity_type, + table.name, + len(transactions_to_process), + str(batch_error), + exc=batch_error, + ) + raise + for key in transactions_to_process: + daily_spend_transactions.pop(key, None) spend_log_error( - "Daily %s spend batch upsert failed. Table: %s, Rows: %d, Error: %s", + "Spend tracking - dropped %d daily %s spend rows: the failed statement may have " + "applied or the database refused the data, so re-sending it is not safe. " + "Table: %s, Error: %s", + len(transactions_to_process), entity_type, table.name, - len(transactions_to_process), str(batch_error), exc=batch_error, ) diff --git a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py index 155bca656d5..20f1fa9d363 100644 --- a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py +++ b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py @@ -1624,6 +1624,33 @@ async def test_update_daily_spend_keeps_failed_transactions_for_retry(): assert daily_spend_transactions == expected +@pytest.mark.asyncio +async def test_update_daily_spend_drops_the_batch_whose_failure_cannot_be_resent(): + """A reply lost after the statement was sent may already have applied, so the batch is + taken out of the caller's dict before the error propagates: whichever requeue the caller + runs afterwards, the Redis restore included, cannot send it a second time.""" + + def lose_the_reply() -> int: + raise httpx.ReadTimeout("no reply") + + prisma_client = _RecordingPrisma(execute_raw=lose_the_reply) + daily_spend_transactions = {"user-key": _daily_txn(user_id="user-1")} + proxy_logging_obj = MagicMock() + proxy_logging_obj.failure_handler = AsyncMock() + + with pytest.raises(httpx.ReadTimeout): + await DBSpendUpdateWriter._update_daily_spend( + n_retry_times=0, + prisma_client=prisma_client, + proxy_logging_obj=proxy_logging_obj, + daily_spend_transactions=daily_spend_transactions, + entity_type="user", + entity_id_field="user_id", + ) + + assert daily_spend_transactions == {} + + @pytest.mark.asyncio async def test_commit_key_spend_updates_includes_last_active(): """ @@ -2875,6 +2902,33 @@ async def test_failed_daily_spend_commit_is_requeued_only_when_the_rows_are_prov assert db_writer.daily_spend_update_queue.update_queue.empty() +@pytest.mark.asyncio +async def test_failed_daily_spend_commit_drops_only_the_batch_that_was_sent(): + """A tick holding more than one batch of 100 rows sends them one statement at a time, and + a reply lost on one statement says nothing about the batches after it: only the batch that + was on the wire is dropped, the ones never sent go back on the queue and land next tick.""" + db_writer = DBSpendUpdateWriter() + await db_writer.daily_spend_update_queue.add_update( + {f"user-{i:03d}": _daily_txn(user_id=f"user-{i:03d}") for i in range(150)} + ) + db = _DailySpendFakeDB(failing_table="LiteLLM_DailyUserSpend", failure=httpx.ReadTimeout("no reply")) + db_writer._flush_tool_discovery_queue = AsyncMock() + proxy_logging_obj = MagicMock() + proxy_logging_obj.failure_handler = AsyncMock() + + await db_writer._commit_spend_updates_to_db_without_redis_buffer( + prisma_client=_WindowSpendFakePrisma(db), n_retry_times=0, proxy_logging_obj=proxy_logging_obj + ) + db.failing_table = None + await db_writer._commit_spend_updates_to_db_without_redis_buffer( + prisma_client=_WindowSpendFakePrisma(db), n_retry_times=0, proxy_logging_obj=proxy_logging_obj + ) + + (upsert,) = _daily_upserts(db, "LiteLLM_DailyUserSpend") + assert _row_values(upsert, "user_id") == [f"user-{i:03d}" for i in range(100, 150)] + assert db_writer.daily_spend_update_queue.update_queue.empty() + + @pytest.mark.asyncio async def test_failed_daily_spend_commit_requeues_the_rows_and_flushes_the_other_tables(): """With the Redis buffer off, a daily batch that failed to commit was discarded along From abb9618971e80649d3db5e1ee85eaec82384ede4 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Fri, 18 Sep 2026 23:28:11 +0000 Subject: [PATCH 080/144] feat(rust): add litellm-http client pool and inject it into the OCR route Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm-rust/Cargo.lock | 13 + litellm-rust/Cargo.toml | 1 + litellm-rust/crates/core/Cargo.toml | 1 + litellm-rust/crates/core/src/ocr/client.rs | 9 +- litellm-rust/crates/core/tests/ocr.rs | 46 ++- litellm-rust/crates/http/Cargo.toml | 14 + litellm-rust/crates/http/src/config.rs | 327 ++++++++++++++++++ litellm-rust/crates/http/src/lib.rs | 11 + litellm-rust/crates/http/src/pool.rs | 172 +++++++++ litellm-rust/crates/http/src/settings.rs | 171 +++++++++ litellm-rust/crates/llms/Cargo.toml | 1 + .../llms/src/base_llm/ocr/transformation.rs | 1 - .../llms/src/custom_httpx/llm_http_handler.rs | 46 +-- .../crates/llms/src/custom_httpx/media.rs | 38 +- .../crates/llms/src/custom_httpx/transport.rs | 6 + litellm-rust/crates/python-bridge/Cargo.toml | 1 + litellm-rust/crates/python-bridge/src/http.rs | 234 +++++++++++++ litellm-rust/crates/python-bridge/src/lib.rs | 1 + .../python-bridge/src/routes/ocr/mod.rs | 6 +- 19 files changed, 1037 insertions(+), 62 deletions(-) create mode 100644 litellm-rust/crates/http/Cargo.toml create mode 100644 litellm-rust/crates/http/src/config.rs create mode 100644 litellm-rust/crates/http/src/lib.rs create mode 100644 litellm-rust/crates/http/src/pool.rs create mode 100644 litellm-rust/crates/http/src/settings.rs create mode 100644 litellm-rust/crates/python-bridge/src/http.rs diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index c359ca19986..ffbdc80efd3 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -2058,6 +2058,7 @@ dependencies = [ "litellm-auth-aws", "litellm-callbacks", "litellm-core-utils", + "litellm-http", "litellm-llms", "litellm-types", "mime_guess", @@ -2125,6 +2126,16 @@ dependencies = [ "tokio", ] +[[package]] +name = "litellm-http" +version = "0.1.0" +dependencies = [ + "reqwest 0.12.28", + "rstest", + "thiserror 2.0.19", + "tokio", +] + [[package]] name = "litellm-llms" version = "0.1.0" @@ -2142,6 +2153,7 @@ dependencies = [ "litellm-callbacks", "litellm-core-utils", "litellm-framing", + "litellm-http", "litellm-types", "reqwest 0.12.28", "rstest", @@ -2166,6 +2178,7 @@ dependencies = [ "litellm-callbacks-legacy", "litellm-core", "litellm-host-python", + "litellm-http", "litellm-llms", "litellm-token-counter", "litellm-types", diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index ffdbf64bb49..eeb473cd2e9 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -17,6 +17,7 @@ litellm-auth = { path = "crates/auth" } litellm-auth-aws = { path = "crates/auth-aws" } litellm-auth-azure = { path = "crates/auth-azure" } litellm-auth-gcp = { path = "crates/auth-gcp" } +litellm-http = { path = "crates/http" } litellm-llms = { path = "crates/llms" } litellm-types = { path = "crates/types" } litellm-core-utils = { path = "crates/core-utils" } diff --git a/litellm-rust/crates/core/Cargo.toml b/litellm-rust/crates/core/Cargo.toml index db6cfc4b340..047188c74f9 100644 --- a/litellm-rust/crates/core/Cargo.toml +++ b/litellm-rust/crates/core/Cargo.toml @@ -15,6 +15,7 @@ futures-util.workspace = true base64.workspace = true litellm-auth.workspace = true litellm-auth-aws.workspace = true +litellm-http.workspace = true litellm-llms.workspace = true moka.workspace = true mime_guess = "2.0.5" diff --git a/litellm-rust/crates/core/src/ocr/client.rs b/litellm-rust/crates/core/src/ocr/client.rs index 03782d91f24..21c81505cff 100644 --- a/litellm-rust/crates/core/src/ocr/client.rs +++ b/litellm-rust/crates/core/src/ocr/client.rs @@ -1,3 +1,4 @@ +use litellm_http::{HttpClientConfig, HttpClientPool}; use litellm_llms::{ base_llm::ocr::{error::Error, transformation::LiteLLMOcrResponse}, custom_httpx::llm_http_handler::OcrClient, @@ -15,6 +16,10 @@ pub async fn perform( litellm_callbacks::run::run(ocr_machine(client.clone()), &LocalOcrHost::new(request)).await } -pub async fn ocr(request: LiteLLMOcrRequest) -> Result { - perform(&OcrClient::shared()?, request).await +pub async fn ocr( + pool: &HttpClientPool, + config: &HttpClientConfig, + request: LiteLLMOcrRequest, +) -> Result { + perform(&OcrClient::new(pool, config)?, request).await } diff --git a/litellm-rust/crates/core/tests/ocr.rs b/litellm-rust/crates/core/tests/ocr.rs index 1f591d74d5d..d59c6179aa5 100644 --- a/litellm-rust/crates/core/tests/ocr.rs +++ b/litellm-rust/crates/core/tests/ocr.rs @@ -5,6 +5,7 @@ use litellm_callbacks::{ host::{Host, HostOp, HostResult}, machine::{HostFailure, Machine, MachineStep}, }; +use litellm_http::{HttpClientConfig, HttpClientPool, HttpSettings, Verify}; use litellm_llms::{ base_llm::ocr::{ error::Error as OcrError, @@ -171,25 +172,44 @@ async fn facade_retains_native_response_when_requested() { } #[tokio::test] -async fn facade_uses_the_injected_http_client() { +async fn facade_uses_the_injected_http_pool_configuration() { let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; - let mut default_headers = reqwest::header::HeaderMap::new(); - default_headers.insert( - "x-transport-owner", - reqwest::header::HeaderValue::from_static("host"), - ); - let provider_http = reqwest::Client::builder() - .default_headers(default_headers) - .build() - .unwrap(); - crate::ocr::client::perform( - &OcrClient::new(provider_http).unwrap(), + let settings = HttpSettings { + user_agent: Some("host-owned/1".into()), + ..HttpSettings::default() + }; + let config = HttpClientConfig::resolve(&settings, None).unwrap(); + crate::ocr::client::ocr( + &HttpClientPool::new(), + &config, wire_request("mistral/model", &base, json!({})), ) .await .unwrap(); server.await.unwrap(); - assert!(seen.lock().unwrap()[0].contains("x-transport-owner: host")); + assert!(seen.lock().unwrap()[0].contains("user-agent: host-owned/1")); +} + +#[tokio::test] +async fn unbuildable_http_configuration_fails_before_dispatch() { + let (base, _seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; + let config = HttpClientConfig { + verify: Verify::CaBundle(std::env::temp_dir().join("litellm-ocr-missing-bundle.pem")), + ..HttpClientConfig::resolve(&HttpSettings::default(), None).unwrap() + }; + let error = crate::ocr::client::ocr( + &HttpClientPool::new(), + &config, + wire_request("mistral/model", &base, json!({})), + ) + .await + .unwrap_err(); + server.abort(); + assert!(matches!( + error, + OcrError::Transport(litellm_llms::custom_httpx::transport::Error::Connect(_)) + )); + assert!(error.to_string().contains("litellm-ocr-missing-bundle.pem")); } fn event_name(event: &CallEvent) -> &'static str { diff --git a/litellm-rust/crates/http/Cargo.toml b/litellm-rust/crates/http/Cargo.toml new file mode 100644 index 00000000000..48ea4e66cef --- /dev/null +++ b/litellm-rust/crates/http/Cargo.toml @@ -0,0 +1,14 @@ +[package] +name = "litellm-http" +version = "0.1.0" +edition.workspace = true +license.workspace = true +repository.workspace = true + +[dependencies] +reqwest.workspace = true +thiserror.workspace = true + +[dev-dependencies] +rstest.workspace = true +tokio.workspace = true diff --git a/litellm-rust/crates/http/src/config.rs b/litellm-rust/crates/http/src/config.rs new file mode 100644 index 00000000000..86f1a9b43b6 --- /dev/null +++ b/litellm-rust/crates/http/src/config.rs @@ -0,0 +1,327 @@ +use std::{ + net::{IpAddr, Ipv4Addr}, + path::{Path, PathBuf}, + time::Duration, +}; + +use crate::settings::{HttpSettings, SslVerify}; + +#[derive(Clone, Debug, thiserror::Error, PartialEq, Eq)] +pub enum Error { + #[error("{setting} cannot be expressed with rustls: {reason}")] + Unsupported { + setting: &'static str, + reason: String, + }, + #[error("could not read {}: {message}", path.display())] + Read { path: PathBuf, message: String }, + #[error("{} is not a PEM file: {message}", path.display())] + InvalidPem { path: PathBuf, message: String }, + #[error("could not build the HTTP client: {0}")] + Client(String), +} + +impl From for Error { + fn from(error: reqwest::Error) -> Self { + Self::Client(error.without_url().to_string()) + } +} + +#[derive(Clone, Debug, PartialEq, Eq, Hash)] +pub enum Verify { + Disabled, + CaBundle(PathBuf), + BuiltInRoots, +} + +/// One fully resolved client configuration. Every field is a plain value so the pool can +/// key cached clients on it. +#[derive(Clone, Debug, PartialEq, Eq, Hash)] +pub struct HttpClientConfig { + pub verify: Verify, + pub client_certificate: Option, + pub force_ipv4: bool, + pub http2: bool, + pub user_agent: Option, + pub trust_proxy_env: bool, + pub connect_timeout: Duration, + pub request_timeout: Option, +} + +impl HttpClientConfig { + /// Port of `get_ssl_verify` + `get_ssl_configuration`: the per-call value wins, then the + /// configured (environment-overlaid) `ssl_verify`, then `SSL_CERT_FILE`, then the built-in + /// roots. Settings rustls has no equivalent for are an error instead of a silent no-op. + pub fn resolve( + settings: &HttpSettings, + per_call_ssl_verify: Option<&SslVerify>, + ) -> Result { + if let Some(level) = &settings.ssl_security_level { + return Err(Error::Unsupported { + setting: "ssl_security_level", + reason: format!("OpenSSL cipher string {level:?} has no rustls equivalent"), + }); + } + if let Some(curve) = &settings.ssl_ecdh_curve { + return Err(Error::Unsupported { + setting: "ssl_ecdh_curve", + reason: format!("key exchange group {curve:?} is fixed by the rustls provider"), + }); + } + let verify = match per_call_ssl_verify.or(settings.ssl_verify.as_ref()) { + Some(SslVerify::Disabled) => Verify::Disabled, + Some(SslVerify::CaBundle(path)) => Verify::CaBundle(path.clone()), + Some(SslVerify::Enabled) | None => settings + .ssl_cert_file + .clone() + .map_or(Verify::BuiltInRoots, Verify::CaBundle), + }; + Ok(Self { + verify, + client_certificate: settings.ssl_certificate.clone(), + force_ipv4: settings.force_ipv4, + http2: settings.http2, + user_agent: settings.user_agent.clone(), + trust_proxy_env: settings.trust_proxy_env, + connect_timeout: settings.connect_timeout, + request_timeout: settings.request_timeout, + }) + } + + /// A builder carrying every shared setting; variants add their own policy on top. + pub fn client_builder(&self) -> Result { + let base = reqwest::Client::builder().connect_timeout(self.connect_timeout); + let with_roots = match &self.verify { + Verify::Disabled => base.danger_accept_invalid_certs(true), + Verify::BuiltInRoots => base, + Verify::CaBundle(path) => { + let pem = read(path)?; + let certificates = + reqwest::Certificate::from_pem_bundle(&pem).map_err(|error| { + Error::InvalidPem { + path: path.clone(), + message: error.without_url().to_string(), + } + })?; + if certificates.is_empty() { + return Err(Error::InvalidPem { + path: path.clone(), + message: "no certificates found".into(), + }); + } + certificates.into_iter().fold( + base.tls_built_in_root_certs(false), + |builder, certificate| builder.add_root_certificate(certificate), + ) + } + }; + let with_identity = match &self.client_certificate { + None => with_roots, + Some(path) => { + let identity = reqwest::Identity::from_pem(&read(path)?).map_err(|error| { + Error::InvalidPem { + path: path.clone(), + message: error.without_url().to_string(), + } + })?; + with_roots.identity(identity) + } + }; + let with_address = if self.force_ipv4 { + with_identity.local_address(IpAddr::V4(Ipv4Addr::UNSPECIFIED)) + } else { + with_identity + }; + let with_protocol = if self.http2 { + with_address + } else { + with_address.http1_only() + }; + let with_agent = match &self.user_agent { + Some(agent) => with_protocol.user_agent(agent), + None => with_protocol, + }; + let with_proxy = if self.trust_proxy_env { + with_agent + } else { + with_agent.no_proxy() + }; + Ok(match self.request_timeout { + Some(timeout) => with_proxy.timeout(timeout), + None => with_proxy, + }) + } +} + +fn read(path: &Path) -> Result, Error> { + std::fs::read(path).map_err(|error| Error::Read { + path: path.to_path_buf(), + message: error.to_string(), + }) +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + + use super::*; + + fn no_env(_: &str) -> Option { + None + } + + fn settings(ssl_verify: Option, ssl_cert_file: Option<&str>) -> HttpSettings { + HttpSettings { + ssl_verify, + ssl_cert_file: ssl_cert_file.map(PathBuf::from), + ..HttpSettings::default() + } + .with_environment(&no_env) + } + + #[rstest] + #[case::default(settings(None, None), None, Verify::BuiltInRoots)] + #[case::setting_disables( + settings(Some(SslVerify::Disabled), Some("/env/roots.pem")), + None, + Verify::Disabled + )] + #[case::setting_bundle( + settings(Some(SslVerify::CaBundle("/configured.pem".into())), Some("/env/roots.pem")), + None, + Verify::CaBundle("/configured.pem".into()) + )] + #[case::enabled_uses_cert_file( + settings(Some(SslVerify::Enabled), Some("/env/roots.pem")), + None, + Verify::CaBundle("/env/roots.pem".into()) + )] + #[case::unset_uses_cert_file(settings(None, Some("/env/roots.pem")), None, Verify::CaBundle("/env/roots.pem".into()))] + #[case::per_call_beats_setting( + settings(Some(SslVerify::Disabled), None), + Some(SslVerify::Enabled), + Verify::BuiltInRoots + )] + #[case::per_call_disables( + settings(Some(SslVerify::CaBundle("/configured.pem".into())), Some("/env/roots.pem")), + Some(SslVerify::Disabled), + Verify::Disabled + )] + #[case::per_call_bundle( + settings(None, Some("/env/roots.pem")), + Some(SslVerify::CaBundle("/call.pem".into())), + Verify::CaBundle("/call.pem".into()) + )] + #[case::per_call_enabled_still_honours_cert_file( + settings(Some(SslVerify::Disabled), Some("/env/roots.pem")), + Some(SslVerify::Enabled), + Verify::CaBundle("/env/roots.pem".into()) + )] + fn verify_follows_per_call_then_setting_then_cert_file( + #[case] settings: HttpSettings, + #[case] per_call: Option, + #[case] expected: Verify, + ) { + let config = HttpClientConfig::resolve(&settings, per_call.as_ref()).unwrap(); + assert_eq!(config.verify, expected); + } + + #[test] + fn ssl_verify_environment_variable_beats_the_configured_setting() { + let settings = HttpSettings { + ssl_verify: Some(SslVerify::Disabled), + ..HttpSettings::default() + } + .with_environment(&|name: &str| (name == "SSL_VERIFY").then(|| "true".to_string())); + let config = HttpClientConfig::resolve(&settings, None).unwrap(); + assert_eq!(config.verify, Verify::BuiltInRoots); + } + + #[test] + fn cipher_strings_are_rejected_rather_than_ignored() { + let settings = HttpSettings { + ssl_security_level: Some("DEFAULT@SECLEVEL=1".into()), + ..HttpSettings::default() + }; + assert!(matches!( + HttpClientConfig::resolve(&settings, None), + Err(Error::Unsupported { + setting: "ssl_security_level", + .. + }) + )); + } + + #[test] + fn ecdh_curves_are_rejected_rather_than_ignored() { + let settings = HttpSettings { + ssl_ecdh_curve: Some("X25519".into()), + ..HttpSettings::default() + }; + assert!(matches!( + HttpClientConfig::resolve(&settings, None), + Err(Error::Unsupported { + setting: "ssl_ecdh_curve", + .. + }) + )); + } + + #[test] + fn connection_settings_carry_over_unchanged() { + let settings = HttpSettings { + ssl_certificate: Some("/client.pem".into()), + force_ipv4: true, + http2: true, + user_agent: Some("litellm/1.0".into()), + trust_proxy_env: true, + connect_timeout: Duration::from_secs(7), + request_timeout: Some(Duration::from_secs(70)), + ..HttpSettings::default() + }; + let config = HttpClientConfig::resolve(&settings, None).unwrap(); + assert_eq!( + config, + HttpClientConfig { + verify: Verify::BuiltInRoots, + client_certificate: Some("/client.pem".into()), + force_ipv4: true, + http2: true, + user_agent: Some("litellm/1.0".into()), + trust_proxy_env: true, + connect_timeout: Duration::from_secs(7), + request_timeout: Some(Duration::from_secs(70)), + } + ); + } + + #[test] + fn missing_ca_bundle_is_a_read_error() { + let path = std::env::temp_dir().join("litellm-http-missing-bundle.pem"); + let config = HttpClientConfig { + verify: Verify::CaBundle(path.clone()), + ..HttpClientConfig::resolve(&HttpSettings::default(), None).unwrap() + }; + assert!(matches!( + config.client_builder(), + Err(Error::Read { path: reported, .. }) if reported == path + )); + } + + #[test] + fn non_pem_ca_bundle_is_an_invalid_pem_error() { + let path = + std::env::temp_dir().join(format!("litellm-http-not-pem-{}.pem", std::process::id())); + std::fs::write(&path, b"not a certificate").unwrap(); + let config = HttpClientConfig { + verify: Verify::CaBundle(path.clone()), + ..HttpClientConfig::resolve(&HttpSettings::default(), None).unwrap() + }; + let result = config.client_builder().map(drop); + std::fs::remove_file(&path).unwrap(); + assert!(matches!( + result, + Err(Error::InvalidPem { path: reported, .. }) if reported == path + )); + } +} diff --git a/litellm-rust/crates/http/src/lib.rs b/litellm-rust/crates/http/src/lib.rs new file mode 100644 index 00000000000..d62dd768fe1 --- /dev/null +++ b/litellm-rust/crates/http/src/lib.rs @@ -0,0 +1,11 @@ +//! Rust counterpart of `litellm/llms/custom_httpx/http_handler.py`: the plain HTTP settings +//! LiteLLM exposes, their resolution into one typed client configuration, and a pool that +//! caches `reqwest::Client`s per resolved configuration. + +mod config; +mod pool; +mod settings; + +pub use config::{Error, HttpClientConfig, Verify}; +pub use pool::{ClientVariant, HttpClientPool}; +pub use settings::{HttpSettings, SslVerify}; diff --git a/litellm-rust/crates/http/src/pool.rs b/litellm-rust/crates/http/src/pool.rs new file mode 100644 index 00000000000..065097556e1 --- /dev/null +++ b/litellm-rust/crates/http/src/pool.rs @@ -0,0 +1,172 @@ +use std::{ + collections::HashMap, + sync::{Mutex, PoisonError}, +}; + +use crate::config::{Error, HttpClientConfig}; + +/// The client shapes routes need; each is the shared base plus one policy. +#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)] +pub enum ClientVariant { + Provider, + NoRedirect, + /// Media downloads: no redirects (the fetcher validates each hop) and never a proxy. + Media, +} + +impl ClientVariant { + fn apply(self, builder: reqwest::ClientBuilder) -> reqwest::ClientBuilder { + match self { + Self::Provider => builder, + Self::NoRedirect => builder.redirect(reqwest::redirect::Policy::none()), + Self::Media => builder + .redirect(reqwest::redirect::Policy::none()) + .no_proxy(), + } + } +} + +/// Counterpart of `get_async_httpx_client`: one `reqwest::Client` per resolved configuration +/// and variant, built on first use and shared afterwards. +#[derive(Default)] +pub struct HttpClientPool { + clients: Mutex>, +} + +impl HttpClientPool { + pub fn new() -> Self { + Self::default() + } + + pub fn client( + &self, + config: &HttpClientConfig, + variant: ClientVariant, + ) -> Result { + self.client_with(config, variant, |builder| builder) + } + + /// Like [`Self::client`], with a caller hook for builder options that are not plain values + /// (a DNS resolver, for example). The hook only runs when the client is first built. + pub fn client_with( + &self, + config: &HttpClientConfig, + variant: ClientVariant, + customize: impl FnOnce(reqwest::ClientBuilder) -> reqwest::ClientBuilder, + ) -> Result { + let key = (config.clone(), variant); + let mut clients = self.clients.lock().unwrap_or_else(PoisonError::into_inner); + if let Some(client) = clients.get(&key) { + return Ok(client.clone()); + } + let client = customize(variant.apply(config.client_builder()?)).build()?; + clients.insert(key, client.clone()); + Ok(client) + } +} + +#[cfg(test)] +mod tests { + use std::{cell::Cell, time::Duration}; + + use tokio::io::{AsyncReadExt, AsyncWriteExt}; + + use super::*; + use crate::{HttpSettings, Verify}; + + fn config(user_agent: &str) -> HttpClientConfig { + HttpClientConfig { + user_agent: Some(user_agent.into()), + ..HttpClientConfig::resolve(&HttpSettings::default(), None).unwrap() + } + } + + #[test] + fn clients_are_built_once_per_config_and_variant() { + let pool = HttpClientPool::new(); + let builds = Cell::new(0); + let build = |config: &HttpClientConfig, variant| { + pool.client_with(config, variant, |builder| { + builds.set(builds.get() + 1); + builder + }) + .unwrap() + }; + build(&config("a"), ClientVariant::Provider); + build(&config("a"), ClientVariant::Provider); + assert_eq!(builds.get(), 1); + build(&config("a"), ClientVariant::NoRedirect); + assert_eq!(builds.get(), 2); + build(&config("b"), ClientVariant::Provider); + assert_eq!(builds.get(), 3); + build(&config("b"), ClientVariant::Provider); + build(&config("a"), ClientVariant::NoRedirect); + assert_eq!(builds.get(), 3); + } + + #[test] + fn build_failures_are_not_cached() { + let pool = HttpClientPool::new(); + let missing = HttpClientConfig { + verify: Verify::CaBundle(std::env::temp_dir().join("litellm-http-absent.pem")), + ..config("a") + }; + assert!(pool.client(&missing, ClientVariant::Provider).is_err()); + assert!(pool.client(&missing, ClientVariant::Provider).is_err()); + assert!(pool.client(&config("a"), ClientVariant::Provider).is_ok()); + } + + async fn serve_once(status_line: &'static str) -> (String, tokio::task::JoinHandle) { + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let base = format!("http://{}", listener.local_addr().unwrap()); + let server = tokio::spawn(async move { + let (mut socket, _) = listener.accept().await.unwrap(); + let mut request = vec![0u8; 4096]; + let read = socket.read(&mut request).await.unwrap(); + socket + .write_all( + format!("{status_line}\r\nLocation: /elsewhere\r\nContent-Length: 0\r\nConnection: close\r\n\r\n") + .as_bytes(), + ) + .await + .unwrap(); + String::from_utf8_lossy(&request[..read]).into_owned() + }); + (base, server) + } + + #[tokio::test] + async fn provider_client_sends_the_configured_user_agent_over_http1() { + let (base, server) = serve_once("HTTP/1.1 204 No Content").await; + let config = HttpClientConfig { + connect_timeout: Duration::from_secs(2), + ..config("litellm-test/9") + }; + let response = HttpClientPool::new() + .client(&config, ClientVariant::Provider) + .unwrap() + .get(&base) + .send() + .await + .unwrap(); + assert_eq!(response.status(), 204); + assert_eq!(response.version(), reqwest::Version::HTTP_11); + let request = server.await.unwrap(); + assert!(request.contains("user-agent: litellm-test/9"), "{request}"); + } + + #[tokio::test] + async fn no_redirect_variant_returns_the_redirect_instead_of_following_it() { + let (base, server) = serve_once("HTTP/1.1 302 Found").await; + let response = HttpClientPool::new() + .client(&config("a"), ClientVariant::NoRedirect) + .unwrap() + .get(&base) + .send() + .await + .unwrap(); + assert_eq!(response.status(), 302); + assert_eq!(response.headers()["location"], "/elsewhere"); + server.await.unwrap(); + } +} diff --git a/litellm-rust/crates/http/src/settings.rs b/litellm-rust/crates/http/src/settings.rs new file mode 100644 index 00000000000..4d936e89046 --- /dev/null +++ b/litellm-rust/crates/http/src/settings.rs @@ -0,0 +1,171 @@ +use std::{path::PathBuf, time::Duration}; + +/// `litellm.ssl_verify` / `SSL_VERIFY`: a bool or a CA bundle path. +#[derive(Clone, Debug, PartialEq, Eq, Hash)] +pub enum SslVerify { + Enabled, + Disabled, + CaBundle(PathBuf), +} + +impl SslVerify { + pub fn parse(value: &str) -> Self { + match value.trim().to_ascii_lowercase().as_str() { + "true" => Self::Enabled, + "false" => Self::Disabled, + _ => Self::CaBundle(PathBuf::from(value)), + } + } +} + +/// The plain inputs `http_handler.py` reads from `litellm.*` globals and the environment. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct HttpSettings { + pub ssl_verify: Option, + pub ssl_cert_file: Option, + pub ssl_certificate: Option, + pub ssl_security_level: Option, + pub ssl_ecdh_curve: Option, + pub force_ipv4: bool, + pub http2: bool, + pub user_agent: Option, + pub trust_proxy_env: bool, + pub connect_timeout: Duration, + pub request_timeout: Option, +} + +impl Default for HttpSettings { + fn default() -> Self { + Self { + ssl_verify: None, + ssl_cert_file: None, + ssl_certificate: None, + ssl_security_level: None, + ssl_ecdh_curve: None, + force_ipv4: false, + http2: false, + user_agent: None, + trust_proxy_env: false, + connect_timeout: Duration::from_secs(5), + request_timeout: None, + } + } +} + +impl HttpSettings { + /// Overlay the environment variables `http_handler.py` consults, with the same precedence: + /// `SSL_VERIFY`, `SSL_CERTIFICATE`, `SSL_SECURITY_LEVEL`, `SSL_ECDH_CURVE` and + /// `LITELLM_USER_AGENT` win over the configured value; `SSL_CERT_FILE` only applies when + /// verification is on without an explicit bundle; `LITELLM_HTTP2` and `AIOHTTP_TRUST_ENV` + /// can only turn their switch on. + pub fn with_environment(self, env: &(dyn Fn(&str) -> Option + Sync)) -> Self { + let enabled = + |name: &str| env(name).is_some_and(|value| value.trim().eq_ignore_ascii_case("true")); + Self { + ssl_verify: env("SSL_VERIFY") + .map(|value| SslVerify::parse(&value)) + .or(self.ssl_verify), + ssl_cert_file: env("SSL_CERT_FILE") + .map(PathBuf::from) + .or(self.ssl_cert_file), + ssl_certificate: env("SSL_CERTIFICATE") + .map(PathBuf::from) + .or(self.ssl_certificate), + ssl_security_level: env("SSL_SECURITY_LEVEL").or(self.ssl_security_level), + ssl_ecdh_curve: env("SSL_ECDH_CURVE").or(self.ssl_ecdh_curve), + http2: self.http2 || enabled("LITELLM_HTTP2"), + user_agent: env("LITELLM_USER_AGENT").or(self.user_agent), + trust_proxy_env: self.trust_proxy_env || enabled("AIOHTTP_TRUST_ENV"), + ..self + } + } +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + + use super::*; + + fn no_env(_: &str) -> Option { + None + } + + fn env_of( + values: &'static [(&'static str, &'static str)], + ) -> impl Fn(&str) -> Option + Sync { + move |name| { + values + .iter() + .find(|(key, _)| *key == name) + .map(|(_, value)| value.to_string()) + } + } + + #[rstest] + #[case("true", SslVerify::Enabled)] + #[case(" True ", SslVerify::Enabled)] + #[case("FALSE", SslVerify::Disabled)] + #[case("/etc/ssl/bundle.pem", SslVerify::CaBundle("/etc/ssl/bundle.pem".into()))] + fn ssl_verify_parses_bools_and_treats_anything_else_as_a_bundle_path( + #[case] value: &str, + #[case] expected: SslVerify, + ) { + assert_eq!(SslVerify::parse(value), expected); + } + + #[test] + fn environment_overrides_configured_ssl_values() { + let settings = HttpSettings { + ssl_verify: Some(SslVerify::Enabled), + ssl_certificate: Some("/configured/client.pem".into()), + ssl_security_level: Some("configured".into()), + user_agent: Some("configured/1".into()), + ..HttpSettings::default() + } + .with_environment(&env_of(&[ + ("SSL_VERIFY", "false"), + ("SSL_CERT_FILE", "/env/roots.pem"), + ("SSL_CERTIFICATE", "/env/client.pem"), + ("SSL_SECURITY_LEVEL", "DEFAULT@SECLEVEL=1"), + ("SSL_ECDH_CURVE", "X25519"), + ("LITELLM_USER_AGENT", "env/2"), + ])); + assert_eq!(settings.ssl_verify, Some(SslVerify::Disabled)); + assert_eq!(settings.ssl_cert_file, Some("/env/roots.pem".into())); + assert_eq!(settings.ssl_certificate, Some("/env/client.pem".into())); + assert_eq!( + settings.ssl_security_level.as_deref(), + Some("DEFAULT@SECLEVEL=1") + ); + assert_eq!(settings.ssl_ecdh_curve.as_deref(), Some("X25519")); + assert_eq!(settings.user_agent.as_deref(), Some("env/2")); + } + + #[test] + fn missing_environment_keeps_configured_values() { + let configured = HttpSettings { + ssl_verify: Some(SslVerify::CaBundle("/configured/roots.pem".into())), + http2: true, + trust_proxy_env: true, + user_agent: Some("configured/1".into()), + ..HttpSettings::default() + }; + assert_eq!(configured.clone().with_environment(&no_env), configured); + } + + #[rstest] + #[case("true", true)] + #[case("True", true)] + #[case("false", false)] + #[case("1", false)] + fn boolean_switches_only_turn_on_for_true(#[case] value: &'static str, #[case] expected: bool) { + let env = move |name: &str| match name { + "LITELLM_HTTP2" | "AIOHTTP_TRUST_ENV" => Some(value.to_string()), + _ => None, + }; + let settings = HttpSettings::default().with_environment(&env); + assert_eq!(settings.http2, expected); + assert_eq!(settings.trust_proxy_env, expected); + } +} diff --git a/litellm-rust/crates/llms/Cargo.toml b/litellm-rust/crates/llms/Cargo.toml index 4ca6c7cb2a5..ca76aaf380e 100644 --- a/litellm-rust/crates/llms/Cargo.toml +++ b/litellm-rust/crates/llms/Cargo.toml @@ -17,6 +17,7 @@ litellm-auth-azure.workspace = true litellm-auth-gcp.workspace = true litellm-callbacks.workspace = true litellm-framing.workspace = true +litellm-http.workspace = true base64.workspace = true bytes.workspace = true data-url = "0.3.2" diff --git a/litellm-rust/crates/llms/src/base_llm/ocr/transformation.rs b/litellm-rust/crates/llms/src/base_llm/ocr/transformation.rs index e6fe5d9556d..f20ec726f61 100644 --- a/litellm-rust/crates/llms/src/base_llm/ocr/transformation.rs +++ b/litellm-rust/crates/llms/src/base_llm/ocr/transformation.rs @@ -21,7 +21,6 @@ use crate::{ pub const OCR_RESPONSE_MAX_BYTES: usize = 64 * 1024 * 1024; pub const OCR_HTTP_TIMEOUT_SECS: u64 = 600; -pub const OCR_CONNECT_TIMEOUT_SECS: u64 = 10; pub const OCR_INLINE_MAX_BYTES: usize = 50 * 1024 * 1024; pub const OCR_DOWNLOAD_MAX_BYTES: u64 = 50 * 1024 * 1024; pub const OCR_MAX_FETCH_REDIRECTS: usize = 10; diff --git a/litellm-rust/crates/llms/src/custom_httpx/llm_http_handler.rs b/litellm-rust/crates/llms/src/custom_httpx/llm_http_handler.rs index e93ddee3c50..38ddec08da9 100644 --- a/litellm-rust/crates/llms/src/custom_httpx/llm_http_handler.rs +++ b/litellm-rust/crates/llms/src/custom_httpx/llm_http_handler.rs @@ -1,9 +1,8 @@ -use std::{sync::OnceLock, time::Duration}; - use bytes::{Bytes, BytesMut}; use futures_util::future::BoxFuture; use litellm_auth_gcp::VertexAuth; use litellm_callbacks::event::{Passthrough, WireRequest}; +use litellm_http::{ClientVariant, HttpClientConfig, HttpClientPool}; use serde::{Serialize, de::DeserializeOwned}; use serde_json::{Map, Value}; @@ -11,9 +10,8 @@ use crate::{ base_llm::ocr::{ error::Error, transformation::{ - BaseOcrConfig, DecodedOcrResponse, LiteLLMOcrResponse, OCR_CONNECT_TIMEOUT_SECS, - OcrDocument, OcrResponseContext, PreparedOcrRequest, decode_request_value, - decode_response, + BaseOcrConfig, DecodedOcrResponse, LiteLLMOcrResponse, OcrDocument, OcrResponseContext, + PreparedOcrRequest, decode_request_value, decode_response, }, }, custom_httpx::{ @@ -44,30 +42,15 @@ pub struct OcrClient { } impl OcrClient { - pub fn new(provider_http: reqwest::Client) -> Result { - let document_fetcher = MediaFetcher::new().map_err(transport::Error::from)?; + pub fn new(pool: &HttpClientPool, config: &HttpClientConfig) -> Result { Ok(Self { - provider_http, - polling_http: no_redirect_http()?, - document_fetcher, + provider_http: pool.client(config, ClientVariant::Provider)?, + polling_http: pool.client(config, ClientVariant::NoRedirect)?, + document_fetcher: MediaFetcher::new(pool, config)?, vertex_auth: VertexAuth::default(), }) } - pub fn shared() -> Result { - static CLIENT: OnceLock> = OnceLock::new(); - let client = CLIENT - .get_or_init(|| { - reqwest::Client::builder() - .connect_timeout(Duration::from_secs(OCR_CONNECT_TIMEOUT_SECS)) - .build() - .map_err(transport::Error::from) - .and_then(OcrClient::new) - }) - .clone()?; - Ok(client) - } - pub fn provider_http(&self) -> &reqwest::Client { &self.provider_http } @@ -88,21 +71,16 @@ impl OcrClient { pub fn for_test(provider_http: reqwest::Client, document_http: reqwest::Client) -> Self { Self { provider_http, - polling_http: no_redirect_http().expect("test polling client builds"), + polling_http: reqwest::Client::builder() + .redirect(reqwest::redirect::Policy::none()) + .build() + .expect("test polling client builds"), document_fetcher: MediaFetcher::for_test(document_http), vertex_auth: VertexAuth::default(), } } } -fn no_redirect_http() -> Result { - reqwest::Client::builder() - .connect_timeout(Duration::from_secs(OCR_CONNECT_TIMEOUT_SECS)) - .redirect(reqwest::redirect::Policy::none()) - .build() - .map_err(transport::Error::from) -} - /// Rust counterpart of `BaseLLMHTTPHandler.async_ocr`: prepare the provider request, /// send it, and hand the response to the config for normalization. pub async fn ocr( @@ -318,6 +296,8 @@ pub fn body_document(body: &Value) -> Result { #[cfg(test)] mod tests { + use std::time::Duration; + use super::*; #[tokio::test] diff --git a/litellm-rust/crates/llms/src/custom_httpx/media.rs b/litellm-rust/crates/llms/src/custom_httpx/media.rs index 0b7fa30e34b..aeac4894683 100644 --- a/litellm-rust/crates/llms/src/custom_httpx/media.rs +++ b/litellm-rust/crates/llms/src/custom_httpx/media.rs @@ -7,13 +7,12 @@ use std::{ time::Duration, }; +use litellm_http::{ClientVariant, HttpClientConfig, HttpClientPool}; use reqwest::{ Url, dns::{Addrs, Name, Resolve, Resolving}, }; -const MEDIA_CONNECT_TIMEOUT_SECS: u64 = 10; - #[derive(Debug, thiserror::Error)] pub enum Error { #[error("media URL rejected by network policy")] @@ -63,23 +62,30 @@ pub struct DownloadedMedia { } impl MediaFetcher { - pub fn new() -> Result { - Self::with_resolvers(Arc::new(PublicDnsResolver), Arc::new(SystemAddressResolver)) + pub fn new( + pool: &HttpClientPool, + config: &HttpClientConfig, + ) -> Result { + Self::with_resolvers( + pool, + config, + Arc::new(PublicDnsResolver), + Arc::new(SystemAddressResolver), + ) } fn with_resolvers( + pool: &HttpClientPool, + config: &HttpClientConfig, transport_resolver: Arc, address_resolver: Arc, - ) -> Result + ) -> Result where R: Resolve + 'static, { - let client = reqwest::Client::builder() - .connect_timeout(Duration::from_secs(MEDIA_CONNECT_TIMEOUT_SECS)) - .redirect(reqwest::redirect::Policy::none()) - .no_proxy() - .dns_resolver(transport_resolver) - .build()?; + let client = pool.client_with(config, ClientVariant::Media, |builder| { + builder.dns_resolver(transport_resolver) + })?; Ok(Self { client, address_resolver, @@ -281,6 +287,7 @@ impl Resolve for PublicDnsResolver { mod tests { use std::collections::HashSet; + use litellm_http::HttpSettings; use tokio::{ io::{AsyncReadExt, AsyncWriteExt}, net::TcpListener, @@ -365,6 +372,8 @@ mod tests { blocked_hosts: HashSet<&'static str>, ) -> MediaFetcher { MediaFetcher::with_resolvers( + &HttpClientPool::new(), + &HttpClientConfig::resolve(&HttpSettings::default(), None).unwrap(), Arc::new(LoopbackDnsResolver(address)), Arc::new(TestAddressResolver { blocked_hosts }), ) @@ -542,7 +551,12 @@ mod tests { #[tokio::test] async fn rejects_url_credentials_before_network_access() { - let fetcher = MediaFetcher::new().expect("media fetcher builds"); + let fetcher = MediaFetcher::new( + &HttpClientPool::new(), + &HttpClientConfig::resolve(&litellm_http::HttpSettings::default(), None) + .expect("default settings resolve"), + ) + .expect("media fetcher builds"); let url = Url::parse("https://user:password@8.8.8.8/document").expect("credentialed URL parses"); assert!(matches!( diff --git a/litellm-rust/crates/llms/src/custom_httpx/transport.rs b/litellm-rust/crates/llms/src/custom_httpx/transport.rs index 172dd96476a..8e5e1a8832d 100644 --- a/litellm-rust/crates/llms/src/custom_httpx/transport.rs +++ b/litellm-rust/crates/llms/src/custom_httpx/transport.rs @@ -26,6 +26,12 @@ impl From for Error { } } +impl From for Error { + fn from(error: litellm_http::Error) -> Self { + Self::Connect(error.to_string()) + } +} + #[cfg(test)] mod tests { #[tokio::test] diff --git a/litellm-rust/crates/python-bridge/Cargo.toml b/litellm-rust/crates/python-bridge/Cargo.toml index e9b7f384406..762e22e433b 100644 --- a/litellm-rust/crates/python-bridge/Cargo.toml +++ b/litellm-rust/crates/python-bridge/Cargo.toml @@ -20,6 +20,7 @@ bytes.workspace = true litellm-auth.workspace = true litellm-callbacks-legacy.workspace = true litellm-core.workspace = true +litellm-http.workspace = true litellm-llms.workspace = true litellm-types.workspace = true litellm-host-python.workspace = true diff --git a/litellm-rust/crates/python-bridge/src/http.rs b/litellm-rust/crates/python-bridge/src/http.rs new file mode 100644 index 00000000000..d2e05c3b949 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/http.rs @@ -0,0 +1,234 @@ +use std::{path::PathBuf, sync::LazyLock}; + +use litellm_http::{HttpClientConfig, HttpClientPool, HttpSettings, SslVerify}; +use pyo3::{prelude::*, types::PyDict}; + +use crate::errors::RustBridgeDeclined; + +static POOL: LazyLock = LazyLock::new(HttpClientPool::new); + +/// Keyword arguments that carry a live Python HTTP client or session. They cannot cross into +/// Rust, so a call that supplies one stays on the Python path. +const LIVE_CLIENT_ARGUMENTS: [&str; 3] = ["client", "shared_session", "aclient_session"]; + +pub(crate) fn pool() -> &'static HttpClientPool { + &POOL +} + +/// The client configuration for one call: the process settings from the `litellm` module and +/// the environment, narrowed by the call's own `ssl_verify`. +pub(crate) fn call_config( + py: Python<'_>, + kwargs: &Bound<'_, PyDict>, +) -> PyResult { + decline_live_clients(kwargs)?; + let settings = settings(py.import("litellm")?.as_any())? + .with_environment(&|name| std::env::var(name).ok()); + let per_call = kwargs + .get_item("ssl_verify")? + .map(|value| ssl_verify(&value, "ssl_verify")) + .transpose()? + .flatten(); + HttpClientConfig::resolve(&settings, per_call.as_ref()) + .map_err(|error| RustBridgeDeclined::new_err(error.to_string())) +} + +pub(crate) fn decline_live_clients(kwargs: &Bound<'_, PyDict>) -> PyResult<()> { + for name in LIVE_CLIENT_ARGUMENTS { + if kwargs.get_item(name)?.is_some_and(|value| !value.is_none()) { + return Err(RustBridgeDeclined::new_err(format!( + "{name} is a live Python HTTP client and cannot be used by the Rust route" + ))); + } + } + Ok(()) +} + +/// Read the `litellm.*` globals `http_handler.py` consults. `globals` is the `litellm` module in +/// production and any attribute holder in tests. +pub(crate) fn settings(globals: &Bound<'_, PyAny>) -> PyResult { + Ok(HttpSettings { + ssl_verify: ssl_verify(&globals.getattr("ssl_verify")?, "litellm.ssl_verify")?, + ssl_certificate: optional_path(globals, "ssl_certificate")?, + ssl_security_level: globals.getattr("ssl_security_level")?.extract()?, + ssl_ecdh_curve: globals.getattr("ssl_ecdh_curve")?.extract()?, + force_ipv4: globals.getattr("force_ipv4")?.extract()?, + http2: globals.getattr("http2")?.extract()?, + trust_proxy_env: globals.getattr("aiohttp_trust_env")?.extract()?, + ..HttpSettings::default() + }) +} + +fn optional_path(globals: &Bound<'_, PyAny>, name: &str) -> PyResult> { + Ok(globals + .getattr(name)? + .extract::>()? + .map(PathBuf::from)) +} + +fn ssl_verify(value: &Bound<'_, PyAny>, name: &str) -> PyResult> { + if value.is_none() { + return Ok(None); + } + if let Ok(enabled) = value.extract::() { + return Ok(Some(if enabled { + SslVerify::Enabled + } else { + SslVerify::Disabled + })); + } + if let Ok(path) = value.extract::() { + return Ok(Some(SslVerify::CaBundle(PathBuf::from(path)))); + } + Err(RustBridgeDeclined::new_err(format!( + "{name} is a live Python object and cannot be used by the Rust route" + ))) +} + +#[cfg(test)] +mod tests { + use litellm_http::Verify; + use rstest::rstest; + + use super::*; + + fn eval<'py>(py: Python<'py>, source: &std::ffi::CStr) -> Bound<'py, PyDict> { + let locals = PyDict::new(py); + py.run(source, Some(&locals), Some(&locals)).unwrap(); + locals + } + + fn globals<'py>(py: Python<'py>, overrides: &str) -> Bound<'py, PyAny> { + let source = format!( + " +import types +globals = types.SimpleNamespace( + ssl_verify=True, + ssl_certificate=None, + ssl_security_level=None, + ssl_ecdh_curve=None, + force_ipv4=False, + http2=False, + aiohttp_trust_env=False, +) +{overrides} +" + ); + let source = std::ffi::CString::new(source).unwrap(); + eval(py, &source).get_item("globals").unwrap().unwrap() + } + + #[test] + fn default_globals_produce_default_settings_with_verification_on() { + Python::initialize(); + Python::attach(|py| { + let settings = settings(&globals(py, "")).unwrap(); + assert_eq!( + settings, + HttpSettings { + ssl_verify: Some(SslVerify::Enabled), + ..HttpSettings::default() + } + ); + }); + } + + #[test] + fn globals_flow_into_settings() { + Python::initialize(); + Python::attach(|py| { + let settings = settings(&globals( + py, + " +globals.ssl_verify = '/etc/ssl/corp.pem' +globals.ssl_certificate = '/etc/ssl/client.pem' +globals.ssl_security_level = '2' +globals.ssl_ecdh_curve = 'X25519' +globals.force_ipv4 = True +globals.http2 = True +globals.aiohttp_trust_env = True +", + )) + .unwrap(); + assert_eq!( + settings, + HttpSettings { + ssl_verify: Some(SslVerify::CaBundle("/etc/ssl/corp.pem".into())), + ssl_certificate: Some("/etc/ssl/client.pem".into()), + ssl_security_level: Some("2".into()), + ssl_ecdh_curve: Some("X25519".into()), + force_ipv4: true, + http2: true, + trust_proxy_env: true, + ..HttpSettings::default() + } + ); + }); + } + + #[test] + fn disabled_verification_global_resolves_to_disabled() { + Python::initialize(); + Python::attach(|py| { + let settings = settings(&globals(py, "globals.ssl_verify = False")).unwrap(); + assert_eq!(settings.ssl_verify, Some(SslVerify::Disabled)); + let config = HttpClientConfig::resolve(&settings, None).unwrap(); + assert_eq!(config.verify, Verify::Disabled); + }); + } + + #[test] + fn ssl_context_global_declines_instead_of_being_dropped() { + Python::initialize(); + Python::attach(|py| { + let error = settings(&globals(py, "globals.ssl_verify = object()")).unwrap_err(); + assert!(error.is_instance_of::(py)); + assert!(error.value(py).to_string().contains("litellm.ssl_verify")); + }); + } + + #[rstest] + #[case::client("client")] + #[case::shared_session("shared_session")] + #[case::aclient_session("aclient_session")] + fn live_python_clients_decline_before_dispatch(#[case] name: &str) { + Python::initialize(); + Python::attach(|py| { + let kwargs = PyDict::new(py); + kwargs + .set_item(name, py.eval(c"object()", None, None).unwrap()) + .unwrap(); + let error = decline_live_clients(&kwargs).unwrap_err(); + assert!(error.is_instance_of::(py)); + assert!(error.value(py).to_string().contains(name)); + }); + } + + #[test] + fn none_valued_client_arguments_are_not_live_clients() { + Python::initialize(); + Python::attach(|py| { + let kwargs = PyDict::new(py); + for name in LIVE_CLIENT_ARGUMENTS { + kwargs.set_item(name, py.None()).unwrap(); + } + decline_live_clients(&kwargs).unwrap(); + }); + } + + #[rstest] + #[case::disabled(c"False", Some(SslVerify::Disabled))] + #[case::enabled(c"True", Some(SslVerify::Enabled))] + #[case::bundle(c"'/tmp/ca.pem'", Some(SslVerify::CaBundle("/tmp/ca.pem".into())))] + #[case::unset(c"None", None)] + fn per_call_ssl_verify_values_project( + #[case] source: &std::ffi::CStr, + #[case] expected: Option, + ) { + Python::initialize(); + Python::attach(|py| { + let value = py.eval(source, None, None).unwrap(); + assert_eq!(ssl_verify(&value, "ssl_verify").unwrap(), expected); + }); + } +} diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs index ca699e7c483..11cb0a7f655 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -1,6 +1,7 @@ mod credentials; mod diagnostics; mod errors; +mod http; mod marshal; mod routes; mod token_counter; diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs b/litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs index b5bb941708d..f252e1dc47b 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs @@ -12,6 +12,8 @@ use pyo3::{ types::{PyDict, PyTuple}, }; +use crate::{errors::RustBridgeDeclined, http}; + const SURFACE: LegacySurface = LegacySurface { call_type: "ocr", input_description: "OCR document processing", @@ -29,7 +31,9 @@ fn run_ocr( kwargs: Bound<'_, PyDict>, asynchronous: bool, ) -> PyResult> { - let client = OcrClient::shared().map_err(errors::to_pyerr)?; + let config = http::call_config(py, &kwargs)?; + let client = OcrClient::new(http::pool(), &config) + .map_err(|error| RustBridgeDeclined::new_err(error.to_string()))?; run_legacy_call( py, if asynchronous { ASYNC_SURFACE } else { SURFACE }, From 5a474fd7996e96f90e18c539108381d811cd5e63 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Fri, 18 Sep 2026 23:31:27 +0000 Subject: [PATCH 081/144] refactor(rust): inject VertexAuth into OcrClient so the bridge keeps one token cache Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm-rust/Cargo.lock | 2 ++ litellm-rust/crates/core/Cargo.toml | 1 + litellm-rust/crates/core/src/ocr/client.rs | 4 +++- litellm-rust/crates/core/tests/ocr.rs | 3 +++ .../crates/llms/src/custom_httpx/llm_http_handler.rs | 8 ++++++-- litellm-rust/crates/python-bridge/Cargo.toml | 1 + litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs | 7 ++++++- 7 files changed, 22 insertions(+), 4 deletions(-) diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index ffbdc80efd3..ed0fee3411f 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -2056,6 +2056,7 @@ dependencies = [ "futures-util", "litellm-auth", "litellm-auth-aws", + "litellm-auth-gcp", "litellm-callbacks", "litellm-core-utils", "litellm-http", @@ -2175,6 +2176,7 @@ dependencies = [ "criterion", "futures-util", "litellm-auth", + "litellm-auth-gcp", "litellm-callbacks-legacy", "litellm-core", "litellm-host-python", diff --git a/litellm-rust/crates/core/Cargo.toml b/litellm-rust/crates/core/Cargo.toml index 047188c74f9..ce8463affc1 100644 --- a/litellm-rust/crates/core/Cargo.toml +++ b/litellm-rust/crates/core/Cargo.toml @@ -15,6 +15,7 @@ futures-util.workspace = true base64.workspace = true litellm-auth.workspace = true litellm-auth-aws.workspace = true +litellm-auth-gcp.workspace = true litellm-http.workspace = true litellm-llms.workspace = true moka.workspace = true diff --git a/litellm-rust/crates/core/src/ocr/client.rs b/litellm-rust/crates/core/src/ocr/client.rs index 21c81505cff..a380037487e 100644 --- a/litellm-rust/crates/core/src/ocr/client.rs +++ b/litellm-rust/crates/core/src/ocr/client.rs @@ -1,3 +1,4 @@ +use litellm_auth_gcp::VertexAuth; use litellm_http::{HttpClientConfig, HttpClientPool}; use litellm_llms::{ base_llm::ocr::{error::Error, transformation::LiteLLMOcrResponse}, @@ -19,7 +20,8 @@ pub async fn perform( pub async fn ocr( pool: &HttpClientPool, config: &HttpClientConfig, + vertex_auth: VertexAuth, request: LiteLLMOcrRequest, ) -> Result { - perform(&OcrClient::new(pool, config)?, request).await + perform(&OcrClient::new(pool, config, vertex_auth)?, request).await } diff --git a/litellm-rust/crates/core/tests/ocr.rs b/litellm-rust/crates/core/tests/ocr.rs index d59c6179aa5..4ce6d4ee8f6 100644 --- a/litellm-rust/crates/core/tests/ocr.rs +++ b/litellm-rust/crates/core/tests/ocr.rs @@ -1,5 +1,6 @@ use std::sync::{Arc, Mutex}; +use litellm_auth_gcp::VertexAuth; use litellm_callbacks::{ event::{CallEvent, WireRequest}, host::{Host, HostOp, HostResult}, @@ -182,6 +183,7 @@ async fn facade_uses_the_injected_http_pool_configuration() { crate::ocr::client::ocr( &HttpClientPool::new(), &config, + VertexAuth::default(), wire_request("mistral/model", &base, json!({})), ) .await @@ -200,6 +202,7 @@ async fn unbuildable_http_configuration_fails_before_dispatch() { let error = crate::ocr::client::ocr( &HttpClientPool::new(), &config, + VertexAuth::default(), wire_request("mistral/model", &base, json!({})), ) .await diff --git a/litellm-rust/crates/llms/src/custom_httpx/llm_http_handler.rs b/litellm-rust/crates/llms/src/custom_httpx/llm_http_handler.rs index 38ddec08da9..9826d73a061 100644 --- a/litellm-rust/crates/llms/src/custom_httpx/llm_http_handler.rs +++ b/litellm-rust/crates/llms/src/custom_httpx/llm_http_handler.rs @@ -42,12 +42,16 @@ pub struct OcrClient { } impl OcrClient { - pub fn new(pool: &HttpClientPool, config: &HttpClientConfig) -> Result { + pub fn new( + pool: &HttpClientPool, + config: &HttpClientConfig, + vertex_auth: VertexAuth, + ) -> Result { Ok(Self { provider_http: pool.client(config, ClientVariant::Provider)?, polling_http: pool.client(config, ClientVariant::NoRedirect)?, document_fetcher: MediaFetcher::new(pool, config)?, - vertex_auth: VertexAuth::default(), + vertex_auth, }) } diff --git a/litellm-rust/crates/python-bridge/Cargo.toml b/litellm-rust/crates/python-bridge/Cargo.toml index 762e22e433b..c66701548d1 100644 --- a/litellm-rust/crates/python-bridge/Cargo.toml +++ b/litellm-rust/crates/python-bridge/Cargo.toml @@ -20,6 +20,7 @@ bytes.workspace = true litellm-auth.workspace = true litellm-callbacks-legacy.workspace = true litellm-core.workspace = true +litellm-auth-gcp.workspace = true litellm-http.workspace = true litellm-llms.workspace = true litellm-types.workspace = true diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs b/litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs index f252e1dc47b..be188434275 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs @@ -3,7 +3,10 @@ mod errors; mod host; mod project; +use std::sync::LazyLock; + use host::OcrRouteHost; +use litellm_auth_gcp::VertexAuth; use litellm_callbacks_legacy::{LegacySurface, PublicCall, run_legacy_call}; use litellm_core::ocr::route::ocr_machine; use litellm_llms::custom_httpx::llm_http_handler::OcrClient; @@ -24,6 +27,8 @@ const ASYNC_SURFACE: LegacySurface = LegacySurface { ..SURFACE }; +static VERTEX_AUTH: LazyLock = LazyLock::new(VertexAuth::default); + fn run_ocr( py: Python<'_>, request: Bound<'_, PyAny>, @@ -32,7 +37,7 @@ fn run_ocr( asynchronous: bool, ) -> PyResult> { let config = http::call_config(py, &kwargs)?; - let client = OcrClient::new(http::pool(), &config) + let client = OcrClient::new(http::pool(), &config, VERTEX_AUTH.clone()) .map_err(|error| RustBridgeDeclined::new_err(error.to_string()))?; run_legacy_call( py, From 3d7a771ea7d2a327a4cf09b5425fdfe4e3fbc69d Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 18 Sep 2026 16:53:06 -0700 Subject: [PATCH 082/144] fix(vertex_ai): apply finals before interims and refresh the token per stream --- .../audio_transcription/realtime_backend.py | 33 ++++++--- .../realtime_transformation.py | 30 ++++---- .../llms/vertex_ai/realtime/transformation.py | 23 +++++++ litellm/realtime_api/main.py | 35 ++++------ .../test_vertex_ai_realtime_backend.py | 68 ++++++++++++++++++- .../test_vertex_ai_realtime_transformation.py | 48 ++++++++++++- 6 files changed, 186 insertions(+), 51 deletions(-) diff --git a/litellm/llms/vertex_ai/audio_transcription/realtime_backend.py b/litellm/llms/vertex_ai/audio_transcription/realtime_backend.py index 87c72c193bc..0c16fea9e9f 100644 --- a/litellm/llms/vertex_ai/audio_transcription/realtime_backend.py +++ b/litellm/llms/vertex_ai/audio_transcription/realtime_backend.py @@ -80,7 +80,7 @@ class _Closed: pass -def open_speech_client(target: SpeechStreamingTarget) -> SpeechStreamingClient: +def open_speech_client(target: SpeechStreamingTarget, access_token: str) -> SpeechStreamingClient: try: from google.api_core.client_options import ClientOptions from google.cloud.speech_v2 import SpeechAsyncClient @@ -88,7 +88,7 @@ def open_speech_client(target: SpeechStreamingTarget) -> SpeechStreamingClient: except ImportError as e: raise ImportError(SPEECH_SDK_INSTALL_HINT) from e return SpeechAsyncClient( - credentials=Credentials(token=target.access_token), + credentials=Credentials(token=access_token), transport="grpc_asyncio", client_options=ClientOptions(api_endpoint=target.api_endpoint), ) @@ -157,6 +157,7 @@ class _RecognizeStream: self.speech_active: bool = False self.billed_seconds: float = 0.0 self._cancelled: bool = False + self._closed: bool = False self._task: asyncio.Task[None] | None = None async def send_audio(self, audio: bytes) -> None: @@ -170,8 +171,15 @@ class _RecognizeStream: if self._task is not None: self._task.cancel() + async def close(self) -> None: + if self._closed: + return + self._closed = True + await self._client.transport.close() + async def relay(self, outbox: "asyncio.Queue[str | _StreamFailure | _Closed]", billed_before: float) -> float: if self._cancelled: + await self.close() return 0.0 task: Final = asyncio.create_task(self._forward(outbox, billed_before)) self._task = task @@ -181,6 +189,8 @@ class _RecognizeStream: task.cancel() await asyncio.wait((task,)) raise + finally: + await self.close() return self.billed_seconds async def _forward(self, outbox: "asyncio.Queue[str | _StreamFailure | _Closed]", billed_before: float) -> None: @@ -209,7 +219,7 @@ class SpeechStreamingBackend: self, target: SpeechStreamingTarget, *, - client_factory: Callable[[SpeechStreamingTarget], SpeechStreamingClient] = open_speech_client, + client_factory: Callable[[SpeechStreamingTarget, str], SpeechStreamingClient] = open_speech_client, clock: Callable[[], float] = time.monotonic, rotation_seconds: float = STREAM_ROTATION_SECONDS, rotation_deadline_seconds: float = STREAM_ROTATION_DEADLINE_SECONDS, @@ -222,7 +232,6 @@ class SpeechStreamingBackend: self._outbox: Final[asyncio.Queue[str | _StreamFailure | _Closed]] = asyncio.Queue(maxsize=OUTBOX_SIZE) self._links: Final[asyncio.Queue[_RecognizeStream | str]] = asyncio.Queue(maxsize=_LINK_QUEUE_SIZE) self._pump: asyncio.Task[None] | None = None - self._client: SpeechStreamingClient | None = None self._config: StreamingRecognitionConfig | None = None self._turn: tuple[_RecognizeStream, ...] = () self._billed_before: float = 0.0 @@ -282,13 +291,16 @@ class SpeechStreamingBackend: if pump is not None: pump.cancel() await asyncio.wait((pump,)) - client: Final = self._client - self._client = None - if client is not None: - await client.transport.close() + await self._close_unrelayed_streams() if not self._outbox.full(): self._outbox.put_nowait(_Closed()) + async def _close_unrelayed_streams(self) -> None: + unrelayed: Final = tuple(self._links.get_nowait() for _ in range(self._links.qsize())) + for link in unrelayed: + if isinstance(link, _RecognizeStream): + await link.close() + async def _link(self, item: _RecognizeStream | str) -> None: if self._pump is None: self._pump = asyncio.create_task(self._pump_links()) @@ -333,10 +345,9 @@ class SpeechStreamingBackend: config: Final = self._config if config is None: raise RuntimeError("audio was sent before the Speech-to-Text stream was configured") - if self._client is None: - self._client = self._client_factory(self._target) + access_token: Final = await self._target.resolve_access_token() stream: Final = _RecognizeStream( - client=self._client, + client=self._client_factory(self._target, access_token), request_type=StreamingRecognizeRequest, first_request=StreamingRecognizeRequest(recognizer=self._target.recognizer, streaming_config=config), opened_at=self._clock(), diff --git a/litellm/llms/vertex_ai/audio_transcription/realtime_transformation.py b/litellm/llms/vertex_ai/audio_transcription/realtime_transformation.py index 6ec7a21a134..dab2e980fd0 100644 --- a/litellm/llms/vertex_ai/audio_transcription/realtime_transformation.py +++ b/litellm/llms/vertex_ai/audio_transcription/realtime_transformation.py @@ -1,4 +1,4 @@ -from collections.abc import Callable, Mapping +from collections.abc import Awaitable, Callable, Mapping from dataclasses import dataclass, replace from typing import Final @@ -6,7 +6,7 @@ from pydantic import JsonValue, TypeAdapter from typing_extensions import assert_never import litellm -from litellm import verbose_logger +from litellm._logging import verbose_logger from litellm._uuid import uuid from litellm.litellm_core_utils.audio_utils.utils import normalize_transcription_language_to_bcp47 from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj @@ -73,7 +73,7 @@ class ChirpProtocolError(RealtimeTranscriptionProtocolError): class SpeechStreamingTarget: api_endpoint: str recognizer: str - access_token: str + resolve_access_token: Callable[[], Awaitable[str]] @dataclass(frozen=True, slots=True) @@ -249,16 +249,14 @@ class ChirpEventTransformer: finals: Final = tuple( result.transcript.strip() for result in frame.results if result.is_final and result.transcript.strip() ) - begin_events: Final = self._begin() if frame.speech_event == "begin" or interim or finals else () - interim_events: Final = self._hypothesis(interim) if interim else () + begin_events: Final = self._begin() if frame.speech_event == "begin" else () final_events: Final = tuple(event for final in finals for event in self._final(final)) + interim_events: Final = self._hypothesis(interim) if interim else () end_events: Final = self._stop() if frame.speech_event == "end" else () - return (*begin_events, *interim_events, *final_events, *end_events) + return (*begin_events, *final_events, *interim_events, *end_events) def _begin(self) -> tuple[OpenAIRealtimeEvents, ...]: - if self._turn is None: - self._turn = _Turn(item_id=self._new_item_id()) - turn: Final = self._turn + turn: Final = self._require_turn() if turn.started_emitted or not self._require_config().server_vad: return () self._turn = replace(turn, started_emitted=True) @@ -272,21 +270,23 @@ class ChirpEventTransformer: return (speech_event("input_audio_buffer.speech_stopped", turn.item_id),) def _hypothesis(self, text: str) -> tuple[OpenAIRealtimeEvents, ...]: + begin_events: Final = self._begin() turn: Final = self._require_turn() hypothesis: Final = _join_transcript(turn.committed, text) delta: Final = new_words(turn.preview, hypothesis) self._turn = replace(turn, preview=hypothesis) - return (delta_event(turn.item_id, delta),) if delta else () + return (*begin_events, delta_event(turn.item_id, delta)) if delta else begin_events def _final(self, text: str) -> tuple[OpenAIRealtimeEvents, ...]: + begin_events: Final = self._begin() turn: Final = self._require_turn() committed: Final = _join_transcript(turn.committed, text) delta: Final = new_words(turn.preview, committed) self._turn = replace(turn, committed=committed, preview=committed) delta_events: Final[tuple[OpenAIRealtimeEvents, ...]] = (delta_event(turn.item_id, delta),) if delta else () if not self._require_config().server_vad: - return delta_events - return (*delta_events, *self._complete()) + return (*begin_events, *delta_events) + return (*begin_events, *delta_events, *self._complete()) def _finish_turn(self) -> tuple[OpenAIRealtimeEvents, ...]: if self._turn is None: @@ -326,12 +326,12 @@ class VertexChirpRealtimeConfig(BaseRealtimeConfig): def __init__( self, *, - access_token: str, + resolve_access_token: Callable[[], Awaitable[str]], project: str, location: str | None, backend_factory: Callable[[SpeechStreamingTarget], RealtimeBackend] = _default_backend_factory, ) -> None: - self._access_token: Final = access_token + self._resolve_access_token: Final = resolve_access_token self._project: Final = validate_vertex_transcription_project_id(project) self._location: Final = validate_vertex_transcription_location(location, DEFAULT_SPEECH_TO_TEXT_LOCATION) self._backend_factory: Final = backend_factory @@ -357,7 +357,7 @@ class VertexChirpRealtimeConfig(BaseRealtimeConfig): SpeechStreamingTarget( api_endpoint=url, recognizer=f"projects/{self._project}/locations/{self._location}/recognizers/_", - access_token=self._access_token, + resolve_access_token=self._resolve_access_token, ) ) diff --git a/litellm/llms/vertex_ai/realtime/transformation.py b/litellm/llms/vertex_ai/realtime/transformation.py index fe59034c27b..9fed6d52f0e 100644 --- a/litellm/llms/vertex_ai/realtime/transformation.py +++ b/litellm/llms/vertex_ai/realtime/transformation.py @@ -12,10 +12,16 @@ Auth: OAuth2 Bearer token (not an API key). """ import json +from collections.abc import Awaitable, Callable from typing import Final from litellm import verbose_logger from litellm.llms.gemini.realtime.transformation import GeminiRealtimeConfig +from litellm.llms.vertex_ai.audio_transcription.realtime_transformation import ( + VertexChirpRealtimeConfig, + is_vertex_speech_to_text_model, +) +from litellm.llms.vertex_ai.vertex_llm_base import VertexBase class VertexAIRealtimeConfig(GeminiRealtimeConfig): @@ -232,3 +238,20 @@ class VertexAIRealtimeConfig(GeminiRealtimeConfig): return [] return super().transform_realtime_request(message, model, session_configuration_request) + + +def vertex_realtime_config( + model: str, + *, + access_token: str, + resolve_access_token: Callable[[], Awaitable[str]], + project: str, + location: str | None, +) -> VertexAIRealtimeConfig | VertexChirpRealtimeConfig: + if is_vertex_speech_to_text_model(model): + return VertexChirpRealtimeConfig(resolve_access_token=resolve_access_token, project=project, location=location) + return VertexAIRealtimeConfig( + access_token=access_token, + project=project, + location=VertexBase.get_vertex_region(vertex_region=location, model=model), + ) diff --git a/litellm/realtime_api/main.py b/litellm/realtime_api/main.py index aed7bf15bdc..0e83edab5e1 100644 --- a/litellm/realtime_api/main.py +++ b/litellm/realtime_api/main.py @@ -38,11 +38,8 @@ from ..llms.azure.realtime.handler import AzureOpenAIRealtime, azure_realtime_pr from ..llms.bedrock.realtime.handler import BedrockRealtime from ..llms.custom_httpx.http_handler import get_shared_realtime_ssl_context from ..llms.openai.realtime.handler import OpenAIRealtime -from ..llms.vertex_ai.audio_transcription.realtime_transformation import ( - VertexChirpRealtimeConfig, - is_vertex_speech_to_text_model, -) -from ..llms.vertex_ai.realtime.transformation import VertexAIRealtimeConfig +from ..llms.vertex_ai.audio_transcription.realtime_transformation import is_vertex_speech_to_text_model +from ..llms.vertex_ai.realtime.transformation import VertexAIRealtimeConfig, vertex_realtime_config from ..llms.vertex_ai.vertex_llm_base import VertexBase from ..llms.xai.realtime.handler import XAIRealtime from ..utils import client as wrapper_client @@ -555,9 +552,19 @@ async def _arealtime( timeout_seconds=REALTIME_CREDENTIAL_RESOLUTION_TIMEOUT_SECONDS, ) - vertex_realtime_config: Final = _vertex_realtime_config( - model=model, + async def resolve_vertex_access_token() -> str: + refreshed_token, _ = await _resolve_vertex_access_token_bounded( + credentials=vertex_credentials, + project_id=resolved_project, + resolver=vertex_access_token_resolver, + timeout_seconds=REALTIME_CREDENTIAL_RESOLUTION_TIMEOUT_SECONDS, + ) + return refreshed_token + + vertex_provider_config: Final = vertex_realtime_config( + model, access_token=access_token, + resolve_access_token=resolve_vertex_access_token, project=resolved_project, location=vertex_location, ) @@ -566,7 +573,7 @@ async def _arealtime( model=model, websocket=websocket, logging_obj=litellm_logging_obj, - provider_config=vertex_realtime_config, + provider_config=vertex_provider_config, api_base=dynamic_api_base or litellm_params.api_base, api_key=None, client=client, @@ -580,18 +587,6 @@ async def _arealtime( raise ValueError(f"Unsupported model: {model}") -def _vertex_realtime_config( - model: str, access_token: str, project: str, location: str | None -) -> VertexAIRealtimeConfig | VertexChirpRealtimeConfig: - if is_vertex_speech_to_text_model(model): - return VertexChirpRealtimeConfig(access_token=access_token, project=project, location=location) - return VertexAIRealtimeConfig( - access_token=access_token, - project=project, - location=vertex_llm_base.get_vertex_region(vertex_region=location, model=model), - ) - - def _is_transcription_only_realtime_model(model: str, custom_llm_provider: str) -> bool: try: model_info: Final = litellm.get_model_info(model=model, custom_llm_provider=custom_llm_provider) diff --git a/tests/test_litellm/llms/vertex_ai/audio_transcription/test_vertex_ai_realtime_backend.py b/tests/test_litellm/llms/vertex_ai/audio_transcription/test_vertex_ai_realtime_backend.py index 49758e62415..d6f65c90806 100644 --- a/tests/test_litellm/llms/vertex_ai/audio_transcription/test_vertex_ai_realtime_backend.py +++ b/tests/test_litellm/llms/vertex_ai/audio_transcription/test_vertex_ai_realtime_backend.py @@ -1,6 +1,7 @@ import asyncio import json from collections.abc import AsyncIterator, Sequence +from dataclasses import replace from datetime import timedelta from typing import Final @@ -17,10 +18,15 @@ from websockets.exceptions import ConnectionClosedError, ConnectionClosedOK from litellm.llms.vertex_ai.audio_transcription.realtime_backend import REQUEST_QUEUE_SIZE, SpeechStreamingBackend from litellm.llms.vertex_ai.audio_transcription.realtime_transformation import SpeechStreamingTarget + +async def _static_token() -> str: + return "token" + + TARGET: Final = SpeechStreamingTarget( api_endpoint="us-speech.googleapis.com", recognizer="projects/proj-1/locations/us/recognizers/_", - access_token="token", + resolve_access_token=_static_token, ) CONFIGURE: Final = json.dumps( {"kind": "configure", "model": "chirp_3", "language_codes": ["en-US"], "sample_rate_hertz": 16_000} @@ -101,7 +107,7 @@ class _FakeSpeechClient: def _backend(client: _FakeSpeechClient, **kwargs: object) -> SpeechStreamingBackend: - return SpeechStreamingBackend(TARGET, client_factory=lambda target: client, **kwargs) + return SpeechStreamingBackend(TARGET, client_factory=lambda target, access_token: client, **kwargs) async def _recv(backend: SpeechStreamingBackend) -> dict[str, object]: @@ -346,6 +352,64 @@ async def test_rotation_is_forced_at_the_deadline_during_continuous_speech(): assert [_audio(stream) for stream in client.streams] == [[b"\x01\x01", b"\x02\x02"], [b"\x03\x03"]] +@pytest.mark.asyncio +async def test_every_stream_opens_its_own_client_with_a_freshly_resolved_token(): + now = [0.0] + tokens = iter(("token-1", "token-2")) + seen_tokens: list[str] = [] + clients = [_FakeSpeechClient([_response("first")]), _FakeSpeechClient([_response("second")])] + unopened = iter(clients) + + async def resolve_access_token() -> str: + return next(tokens) + + def open_client(target: SpeechStreamingTarget, access_token: str) -> _FakeSpeechClient: + seen_tokens.append(access_token) + return next(unopened) + + backend = SpeechStreamingBackend( + replace(TARGET, resolve_access_token=resolve_access_token), + client_factory=open_client, + clock=lambda: now[0], + rotation_seconds=240.0, + ) + async with backend: + await _configure(backend) + await backend.send(b"\x01\x01") + assert await _transcript(backend) == "first" + now[0] = 240.0 + await backend.send(b"\x02\x02") + assert await _transcript(backend) == "second" + assert clients[0].transport.closed + assert not clients[1].transport.closed + assert seen_tokens == ["token-1", "token-2"] + assert [len(client.streams) for client in clients] == [1, 1] + assert clients[1].transport.closed + + +@pytest.mark.asyncio +async def test_close_releases_a_rotated_stream_that_never_started_relaying(): + now = [0.0] + hold = asyncio.Event() + clients = [_FakeSpeechClient([_response("first"), hold]), _FakeSpeechClient([_response("never")])] + unopened = iter(clients) + backend = SpeechStreamingBackend( + TARGET, + client_factory=lambda target, access_token: next(unopened), + clock=lambda: now[0], + rotation_seconds=240.0, + ) + await _configure(backend) + await backend.send(b"\x01\x01") + assert await _transcript(backend) == "first" + now[0] = 240.0 + await backend.send(b"\x02\x02") + await asyncio.sleep(0) + assert clients[1].streams == [] + await backend.close() + assert [client.transport.closed for client in clients] == [True, True] + + @pytest.mark.asyncio async def test_discard_turn_cancels_every_stream_of_the_turn(): now = [0.0] diff --git a/tests/test_litellm/llms/vertex_ai/audio_transcription/test_vertex_ai_realtime_transformation.py b/tests/test_litellm/llms/vertex_ai/audio_transcription/test_vertex_ai_realtime_transformation.py index fcbfad8bf12..719dd621c82 100644 --- a/tests/test_litellm/llms/vertex_ai/audio_transcription/test_vertex_ai_realtime_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/audio_transcription/test_vertex_ai_realtime_transformation.py @@ -63,8 +63,12 @@ def _ga_session_update( ) +async def _token() -> str: + return "token" + + def _config(location: str | None = "us") -> VertexChirpRealtimeConfig: - return VertexChirpRealtimeConfig(access_token="token", project="proj-1", location=location) + return VertexChirpRealtimeConfig(resolve_access_token=_token, project="proj-1", location=location) def _configured( @@ -261,6 +265,41 @@ def test_server_vad_turn_streams_new_words_then_completes_with_usage(): assert _backend_events(config, _response(speech_event="end")) == [] +def test_server_vad_final_result_completes_before_the_interim_that_follows_it(): + config = _configured() + _backend_events(config, _response(speech_event="begin")) + events = _backend_events(config, _response(("four score", True), ("and seven", False))) + assert _types(events) == [ + DELTA, + "input_audio_buffer.speech_stopped", + COMPLETED, + "input_audio_buffer.speech_started", + DELTA, + ] + assert events[2]["transcript"] == "four score" + assert events[4]["delta"] == "and seven" + assert events[4]["item_id"] != events[2]["item_id"] + assert events[4]["item_id"] == events[3]["item_id"] + finished = _backend_events(config, _response(("and seven years", True))) + assert [(event["type"], event.get("delta", event.get("transcript"))) for event in finished] == [ + (DELTA, " years"), + ("input_audio_buffer.speech_stopped", None), + (COMPLETED, "and seven years"), + ] + assert {event["item_id"] for event in finished} == {events[4]["item_id"]} + + +def test_manual_turn_keeps_the_interim_that_follows_a_final_in_the_same_frame(): + config = _configured(turn_detection=None) + first = _backend_events(config, _response(("four score", True), ("and seven", False))) + assert [(event["type"], event["delta"]) for event in first] == [(DELTA, "four score"), (DELTA, " and seven")] + second = _backend_events(config, _response(("and seven years", True))) + assert [event["delta"] for event in second] == [" years"] + completed = _backend_events(config, VertexSpeechStreamingTurnFinished()) + assert [(event["type"], event["transcript"]) for event in completed] == [(COMPLETED, "four score and seven years")] + assert {event["item_id"] for event in (*first, *second, *completed)} == {first[0]["item_id"]} + + def test_manual_turns_complete_on_commit_without_speech_events(): config = _configured(turn_detection=None) assert _backend_events(config, _response(speech_event="begin")) == [] @@ -322,7 +361,9 @@ async def test_open_backend_targets_the_regional_speech_endpoint(): targets.append(target) return _NullBackend() - config = VertexChirpRealtimeConfig(access_token="token", project="proj-1", location=None, backend_factory=factory) + config = VertexChirpRealtimeConfig( + resolve_access_token=_token, project="proj-1", location=None, backend_factory=factory + ) url = config.get_complete_url(None, "vertex_ai/chirp_3") assert url == "us-speech.googleapis.com" assert config.validate_environment({}, MODEL, "https://" + url) == {} @@ -332,9 +373,10 @@ async def test_open_backend_targets_the_regional_speech_endpoint(): SpeechStreamingTarget( api_endpoint="us-speech.googleapis.com", recognizer="projects/proj-1/locations/us/recognizers/_", - access_token="token", + resolve_access_token=_token, ) ] + assert await targets[0].resolve_access_token() == "token" @pytest.mark.parametrize( From febe9aec6582f3aa47a9e0fcd405b4c2cb6c86fc Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 18 Sep 2026 17:10:01 -0700 Subject: [PATCH 083/144] fix(responses): book a rejected WebSocket connection as a failed request --- litellm/llms/custom_httpx/llm_http_handler.py | 7 +- litellm/proxy/_lazy_openapi_snapshot.json | 2 +- .../proxy/response_api_endpoints/endpoints.py | 8 +- litellm/responses/main.py | 6 +- litellm/responses/streaming_iterator.py | 24 ++-- litellm/utils.py | 2 +- .../test_litellm_logging.py | 32 +++++ .../response_api_endpoints/test_endpoints.py | 68 +++++++++++ .../test_responses_websocket_all_providers.py | 112 ++++++++++++++++++ 9 files changed, 243 insertions(+), 18 deletions(-) diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index ab327299243..221bc241999 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -6589,7 +6589,7 @@ class BaseLLMHTTPHandler: custom_llm_provider: str | None = None, first_message: str | None = None, **kwargs: Any, - ): + ) -> Exception | None: """ Handles Responses API WebSocket mode. @@ -6623,7 +6623,7 @@ class BaseLLMHTTPHandler: **kwargs, ) await handler.run() - return + return None import websockets from websockets.asyncio.client import ClientConnection @@ -6744,7 +6744,7 @@ class BaseLLMHTTPHandler: authorized_model=model, custom_llm_provider=custom_llm_provider, ) - await streaming.bidirectional_forward() + return await streaming.bidirectional_forward() except websockets.exceptions.InvalidStatusCode as e: verbose_logger.exception("Error connecting to responses WS backend: %s", e) @@ -6758,6 +6758,7 @@ class BaseLLMHTTPHandler: pass else: raise Exception(f"Unexpected error while closing WebSocket: {close_error}") + return None def image_edit_handler( self, diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index 40b64160b71..213cd88b6ce 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -19616,7 +19616,7 @@ } } }, - "description": "\nUnified rate-limit error.\n\nEvery rate-limit condition surfaced by litellm \u2014 whether it originated from\nan upstream LLM provider, a vendor batch endpoint, or one of litellm's own\nproxy-side limiters (parallel-requests, dynamic-rate, batch-rate, budget,\nmax-iterations, etc.) \u2014 is raised as an instance of this class.\n\nThe :attr:`category` attribute lets callers distinguish the source. See\n:class:`RateLimitErrorCategory` for the available values.\n" + "description": "\n Unified rate-limit error.\n\n Every rate-limit condition surfaced by litellm \u2014 whether it originated from\n an upstream LLM provider, a vendor batch endpoint, or one of litellm's own\n proxy-side limiters (parallel-requests, dynamic-rate, batch-rate, budget,\n max-iterations, etc.) \u2014 is raised as an instance of this class.\n\n The :attr:`category` attribute lets callers distinguish the source. See\n :class:`RateLimitErrorCategory` for the available values.\n " }, "500": { "content": { diff --git a/litellm/proxy/response_api_endpoints/endpoints.py b/litellm/proxy/response_api_endpoints/endpoints.py index ea6b67fa026..4b178c52de8 100644 --- a/litellm/proxy/response_api_endpoints/endpoints.py +++ b/litellm/proxy/response_api_endpoints/endpoints.py @@ -1567,7 +1567,13 @@ async def responses_websocket_endpoint( llm_router=llm_router, user_model=user_model, ) - await llm_call + failure: Final = await llm_call + if isinstance(failure, Exception): + await proxy_logging_obj.post_call_failure_hook( + user_api_key_dict=user_api_key_dict, + original_exception=failure, + request_data=data, + ) except Exception: verbose_proxy_logger.exception("Responses WebSocket error") await websocket.close(code=1011, reason="Internal server error") diff --git a/litellm/responses/main.py b/litellm/responses/main.py index 9705794d01d..3a4be06add9 100644 --- a/litellm/responses/main.py +++ b/litellm/responses/main.py @@ -2269,11 +2269,11 @@ async def _aresponses_websocket( api_key: str | None = None, timeout: float | None = None, **kwargs, -): +) -> Exception | None: """ Private function to handle the Responses API WebSocket mode. - For PROXY use only. + For PROXY use only. Returns the provider failure that ended the connection, if any. Resolves the LLM provider from ``model``, looks up the matching ``BaseResponsesAPIConfig``, and hands off to @@ -2343,7 +2343,7 @@ async def _aresponses_websocket( } remaining_kwargs: Final = {k: v for k, v in kwargs.items() if k not in _explicit_keys} - await base_llm_http_handler.async_responses_websocket( + return await base_llm_http_handler.async_responses_websocket( model=resolved_model, websocket=websocket, logging_obj=litellm_logging_obj, diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index c99a481db7d..b9abdffce94 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -1857,6 +1857,16 @@ class ResponsesWebSocketStreaming: if self.logging_obj: self.logging_obj.pre_call(input=message, api_key="") + def _failure_exception(self) -> Exception | None: + failed_event: Final = next( + (event for event in self.messages if event.get("type") in _RESPONSES_WS_FAILURE_EVENT_TYPES), None + ) + if failed_event is None: + return None + return _map_stream_error_to_exception( + _ws_event_error(failed_event), self.authorized_model or "", self.custom_llm_provider or "" + ) + async def _log_messages(self) -> None: if not self.logging_obj: return @@ -1864,16 +1874,11 @@ class ResponsesWebSocketStreaming: self.logging_obj.model_call_details["messages"] = self.input_messages if not self.messages: return - failed_event: Final = next( - (event for event in self.messages if event.get("type") in _RESPONSES_WS_FAILURE_EVENT_TYPES), None - ) - if failed_event is None: + exception: Final = self._failure_exception() + if exception is None: asyncio.create_task(self.logging_obj.dispatch_success_handlers(self.messages, prefer_async_handlers=True)) return self._record_usage_for_failure() - exception: Final = _map_stream_error_to_exception( - _ws_event_error(failed_event), self.authorized_model or "", self.custom_llm_provider or "" - ) traceback_exception: Final = "".join(traceback.format_exception(exception)) asyncio.create_task( self.logging_obj.dispatch_failure_handlers(exception, traceback_exception, prefer_async_handlers=True) @@ -2306,8 +2311,8 @@ class ResponsesWebSocketStreaming: except Exception as e: verbose_logger.debug("Responses WS client_to_backend ended: %s", e) - async def bidirectional_forward(self) -> None: - """Run both forwarding directions concurrently.""" + async def bidirectional_forward(self) -> Exception | None: + """Run both forwarding directions concurrently and return the provider failure that ended the connection.""" forward_task: Final = asyncio.create_task(self.backend_to_client()) try: await self.client_to_backend() @@ -2324,6 +2329,7 @@ class ResponsesWebSocketStreaming: await self.backend_ws.close() except Exception: pass + return self._failure_exception() # --------------------------------------------------------------------------- diff --git a/litellm/utils.py b/litellm/utils.py index 2c9200fbad7..298c5471b48 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -2008,7 +2008,7 @@ def client(original_function): result=result, call_type=call_type, ) - elif call_type == CallTypes.arealtime.value: + elif call_type in (CallTypes.arealtime.value, CallTypes.aresponses_websocket.value): return result ### POST-CALL RULES ### post_call_processing( diff --git a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py index 8ce5357dc94..0d2d600a7fc 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -1068,6 +1068,38 @@ async def test_arealtime_marks_litellm_params_async(monkeypatch): assert LitellmLogging._is_sync_litellm_request(captured["litellm_params"]) is False +@pytest.mark.asyncio +async def test_aresponses_websocket_hands_back_the_provider_failure_without_a_success_log(monkeypatch): + """A native Responses WebSocket connection the provider rejected comes back from the ``@client`` + wrapper as the mapped failure, and the wrapper books no success for it: the relay's own dispatch + is the connection's single log, so the proxy can record the connection as a failed request.""" + from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER + from litellm.responses.main import base_llm_http_handler + + success_events = [] + + class CaptureLogger(CustomLogger): + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + success_events.append(response_obj) + + monkeypatch.setattr(litellm, "callbacks", [CaptureLogger()]) + monkeypatch.setattr(litellm, "failure_callback", []) + monkeypatch.setattr(litellm, "_async_failure_callback", []) + monkeypatch.setattr(litellm, "success_callback", []) + monkeypatch.setattr(litellm, "_async_success_callback", []) + failure = litellm.BadRequestError(message="invalid_encrypted_content", model="gpt-4o", llm_provider="openai") + with patch.object( # test-quality-ok: the provider socket is the seam; how the wrapper treats the relay's outcome is under test + base_llm_http_handler, "async_responses_websocket", AsyncMock(return_value=failure) + ): + outcome = await litellm._aresponses_websocket(model="openai/gpt-4o", websocket=MagicMock(), api_key="sk-test") + await asyncio.sleep(0) + with contextlib.suppress(asyncio.TimeoutError): + await asyncio.wait_for(GLOBAL_LOGGING_WORKER.flush(), timeout=10.0) + + assert outcome is failure + assert success_events == [] + + @pytest.mark.asyncio async def test_agenerate_content_marks_litellm_params_async(): """LIT-4475: the async ``agenerate_content`` entrypoint must plant diff --git a/tests/test_litellm/proxy/response_api_endpoints/test_endpoints.py b/tests/test_litellm/proxy/response_api_endpoints/test_endpoints.py index 91d688fbacf..45ec529ce7d 100644 --- a/tests/test_litellm/proxy/response_api_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/response_api_endpoints/test_endpoints.py @@ -570,6 +570,74 @@ class TestResponsesWSFirstFrameModelAuth: assert mock_route_request.await_args.kwargs["route_type"] == "_aresponses_websocket" ws.close.assert_not_awaited() + @pytest.mark.asyncio + @pytest.mark.parametrize("provider_rejected", [True, False]) + async def test_endpoint_books_a_provider_rejected_connection_as_a_failed_request(self, provider_rejected): + from litellm.proxy.response_api_endpoints.endpoints import ( + responses_websocket_endpoint, + ) + + ws = MagicMock() + ws.headers = {} + ws.query_params = {} + ws.scope = {"headers": []} + ws.url = "ws://testserver/v1/responses" + ws.accept = AsyncMock() + ws.receive_text = AsyncMock( + return_value=json.dumps({"type": "response.create", "model": "gpt-4o-mini", "input": []}) + ) + ws.close = AsyncMock() + + processor = MagicMock() + processor.common_processing_pre_call_logic = AsyncMock( + return_value=({"model": "gpt-4o-mini", "litellm_metadata": {}}, MagicMock()) + ) + failure = litellm.BadRequestError( + message="invalid_encrypted_content", model="gpt-4o-mini", llm_provider="openai" + ) + + async def fake_llm_call(): + return failure if provider_rejected else None + + proxy_logging_obj = MagicMock() + proxy_logging_obj.post_call_failure_hook = AsyncMock(return_value=None) + user_api_key_dict = MagicMock() + + with ( + patch( # test-quality-ok: first-frame model auth needs a live router and key table and has its own tests above + "litellm.proxy.response_api_endpoints.endpoints._enforce_responses_ws_first_frame_model_auth", + new_callable=AsyncMock, + ), + patch( # test-quality-ok: the pre-call processor needs a live proxy; what the endpoint does with the relay's outcome is under test + "litellm.proxy.response_api_endpoints.endpoints.ProxyBaseLLMRequestProcessing", + return_value=processor, + ), + patch( # test-quality-ok: routing is the seam that hands back the relay's outcome + "litellm.proxy.route_llm_request.route_request", + new_callable=AsyncMock, + return_value=fake_llm_call(), + ), + patch( # test-quality-ok: the failure hook is the proxy's only path to a failed spend log row + "litellm.proxy.proxy_server.proxy_logging_obj", + proxy_logging_obj, + ), + ): + await responses_websocket_endpoint( + websocket=ws, + model=None, + user_api_key_dict=user_api_key_dict, + ) + + ws.close.assert_not_awaited() + if not provider_rejected: + proxy_logging_obj.post_call_failure_hook.assert_not_awaited() + return + proxy_logging_obj.post_call_failure_hook.assert_awaited_once() + booked = proxy_logging_obj.post_call_failure_hook.await_args.kwargs + assert booked["original_exception"] is failure + assert booked["user_api_key_dict"] is user_api_key_dict + assert booked["request_data"]["model"] == "gpt-4o-mini" + @pytest.mark.asyncio async def test_reruns_model_auth_for_first_frame_model(self): from starlette.requests import Request diff --git a/tests/test_litellm/responses/test_responses_websocket_all_providers.py b/tests/test_litellm/responses/test_responses_websocket_all_providers.py index b671e60438e..43946c8907d 100644 --- a/tests/test_litellm/responses/test_responses_websocket_all_providers.py +++ b/tests/test_litellm/responses/test_responses_websocket_all_providers.py @@ -2894,3 +2894,115 @@ class TestNativeWebSocketEncryptedContentAffinity: assert response_cost == 0.01 logging_obj.dispatch_success_handlers.assert_not_awaited() logging_obj.dispatch_failure_handlers.assert_awaited_once() + + @pytest.mark.asyncio + async def test_bidirectional_forward_returns_the_provider_failure(self): + import asyncio + from unittest.mock import AsyncMock + + import websockets.exceptions # noqa: F401 (lazy submodule must be importable) + + backend_drained = asyncio.Event() + backend_events = [ + json.dumps({"type": "response.created", "response": {"id": "resp_1", "status": "in_progress"}}), + json.dumps( + { + "type": "error", + "status": 400, + "error": { + "type": "invalid_request_error", + "code": "invalid_encrypted_content", + "message": "could not be verified", + }, + } + ), + ] + + async def recv(decode=False): + if backend_events: + return backend_events.pop(0) + backend_drained.set() + raise Exception("stop") + + async def receive_text(): + await backend_drained.wait() + raise Exception("client gone") + + websocket = MagicMock() + websocket.send_text = AsyncMock() + websocket.receive_text = receive_text + backend_ws = MagicMock() + backend_ws.recv = recv + backend_ws.send = AsyncMock() + backend_ws.close = AsyncMock() + logging_obj = MagicMock() + logging_obj.dispatch_success_handlers = AsyncMock() + logging_obj.dispatch_failure_handlers = AsyncMock() + logging_obj._response_cost_calculator = MagicMock(return_value=0.0) + handler = _make_streaming( + websocket=websocket, + backend_ws=backend_ws, + logging_obj=logging_obj, + request_data={}, + authorized_model="gpt-5.6", + custom_llm_provider="openai", + ) + + failure = await handler.bidirectional_forward() + + assert isinstance(failure, Exception) + assert failure.status_code == 400 + assert "could not be verified" in str(failure) + + @pytest.mark.asyncio + async def test_bidirectional_forward_returns_none_after_a_completed_turn(self): + import asyncio + from unittest.mock import AsyncMock + + import websockets.exceptions # noqa: F401 (lazy submodule must be importable) + + backend_drained = asyncio.Event() + backend_events = [ + json.dumps( + { + "type": "response.completed", + "response": { + "id": "resp_1", + "status": "completed", + "output": [], + "usage": {"input_tokens": 10, "output_tokens": 5, "total_tokens": 15}, + }, + } + ), + ] + + async def recv(decode=False): + if backend_events: + return backend_events.pop(0) + backend_drained.set() + raise Exception("stop") + + async def receive_text(): + await backend_drained.wait() + raise Exception("client gone") + + websocket = MagicMock() + websocket.send_text = AsyncMock() + websocket.receive_text = receive_text + backend_ws = MagicMock() + backend_ws.recv = recv + backend_ws.send = AsyncMock() + backend_ws.close = AsyncMock() + logging_obj = MagicMock() + logging_obj.dispatch_success_handlers = AsyncMock() + logging_obj.dispatch_failure_handlers = AsyncMock() + handler = _make_streaming( + websocket=websocket, + backend_ws=backend_ws, + logging_obj=logging_obj, + request_data={}, + authorized_model="gpt-5.6", + custom_llm_provider="openai", + ) + + assert await handler.bidirectional_forward() is None From 37da5b6f4d8ebf6e33379f09e36fcccc9bf11c4f Mon Sep 17 00:00:00 2001 From: kerry Date: Sat, 19 Sep 2026 00:10:17 +0000 Subject: [PATCH 084/144] fix(bedrock): sign batch retrieve and cancel with deployment credentials when AWS_BEARER_TOKEN_BEDROCK is set Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/llms/bedrock/batches/handler.py | 10 ++++ .../llms/bedrock/batches/test_handler.py | 50 +++++++++++++++++++ 2 files changed, 60 insertions(+) diff --git a/litellm/llms/bedrock/batches/handler.py b/litellm/llms/bedrock/batches/handler.py index 6239973eb7c..fd4c3dc1659 100644 --- a/litellm/llms/bedrock/batches/handler.py +++ b/litellm/llms/bedrock/batches/handler.py @@ -10,6 +10,8 @@ from litellm.types.llms.bedrock import AwsAuthParams, AwsSessionTag from litellm.types.utils import LiteLLMBatch if TYPE_CHECKING: + from botocore.config import Config + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj # AWS Bedrock model-invocation-job statuses → OpenAI Batch statuses. @@ -31,6 +33,12 @@ _BEDROCK_MIJ_STATUS_TO_OPENAI: Final = { _CANCEL_IDEMPOTENT_STATUSES: Final = frozenset({"cancelling", "cancelled", "completed", "failed", "expired"}) +def _sigv4_config() -> "Config": + from botocore.config import Config + + return Config(signature_version="v4") + + def _extract_region_from_bedrock_arn(arn: str) -> str | None: """ARN shape: ``arn:aws:bedrock:::/``""" try: @@ -150,6 +158,7 @@ class BedrockBatchesHandler: aws_access_key_id=creds.access_key, aws_secret_access_key=creds.secret_key, aws_session_token=creds.token, + config=_sigv4_config(), ) def job_status() -> "LiteLLMBatch": @@ -309,6 +318,7 @@ class BedrockBatchesHandler: aws_access_key_id=creds.access_key, aws_secret_access_key=creds.secret_key, aws_session_token=creds.token, + config=_sigv4_config(), ) if logging_obj is not None: diff --git a/tests/test_litellm/llms/bedrock/batches/test_handler.py b/tests/test_litellm/llms/bedrock/batches/test_handler.py index 03daafcad72..056378f97c9 100644 --- a/tests/test_litellm/llms/bedrock/batches/test_handler.py +++ b/tests/test_litellm/llms/bedrock/batches/test_handler.py @@ -570,3 +570,53 @@ def test_cancel_batch_stops_and_polls_the_job_with_the_tagged_session(monkeypatc fake_bedrock.stop_model_invocation_job.assert_called_once_with(jobIdentifier=JOB_ARN) assert batch.status == "cancelled" assert [kwargs["aws_access_key_id"] for kwargs in bedrock_client_kwargs] == ["ASIABATCHCANCELTAGGED"] * 2 + + +def _sigv4_capture_send(sent_headers: list[dict[str, str]], body: dict): + import json + + from botocore.awsrequest import AWSResponse + + def send(_self, request): + sent_headers.append({k: v.decode() if isinstance(v, bytes) else v for k, v in request.headers.items()}) + raw = MagicMock() + raw.stream.return_value = iter([json.dumps(body, default=str).encode()]) + return AWSResponse(request.url, 200, {"content-type": "application/json"}, raw) + + return send + + +def test_retrieve_signs_with_deployment_credentials_when_env_bearer_token_is_set(monkeypatch): + """A proxy-wide AWS_BEARER_TOKEN_BEDROCK must not override the deployment's own SigV4 credentials.""" + monkeypatch.setenv("AWS_BEARER_TOKEN_BEDROCK", "env-bearer-token") + sent_headers: list[dict[str, str]] = [] + + with patch("botocore.httpsession.URLLib3Session.send", _sigv4_capture_send(sent_headers, _fake_boto3_response())): + batch = BedrockBatchesHandler._handle_model_invocation_job_status( + batch_id=JOB_ARN, + aws_access_key_id="AKIADEPLOYMENTKEY", + aws_secret_access_key="deployment-secret", + ) + + assert batch.status == "completed" + assert len(sent_headers) == 1 + assert sent_headers[0]["Authorization"].startswith("AWS4-HMAC-SHA256 Credential=AKIADEPLOYMENTKEY/") + + +def test_cancel_signs_with_deployment_credentials_when_env_bearer_token_is_set(monkeypatch): + monkeypatch.setenv("AWS_BEARER_TOKEN_BEDROCK", "env-bearer-token") + sent_headers: list[dict[str, str]] = [] + + with patch( + "botocore.httpsession.URLLib3Session.send", + _sigv4_capture_send(sent_headers, _fake_boto3_response(status="Stopped")), + ): + batch = BedrockBatchesHandler.cancel_batch( + batch_id=JOB_ARN, + aws_access_key_id="AKIADEPLOYMENTKEY", + aws_secret_access_key="deployment-secret", + ) + + assert batch.status == "cancelled" + assert len(sent_headers) == 2 + assert all(h["Authorization"].startswith("AWS4-HMAC-SHA256 Credential=AKIADEPLOYMENTKEY/") for h in sent_headers) From 83d89aa134bbeac391ce4a49dd63223f52b97019 Mon Sep 17 00:00:00 2001 From: kerry Date: Sat, 19 Sep 2026 00:08:54 +0000 Subject: [PATCH 085/144] test(integration): cover off-peak pricing on a live proxy Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/integration/_support/client.py | 4 +- tests/integration/contracts.json | 6 ++ .../pricing/test_off_peak_pricing.py | 82 +++++++++++++++++++ 3 files changed, 90 insertions(+), 2 deletions(-) create mode 100644 tests/integration/pricing/test_off_peak_pricing.py diff --git a/tests/integration/_support/client.py b/tests/integration/_support/client.py index 97522e5728c..9f1118ab1e3 100644 --- a/tests/integration/_support/client.py +++ b/tests/integration/_support/client.py @@ -161,7 +161,7 @@ class Scenario: assert all(object_value(object_value(entry)["model_info"])["id"] != identity for entry in entries) assert read_rows('SELECT model_id FROM "LiteLLM_ProxyModelTable" WHERE model_id = %s', (identity,)) == [] - def model(self, **parameters: JsonValue) -> str: + def model(self, *, model_info: Mapping[str, JsonValue] | None = None, **parameters: JsonValue) -> str: name: Final = f"integration-{uuid.uuid4().hex}" created: Final = self.gateway.post( "/model/new", @@ -173,7 +173,7 @@ class Scenario: "api_base": f"{self.gateway.upstream_url}/v1", **parameters, }, - "model_info": {}, + "model_info": dict(model_info) if model_info is not None else {}, }, ) identity: Final = string_value(object_value(created["model_info"])["id"]) diff --git a/tests/integration/contracts.json b/tests/integration/contracts.json index 91b1bd86954..6958ade50f7 100644 --- a/tests/integration/contracts.json +++ b/tests/integration/contracts.json @@ -92,6 +92,12 @@ "tests/integration/pricing/test_price_precedence.py::test_same_upstream_aliases_keep_distinct_prices_after_reload": [ "quota_management.spend_tracking.alias_prices.remain_independent_on_reload" ], + "tests/integration/pricing/test_off_peak_pricing.py::test_open_off_peak_window_bills_off_peak_rates": [ + "quota_management.spend_tracking.off_peak_pricing.open_window_bills_off_peak_rates" + ], + "tests/integration/pricing/test_off_peak_pricing.py::test_closed_off_peak_window_bills_standard_rates": [ + "quota_management.spend_tracking.off_peak_pricing.closed_window_bills_standard_rates" + ], "tests/integration/spend/test_cache_and_quota.py::test_generated_cache_sequences_preserve_content_usage_and_zero_hit_cost": [ "quota_management.response_cache.generated_sequences_preserve_content_and_accounting" ], diff --git a/tests/integration/pricing/test_off_peak_pricing.py b/tests/integration/pricing/test_off_peak_pricing.py new file mode 100644 index 00000000000..5623356c078 --- /dev/null +++ b/tests/integration/pricing/test_off_peak_pricing.py @@ -0,0 +1,82 @@ +import json +from collections.abc import Mapping +from datetime import datetime, timedelta, timezone +from typing import Final + +import pytest +from pydantic import JsonValue + +from tests.integration._support.client import Gateway, Scenario, eventually, object_value, string_value +from tests.integration._support.database import read_rows + +STANDARD_INPUT_RATE: Final = 0.001 +STANDARD_OUTPUT_RATE: Final = 0.002 +OFF_PEAK_INPUT_RATE: Final = 0.0001 +OFF_PEAK_OUTPUT_RATE: Final = 0.0002 + + +def off_peak_window(start_offset_hours: int, end_offset_hours: int) -> Mapping[str, JsonValue]: + now: Final = datetime.now(timezone.utc) + start: Final = now + timedelta(hours=start_offset_hours) + end: Final = now + timedelta(hours=end_offset_hours) + return { + "hours_utc": f"{start:%H:%M}-{end:%H:%M}", + "input_cost_per_token": OFF_PEAK_INPUT_RATE, + "output_cost_per_token": OFF_PEAK_OUTPUT_RATE, + } + + +def billed_model(scenario: Scenario, off_peak: Mapping[str, JsonValue]) -> str: + return scenario.model( + input_cost_per_token=STANDARD_INPUT_RATE, + output_cost_per_token=STANDARD_OUTPUT_RATE, + model_info={"off_peak_pricing": dict(off_peak)}, + ) + + +def assert_chat_bills_rates(gateway: Gateway, model: str, input_rate: float, output_rate: float) -> None: + response: Final = gateway.request( + "POST", "/v1/chat/completions", {"model": model, "messages": [{"role": "user", "content": "off peak control"}]} + ) + assert response.status_code == 200, response.text + expected: Final = 20 * input_rate + 20 * output_rate + assert float(response.headers["x-litellm-response-cost"]) == pytest.approx(expected, rel=1e-6) + request_id: Final = string_value(object_value(response.json())["id"]) + rows: Final = eventually( + lambda: read_rows( + 'SELECT spend, metadata, prompt_tokens, completion_tokens FROM "LiteLLM_SpendLogs" WHERE request_id = %s', + (request_id,), + ), + lambda values: len(values) == 1, + seconds=70, + ) + assert rows[0]["prompt_tokens"] == 20 + assert rows[0]["completion_tokens"] == 20 + assert float(rows[0]["spend"]) == pytest.approx(expected, rel=1e-6) + metadata: Final = rows[0]["metadata"] + parsed: Final = json.loads(metadata) if isinstance(metadata, str) else object_value(metadata) + breakdown: Final = object_value(parsed["cost_breakdown"]) + assert float(breakdown["input_cost"]) == pytest.approx(20 * input_rate, rel=1e-6) + assert float(breakdown["output_cost"]) == pytest.approx(20 * output_rate, rel=1e-6) + + +@pytest.mark.covers("quota_management.spend_tracking.off_peak_pricing.open_window_bills_off_peak_rates") +def test_open_off_peak_window_bills_off_peak_rates(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + model: Final = billed_model(scenario, off_peak_window(-1, 1)) + entries: Final = gateway.get("/model/info")["data"] + assert isinstance(entries, list) + matching: Final = tuple(object_value(entry) for entry in entries if object_value(entry)["model_name"] == model) + assert len(matching) == 1 + info: Final = object_value(matching[0]["model_info"]) + off_peak: Final = object_value(info["off_peak_pricing"]) + assert off_peak["input_cost_per_token"] == OFF_PEAK_INPUT_RATE + assert off_peak["output_cost_per_token"] == OFF_PEAK_OUTPUT_RATE + assert_chat_bills_rates(gateway, model, OFF_PEAK_INPUT_RATE, OFF_PEAK_OUTPUT_RATE) + + +@pytest.mark.covers("quota_management.spend_tracking.off_peak_pricing.closed_window_bills_standard_rates") +def test_closed_off_peak_window_bills_standard_rates(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + model: Final = billed_model(scenario, off_peak_window(2, 3)) + assert_chat_bills_rates(gateway, model, STANDARD_INPUT_RATE, STANDARD_OUTPUT_RATE) From 12120fe59bd9dd36486fa683f84b06fb91bd9c9f Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 18 Sep 2026 17:11:29 -0700 Subject: [PATCH 086/144] refactor(bedrock): inline maxTokens clamp and cover inference-profile ARNs in tests --- litellm/llms/bedrock/chat/converse_transformation.py | 10 ++-------- .../llms/bedrock/chat/test_converse_transformation.py | 3 +++ 2 files changed, 5 insertions(+), 8 deletions(-) diff --git a/litellm/llms/bedrock/chat/converse_transformation.py b/litellm/llms/bedrock/chat/converse_transformation.py index 801ec571376..176819c0dab 100644 --- a/litellm/llms/bedrock/chat/converse_transformation.py +++ b/litellm/llms/bedrock/chat/converse_transformation.py @@ -382,12 +382,6 @@ class AmazonConverseConfig(BaseConfig): def _requires_min_max_tokens(model: str) -> bool: return re.search(r"openai\.gpt-\d|xai\.grok-", model) is not None - @staticmethod - def _enforce_min_max_tokens(max_tokens: object) -> object: - if isinstance(max_tokens, int) and max_tokens < BEDROCK_OPENAI_COMPAT_MIN_MAX_TOKENS: - return BEDROCK_OPENAI_COMPAT_MIN_MAX_TOKENS - return max_tokens - def _is_nova_2_model(self, model: str) -> bool: """ Check if the model is a Nova 2 model that supports reasoningConfig. @@ -1011,8 +1005,8 @@ class AmazonConverseConfig(BaseConfig): ) if param == "max_tokens" or param == "max_completion_tokens": optional_params["maxTokens"] = ( - self._enforce_min_max_tokens(value) - if self._requires_min_max_tokens(model) and isinstance(value, int) + max(value, BEDROCK_OPENAI_COMPAT_MIN_MAX_TOKENS) + if isinstance(value, int) and self._requires_min_max_tokens(model) else value ) if param == "stream": diff --git a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py index 9d8bf786829..086a7e59f56 100644 --- a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py +++ b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py @@ -467,6 +467,9 @@ def test_reasoning_with_forced_tool_choice_switches_to_auto(): ("global.xai.grok-4.6", "max_completion_tokens", 1, 16), ("us.xai.grok-4.6", "max_tokens", 32, 32), ("anthropic.claude-sonnet-4-5-20250929-v1:0", "max_tokens", 1, 1), + ("arn:aws:bedrock:us-east-1:123456789012:inference-profile/us.openai.gpt-6-astra", "max_tokens", 1, 16), + ("arn:aws:bedrock:us-east-1:123456789012:inference-profile/global.xai.grok-4.6", "max_tokens", 1, 16), + ("arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/abc123xyz", "max_tokens", 1, 1), ], ) def test_map_openai_params_enforces_minimum_max_tokens_for_openai_compat_models( From c32309fb2de108768fa8704ee0696b29d383733c Mon Sep 17 00:00:00 2001 From: yassin Date: Sat, 19 Sep 2026 00:13:44 +0000 Subject: [PATCH 087/144] feat(ui): show MCP allowed clients as cards edited in a dialog Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../_components/MCPNetworkSettings.test.tsx | 114 +++++++++--- .../_components/MCPNetworkSettings.tsx | 167 ++++++++++++------ 2 files changed, 201 insertions(+), 80 deletions(-) diff --git a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPNetworkSettings.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPNetworkSettings.test.tsx index d27c18c5ae3..4cc87f1455c 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPNetworkSettings.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPNetworkSettings.test.tsx @@ -1,4 +1,4 @@ -import { fireEvent, render, screen, waitFor } from "@testing-library/react"; +import { fireEvent, render, screen, waitFor, within } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; import { describe, it, expect, vi, beforeEach } from "vitest"; import MCPNetworkSettings from "./MCPNetworkSettings"; @@ -26,12 +26,20 @@ const renderSettings = () => render(); const ANTIGRAVITY = { alias: "Antigravity CLI", value: "antigravity-cli" }; const CODEX = { alias: "Codex", value: "codex-mcp-client" }; +const clientCard = (alias: string) => screen.getByRole("button", { name: new RegExp(`^${alias}`) }); + +const fillClientDialog = async (alias: string, value: string) => { + const dialog = await screen.findByRole("dialog"); + fireEvent.change(within(dialog).getByRole("textbox", { name: "Alias" }), { target: { value: alias } }); + fireEvent.change(within(dialog).getByRole("textbox", { name: "Value" }), { target: { value } }); + return dialog; +}; + const addClient = async (alias: string, value: string) => { await userEvent.click(screen.getByRole("button", { name: "Add client" })); - const aliases = screen.getAllByRole("textbox", { name: /^Client \d+ alias$/ }); - const values = screen.getAllByRole("textbox", { name: /^Client \d+ value$/ }); - fireEvent.change(aliases[aliases.length - 1], { target: { value: alias } }); - fireEvent.change(values[values.length - 1], { target: { value } }); + const dialog = await fillClientDialog(alias, value); + await userEvent.click(within(dialog).getByRole("button", { name: "Add" })); + await waitFor(() => expect(screen.queryByRole("dialog")).not.toBeInTheDocument()); }; describe("MCPNetworkSettings", () => { @@ -139,7 +147,7 @@ describe("MCPNetworkSettings", () => { expect(updateConfigFieldSetting).not.toHaveBeenCalled(); }); - it("labels the section Allowed Clients and renders each stored client as an alias and value row", async () => { + it("labels the section Allowed Clients and renders each stored client as a card showing alias and value", async () => { vi.mocked(getGeneralSettingsCall).mockResolvedValue([ { field_name: "mcp_allowed_clients", field_value: [ANTIGRAVITY, CODEX] }, ]); @@ -148,10 +156,24 @@ describe("MCPNetworkSettings", () => { expect(await screen.findByText("Allowed Clients")).toBeVisible(); expect(screen.queryByText(/Allowed Client IDs/)).not.toBeInTheDocument(); - expect(screen.getByRole("textbox", { name: "Client 1 alias" })).toHaveValue("Antigravity CLI"); - expect(screen.getByRole("textbox", { name: "Client 1 value" })).toHaveValue("antigravity-cli"); - expect(screen.getByRole("textbox", { name: "Client 2 alias" })).toHaveValue("Codex"); - expect(screen.getByRole("textbox", { name: "Client 2 value" })).toHaveValue("codex-mcp-client"); + expect(screen.queryByText(/Allowed Client Applications/)).not.toBeInTheDocument(); + expect(clientCard("Antigravity CLI")).toHaveTextContent("antigravity-cli"); + expect(clientCard("Codex")).toHaveTextContent("codex-mcp-client"); + expect(screen.queryByRole("dialog")).not.toBeInTheDocument(); + }); + + it("opens an edit dialog when a client card is clicked, prefilled with that client's alias and value", async () => { + vi.mocked(getGeneralSettingsCall).mockResolvedValue([ + { field_name: "mcp_allowed_clients", field_value: [ANTIGRAVITY, CODEX] }, + ]); + + renderSettings(); + await screen.findByText("Allowed Clients"); + await userEvent.click(clientCard("Codex")); + + const dialog = await screen.findByRole("dialog", { name: "Edit client" }); + expect(within(dialog).getByRole("textbox", { name: "Alias" })).toHaveValue("Codex"); + expect(within(dialog).getByRole("textbox", { name: "Value" })).toHaveValue("codex-mcp-client"); }); it("warns that a stored allowlist in the old plain-string shape denies every client and lets Save remove it", async () => { @@ -162,7 +184,7 @@ describe("MCPNetworkSettings", () => { renderSettings(); expect(await screen.findByText(/stored allowlist is not a list of alias and value pairs/)).toBeVisible(); - expect(screen.queryByRole("textbox", { name: "Client 1 value" })).not.toBeInTheDocument(); + expect(screen.queryByText("antigravity-cli")).not.toBeInTheDocument(); await userEvent.click(screen.getByRole("button", { name: /Save/ })); @@ -203,15 +225,22 @@ describe("MCPNetworkSettings", () => { expect(deleteConfigFieldSetting).not.toHaveBeenCalledWith("tok", "mcp_allowed_clients"); }); - it("edits a stored client's value in place and saves the new value", async () => { + it("edits a stored client's value through its dialog and saves the new value", async () => { vi.mocked(getGeneralSettingsCall).mockResolvedValue([ { field_name: "mcp_allowed_clients", field_value: [ANTIGRAVITY] }, ]); renderSettings(); - fireEvent.change(await screen.findByRole("textbox", { name: "Client 1 value" }), { - target: { value: "0oa1b2c3d4e5f6g7h8i9" }, + await screen.findByText("Allowed Clients"); + await userEvent.click(clientCard("Antigravity CLI")); + const dialog = await screen.findByRole("dialog"); + fireEvent.change(within(dialog).getByRole("textbox", { name: "Value" }), { + target: { value: " 0oa1b2c3d4e5f6g7h8i9 " }, }); + await userEvent.click(within(dialog).getByRole("button", { name: "Done" })); + + await waitFor(() => expect(screen.queryByRole("dialog")).not.toBeInTheDocument()); + expect(clientCard("Antigravity CLI")).toHaveTextContent("0oa1b2c3d4e5f6g7h8i9"); await userEvent.click(screen.getByRole("button", { name: /Save/ })); await waitFor(() => @@ -221,22 +250,37 @@ describe("MCPNetworkSettings", () => { ); }); - it("refuses to save a client that has an alias but no value, and reports why", async () => { + it("keeps a stored client untouched when its dialog is cancelled", async () => { + vi.mocked(getGeneralSettingsCall).mockResolvedValue([ + { field_name: "mcp_allowed_clients", field_value: [ANTIGRAVITY] }, + ]); + + renderSettings(); + await screen.findByText("Allowed Clients"); + await userEvent.click(clientCard("Antigravity CLI")); + const dialog = await screen.findByRole("dialog"); + fireEvent.change(within(dialog).getByRole("textbox", { name: "Value" }), { target: { value: "changed" } }); + await userEvent.click(within(dialog).getByRole("button", { name: "Cancel" })); + + await waitFor(() => expect(screen.queryByRole("dialog")).not.toBeInTheDocument()); + expect(clientCard("Antigravity CLI")).toHaveTextContent("antigravity-cli"); + await userEvent.click(screen.getByRole("button", { name: /Save/ })); + + await waitFor(() => expect(toast.success).toHaveBeenCalledWith("MCP network settings saved")); + expect(updateConfigFieldSetting).not.toHaveBeenCalled(); + }); + + it("will not add a client that has an alias but no value", async () => { renderSettings(); await screen.findByText("Allowed Clients"); - await addClient("Antigravity CLI", ""); - await userEvent.click(screen.getByRole("button", { name: /Save/ })); + await userEvent.click(screen.getByRole("button", { name: "Add client" })); + const dialog = await fillClientDialog("Antigravity CLI", " "); - await waitFor(() => - expect(toast.fromError).toHaveBeenCalledWith(new Error("Every allowed client needs both an alias and a value")), - ); - expect(updateConfigFieldSetting).not.toHaveBeenCalled(); - expect(deleteConfigFieldSetting).not.toHaveBeenCalled(); - expect(toast.success).not.toHaveBeenCalled(); + expect(within(dialog).getByRole("button", { name: "Add" })).toBeDisabled(); }); - it("drops rows left completely blank instead of saving or failing on them", async () => { + it("adds nothing when the add dialog is cancelled", async () => { vi.mocked(getGeneralSettingsCall).mockResolvedValue([ { field_name: "mcp_allowed_clients", field_value: [ANTIGRAVITY] }, ]); @@ -244,6 +288,11 @@ describe("MCPNetworkSettings", () => { renderSettings(); await screen.findByText("Allowed Clients"); await userEvent.click(screen.getByRole("button", { name: "Add client" })); + const dialog = await fillClientDialog("Codex", "codex-mcp-client"); + await userEvent.click(within(dialog).getByRole("button", { name: "Cancel" })); + + await waitFor(() => expect(screen.queryByRole("dialog")).not.toBeInTheDocument()); + expect(screen.queryByText("Codex")).not.toBeInTheDocument(); await userEvent.click(screen.getByRole("button", { name: /Save/ })); await waitFor(() => expect(toast.success).toHaveBeenCalledWith("MCP network settings saved")); @@ -260,13 +309,17 @@ describe("MCPNetworkSettings", () => { ]); renderSettings(); - await userEvent.click(await screen.findByRole("button", { name: "Remove client Claude Code" })); + await screen.findByText("Allowed Clients"); + await userEvent.click(clientCard("Claude Code")); + await userEvent.click(within(await screen.findByRole("dialog")).getByRole("button", { name: "Remove client" })); + + await waitFor(() => expect(screen.queryByRole("dialog")).not.toBeInTheDocument()); + expect(screen.queryByText("claude-code")).not.toBeInTheDocument(); await userEvent.click(screen.getByRole("button", { name: /Save/ })); await waitFor(() => expect(updateConfigFieldSetting).toHaveBeenCalledWith("tok", "mcp_allowed_clients", [ANTIGRAVITY, CODEX]), ); - expect(screen.queryByDisplayValue("claude-code")).not.toBeInTheDocument(); }); it("removes a client and clears the setting when the list becomes empty", async () => { @@ -275,9 +328,12 @@ describe("MCPNetworkSettings", () => { ]); renderSettings(); - await userEvent.click(await screen.findByRole("button", { name: "Remove client Claude Code" })); + await screen.findByText("Allowed Clients"); + await userEvent.click(clientCard("Claude Code")); + await userEvent.click(within(await screen.findByRole("dialog")).getByRole("button", { name: "Remove client" })); - expect(screen.queryByDisplayValue("claude-code")).not.toBeInTheDocument(); + await waitFor(() => expect(screen.queryByRole("dialog")).not.toBeInTheDocument()); + expect(screen.queryByText("claude-code")).not.toBeInTheDocument(); await userEvent.click(screen.getByRole("button", { name: /Save/ })); @@ -304,7 +360,7 @@ describe("MCPNetworkSettings", () => { renderSettings(); - await screen.findByText("Allowed Client Applications"); + await screen.findByText("Allowed Clients"); expect(screen.queryByText(/every client is denied/)).not.toBeInTheDocument(); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPNetworkSettings.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPNetworkSettings.tsx index ae1fad36599..db45cfdedd1 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPNetworkSettings.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPNetworkSettings.tsx @@ -1,9 +1,18 @@ -import React, { useState, useEffect } from "react"; +import React, { useState, useEffect, useId } from "react"; import { Save, Plus, X } from "lucide-react"; import { Badge } from "@/components/ui/badge"; import { Button } from "@/components/ui/button"; import { Card } from "@/components/ui/card"; +import { + Dialog, + DialogContent, + DialogDescription, + DialogFooter, + DialogHeader, + DialogTitle, +} from "@/components/ui/dialog"; import { Input } from "@/components/ui/input"; +import { Label } from "@/components/ui/label"; import { UiLoadingSpinner } from "@/components/ui/ui-loading-spinner"; import { DeprecationBanner } from "@/components/DeprecationBanner"; import { toast } from "@/lib/toast"; @@ -36,6 +45,10 @@ interface AllowedClientRow extends AllowedClient { readonly key: string; } +interface ClientDraft extends AllowedClient { + readonly key: string | null; +} + const isAllowedClient = (entry: unknown): entry is AllowedClient => { if (typeof entry !== "object" || entry === null) return false; const { alias, value } = entry as Partial>; @@ -58,14 +71,10 @@ const parseStoredClients = (fieldValue: unknown): StoredAllowlist => { }; let nextRowKey = 0; -const newRow = (client: AllowedClient = { alias: "", value: "" }): AllowedClientRow => ({ - ...client, - key: `client-${nextRowKey++}`, -}); +const newRow = (client: AllowedClient): AllowedClientRow => ({ ...client, key: `client-${nextRowKey++}` }); const trimClient = ({ alias, value }: AllowedClient): AllowedClient => ({ alias: alias.trim(), value: value.trim() }); -const isBlank = ({ alias, value }: AllowedClient) => alias === "" && value === ""; const isIncomplete = ({ alias, value }: AllowedClient) => alias === "" || value === ""; const sameList = (a: string[], b: string[]) => a.length === b.length && a.every((value, i) => value === b[i]); @@ -90,6 +99,67 @@ const clientsUnchangedSinceLoad = (value: AllowedClient[], stored: StoredAllowli const headerUnchangedSinceLoad = (value: string, stored: string | null) => stored === null ? value === "" : value !== "" && value === stored; +interface AllowedClientDialogProps { + readonly draft: ClientDraft | null; + readonly onChange: (draft: ClientDraft) => void; + readonly onCommit: () => void; + readonly onRemove: () => void; + readonly onClose: () => void; +} + +const AllowedClientDialog: React.FC = ({ draft, onChange, onCommit, onRemove, onClose }) => { + const aliasId = useId(); + const valueId = useId(); + if (draft === null) return null; + return ( +

!open && onClose()}> + + + {draft.key === null ? "Add client" : "Edit client"} + + The alias is the name shown in the dashboard and gateway logs. The value is the exact JWT claim or header + value that identifies the client, such as the OAuth client ID your identity provider issues. + + +
+
+ + onChange({ ...draft, alias: e.target.value })} + /> +
+
+ + onChange({ ...draft, value: e.target.value })} + /> +
+
+ + {draft.key !== null && ( + + )} + + + + + + ); +}; + const MCPNetworkSettings: React.FC = ({ accessToken }) => { const [loading, setLoading] = useState(true); const [saving, setSaving] = useState(false); @@ -101,6 +171,7 @@ const MCPNetworkSettings: React.FC = ({ accessToken }) const [storedClientIdHeader, setStoredClientIdHeader] = useState(null); const [currentIp, setCurrentIp] = useState(null); const [rangeDraft, setRangeDraft] = useState(""); + const [clientDraft, setClientDraft] = useState(null); useEffect(() => { loadSettings(); @@ -154,10 +225,7 @@ const MCPNetworkSettings: React.FC = ({ accessToken }) }; const persistAllowedClients = async (token: string) => { - const clients = allowedClients.map(trimClient).filter((client) => !isBlank(client)); - if (clients.some(isIncomplete)) { - throw new Error("Every allowed client needs both an alias and a value"); - } + const clients = allowedClients.map(({ alias, value }) => ({ alias, value })); if (clientsUnchangedSinceLoad(clients, storedClients)) return; if (clients.length > 0) { await updateConfigFieldSetting(token, "mcp_allowed_clients", clients); @@ -218,10 +286,22 @@ const MCPNetworkSettings: React.FC = ({ accessToken }) setRangeDraft(""); }; - const updateClient = (key: string, patch: Partial) => - setAllowedClients(allowedClients.map((row) => (row.key === key ? { ...row, ...patch } : row))); + const commitClientDraft = () => { + if (clientDraft === null) return; + const client = trimClient(clientDraft); + setAllowedClients( + clientDraft.key === null + ? [...allowedClients, newRow(client)] + : allowedClients.map((row) => (row.key === clientDraft.key ? { ...row, ...client } : row)), + ); + setClientDraft(null); + }; - const removeClient = (key: string) => setAllowedClients(allowedClients.filter((row) => row.key !== key)); + const removeDraftedClient = () => { + if (clientDraft === null) return; + setAllowedClients(allowedClients.filter((row) => row.key !== clientDraft.key)); + setClientDraft(null); + }; if (loading) { return ( @@ -307,7 +387,7 @@ const MCPNetworkSettings: React.FC = ({ accessToken })
-

Allowed Client Applications

+

Allowed Clients

Only the MCP client applications listed here can use the gateway. Leave empty to allow every client. A client that authenticates with a JWT is identified by the claim named in litellm_jwtauth.mcp_client_id_jwt_field in @@ -317,9 +397,6 @@ const MCPNetworkSettings: React.FC = ({ accessToken })

-
-

Allowed Clients

-
{storedAllowlistIsMalformed && (

The stored allowlist is not a list of alias and value pairs, so every client is denied. Add the clients you @@ -333,35 +410,17 @@ const MCPNetworkSettings: React.FC = ({ accessToken })

)} {allowedClients.length > 0 && ( -
-

Alias

-

Value

- - {allowedClients.map((row, index) => ( - - updateClient(row.key, { alias: e.target.value })} - /> - updateClient(row.key, { value: e.target.value })} - /> - - +
+ {allowedClients.map((row) => ( + ))}
)} @@ -369,16 +428,14 @@ const MCPNetworkSettings: React.FC = ({ accessToken }) type="button" variant="outline" size="sm" - onClick={() => setAllowedClients([...allowedClients, newRow()])} + onClick={() => setClientDraft({ key: null, alias: "", value: "" })} > Add client

- The alias is the name shown here and in gateway logs. The value is the exact JWT claim or header value that - identifies the client, such as the OAuth client ID your identity provider issues. Leave the list empty to - allow every client. Every MCP request from an unlisted client, or from one with no resolvable identity, gets a - 403. + Click a client to edit or remove it. Leave the list empty to allow every client. Every MCP request from an + unlisted client, or from one with no resolvable identity, gets a 403.

@@ -403,6 +460,14 @@ const MCPNetworkSettings: React.FC = ({ accessToken }) Save
+ + setClientDraft(null)} + />
); }; From 92d82841fd6d00f309e8afcc7044938e598f25bf Mon Sep 17 00:00:00 2001 From: kerry Date: Sat, 19 Sep 2026 00:14:40 +0000 Subject: [PATCH 088/144] chore(model_info): backfill reseller Gemini entries from provider catalogs and prune retired ids Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ...odel_prices_and_context_window_backup.json | 268 ++++++++++++------ model_prices_and_context_window.json | 268 ++++++++++++------ 2 files changed, 370 insertions(+), 166 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 30b08e54410..0e63653e2f2 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -19244,7 +19244,20 @@ "supports_function_calling": true, "supports_anthropic_thinking_payload": true, "supports_prompt_caching": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_reasoning": true, + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text" + ], + "supports_vision": true, + "supports_audio_input": true, + "supports_video_input": true }, "databricks/databricks-gemini-2-5-pro": { "cache_creation_input_token_cost": 1.24999e-06, @@ -19265,7 +19278,21 @@ "supports_function_calling": true, "supports_anthropic_thinking_payload": true, "supports_prompt_caching": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "deprecation_date": "2026-10-02", + "supports_reasoning": true, + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text" + ], + "supports_vision": true, + "supports_audio_input": true, + "supports_video_input": true }, "databricks/databricks-gemini-3-1-flash-lite": { "cache_creation_input_token_cost": 3.1248e-07, @@ -19285,7 +19312,19 @@ "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_function_calling": true, "supports_prompt_caching": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text" + ], + "supports_vision": true, + "supports_audio_input": true, + "supports_video_input": true }, "databricks/databricks-gemini-3-1-flash-image": { "litellm_provider": "databricks", @@ -19347,7 +19386,20 @@ "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_function_calling": true, "supports_prompt_caching": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_reasoning": true, + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text" + ], + "supports_vision": true, + "supports_audio_input": true, + "supports_video_input": true }, "databricks/databricks-gemini-3-flash": { "cache_creation_input_token_cost": 6.2503e-07, @@ -19367,7 +19419,19 @@ "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_function_calling": true, "supports_prompt_caching": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text" + ], + "supports_vision": true, + "supports_audio_input": true, + "supports_video_input": true }, "databricks/databricks-gemini-3-pro": { "cache_creation_input_token_cost": 2.49998e-06, @@ -21433,7 +21497,11 @@ "mode": "chat", "supports_tool_choice": true, "supports_function_calling": true, - "supports_image_size": false + "supports_image_size": false, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_vision": true, + "supports_audio_input": true }, "deepinfra/google/gemini-2.5-pro": { "max_tokens": 1000000, @@ -21444,7 +21512,11 @@ "litellm_provider": "deepinfra", "mode": "chat", "supports_tool_choice": true, - "supports_function_calling": true + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_vision": true, + "supports_audio_input": true }, "deepinfra/google/gemma-3-12b-it": { "max_tokens": 131072, @@ -29340,26 +29412,6 @@ "supports_parallel_function_calling": true, "supports_vision": true }, - "github_copilot/gemini-2.5-pro": { - "litellm_provider": "github_copilot", - "max_input_tokens": 128000, - "max_output_tokens": 64000, - "max_tokens": 64000, - "mode": "chat", - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_vision": true - }, - "github_copilot/gemini-3-pro-preview": { - "litellm_provider": "github_copilot", - "max_input_tokens": 128000, - "max_output_tokens": 64000, - "max_tokens": 64000, - "mode": "chat", - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_vision": true - }, "github_copilot/gpt-3.5-turbo": { "litellm_provider": "github_copilot", "max_input_tokens": 16384, @@ -30014,17 +30066,6 @@ "output_cost_per_token": 8.8e-07, "supports_function_calling": true }, - "gmi/google/gemini-3-pro-preview": { - "input_cost_per_token": 2e-06, - "litellm_provider": "gmi", - "max_input_tokens": 1048576, - "max_output_tokens": 65536, - "max_tokens": 65536, - "mode": "chat", - "output_cost_per_token": 1.2e-05, - "supports_function_calling": true, - "supports_vision": true - }, "gmi/google/gemini-3-flash-preview": { "input_cost_per_token": 5e-07, "litellm_provider": "gmi", @@ -30034,7 +30075,8 @@ "mode": "chat", "output_cost_per_token": 3e-06, "supports_function_calling": true, - "supports_vision": true + "supports_vision": true, + "supports_system_messages": true }, "gmi/moonshotai/Kimi-K2-Thinking": { "input_cost_per_token": 8e-07, @@ -40163,7 +40205,12 @@ "supports_response_schema": true, "supports_vision": true, "supports_native_streaming": true, - "supports_image_size": false + "supports_image_size": false, + "supports_reasoning": true, + "supports_system_messages": true, + "supports_pdf_input": true, + "supports_audio_input": true, + "supports_video_input": true }, "oci/google.gemini-2.5-pro": { "input_cost_per_token": 1.25e-06, @@ -40177,7 +40224,12 @@ "supports_function_calling": true, "supports_response_schema": true, "supports_vision": true, - "supports_native_streaming": true + "supports_native_streaming": true, + "supports_reasoning": true, + "supports_system_messages": true, + "supports_pdf_input": true, + "supports_audio_input": true, + "supports_video_input": true }, "oci/google.gemini-2.5-flash-lite": { "input_cost_per_token": 7.5e-08, @@ -40192,7 +40244,12 @@ "supports_response_schema": true, "supports_vision": true, "supports_native_streaming": true, - "supports_image_size": false + "supports_image_size": false, + "supports_reasoning": false, + "supports_system_messages": true, + "supports_pdf_input": true, + "supports_audio_input": true, + "supports_video_input": true }, "oci/cohere.command-a-vision": { "input_cost_per_token": 1.56e-06, @@ -41435,7 +41492,7 @@ "max_tokens": 65535, "mode": "chat", "output_cost_per_token": 2.5e-06, - "supports_audio_output": true, + "supports_audio_output": false, "supports_function_calling": true, "supports_response_schema": true, "supports_system_messages": true, @@ -41449,7 +41506,8 @@ "supports_audio_input": true, "supports_pdf_input": true, "supports_reasoning": true, - "supports_web_search": false + "supports_web_search": false, + "supports_video_input": true }, "openrouter/google/gemini-2.5-pro": { "cache_creation_input_token_cost": 3.75e-07, @@ -41462,7 +41520,7 @@ "max_tokens": 65536, "mode": "chat", "output_cost_per_token": 1e-05, - "supports_audio_output": true, + "supports_audio_output": false, "supports_function_calling": true, "supports_response_schema": true, "supports_system_messages": true, @@ -41478,7 +41536,8 @@ "supports_audio_input": true, "supports_pdf_input": true, "supports_reasoning": true, - "supports_web_search": false + "supports_web_search": false, + "supports_video_input": true }, "openrouter/google/gemini-3-pro-preview": { "cache_read_input_token_cost": 2e-07, @@ -41563,7 +41622,8 @@ "supports_url_context": true, "supports_vision": true, "supports_web_search": false, - "tpm": 800000 + "tpm": 800000, + "supports_video_input": true }, "openrouter/google/gemini-3.1-flash-lite-preview": { "cache_creation_input_token_cost": 8.33333333333333e-08, @@ -41690,7 +41750,8 @@ "supports_system_messages": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": false, + "supports_video_input": true }, "openrouter/gryphe/mythomax-l2-13b": { "input_cost_per_token": 8e-08, @@ -44413,12 +44474,16 @@ "output_cost_per_token": 1.2e-05, "litellm_provider": "replicate", "mode": "chat", - "supports_function_calling": true, - "supports_parallel_function_calling": true, + "supports_function_calling": false, + "supports_parallel_function_calling": false, "supports_vision": true, "supports_system_messages": true, - "supports_tool_choice": true, - "supports_response_schema": true + "supports_tool_choice": false, + "supports_response_schema": false, + "input_cost_per_token_above_200k_tokens": 4e-06, + "output_cost_per_token_above_200k_tokens": 1.8e-05, + "supports_audio_input": true, + "supports_video_input": true }, "replicate/anthropic/claude-4.5-sonnet": { "input_cost_per_token": 3e-06, @@ -44487,17 +44552,19 @@ "supports_response_schema": true }, "replicate/google/gemini-2.5-flash": { - "input_cost_per_token": 2.5e-06, + "input_cost_per_token": 3e-07, "output_cost_per_token": 2.5e-06, "litellm_provider": "replicate", "mode": "chat", - "supports_function_calling": true, - "supports_parallel_function_calling": true, + "supports_function_calling": false, + "supports_parallel_function_calling": false, "supports_vision": true, "supports_system_messages": true, - "supports_tool_choice": true, - "supports_response_schema": true, - "supports_image_size": false + "supports_tool_choice": false, + "supports_response_schema": false, + "supports_image_size": false, + "supports_reasoning": true, + "supports_video_input": true }, "replicate/openai/gpt-oss-120b": { "input_cost_per_token": 1.8e-07, @@ -48076,10 +48143,15 @@ "supports_function_calling": true, "supports_tool_choice": true, "supports_response_schema": true, - "supports_image_size": false + "supports_image_size": false, + "cache_read_input_token_cost": 3e-08, + "supports_reasoning": true, + "supports_pdf_input": true, + "supports_web_search": true, + "supports_prompt_caching": true }, "vercel_ai_gateway/google/gemini-2.5-pro": { - "input_cost_per_token": 2.5e-06, + "input_cost_per_token": 1.25e-06, "litellm_provider": "vercel_ai_gateway", "max_input_tokens": 1048576, "max_output_tokens": 65536, @@ -48089,7 +48161,15 @@ "supports_vision": true, "supports_function_calling": true, "supports_tool_choice": true, - "supports_response_schema": true + "supports_response_schema": true, + "input_cost_per_token_above_200k_tokens": 2.5e-06, + "output_cost_per_token_above_200k_tokens": 1.5e-05, + "cache_read_input_token_cost": 1.25e-07, + "cache_read_input_token_cost_above_200k_tokens": 2.5e-07, + "supports_reasoning": true, + "supports_pdf_input": true, + "supports_web_search": true, + "supports_prompt_caching": true }, "vercel_ai_gateway/google/gemini-embedding-001": { "input_cost_per_token": 1.5e-07, @@ -62224,7 +62304,8 @@ "supports_response_schema": true, "supports_reasoning": true, "supports_vision": true, - "source": "https://deepinfra.com/pricing" + "source": "https://deepinfra.com/pricing", + "supports_audio_input": true }, "deepinfra/XiaomiMiMo/MiMo-V2.5": { "max_tokens": 262144, @@ -62524,7 +62605,8 @@ "supports_response_schema": true, "supports_reasoning": true, "supports_vision": true, - "source": "https://deepinfra.com/pricing" + "source": "https://deepinfra.com/pricing", + "supports_audio_input": true }, "deepinfra/google/gemini-3.7-flash": { "max_tokens": 1000000, @@ -62538,7 +62620,8 @@ "supports_response_schema": true, "supports_reasoning": true, "supports_vision": true, - "source": "https://deepinfra.com/pricing" + "source": "https://deepinfra.com/pricing", + "supports_audio_input": true }, "deepinfra/inclusionAI/Ling-3.0-flash": { "max_tokens": 131072, @@ -62970,7 +63053,8 @@ "supports_response_schema": true, "supports_reasoning": true, "supports_vision": true, - "source": "https://deepinfra.com/pricing" + "source": "https://deepinfra.com/pricing", + "supports_audio_input": true }, "deepinfra/XiaomiMiMo/MiMo-V2.5-Pro": { "max_tokens": 1048576, @@ -65365,7 +65449,8 @@ "deprecation_date": "2026-10-20", "input_cost_per_audio_token": 3e-07, "supports_prompt_caching": true, - "supports_web_search": false + "supports_web_search": false, + "supports_video_input": true }, "openrouter/google/gemini-3.5-flash": { "cache_creation_input_token_cost": 8.33333333333333e-08, @@ -65388,7 +65473,8 @@ "cache_read_input_token_cost": 1.5e-07, "input_cost_per_audio_token": 3e-06, "supports_prompt_caching": true, - "supports_web_search": false + "supports_web_search": false, + "supports_video_input": true }, "openrouter/google/gemini-3.5-flash-lite": { "cache_creation_input_token_cost": 8.33333333333333e-08, @@ -65411,7 +65497,8 @@ "cache_read_input_token_cost": 3e-08, "input_cost_per_audio_token": 3e-07, "supports_prompt_caching": true, - "supports_web_search": false + "supports_web_search": false, + "supports_video_input": true }, "openrouter/google/gemini-3.6-flash": { "cache_creation_input_token_cost": 4.16666666666667e-08, @@ -65434,7 +65521,8 @@ "cache_read_input_token_cost": 7.5e-08, "input_cost_per_audio_token": 7.5e-07, "supports_prompt_caching": true, - "supports_web_search": false + "supports_web_search": false, + "supports_video_input": true }, "openrouter/google/gemini-3.7-flash": { "cache_creation_input_token_cost": 4.16666666666667e-08, @@ -65457,7 +65545,8 @@ "cache_read_input_token_cost": 7.5e-08, "input_cost_per_audio_token": 7.5e-07, "supports_prompt_caching": true, - "supports_web_search": false + "supports_web_search": false, + "supports_video_input": true }, "openrouter/google/gemini-3.8-flash": { "cache_creation_input_token_cost": 4.16666666666667e-08, @@ -65480,7 +65569,8 @@ "cache_read_input_token_cost": 7.5e-08, "input_cost_per_audio_token": 7.5e-07, "supports_prompt_caching": true, - "supports_web_search": false + "supports_web_search": false, + "supports_video_input": true }, "openrouter/openai/gpt-4o-mini": { "input_cost_per_token": 1.5e-07, @@ -67108,7 +67198,8 @@ "supports_pdf_input": true, "supports_audio_input": true, "supports_prompt_caching": true, - "supports_web_search": false + "supports_web_search": false, + "supports_video_input": true }, "openrouter/qwen/qwen3-max-thinking": { "input_cost_per_token": 7.8e-07, @@ -71974,7 +72065,8 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": false, + "supports_video_input": true }, "openrouter/google/gemini-2.5-flash:batch": { "cache_read_input_audio_token_cost": 1e-07, @@ -71997,7 +72089,8 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": false, + "supports_video_input": true }, "openrouter/google/gemini-2.5-pro:batch": { "cache_read_input_audio_token_cost": 1.25e-07, @@ -72023,7 +72116,8 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": false, + "supports_video_input": true }, "openrouter/google/gemini-3-flash-preview:batch": { "input_cost_per_audio_token": 5e-07, @@ -72043,7 +72137,8 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": false, + "supports_video_input": true }, "openrouter/google/gemini-3.1-flash-lite:batch": { "cache_read_input_audio_token_cost": 2.5e-08, @@ -72065,7 +72160,8 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": false, + "supports_video_input": true }, "openrouter/google/gemini-3.1-pro-preview:batch": { "input_cost_per_audio_token": 1e-06, @@ -72087,7 +72183,8 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": false, + "supports_video_input": true }, "openrouter/google/gemini-3.5-flash-lite:batch": { "cache_read_input_audio_token_cost": 1.5e-08, @@ -72109,7 +72206,8 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": false, + "supports_video_input": true }, "openrouter/google/gemini-3.5-flash:batch": { "cache_read_input_audio_token_cost": 1.5e-07, @@ -72131,7 +72229,8 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": false, + "supports_video_input": true }, "openrouter/google/gemini-3.6-flash:batch": { "cache_creation_input_token_cost": 4.16666666666667e-08, @@ -72154,7 +72253,8 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": false, + "supports_video_input": true }, "openrouter/google/gemini-3.7-flash:batch": { "cache_creation_input_token_cost": 4.16666666666667e-08, @@ -72177,7 +72277,8 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": false, + "supports_video_input": true }, "openrouter/google/gemini-3.8-flash:batch": { "cache_creation_input_token_cost": 4.16666666666667e-08, @@ -72200,7 +72301,8 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": false, + "supports_video_input": true }, "openrouter/ibm-granite/granite-4.0-h-micro": { "input_cost_per_token": 1.7e-08, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 30b08e54410..0e63653e2f2 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -19244,7 +19244,20 @@ "supports_function_calling": true, "supports_anthropic_thinking_payload": true, "supports_prompt_caching": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_reasoning": true, + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text" + ], + "supports_vision": true, + "supports_audio_input": true, + "supports_video_input": true }, "databricks/databricks-gemini-2-5-pro": { "cache_creation_input_token_cost": 1.24999e-06, @@ -19265,7 +19278,21 @@ "supports_function_calling": true, "supports_anthropic_thinking_payload": true, "supports_prompt_caching": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "deprecation_date": "2026-10-02", + "supports_reasoning": true, + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text" + ], + "supports_vision": true, + "supports_audio_input": true, + "supports_video_input": true }, "databricks/databricks-gemini-3-1-flash-lite": { "cache_creation_input_token_cost": 3.1248e-07, @@ -19285,7 +19312,19 @@ "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_function_calling": true, "supports_prompt_caching": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text" + ], + "supports_vision": true, + "supports_audio_input": true, + "supports_video_input": true }, "databricks/databricks-gemini-3-1-flash-image": { "litellm_provider": "databricks", @@ -19347,7 +19386,20 @@ "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_function_calling": true, "supports_prompt_caching": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_reasoning": true, + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text" + ], + "supports_vision": true, + "supports_audio_input": true, + "supports_video_input": true }, "databricks/databricks-gemini-3-flash": { "cache_creation_input_token_cost": 6.2503e-07, @@ -19367,7 +19419,19 @@ "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_function_calling": true, "supports_prompt_caching": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text" + ], + "supports_vision": true, + "supports_audio_input": true, + "supports_video_input": true }, "databricks/databricks-gemini-3-pro": { "cache_creation_input_token_cost": 2.49998e-06, @@ -21433,7 +21497,11 @@ "mode": "chat", "supports_tool_choice": true, "supports_function_calling": true, - "supports_image_size": false + "supports_image_size": false, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_vision": true, + "supports_audio_input": true }, "deepinfra/google/gemini-2.5-pro": { "max_tokens": 1000000, @@ -21444,7 +21512,11 @@ "litellm_provider": "deepinfra", "mode": "chat", "supports_tool_choice": true, - "supports_function_calling": true + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_vision": true, + "supports_audio_input": true }, "deepinfra/google/gemma-3-12b-it": { "max_tokens": 131072, @@ -29340,26 +29412,6 @@ "supports_parallel_function_calling": true, "supports_vision": true }, - "github_copilot/gemini-2.5-pro": { - "litellm_provider": "github_copilot", - "max_input_tokens": 128000, - "max_output_tokens": 64000, - "max_tokens": 64000, - "mode": "chat", - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_vision": true - }, - "github_copilot/gemini-3-pro-preview": { - "litellm_provider": "github_copilot", - "max_input_tokens": 128000, - "max_output_tokens": 64000, - "max_tokens": 64000, - "mode": "chat", - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_vision": true - }, "github_copilot/gpt-3.5-turbo": { "litellm_provider": "github_copilot", "max_input_tokens": 16384, @@ -30014,17 +30066,6 @@ "output_cost_per_token": 8.8e-07, "supports_function_calling": true }, - "gmi/google/gemini-3-pro-preview": { - "input_cost_per_token": 2e-06, - "litellm_provider": "gmi", - "max_input_tokens": 1048576, - "max_output_tokens": 65536, - "max_tokens": 65536, - "mode": "chat", - "output_cost_per_token": 1.2e-05, - "supports_function_calling": true, - "supports_vision": true - }, "gmi/google/gemini-3-flash-preview": { "input_cost_per_token": 5e-07, "litellm_provider": "gmi", @@ -30034,7 +30075,8 @@ "mode": "chat", "output_cost_per_token": 3e-06, "supports_function_calling": true, - "supports_vision": true + "supports_vision": true, + "supports_system_messages": true }, "gmi/moonshotai/Kimi-K2-Thinking": { "input_cost_per_token": 8e-07, @@ -40163,7 +40205,12 @@ "supports_response_schema": true, "supports_vision": true, "supports_native_streaming": true, - "supports_image_size": false + "supports_image_size": false, + "supports_reasoning": true, + "supports_system_messages": true, + "supports_pdf_input": true, + "supports_audio_input": true, + "supports_video_input": true }, "oci/google.gemini-2.5-pro": { "input_cost_per_token": 1.25e-06, @@ -40177,7 +40224,12 @@ "supports_function_calling": true, "supports_response_schema": true, "supports_vision": true, - "supports_native_streaming": true + "supports_native_streaming": true, + "supports_reasoning": true, + "supports_system_messages": true, + "supports_pdf_input": true, + "supports_audio_input": true, + "supports_video_input": true }, "oci/google.gemini-2.5-flash-lite": { "input_cost_per_token": 7.5e-08, @@ -40192,7 +40244,12 @@ "supports_response_schema": true, "supports_vision": true, "supports_native_streaming": true, - "supports_image_size": false + "supports_image_size": false, + "supports_reasoning": false, + "supports_system_messages": true, + "supports_pdf_input": true, + "supports_audio_input": true, + "supports_video_input": true }, "oci/cohere.command-a-vision": { "input_cost_per_token": 1.56e-06, @@ -41435,7 +41492,7 @@ "max_tokens": 65535, "mode": "chat", "output_cost_per_token": 2.5e-06, - "supports_audio_output": true, + "supports_audio_output": false, "supports_function_calling": true, "supports_response_schema": true, "supports_system_messages": true, @@ -41449,7 +41506,8 @@ "supports_audio_input": true, "supports_pdf_input": true, "supports_reasoning": true, - "supports_web_search": false + "supports_web_search": false, + "supports_video_input": true }, "openrouter/google/gemini-2.5-pro": { "cache_creation_input_token_cost": 3.75e-07, @@ -41462,7 +41520,7 @@ "max_tokens": 65536, "mode": "chat", "output_cost_per_token": 1e-05, - "supports_audio_output": true, + "supports_audio_output": false, "supports_function_calling": true, "supports_response_schema": true, "supports_system_messages": true, @@ -41478,7 +41536,8 @@ "supports_audio_input": true, "supports_pdf_input": true, "supports_reasoning": true, - "supports_web_search": false + "supports_web_search": false, + "supports_video_input": true }, "openrouter/google/gemini-3-pro-preview": { "cache_read_input_token_cost": 2e-07, @@ -41563,7 +41622,8 @@ "supports_url_context": true, "supports_vision": true, "supports_web_search": false, - "tpm": 800000 + "tpm": 800000, + "supports_video_input": true }, "openrouter/google/gemini-3.1-flash-lite-preview": { "cache_creation_input_token_cost": 8.33333333333333e-08, @@ -41690,7 +41750,8 @@ "supports_system_messages": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": false, + "supports_video_input": true }, "openrouter/gryphe/mythomax-l2-13b": { "input_cost_per_token": 8e-08, @@ -44413,12 +44474,16 @@ "output_cost_per_token": 1.2e-05, "litellm_provider": "replicate", "mode": "chat", - "supports_function_calling": true, - "supports_parallel_function_calling": true, + "supports_function_calling": false, + "supports_parallel_function_calling": false, "supports_vision": true, "supports_system_messages": true, - "supports_tool_choice": true, - "supports_response_schema": true + "supports_tool_choice": false, + "supports_response_schema": false, + "input_cost_per_token_above_200k_tokens": 4e-06, + "output_cost_per_token_above_200k_tokens": 1.8e-05, + "supports_audio_input": true, + "supports_video_input": true }, "replicate/anthropic/claude-4.5-sonnet": { "input_cost_per_token": 3e-06, @@ -44487,17 +44552,19 @@ "supports_response_schema": true }, "replicate/google/gemini-2.5-flash": { - "input_cost_per_token": 2.5e-06, + "input_cost_per_token": 3e-07, "output_cost_per_token": 2.5e-06, "litellm_provider": "replicate", "mode": "chat", - "supports_function_calling": true, - "supports_parallel_function_calling": true, + "supports_function_calling": false, + "supports_parallel_function_calling": false, "supports_vision": true, "supports_system_messages": true, - "supports_tool_choice": true, - "supports_response_schema": true, - "supports_image_size": false + "supports_tool_choice": false, + "supports_response_schema": false, + "supports_image_size": false, + "supports_reasoning": true, + "supports_video_input": true }, "replicate/openai/gpt-oss-120b": { "input_cost_per_token": 1.8e-07, @@ -48076,10 +48143,15 @@ "supports_function_calling": true, "supports_tool_choice": true, "supports_response_schema": true, - "supports_image_size": false + "supports_image_size": false, + "cache_read_input_token_cost": 3e-08, + "supports_reasoning": true, + "supports_pdf_input": true, + "supports_web_search": true, + "supports_prompt_caching": true }, "vercel_ai_gateway/google/gemini-2.5-pro": { - "input_cost_per_token": 2.5e-06, + "input_cost_per_token": 1.25e-06, "litellm_provider": "vercel_ai_gateway", "max_input_tokens": 1048576, "max_output_tokens": 65536, @@ -48089,7 +48161,15 @@ "supports_vision": true, "supports_function_calling": true, "supports_tool_choice": true, - "supports_response_schema": true + "supports_response_schema": true, + "input_cost_per_token_above_200k_tokens": 2.5e-06, + "output_cost_per_token_above_200k_tokens": 1.5e-05, + "cache_read_input_token_cost": 1.25e-07, + "cache_read_input_token_cost_above_200k_tokens": 2.5e-07, + "supports_reasoning": true, + "supports_pdf_input": true, + "supports_web_search": true, + "supports_prompt_caching": true }, "vercel_ai_gateway/google/gemini-embedding-001": { "input_cost_per_token": 1.5e-07, @@ -62224,7 +62304,8 @@ "supports_response_schema": true, "supports_reasoning": true, "supports_vision": true, - "source": "https://deepinfra.com/pricing" + "source": "https://deepinfra.com/pricing", + "supports_audio_input": true }, "deepinfra/XiaomiMiMo/MiMo-V2.5": { "max_tokens": 262144, @@ -62524,7 +62605,8 @@ "supports_response_schema": true, "supports_reasoning": true, "supports_vision": true, - "source": "https://deepinfra.com/pricing" + "source": "https://deepinfra.com/pricing", + "supports_audio_input": true }, "deepinfra/google/gemini-3.7-flash": { "max_tokens": 1000000, @@ -62538,7 +62620,8 @@ "supports_response_schema": true, "supports_reasoning": true, "supports_vision": true, - "source": "https://deepinfra.com/pricing" + "source": "https://deepinfra.com/pricing", + "supports_audio_input": true }, "deepinfra/inclusionAI/Ling-3.0-flash": { "max_tokens": 131072, @@ -62970,7 +63053,8 @@ "supports_response_schema": true, "supports_reasoning": true, "supports_vision": true, - "source": "https://deepinfra.com/pricing" + "source": "https://deepinfra.com/pricing", + "supports_audio_input": true }, "deepinfra/XiaomiMiMo/MiMo-V2.5-Pro": { "max_tokens": 1048576, @@ -65365,7 +65449,8 @@ "deprecation_date": "2026-10-20", "input_cost_per_audio_token": 3e-07, "supports_prompt_caching": true, - "supports_web_search": false + "supports_web_search": false, + "supports_video_input": true }, "openrouter/google/gemini-3.5-flash": { "cache_creation_input_token_cost": 8.33333333333333e-08, @@ -65388,7 +65473,8 @@ "cache_read_input_token_cost": 1.5e-07, "input_cost_per_audio_token": 3e-06, "supports_prompt_caching": true, - "supports_web_search": false + "supports_web_search": false, + "supports_video_input": true }, "openrouter/google/gemini-3.5-flash-lite": { "cache_creation_input_token_cost": 8.33333333333333e-08, @@ -65411,7 +65497,8 @@ "cache_read_input_token_cost": 3e-08, "input_cost_per_audio_token": 3e-07, "supports_prompt_caching": true, - "supports_web_search": false + "supports_web_search": false, + "supports_video_input": true }, "openrouter/google/gemini-3.6-flash": { "cache_creation_input_token_cost": 4.16666666666667e-08, @@ -65434,7 +65521,8 @@ "cache_read_input_token_cost": 7.5e-08, "input_cost_per_audio_token": 7.5e-07, "supports_prompt_caching": true, - "supports_web_search": false + "supports_web_search": false, + "supports_video_input": true }, "openrouter/google/gemini-3.7-flash": { "cache_creation_input_token_cost": 4.16666666666667e-08, @@ -65457,7 +65545,8 @@ "cache_read_input_token_cost": 7.5e-08, "input_cost_per_audio_token": 7.5e-07, "supports_prompt_caching": true, - "supports_web_search": false + "supports_web_search": false, + "supports_video_input": true }, "openrouter/google/gemini-3.8-flash": { "cache_creation_input_token_cost": 4.16666666666667e-08, @@ -65480,7 +65569,8 @@ "cache_read_input_token_cost": 7.5e-08, "input_cost_per_audio_token": 7.5e-07, "supports_prompt_caching": true, - "supports_web_search": false + "supports_web_search": false, + "supports_video_input": true }, "openrouter/openai/gpt-4o-mini": { "input_cost_per_token": 1.5e-07, @@ -67108,7 +67198,8 @@ "supports_pdf_input": true, "supports_audio_input": true, "supports_prompt_caching": true, - "supports_web_search": false + "supports_web_search": false, + "supports_video_input": true }, "openrouter/qwen/qwen3-max-thinking": { "input_cost_per_token": 7.8e-07, @@ -71974,7 +72065,8 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": false, + "supports_video_input": true }, "openrouter/google/gemini-2.5-flash:batch": { "cache_read_input_audio_token_cost": 1e-07, @@ -71997,7 +72089,8 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": false, + "supports_video_input": true }, "openrouter/google/gemini-2.5-pro:batch": { "cache_read_input_audio_token_cost": 1.25e-07, @@ -72023,7 +72116,8 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": false, + "supports_video_input": true }, "openrouter/google/gemini-3-flash-preview:batch": { "input_cost_per_audio_token": 5e-07, @@ -72043,7 +72137,8 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": false, + "supports_video_input": true }, "openrouter/google/gemini-3.1-flash-lite:batch": { "cache_read_input_audio_token_cost": 2.5e-08, @@ -72065,7 +72160,8 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": false, + "supports_video_input": true }, "openrouter/google/gemini-3.1-pro-preview:batch": { "input_cost_per_audio_token": 1e-06, @@ -72087,7 +72183,8 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": false, + "supports_video_input": true }, "openrouter/google/gemini-3.5-flash-lite:batch": { "cache_read_input_audio_token_cost": 1.5e-08, @@ -72109,7 +72206,8 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": false, + "supports_video_input": true }, "openrouter/google/gemini-3.5-flash:batch": { "cache_read_input_audio_token_cost": 1.5e-07, @@ -72131,7 +72229,8 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": false, + "supports_video_input": true }, "openrouter/google/gemini-3.6-flash:batch": { "cache_creation_input_token_cost": 4.16666666666667e-08, @@ -72154,7 +72253,8 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": false, + "supports_video_input": true }, "openrouter/google/gemini-3.7-flash:batch": { "cache_creation_input_token_cost": 4.16666666666667e-08, @@ -72177,7 +72277,8 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": false, + "supports_video_input": true }, "openrouter/google/gemini-3.8-flash:batch": { "cache_creation_input_token_cost": 4.16666666666667e-08, @@ -72200,7 +72301,8 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": false, + "supports_video_input": true }, "openrouter/ibm-granite/granite-4.0-h-micro": { "input_cost_per_token": 1.7e-08, From 1961cbcb6c9b2f82ac7379ca45274e3cce855923 Mon Sep 17 00:00:00 2001 From: kerry Date: Sat, 19 Sep 2026 00:18:13 +0000 Subject: [PATCH 089/144] fix(timing): subtract every provider attempt from receive-anchored overhead Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../llm_response_utils/response_metadata.py | 31 +++++++-- litellm/litellm_core_utils/logging_utils.py | 10 ++- .../test_response_metadata.py | 56 ++++++++++++---- .../litellm_core_utils/test_logging_utils.py | 28 ++++++-- .../test_router_retry_non_retryable_errors.py | 65 +++++++++++++++++++ 5 files changed, 168 insertions(+), 22 deletions(-) diff --git a/litellm/litellm_core_utils/llm_response_utils/response_metadata.py b/litellm/litellm_core_utils/llm_response_utils/response_metadata.py index 9a007489473..3778ae1281f 100644 --- a/litellm/litellm_core_utils/llm_response_utils/response_metadata.py +++ b/litellm/litellm_core_utils/llm_response_utils/response_metadata.py @@ -1,6 +1,6 @@ import datetime from collections.abc import Mapping -from typing import Any, Final +from typing import Any, Final, cast import httpx @@ -16,9 +16,13 @@ from litellm.types.utils import ( ) -def _timing_window_start(start_time: datetime.datetime, logging_obj: LiteLLMLoggingObject) -> datetime.datetime: +def _timing_window_start( + start_time: datetime.datetime, logging_obj: LiteLLMLoggingObject +) -> tuple[datetime.datetime, bool]: received_at: Final = get_litellm_metadata_from_kwargs(logging_obj.model_call_details).get("litellm_received_at") - return received_at if isinstance(received_at, datetime.datetime) else start_time + if isinstance(received_at, datetime.datetime): + return received_at, True + return start_time, False def response_timing_metrics( @@ -33,7 +37,9 @@ def response_timing_metrics( the provider call (``llm_api_duration_ms``). It is omitted when neither duration was recorded, and when ``include_overhead`` is False because the two durations cover different windows. """ - window_start: Final = _timing_window_start(start_time, logging_obj) + timing_window: Final = _timing_window_start(start_time, logging_obj) + window_start: Final = timing_window[0] + receive_anchored: Final = timing_window[1] total_response_time_ms: Final = (end_time.timestamp() - window_start.timestamp()) * 1000 if not include_overhead: return {"_response_ms": total_response_time_ms} # mutable-ok: read-only timing result @@ -43,11 +49,26 @@ def response_timing_metrics( if caching_details is not None and caching_details.get("cache_hit") is True else None ) + metadata_value: Final = get_litellm_metadata_from_kwargs(logging_obj.model_call_details) + metadata: Final = cast(dict[str, object], metadata_value) if isinstance(metadata_value, dict) else {} llm_api_duration_ms: Final = logging_obj.model_call_details.get("llm_api_duration_ms") if cache_duration_ms is not None: overhead_ms: float | None = total_response_time_ms - cache_duration_ms elif llm_api_duration_ms is not None: - overhead_ms = round(total_response_time_ms - llm_api_duration_ms, 4) + total_provider_duration_ms: Final = metadata.get("llm_api_duration_ms_total") + provider_duration_ms: Final = ( + total_provider_duration_ms + if receive_anchored + and isinstance(total_provider_duration_ms, float) + and isinstance(llm_api_duration_ms, (int, float)) + and total_provider_duration_ms >= llm_api_duration_ms + else llm_api_duration_ms + ) + overhead_ms = ( + round(total_response_time_ms - provider_duration_ms, 4) + if isinstance(provider_duration_ms, (int, float)) + else None + ) else: overhead_ms = None if overhead_ms is None: diff --git a/litellm/litellm_core_utils/logging_utils.py b/litellm/litellm_core_utils/logging_utils.py index 0f14b461d3d..82bfb0efdb1 100644 --- a/litellm/litellm_core_utils/logging_utils.py +++ b/litellm/litellm_core_utils/logging_utils.py @@ -5,13 +5,14 @@ import re import time from collections.abc import Iterator, Mapping, Sequence from datetime import datetime -from typing import TYPE_CHECKING, Any, Final +from typing import TYPE_CHECKING, Any, Final, cast from litellm._logging import format_base64_size, verbose_logger from litellm.constants import ( BASE64_TRUNCATION_OFFLOAD_THRESHOLD_CHARS, MAX_BASE64_LENGTH_FOR_LOGGING, ) +from litellm.litellm_core_utils.core_helpers import get_litellm_metadata_from_kwargs from litellm.types.utils import ( ModelResponse, ModelResponseStream, @@ -286,6 +287,13 @@ def _set_duration_in_model_call_details( duration_ms: Final = (end_time - start_time).total_seconds() * 1000 if logging_obj and hasattr(logging_obj, "model_call_details"): logging_obj.model_call_details["llm_api_duration_ms"] = duration_ms + metadata_value: Final = get_litellm_metadata_from_kwargs(logging_obj.model_call_details) + if isinstance(metadata_value, dict): + metadata: Final = cast(dict[str, object], metadata_value) + existing_total: Final = metadata.get("llm_api_duration_ms_total") + metadata["llm_api_duration_ms_total"] = ( + existing_total if isinstance(existing_total, float) else 0.0 + ) + duration_ms else: verbose_logger.debug("`logging_obj` not found - unable to track `llm_api_duration_ms") except Exception as e: diff --git a/tests/test_litellm/litellm_core_utils/llm_response_utils/test_response_metadata.py b/tests/test_litellm/litellm_core_utils/llm_response_utils/test_response_metadata.py index eeccbc719d3..832be1a12d9 100644 --- a/tests/test_litellm/litellm_core_utils/llm_response_utils/test_response_metadata.py +++ b/tests/test_litellm/litellm_core_utils/llm_response_utils/test_response_metadata.py @@ -72,9 +72,7 @@ class TestCallbackDurationMs: def test_update_response_metadata_includes_callback_duration(self): """End-to-end: update_response_metadata should propagate callback_duration_ms.""" result = ModelResponse() - logging_obj = self._make_logging_obj( - callback_duration_ms=5.5, llm_api_duration_ms=800.0 - ) + logging_obj = self._make_logging_obj(callback_duration_ms=5.5, llm_api_duration_ms=800.0) logging_obj._response_cost_calculator = MagicMock(return_value=0.001) logging_obj.litellm_call_id = "test-call-id" @@ -236,6 +234,7 @@ class TestResponseTimingMetrics: def _make_logging_obj( self, llm_api_duration_ms: float | None = None, + llm_api_duration_ms_total: float | None = None, caching_details: dict[str, object] | None = None, received_at: datetime.datetime | str | None = None, ) -> MagicMock: @@ -243,8 +242,13 @@ class TestResponseTimingMetrics: logging_obj.model_call_details = {} if llm_api_duration_ms is not None: logging_obj.model_call_details["llm_api_duration_ms"] = llm_api_duration_ms - if received_at is not None: - logging_obj.model_call_details["litellm_params"] = {"metadata": {"litellm_received_at": received_at}} + if received_at is not None or llm_api_duration_ms_total is not None: + metadata = {} + if received_at is not None: + metadata["litellm_received_at"] = received_at + if llm_api_duration_ms_total is not None: + metadata["llm_api_duration_ms_total"] = llm_api_duration_ms_total + logging_obj.model_call_details["litellm_params"] = {"metadata": metadata} logging_obj.caching_details = caching_details return logging_obj @@ -264,6 +268,40 @@ class TestResponseTimingMetrics: assert result["_response_ms"] == pytest.approx(4000.0) assert result["litellm_overhead_time_ms"] == pytest.approx(3100.0) + def test_receive_anchored_window_subtracts_all_provider_attempts(self): + logging_obj = self._make_logging_obj( + llm_api_duration_ms=300.0, + llm_api_duration_ms_total=700.0, + received_at=self.START, + ) + + result = response_timing_metrics(self.START, self.END, logging_obj) + + assert result["_response_ms"] == pytest.approx(1000.0) + assert result["litellm_overhead_time_ms"] == pytest.approx(300.0) + + def test_sdk_window_subtracts_current_provider_attempt(self): + logging_obj = self._make_logging_obj( + llm_api_duration_ms=300.0, + llm_api_duration_ms_total=700.0, + ) + + result = response_timing_metrics(self.START, self.END, logging_obj) + + assert result["_response_ms"] == pytest.approx(1000.0) + assert result["litellm_overhead_time_ms"] == pytest.approx(700.0) + + def test_receive_anchored_window_falls_back_to_current_provider_attempt(self): + logging_obj = self._make_logging_obj( + llm_api_duration_ms=300.0, + received_at=self.START, + ) + + result = response_timing_metrics(self.START, self.END, logging_obj) + + assert result["_response_ms"] == pytest.approx(1000.0) + assert result["litellm_overhead_time_ms"] == pytest.approx(700.0) + def test_cache_hit_window_starts_at_proxy_receive_when_stamped(self): received_at = self.START.astimezone(datetime.timezone.utc) - datetime.timedelta(seconds=3) logging_obj = self._make_logging_obj( @@ -415,9 +453,7 @@ class TestDetailedTiming: def test_detailed_timing_headers_in_custom_headers(self, monkeypatch): """When LITELLM_DETAILED_TIMING is true, headers flow to get_custom_headers.""" - monkeypatch.setattr( - common_request_processing_mod, "LITELLM_DETAILED_TIMING", True - ) + monkeypatch.setattr(common_request_processing_mod, "LITELLM_DETAILED_TIMING", True) user_api_key_dict = UserAPIKeyAuth(api_key="sk-test") hidden_params = { @@ -440,9 +476,7 @@ class TestDetailedTiming: def test_detailed_timing_headers_absent_when_disabled(self, monkeypatch): """When LITELLM_DETAILED_TIMING is false, no timing headers emitted.""" - monkeypatch.setattr( - common_request_processing_mod, "LITELLM_DETAILED_TIMING", False - ) + monkeypatch.setattr(common_request_processing_mod, "LITELLM_DETAILED_TIMING", False) user_api_key_dict = UserAPIKeyAuth(api_key="sk-test") hidden_params = { diff --git a/tests/test_litellm/litellm_core_utils/test_logging_utils.py b/tests/test_litellm/litellm_core_utils/test_logging_utils.py index b446021a7dc..f669ff86c13 100644 --- a/tests/test_litellm/litellm_core_utils/test_logging_utils.py +++ b/tests/test_litellm/litellm_core_utils/test_logging_utils.py @@ -2,18 +2,39 @@ Tests for litellm.litellm_core_utils.logging_utils — base64 truncation helpers. """ +import datetime import threading +from unittest.mock import MagicMock import pytest from litellm.litellm_core_utils import logging_utils from litellm.litellm_core_utils.logging_utils import ( - format_base64_size, + _set_duration_in_model_call_details, _truncate_base64_in_string, + format_base64_size, truncate_base64_in_messages, truncate_base64_in_messages_async, ) + +class TestSetDurationInModelCallDetails: + def test_accumulates_provider_attempts_in_shared_metadata(self): + metadata = {"request_id": "test"} + logging_obj = MagicMock() + logging_obj.model_call_details = {"litellm_params": {"metadata": metadata}} + first_start = datetime.datetime(2025, 1, 1, 0, 0, 0) + first_end = first_start + datetime.timedelta(milliseconds=300) + second_start = datetime.datetime(2025, 1, 1, 0, 0, 1) + second_end = second_start + datetime.timedelta(milliseconds=700) + + _set_duration_in_model_call_details(logging_obj, first_start, first_end) + _set_duration_in_model_call_details(logging_obj, second_start, second_end) + + assert metadata["llm_api_duration_ms_total"] == pytest.approx(1000.0) + assert logging_obj.model_call_details["llm_api_duration_ms"] == pytest.approx(700.0) + + # --------------------------------------------------------------------------- # format_base64_size # --------------------------------------------------------------------------- @@ -157,10 +178,7 @@ class TestTruncateBase64InMessages: } ] result = truncate_base64_in_messages(messages) - assert ( - result[0]["content"][0]["image_url"]["url"] - == f"data:image/png;base64,{short}" - ) + assert result[0]["content"][0]["image_url"]["url"] == f"data:image/png;base64,{short}" # --------------------------------------------------------------------------- diff --git a/tests/test_litellm/test_router_retry_non_retryable_errors.py b/tests/test_litellm/test_router_retry_non_retryable_errors.py index 0728947eafe..c797f0f96a6 100644 --- a/tests/test_litellm/test_router_retry_non_retryable_errors.py +++ b/tests/test_litellm/test_router_retry_non_retryable_errors.py @@ -10,12 +10,20 @@ Verifies that: Regression tests for https://github.com/BerriAI/litellm/issues/21343 """ +import asyncio +import datetime +from collections.abc import Awaitable, Callable +from typing import Final, cast from unittest.mock import AsyncMock, patch import pytest import litellm from litellm import Router +from litellm.litellm_core_utils.litellm_logging import Logging +from litellm.litellm_core_utils.logging_utils import track_llm_api_timing +from litellm.litellm_core_utils.rules import Rules +from litellm.utils import function_setup def _make_rate_limit_error(message="Rate limited"): @@ -274,3 +282,60 @@ async def test_not_found_error_in_retry_loop_raises_immediately(): # Only 2 calls: initial + first retry that hits non-retryable assert call_count == 2 + + +@pytest.mark.asyncio +async def test_retry_attempts_accumulate_timing_in_shared_request_metadata(): + metadata: dict[str, object] = {"model_group": "test-model"} + logging_obj_raw, _ = function_setup( + "acompletion", + Rules(), + datetime.datetime.now(), + model="test-model", + messages=[{"role": "user", "content": "test"}], + metadata=metadata, + litellm_call_id="retry-timing-test", + is_async_call=True, + ) + logging_obj: Final[Logging] = cast(Logging, logging_obj_raw) + attempt_numbers: list[int] = [] + metadata_ids: list[int] = [] + + @track_llm_api_timing() + async def timed_attempt(*, logging_obj: Logging, **kwargs: object) -> str: + del kwargs + attempt_numbers.append(len(attempt_numbers) + 1) + metadata_ids.append(id(logging_obj.model_call_details["litellm_params"]["metadata"])) + await asyncio.sleep(0.01) + if len(attempt_numbers) == 1: + raise _make_rate_limit_error() + return "success" + + async def invoke(original_function: Callable[..., Awaitable[str]], *args: object, **kwargs: object) -> str: + return await original_function(*args, **kwargs) + + router = _create_router(num_retries=1) + with ( + patch.object(router, "make_call", new=AsyncMock(side_effect=invoke)), + patch.object( + router, + "_async_get_healthy_deployments", + new=AsyncMock(return_value=(["d1"], ["d1"])), + ), + patch.object(router, "_time_to_sleep_before_retry", return_value=0), + ): + result = await router.async_function_with_retries( + original_function=timed_attempt, + model="test-model", + messages=[{"role": "user", "content": "test"}], + metadata=metadata, + logging_obj=logging_obj, + num_retries=1, + ) + + request_metadata: Final = logging_obj.model_call_details["litellm_params"]["metadata"] + assert result == "success" + assert attempt_numbers == [1, 2] + assert request_metadata is metadata + assert metadata_ids == [id(metadata), id(metadata)] + assert request_metadata["llm_api_duration_ms_total"] > logging_obj.model_call_details["llm_api_duration_ms"] From e2141da81ea059c6946c7c7c674babd7ef62443a Mon Sep 17 00:00:00 2001 From: yassin Date: Sat, 19 Sep 2026 00:23:11 +0000 Subject: [PATCH 090/144] fix(ui): treat MCP allowed clients with an empty alias or value as malformed Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../_components/MCPNetworkSettings.test.tsx | 11 +++++++++++ .../mcp-servers/_components/MCPNetworkSettings.tsx | 2 +- 2 files changed, 12 insertions(+), 1 deletion(-) diff --git a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPNetworkSettings.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPNetworkSettings.test.tsx index 4cc87f1455c..4a486bfc648 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPNetworkSettings.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPNetworkSettings.test.tsx @@ -211,6 +211,17 @@ describe("MCPNetworkSettings", () => { await waitFor(() => expect(screen.queryByText(/stored allowlist is not a list/)).not.toBeInTheDocument()); }); + it("treats a stored entry with an empty alias or value as denying every client, like the gateway does", async () => { + vi.mocked(getGeneralSettingsCall).mockResolvedValue([ + { field_name: "mcp_allowed_clients", field_value: [ANTIGRAVITY, { alias: "", value: "claude-code" }] }, + ]); + + renderSettings(); + + expect(await screen.findByText(/stored allowlist is not a list of alias and value pairs/)).toBeVisible(); + expect(screen.queryByRole("button", { name: /^Antigravity CLI/ })).not.toBeInTheDocument(); + }); + it("adds clients as alias and value pairs and saves them under mcp_allowed_clients", async () => { renderSettings(); await screen.findByText("Allowed Clients"); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPNetworkSettings.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPNetworkSettings.tsx index db45cfdedd1..2fd62c7f1ef 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPNetworkSettings.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPNetworkSettings.tsx @@ -52,7 +52,7 @@ interface ClientDraft extends AllowedClient { const isAllowedClient = (entry: unknown): entry is AllowedClient => { if (typeof entry !== "object" || entry === null) return false; const { alias, value } = entry as Partial>; - return typeof alias === "string" && typeof value === "string"; + return typeof alias === "string" && typeof value === "string" && !isIncomplete({ alias, value }); }; type StoredAllowlist = From b2d6cd1fcfde4bffb48473ff14b84fa221733864 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Fri, 18 Sep 2026 17:23:26 -0700 Subject: [PATCH 091/144] refactor(rust): read litellm HTTP globals through one Python shim and tighten the http pool Drop the core ocr() facade so VertexAuth and the http pool stay out of litellm-core's public API, move the http Error enum to error.rs, and inject the media DNS resolver into HttpClientPool instead of a per-call builder hook the cache key ignored. The bridge now reads litellm.* HTTP settings only through litellm/rust_bridge/settings.py, pinned by python_settings.json, while env overrides stay in Rust. This adds the Python default User-Agent, parses string ssl_verify globals like get_ssl_verify, drops per-call ssl_verify that Python OCR never honored, and removes the unused request_timeout. Co-Authored-By: Claude Opus 5 --- litellm-rust/crates/core/Cargo.toml | 4 +- litellm-rust/crates/core/src/ocr/client.rs | 11 - litellm-rust/crates/core/tests/ocr.rs | 41 +-- litellm-rust/crates/http/src/config.rs | 93 ++----- litellm-rust/crates/http/src/error.rs | 22 ++ litellm-rust/crates/http/src/lib.rs | 4 +- litellm-rust/crates/http/src/pool.rs | 253 +++++++++++------- litellm-rust/crates/http/src/settings.rs | 2 - .../crates/llms/src/custom_httpx/media.rs | 33 +-- .../crates/python-bridge/python_settings.json | 12 + litellm-rust/crates/python-bridge/src/http.rs | 175 ++++++------ litellm-rust/crates/python-bridge/src/lib.rs | 1 + .../python-bridge/src/python_settings.rs | 48 ++++ litellm/llms/custom_httpx/http_handler.py | 6 +- litellm/rust_bridge/settings.py | 37 +++ .../test_litellm/rust_bridge/test_settings.py | 50 ++++ 16 files changed, 467 insertions(+), 325 deletions(-) create mode 100644 litellm-rust/crates/http/src/error.rs create mode 100644 litellm-rust/crates/python-bridge/python_settings.json create mode 100644 litellm-rust/crates/python-bridge/src/python_settings.rs create mode 100644 litellm/rust_bridge/settings.py create mode 100644 tests/test_litellm/rust_bridge/test_settings.py diff --git a/litellm-rust/crates/core/Cargo.toml b/litellm-rust/crates/core/Cargo.toml index 1a769ec9708..ab04fb8d4ae 100644 --- a/litellm-rust/crates/core/Cargo.toml +++ b/litellm-rust/crates/core/Cargo.toml @@ -15,8 +15,6 @@ futures-util.workspace = true base64.workspace = true litellm-auth.workspace = true litellm-auth-aws.workspace = true -litellm-auth-gcp.workspace = true -litellm-http.workspace = true litellm-llms.workspace = true moka.workspace = true mime_guess = "2.0.5" @@ -37,6 +35,8 @@ url.workspace = true veil.workspace = true [dev-dependencies] +litellm-auth-gcp.workspace = true +litellm-http.workspace = true litellm-llms = { workspace = true, features = ["test-support"] } rstest.workspace = true rstest_reuse.workspace = true diff --git a/litellm-rust/crates/core/src/ocr/client.rs b/litellm-rust/crates/core/src/ocr/client.rs index f0b24623a88..c7b4751bd9e 100644 --- a/litellm-rust/crates/core/src/ocr/client.rs +++ b/litellm-rust/crates/core/src/ocr/client.rs @@ -1,5 +1,3 @@ -use litellm_auth_gcp::VertexAuth; -use litellm_http::{HttpClientConfig, HttpClientPool}; use litellm_llms::{ base_llm::ocr::{error::Error, transformation::LiteLLMOcrResponse}, custom_httpx::llm_http_handler::OcrClient, @@ -16,12 +14,3 @@ pub async fn perform( ) -> Result { litellm_host::run::run(ocr_machine(client.clone()), &LocalOcrHost::new(request)).await } - -pub async fn ocr( - pool: &HttpClientPool, - config: &HttpClientConfig, - vertex_auth: VertexAuth, - request: LiteLLMOcrRequest, -) -> Result { - perform(&OcrClient::new(pool, config, vertex_auth)?, request).await -} diff --git a/litellm-rust/crates/core/tests/ocr.rs b/litellm-rust/crates/core/tests/ocr.rs index 2e414c58541..2ae162d964f 100644 --- a/litellm-rust/crates/core/tests/ocr.rs +++ b/litellm-rust/crates/core/tests/ocr.rs @@ -6,13 +6,13 @@ use litellm_host::{ host::{Host, HostOp, HostResult}, machine::{HostFailure, Machine, MachineStep}, }; -use litellm_http::{HttpClientConfig, HttpClientPool, HttpSettings, Verify}; +use litellm_http::{HttpClientConfig, HttpClientPool, HttpSettings}; use litellm_llms::{ base_llm::ocr::{ error::Error as OcrError, transformation::{LiteLLMOcrResponse, OCR_RESPONSE_MAX_BYTES, OcrTransportConfig}, }, - custom_httpx::llm_http_handler::OcrClient, + custom_httpx::{llm_http_handler::OcrClient, media::PublicDnsResolver}, }; use rstest::rstest; use serde_json::{Value, json}; @@ -173,48 +173,25 @@ async fn facade_retains_native_response_when_requested() { } #[tokio::test] -async fn facade_uses_the_injected_http_pool_configuration() { +async fn ocr_client_uses_the_injected_http_pool_configuration() { let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; let settings = HttpSettings { user_agent: Some("host-owned/1".into()), ..HttpSettings::default() }; - let config = HttpClientConfig::resolve(&settings, None).unwrap(); - crate::ocr::client::ocr( - &HttpClientPool::new(), - &config, + let client = OcrClient::new( + &HttpClientPool::new(Arc::new(PublicDnsResolver)), + &HttpClientConfig::resolve(&settings).unwrap(), VertexAuth::default(), - wire_request("mistral/model", &base, json!({})), ) - .await .unwrap(); + crate::ocr::client::perform(&client, wire_request("mistral/model", &base, json!({}))) + .await + .unwrap(); server.await.unwrap(); assert!(seen.lock().unwrap()[0].contains("user-agent: host-owned/1")); } -#[tokio::test] -async fn unbuildable_http_configuration_fails_before_dispatch() { - let (base, _seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; - let config = HttpClientConfig { - verify: Verify::CaBundle(std::env::temp_dir().join("litellm-ocr-missing-bundle.pem")), - ..HttpClientConfig::resolve(&HttpSettings::default(), None).unwrap() - }; - let error = crate::ocr::client::ocr( - &HttpClientPool::new(), - &config, - VertexAuth::default(), - wire_request("mistral/model", &base, json!({})), - ) - .await - .unwrap_err(); - server.abort(); - assert!(matches!( - error, - OcrError::Transport(litellm_llms::custom_httpx::transport::Error::Connect(_)) - )); - assert!(error.to_string().contains("litellm-ocr-missing-bundle.pem")); -} - fn event_name(event: &CallEvent) -> &'static str { match event { CallEvent::Started { .. } => "started", diff --git a/litellm-rust/crates/http/src/config.rs b/litellm-rust/crates/http/src/config.rs index 86f1a9b43b6..f27092c1fb5 100644 --- a/litellm-rust/crates/http/src/config.rs +++ b/litellm-rust/crates/http/src/config.rs @@ -4,28 +4,10 @@ use std::{ time::Duration, }; -use crate::settings::{HttpSettings, SslVerify}; - -#[derive(Clone, Debug, thiserror::Error, PartialEq, Eq)] -pub enum Error { - #[error("{setting} cannot be expressed with rustls: {reason}")] - Unsupported { - setting: &'static str, - reason: String, - }, - #[error("could not read {}: {message}", path.display())] - Read { path: PathBuf, message: String }, - #[error("{} is not a PEM file: {message}", path.display())] - InvalidPem { path: PathBuf, message: String }, - #[error("could not build the HTTP client: {0}")] - Client(String), -} - -impl From for Error { - fn from(error: reqwest::Error) -> Self { - Self::Client(error.without_url().to_string()) - } -} +use crate::{ + error::Error, + settings::{HttpSettings, SslVerify}, +}; #[derive(Clone, Debug, PartialEq, Eq, Hash)] pub enum Verify { @@ -45,17 +27,13 @@ pub struct HttpClientConfig { pub user_agent: Option, pub trust_proxy_env: bool, pub connect_timeout: Duration, - pub request_timeout: Option, } impl HttpClientConfig { - /// Port of `get_ssl_verify` + `get_ssl_configuration`: the per-call value wins, then the - /// configured (environment-overlaid) `ssl_verify`, then `SSL_CERT_FILE`, then the built-in - /// roots. Settings rustls has no equivalent for are an error instead of a silent no-op. - pub fn resolve( - settings: &HttpSettings, - per_call_ssl_verify: Option<&SslVerify>, - ) -> Result { + /// Port of `get_ssl_verify` + `get_ssl_configuration`: the configured (environment-overlaid) + /// `ssl_verify`, then `SSL_CERT_FILE`, then the built-in roots. Settings rustls has no + /// equivalent for are an error instead of a silent no-op. + pub fn resolve(settings: &HttpSettings) -> Result { if let Some(level) = &settings.ssl_security_level { return Err(Error::Unsupported { setting: "ssl_security_level", @@ -68,7 +46,7 @@ impl HttpClientConfig { reason: format!("key exchange group {curve:?} is fixed by the rustls provider"), }); } - let verify = match per_call_ssl_verify.or(settings.ssl_verify.as_ref()) { + let verify = match &settings.ssl_verify { Some(SslVerify::Disabled) => Verify::Disabled, Some(SslVerify::CaBundle(path)) => Verify::CaBundle(path.clone()), Some(SslVerify::Enabled) | None => settings @@ -84,7 +62,6 @@ impl HttpClientConfig { user_agent: settings.user_agent.clone(), trust_proxy_env: settings.trust_proxy_env, connect_timeout: settings.connect_timeout, - request_timeout: settings.request_timeout, }) } @@ -141,14 +118,10 @@ impl HttpClientConfig { Some(agent) => with_protocol.user_agent(agent), None => with_protocol, }; - let with_proxy = if self.trust_proxy_env { + Ok(if self.trust_proxy_env { with_agent } else { with_agent.no_proxy() - }; - Ok(match self.request_timeout { - Some(timeout) => with_proxy.timeout(timeout), - None => with_proxy, }) } } @@ -180,49 +153,25 @@ mod tests { } #[rstest] - #[case::default(settings(None, None), None, Verify::BuiltInRoots)] + #[case::default(settings(None, None), Verify::BuiltInRoots)] #[case::setting_disables( settings(Some(SslVerify::Disabled), Some("/env/roots.pem")), - None, Verify::Disabled )] #[case::setting_bundle( settings(Some(SslVerify::CaBundle("/configured.pem".into())), Some("/env/roots.pem")), - None, Verify::CaBundle("/configured.pem".into()) )] #[case::enabled_uses_cert_file( settings(Some(SslVerify::Enabled), Some("/env/roots.pem")), - None, Verify::CaBundle("/env/roots.pem".into()) )] - #[case::unset_uses_cert_file(settings(None, Some("/env/roots.pem")), None, Verify::CaBundle("/env/roots.pem".into()))] - #[case::per_call_beats_setting( - settings(Some(SslVerify::Disabled), None), - Some(SslVerify::Enabled), - Verify::BuiltInRoots - )] - #[case::per_call_disables( - settings(Some(SslVerify::CaBundle("/configured.pem".into())), Some("/env/roots.pem")), - Some(SslVerify::Disabled), - Verify::Disabled - )] - #[case::per_call_bundle( - settings(None, Some("/env/roots.pem")), - Some(SslVerify::CaBundle("/call.pem".into())), - Verify::CaBundle("/call.pem".into()) - )] - #[case::per_call_enabled_still_honours_cert_file( - settings(Some(SslVerify::Disabled), Some("/env/roots.pem")), - Some(SslVerify::Enabled), - Verify::CaBundle("/env/roots.pem".into()) - )] - fn verify_follows_per_call_then_setting_then_cert_file( + #[case::unset_uses_cert_file(settings(None, Some("/env/roots.pem")), Verify::CaBundle("/env/roots.pem".into()))] + fn verify_follows_setting_then_cert_file( #[case] settings: HttpSettings, - #[case] per_call: Option, #[case] expected: Verify, ) { - let config = HttpClientConfig::resolve(&settings, per_call.as_ref()).unwrap(); + let config = HttpClientConfig::resolve(&settings).unwrap(); assert_eq!(config.verify, expected); } @@ -233,7 +182,7 @@ mod tests { ..HttpSettings::default() } .with_environment(&|name: &str| (name == "SSL_VERIFY").then(|| "true".to_string())); - let config = HttpClientConfig::resolve(&settings, None).unwrap(); + let config = HttpClientConfig::resolve(&settings).unwrap(); assert_eq!(config.verify, Verify::BuiltInRoots); } @@ -244,7 +193,7 @@ mod tests { ..HttpSettings::default() }; assert!(matches!( - HttpClientConfig::resolve(&settings, None), + HttpClientConfig::resolve(&settings), Err(Error::Unsupported { setting: "ssl_security_level", .. @@ -259,7 +208,7 @@ mod tests { ..HttpSettings::default() }; assert!(matches!( - HttpClientConfig::resolve(&settings, None), + HttpClientConfig::resolve(&settings), Err(Error::Unsupported { setting: "ssl_ecdh_curve", .. @@ -276,10 +225,9 @@ mod tests { user_agent: Some("litellm/1.0".into()), trust_proxy_env: true, connect_timeout: Duration::from_secs(7), - request_timeout: Some(Duration::from_secs(70)), ..HttpSettings::default() }; - let config = HttpClientConfig::resolve(&settings, None).unwrap(); + let config = HttpClientConfig::resolve(&settings).unwrap(); assert_eq!( config, HttpClientConfig { @@ -290,7 +238,6 @@ mod tests { user_agent: Some("litellm/1.0".into()), trust_proxy_env: true, connect_timeout: Duration::from_secs(7), - request_timeout: Some(Duration::from_secs(70)), } ); } @@ -300,7 +247,7 @@ mod tests { let path = std::env::temp_dir().join("litellm-http-missing-bundle.pem"); let config = HttpClientConfig { verify: Verify::CaBundle(path.clone()), - ..HttpClientConfig::resolve(&HttpSettings::default(), None).unwrap() + ..HttpClientConfig::resolve(&HttpSettings::default()).unwrap() }; assert!(matches!( config.client_builder(), @@ -315,7 +262,7 @@ mod tests { std::fs::write(&path, b"not a certificate").unwrap(); let config = HttpClientConfig { verify: Verify::CaBundle(path.clone()), - ..HttpClientConfig::resolve(&HttpSettings::default(), None).unwrap() + ..HttpClientConfig::resolve(&HttpSettings::default()).unwrap() }; let result = config.client_builder().map(drop); std::fs::remove_file(&path).unwrap(); diff --git a/litellm-rust/crates/http/src/error.rs b/litellm-rust/crates/http/src/error.rs new file mode 100644 index 00000000000..27899f06cf1 --- /dev/null +++ b/litellm-rust/crates/http/src/error.rs @@ -0,0 +1,22 @@ +use std::path::PathBuf; + +#[derive(Clone, Debug, thiserror::Error, PartialEq, Eq)] +pub enum Error { + #[error("{setting} cannot be expressed with rustls: {reason}")] + Unsupported { + setting: &'static str, + reason: String, + }, + #[error("could not read {}: {message}", path.display())] + Read { path: PathBuf, message: String }, + #[error("{} is not a PEM file: {message}", path.display())] + InvalidPem { path: PathBuf, message: String }, + #[error("could not build the HTTP client: {0}")] + Client(String), +} + +impl From for Error { + fn from(error: reqwest::Error) -> Self { + Self::Client(error.without_url().to_string()) + } +} diff --git a/litellm-rust/crates/http/src/lib.rs b/litellm-rust/crates/http/src/lib.rs index d62dd768fe1..9c88e3101a7 100644 --- a/litellm-rust/crates/http/src/lib.rs +++ b/litellm-rust/crates/http/src/lib.rs @@ -3,9 +3,11 @@ //! caches `reqwest::Client`s per resolved configuration. mod config; +mod error; mod pool; mod settings; -pub use config::{Error, HttpClientConfig, Verify}; +pub use config::{HttpClientConfig, Verify}; +pub use error::Error; pub use pool::{ClientVariant, HttpClientPool}; pub use settings::{HttpSettings, SslVerify}; diff --git a/litellm-rust/crates/http/src/pool.rs b/litellm-rust/crates/http/src/pool.rs index 065097556e1..03c01e968ac 100644 --- a/litellm-rust/crates/http/src/pool.rs +++ b/litellm-rust/crates/http/src/pool.rs @@ -1,112 +1,174 @@ use std::{ collections::HashMap, - sync::{Mutex, PoisonError}, + sync::{Arc, Mutex, PoisonError}, }; -use crate::config::{Error, HttpClientConfig}; +use reqwest::dns::Resolve; + +use crate::{config::HttpClientConfig, error::Error}; /// The client shapes routes need; each is the shared base plus one policy. #[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)] pub enum ClientVariant { Provider, NoRedirect, - /// Media downloads: no redirects (the fetcher validates each hop) and never a proxy. + /// Media downloads: no redirects (the fetcher validates each hop), never a proxy, and the + /// pool's media resolver. Media, } -impl ClientVariant { - fn apply(self, builder: reqwest::ClientBuilder) -> reqwest::ClientBuilder { - match self { - Self::Provider => builder, - Self::NoRedirect => builder.redirect(reqwest::redirect::Policy::none()), - Self::Media => builder - .redirect(reqwest::redirect::Policy::none()) - .no_proxy(), - } - } -} - /// Counterpart of `get_async_httpx_client`: one `reqwest::Client` per resolved configuration /// and variant, built on first use and shared afterwards. -#[derive(Default)] pub struct HttpClientPool { + media_resolver: Arc, clients: Mutex>, } impl HttpClientPool { - pub fn new() -> Self { - Self::default() + pub fn new(media_resolver: Arc) -> Self { + Self { + media_resolver, + clients: Mutex::default(), + } } pub fn client( &self, config: &HttpClientConfig, variant: ClientVariant, - ) -> Result { - self.client_with(config, variant, |builder| builder) - } - - /// Like [`Self::client`], with a caller hook for builder options that are not plain values - /// (a DNS resolver, for example). The hook only runs when the client is first built. - pub fn client_with( - &self, - config: &HttpClientConfig, - variant: ClientVariant, - customize: impl FnOnce(reqwest::ClientBuilder) -> reqwest::ClientBuilder, ) -> Result { let key = (config.clone(), variant); let mut clients = self.clients.lock().unwrap_or_else(PoisonError::into_inner); if let Some(client) = clients.get(&key) { return Ok(client.clone()); } - let client = customize(variant.apply(config.client_builder()?)).build()?; + let client = self.apply(variant, config.client_builder()?).build()?; clients.insert(key, client.clone()); Ok(client) } + + fn apply( + &self, + variant: ClientVariant, + builder: reqwest::ClientBuilder, + ) -> reqwest::ClientBuilder { + match variant { + ClientVariant::Provider => builder, + ClientVariant::NoRedirect => builder.redirect(reqwest::redirect::Policy::none()), + ClientVariant::Media => builder + .redirect(reqwest::redirect::Policy::none()) + .no_proxy() + .dns_resolver2(Arc::clone(&self.media_resolver)), + } + } } #[cfg(test)] mod tests { - use std::{cell::Cell, time::Duration}; + use std::{ + net::SocketAddr, + sync::atomic::{AtomicUsize, Ordering}, + time::Duration, + }; - use tokio::io::{AsyncReadExt, AsyncWriteExt}; + use reqwest::dns::{Addrs, Name, Resolving}; + use tokio::{ + io::{AsyncReadExt, AsyncWriteExt}, + net::TcpListener, + }; use super::*; use crate::{HttpSettings, Verify}; - fn config(user_agent: &str) -> HttpClientConfig { - HttpClientConfig { - user_agent: Some(user_agent.into()), - ..HttpClientConfig::resolve(&HttpSettings::default(), None).unwrap() + struct FixedResolver(SocketAddr); + + impl Resolve for FixedResolver { + fn resolve(&self, _: Name) -> Resolving { + let addrs: Addrs = Box::new(std::iter::once(self.0)); + Box::pin(std::future::ready(Ok(addrs))) } } - #[test] - fn clients_are_built_once_per_config_and_variant() { - let pool = HttpClientPool::new(); - let builds = Cell::new(0); - let build = |config: &HttpClientConfig, variant| { - pool.client_with(config, variant, |builder| { - builds.set(builds.get() + 1); - builder - }) + fn pool() -> HttpClientPool { + HttpClientPool::new(Arc::new(FixedResolver(([192, 0, 2, 1], 80).into()))) + } + + fn config(user_agent: &str) -> HttpClientConfig { + HttpClientConfig { + user_agent: Some(user_agent.into()), + ..HttpClientConfig::resolve(&HttpSettings::default()).unwrap() + } + } + + /// Answers every request on every connection with `status_line` and counts connections, + /// so a reused client shows up as a reused keep-alive connection. + async fn serve( + status_line: &'static str, + ) -> (SocketAddr, Arc, Arc>>) { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let connections = Arc::new(AtomicUsize::new(0)); + let requests = Arc::new(Mutex::new(Vec::new())); + let (accepted, seen) = (Arc::clone(&connections), Arc::clone(&requests)); + tokio::spawn(async move { + loop { + let (mut socket, _) = listener.accept().await.unwrap(); + accepted.fetch_add(1, Ordering::SeqCst); + let seen = Arc::clone(&seen); + tokio::spawn(async move { + let mut buffer = vec![0u8; 4096]; + while let Ok(read) = socket.read(&mut buffer).await { + if read == 0 { + return; + } + seen.lock() + .unwrap() + .push(String::from_utf8_lossy(&buffer[..read]).into_owned()); + let response = format!( + "{status_line}\r\nLocation: /elsewhere\r\nContent-Length: 0\r\n\r\n" + ); + if socket.write_all(response.as_bytes()).await.is_err() { + return; + } + } + }); + } + }); + (address, connections, requests) + } + + async fn get( + pool: &HttpClientPool, + config: &HttpClientConfig, + variant: ClientVariant, + url: &str, + ) -> reqwest::Response { + pool.client(config, variant) .unwrap() - }; - build(&config("a"), ClientVariant::Provider); - build(&config("a"), ClientVariant::Provider); - assert_eq!(builds.get(), 1); - build(&config("a"), ClientVariant::NoRedirect); - assert_eq!(builds.get(), 2); - build(&config("b"), ClientVariant::Provider); - assert_eq!(builds.get(), 3); - build(&config("b"), ClientVariant::Provider); - build(&config("a"), ClientVariant::NoRedirect); - assert_eq!(builds.get(), 3); + .get(url) + .timeout(Duration::from_secs(5)) + .send() + .await + .unwrap() + } + + #[tokio::test] + async fn clients_are_shared_per_config_and_variant() { + let (address, connections, _) = serve("HTTP/1.1 204 No Content").await; + let url = format!("http://{address}"); + let pool = pool(); + get(&pool, &config("a"), ClientVariant::Provider, &url).await; + get(&pool, &config("a"), ClientVariant::Provider, &url).await; + assert_eq!(connections.load(Ordering::SeqCst), 1); + get(&pool, &config("a"), ClientVariant::NoRedirect, &url).await; + assert_eq!(connections.load(Ordering::SeqCst), 2); + get(&pool, &config("b"), ClientVariant::Provider, &url).await; + assert_eq!(connections.load(Ordering::SeqCst), 3); } #[test] fn build_failures_are_not_cached() { - let pool = HttpClientPool::new(); + let pool = pool(); let missing = HttpClientConfig { verify: Verify::CaBundle(std::env::temp_dir().join("litellm-http-absent.pem")), ..config("a") @@ -116,57 +178,52 @@ mod tests { assert!(pool.client(&config("a"), ClientVariant::Provider).is_ok()); } - async fn serve_once(status_line: &'static str) -> (String, tokio::task::JoinHandle) { - let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); - let base = format!("http://{}", listener.local_addr().unwrap()); - let server = tokio::spawn(async move { - let (mut socket, _) = listener.accept().await.unwrap(); - let mut request = vec![0u8; 4096]; - let read = socket.read(&mut request).await.unwrap(); - socket - .write_all( - format!("{status_line}\r\nLocation: /elsewhere\r\nContent-Length: 0\r\nConnection: close\r\n\r\n") - .as_bytes(), - ) - .await - .unwrap(); - String::from_utf8_lossy(&request[..read]).into_owned() - }); - (base, server) - } - #[tokio::test] async fn provider_client_sends_the_configured_user_agent_over_http1() { - let (base, server) = serve_once("HTTP/1.1 204 No Content").await; - let config = HttpClientConfig { - connect_timeout: Duration::from_secs(2), - ..config("litellm-test/9") - }; - let response = HttpClientPool::new() - .client(&config, ClientVariant::Provider) - .unwrap() - .get(&base) - .send() - .await - .unwrap(); + let (address, _, requests) = serve("HTTP/1.1 204 No Content").await; + let response = get( + &pool(), + &config("litellm-test/9"), + ClientVariant::Provider, + &format!("http://{address}"), + ) + .await; assert_eq!(response.status(), 204); assert_eq!(response.version(), reqwest::Version::HTTP_11); - let request = server.await.unwrap(); + let request = requests.lock().unwrap()[0].clone(); assert!(request.contains("user-agent: litellm-test/9"), "{request}"); } #[tokio::test] async fn no_redirect_variant_returns_the_redirect_instead_of_following_it() { - let (base, server) = serve_once("HTTP/1.1 302 Found").await; - let response = HttpClientPool::new() - .client(&config("a"), ClientVariant::NoRedirect) - .unwrap() - .get(&base) - .send() - .await - .unwrap(); + let (address, _, _) = serve("HTTP/1.1 302 Found").await; + let response = get( + &pool(), + &config("a"), + ClientVariant::NoRedirect, + &format!("http://{address}"), + ) + .await; assert_eq!(response.status(), 302); assert_eq!(response.headers()["location"], "/elsewhere"); - server.await.unwrap(); + } + + #[tokio::test] + async fn media_variant_resolves_through_the_injected_resolver() { + let (address, _, requests) = serve("HTTP/1.1 204 No Content").await; + let pool = HttpClientPool::new(Arc::new(FixedResolver(address))); + let url = format!("http://media.invalid:{}/doc", address.port()); + let response = get(&pool, &config("a"), ClientVariant::Media, &url).await; + assert_eq!(response.status(), 204); + assert!(requests.lock().unwrap()[0].contains("host: media.invalid")); + assert!( + pool.client(&config("a"), ClientVariant::Provider) + .unwrap() + .get(&url) + .timeout(Duration::from_secs(5)) + .send() + .await + .is_err() + ); } } diff --git a/litellm-rust/crates/http/src/settings.rs b/litellm-rust/crates/http/src/settings.rs index 4d936e89046..45aab0d6fa5 100644 --- a/litellm-rust/crates/http/src/settings.rs +++ b/litellm-rust/crates/http/src/settings.rs @@ -31,7 +31,6 @@ pub struct HttpSettings { pub user_agent: Option, pub trust_proxy_env: bool, pub connect_timeout: Duration, - pub request_timeout: Option, } impl Default for HttpSettings { @@ -47,7 +46,6 @@ impl Default for HttpSettings { user_agent: None, trust_proxy_env: false, connect_timeout: Duration::from_secs(5), - request_timeout: None, } } } diff --git a/litellm-rust/crates/llms/src/custom_httpx/media.rs b/litellm-rust/crates/llms/src/custom_httpx/media.rs index aeac4894683..a1c4fe68734 100644 --- a/litellm-rust/crates/llms/src/custom_httpx/media.rs +++ b/litellm-rust/crates/llms/src/custom_httpx/media.rs @@ -66,26 +66,15 @@ impl MediaFetcher { pool: &HttpClientPool, config: &HttpClientConfig, ) -> Result { - Self::with_resolvers( - pool, - config, - Arc::new(PublicDnsResolver), - Arc::new(SystemAddressResolver), - ) + Self::with_address_resolver(pool, config, Arc::new(SystemAddressResolver)) } - fn with_resolvers( + fn with_address_resolver( pool: &HttpClientPool, config: &HttpClientConfig, - transport_resolver: Arc, address_resolver: Arc, - ) -> Result - where - R: Resolve + 'static, - { - let client = pool.client_with(config, ClientVariant::Media, |builder| { - builder.dns_resolver(transport_resolver) - })?; + ) -> Result { + let client = pool.client(config, ClientVariant::Media)?; Ok(Self { client, address_resolver, @@ -242,7 +231,7 @@ fn is_blocked_ip(ip: IpAddr) -> bool { } #[derive(Default)] -struct PublicDnsResolver; +pub struct PublicDnsResolver; struct SystemAddressResolver; @@ -371,10 +360,9 @@ mod tests { address: SocketAddr, blocked_hosts: HashSet<&'static str>, ) -> MediaFetcher { - MediaFetcher::with_resolvers( - &HttpClientPool::new(), - &HttpClientConfig::resolve(&HttpSettings::default(), None).unwrap(), - Arc::new(LoopbackDnsResolver(address)), + MediaFetcher::with_address_resolver( + &HttpClientPool::new(Arc::new(LoopbackDnsResolver(address))), + &HttpClientConfig::resolve(&HttpSettings::default()).unwrap(), Arc::new(TestAddressResolver { blocked_hosts }), ) .expect("test fetcher builds") @@ -552,9 +540,8 @@ mod tests { #[tokio::test] async fn rejects_url_credentials_before_network_access() { let fetcher = MediaFetcher::new( - &HttpClientPool::new(), - &HttpClientConfig::resolve(&litellm_http::HttpSettings::default(), None) - .expect("default settings resolve"), + &HttpClientPool::new(Arc::new(PublicDnsResolver)), + &HttpClientConfig::resolve(&HttpSettings::default()).expect("default settings resolve"), ) .expect("media fetcher builds"); let url = diff --git a/litellm-rust/crates/python-bridge/python_settings.json b/litellm-rust/crates/python-bridge/python_settings.json new file mode 100644 index 00000000000..a6cd959c6de --- /dev/null +++ b/litellm-rust/crates/python-bridge/python_settings.json @@ -0,0 +1,12 @@ +{ + "http_settings": [ + "ssl_verify", + "ssl_certificate", + "ssl_security_level", + "ssl_ecdh_curve", + "force_ipv4", + "http2", + "aiohttp_trust_env", + "user_agent" + ] +} diff --git a/litellm-rust/crates/python-bridge/src/http.rs b/litellm-rust/crates/python-bridge/src/http.rs index d2e05c3b949..de6fb5bb96d 100644 --- a/litellm-rust/crates/python-bridge/src/http.rs +++ b/litellm-rust/crates/python-bridge/src/http.rs @@ -1,11 +1,16 @@ -use std::{path::PathBuf, sync::LazyLock}; +use std::{ + path::PathBuf, + sync::{Arc, LazyLock}, +}; use litellm_http::{HttpClientConfig, HttpClientPool, HttpSettings, SslVerify}; +use litellm_llms::custom_httpx::media::PublicDnsResolver; use pyo3::{prelude::*, types::PyDict}; -use crate::errors::RustBridgeDeclined; +use crate::{errors::RustBridgeDeclined, python_settings::PythonSettings}; -static POOL: LazyLock = LazyLock::new(HttpClientPool::new); +static POOL: LazyLock = + LazyLock::new(|| HttpClientPool::new(Arc::new(PublicDnsResolver))); /// Keyword arguments that carry a live Python HTTP client or session. They cannot cross into /// Rust, so a call that supplies one stays on the Python path. @@ -15,21 +20,16 @@ pub(crate) fn pool() -> &'static HttpClientPool { &POOL } -/// The client configuration for one call: the process settings from the `litellm` module and -/// the environment, narrowed by the call's own `ssl_verify`. +/// The client configuration for one call: the `litellm.*` HTTP settings with the environment +/// overlaid, the same way `http_handler.py` combines them. pub(crate) fn call_config( py: Python<'_>, kwargs: &Bound<'_, PyDict>, ) -> PyResult { decline_live_clients(kwargs)?; - let settings = settings(py.import("litellm")?.as_any())? + let settings = settings(&PythonSettings::Http.read(py)?)? .with_environment(&|name| std::env::var(name).ok()); - let per_call = kwargs - .get_item("ssl_verify")? - .map(|value| ssl_verify(&value, "ssl_verify")) - .transpose()? - .flatten(); - HttpClientConfig::resolve(&settings, per_call.as_ref()) + HttpClientConfig::resolve(&settings) .map_err(|error| RustBridgeDeclined::new_err(error.to_string())) } @@ -44,45 +44,47 @@ pub(crate) fn decline_live_clients(kwargs: &Bound<'_, PyDict>) -> PyResult<()> { Ok(()) } -/// Read the `litellm.*` globals `http_handler.py` consults. `globals` is the `litellm` module in -/// production and any attribute holder in tests. -pub(crate) fn settings(globals: &Bound<'_, PyAny>) -> PyResult { +#[derive(FromPyObject)] +struct PythonHttpSettings<'py> { + ssl_verify: Bound<'py, PyAny>, + ssl_certificate: Option, + ssl_security_level: Option, + ssl_ecdh_curve: Option, + force_ipv4: bool, + http2: bool, + aiohttp_trust_env: bool, + user_agent: String, +} + +fn settings(value: &Bound<'_, PyAny>) -> PyResult { + let python: PythonHttpSettings = value.extract()?; Ok(HttpSettings { - ssl_verify: ssl_verify(&globals.getattr("ssl_verify")?, "litellm.ssl_verify")?, - ssl_certificate: optional_path(globals, "ssl_certificate")?, - ssl_security_level: globals.getattr("ssl_security_level")?.extract()?, - ssl_ecdh_curve: globals.getattr("ssl_ecdh_curve")?.extract()?, - force_ipv4: globals.getattr("force_ipv4")?.extract()?, - http2: globals.getattr("http2")?.extract()?, - trust_proxy_env: globals.getattr("aiohttp_trust_env")?.extract()?, + ssl_verify: Some(ssl_verify(&python.ssl_verify)?), + ssl_certificate: python.ssl_certificate.map(PathBuf::from), + ssl_security_level: python.ssl_security_level, + ssl_ecdh_curve: python.ssl_ecdh_curve, + force_ipv4: python.force_ipv4, + http2: python.http2, + user_agent: Some(python.user_agent), + trust_proxy_env: python.aiohttp_trust_env, ..HttpSettings::default() }) } -fn optional_path(globals: &Bound<'_, PyAny>, name: &str) -> PyResult> { - Ok(globals - .getattr(name)? - .extract::>()? - .map(PathBuf::from)) -} - -fn ssl_verify(value: &Bound<'_, PyAny>, name: &str) -> PyResult> { - if value.is_none() { - return Ok(None); - } +fn ssl_verify(value: &Bound<'_, PyAny>) -> PyResult { if let Ok(enabled) = value.extract::() { - return Ok(Some(if enabled { + return Ok(if enabled { SslVerify::Enabled } else { SslVerify::Disabled - })); + }); } if let Ok(path) = value.extract::() { - return Ok(Some(SslVerify::CaBundle(PathBuf::from(path)))); + return Ok(SslVerify::parse(&path)); } - Err(RustBridgeDeclined::new_err(format!( - "{name} is a live Python object and cannot be used by the Rust route" - ))) + Err(RustBridgeDeclined::new_err( + "litellm.ssl_verify is a live Python object and cannot be used by the Rust route", + )) } #[cfg(test)] @@ -91,18 +93,16 @@ mod tests { use rstest::rstest; use super::*; + use crate::python_settings::CONTRACT; - fn eval<'py>(py: Python<'py>, source: &std::ffi::CStr) -> Bound<'py, PyDict> { - let locals = PyDict::new(py); - py.run(source, Some(&locals), Some(&locals)).unwrap(); - locals - } - - fn globals<'py>(py: Python<'py>, overrides: &str) -> Bound<'py, PyAny> { + /// A stand-in for `http_settings()` carrying exactly the fields the contract declares, so a + /// field Rust reads but Python does not return fails here. + fn python_settings<'py>(py: Python<'py>, overrides: &str) -> Bound<'py, PyAny> { let source = format!( " +import json import types -globals = types.SimpleNamespace( +defaults = dict( ssl_verify=True, ssl_certificate=None, ssl_security_level=None, @@ -110,23 +110,29 @@ globals = types.SimpleNamespace( force_ipv4=False, http2=False, aiohttp_trust_env=False, + user_agent='litellm/test', ) -{overrides} +defaults.update(dict({overrides})) +settings = types.SimpleNamespace(**{{name: defaults[name] for name in json.loads(contract)['http_settings']}}) " ); + let locals = PyDict::new(py); + locals.set_item("contract", CONTRACT).unwrap(); let source = std::ffi::CString::new(source).unwrap(); - eval(py, &source).get_item("globals").unwrap().unwrap() + py.run(&source, Some(&locals), Some(&locals)).unwrap(); + locals.get_item("settings").unwrap().unwrap() } #[test] - fn default_globals_produce_default_settings_with_verification_on() { + fn default_python_settings_produce_default_settings_with_verification_on() { Python::initialize(); Python::attach(|py| { - let settings = settings(&globals(py, "")).unwrap(); + let settings = settings(&python_settings(py, "")).unwrap(); assert_eq!( settings, HttpSettings { ssl_verify: Some(SslVerify::Enabled), + user_agent: Some("litellm/test".into()), ..HttpSettings::default() } ); @@ -134,19 +140,20 @@ globals = types.SimpleNamespace( } #[test] - fn globals_flow_into_settings() { + fn python_settings_flow_into_settings() { Python::initialize(); Python::attach(|py| { - let settings = settings(&globals( + let settings = settings(&python_settings( py, " -globals.ssl_verify = '/etc/ssl/corp.pem' -globals.ssl_certificate = '/etc/ssl/client.pem' -globals.ssl_security_level = '2' -globals.ssl_ecdh_curve = 'X25519' -globals.force_ipv4 = True -globals.http2 = True -globals.aiohttp_trust_env = True +ssl_verify='/etc/ssl/corp.pem', +ssl_certificate='/etc/ssl/client.pem', +ssl_security_level='2', +ssl_ecdh_curve='X25519', +force_ipv4=True, +http2=True, +aiohttp_trust_env=True, +user_agent='litellm/9.9.9', ", )) .unwrap(); @@ -159,6 +166,7 @@ globals.aiohttp_trust_env = True ssl_ecdh_curve: Some("X25519".into()), force_ipv4: true, http2: true, + user_agent: Some("litellm/9.9.9".into()), trust_proxy_env: true, ..HttpSettings::default() } @@ -167,13 +175,32 @@ globals.aiohttp_trust_env = True } #[test] - fn disabled_verification_global_resolves_to_disabled() { + fn user_agent_environment_variable_beats_the_python_default() { Python::initialize(); Python::attach(|py| { - let settings = settings(&globals(py, "globals.ssl_verify = False")).unwrap(); - assert_eq!(settings.ssl_verify, Some(SslVerify::Disabled)); - let config = HttpClientConfig::resolve(&settings, None).unwrap(); - assert_eq!(config.verify, Verify::Disabled); + let settings = settings(&python_settings(py, "")) + .unwrap() + .with_environment(&|name| { + (name == "LITELLM_USER_AGENT").then(|| "operator/1".to_string()) + }); + assert_eq!(settings.user_agent.as_deref(), Some("operator/1")); + }); + } + + #[rstest] + #[case::disabled("ssl_verify=False", Verify::Disabled)] + #[case::disabled_string("ssl_verify='False'", Verify::Disabled)] + #[case::enabled_string("ssl_verify='true'", Verify::BuiltInRoots)] + #[case::bundle("ssl_verify='/tmp/ca.pem'", Verify::CaBundle("/tmp/ca.pem".into()))] + fn ssl_verify_global_resolves_like_get_ssl_verify( + #[case] overrides: &str, + #[case] expected: Verify, + ) { + Python::initialize(); + Python::attach(|py| { + let settings = settings(&python_settings(py, overrides)).unwrap(); + let config = HttpClientConfig::resolve(&settings).unwrap(); + assert_eq!(config.verify, expected); }); } @@ -181,7 +208,7 @@ globals.aiohttp_trust_env = True fn ssl_context_global_declines_instead_of_being_dropped() { Python::initialize(); Python::attach(|py| { - let error = settings(&globals(py, "globals.ssl_verify = object()")).unwrap_err(); + let error = settings(&python_settings(py, "ssl_verify=object()")).unwrap_err(); assert!(error.is_instance_of::(py)); assert!(error.value(py).to_string().contains("litellm.ssl_verify")); }); @@ -215,20 +242,4 @@ globals.aiohttp_trust_env = True decline_live_clients(&kwargs).unwrap(); }); } - - #[rstest] - #[case::disabled(c"False", Some(SslVerify::Disabled))] - #[case::enabled(c"True", Some(SslVerify::Enabled))] - #[case::bundle(c"'/tmp/ca.pem'", Some(SslVerify::CaBundle("/tmp/ca.pem".into())))] - #[case::unset(c"None", None)] - fn per_call_ssl_verify_values_project( - #[case] source: &std::ffi::CStr, - #[case] expected: Option, - ) { - Python::initialize(); - Python::attach(|py| { - let value = py.eval(source, None, None).unwrap(); - assert_eq!(ssl_verify(&value, "ssl_verify").unwrap(), expected); - }); - } } diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs index 11cb0a7f655..7eba0d201be 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -3,6 +3,7 @@ mod diagnostics; mod errors; mod http; mod marshal; +mod python_settings; mod routes; mod token_counter; diff --git a/litellm-rust/crates/python-bridge/src/python_settings.rs b/litellm-rust/crates/python-bridge/src/python_settings.rs new file mode 100644 index 00000000000..0c6554f0970 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/python_settings.rs @@ -0,0 +1,48 @@ +use pyo3::prelude::*; + +const MODULE: &str = "litellm.rust_bridge.settings"; + +/// Every group of `litellm.*` module globals the native routes read. Environment overrides are +/// applied on the Rust side, so each function returns only what the Python process configured. +/// A group is deleted once Rust owns loading that configuration, so this enum only shrinks. +/// +/// `litellm/rust_bridge/settings.py` is the only Python module behind it, and +/// `python_settings.json` pins the fields each function returns on both sides. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(crate) enum PythonSettings { + Http, +} + +impl PythonSettings { + #[cfg(test)] + pub(crate) const ALL: [Self; 1] = [Self::Http]; + + pub(crate) fn name(self) -> &'static str { + match self { + Self::Http => "http_settings", + } + } + + pub(crate) fn read(self, py: Python<'_>) -> PyResult> { + py.import(MODULE)?.getattr(self.name())?.call0() + } +} + +#[cfg(test)] +pub(crate) const CONTRACT: &str = include_str!("../python_settings.json"); + +#[cfg(test)] +mod tests { + use std::collections::BTreeSet; + + use super::{CONTRACT, PythonSettings}; + + #[test] + fn every_settings_group_is_in_the_python_contract() { + let contract: serde_json::Map = + serde_json::from_str(CONTRACT).unwrap(); + let declared: BTreeSet<&str> = contract.keys().map(String::as_str).collect(); + let read: BTreeSet<&str> = PythonSettings::ALL.map(PythonSettings::name).into(); + assert_eq!(read, declared); + } +} diff --git a/litellm/llms/custom_httpx/http_handler.py b/litellm/llms/custom_httpx/http_handler.py index 05dff0cb9d8..6b90394043f 100644 --- a/litellm/llms/custom_httpx/http_handler.py +++ b/litellm/llms/custom_httpx/http_handler.py @@ -150,7 +150,11 @@ def get_default_headers() -> dict: if user_agent is not None: return {"User-Agent": user_agent} - return {"User-Agent": f"litellm/{version}"} + return {"User-Agent": default_user_agent()} + + +def default_user_agent() -> str: + return f"litellm/{version}" # Initialize headers (User-Agent) diff --git a/litellm/rust_bridge/settings.py b/litellm/rust_bridge/settings.py new file mode 100644 index 00000000000..a8229b12d13 --- /dev/null +++ b/litellm/rust_bridge/settings.py @@ -0,0 +1,37 @@ +"""The `litellm.*` module globals the native routes read. + +Environment variables that override these are applied in Rust, so nothing here reads `os.environ`. +`litellm-rust/crates/python-bridge/python_settings.json` pins the fields each function returns. +""" + +from __future__ import annotations + +from dataclasses import dataclass + + +@dataclass(frozen=True, slots=True) +class HttpSettings: + ssl_verify: bool | str + ssl_certificate: str | None + ssl_security_level: str | None + ssl_ecdh_curve: str | None + force_ipv4: bool + http2: bool + aiohttp_trust_env: bool + user_agent: str + + +def http_settings() -> HttpSettings: + import litellm + from litellm.llms.custom_httpx.http_handler import default_user_agent + + return HttpSettings( + ssl_verify=litellm.ssl_verify, + ssl_certificate=litellm.ssl_certificate, + ssl_security_level=litellm.ssl_security_level, + ssl_ecdh_curve=litellm.ssl_ecdh_curve, + force_ipv4=litellm.force_ipv4, + http2=litellm.http2, + aiohttp_trust_env=litellm.aiohttp_trust_env, + user_agent=default_user_agent(), + ) diff --git a/tests/test_litellm/rust_bridge/test_settings.py b/tests/test_litellm/rust_bridge/test_settings.py new file mode 100644 index 00000000000..618fa400136 --- /dev/null +++ b/tests/test_litellm/rust_bridge/test_settings.py @@ -0,0 +1,50 @@ +import dataclasses +from pathlib import Path +from typing import Final + +import pytest +from pydantic import TypeAdapter + +import litellm +from litellm.llms.custom_httpx.http_handler import default_user_agent +from litellm.rust_bridge import settings + +CONTRACT_PATH: Final = Path(__file__).parents[3] / "litellm-rust/crates/python-bridge/python_settings.json" + + +def test_the_rust_contract_matches_the_returned_fields() -> None: + contract: Final = TypeAdapter(dict[str, list[str]]).validate_json(CONTRACT_PATH.read_text()) + + assert contract == {"http_settings": [field.name for field in dataclasses.fields(settings.http_settings())]} + + +def test_http_settings_reads_the_litellm_globals(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "ssl_verify", "/etc/ssl/corp.pem") + monkeypatch.setattr(litellm, "ssl_certificate", "/etc/ssl/client.pem") + monkeypatch.setattr(litellm, "ssl_security_level", "DEFAULT@SECLEVEL=1") + monkeypatch.setattr(litellm, "ssl_ecdh_curve", "X25519") + monkeypatch.setattr(litellm, "force_ipv4", True) + monkeypatch.setattr(litellm, "http2", True) + monkeypatch.setattr(litellm, "aiohttp_trust_env", True) + + assert settings.http_settings() == settings.HttpSettings( + ssl_verify="/etc/ssl/corp.pem", + ssl_certificate="/etc/ssl/client.pem", + ssl_security_level="DEFAULT@SECLEVEL=1", + ssl_ecdh_curve="X25519", + force_ipv4=True, + http2=True, + aiohttp_trust_env=True, + user_agent=settings.http_settings().user_agent, + ) + + +def test_http_settings_ignores_environment_overrides(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("LITELLM_USER_AGENT", "operator/1") + monkeypatch.setenv("SSL_VERIFY", "false") + monkeypatch.setattr(litellm, "ssl_verify", True) + + result: Final = settings.http_settings() + + assert result.user_agent == default_user_agent() + assert result.ssl_verify is True From 59f7a00cf62fc40aa281eac50a57e02cedd3d0f7 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 18 Sep 2026 17:23:59 -0700 Subject: [PATCH 092/144] fix(claude_code_gateway): scope the protobuf body skip to the OTLP routes and match the metrics middleware on the route path --- .../anthropic_endpoints/gateway_endpoints.py | 14 +++- .../proxy/common_utils/http_parsing_utils.py | 10 +-- .../middleware/prometheus_auth_middleware.py | 9 ++- .../test_gateway_endpoints.py | 7 +- .../common_utils/test_http_parsing_utils.py | 8 +-- .../test_prometheus_auth_middleware.py | 68 +++++++++++++++++++ 6 files changed, 97 insertions(+), 19 deletions(-) diff --git a/litellm/proxy/anthropic_endpoints/gateway_endpoints.py b/litellm/proxy/anthropic_endpoints/gateway_endpoints.py index cc3106fce53..08579186f5e 100644 --- a/litellm/proxy/anthropic_endpoints/gateway_endpoints.py +++ b/litellm/proxy/anthropic_endpoints/gateway_endpoints.py @@ -34,6 +34,7 @@ from litellm.constants import ( ) from litellm.proxy.anthropic_endpoints.endpoints import anthropic_response, count_tokens from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.common_utils.http_parsing_utils import _safe_set_request_parsed_body GATEWAY_PREFIX: Final = "/claude_code_gateway" _DEVICE_CODE_GRANT: Final = "urn:ietf:params:oauth:grant-type:device_code" @@ -349,21 +350,28 @@ async def managed_settings(request: Request) -> Response: return Response(content=body.model_dump_json(), media_type="application/json", headers=headers) +async def _skip_otlp_body_parsing(request: Request) -> None: + _safe_set_request_parsed_body(request=request, parsed_body={}) + + +_OTLP_AUTHENTICATED: Final = (Depends(_skip_otlp_body_parsing), *_AUTHENTICATED) + + def _accept_otlp() -> Response: ensure_gateway_enabled() return Response(status_code=200) -@router.post("/v1/metrics", include_in_schema=False, dependencies=_AUTHENTICATED) +@router.post("/v1/metrics", include_in_schema=False, dependencies=_OTLP_AUTHENTICATED) async def otlp_metrics() -> Response: return _accept_otlp() -@router.post("/v1/logs", include_in_schema=False, dependencies=_AUTHENTICATED) +@router.post("/v1/logs", include_in_schema=False, dependencies=_OTLP_AUTHENTICATED) async def otlp_logs() -> Response: return _accept_otlp() -@router.post("/v1/traces", include_in_schema=False, dependencies=_AUTHENTICATED) +@router.post("/v1/traces", include_in_schema=False, dependencies=_OTLP_AUTHENTICATED) async def otlp_traces() -> Response: return _accept_otlp() diff --git a/litellm/proxy/common_utils/http_parsing_utils.py b/litellm/proxy/common_utils/http_parsing_utils.py index 592060e84ee..f5b6a0a766d 100644 --- a/litellm/proxy/common_utils/http_parsing_utils.py +++ b/litellm/proxy/common_utils/http_parsing_utils.py @@ -18,8 +18,6 @@ from litellm.types.router import Deployment _FORM_CONTENT_TYPES: Final[frozenset[str]] = frozenset({"application/x-www-form-urlencoded", "multipart/form-data"}) -_PROTOBUF_CONTENT_TYPES: Final[frozenset[str]] = frozenset({"application/x-protobuf", "application/protobuf"}) - _ANNOTATION_QUALIFIERS: Final[frozenset[object]] = frozenset({Annotated, NotRequired, ReadOnly, Required}) @@ -46,10 +44,6 @@ def is_json_content_type(content_type: str) -> bool: return _normalize_media_type(content_type) == "application/json" -def _is_protobuf_content_type(content_type: str) -> bool: - return _normalize_media_type(content_type) in _PROTOBUF_CONTENT_TYPES - - def _unqualified(annotation: object) -> object: """Which qualifiers ``get_type_hints`` already stripped varies by interpreter version, so peel them all.""" if get_origin(annotation) not in _ANNOTATION_QUALIFIERS: @@ -139,9 +133,7 @@ async def _read_request_body(request: Request | None) -> dict: _request_headers: Final[dict] = _safe_get_request_headers(request=request) content_type: Final = _request_headers.get("content-type", "") - if _is_protobuf_content_type(content_type): - parsed_body = {} - elif _is_form_content_type(content_type): + if _is_form_content_type(content_type): try: form_data: Final = await request.form() except Exception as e: diff --git a/litellm/proxy/middleware/prometheus_auth_middleware.py b/litellm/proxy/middleware/prometheus_auth_middleware.py index 36818a8cfbd..ebdd3e92bb2 100644 --- a/litellm/proxy/middleware/prometheus_auth_middleware.py +++ b/litellm/proxy/middleware/prometheus_auth_middleware.py @@ -7,6 +7,7 @@ from collections.abc import MutableMapping from typing import Any, Final from fastapi import Request +from starlette.routing import get_route_path from starlette.types import ASGIApp, Receive, Scope, Send import litellm @@ -15,6 +16,12 @@ from litellm.proxy.auth.user_api_key_auth import user_api_key_auth # Cache the header name at module level to avoid repeated enum attribute access _AUTHORIZATION_HEADER: Final = SpecialHeaders.openai_authorization.value # "Authorization" +_METRICS_MOUNT: Final = "/metrics" + + +def _is_metrics_route(scope: Scope) -> bool: + route_path: Final = get_route_path(scope) + return route_path == _METRICS_MOUNT or route_path.startswith(_METRICS_MOUNT + "/") class PrometheusAuthMiddleware: @@ -36,7 +43,7 @@ class PrometheusAuthMiddleware: async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None: # Fast path: only inspect HTTP requests; pass through websocket/lifespan immediately - if scope["type"] != "http" or "/metrics" not in scope.get("path", ""): + if scope["type"] != "http" or not _is_metrics_route(scope): await self.app(scope, receive, send) return diff --git a/tests/test_litellm/proxy/anthropic_endpoints/test_gateway_endpoints.py b/tests/test_litellm/proxy/anthropic_endpoints/test_gateway_endpoints.py index 158fe253796..e49047634bc 100644 --- a/tests/test_litellm/proxy/anthropic_endpoints/test_gateway_endpoints.py +++ b/tests/test_litellm/proxy/anthropic_endpoints/test_gateway_endpoints.py @@ -21,6 +21,7 @@ from litellm.caching.dual_cache import DualCache from litellm.proxy._types import ProxyException from litellm.proxy.anthropic_endpoints import gateway_endpoints from litellm.proxy.management_endpoints.ui_sso import _get_cli_sso_flow_cache_key, _set_cli_sso_flow +from litellm.proxy.middleware.prometheus_auth_middleware import PrometheusAuthMiddleware _DEVICE_CODE_GRANT: Final = "urn:ietf:params:oauth:grant-type:device_code" _MASTER_KEY: Final = "sk-master-key" @@ -107,6 +108,7 @@ def _gateway_env( session_cache: Final = cache or DualCache(default_in_memory_ttl=600) app: Final = FastAPI() + app.add_middleware(PrometheusAuthMiddleware) app.include_router(gateway_endpoints.router) async def _fake_auth() -> object: @@ -378,10 +380,11 @@ def test_otlp_endpoints_404_when_disabled(signal: str): assert resp.status_code == 404 -def test_otlp_protobuf_body_is_accepted_through_real_auth(): +@pytest.mark.parametrize("signal", ["metrics", "logs", "traces"]) +def test_otlp_protobuf_body_is_accepted_through_real_auth(signal: str): with _gateway_env(real_auth=True) as (client, _): resp = client.post( - "/claude_code_gateway/v1/metrics", + f"/claude_code_gateway/v1/{signal}", content=_PROTOBUF_BODY, headers={"Authorization": f"Bearer {_MASTER_KEY}", "Content-Type": "application/x-protobuf"}, ) diff --git a/tests/test_litellm/proxy/common_utils/test_http_parsing_utils.py b/tests/test_litellm/proxy/common_utils/test_http_parsing_utils.py index bd9912a96ac..7929a0b21af 100644 --- a/tests/test_litellm/proxy/common_utils/test_http_parsing_utils.py +++ b/tests/test_litellm/proxy/common_utils/test_http_parsing_utils.py @@ -574,10 +574,10 @@ async def test_lone_surrogate_escape_is_rejected_with_400(content: bytes): @pytest.mark.asyncio -@pytest.mark.parametrize("media_type", ["application/x-protobuf", "application/protobuf"]) -async def test_protobuf_body_is_left_unparsed(media_type: str): - request = _starlette_request(b"\x0a\x05hello\x12\x03{{{", media_type) - assert await _read_request_body(request) == {} +@pytest.mark.parametrize("media_type", ["application/x-protobuf", "application/protobuf", "application/octet-stream"]) +async def test_json_body_under_a_binary_content_type_is_still_parsed(media_type: str): + request = _starlette_request(b'{"model": "claude-sonnet-5"}', media_type) + assert await _read_request_body(request) == {"model": "claude-sonnet-5"} @pytest.mark.asyncio diff --git a/tests/test_litellm/proxy/middleware/test_prometheus_auth_middleware.py b/tests/test_litellm/proxy/middleware/test_prometheus_auth_middleware.py index 1d0c0f90fd1..beb841878d5 100644 --- a/tests/test_litellm/proxy/middleware/test_prometheus_auth_middleware.py +++ b/tests/test_litellm/proxy/middleware/test_prometheus_auth_middleware.py @@ -51,6 +51,14 @@ def app_with_middleware(): async def embeddings(): return {"msg": "embeddings OK"} + @app.post("/claude_code_gateway/v1/metrics") + async def gateway_telemetry(): + return {"msg": "gateway telemetry OK"} + + @app.get("/metrics/detail") + async def metrics_detail(): + return {"msg": "metrics detail OK"} + return app @@ -240,3 +248,63 @@ def test_non_metrics_requests_dont_trigger_auth(app_with_middleware, monkeypatch response = client.get("/embeddings") assert response.status_code == 200, response.text assert response.json() == {"msg": "embeddings OK"} + + +def test_gateway_telemetry_path_is_not_treated_as_the_metrics_endpoint(app_with_middleware, monkeypatch): + monkeypatch.setattr(litellm, "require_auth_for_metrics_endpoint", True) + + def should_not_be_called(*args, **kwargs): + raise Exception("Auth should not be called for the gateway telemetry route") + + monkeypatch.setattr( + "litellm.proxy.middleware.prometheus_auth_middleware.user_api_key_auth", + should_not_be_called, + ) + + client = TestClient(app_with_middleware) + + response = client.post("/claude_code_gateway/v1/metrics", content=b"\x0a\x05hello") + assert response.status_code == 200, response.text + assert response.json() == {"msg": "gateway telemetry OK"} + + +@pytest.mark.parametrize("path", ["/metrics", "/metrics/", "/metrics/detail"]) +def test_metrics_paths_still_require_auth(app_with_middleware, monkeypatch, path): + monkeypatch.setattr(litellm, "require_auth_for_metrics_endpoint", True) + + async def reject(*args, **kwargs): + raise Exception("Invalid API key") + + monkeypatch.setattr( + "litellm.proxy.middleware.prometheus_auth_middleware.user_api_key_auth", + reject, + ) + + client = TestClient(app_with_middleware) + + response = client.get(path) + assert response.status_code == 401, response.text + + +def test_metrics_under_a_root_path_still_requires_auth(monkeypatch): + monkeypatch.setattr(litellm, "require_auth_for_metrics_endpoint", True) + + async def reject(*args, **kwargs): + raise Exception("Invalid API key") + + monkeypatch.setattr( + "litellm.proxy.middleware.prometheus_auth_middleware.user_api_key_auth", + reject, + ) + + app = FastAPI(root_path="/litellm") + app.add_middleware(PrometheusAuthMiddleware) + + @app.get("/metrics") + async def metrics(): + return {"msg": "metrics OK"} + + client = TestClient(app, root_path="/litellm") + + response = client.get("/metrics") + assert response.status_code == 401, response.text From 0f54d76079003cc16c23f3570f459c9b7146a53a Mon Sep 17 00:00:00 2001 From: kerry Date: Sat, 19 Sep 2026 00:24:34 +0000 Subject: [PATCH 093/144] fix(timing): drop banned typing.cast from provider duration accounting Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../llm_response_utils/response_metadata.py | 5 ++--- litellm/litellm_core_utils/logging_utils.py | 14 ++++++-------- 2 files changed, 8 insertions(+), 11 deletions(-) diff --git a/litellm/litellm_core_utils/llm_response_utils/response_metadata.py b/litellm/litellm_core_utils/llm_response_utils/response_metadata.py index 3778ae1281f..780691b4696 100644 --- a/litellm/litellm_core_utils/llm_response_utils/response_metadata.py +++ b/litellm/litellm_core_utils/llm_response_utils/response_metadata.py @@ -1,6 +1,6 @@ import datetime from collections.abc import Mapping -from typing import Any, Final, cast +from typing import Any, Final import httpx @@ -49,8 +49,7 @@ def response_timing_metrics( if caching_details is not None and caching_details.get("cache_hit") is True else None ) - metadata_value: Final = get_litellm_metadata_from_kwargs(logging_obj.model_call_details) - metadata: Final = cast(dict[str, object], metadata_value) if isinstance(metadata_value, dict) else {} + metadata: Final[Mapping[str, object]] = get_litellm_metadata_from_kwargs(logging_obj.model_call_details) llm_api_duration_ms: Final = logging_obj.model_call_details.get("llm_api_duration_ms") if cache_duration_ms is not None: overhead_ms: float | None = total_response_time_ms - cache_duration_ms diff --git a/litellm/litellm_core_utils/logging_utils.py b/litellm/litellm_core_utils/logging_utils.py index 82bfb0efdb1..91cc13c8315 100644 --- a/litellm/litellm_core_utils/logging_utils.py +++ b/litellm/litellm_core_utils/logging_utils.py @@ -5,7 +5,7 @@ import re import time from collections.abc import Iterator, Mapping, Sequence from datetime import datetime -from typing import TYPE_CHECKING, Any, Final, cast +from typing import TYPE_CHECKING, Any, Final from litellm._logging import format_base64_size, verbose_logger from litellm.constants import ( @@ -287,13 +287,11 @@ def _set_duration_in_model_call_details( duration_ms: Final = (end_time - start_time).total_seconds() * 1000 if logging_obj and hasattr(logging_obj, "model_call_details"): logging_obj.model_call_details["llm_api_duration_ms"] = duration_ms - metadata_value: Final = get_litellm_metadata_from_kwargs(logging_obj.model_call_details) - if isinstance(metadata_value, dict): - metadata: Final = cast(dict[str, object], metadata_value) - existing_total: Final = metadata.get("llm_api_duration_ms_total") - metadata["llm_api_duration_ms_total"] = ( - existing_total if isinstance(existing_total, float) else 0.0 - ) + duration_ms + metadata: Final[dict[str, object]] = get_litellm_metadata_from_kwargs(logging_obj.model_call_details) + existing_total: Final = metadata.get("llm_api_duration_ms_total") + metadata["llm_api_duration_ms_total"] = ( + existing_total if isinstance(existing_total, float) else 0.0 + ) + duration_ms else: verbose_logger.debug("`logging_obj` not found - unable to track `llm_api_duration_ms") except Exception as e: From cdc0e57e93fc6f57b54268b0ca7ce2b783ce29b9 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 18 Sep 2026 17:24:52 -0700 Subject: [PATCH 094/144] fix(websearch_interception): surface a failed search as a web_search_tool_result_error block and end the turn --- .../websearch_interception/handler.py | 107 +++++------ .../websearch_interception/transformation.py | 77 +++++++- litellm/llms/anthropic/common_utils.py | 25 ++- litellm/llms/custom_httpx/llm_http_handler.py | 10 +- .../integrations/websearch_interception.py | 31 ++- .../test_websearch_agentic_loop_cap.py | 137 ++++++++++++- .../test_websearch_native_blocks.py | 181 ++++++++++++++++-- .../anthropic/test_anthropic_common_utils.py | 17 +- 8 files changed, 492 insertions(+), 93 deletions(-) diff --git a/litellm/integrations/websearch_interception/handler.py b/litellm/integrations/websearch_interception/handler.py index 587da997f94..29a586eaf20 100644 --- a/litellm/integrations/websearch_interception/handler.py +++ b/litellm/integrations/websearch_interception/handler.py @@ -44,6 +44,8 @@ from litellm.types.integrations.custom_logger import ( from litellm.types.integrations.websearch_interception import ( AnthropicSearchQuery, AnthropicServerToolUseBlock, + SearchFailed, + SearchOutcome, WebSearchInterceptionConfig, ) from litellm.types.llms.anthropic import AnthropicThinkingParam @@ -332,16 +334,8 @@ class WebSearchInterceptionLogger(CustomLogger): None, ) - # Execute search — keep the structured SearchResponse so the native - # block can carry per-result url/title/page_age. - try: - if kwargs is None: - search_result_text, structured = await self._execute_search(query) - else: - search_result_text, structured = await self._execute_search(query, kwargs=kwargs) - except Exception as e: - verbose_logger.error("WebSearchInterception: Short-circuit search failed: %s", e) - search_result_text, structured = f"Search failed: {e}", None + outcome: Final = await self._short_circuit_search_outcome(query, kwargs=kwargs) + search_result_text: Final = WebSearchTransformation.search_outcome_text(outcome) content: Final[list[dict[str, object]]] = [] if native_tool is not None: @@ -355,12 +349,7 @@ class WebSearchInterceptionLogger(CustomLogger): "input": {"query": query}, } ) - content.append( - WebSearchTransformation.build_web_search_tool_result_block( - tool_use_id=tool_use_id, - search_response=structured, - ) - ) + content.append(WebSearchTransformation.build_web_search_outcome_block(tool_use_id=tool_use_id, outcome=outcome)) # Keep the text block so non-native short-circuit callers (Claude Code, # github_copilot, etc.) see the same payload they always have. content.append({"type": "text", "text": search_result_text}) @@ -934,7 +923,7 @@ class WebSearchInterceptionLogger(CustomLogger): tool_calls: Final = tools["tool_calls"] thinking_blocks: Final = tools.get("thinking_blocks", []) - request_patch, structured_results = await self._build_anthropic_request_patch( + request_patch, search_outcomes = await self._build_anthropic_request_patch( model=model, messages=messages, tool_calls=tool_calls, @@ -953,17 +942,19 @@ class WebSearchInterceptionLogger(CustomLogger): # pre-build the Anthropic-native ``web_search_tool_result`` blocks now # (while we still have the structured SearchResponse list) and stash # them on plan metadata for the post-hook to inject. - if kwargs.get(WEBSEARCH_EMIT_NATIVE_BLOCKS_KEY): - metadata[WEBSEARCH_NATIVE_BLOCKS_METADATA_KEY] = self._build_native_result_blocks( - tool_calls=tool_calls, - structured_results=structured_results, - ) + if not kwargs.get(WEBSEARCH_EMIT_NATIVE_BLOCKS_KEY): + return AgenticLoopPlan(run_agentic_loop=True, request_patch=request_patch, metadata=metadata) - return AgenticLoopPlan( - run_agentic_loop=True, - request_patch=request_patch, - metadata=metadata, + metadata[WEBSEARCH_NATIVE_BLOCKS_METADATA_KEY] = self._build_native_result_blocks( + tool_calls=tool_calls, + search_outcomes=search_outcomes, ) + every_search_failed: Final = bool(search_outcomes) and all( + isinstance(outcome, SearchFailed) for outcome in search_outcomes + ) + if every_search_failed: + return AgenticLoopPlan(run_agentic_loop=False, terminate=True, stop_reason="web_search_failed", metadata=metadata) + return AgenticLoopPlan(run_agentic_loop=True, request_patch=request_patch, metadata=metadata) async def async_post_agentic_loop_response_hook( self, @@ -992,7 +983,7 @@ class WebSearchInterceptionLogger(CustomLogger): @staticmethod def _build_native_result_blocks( tool_calls: list[dict], - structured_results: list[SearchResponse | None], + search_outcomes: Sequence[SearchOutcome], ) -> tuple[Mapping[str, object], ...]: """ Build a ``server_tool_use`` + ``web_search_tool_result`` pair per tool_call. @@ -1004,10 +995,10 @@ class WebSearchInterceptionLogger(CustomLogger): """ return tuple( block - for i, tool_call in enumerate(tool_calls) + for tool_call, outcome in zip(tool_calls, search_outcomes, strict=True) for block in WebSearchInterceptionLogger._native_result_pair( query=WebSearchInterceptionLogger._tool_call_query(tool_call), - search_response=structured_results[i] if i < len(structured_results) else None, + outcome=outcome, ) ) @@ -1022,15 +1013,12 @@ class WebSearchInterceptionLogger(CustomLogger): @staticmethod def _native_result_pair( query: str, - search_response: SearchResponse | None, + outcome: SearchOutcome, ) -> tuple[Mapping[str, object], Mapping[str, object]]: tool_use_id: Final = f"srvtoolu_{uuid.uuid4().hex}" return ( AnthropicServerToolUseBlock(id=tool_use_id, input=AnthropicSearchQuery(query=query)).model_dump(), - WebSearchTransformation.build_web_search_tool_result_block( - tool_use_id=tool_use_id, - search_response=search_response, - ), + WebSearchTransformation.build_web_search_outcome_block(tool_use_id=tool_use_id, outcome=outcome), ) @staticmethod @@ -1306,7 +1294,7 @@ class WebSearchInterceptionLogger(CustomLogger): kwargs: Mapping[str, object], ) -> "AnthropicMessagesResponse | AsyncIterator[object]": """Legacy path: execute search + build patch + run follow-up call.""" - request_patch, structured_results = await self._build_anthropic_request_patch( + request_patch, search_outcomes = await self._build_anthropic_request_patch( model=model, messages=messages, tool_calls=tool_calls, @@ -1344,7 +1332,7 @@ class WebSearchInterceptionLogger(CustomLogger): if kwargs.get(WEBSEARCH_EMIT_NATIVE_BLOCKS_KEY): native_blocks: Final = self._build_native_result_blocks( tool_calls=tool_calls, - structured_results=structured_results, + search_outcomes=search_outcomes, ) response = self._inject_native_blocks(response, native_blocks) @@ -1359,15 +1347,14 @@ class WebSearchInterceptionLogger(CustomLogger): anthropic_messages_optional_request_params: dict, logging_obj: "LiteLLMLoggingObj | None", kwargs: dict, - ) -> tuple[AgenticLoopRequestPatch, list[SearchResponse | None]]: + ) -> tuple[AgenticLoopRequestPatch, tuple[SearchOutcome, ...]]: """ Execute litellm.search() and build follow-up request patch. - Returns the patch alongside the parallel list of structured - ``SearchResponse`` objects (one per tool_call, ``None`` when the - search failed or the tool_call had no query). The caller uses these - to optionally build Anthropic-native ``web_search_tool_result`` - content blocks for the final response. + Returns the patch alongside the parallel tuple of search outcomes (one + per tool_call). The caller uses these to optionally build + Anthropic-native ``web_search_tool_result`` content blocks for the + final response and to decide whether a follow-up call is worth making. """ # Extract search queries from tool_use blocks @@ -1385,27 +1372,10 @@ class WebSearchInterceptionLogger(CustomLogger): # Execute searches in parallel verbose_logger.debug("WebSearchInterception: Executing %s search(es) in parallel", len(search_tasks)) search_results: Final = await asyncio.gather(*search_tasks, return_exceptions=True) - - # Split the gathered (text, structured) tuples into two parallel lists. - # The text list feeds the follow-up model call; the structured list - # is returned to the caller for native-block emission. - final_search_results: Final[list[str]] = [] - structured_results: Final[list[SearchResponse | None]] = [] - for i, result in enumerate(search_results): - if isinstance(result, Exception): - verbose_logger.error("WebSearchInterception: Search %s failed with error: %s", i, result) - final_search_results.append(f"Search failed: {result}") - structured_results.append(None) - elif isinstance(result, tuple) and len(result) == 2: - text_value, structured_value = result - final_search_results.append(cast(str, text_value) if isinstance(text_value, str) else str(text_value)) - structured_results.append(structured_value if isinstance(structured_value, SearchResponse) else None) - else: - # Defensive: legacy callers / unexpected shape — preserve text, - # drop structure. - verbose_logger.debug("WebSearchInterception: Unexpected result type %s at index %s", type(result), i) - final_search_results.append(str(result)) - structured_results.append(None) + search_outcomes: Final = tuple(WebSearchTransformation.search_outcome(result) for result in search_results) + final_search_results: Final = tuple( + WebSearchTransformation.search_outcome_text(outcome) for outcome in search_outcomes + ) # Build assistant and user messages using transformation assistant_message, user_message = WebSearchTransformation.transform_response( @@ -1449,7 +1419,16 @@ class WebSearchInterceptionLogger(CustomLogger): optional_params=optional_params_without_max_tokens, kwargs=kwargs_for_followup, ) - return patch, structured_results + return patch, search_outcomes + + async def _short_circuit_search_outcome(self, query: str, kwargs: Mapping[str, object] | None) -> SearchOutcome: + try: + result: Final = ( + await self._execute_search(query) if kwargs is None else await self._execute_search(query, kwargs=kwargs) + ) + except Exception as e: + return WebSearchTransformation.search_outcome(e) + return WebSearchTransformation.search_outcome(result) async def _execute_search( self, query: str, kwargs: Mapping[str, object] | None = None diff --git a/litellm/integrations/websearch_interception/transformation.py b/litellm/integrations/websearch_interception/transformation.py index fe4b6583c55..47af73570fc 100644 --- a/litellm/integrations/websearch_interception/transformation.py +++ b/litellm/integrations/websearch_interception/transformation.py @@ -5,11 +5,21 @@ Transforms between Anthropic/OpenAI tool_use format and LiteLLM search format. """ import json +from collections.abc import Sequence from typing import Any, Final +from typing_extensions import assert_never + from litellm._logging import verbose_logger from litellm.constants import LITELLM_WEB_SEARCH_TOOL_NAME +from litellm.exceptions import BadRequestError, RateLimitError from litellm.llms.base_llm.search.transformation import SearchResponse +from litellm.types.integrations.websearch_interception import ( + SearchFailed, + SearchOutcome, + SearchSucceeded, + WebSearchToolResultErrorCode, +) class WebSearchTransformation: @@ -280,7 +290,7 @@ class WebSearchTransformation: @staticmethod def transform_response( tool_calls: list[dict], - search_results: list[str], + search_results: Sequence[str], response_format: str = "anthropic", thinking_blocks: list[dict] | None = None, ) -> tuple[dict, dict | list[dict]]: @@ -314,7 +324,7 @@ class WebSearchTransformation: @staticmethod def _transform_response_anthropic( tool_calls: list[dict], - search_results: list[str], + search_results: Sequence[str], thinking_blocks: list[dict] | None = None, ) -> tuple[dict, dict]: """Transform to Anthropic format (single user message with tool_result blocks)""" @@ -364,7 +374,7 @@ class WebSearchTransformation: @staticmethod def _transform_response_openai( tool_calls: list[dict], - search_results: list[str], + search_results: Sequence[str], ) -> tuple[dict, list[dict]]: """Transform to OpenAI format (assistant with tool_calls, separate tool messages)""" # Build assistant message with tool_calls @@ -456,6 +466,67 @@ class WebSearchTransformation: "content": items, } + @staticmethod + def build_web_search_tool_result_error_block( + tool_use_id: str, + error_code: WebSearchToolResultErrorCode, + ) -> dict[str, object]: + return { + "type": "web_search_tool_result", + "tool_use_id": tool_use_id, + "content": {"type": "web_search_tool_result_error", "error_code": error_code}, + } + + @staticmethod + def build_web_search_outcome_block(tool_use_id: str, outcome: SearchOutcome) -> dict[str, object]: + match outcome: + case SearchSucceeded(response=response): + return WebSearchTransformation.build_web_search_tool_result_block( + tool_use_id=tool_use_id, + search_response=response, + ) + case SearchFailed(error_code=error_code): + return WebSearchTransformation.build_web_search_tool_result_error_block( + tool_use_id=tool_use_id, + error_code=error_code, + ) + case _: + assert_never(outcome) + + @staticmethod + def search_error_code(error: BaseException) -> WebSearchToolResultErrorCode: + match error: + case RateLimitError(): + return "too_many_requests" + case BadRequestError(): + return "invalid_tool_input" + case _: + return "unavailable" + + @staticmethod + def search_outcome(result: object) -> SearchOutcome: + match result: + case BaseException(): + verbose_logger.error("WebSearchInterception: Search failed with error: %s", result) + return SearchFailed(error_code=WebSearchTransformation.search_error_code(result), message=str(result)) + case (str() as text, SearchResponse() as response): + return SearchSucceeded(text=text, response=response) + case (str() as text, None): + return SearchSucceeded(text=text, response=None) + case _: + verbose_logger.debug("WebSearchInterception: Unexpected search result type %s", type(result)) + return SearchSucceeded(text=str(result), response=None) + + @staticmethod + def search_outcome_text(outcome: SearchOutcome) -> str: + match outcome: + case SearchSucceeded(text=text): + return text + case SearchFailed(message=message): + return f"Search failed: {message}" + case _: + assert_never(outcome) + @staticmethod def format_search_response(result: SearchResponse) -> str: """ diff --git a/litellm/llms/anthropic/common_utils.py b/litellm/llms/anthropic/common_utils.py index d35a9372058..06561faae8c 100644 --- a/litellm/llms/anthropic/common_utils.py +++ b/litellm/llms/anthropic/common_utils.py @@ -1410,12 +1410,19 @@ class _ReplayedWebSearchResult(BaseModel): encrypted_content: str = "" +class _ReplayedWebSearchToolResultError(BaseModel): + model_config = ConfigDict(extra="allow") + + type: Literal["web_search_tool_result_error"] + error_code: str = "" + + class _ReplayedWebSearchToolResult(BaseModel): model_config = ConfigDict(extra="allow") type: Literal["web_search_tool_result"] tool_use_id: str - content: tuple[_ReplayedWebSearchResult, ...] + content: tuple[_ReplayedWebSearchResult, ...] | _ReplayedWebSearchToolResultError class _ReplayedServerToolUse(BaseModel): @@ -1441,15 +1448,17 @@ def _flattenable_web_search_tool_result(block: object) -> _ReplayedWebSearchTool ``encrypted_content``, else None for anything Anthropic itself issued. An empty ``content`` list is flattenable too. It is what the interceptor emits - when a search legitimately returns nothing and when a search raises, and it - carries neither evidence to preserve nor an ``encrypted_content`` to respect, - so leaving it in place only buys the 400 this whole function exists to avoid. + when a search legitimately returns nothing, and it carries neither evidence to + preserve nor an ``encrypted_content`` to respect, so leaving it in place only + buys the 400 this whole function exists to avoid. The same goes for the + ``web_search_tool_result_error`` object the interceptor emits when a search + raises: it never carries ``encrypted_content``, so it is flattened as well. """ try: parsed: Final = _WEB_SEARCH_TOOL_RESULT_ADAPTER.validate_python(block) except ValidationError: return None - if any(result.encrypted_content for result in parsed.content): + if isinstance(parsed.content, tuple) and any(result.encrypted_content for result in parsed.content): return None return parsed @@ -1461,8 +1470,12 @@ def _replayed_server_tool_use(block: object) -> _ReplayedServerToolUse | None: return None -def _render_web_search_results(query: str, results: tuple[_ReplayedWebSearchResult, ...]) -> str: +def _render_web_search_results( + query: str, results: tuple[_ReplayedWebSearchResult, ...] | _ReplayedWebSearchToolResultError +) -> str: header: Final = f"Web search results for '{query}':" if query else "Web search results:" + if isinstance(results, _ReplayedWebSearchToolResultError): + return f"{header}\n\nSearch failed: {results.error_code or 'unavailable'}" if not results: return f"{header}\n\nNo results were returned." body: Final = "\n\n".join( diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 98fe0014386..7743bfad56c 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -5905,7 +5905,15 @@ class BaseLLMHTTPHandler: callback.__class__.__name__, plan.stop_reason, ) - return self._maybe_wrap_in_fake_stream(response, logging_obj, api_surface) + return self._maybe_wrap_in_fake_stream( + await callback.async_post_agentic_loop_response_hook( + response=self._finalize_refused_agentic_response(response=response, tool_calls=tool_calls), + plan=plan, + kwargs=kwargs_with_provider, + ), + logging_obj, + api_surface, + ) if not plan.run_agentic_loop: continue diff --git a/litellm/types/integrations/websearch_interception.py b/litellm/types/integrations/websearch_interception.py index 7926b9eee0a..a7c2e7f2315 100644 --- a/litellm/types/integrations/websearch_interception.py +++ b/litellm/types/integrations/websearch_interception.py @@ -2,11 +2,15 @@ Type definitions for WebSearch Interception integration. """ -from typing import Literal, TypedDict +from dataclasses import dataclass +from typing import TYPE_CHECKING, Literal, TypeAlias, TypedDict from pydantic import BaseModel from typing_extensions import ReadOnly +if TYPE_CHECKING: + from litellm.llms.base_llm.search.transformation import SearchResponse + class AnthropicSearchQuery(BaseModel): """``input`` of an Anthropic ``server_tool_use`` block for a web search.""" @@ -27,6 +31,31 @@ class AnthropicServerToolUseBlock(BaseModel): input: AnthropicSearchQuery +WebSearchToolResultErrorCode: TypeAlias = Literal[ + "invalid_tool_input", + "unavailable", + "max_uses_exceeded", + "too_many_requests", + "query_too_long", + "request_too_large", +] + + +@dataclass(frozen=True, slots=True) +class SearchSucceeded: + text: str + response: "SearchResponse | None" + + +@dataclass(frozen=True, slots=True) +class SearchFailed: + error_code: WebSearchToolResultErrorCode + message: str + + +SearchOutcome: TypeAlias = SearchSucceeded | SearchFailed + + class WebSearchInterceptionConfig(TypedDict, total=False): """ Configuration parameters for WebSearchInterceptionLogger. diff --git a/tests/test_litellm/integrations/websearch_interception/test_websearch_agentic_loop_cap.py b/tests/test_litellm/integrations/websearch_interception/test_websearch_agentic_loop_cap.py index 40fd8c4e9e6..57e3ba59456 100644 --- a/tests/test_litellm/integrations/websearch_interception/test_websearch_agentic_loop_cap.py +++ b/tests/test_litellm/integrations/websearch_interception/test_websearch_agentic_loop_cap.py @@ -12,19 +12,23 @@ config.yaml through to the settings the loop actually reads. """ import json -from unittest.mock import MagicMock +from unittest.mock import MagicMock, patch import pytest import litellm +from litellm.exceptions import AuthenticationError, RateLimitError from litellm.integrations.custom_logger import CustomLogger from litellm.integrations.websearch_interception.handler import ( + WEBSEARCH_EMIT_NATIVE_BLOCKS_KEY, WebSearchInterceptionLogger, ) +from litellm.integrations.websearch_interception.tools import get_litellm_web_search_tool +from litellm.litellm_core_utils.agentic_loop_settings import DEFAULT_MAX_AGENTIC_LOOPS from litellm.llms.anthropic.experimental_pass_through.messages.fake_stream_iterator import ( FakeAnthropicMessagesStreamIterator, ) -from litellm.litellm_core_utils.agentic_loop_settings import DEFAULT_MAX_AGENTIC_LOOPS +from litellm.llms.base_llm.search.transformation import SearchResponse, SearchResult from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler from litellm.secret_managers.main import get_secret from litellm.types.integrations.custom_logger import ( @@ -490,6 +494,135 @@ class TestOuterFramePostHookStillRuns: assert result["stop_reason"] == "end_turn" +def _response_asking_for_searches(*queries: str) -> dict: + return { + "id": "msg_123", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-4-5", + "content": [ + {"id": f"toolu_internal_{index}", "type": "tool_use", "name": INTERNAL_TOOL_NAME, "input": {"query": query}} + for index, query in enumerate(queries, start=1) + ], + "stop_reason": "tool_use", + "usage": {"input_tokens": 10, "output_tokens": 5}, + } + + +class TestFailedSearchEndsTheTurn: + """ + A search that failed used to come back to the client as an empty successful + ``web_search_tool_result`` while the model was re-asked the same query until + the loop cap tripped. When the client sent a native web search tool, the + turn now ends after the first failed search, with Anthropic's + ``web_search_tool_result_error`` object in the tool result and no follow-up + model call. An iteration where some search still succeeded keeps its + follow-up call. + """ + + def setup_method(self): + self.handler = BaseLLMHTTPHandler() + self.logger = WebSearchInterceptionLogger(enabled_providers=["anthropic"]) + self.followup_calls: list[dict] = [] + + async def _fake_acreate(self, **call_kwargs): + self.followup_calls.append(call_kwargs) + return { + "id": "msg_followup", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-4-5", + "content": [{"type": "text", "text": "final answer"}], + "stop_reason": "end_turn", + "usage": {"input_tokens": 20, "output_tokens": 5}, + } + + async def _run(self, response: dict, converted_stream: bool = False): + return await self.handler._call_agentic_completion_hooks( + response=response, + model="claude-sonnet-4-5", + messages=[{"role": "user", "content": "who won the world cup"}], + anthropic_messages_provider_config=MagicMock(), + anthropic_messages_optional_request_params={"tools": [get_litellm_web_search_tool()]}, + logging_obj=_logging_obj(self.logger, converted_stream=converted_stream), + stream=False, + custom_llm_provider="anthropic", + kwargs={"_agentic_loop_depth": 0, "max_agentic_loops": 3, WEBSEARCH_EMIT_NATIVE_BLOCKS_KEY: True}, + ) + + @pytest.mark.asyncio + async def test_all_failed_iteration_ends_the_turn_without_a_follow_up_call(self, monkeypatch): + monkeypatch.setattr("litellm.anthropic_interface.messages.acreate", self._fake_acreate) + + with patch.object( + self.logger, + "_execute_search", + side_effect=AuthenticationError("401 Unauthorized", llm_provider="tavily", model="tavily"), + ): + result = await self._run(_response_asking_for_searches("who won the world cup")) + + assert self.followup_calls == [] + assert result["stop_reason"] == "end_turn" + assert INTERNAL_TOOL_NAME not in _tool_use_names(result) + assert _block_types(result) == ["server_tool_use", "web_search_tool_result"] + server_tool_use, tool_result = result["content"] + assert server_tool_use["id"].startswith("srvtoolu_") + assert server_tool_use["input"] == {"query": "who won the world cup"} + assert tool_result["tool_use_id"] == server_tool_use["id"] + assert tool_result["content"] == {"type": "web_search_tool_result_error", "error_code": "unavailable"} + + @pytest.mark.asyncio + async def test_all_failed_iteration_streams_the_error_block(self, monkeypatch): + monkeypatch.setattr("litellm.anthropic_interface.messages.acreate", self._fake_acreate) + + with patch.object( + self.logger, + "_execute_search", + side_effect=AuthenticationError("401 Unauthorized", llm_provider="tavily", model="tavily"), + ): + result = await self._run(_response_asking_for_searches("who won the world cup"), converted_stream=True) + + assert self.followup_calls == [] + assert isinstance(result, FakeAnthropicMessagesStreamIterator) + events = _stream_events(result.response) + started = [event["content_block"] for event in events if event["type"] == "content_block_start"] + assert [block["type"] for block in started] == ["server_tool_use", "web_search_tool_result"] + assert started[1]["tool_use_id"] == started[0]["id"] + assert started[1]["content"] == {"type": "web_search_tool_result_error", "error_code": "unavailable"} + assert [event["delta"]["stop_reason"] for event in events if event["type"] == "message_delta"] == ["end_turn"] + + @pytest.mark.asyncio + async def test_mixed_iteration_keeps_the_follow_up_call(self, monkeypatch): + monkeypatch.setattr("litellm.anthropic_interface.messages.acreate", self._fake_acreate) + + async def search(query, kwargs=None): + if query == "fails": + raise RateLimitError("slow down", llm_provider="tavily", model="tavily") + found = SearchResult(title="Result", url="https://example.com", snippet="A result.", date=None) + return ("Title: Result\nURL: https://example.com", SearchResponse(results=[found])) + + with patch.object(self.logger, "_execute_search", side_effect=search): + result = await self._run(_response_asking_for_searches("fails", "works")) + + assert len(self.followup_calls) == 1 + tool_results = self.followup_calls[0]["messages"][-1]["content"] + assert [block["type"] for block in tool_results] == ["tool_result", "tool_result"] + assert tool_results[0]["content"] == "Search failed: litellm.RateLimitError: slow down" + assert tool_results[1]["content"] == "Title: Result\nURL: https://example.com" + assert result["stop_reason"] == "end_turn" + assert _block_types(result) == [ + "server_tool_use", + "web_search_tool_result", + "server_tool_use", + "web_search_tool_result", + "text", + ] + assert result["content"][0]["input"] == {"query": "fails"} + assert result["content"][1]["content"] == {"type": "web_search_tool_result_error", "error_code": "too_many_requests"} + assert result["content"][2]["input"] == {"query": "works"} + assert result["content"][3]["content"][0]["url"] == "https://example.com" + + class TestMaxAgenticLoopsConfigKnob: def test_from_config_yaml_reads_the_knob(self): logger = WebSearchInterceptionLogger.from_config_yaml( diff --git a/tests/test_litellm/integrations/websearch_interception/test_websearch_native_blocks.py b/tests/test_litellm/integrations/websearch_interception/test_websearch_native_blocks.py index c859f9b2f55..291fb5a5941 100644 --- a/tests/test_litellm/integrations/websearch_interception/test_websearch_native_blocks.py +++ b/tests/test_litellm/integrations/websearch_interception/test_websearch_native_blocks.py @@ -10,6 +10,13 @@ from unittest.mock import AsyncMock, MagicMock, patch import pytest +from litellm.exceptions import ( + APIConnectionError, + AuthenticationError, + BadRequestError, + RateLimitError, + Timeout, +) from litellm.integrations.websearch_interception.handler import ( WEBSEARCH_EMIT_NATIVE_BLOCKS_KEY, WEBSEARCH_NATIVE_BLOCKS_METADATA_KEY, @@ -27,6 +34,10 @@ from litellm.types.integrations.custom_logger import ( AgenticLoopPlan, AgenticLoopRequestPatch, ) +from litellm.types.integrations.websearch_interception import ( + SearchFailed, + SearchSucceeded, +) def _make_search_response() -> SearchResponse: @@ -48,6 +59,10 @@ def _make_search_response() -> SearchResponse: ) +def _succeeded_outcome() -> SearchSucceeded: + return SearchSucceeded(text="Title: LiteLLM Docs\nURL: https://docs.litellm.ai/", response=_make_search_response()) + + class TestIsAnthropicNativeWebSearchTool: """The detector must match native tools without catching look-alikes.""" @@ -227,12 +242,10 @@ class TestBuildPlanAttachesBlocks: messages=[{"role": "user", "content": "hi"}], max_tokens=1024, ) - structured = [_make_search_response()] - with patch.object( logger, "_build_anthropic_request_patch", - new=AsyncMock(return_value=(patch_obj, structured)), + new=AsyncMock(return_value=(patch_obj, (_succeeded_outcome(),))), ): plan = await logger.async_build_agentic_loop_plan( tools={"tool_calls": tool_calls, "thinking_blocks": []}, @@ -277,7 +290,7 @@ class TestBuildPlanAttachesBlocks: with patch.object( logger, "_build_anthropic_request_patch", - new=AsyncMock(return_value=(patch_obj, [_make_search_response()])), + new=AsyncMock(return_value=(patch_obj, (_succeeded_outcome(),))), ): plan = await logger.async_build_agentic_loop_plan( tools={"tool_calls": tool_calls, "thinking_blocks": []}, @@ -294,6 +307,145 @@ class TestBuildPlanAttachesBlocks: assert WEBSEARCH_NATIVE_BLOCKS_METADATA_KEY not in plan.metadata +class TestFailedSearchOutcome: + """A search that raises becomes a ``web_search_tool_result_error`` block, coded by exception type.""" + + @pytest.mark.parametrize( + ("error", "expected_code"), + [ + (RateLimitError("slow down", llm_provider="tavily", model="tavily"), "too_many_requests"), + (BadRequestError("bad query", model="tavily", llm_provider="tavily"), "invalid_tool_input"), + (AuthenticationError("401 Unauthorized", llm_provider="tavily", model="tavily"), "unavailable"), + (APIConnectionError("connection refused", llm_provider="tavily", model="tavily"), "unavailable"), + (Timeout("timed out", model="tavily", llm_provider="tavily"), "unavailable"), + (RuntimeError("boom"), "unavailable"), + ], + ) + def test_error_block_carries_the_mapped_error_code(self, error, expected_code): + outcome = WebSearchTransformation.search_outcome(error) + + assert outcome == SearchFailed(error_code=expected_code, message=str(error)) + assert WebSearchTransformation.build_web_search_outcome_block("srvtoolu_x", outcome) == { + "type": "web_search_tool_result", + "tool_use_id": "srvtoolu_x", + "content": {"type": "web_search_tool_result_error", "error_code": expected_code}, + } + assert WebSearchTransformation.search_outcome_text(outcome) == f"Search failed: {error}" + + def test_succeeded_outcome_still_yields_result_items(self): + outcome = WebSearchTransformation.search_outcome(("Title: x", _make_search_response())) + + assert outcome == SearchSucceeded(text="Title: x", response=_make_search_response()) + block = WebSearchTransformation.build_web_search_outcome_block("srvtoolu_x", outcome) + assert [item["type"] for item in block["content"]] == ["web_search_result", "web_search_result"] + assert block["content"][0]["url"] == "https://docs.litellm.ai/" + assert WebSearchTransformation.search_outcome_text(outcome) == "Title: x" + + @pytest.mark.asyncio + async def test_all_failed_iteration_terminates_when_native_blocks_are_emitted(self): + logger = WebSearchInterceptionLogger(enabled_providers=["bedrock"]) + tool_calls = [ + {"id": "toolu_one", "type": "tool_use", "name": "litellm_web_search", "input": {"query": "q1"}}, + {"id": "toolu_two", "type": "tool_use", "name": "litellm_web_search", "input": {"query": "q2"}}, + ] + + with patch.object( + logger, + "_execute_search", + side_effect=AuthenticationError("401 Unauthorized", llm_provider="tavily", model="tavily"), + ): + plan = await logger.async_build_agentic_loop_plan( + tools={"tool_calls": tool_calls, "thinking_blocks": []}, + model="bedrock/claude", + messages=[{"role": "user", "content": "hi"}], + response=MagicMock(), + anthropic_messages_provider_config=None, + anthropic_messages_optional_request_params={}, + logging_obj=MagicMock(model_call_details={}), + stream=False, + kwargs={WEBSEARCH_EMIT_NATIVE_BLOCKS_KEY: True}, + ) + + assert plan.run_agentic_loop is False + assert plan.terminate is True + assert plan.stop_reason == "web_search_failed" + blocks = plan.metadata[WEBSEARCH_NATIVE_BLOCKS_METADATA_KEY] + assert [b["type"] for b in blocks] == [ + "server_tool_use", + "web_search_tool_result", + "server_tool_use", + "web_search_tool_result", + ] + assert blocks[1]["tool_use_id"] == blocks[0]["id"] + assert blocks[1]["content"] == {"type": "web_search_tool_result_error", "error_code": "unavailable"} + assert blocks[3]["tool_use_id"] == blocks[2]["id"] + assert blocks[3]["content"] == {"type": "web_search_tool_result_error", "error_code": "unavailable"} + + @pytest.mark.asyncio + async def test_all_failed_iteration_keeps_the_follow_up_without_native_blocks(self): + logger = WebSearchInterceptionLogger(enabled_providers=["bedrock"]) + tool_calls = [ + {"id": "toolu_one", "type": "tool_use", "name": "litellm_web_search", "input": {"query": "q1"}}, + ] + + with patch.object( + logger, + "_execute_search", + side_effect=AuthenticationError("401 Unauthorized", llm_provider="tavily", model="tavily"), + ): + plan = await logger.async_build_agentic_loop_plan( + tools={"tool_calls": tool_calls, "thinking_blocks": []}, + model="bedrock/claude", + messages=[{"role": "user", "content": "hi"}], + response=MagicMock(), + anthropic_messages_provider_config=None, + anthropic_messages_optional_request_params={}, + logging_obj=MagicMock(model_call_details={}), + stream=False, + kwargs={}, + ) + + assert plan.run_agentic_loop is True + assert plan.terminate is False + assert plan.request_patch is not None + tool_results = plan.request_patch.messages[-1]["content"] + assert "Search failed: litellm.AuthenticationError: 401 Unauthorized" in tool_results[0]["content"] + + @pytest.mark.asyncio + async def test_mixed_iteration_keeps_the_follow_up_and_pairs_each_block(self): + logger = WebSearchInterceptionLogger(enabled_providers=["bedrock"]) + tool_calls = [ + {"id": "toolu_one", "type": "tool_use", "name": "litellm_web_search", "input": {"query": "fails"}}, + {"id": "toolu_two", "type": "tool_use", "name": "litellm_web_search", "input": {"query": "works"}}, + ] + + async def search(query, kwargs=None): + if query == "fails": + raise RateLimitError("slow down", llm_provider="tavily", model="tavily") + return ("Title: x", _make_search_response()) + + with patch.object(logger, "_execute_search", side_effect=search): + plan = await logger.async_build_agentic_loop_plan( + tools={"tool_calls": tool_calls, "thinking_blocks": []}, + model="bedrock/claude", + messages=[{"role": "user", "content": "hi"}], + response=MagicMock(), + anthropic_messages_provider_config=None, + anthropic_messages_optional_request_params={}, + logging_obj=MagicMock(model_call_details={}), + stream=False, + kwargs={WEBSEARCH_EMIT_NATIVE_BLOCKS_KEY: True}, + ) + + assert plan.run_agentic_loop is True + assert plan.terminate is False + blocks = plan.metadata[WEBSEARCH_NATIVE_BLOCKS_METADATA_KEY] + assert blocks[0]["input"] == {"query": "fails"} + assert blocks[1]["content"] == {"type": "web_search_tool_result_error", "error_code": "too_many_requests"} + assert blocks[2]["input"] == {"query": "works"} + assert blocks[3]["content"][0]["url"] == "https://docs.litellm.ai/" + + class TestPostHookInjectsBlocks: """The post-hook must prepend blocks; absent metadata is a no-op.""" @@ -437,13 +589,17 @@ class TestShortCircuitEmitsNativeBlocks: assert block_types == ["text"] @pytest.mark.asyncio - async def test_native_short_circuit_failure_still_emits_blocks(self): - """Search failure on native path: emit blocks with empty results + - the legacy text-error block, so the client gets a well-formed - response instead of a malformed half-shape.""" + async def test_native_short_circuit_failure_emits_the_error_block(self): + """Search failure on native path: the tool result carries Anthropic's + error object (rendered as "Web search error: " by the client) + next to the legacy text-error block.""" logger = WebSearchInterceptionLogger(enabled_providers=["github_copilot"]) - with patch.object(logger, "_execute_search", side_effect=RuntimeError("boom")): + with patch.object( + logger, + "_execute_search", + side_effect=RateLimitError("slow down", llm_provider="tavily", model="tavily"), + ): result = await logger.try_short_circuit_search( model="github_copilot/claude-sonnet-4", messages=[{"role": "user", "content": "search query"}], @@ -455,9 +611,10 @@ class TestShortCircuitEmitsNativeBlocks: block_types = [b["type"] for b in result["content"]] assert block_types == ["server_tool_use", "web_search_tool_result", "text"] tool_result = result["content"][1] - assert tool_result["content"] == [] + assert tool_result["tool_use_id"] == result["content"][0]["id"] + assert tool_result["content"] == {"type": "web_search_tool_result_error", "error_code": "too_many_requests"} text_block = result["content"][2] - assert "Search failed" in text_block["text"] + assert text_block["text"] == "Search failed: litellm.RateLimitError: slow down" class TestLegacyPathMatchesNewPath: @@ -489,7 +646,7 @@ class TestLegacyPathMatchesNewPath: patch.object( logger, "_build_anthropic_request_patch", - new=AsyncMock(return_value=(patch_obj, [_make_search_response()])), + new=AsyncMock(return_value=(patch_obj, (_succeeded_outcome(),))), ), patch( "litellm.integrations.websearch_interception.handler.anthropic_messages.acreate", diff --git a/tests/test_litellm/llms/anthropic/test_anthropic_common_utils.py b/tests/test_litellm/llms/anthropic/test_anthropic_common_utils.py index 133d6e502f4..945033c5cac 100644 --- a/tests/test_litellm/llms/anthropic/test_anthropic_common_utils.py +++ b/tests/test_litellm/llms/anthropic/test_anthropic_common_utils.py @@ -1828,7 +1828,11 @@ class TestAnthropicThinkingSignatureSelfHeal: assert out[0] is msgs[0] - def test_flatten_unencrypted_web_search_results_leaves_error_blocks_alone(self): + def test_flatten_unencrypted_web_search_results_flattens_error_blocks(self): + """A failed intercepted search is replayed by the client as the error + object LiteLLM emitted. Anthropic rejects a replayed ``server_tool_use`` + it never issued, so the pair is flattened to text the same way a + successful unencrypted result is.""" from litellm.llms.anthropic.common_utils import ( flatten_unencrypted_web_search_results_in_anthropic_messages, ) @@ -1837,6 +1841,7 @@ class TestAnthropicThinkingSignatureSelfHeal: { "role": "assistant", "content": [ + {"type": "server_tool_use", "id": "srvtoolu_1", "name": "web_search", "input": {"query": "q"}}, { "type": "web_search_tool_result", "tool_use_id": "srvtoolu_1", @@ -1844,14 +1849,18 @@ class TestAnthropicThinkingSignatureSelfHeal: "type": "web_search_tool_result_error", "error_code": "max_uses_exceeded", }, - } + }, ], } ] - out = flatten_unencrypted_web_search_results_in_anthropic_messages(msgs) + once = flatten_unencrypted_web_search_results_in_anthropic_messages(msgs) + twice = flatten_unencrypted_web_search_results_in_anthropic_messages(once) - assert out[0] is msgs[0] + assert once[0]["content"] == [ + {"type": "text", "text": "Web search results for 'q':\n\nSearch failed: max_uses_exceeded"} + ] + assert json.dumps(twice) == json.dumps(once) def test_sanitize_tool_use_ids_in_anthropic_messages(self): from litellm.llms.anthropic.common_utils import ( From 6b082d3a018b9956425babb47ce8cb2aee260a6f Mon Sep 17 00:00:00 2001 From: kerry Date: Sat, 19 Sep 2026 00:25:01 +0000 Subject: [PATCH 095/144] test(bedrock): type the SigV4 request recorder and drop caller-owned mutation Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../llms/bedrock/batches/test_handler.py | 51 +++++++++++-------- 1 file changed, 31 insertions(+), 20 deletions(-) diff --git a/tests/test_litellm/llms/bedrock/batches/test_handler.py b/tests/test_litellm/llms/bedrock/batches/test_handler.py index 056378f97c9..d328e09056b 100644 --- a/tests/test_litellm/llms/bedrock/batches/test_handler.py +++ b/tests/test_litellm/llms/bedrock/batches/test_handler.py @@ -8,10 +8,14 @@ the tests don't hit AWS. from __future__ import annotations +import json +from collections.abc import Iterator, Mapping from datetime import datetime, timedelta, timezone +from typing import Final from unittest.mock import MagicMock, patch import pytest +from botocore.awsrequest import AWSPreparedRequest, AWSResponse from litellm.llms.bedrock.batches.handler import ( # noqa: E402 @@ -572,26 +576,36 @@ def test_cancel_batch_stops_and_polls_the_job_with_the_tagged_session(monkeypatc assert [kwargs["aws_access_key_id"] for kwargs in bedrock_client_kwargs] == ["ASIABATCHCANCELTAGGED"] * 2 -def _sigv4_capture_send(sent_headers: list[dict[str, str]], body: dict): - import json +class _JsonBody: + def __init__(self, payload: bytes) -> None: + self._payload: Final = payload - from botocore.awsrequest import AWSResponse + def stream(self) -> Iterator[bytes]: + return iter((self._payload,)) - def send(_self, request): - sent_headers.append({k: v.decode() if isinstance(v, bytes) else v for k, v in request.headers.items()}) - raw = MagicMock() - raw.stream.return_value = iter([json.dumps(body, default=str).encode()]) - return AWSResponse(request.url, 200, {"content-type": "application/json"}, raw) - return send +class _AuthorizationRecorder: + """Stands in for botocore's HTTP session and records the Authorization header of every request it receives.""" + + def __init__(self, body: Mapping[str, object]) -> None: + self._payload: Final = json.dumps(body, default=str).encode() + self.authorization_headers: tuple[str, ...] = () + + def send(self, request: AWSPreparedRequest) -> AWSResponse: + raw_authorization: Final = request.headers["Authorization"] + authorization: Final = ( + raw_authorization.decode() if isinstance(raw_authorization, bytes) else str(raw_authorization) + ) + self.authorization_headers = (*self.authorization_headers, authorization) + return AWSResponse(request.url, 200, {"content-type": "application/json"}, _JsonBody(self._payload)) def test_retrieve_signs_with_deployment_credentials_when_env_bearer_token_is_set(monkeypatch): """A proxy-wide AWS_BEARER_TOKEN_BEDROCK must not override the deployment's own SigV4 credentials.""" monkeypatch.setenv("AWS_BEARER_TOKEN_BEDROCK", "env-bearer-token") - sent_headers: list[dict[str, str]] = [] + recorder: Final = _AuthorizationRecorder(_fake_boto3_response()) - with patch("botocore.httpsession.URLLib3Session.send", _sigv4_capture_send(sent_headers, _fake_boto3_response())): + with patch("botocore.httpsession.URLLib3Session.send", recorder.send): batch = BedrockBatchesHandler._handle_model_invocation_job_status( batch_id=JOB_ARN, aws_access_key_id="AKIADEPLOYMENTKEY", @@ -599,18 +613,15 @@ def test_retrieve_signs_with_deployment_credentials_when_env_bearer_token_is_set ) assert batch.status == "completed" - assert len(sent_headers) == 1 - assert sent_headers[0]["Authorization"].startswith("AWS4-HMAC-SHA256 Credential=AKIADEPLOYMENTKEY/") + assert len(recorder.authorization_headers) == 1 + assert recorder.authorization_headers[0].startswith("AWS4-HMAC-SHA256 Credential=AKIADEPLOYMENTKEY/") def test_cancel_signs_with_deployment_credentials_when_env_bearer_token_is_set(monkeypatch): monkeypatch.setenv("AWS_BEARER_TOKEN_BEDROCK", "env-bearer-token") - sent_headers: list[dict[str, str]] = [] + recorder: Final = _AuthorizationRecorder(_fake_boto3_response(status="Stopped")) - with patch( - "botocore.httpsession.URLLib3Session.send", - _sigv4_capture_send(sent_headers, _fake_boto3_response(status="Stopped")), - ): + with patch("botocore.httpsession.URLLib3Session.send", recorder.send): batch = BedrockBatchesHandler.cancel_batch( batch_id=JOB_ARN, aws_access_key_id="AKIADEPLOYMENTKEY", @@ -618,5 +629,5 @@ def test_cancel_signs_with_deployment_credentials_when_env_bearer_token_is_set(m ) assert batch.status == "cancelled" - assert len(sent_headers) == 2 - assert all(h["Authorization"].startswith("AWS4-HMAC-SHA256 Credential=AKIADEPLOYMENTKEY/") for h in sent_headers) + assert len(recorder.authorization_headers) == 2 + assert all(h.startswith("AWS4-HMAC-SHA256 Credential=AKIADEPLOYMENTKEY/") for h in recorder.authorization_headers) From 542ad7dbacb4448878da75432fd837cec4885b56 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Sat, 19 Sep 2026 00:27:54 +0000 Subject: [PATCH 096/144] fix(ocr): forward the supplied client on the Python path and build pooled clients outside the lock Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm-rust/crates/http/src/pool.rs | 12 +++++++----- litellm/ocr/main.py | 8 ++++++++ tests/test_litellm/ocr/test_main.py | 29 ++++++++++++++++++++++++++++ 3 files changed, 44 insertions(+), 5 deletions(-) diff --git a/litellm-rust/crates/http/src/pool.rs b/litellm-rust/crates/http/src/pool.rs index 03c01e968ac..613ce9c2831 100644 --- a/litellm-rust/crates/http/src/pool.rs +++ b/litellm-rust/crates/http/src/pool.rs @@ -1,6 +1,6 @@ use std::{ collections::HashMap, - sync::{Arc, Mutex, PoisonError}, + sync::{Arc, Mutex, MutexGuard, PoisonError}, }; use reqwest::dns::Resolve; @@ -38,13 +38,15 @@ impl HttpClientPool { variant: ClientVariant, ) -> Result { let key = (config.clone(), variant); - let mut clients = self.clients.lock().unwrap_or_else(PoisonError::into_inner); - if let Some(client) = clients.get(&key) { + if let Some(client) = self.lock().get(&key) { return Ok(client.clone()); } let client = self.apply(variant, config.client_builder()?).build()?; - clients.insert(key, client.clone()); - Ok(client) + Ok(self.lock().entry(key).or_insert(client).clone()) + } + + fn lock(&self) -> MutexGuard<'_, HashMap<(HttpClientConfig, ClientVariant), reqwest::Client>> { + self.clients.lock().unwrap_or_else(PoisonError::into_inner) } fn apply( diff --git a/litellm/ocr/main.py b/litellm/ocr/main.py index 06830ed4b53..851d9162964 100644 --- a/litellm/ocr/main.py +++ b/litellm/ocr/main.py @@ -25,6 +25,7 @@ from litellm.llms.base_llm.ocr.transformation import ( OCRResponse, parse_ocr_request_format, ) +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import CustomPricingLiteLLMParams @@ -52,6 +53,11 @@ class _PreparedOCRRequest: litellm_logging_obj: LiteLLMLoggingObj +def _supplied_client(kwargs: Mapping[str, object]) -> HTTPHandler | AsyncHTTPHandler | None: + candidate: Final = kwargs.get("client") + return candidate if isinstance(candidate, (HTTPHandler, AsyncHTTPHandler)) else None + + def _prepare_ocr_request( model: str, document: Mapping[str, object], @@ -238,6 +244,7 @@ async def aocr( api_key=prepared.api_key, api_base=prepared.api_base, custom_llm_provider=prepared.custom_llm_provider, + client=_supplied_client(kwargs), aocr=True, headers=prepared.extra_headers, provider_config=prepared.provider_config, @@ -404,6 +411,7 @@ def ocr( api_key=prepared.api_key, api_base=prepared.api_base, custom_llm_provider=prepared.custom_llm_provider, + client=_supplied_client(kwargs), aocr=_is_async, headers=prepared.extra_headers, provider_config=prepared.provider_config, diff --git a/tests/test_litellm/ocr/test_main.py b/tests/test_litellm/ocr/test_main.py index 5531a2639c0..32e5637ee09 100644 --- a/tests/test_litellm/ocr/test_main.py +++ b/tests/test_litellm/ocr/test_main.py @@ -113,6 +113,35 @@ async def test_python_request_response_and_callbacks( assert logger.log_pre_api_call.call_count == 1 +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", [False, True]) +async def test_python_uses_the_supplied_client(provider: Mock, asynchronous: bool) -> None: + supplied: Final = Mock(return_value=provider.return_value) + transport: Final = httpx.MockTransport(supplied) + arguments: Final = { + "model": "mistral/mistral-ocr-latest", + "document": dict(PRICING_DOCUMENT), + "api_key": "test-key", + "api_base": "https://ocr.test/v1", + } + + async def call() -> OCRResponse: + if not asynchronous: + with httpx.Client(transport=transport) as sync_client: + return litellm.ocr(**arguments, client=HTTPHandler(client=sync_client)) + async with httpx.AsyncClient(transport=transport) as async_client: + handler: Final = AsyncHTTPHandler() + await handler.client.aclose() + handler.client = async_client + return await litellm.aocr(**arguments, client=handler) + + response: Final = await call() + assert response.pages[0].markdown == "parsed document" + assert supplied.call_count == 1 + assert str(supplied.call_args.args[0].url) == "https://ocr.test/v1/ocr" + assert provider.call_count == 0 + + @pytest.mark.asyncio @pytest.mark.parametrize("asynchronous", [False, True]) async def test_python_provider_errors_keep_public_exception(provider: Mock, asynchronous: bool) -> None: From 988676a0b81370c645f20a12547164b6eba40585 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Sat, 19 Sep 2026 00:30:43 +0000 Subject: [PATCH 097/144] test(rust): parse the settings contract through Python so the bridge keeps to the interop boundary Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../python-bridge/src/python_settings.rs | 28 +++++++++++++++---- 1 file changed, 22 insertions(+), 6 deletions(-) diff --git a/litellm-rust/crates/python-bridge/src/python_settings.rs b/litellm-rust/crates/python-bridge/src/python_settings.rs index 0c6554f0970..c5f9f309615 100644 --- a/litellm-rust/crates/python-bridge/src/python_settings.rs +++ b/litellm-rust/crates/python-bridge/src/python_settings.rs @@ -33,16 +33,32 @@ pub(crate) const CONTRACT: &str = include_str!("../python_settings.json"); #[cfg(test)] mod tests { - use std::collections::BTreeSet; + use std::{collections::BTreeSet, ffi::CString}; + + use pyo3::{prelude::*, types::PyDict}; use super::{CONTRACT, PythonSettings}; #[test] fn every_settings_group_is_in_the_python_contract() { - let contract: serde_json::Map = - serde_json::from_str(CONTRACT).unwrap(); - let declared: BTreeSet<&str> = contract.keys().map(String::as_str).collect(); - let read: BTreeSet<&str> = PythonSettings::ALL.map(PythonSettings::name).into(); - assert_eq!(read, declared); + Python::initialize(); + Python::attach(|py| { + let locals = PyDict::new(py); + locals.set_item("contract", CONTRACT).unwrap(); + let source = CString::new("import json\nkeys = list(json.loads(contract))").unwrap(); + py.run(&source, Some(&locals), Some(&locals)).unwrap(); + let declared: BTreeSet = locals + .get_item("keys") + .unwrap() + .unwrap() + .extract::>() + .unwrap() + .into_iter() + .collect(); + let read: BTreeSet = PythonSettings::ALL + .map(|group| group.name().to_owned()) + .into(); + assert_eq!(read, declared); + }); } } From 8b1f78fa08fb6322398ff78b27646ce15e6684d9 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 18 Sep 2026 17:30:57 -0700 Subject: [PATCH 098/144] fix(vertex_ai): drop a cleared turn's queued transcripts and carry its billed seconds --- .../audio_transcription/realtime_backend.py | 59 +++++++++++++++---- .../realtime_transformation.py | 1 + .../types/llms/vertex_ai_speech_to_text.py | 1 + .../test_vertex_ai_realtime_backend.py | 51 ++++++++++++++-- .../test_vertex_ai_realtime_transformation.py | 9 ++- 5 files changed, 104 insertions(+), 17 deletions(-) diff --git a/litellm/llms/vertex_ai/audio_transcription/realtime_backend.py b/litellm/llms/vertex_ai/audio_transcription/realtime_backend.py index 0c16fea9e9f..4c8338c027e 100644 --- a/litellm/llms/vertex_ai/audio_transcription/realtime_backend.py +++ b/litellm/llms/vertex_ai/audio_transcription/realtime_backend.py @@ -45,7 +45,6 @@ _LINK_QUEUE_SIZE: Final = 64 _CLOSE_REASON_MAX_CHARS: Final = 120 _CONFIGURED_EVENT: Final = VertexSpeechStreamingConfigured().model_dump_json() _TURN_FINISHED_EVENT: Final = VertexSpeechStreamingTurnFinished().model_dump_json() -_TURN_DISCARDED_EVENT: Final = VertexSpeechStreamingTurnDiscarded().model_dump_json() _COMMAND_ADAPTER: Final = TypeAdapter[VertexSpeechStreamingCommandUnion](VertexSpeechStreamingCommand) _TIMEDELTA_ADAPTER: Final = TypeAdapter(timedelta) _SPEECH_EVENTS: Final[MappingProxyType[str, Literal["begin", "end"]]] = MappingProxyType( @@ -80,6 +79,20 @@ class _Closed: pass +@dataclass(frozen=True, slots=True) +class _TurnResult: + turn: int + event: str + + +@dataclass(frozen=True, slots=True) +class _TurnDiscarded: + pass + + +_OutboxItem = str | _TurnResult | _StreamFailure | _Closed + + def open_speech_client(target: SpeechStreamingTarget, access_token: str) -> SpeechStreamingClient: try: from google.api_core.client_options import ClientOptions @@ -146,10 +159,12 @@ class _RecognizeStream: request_type: "type[StreamingRecognizeRequest]", first_request: "StreamingRecognizeRequest", opened_at: float, + turn: int, ) -> None: self._client: Final = client self._request_type: Final = request_type self.opened_at: Final = opened_at + self.turn: Final = turn self._requests: Final[asyncio.Queue[StreamingRecognizeRequest | None]] = asyncio.Queue( maxsize=REQUEST_QUEUE_SIZE ) @@ -177,7 +192,7 @@ class _RecognizeStream: self._closed = True await self._client.transport.close() - async def relay(self, outbox: "asyncio.Queue[str | _StreamFailure | _Closed]", billed_before: float) -> float: + async def relay(self, outbox: asyncio.Queue[_OutboxItem], billed_before: float) -> float: if self._cancelled: await self.close() return 0.0 @@ -193,12 +208,14 @@ class _RecognizeStream: await self.close() return self.billed_seconds - async def _forward(self, outbox: "asyncio.Queue[str | _StreamFailure | _Closed]", billed_before: float) -> None: + async def _forward(self, outbox: asyncio.Queue[_OutboxItem], billed_before: float) -> None: try: responses: Final = await self._client.streaming_recognize(self._drain()) async for response in responses: self._note(response) - await outbox.put(_response_event(response, billed_before + self.billed_seconds)) + await outbox.put( + _TurnResult(turn=self.turn, event=_response_event(response, billed_before + self.billed_seconds)) + ) except Exception as e: # noqa: BLE001 # task boundary: a swallowed failure would hang the client session verbose_logger.warning("Google Speech-to-Text streaming failed: %s", e) await outbox.put(_StreamFailure(reason=f"Google Speech-to-Text streaming failed: {e}")) @@ -214,6 +231,9 @@ class _RecognizeStream: yield request +_Link = _RecognizeStream | str | _TurnDiscarded + + class SpeechStreamingBackend: def __init__( self, @@ -229,11 +249,13 @@ class SpeechStreamingBackend: self._clock: Final = clock self._rotation_seconds: Final = rotation_seconds self._rotation_deadline_seconds: Final = rotation_deadline_seconds - self._outbox: Final[asyncio.Queue[str | _StreamFailure | _Closed]] = asyncio.Queue(maxsize=OUTBOX_SIZE) - self._links: Final[asyncio.Queue[_RecognizeStream | str]] = asyncio.Queue(maxsize=_LINK_QUEUE_SIZE) + self._outbox: Final[asyncio.Queue[_OutboxItem]] = asyncio.Queue(maxsize=OUTBOX_SIZE) + self._links: Final[asyncio.Queue[_Link]] = asyncio.Queue(maxsize=_LINK_QUEUE_SIZE) self._pump: asyncio.Task[None] | None = None self._config: StreamingRecognitionConfig | None = None self._turn: tuple[_RecognizeStream, ...] = () + self._turn_index: int = 0 + self._discarded_turns: frozenset[int] = frozenset() self._billed_before: float = 0.0 self._closed: bool = False @@ -267,9 +289,12 @@ class SpeechStreamingBackend: assert_never(command) async def recv(self, decode: bool | None = None) -> str | bytes: - if self._closed and self._outbox.empty(): - raise _normal_closure() - item: Final = await self._outbox.get() + while not (self._closed and self._outbox.empty()): + if (event := self._deliverable(await self._outbox.get())) is not None: + return event + raise _normal_closure() + + def _deliverable(self, item: _OutboxItem) -> str | None: match item: case _StreamFailure(): raise ConnectionClosedError( @@ -277,6 +302,8 @@ class SpeechStreamingBackend: ) case _Closed(): raise _normal_closure() + case _TurnResult(): + return None if item.turn in self._discarded_turns else item.event case str(): return item case _: @@ -301,7 +328,7 @@ class SpeechStreamingBackend: if isinstance(link, _RecognizeStream): await link.close() - async def _link(self, item: _RecognizeStream | str) -> None: + async def _link(self, item: _Link) -> None: if self._pump is None: self._pump = asyncio.create_task(self._pump_links()) await self._links.put(item) @@ -310,12 +337,16 @@ class SpeechStreamingBackend: while True: await self._relay(await self._links.get()) - async def _relay(self, link: _RecognizeStream | str) -> None: + async def _relay(self, link: _Link) -> None: match link: case str(): await self._outbox.put(link) case _RecognizeStream(): self._billed_before += await link.relay(self._outbox, self._billed_before) + case _TurnDiscarded(): + await self._outbox.put( + VertexSpeechStreamingTurnDiscarded(billed_seconds=self._billed_before).model_dump_json() + ) case _: assert_never(link) @@ -351,6 +382,7 @@ class SpeechStreamingBackend: request_type=StreamingRecognizeRequest, first_request=StreamingRecognizeRequest(recognizer=self._target.recognizer, streaming_config=config), opened_at=self._clock(), + turn=self._turn_index, ) await self._link(stream) return stream @@ -358,6 +390,7 @@ class SpeechStreamingBackend: async def _finish_turn(self) -> None: turn: Final = self._turn self._turn = () + self._turn_index += 1 if turn: await turn[-1].half_close() await self._link(_TURN_FINISHED_EVENT) @@ -365,6 +398,8 @@ class SpeechStreamingBackend: async def _discard_turn(self) -> None: turn: Final = self._turn self._turn = () + self._discarded_turns |= {self._turn_index} + self._turn_index += 1 for stream in turn: stream.cancel() - await self._link(_TURN_DISCARDED_EVENT) + await self._link(_TurnDiscarded()) diff --git a/litellm/llms/vertex_ai/audio_transcription/realtime_transformation.py b/litellm/llms/vertex_ai/audio_transcription/realtime_transformation.py index dab2e980fd0..ac23901accb 100644 --- a/litellm/llms/vertex_ai/audio_transcription/realtime_transformation.py +++ b/litellm/llms/vertex_ai/audio_transcription/realtime_transformation.py @@ -236,6 +236,7 @@ class ChirpEventTransformer: case VertexSpeechStreamingTurnFinished(): return self._finish_turn() case VertexSpeechStreamingTurnDiscarded(): + self._billed_seconds = max(self._billed_seconds, frame.billed_seconds) self._turn = None return () case _: diff --git a/litellm/types/llms/vertex_ai_speech_to_text.py b/litellm/types/llms/vertex_ai_speech_to_text.py index e954960d41f..d07a5bbc192 100644 --- a/litellm/types/llms/vertex_ai_speech_to_text.py +++ b/litellm/types/llms/vertex_ai_speech_to_text.py @@ -93,6 +93,7 @@ class VertexSpeechStreamingTurnFinished(BaseModel): class VertexSpeechStreamingTurnDiscarded(BaseModel): model_config = ConfigDict(frozen=True) kind: Literal["turn_discarded"] = "turn_discarded" + billed_seconds: float VertexSpeechStreamingEventUnion = ( diff --git a/tests/test_litellm/llms/vertex_ai/audio_transcription/test_vertex_ai_realtime_backend.py b/tests/test_litellm/llms/vertex_ai/audio_transcription/test_vertex_ai_realtime_backend.py index d6f65c90806..15601c5ca6c 100644 --- a/tests/test_litellm/llms/vertex_ai/audio_transcription/test_vertex_ai_realtime_backend.py +++ b/tests/test_litellm/llms/vertex_ai/audio_transcription/test_vertex_ai_realtime_backend.py @@ -1,6 +1,6 @@ import asyncio import json -from collections.abc import AsyncIterator, Sequence +from collections.abc import AsyncIterator, Callable, Sequence from dataclasses import replace from datetime import timedelta from typing import Final @@ -128,6 +128,14 @@ async def _configure(backend: SpeechStreamingBackend) -> None: assert await _recv(backend) == {"kind": "configured"} +async def _until(condition: Callable[[], bool]) -> None: + async def poll() -> None: + while not condition(): + await asyncio.sleep(0) + + await asyncio.wait_for(poll(), timeout=2) + + def _audio(stream: list[StreamingRecognizeRequest]) -> list[bytes]: return [bytes(request.audio) for request in stream[1:]] @@ -221,7 +229,7 @@ async def test_turn_commands_without_audio_answer_immediately(): await backend.send(FINISH_TURN) assert await _recv(backend) == {"kind": "turn_finished"} await backend.send(DISCARD_TURN) - assert await _recv(backend) == {"kind": "turn_discarded"} + assert await _recv(backend) == {"kind": "turn_discarded", "billed_seconds": 0.0} @pytest.mark.asyncio @@ -232,12 +240,47 @@ async def test_discard_turn_cancels_the_open_stream_and_the_next_turn_starts_fre await backend.send(b"\x01\x01") assert await _transcript(backend) == "draft" await backend.send(DISCARD_TURN) - assert await _recv(backend) == {"kind": "turn_discarded"} + assert await _recv(backend) == {"kind": "turn_discarded", "billed_seconds": 0.0} await backend.send(b"\x02\x02") assert await _transcript(backend) == "again" assert [_audio(stream) for stream in client.streams] == [[b"\x01\x01"], [b"\x02\x02"]] +@pytest.mark.asyncio +async def test_discard_turn_drops_its_queued_results_and_keeps_google_billed_seconds(): + client = _FakeSpeechClient( + [_response("draft"), _response("leftover", is_final=True, billed=2.0)], + [_response("fresh", is_final=True, billed=1.0)], + ) + async with _backend(client) as backend: + await _configure(backend) + await backend.send(b"\x01\x01") + assert await _transcript(backend) == "draft" + await backend.send(b"\x02\x02") + await _until(lambda: len(client.streams[0]) == 3) + await backend.send(DISCARD_TURN) + assert await _recv(backend) == {"kind": "turn_discarded", "billed_seconds": 2.0} + await backend.send(b"\x03\x03") + fresh = await _recv(backend) + assert fresh["results"] == [{"transcript": "fresh", "is_final": True}] + assert fresh["billed_seconds"] == 3.0 + + +@pytest.mark.asyncio +async def test_discard_turn_keeps_the_queued_results_of_the_turn_finished_before_it(): + client = _FakeSpeechClient([_response("one", is_final=True, billed=2.0)], [_response("two")]) + async with _backend(client) as backend: + await _configure(backend) + await backend.send(b"\x01\x01") + await backend.send(FINISH_TURN) + await backend.send(b"\x02\x02") + await _until(lambda: len(client.streams) == 2 and len(client.streams[1]) == 2) + await backend.send(DISCARD_TURN) + assert await _transcript(backend) == "one" + assert await _recv(backend) == {"kind": "turn_finished"} + assert await _recv(backend) == {"kind": "turn_discarded", "billed_seconds": 2.0} + + @pytest.mark.asyncio async def test_billed_seconds_accumulate_across_turns(): client = _FakeSpeechClient( @@ -425,7 +468,7 @@ async def test_discard_turn_cancels_every_stream_of_the_turn(): now[0] = 240.0 await backend.send(b"\x02\x02") await backend.send(DISCARD_TURN) - assert await _recv(backend) == {"kind": "turn_discarded"} + assert await _recv(backend) == {"kind": "turn_discarded", "billed_seconds": 0.0} await backend.send(b"\x03\x03") assert await _transcript(backend) == "fresh" assert [_audio(stream) for stream in client.streams] == [[b"\x01\x01"], [b"\x03\x03"]] diff --git a/tests/test_litellm/llms/vertex_ai/audio_transcription/test_vertex_ai_realtime_transformation.py b/tests/test_litellm/llms/vertex_ai/audio_transcription/test_vertex_ai_realtime_transformation.py index 719dd621c82..84c3a4e244a 100644 --- a/tests/test_litellm/llms/vertex_ai/audio_transcription/test_vertex_ai_realtime_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/audio_transcription/test_vertex_ai_realtime_transformation.py @@ -317,13 +317,20 @@ def test_manual_turns_complete_on_commit_without_speech_events(): def test_clear_discards_the_open_turn(): config = _configured(turn_detection=None) draft = _backend_events(config, _response(("draft", False))) - assert _backend_events(config, VertexSpeechStreamingTurnDiscarded()) == [] + assert _backend_events(config, VertexSpeechStreamingTurnDiscarded(billed_seconds=0.0)) == [] assert _backend_events(config, VertexSpeechStreamingTurnFinished()) == [] fresh = _backend_events(config, _response(("again", False))) assert fresh[0]["delta"] == "again" assert fresh[0]["item_id"] != draft[0]["item_id"] +def test_cleared_audio_keeps_google_billed_seconds_for_the_close_flush(): + config = _configured(turn_detection=None) + assert _backend_events(config, _response(("draft", False), billed_seconds=1.0)) != [] + assert _backend_events(config, VertexSpeechStreamingTurnDiscarded(billed_seconds=2.5)) == [] + assert config.unbilled_usage_on_session_close(MODEL) == {"type": "duration", "seconds": 2.5} + + def test_usage_is_billed_once_across_turns_and_flushed_on_close(): config = _configured() first = _backend_events(config, _response(("one", True), billed_seconds=2.0)) From a20698f802585b7bef5d3abe6ebc32935e5214d2 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Fri, 18 Sep 2026 17:37:02 -0700 Subject: [PATCH 099/144] ci(issues): comment which release carries the fix when a pull request closes an issue --- .github/workflows/issue_fixed_comment.yml | 71 +++++++ scripts/comment-fixed-issue.test.ts | 234 ++++++++++++++++++++++ scripts/comment-fixed-issue.ts | 224 +++++++++++++++++++++ 3 files changed, 529 insertions(+) create mode 100644 .github/workflows/issue_fixed_comment.yml create mode 100644 scripts/comment-fixed-issue.test.ts create mode 100644 scripts/comment-fixed-issue.ts diff --git a/.github/workflows/issue_fixed_comment.yml b/.github/workflows/issue_fixed_comment.yml new file mode 100644 index 00000000000..92993d319a7 --- /dev/null +++ b/.github/workflows/issue_fixed_comment.yml @@ -0,0 +1,71 @@ +name: Issue fixed comment + +on: + issues: + types: [closed] + workflow_dispatch: + inputs: + issue_number: + description: "Closed issue number to comment on manually." + required: true + pull_request: + paths: + - .github/workflows/issue_fixed_comment.yml + - scripts/comment-fixed-issue.ts + - scripts/comment-fixed-issue.test.ts + - scripts/auto-close-duplicates.ts + +permissions: {} + +concurrency: + group: issue-fixed-comment-${{ github.event.issue.number || github.event.inputs.issue_number || github.run_id }} + cancel-in-progress: false + +jobs: + comment-fixed-issue-tests: + if: github.event_name == 'pull_request' + runs-on: ubuntu-latest + timeout-minutes: 5 + permissions: + contents: read + steps: + - name: Checkout repository + uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 + with: + persist-credentials: false + + - name: Setup Bun + uses: oven-sh/setup-bun@0c5077e51419868618aeaa5fe8019c62421857d6 # v2.2.0 + with: + bun-version: "1.4.0" + + - name: Test the closer lookup, the release placement and the comment + run: bun test scripts/comment-fixed-issue.test.ts + + comment-fixed-issue: + if: github.event_name != 'pull_request' && github.repository == 'BerriAI/litellm' + runs-on: ubuntu-latest + timeout-minutes: 5 + permissions: + contents: read + issues: write + steps: + - name: Checkout scripts + uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 + with: + sparse-checkout: scripts + persist-credentials: false + + - name: Setup Bun + uses: oven-sh/setup-bun@0c5077e51419868618aeaa5fe8019c62421857d6 # v2.2.0 + with: + # Exact version, never latest: the next step holds an issues: write token + bun-version: "1.4.0" + + - name: Name the release that carries the fix + run: bun run scripts/comment-fixed-issue.ts | tee -a "${GITHUB_STEP_SUMMARY}" + env: + GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} + ISSUE_NUMBER: ${{ github.event.issue.number || github.event.inputs.issue_number }} + DEFAULT_BRANCH: ${{ github.event.repository.default_branch }} + DRY_RUN: ${{ vars.ISSUE_FIXED_COMMENT_ENABLED != 'true' }} diff --git a/scripts/comment-fixed-issue.test.ts b/scripts/comment-fixed-issue.test.ts new file mode 100644 index 00000000000..f9cd41d96d8 --- /dev/null +++ b/scripts/comment-fixed-issue.test.ts @@ -0,0 +1,234 @@ +import { describe, expect, test } from "bun:test"; + +import type { Comment, GitHubApi } from "./auto-close-duplicates"; +import { + FIXED_MARKER, + closerOf, + commentFixedIssue, + fixedBody, + nextMinor, + parseVersion, + placement, + readConfig, + releaseCandidate, + type ClosedIssue, + type FixedConfig, +} from "./comment-fixed-issue"; + +const MERGE_COMMIT = "68c4c82ac977b48b2b81ee8d633d5771307c6162"; + +const mergedPr = { + __typename: "PullRequest" as const, + number: 41767, + merged: true, + baseRefName: "main", + mergeCommit: { oid: MERGE_COMMIT }, +}; + +type Closer = ClosedIssue["timelineItems"]["nodes"][number]["closer"]; + +const closedBy = (closer: Closer, state: ClosedIssue["state"] = "CLOSED"): ClosedIssue => ({ + state, + timelineItems: { nodes: [{ closer }] }, +}); + +const pyproject = (version: string): string => + `[project]\nname = "litellm"\nversion = "${version}"\n\n[tool.commitizen]\nversion = "${version}"\n`; + +const config: FixedConfig = { repo: "BerriAI/litellm", issueNumber: 41750, defaultBranch: "main", dryRun: false }; + +interface World { + readonly issue?: ClosedIssue | null; + readonly comments?: readonly Comment[]; + readonly version?: string; + // Which existing rc.1 tags contain the merge commit; a tag absent from the map does not exist + readonly tags?: Readonly>; +} + +function fakeApi(world: World = {}): { readonly api: GitHubApi; readonly writes: string[] } { + const writes: string[] = []; + const tags = world.tags ?? {}; + const api: GitHubApi = { + request: async (method: string, path: string, body?: object): Promise => { + if (method === "POST" && path === "/graphql") { + return { data: { repository: { issue: world.issue === undefined ? closedBy(mergedPr) : world.issue } } } as T; + } + if (method !== "GET") { + writes.push(`${method} ${path} ${JSON.stringify(body)}`); + return {} as T; + } + if (path.startsWith("/repos/BerriAI/litellm/issues/41750/comments")) { + return (world.comments ?? []) as T; + } + if (path === `/repos/BerriAI/litellm/contents/pyproject.toml?ref=${MERGE_COMMIT}`) { + return { content: btoa(pyproject(world.version ?? "1.103.0")).replace(/(.{60})/g, "$1\n") } as T; + } + const matching = /^\/repos\/BerriAI\/litellm\/git\/matching-refs\/tags\/(.+)$/.exec(path); + if (matching !== null) { + return (matching[1] in tags ? [{ ref: `refs/tags/${matching[1]}` }] : []) as T; + } + const compare = /^\/repos\/BerriAI\/litellm\/compare\/(.+)\.\.\.(.+)$/.exec(path); + if (compare !== null && compare[2] === MERGE_COMMIT) { + return { status: tags[compare[1]] ? "behind" : "ahead" } as T; + } + throw new Error(`unexpected ${method} ${path}`); + }, + }; + return { api, writes }; +} + +describe("closerOf", () => { + test("a pull request merged into the default branch is the fix", () => { + expect(closerOf(closedBy(mergedPr), "main")).toEqual({ kind: "pull_request", number: 41767, mergeCommit: MERGE_COMMIT }); + }); + + test("an issue closed by hand, by a commit, or by an unmerged pull request gets no comment", () => { + expect(closerOf(closedBy(null), "main")).toEqual({ kind: "skip", reason: "closed by hand, not by a pull request" }); + expect(closerOf(closedBy({ __typename: "Commit", oid: MERGE_COMMIT }), "main").kind).toBe("skip"); + expect(closerOf(closedBy({ ...mergedPr, merged: false }), "main").kind).toBe("skip"); + expect(closerOf(closedBy({ ...mergedPr, mergeCommit: null }), "main").kind).toBe("skip"); + }); + + test("a pull request merged into a release branch is not a fix on main", () => { + const verdict = closerOf(closedBy({ ...mergedPr, baseRefName: "release/1.102.0rc2" }), "main"); + expect(verdict).toEqual({ kind: "skip", reason: "#41767 merged into release/1.102.0rc2, not main" }); + }); + + test("an issue reopened after the close event is left alone", () => { + expect(closerOf(closedBy(mergedPr, "OPEN"), "main")).toEqual({ kind: "skip", reason: "the issue is open again" }); + }); +}); + +describe("version helpers", () => { + test("parseVersion reads the project version and ignores everything else", () => { + expect(parseVersion(pyproject("1.103.0"))).toBe("1.103.0"); + expect(parseVersion('[project]\nversion = "1.103.0rc1"\n')).toBeUndefined(); + expect(parseVersion("[project]\nname = 'litellm'\n")).toBeUndefined(); + }); + + test("the first rc of a version is the release that carries a fix merged under it", () => { + expect(releaseCandidate("1.103.0")).toBe("v1.103.0-rc.1"); + }); + + test("nextMinor bumps the minor and resets the patch", () => { + expect(nextMinor("1.103.0")).toBe("1.104.0"); + expect(nextMinor("1.99.4")).toBe("1.100.0"); + }); +}); + +describe("placement", () => { + test("no rc yet: the fix ships in the rc.1 of the version at the merge commit", async () => { + const { api } = fakeApi({ version: "1.103.0" }); + expect(await placement(api, "BerriAI/litellm", MERGE_COMMIT)).toEqual({ kind: "release", tag: "v1.103.0-rc.1", shipped: false }); + }); + + test("rc.1 already cut with the commit in it: the fix is out", async () => { + const { api } = fakeApi({ version: "1.102.0", tags: { "v1.102.0-rc.1": true } }); + expect(await placement(api, "BerriAI/litellm", MERGE_COMMIT)).toEqual({ kind: "release", tag: "v1.102.0-rc.1", shipped: true }); + }); + + test("rc.1 cut before the merge while main still said that version: the fix waits for the next minor", async () => { + const { api } = fakeApi({ version: "1.102.0", tags: { "v1.102.0-rc.1": false } }); + expect(await placement(api, "BerriAI/litellm", MERGE_COMMIT)).toEqual({ kind: "release", tag: "v1.103.0-rc.1", shipped: false }); + }); + + test("keeps walking minors while each rc.1 exists without the commit, then gives up", async () => { + const twoTaken = fakeApi({ version: "1.102.0", tags: { "v1.102.0-rc.1": false, "v1.103.0-rc.1": false } }); + expect(await placement(twoTaken.api, "BerriAI/litellm", MERGE_COMMIT)).toEqual({ kind: "release", tag: "v1.104.0-rc.1", shipped: false }); + + const allTaken = fakeApi({ + version: "1.102.0", + tags: { "v1.102.0-rc.1": false, "v1.103.0-rc.1": false, "v1.104.0-rc.1": false, "v1.105.0-rc.1": false }, + }); + expect((await placement(allTaken.api, "BerriAI/litellm", MERGE_COMMIT)).kind).toBe("skip"); + }); + + test("a pyproject without a version line is a skip, not a comment", async () => { + const { api } = fakeApi({ version: "not-a-version" }); + expect((await placement(api, "BerriAI/litellm", MERGE_COMMIT)).kind).toBe("skip"); + }); +}); + +describe("fixedBody", () => { + test("names the pull request and the first release, and carries the marker the rerun looks for", () => { + const body = fixedBody(41767, { tag: "v1.103.0-rc.1", shipped: false }); + expect(body.startsWith(FIXED_MARKER)).toBe(true); + expect(body).toContain("Fixed by #41767."); + expect(body).toContain("ships in v1.103.0-rc.1 and up"); + expect(body).toContain("dev pre-release"); + }); + + test("a release that is already out says so instead of promising one", () => { + const body = fixedBody(41767, { tag: "v1.102.0-rc.1", shipped: true }); + expect(body).toContain("is in v1.102.0-rc.1 and up"); + expect(body).not.toContain("ships in"); + }); + + test("stays within the 25-word comment rule either way", () => { + for (const shipped of [true, false]) { + const words = fixedBody(41767, { tag: "v1.103.0-rc.1", shipped }).replace(FIXED_MARKER, "").trim().split(/\s+/); + expect(words.length).toBeGreaterThanOrEqual(15); + expect(words.length).toBeLessThanOrEqual(25); + } + }); +}); + +describe("commentFixedIssue", () => { + test("a real run posts one comment naming the pull request and the release", async () => { + const { api, writes } = fakeApi(); + const verdict = await commentFixedIssue(api, config); + expect(verdict).toMatchObject({ kind: "commented", pullRequest: 41767, tag: "v1.103.0-rc.1" }); + expect(writes).toHaveLength(1); + expect(writes[0]).toContain("POST /repos/BerriAI/litellm/issues/41750/comments"); + expect(writes[0]).toContain("Fixed by #41767. This ships in v1.103.0-rc.1 and up"); + }); + + test("a dry run renders the comment and writes nothing", async () => { + const { api, writes } = fakeApi(); + const verdict = await commentFixedIssue(api, { ...config, dryRun: true }); + expect(verdict.kind).toBe("commented"); + expect(writes).toEqual([]); + }); + + test("an issue that already carries the comment is not commented twice", async () => { + const existing: Comment = { + id: 1, + body: `${FIXED_MARKER}\nFixed by #41767. This ships in v1.103.0-rc.1 and up.`, + created_at: "2026-09-18T00:00:00Z", + user: { type: "Bot", login: "github-actions[bot]" }, + }; + const { api, writes } = fakeApi({ comments: [existing] }); + expect(await commentFixedIssue(api, config)).toEqual({ kind: "skip", reason: "already carries a fixed-in comment" }); + expect(writes).toEqual([]); + }); + + test("a hand-closed issue never reaches the release lookup or the API writes", async () => { + const { api, writes } = fakeApi({ issue: closedBy(null) }); + expect((await commentFixedIssue(api, config)).kind).toBe("skip"); + expect(writes).toEqual([]); + }); + + test("a number that is not an issue in the repository is a skip", async () => { + const { api, writes } = fakeApi({ issue: null }); + expect(await commentFixedIssue(api, config)).toEqual({ kind: "skip", reason: "not an issue in this repository" }); + expect(writes).toEqual([]); + }); +}); + +describe("readConfig", () => { + const env = { GITHUB_TOKEN: "t", GITHUB_REPOSITORY: "BerriAI/litellm", ISSUE_NUMBER: "41750", DEFAULT_BRANCH: "main" }; + + test("reads the four inputs and treats anything but the literal true as a real run", () => { + expect(readConfig(env)).toEqual({ token: "t", repo: "BerriAI/litellm", issueNumber: 41750, defaultBranch: "main", dryRun: false }); + expect(readConfig({ ...env, DRY_RUN: "true" }).dryRun).toBe(true); + expect(readConfig({ ...env, DRY_RUN: "false" }).dryRun).toBe(false); + }); + + test("refuses a missing token, repo, branch or a bad issue number", () => { + expect(() => readConfig({ ...env, GITHUB_TOKEN: undefined })).toThrow("GITHUB_TOKEN"); + expect(() => readConfig({ ...env, GITHUB_REPOSITORY: "litellm" })).toThrow("owner/repo"); + expect(() => readConfig({ ...env, DEFAULT_BRANCH: "" })).toThrow("DEFAULT_BRANCH"); + expect(() => readConfig({ ...env, ISSUE_NUMBER: "0" })).toThrow("ISSUE_NUMBER"); + expect(() => readConfig({ ...env, ISSUE_NUMBER: "abc" })).toThrow("ISSUE_NUMBER"); + }); +}); diff --git a/scripts/comment-fixed-issue.ts b/scripts/comment-fixed-issue.ts new file mode 100644 index 00000000000..480b5e90249 --- /dev/null +++ b/scripts/comment-fixed-issue.ts @@ -0,0 +1,224 @@ +#!/usr/bin/env bun + +import { githubApi, listAll, type Comment, type GitHubApi } from "./auto-close-duplicates"; + +declare const process: { readonly env: Readonly> }; + +export interface FixedConfig { + readonly repo: string; + readonly issueNumber: number; + readonly defaultBranch: string; + readonly dryRun: boolean; +} + +interface PullRequestCloser { + readonly __typename: "PullRequest"; + readonly number: number; + readonly merged: boolean; + readonly baseRefName: string; + readonly mergeCommit: { readonly oid: string } | null; +} + +interface CommitCloser { + readonly __typename: "Commit"; + readonly oid: string; +} + +export interface ClosedIssue { + readonly state: "OPEN" | "CLOSED"; + readonly timelineItems: { + readonly nodes: readonly { readonly closer: PullRequestCloser | CommitCloser | null }[]; + }; +} + +interface TimelineResponse { + readonly data?: { readonly repository?: { readonly issue: ClosedIssue | null } }; +} + +interface MatchingRef { + readonly ref: string; +} + +interface Comparison { + readonly status: "ahead" | "behind" | "identical" | "diverged"; +} + +interface FileContent { + readonly content: string; +} + +export type Closer = + | { readonly kind: "pull_request"; readonly number: number; readonly mergeCommit: string } + | { readonly kind: "skip"; readonly reason: string }; + +export type Placement = + | { readonly kind: "release"; readonly tag: string; readonly shipped: boolean } + | { readonly kind: "skip"; readonly reason: string }; + +export type FixedVerdict = + | { readonly kind: "commented"; readonly pullRequest: number; readonly tag: string; readonly body: string } + | { readonly kind: "skip"; readonly reason: string }; + +export const FIXED_MARKER = ""; +const MAX_MINOR_BUMPS = 3; + +export const CLOSER_QUERY = `query($owner: String!, $name: String!, $number: Int!) { + repository(owner: $owner, name: $name) { + issue(number: $number) { + state + timelineItems(last: 1, itemTypes: [CLOSED_EVENT]) { + nodes { + ... on ClosedEvent { + closer { + __typename + ... on PullRequest { number merged baseRefName mergeCommit { oid } } + ... on Commit { oid } + } + } + } + } + } + } +}`; + +const skip = (reason: string): { readonly kind: "skip"; readonly reason: string } => ({ kind: "skip", reason }); + +export function closerOf(issue: ClosedIssue, defaultBranch: string): Closer { + if (issue.state !== "CLOSED") { + return skip("the issue is open again"); + } + const closer = issue.timelineItems.nodes[0]?.closer ?? null; + if (closer === null) { + return skip("closed by hand, not by a pull request"); + } + if (closer.__typename === "Commit") { + return skip(`closed by commit ${closer.oid.slice(0, 10)}, not by a pull request`); + } + if (!closer.merged || closer.mergeCommit === null) { + return skip(`closed by #${closer.number}, which is not merged`); + } + if (closer.baseRefName !== defaultBranch) { + return skip(`#${closer.number} merged into ${closer.baseRefName}, not ${defaultBranch}`); + } + return { kind: "pull_request", number: closer.number, mergeCommit: closer.mergeCommit.oid }; +} + +export function parseVersion(pyproject: string): string | undefined { + return /^version = "(\d+\.\d+\.\d+)"$/m.exec(pyproject)?.[1]; +} + +export function releaseCandidate(version: string): string { + return `v${version}-rc.1`; +} + +export function nextMinor(version: string): string { + const [major, minor] = version.split(".").map(Number); + return `${major}.${minor + 1}.0`; +} + +async function tagExists(api: GitHubApi, repo: string, tag: string): Promise { + const refs = await api.request("GET", `/repos/${repo}/git/matching-refs/tags/${tag}`); + return refs.some((ref) => ref.ref === `refs/tags/${tag}`); +} + +async function tagContains(api: GitHubApi, repo: string, tag: string, sha: string): Promise { + const comparison = await api.request("GET", `/repos/${repo}/compare/${tag}...${sha}`); + return comparison.status === "behind" || comparison.status === "identical"; +} + +// The first rc of a version is cut straight from main, so a fix merged while pyproject says X.Y.Z ships in +// vX.Y.Z-rc.1 unless that rc was already cut without it, in which case it waits for the next minor's rc.1 +async function firstReleaseWith( + api: GitHubApi, + repo: string, + sha: string, + version: string, + bumpsLeft: number, +): Promise { + const tag = releaseCandidate(version); + if (!(await tagExists(api, repo, tag))) { + return { kind: "release", tag, shipped: false }; + } + if (await tagContains(api, repo, tag, sha)) { + return { kind: "release", tag, shipped: true }; + } + if (bumpsLeft === 0) { + return skip(`${tag} exists without ${sha.slice(0, 10)} and the next ${MAX_MINOR_BUMPS} rc.1 tags are taken too`); + } + return firstReleaseWith(api, repo, sha, nextMinor(version), bumpsLeft - 1); +} + +export async function placement(api: GitHubApi, repo: string, mergeCommit: string): Promise { + const file = await api.request("GET", `/repos/${repo}/contents/pyproject.toml?ref=${mergeCommit}`); + const version = parseVersion(atob(file.content.replace(/\n/g, ""))); + if (version === undefined) { + return skip(`pyproject.toml at ${mergeCommit.slice(0, 10)} has no version line`); + } + return firstReleaseWith(api, repo, mergeCommit, version, MAX_MINOR_BUMPS); +} + +export function fixedBody(pullRequest: number, release: { readonly tag: string; readonly shipped: boolean }): string { + const availability = release.shipped + ? `This is in ${release.tag} and up, so upgrading to that release or any newer one picks it up.` + : `This ships in ${release.tag} and up, and the next dev pre-release cut from main will carry it too.`; + return `${FIXED_MARKER}\nFixed by #${pullRequest}. ${availability}`; +} + +export async function commentFixedIssue(api: GitHubApi, config: FixedConfig): Promise { + const [owner, name] = config.repo.split("/"); + const response = await api.request("POST", "/graphql", { + query: CLOSER_QUERY, + variables: { owner, name, number: config.issueNumber }, + }); + const issue = response.data?.repository?.issue ?? null; + if (issue === null) { + return skip("not an issue in this repository"); + } + const closer = closerOf(issue, config.defaultBranch); + if (closer.kind === "skip") { + return closer; + } + const issuePath = `/repos/${config.repo}/issues/${config.issueNumber}`; + const comments = await listAll(api, `${issuePath}/comments`); + if (comments.some((comment) => comment.body.includes(FIXED_MARKER))) { + return skip("already carries a fixed-in comment"); + } + const release = await placement(api, config.repo, closer.mergeCommit); + if (release.kind === "skip") { + return release; + } + const body = fixedBody(closer.number, release); + if (!config.dryRun) { + await api.request("POST", `${issuePath}/comments`, { body }); + } + return { kind: "commented", pullRequest: closer.number, tag: release.tag, body }; +} + +export function readConfig(env: Readonly>): FixedConfig & { readonly token: string } { + const token = env.GITHUB_TOKEN; + const repo = env.GITHUB_REPOSITORY; + const defaultBranch = env.DEFAULT_BRANCH; + if (!token || !repo || !/^[\w.-]+\/[\w.-]+$/.test(repo) || !defaultBranch) { + throw new Error("GITHUB_TOKEN, GITHUB_REPOSITORY (owner/repo) and DEFAULT_BRANCH are required"); + } + const issueNumber = Number(env.ISSUE_NUMBER); + if (!Number.isInteger(issueNumber) || issueNumber <= 0) { + throw new Error(`ISSUE_NUMBER must be a positive integer, got "${env.ISSUE_NUMBER}"`); + } + return { token, repo, issueNumber, defaultBranch, dryRun: env.DRY_RUN === "true" }; +} + +function describe(config: FixedConfig, verdict: FixedVerdict): string { + if (verdict.kind === "skip") { + return `#${config.issueNumber}: skipped, ${verdict.reason}`; + } + if (config.dryRun) { + return `#${config.issueNumber}: DRY RUN, set the ISSUE_FIXED_COMMENT_ENABLED repo variable to true to post this:\n\n${verdict.body}`; + } + return `#${config.issueNumber}: commented, fixed by #${verdict.pullRequest} in ${verdict.tag}`; +} + +if (import.meta.main) { + const { token, ...config } = readConfig(process.env); + console.log(describe(config, await commentFixedIssue(githubApi(token), config))); +} From a3aceec2f865b30c800b8ea9587582e13d6398ef Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Fri, 18 Sep 2026 17:37:08 -0700 Subject: [PATCH 100/144] fix(rust): match Python proxy, ssl_verify and client expiry behavior in the http pool Honor environment proxies whenever Python would use httpx (sync calls, HTTP/2, aiohttp disabled), apply the per-call ssl_verify argument, ignore empty or missing SSL env values the way http_handler.py does, expire pooled clients after an hour so rotated certificates reload, keep the client certificate off media downloads, and decline instead of raising when a litellm global has an unexpected type --- litellm-rust/crates/http/src/config.rs | 21 ++-- litellm-rust/crates/http/src/lib.rs | 4 - litellm-rust/crates/http/src/pool.rs | 80 +++++++++--- litellm-rust/crates/http/src/settings.rs | 82 ++++++++++-- .../llms/src/custom_httpx/llm_http_handler.rs | 2 +- .../crates/llms/src/custom_httpx/transport.rs | 6 - .../crates/python-bridge/python_settings.json | 1 + litellm-rust/crates/python-bridge/src/http.rs | 118 +++++++++++++++--- .../python-bridge/src/python_settings.rs | 6 - .../python-bridge/src/routes/ocr/mod.rs | 2 +- litellm/rust_bridge/settings.py | 8 +- .../test_litellm/rust_bridge/test_settings.py | 4 +- 12 files changed, 262 insertions(+), 72 deletions(-) diff --git a/litellm-rust/crates/http/src/config.rs b/litellm-rust/crates/http/src/config.rs index f27092c1fb5..7772e6cce5b 100644 --- a/litellm-rust/crates/http/src/config.rs +++ b/litellm-rust/crates/http/src/config.rs @@ -16,8 +16,6 @@ pub enum Verify { BuiltInRoots, } -/// One fully resolved client configuration. Every field is a plain value so the pool can -/// key cached clients on it. #[derive(Clone, Debug, PartialEq, Eq, Hash)] pub struct HttpClientConfig { pub verify: Verify, @@ -30,9 +28,6 @@ pub struct HttpClientConfig { } impl HttpClientConfig { - /// Port of `get_ssl_verify` + `get_ssl_configuration`: the configured (environment-overlaid) - /// `ssl_verify`, then `SSL_CERT_FILE`, then the built-in roots. Settings rustls has no - /// equivalent for are an error instead of a silent no-op. pub fn resolve(settings: &HttpSettings) -> Result { if let Some(level) = &settings.ssl_security_level { return Err(Error::Unsupported { @@ -60,12 +55,11 @@ impl HttpClientConfig { force_ipv4: settings.force_ipv4, http2: settings.http2, user_agent: settings.user_agent.clone(), - trust_proxy_env: settings.trust_proxy_env, + trust_proxy_env: settings.trust_proxy_env || settings.http2 || settings.httpx_transport, connect_timeout: settings.connect_timeout, }) } - /// A builder carrying every shared setting; variants add their own policy on top. pub fn client_builder(&self) -> Result { let base = reqwest::Client::builder().connect_timeout(self.connect_timeout); let with_roots = match &self.verify { @@ -242,6 +236,19 @@ mod tests { ); } + #[rstest] + #[case::aiohttp_default(HttpSettings::default(), false)] + #[case::aiohttp_trust_env(HttpSettings { trust_proxy_env: true, ..HttpSettings::default() }, true)] + #[case::http2_uses_httpx(HttpSettings { http2: true, ..HttpSettings::default() }, true)] + #[case::aiohttp_disabled(HttpSettings { httpx_transport: true, ..HttpSettings::default() }, true)] + fn environment_proxies_apply_whenever_python_would_use_httpx( + #[case] settings: HttpSettings, + #[case] expected: bool, + ) { + let config = HttpClientConfig::resolve(&settings).unwrap(); + assert_eq!(config.trust_proxy_env, expected); + } + #[test] fn missing_ca_bundle_is_a_read_error() { let path = std::env::temp_dir().join("litellm-http-missing-bundle.pem"); diff --git a/litellm-rust/crates/http/src/lib.rs b/litellm-rust/crates/http/src/lib.rs index 9c88e3101a7..c02a82539ff 100644 --- a/litellm-rust/crates/http/src/lib.rs +++ b/litellm-rust/crates/http/src/lib.rs @@ -1,7 +1,3 @@ -//! Rust counterpart of `litellm/llms/custom_httpx/http_handler.py`: the plain HTTP settings -//! LiteLLM exposes, their resolution into one typed client configuration, and a pool that -//! caches `reqwest::Client`s per resolved configuration. - mod config; mod error; mod pool; diff --git a/litellm-rust/crates/http/src/pool.rs b/litellm-rust/crates/http/src/pool.rs index 613ce9c2831..0d9b1abf504 100644 --- a/litellm-rust/crates/http/src/pool.rs +++ b/litellm-rust/crates/http/src/pool.rs @@ -1,33 +1,44 @@ use std::{ collections::HashMap, sync::{Arc, Mutex, MutexGuard, PoisonError}, + time::{Duration, Instant}, }; use reqwest::dns::Resolve; use crate::{config::HttpClientConfig, error::Error}; -/// The client shapes routes need; each is the shared base plus one policy. #[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)] pub enum ClientVariant { Provider, NoRedirect, - /// Media downloads: no redirects (the fetcher validates each hop), never a proxy, and the - /// pool's media resolver. Media, } -/// Counterpart of `get_async_httpx_client`: one `reqwest::Client` per resolved configuration -/// and variant, built on first use and shared afterwards. +const CLIENT_TTL: Duration = Duration::from_secs(3600); + +struct PooledClient { + client: reqwest::Client, + built_at: Instant, +} + +type Clients = HashMap<(HttpClientConfig, ClientVariant), PooledClient>; + pub struct HttpClientPool { media_resolver: Arc, - clients: Mutex>, + ttl: Duration, + clients: Mutex, } impl HttpClientPool { pub fn new(media_resolver: Arc) -> Self { + Self::with_ttl(media_resolver, CLIENT_TTL) + } + + pub fn with_ttl(media_resolver: Arc, ttl: Duration) -> Self { Self { media_resolver, + ttl, clients: Mutex::default(), } } @@ -37,15 +48,31 @@ impl HttpClientPool { config: &HttpClientConfig, variant: ClientVariant, ) -> Result { - let key = (config.clone(), variant); - if let Some(client) = self.lock().get(&key) { - return Ok(client.clone()); + let effective = match variant { + ClientVariant::Media => HttpClientConfig { + client_certificate: None, + ..config.clone() + }, + ClientVariant::Provider | ClientVariant::NoRedirect => config.clone(), + }; + let key = (effective, variant); + if let Some(pooled) = self.lock().get(&key) + && pooled.built_at.elapsed() < self.ttl + { + return Ok(pooled.client.clone()); } - let client = self.apply(variant, config.client_builder()?).build()?; - Ok(self.lock().entry(key).or_insert(client).clone()) + let client = self.apply(variant, key.0.client_builder()?).build()?; + self.lock().insert( + key, + PooledClient { + client: client.clone(), + built_at: Instant::now(), + }, + ); + Ok(client) } - fn lock(&self) -> MutexGuard<'_, HashMap<(HttpClientConfig, ClientVariant), reqwest::Client>> { + fn lock(&self) -> MutexGuard<'_, Clients> { self.clients.lock().unwrap_or_else(PoisonError::into_inner) } @@ -102,8 +129,6 @@ mod tests { } } - /// Answers every request on every connection with `status_line` and counts connections, - /// so a reused client shows up as a reused keep-alive connection. async fn serve( status_line: &'static str, ) -> (SocketAddr, Arc, Arc>>) { @@ -168,6 +193,33 @@ mod tests { assert_eq!(connections.load(Ordering::SeqCst), 3); } + #[tokio::test] + async fn expired_clients_are_rebuilt() { + let (address, connections, _) = serve("HTTP/1.1 204 No Content").await; + let url = format!("http://{address}"); + let pool = HttpClientPool::with_ttl( + Arc::new(FixedResolver(([192, 0, 2, 1], 80).into())), + Duration::ZERO, + ); + get(&pool, &config("a"), ClientVariant::Provider, &url).await; + get(&pool, &config("a"), ClientVariant::Provider, &url).await; + assert_eq!(connections.load(Ordering::SeqCst), 2); + } + + #[test] + fn media_variant_never_loads_the_client_certificate() { + let pool = pool(); + let with_identity = HttpClientConfig { + client_certificate: Some(std::env::temp_dir().join("litellm-http-absent-client.pem")), + ..config("a") + }; + assert!( + pool.client(&with_identity, ClientVariant::Provider) + .is_err() + ); + assert!(pool.client(&with_identity, ClientVariant::Media).is_ok()); + } + #[test] fn build_failures_are_not_cached() { let pool = pool(); diff --git a/litellm-rust/crates/http/src/settings.rs b/litellm-rust/crates/http/src/settings.rs index 45aab0d6fa5..55ac471bba9 100644 --- a/litellm-rust/crates/http/src/settings.rs +++ b/litellm-rust/crates/http/src/settings.rs @@ -1,6 +1,8 @@ -use std::{path::PathBuf, time::Duration}; +use std::{ + path::{Path, PathBuf}, + time::Duration, +}; -/// `litellm.ssl_verify` / `SSL_VERIFY`: a bool or a CA bundle path. #[derive(Clone, Debug, PartialEq, Eq, Hash)] pub enum SslVerify { Enabled, @@ -18,7 +20,6 @@ impl SslVerify { } } -/// The plain inputs `http_handler.py` reads from `litellm.*` globals and the environment. #[derive(Clone, Debug, PartialEq, Eq)] pub struct HttpSettings { pub ssl_verify: Option, @@ -28,6 +29,7 @@ pub struct HttpSettings { pub ssl_ecdh_curve: Option, pub force_ipv4: bool, pub http2: bool, + pub httpx_transport: bool, pub user_agent: Option, pub trust_proxy_env: bool, pub connect_timeout: Duration, @@ -43,6 +45,7 @@ impl Default for HttpSettings { ssl_ecdh_curve: None, force_ipv4: false, http2: false, + httpx_transport: false, user_agent: None, trust_proxy_env: false, connect_timeout: Duration::from_secs(5), @@ -51,11 +54,6 @@ impl Default for HttpSettings { } impl HttpSettings { - /// Overlay the environment variables `http_handler.py` consults, with the same precedence: - /// `SSL_VERIFY`, `SSL_CERTIFICATE`, `SSL_SECURITY_LEVEL`, `SSL_ECDH_CURVE` and - /// `LITELLM_USER_AGENT` win over the configured value; `SSL_CERT_FILE` only applies when - /// verification is on without an explicit bundle; `LITELLM_HTTP2` and `AIOHTTP_TRUST_ENV` - /// can only turn their switch on. pub fn with_environment(self, env: &(dyn Fn(&str) -> Option + Sync)) -> Self { let enabled = |name: &str| env(name).is_some_and(|value| value.trim().eq_ignore_ascii_case("true")); @@ -68,15 +66,32 @@ impl HttpSettings { .or(self.ssl_cert_file), ssl_certificate: env("SSL_CERTIFICATE") .map(PathBuf::from) - .or(self.ssl_certificate), - ssl_security_level: env("SSL_SECURITY_LEVEL").or(self.ssl_security_level), - ssl_ecdh_curve: env("SSL_ECDH_CURVE").or(self.ssl_ecdh_curve), + .or(self.ssl_certificate) + .filter(|path| !path.as_os_str().is_empty()), + ssl_security_level: env("SSL_SECURITY_LEVEL") + .or(self.ssl_security_level) + .filter(|level| !level.is_empty()), + ssl_ecdh_curve: env("SSL_ECDH_CURVE") + .or(self.ssl_ecdh_curve) + .filter(|curve| !curve.is_empty()), http2: self.http2 || enabled("LITELLM_HTTP2"), + httpx_transport: self.httpx_transport || enabled("DISABLE_AIOHTTP_TRANSPORT"), user_agent: env("LITELLM_USER_AGENT").or(self.user_agent), trust_proxy_env: self.trust_proxy_env || enabled("AIOHTTP_TRUST_ENV"), ..self } } + + pub fn without_missing_files(self, exists: &dyn Fn(&Path) -> bool) -> Self { + Self { + ssl_verify: match self.ssl_verify { + Some(SslVerify::CaBundle(path)) if !exists(&path) => Some(SslVerify::Enabled), + other => other, + }, + ssl_cert_file: self.ssl_cert_file.filter(|path| exists(path)), + ..self + } + } } #[cfg(test)] @@ -152,6 +167,46 @@ mod tests { assert_eq!(configured.clone().with_environment(&no_env), configured); } + #[test] + fn empty_environment_values_clear_the_setting_like_python_truthiness() { + let settings = HttpSettings { + ssl_certificate: Some("/configured/client.pem".into()), + ssl_security_level: Some("configured".into()), + ssl_ecdh_curve: Some("X25519".into()), + ..HttpSettings::default() + } + .with_environment(&env_of(&[ + ("SSL_CERTIFICATE", ""), + ("SSL_SECURITY_LEVEL", ""), + ("SSL_ECDH_CURVE", ""), + ])); + assert_eq!(settings.ssl_certificate, None); + assert_eq!(settings.ssl_security_level, None); + assert_eq!(settings.ssl_ecdh_curve, None); + } + + #[test] + fn missing_files_fall_back_to_default_verification() { + let settings = HttpSettings { + ssl_verify: Some(SslVerify::CaBundle("/absent/roots.pem".into())), + ssl_cert_file: Some("/absent/env.pem".into()), + ..HttpSettings::default() + } + .without_missing_files(&|_| false); + assert_eq!(settings.ssl_verify, Some(SslVerify::Enabled)); + assert_eq!(settings.ssl_cert_file, None); + } + + #[test] + fn existing_files_are_kept() { + let settings = HttpSettings { + ssl_verify: Some(SslVerify::CaBundle("/present/roots.pem".into())), + ssl_cert_file: Some("/present/env.pem".into()), + ..HttpSettings::default() + }; + assert_eq!(settings.clone().without_missing_files(&|_| true), settings); + } + #[rstest] #[case("true", true)] #[case("True", true)] @@ -159,11 +214,14 @@ mod tests { #[case("1", false)] fn boolean_switches_only_turn_on_for_true(#[case] value: &'static str, #[case] expected: bool) { let env = move |name: &str| match name { - "LITELLM_HTTP2" | "AIOHTTP_TRUST_ENV" => Some(value.to_string()), + "LITELLM_HTTP2" | "AIOHTTP_TRUST_ENV" | "DISABLE_AIOHTTP_TRANSPORT" => { + Some(value.to_string()) + } _ => None, }; let settings = HttpSettings::default().with_environment(&env); assert_eq!(settings.http2, expected); + assert_eq!(settings.httpx_transport, expected); assert_eq!(settings.trust_proxy_env, expected); } } diff --git a/litellm-rust/crates/llms/src/custom_httpx/llm_http_handler.rs b/litellm-rust/crates/llms/src/custom_httpx/llm_http_handler.rs index 425b71efc78..876fa0aae87 100644 --- a/litellm-rust/crates/llms/src/custom_httpx/llm_http_handler.rs +++ b/litellm-rust/crates/llms/src/custom_httpx/llm_http_handler.rs @@ -42,7 +42,7 @@ impl OcrClient { pool: &HttpClientPool, config: &HttpClientConfig, vertex_auth: VertexAuth, - ) -> Result { + ) -> Result { Ok(Self { provider_http: pool.client(config, ClientVariant::Provider)?, polling_http: pool.client(config, ClientVariant::NoRedirect)?, diff --git a/litellm-rust/crates/llms/src/custom_httpx/transport.rs b/litellm-rust/crates/llms/src/custom_httpx/transport.rs index 8e5e1a8832d..172dd96476a 100644 --- a/litellm-rust/crates/llms/src/custom_httpx/transport.rs +++ b/litellm-rust/crates/llms/src/custom_httpx/transport.rs @@ -26,12 +26,6 @@ impl From for Error { } } -impl From for Error { - fn from(error: litellm_http::Error) -> Self { - Self::Connect(error.to_string()) - } -} - #[cfg(test)] mod tests { #[tokio::test] diff --git a/litellm-rust/crates/python-bridge/python_settings.json b/litellm-rust/crates/python-bridge/python_settings.json index a6cd959c6de..64dd01a0a84 100644 --- a/litellm-rust/crates/python-bridge/python_settings.json +++ b/litellm-rust/crates/python-bridge/python_settings.json @@ -7,6 +7,7 @@ "force_ipv4", "http2", "aiohttp_trust_env", + "disable_aiohttp_transport", "user_agent" ] } diff --git a/litellm-rust/crates/python-bridge/src/http.rs b/litellm-rust/crates/python-bridge/src/http.rs index de6fb5bb96d..cf7ca05515e 100644 --- a/litellm-rust/crates/python-bridge/src/http.rs +++ b/litellm-rust/crates/python-bridge/src/http.rs @@ -1,5 +1,5 @@ use std::{ - path::PathBuf, + path::{Path, PathBuf}, sync::{Arc, LazyLock}, }; @@ -12,27 +12,46 @@ use crate::{errors::RustBridgeDeclined, python_settings::PythonSettings}; static POOL: LazyLock = LazyLock::new(|| HttpClientPool::new(Arc::new(PublicDnsResolver))); -/// Keyword arguments that carry a live Python HTTP client or session. They cannot cross into -/// Rust, so a call that supplies one stays on the Python path. const LIVE_CLIENT_ARGUMENTS: [&str; 3] = ["client", "shared_session", "aclient_session"]; pub(crate) fn pool() -> &'static HttpClientPool { &POOL } -/// The client configuration for one call: the `litellm.*` HTTP settings with the environment -/// overlaid, the same way `http_handler.py` combines them. pub(crate) fn call_config( py: Python<'_>, kwargs: &Bound<'_, PyDict>, + asynchronous: bool, ) -> PyResult { decline_live_clients(kwargs)?; - let settings = settings(&PythonSettings::Http.read(py)?)? + let configured = settings(&PythonSettings::Http.read(py)?)? .with_environment(&|name| std::env::var(name).ok()); + let settings = for_call(configured, call_ssl_verify(kwargs)?, asynchronous) + .without_missing_files(&|path: &Path| path.exists()); HttpClientConfig::resolve(&settings) .map_err(|error| RustBridgeDeclined::new_err(error.to_string())) } +fn call_ssl_verify(kwargs: &Bound<'_, PyDict>) -> PyResult> { + kwargs + .get_item("ssl_verify")? + .filter(|value| !value.is_none()) + .map(|value| ssl_verify(&value, "the ssl_verify argument")) + .transpose() +} + +fn for_call( + configured: HttpSettings, + call_ssl_verify: Option, + asynchronous: bool, +) -> HttpSettings { + HttpSettings { + ssl_verify: call_ssl_verify.or(configured.ssl_verify), + httpx_transport: configured.httpx_transport || !asynchronous, + ..configured + } +} + pub(crate) fn decline_live_clients(kwargs: &Bound<'_, PyDict>) -> PyResult<()> { for name in LIVE_CLIENT_ARGUMENTS { if kwargs.get_item(name)?.is_some_and(|value| !value.is_none()) { @@ -53,25 +72,31 @@ struct PythonHttpSettings<'py> { force_ipv4: bool, http2: bool, aiohttp_trust_env: bool, + disable_aiohttp_transport: bool, user_agent: String, } fn settings(value: &Bound<'_, PyAny>) -> PyResult { - let python: PythonHttpSettings = value.extract()?; + let python: PythonHttpSettings = value.extract().map_err(|error: PyErr| { + RustBridgeDeclined::new_err(format!( + "litellm HTTP settings cannot be used by the Rust route: {error}" + )) + })?; Ok(HttpSettings { - ssl_verify: Some(ssl_verify(&python.ssl_verify)?), + ssl_verify: Some(ssl_verify(&python.ssl_verify, "litellm.ssl_verify")?), ssl_certificate: python.ssl_certificate.map(PathBuf::from), ssl_security_level: python.ssl_security_level, ssl_ecdh_curve: python.ssl_ecdh_curve, force_ipv4: python.force_ipv4, http2: python.http2, + httpx_transport: python.disable_aiohttp_transport, user_agent: Some(python.user_agent), trust_proxy_env: python.aiohttp_trust_env, ..HttpSettings::default() }) } -fn ssl_verify(value: &Bound<'_, PyAny>) -> PyResult { +fn ssl_verify(value: &Bound<'_, PyAny>, source: &str) -> PyResult { if let Ok(enabled) = value.extract::() { return Ok(if enabled { SslVerify::Enabled @@ -82,9 +107,9 @@ fn ssl_verify(value: &Bound<'_, PyAny>) -> PyResult { if let Ok(path) = value.extract::() { return Ok(SslVerify::parse(&path)); } - Err(RustBridgeDeclined::new_err( - "litellm.ssl_verify is a live Python object and cannot be used by the Rust route", - )) + Err(RustBridgeDeclined::new_err(format!( + "{source} is a live Python object and cannot be used by the Rust route" + ))) } #[cfg(test)] @@ -95,8 +120,6 @@ mod tests { use super::*; use crate::python_settings::CONTRACT; - /// A stand-in for `http_settings()` carrying exactly the fields the contract declares, so a - /// field Rust reads but Python does not return fails here. fn python_settings<'py>(py: Python<'py>, overrides: &str) -> Bound<'py, PyAny> { let source = format!( " @@ -110,6 +133,7 @@ defaults = dict( force_ipv4=False, http2=False, aiohttp_trust_env=False, + disable_aiohttp_transport=False, user_agent='litellm/test', ) defaults.update(dict({overrides})) @@ -153,6 +177,7 @@ ssl_ecdh_curve='X25519', force_ipv4=True, http2=True, aiohttp_trust_env=True, +disable_aiohttp_transport=True, user_agent='litellm/9.9.9', ", )) @@ -166,6 +191,7 @@ user_agent='litellm/9.9.9', ssl_ecdh_curve: Some("X25519".into()), force_ipv4: true, http2: true, + httpx_transport: true, user_agent: Some("litellm/9.9.9".into()), trust_proxy_env: true, ..HttpSettings::default() @@ -214,6 +240,70 @@ user_agent='litellm/9.9.9', }); } + #[test] + fn mistyped_python_settings_decline_instead_of_raising() { + Python::initialize(); + Python::attach(|py| { + let error = settings(&python_settings(py, "force_ipv4='yes'")).unwrap_err(); + assert!(error.is_instance_of::(py)); + }); + } + + #[test] + fn call_ssl_verify_beats_the_configured_and_environment_value() { + Python::initialize(); + Python::attach(|py| { + let kwargs = PyDict::new(py); + kwargs.set_item("ssl_verify", false).unwrap(); + let configured = HttpSettings { + ssl_verify: Some(SslVerify::Enabled), + ..HttpSettings::default() + }; + let settings = for_call(configured, call_ssl_verify(&kwargs).unwrap(), true); + assert_eq!(settings.ssl_verify, Some(SslVerify::Disabled)); + }); + } + + #[test] + fn absent_call_ssl_verify_keeps_the_configured_value() { + Python::initialize(); + Python::attach(|py| { + let kwargs = PyDict::new(py); + kwargs.set_item("ssl_verify", py.None()).unwrap(); + let configured = HttpSettings { + ssl_verify: Some(SslVerify::Disabled), + ..HttpSettings::default() + }; + let settings = for_call(configured.clone(), call_ssl_verify(&kwargs).unwrap(), true); + assert_eq!(settings, configured); + }); + } + + #[test] + fn live_ssl_context_argument_declines() { + Python::initialize(); + Python::attach(|py| { + let kwargs = PyDict::new(py); + kwargs + .set_item("ssl_verify", py.eval(c"object()", None, None).unwrap()) + .unwrap(); + let error = call_ssl_verify(&kwargs).unwrap_err(); + assert!(error.is_instance_of::(py)); + }); + } + + #[rstest] + #[case::asynchronous(true, false)] + #[case::synchronous(false, true)] + fn synchronous_calls_honor_environment_proxies_like_httpx( + #[case] asynchronous: bool, + #[case] expected: bool, + ) { + let settings = for_call(HttpSettings::default(), None, asynchronous); + let config = HttpClientConfig::resolve(&settings).unwrap(); + assert_eq!(config.trust_proxy_env, expected); + } + #[rstest] #[case::client("client")] #[case::shared_session("shared_session")] diff --git a/litellm-rust/crates/python-bridge/src/python_settings.rs b/litellm-rust/crates/python-bridge/src/python_settings.rs index c5f9f309615..dcb46e7d2b5 100644 --- a/litellm-rust/crates/python-bridge/src/python_settings.rs +++ b/litellm-rust/crates/python-bridge/src/python_settings.rs @@ -2,12 +2,6 @@ use pyo3::prelude::*; const MODULE: &str = "litellm.rust_bridge.settings"; -/// Every group of `litellm.*` module globals the native routes read. Environment overrides are -/// applied on the Rust side, so each function returns only what the Python process configured. -/// A group is deleted once Rust owns loading that configuration, so this enum only shrinks. -/// -/// `litellm/rust_bridge/settings.py` is the only Python module behind it, and -/// `python_settings.json` pins the fields each function returns on both sides. #[derive(Clone, Copy, Debug, PartialEq, Eq)] pub(crate) enum PythonSettings { Http, diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs b/litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs index d9aeeb234f7..174d0ff18c8 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs @@ -37,7 +37,7 @@ fn run_ocr( kwargs: Bound<'_, PyDict>, asynchronous: bool, ) -> PyResult> { - let config = http::call_config(py, &kwargs)?; + let config = http::call_config(py, &kwargs, asynchronous)?; let client = OcrClient::new(http::pool(), &config, VERTEX_AUTH.clone()) .map_err(|error| RustBridgeDeclined::new_err(error.to_string()))?; run_legacy_call( diff --git a/litellm/rust_bridge/settings.py b/litellm/rust_bridge/settings.py index a8229b12d13..ad478fb28b5 100644 --- a/litellm/rust_bridge/settings.py +++ b/litellm/rust_bridge/settings.py @@ -1,9 +1,3 @@ -"""The `litellm.*` module globals the native routes read. - -Environment variables that override these are applied in Rust, so nothing here reads `os.environ`. -`litellm-rust/crates/python-bridge/python_settings.json` pins the fields each function returns. -""" - from __future__ import annotations from dataclasses import dataclass @@ -18,6 +12,7 @@ class HttpSettings: force_ipv4: bool http2: bool aiohttp_trust_env: bool + disable_aiohttp_transport: bool user_agent: str @@ -33,5 +28,6 @@ def http_settings() -> HttpSettings: force_ipv4=litellm.force_ipv4, http2=litellm.http2, aiohttp_trust_env=litellm.aiohttp_trust_env, + disable_aiohttp_transport=litellm.disable_aiohttp_transport, user_agent=default_user_agent(), ) diff --git a/tests/test_litellm/rust_bridge/test_settings.py b/tests/test_litellm/rust_bridge/test_settings.py index 618fa400136..dce2324de08 100644 --- a/tests/test_litellm/rust_bridge/test_settings.py +++ b/tests/test_litellm/rust_bridge/test_settings.py @@ -26,6 +26,7 @@ def test_http_settings_reads_the_litellm_globals(monkeypatch: pytest.MonkeyPatch monkeypatch.setattr(litellm, "force_ipv4", True) monkeypatch.setattr(litellm, "http2", True) monkeypatch.setattr(litellm, "aiohttp_trust_env", True) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) assert settings.http_settings() == settings.HttpSettings( ssl_verify="/etc/ssl/corp.pem", @@ -35,7 +36,8 @@ def test_http_settings_reads_the_litellm_globals(monkeypatch: pytest.MonkeyPatch force_ipv4=True, http2=True, aiohttp_trust_env=True, - user_agent=settings.http_settings().user_agent, + disable_aiohttp_transport=True, + user_agent=default_user_agent(), ) From 7208e310f038b6559f8ed99f4ddd295ebbc9987c Mon Sep 17 00:00:00 2001 From: kerry Date: Sat, 19 Sep 2026 00:39:57 +0000 Subject: [PATCH 101/144] fix(schema): annotate new off_peak_pricing constants with Final Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ci_cd/generate_model_prices_schema.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/ci_cd/generate_model_prices_schema.py b/ci_cd/generate_model_prices_schema.py index 4ba4368e33c..8eec07dadda 100644 --- a/ci_cd/generate_model_prices_schema.py +++ b/ci_cd/generate_model_prices_schema.py @@ -3,7 +3,7 @@ from __future__ import annotations import json import sys from pathlib import Path -from typing import Optional +from typing import Final, Optional import jsonschema @@ -19,8 +19,8 @@ NONNEG_NUMBER: JsonSchema = {"type": "number", "minimum": 0} NONNEG_INTEGER: JsonSchema = {"type": "integer", "minimum": 0} BOOLEAN: JsonSchema = {"type": "boolean"} STRING: JsonSchema = {"type": "string"} -TIME_WINDOW: JsonSchema = {"type": "string", "pattern": r"^([01]\d|2[0-3]):[0-5]\d-([01]\d|2[0-3]):[0-5]\d$"} -WEEKDAY_PATTERN = ( +TIME_WINDOW: Final[JsonSchema] = {"type": "string", "pattern": r"^([01]\d|2[0-3]):[0-5]\d-([01]\d|2[0-3]):[0-5]\d$"} +WEEKDAY_PATTERN: Final = ( r"(?i)^(mon|monday|tue|tues|tuesday|wed|wednesday|thu|thur|thurs|thursday|fri|friday|sat|saturday|sun|sunday)$" ) @@ -35,12 +35,12 @@ EXTRA_BOOLEAN_KEYS = frozenset( } ) -HOURS_UTC: JsonSchema = { +HOURS_UTC: Final[JsonSchema] = { "description": 'UTC "HH:MM-HH:MM" window, or a list of them; a window may wrap past midnight.', "oneOf": [TIME_WINDOW, {"type": "array", "items": TIME_WINDOW, "minItems": 1}], } -OFF_PEAK_WINDOW: JsonSchema = { +OFF_PEAK_WINDOW: Final[JsonSchema] = { "type": "object", "properties": { "hours_utc": HOURS_UTC, From aa0fb915d0b0ba6d478c725e9034b10c4251fb0d Mon Sep 17 00:00:00 2001 From: yassin Date: Sat, 19 Sep 2026 00:40:51 +0000 Subject: [PATCH 102/144] fix(rate_limiter): render the 429 reset time in UTC as labelled The proxy rate limiters formatted the reset epoch with a naive datetime.fromtimestamp, which reads the process timezone, and then appended a literal UTC suffix. A proxy running outside UTC returned a local wall-clock time labelled as UTC in the 429 body and reset_at header. Convert with tz=timezone.utc in both the request limiter and the batch limiter so the label is true Co-authored-by: Priyansh Nandwana Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/hooks/batch_rate_limiter.py | 6 +- .../hooks/parallel_request_limiter_v3.py | 6 +- .../proxy/hooks/test_batch_rate_limiter.py | 41 +++++++++++++- .../hooks/test_parallel_request_limiter_v3.py | 56 ++++++++++++++++++- 4 files changed, 101 insertions(+), 8 deletions(-) diff --git a/litellm/proxy/hooks/batch_rate_limiter.py b/litellm/proxy/hooks/batch_rate_limiter.py index ab6e10ca76b..a5b6cabf519 100644 --- a/litellm/proxy/hooks/batch_rate_limiter.py +++ b/litellm/proxy/hooks/batch_rate_limiter.py @@ -19,7 +19,7 @@ Quick summary: import json from collections.abc import Callable, Iterable, Mapping, Sequence -from datetime import datetime +from datetime import datetime, timezone from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, NoReturn, TypeAlias @@ -661,7 +661,9 @@ class _PROXY_BatchRateLimiter(CustomLogger): ) or self.parallel_request_limiter.window_size reset_time: Final = now + window_size if window_start is None else window_start + window_size retry_after: Final = max(0, int(reset_time - now)) - reset_time_formatted: Final = datetime.fromtimestamp(reset_time).strftime("%Y-%m-%d %H:%M:%S UTC") + reset_time_formatted: Final = datetime.fromtimestamp(reset_time, tz=timezone.utc).strftime( + "%Y-%m-%d %H:%M:%S UTC" + ) remaining_display: Final = max(0, status["limit_remaining"]) current_limit: Final = status["current_limit"] diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index cdec5922fff..a6b00be1091 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -13,7 +13,7 @@ from collections.abc import AsyncGenerator, Awaitable, Callable, Mapping, Sequen from contextlib import asynccontextmanager from contextvars import ContextVar from dataclasses import dataclass, field -from datetime import datetime +from datetime import datetime, timezone from types import MappingProxyType from typing import ( TYPE_CHECKING, @@ -3124,7 +3124,9 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): now = self._get_current_time().timestamp() reset_time = now + self.window_size - reset_time_formatted = datetime.fromtimestamp(reset_time).strftime("%Y-%m-%d %H:%M:%S UTC") + reset_time_formatted = datetime.fromtimestamp(reset_time, tz=timezone.utc).strftime( + "%Y-%m-%d %H:%M:%S UTC" + ) remaining_display = max(0, status["limit_remaining"]) rate_limit_type = status["rate_limit_type"] diff --git a/tests/test_litellm/proxy/hooks/test_batch_rate_limiter.py b/tests/test_litellm/proxy/hooks/test_batch_rate_limiter.py index 919e9c79828..930f62fcd10 100644 --- a/tests/test_litellm/proxy/hooks/test_batch_rate_limiter.py +++ b/tests/test_litellm/proxy/hooks/test_batch_rate_limiter.py @@ -6,7 +6,10 @@ batch under a per-minute RPM/TPM budget. Scopes that configure `tpd_limit` are charged against a 24h token window instead of their minute counters. """ -from datetime import datetime +import time +from collections.abc import Iterator +from datetime import datetime, timezone +from typing import Final import pytest from fastapi import HTTPException @@ -257,3 +260,39 @@ def test_online_descriptors_ignore_tpd_limit(): model_has_failures=False, ) assert [(d["key"], d["rate_limit"]["window_size"]) for d in descriptors] == [("api_key", rate_limiter.window_size)] + + +@pytest.fixture(params=["Europe/Paris", "Asia/Kolkata", "America/Los_Angeles"]) +def process_timezone(request: pytest.FixtureRequest, monkeypatch: pytest.MonkeyPatch) -> Iterator[str]: + monkeypatch.setenv("TZ", request.param) + time.tzset() + yield request.param + monkeypatch.undo() + time.tzset() + + +@pytest.mark.skipif(not hasattr(time, "tzset"), reason="switching the process timezone needs time.tzset()") +@pytest.mark.asyncio +async def test_batch_rate_limit_error_reports_reset_time_in_utc_on_a_non_utc_proxy(process_timezone: str) -> None: + window_start: Final = datetime(2026, 9, 13, 8, 0, 0, tzinfo=timezone.utc) + clock: Final = _Clock(window_start) + _internal_usage_cache, _rate_limiter, batch_limiter = _make_limiters(clock) + user_api_key_dict: Final = UserAPIKeyAuth(api_key=hash_token("tpd-key-utc"), rpm_limit=1, tpd_limit=1000) + + await batch_limiter._check_and_increment_batch_counters( + user_api_key_dict=user_api_key_dict, + data={}, + batch_usage=BatchFileUsage(total_tokens=600, request_count=6), + ) + clock.now = datetime(2026, 9, 13, 11, 0, 0, tzinfo=timezone.utc) + with pytest.raises(HTTPException) as exc: + await batch_limiter._check_and_increment_batch_counters( + user_api_key_dict=user_api_key_dict, + data={}, + batch_usage=BatchFileUsage(total_tokens=600, request_count=6), + ) + + assert exc.value.status_code == 429 + assert exc.value.headers["retry-after"] == str(BATCH_TPD_WINDOW_SECONDS - 3 * 3600) + assert exc.value.headers["reset_at"] == "2026-09-14 08:00:00 UTC" + assert str(exc.value.detail).endswith("Limit resets at: 2026-09-14 08:00:00 UTC") diff --git a/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py b/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py index 5889b1b513f..4907b4ea054 100644 --- a/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py +++ b/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py @@ -7,10 +7,10 @@ import logging import os import sys import time -from collections.abc import Sequence +from collections.abc import Iterator, Sequence from contextlib import contextmanager -from datetime import datetime, timedelta -from typing import Any, Dict, List, Optional +from datetime import datetime, timedelta, timezone +from typing import Any, Dict, Final, List, Optional import pytest from fastapi import HTTPException @@ -21,10 +21,12 @@ from litellm.caching.caching import DualCache from litellm.caching.in_memory_cache import InMemoryCache from litellm.constants import INTERNAL_CALL_ORIGIN_METADATA_KEY from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError from litellm.proxy.hooks.parallel_request_limiter_v3 import ( PARALLEL_REQUEST_SLOT_TTL_SECONDS, ParallelSlotAcquisition, RateLimitDescriptor, + RateLimitResponse, RequestRateLimiterStash, _request_stash, get_or_create_request_stash, @@ -6911,3 +6913,51 @@ def test_success_tpm_accounting_skips_team_model_pool_when_key_owns_model_tpm_li assert handler.create_rate_limit_keys("model_per_key", f"{hash_token('sk-pool')}:test-model", "tokens") in charged_keys team_pool_key = handler.create_rate_limit_keys("model_per_team", "t:test-model", "tokens") assert (team_pool_key in charged_keys) is charges_team_model_pool + + +@pytest.fixture(params=["Europe/Paris", "Asia/Kolkata", "America/Los_Angeles"]) +def process_timezone(request: pytest.FixtureRequest, monkeypatch: pytest.MonkeyPatch) -> Iterator[str]: + monkeypatch.setenv("TZ", request.param) + time.tzset() + yield request.param + monkeypatch.undo() + time.tzset() + + +@pytest.mark.skipif(not hasattr(time, "tzset"), reason="switching the process timezone needs time.tzset()") +def test_rate_limit_error_reports_reset_time_in_utc_on_a_non_utc_proxy(process_timezone: str) -> None: + now: Final = datetime(2026, 9, 4, 21, 53, 21, tzinfo=timezone.utc) + handler: Final = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(DualCache()), time_provider=lambda: now + ) + expected_reset: Final = (now + timedelta(seconds=handler.window_size)).strftime("%Y-%m-%d %H:%M:%S UTC") + over_limit: Final[RateLimitResponse] = { + "overall_code": "OVER_LIMIT", + "statuses": [ + { + "code": "OVER_LIMIT", + "descriptor_key": "api_key", + "limit_remaining": 0, + "rate_limit_type": "requests", + "current_limit": 2, + } + ], + } + + with pytest.raises(ProxyRateLimitError) as exc_info: + handler._handle_rate_limit_error( + response=over_limit, + descriptors=[{"key": "api_key", "value": "sk-test", "rate_limit": None}], + requested_model="gpt-4o-mini", + ) + + assert exc_info.value.status_code == 429 + assert exc_info.value.headers == { + "retry-after": str(handler.window_size), + "rate_limit_type": "requests", + "reset_at": expected_reset, + } + assert exc_info.value.detail == ( + "Rate limit exceeded for api_key: sk-test. Limit type: requests. " + f"Current limit: 2, Remaining: 0. Limit resets at: {expected_reset}" + ) From 99659e9e7e50b01c86e88b2e570179ff0f943acb Mon Sep 17 00:00:00 2001 From: kerry Date: Sat, 19 Sep 2026 00:43:51 +0000 Subject: [PATCH 103/144] test(bedrock): drop redundant recorder docstring Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/test_litellm/llms/bedrock/batches/test_handler.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/tests/test_litellm/llms/bedrock/batches/test_handler.py b/tests/test_litellm/llms/bedrock/batches/test_handler.py index d328e09056b..e69098a460d 100644 --- a/tests/test_litellm/llms/bedrock/batches/test_handler.py +++ b/tests/test_litellm/llms/bedrock/batches/test_handler.py @@ -585,8 +585,6 @@ class _JsonBody: class _AuthorizationRecorder: - """Stands in for botocore's HTTP session and records the Authorization header of every request it receives.""" - def __init__(self, body: Mapping[str, object]) -> None: self._payload: Final = json.dumps(body, default=str).encode() self.authorization_headers: tuple[str, ...] = () From 3911d62bbeee55f940ac4294275ada2c5580bf88 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 18 Sep 2026 17:44:28 -0700 Subject: [PATCH 104/144] fix(vertex_ai): prune a discarded turn's id once its marker is delivered --- .../audio_transcription/realtime_backend.py | 27 ++++++++++++++----- .../test_vertex_ai_realtime_backend.py | 1 + 2 files changed, 21 insertions(+), 7 deletions(-) diff --git a/litellm/llms/vertex_ai/audio_transcription/realtime_backend.py b/litellm/llms/vertex_ai/audio_transcription/realtime_backend.py index 4c8338c027e..859a883463c 100644 --- a/litellm/llms/vertex_ai/audio_transcription/realtime_backend.py +++ b/litellm/llms/vertex_ai/audio_transcription/realtime_backend.py @@ -87,10 +87,16 @@ class _TurnResult: @dataclass(frozen=True, slots=True) class _TurnDiscarded: - pass + turn: int -_OutboxItem = str | _TurnResult | _StreamFailure | _Closed +@dataclass(frozen=True, slots=True) +class _TurnDiscardedEvent: + turn: int + event: str + + +_OutboxItem = str | _TurnResult | _TurnDiscardedEvent | _StreamFailure | _Closed def open_speech_client(target: SpeechStreamingTarget, access_token: str) -> SpeechStreamingClient: @@ -304,6 +310,9 @@ class SpeechStreamingBackend: raise _normal_closure() case _TurnResult(): return None if item.turn in self._discarded_turns else item.event + case _TurnDiscardedEvent(): + self._discarded_turns -= {item.turn} + return item.event case str(): return item case _: @@ -345,7 +354,10 @@ class SpeechStreamingBackend: self._billed_before += await link.relay(self._outbox, self._billed_before) case _TurnDiscarded(): await self._outbox.put( - VertexSpeechStreamingTurnDiscarded(billed_seconds=self._billed_before).model_dump_json() + _TurnDiscardedEvent( + turn=link.turn, + event=VertexSpeechStreamingTurnDiscarded(billed_seconds=self._billed_before).model_dump_json(), + ) ) case _: assert_never(link) @@ -396,10 +408,11 @@ class SpeechStreamingBackend: await self._link(_TURN_FINISHED_EVENT) async def _discard_turn(self) -> None: - turn: Final = self._turn + streams: Final = self._turn + turn: Final = self._turn_index self._turn = () - self._discarded_turns |= {self._turn_index} + self._discarded_turns |= {turn} self._turn_index += 1 - for stream in turn: + for stream in streams: stream.cancel() - await self._link(_TurnDiscarded()) + await self._link(_TurnDiscarded(turn=turn)) diff --git a/tests/test_litellm/llms/vertex_ai/audio_transcription/test_vertex_ai_realtime_backend.py b/tests/test_litellm/llms/vertex_ai/audio_transcription/test_vertex_ai_realtime_backend.py index 15601c5ca6c..d5e88706e23 100644 --- a/tests/test_litellm/llms/vertex_ai/audio_transcription/test_vertex_ai_realtime_backend.py +++ b/tests/test_litellm/llms/vertex_ai/audio_transcription/test_vertex_ai_realtime_backend.py @@ -260,6 +260,7 @@ async def test_discard_turn_drops_its_queued_results_and_keeps_google_billed_sec await _until(lambda: len(client.streams[0]) == 3) await backend.send(DISCARD_TURN) assert await _recv(backend) == {"kind": "turn_discarded", "billed_seconds": 2.0} + assert backend._discarded_turns == frozenset() await backend.send(b"\x03\x03") fresh = await _recv(backend) assert fresh["results"] == [{"transcript": "fresh", "is_final": True}] From f2138555586f6a5316eb69d35c3436c4f1c98d43 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 18 Sep 2026 17:44:32 -0700 Subject: [PATCH 105/144] refactor: drop the docstrings from the websocket relay and its tests --- litellm/responses/streaming_iterator.py | 1 - tests/test_litellm/litellm_core_utils/test_litellm_logging.py | 3 --- .../responses/test_responses_websocket_all_providers.py | 1 - 3 files changed, 5 deletions(-) diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 33d09efadf3..32a36ffe4e8 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -2329,7 +2329,6 @@ class ResponsesWebSocketStreaming: verbose_logger.debug("Responses WS client_to_backend ended: %s", e) async def bidirectional_forward(self) -> Exception | None: - """Run both forwarding directions concurrently and return the provider failure that ended the connection.""" forward_task: Final = asyncio.create_task(self.backend_to_client()) try: await self.client_to_backend() diff --git a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py index 0d2d600a7fc..91c334692ee 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -1070,9 +1070,6 @@ async def test_arealtime_marks_litellm_params_async(monkeypatch): @pytest.mark.asyncio async def test_aresponses_websocket_hands_back_the_provider_failure_without_a_success_log(monkeypatch): - """A native Responses WebSocket connection the provider rejected comes back from the ``@client`` - wrapper as the mapped failure, and the wrapper books no success for it: the relay's own dispatch - is the connection's single log, so the proxy can record the connection as a failed request.""" from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER from litellm.responses.main import base_llm_http_handler diff --git a/tests/test_litellm/responses/test_responses_websocket_all_providers.py b/tests/test_litellm/responses/test_responses_websocket_all_providers.py index 6bd137be788..b6d4d9e93a6 100644 --- a/tests/test_litellm/responses/test_responses_websocket_all_providers.py +++ b/tests/test_litellm/responses/test_responses_websocket_all_providers.py @@ -2973,7 +2973,6 @@ def _wrapped_reasoning_item(): class TestNativeWebSocketEncryptedContentAffinity: - """The native relay must restore and wrap ids the same way the HTTP /v1/responses path does.""" @pytest.mark.asyncio @pytest.mark.parametrize("nested", [False, True]) From d74e1bb4453b65b50e804b5b2e71ba3114619425 Mon Sep 17 00:00:00 2001 From: kerry Date: Sat, 19 Sep 2026 00:48:17 +0000 Subject: [PATCH 106/144] fix(timing): union provider timing windows and anchor detailed pre-processing at receive time Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../llm_response_utils/response_metadata.py | 51 ++++++++++---- litellm/litellm_core_utils/logging_utils.py | 17 +++-- .../test_response_metadata.py | 69 +++++++++++++++++-- .../litellm_core_utils/test_logging_utils.py | 7 +- .../test_router_retry_non_retryable_errors.py | 26 +++++-- 5 files changed, 142 insertions(+), 28 deletions(-) diff --git a/litellm/litellm_core_utils/llm_response_utils/response_metadata.py b/litellm/litellm_core_utils/llm_response_utils/response_metadata.py index 780691b4696..cc0d10ee7a6 100644 --- a/litellm/litellm_core_utils/llm_response_utils/response_metadata.py +++ b/litellm/litellm_core_utils/llm_response_utils/response_metadata.py @@ -1,5 +1,6 @@ import datetime from collections.abc import Mapping +from functools import reduce from typing import Any, Final import httpx @@ -25,6 +26,34 @@ def _timing_window_start( return start_time, False +def _union_duration_ms(windows: object, lower: float, upper: float) -> float | None: + if not isinstance(windows, (list, tuple)): + return None + clipped: Final[tuple[tuple[float, float], ...]] = tuple( + (max(lower, float(window[0])), min(upper, float(window[1]))) + for window in windows + if isinstance(window, (list, tuple)) + and len(window) == 2 + and isinstance(window[0], (int, float)) + and isinstance(window[1], (int, float)) + and max(lower, float(window[0])) < min(upper, float(window[1])) + ) + if not clipped: + return None + + ordered: Final[tuple[tuple[float, float], ...]] = tuple(sorted(clipped)) + + def merge_window( + merged: tuple[tuple[float, float], ...], current: tuple[float, float] + ) -> tuple[tuple[float, float], ...]: + if not merged or current[0] > merged[-1][1]: + return (*merged, current) + return (*merged[:-1], (merged[-1][0], max(merged[-1][1], current[1]))) + + merged: Final[tuple[tuple[float, float], ...]] = reduce(merge_window, ordered, ()) + return sum(end - start for start, end in merged) * 1000 + + def response_timing_metrics( start_time: datetime.datetime, end_time: datetime.datetime, @@ -54,20 +83,17 @@ def response_timing_metrics( if cache_duration_ms is not None: overhead_ms: float | None = total_response_time_ms - cache_duration_ms elif llm_api_duration_ms is not None: - total_provider_duration_ms: Final = metadata.get("llm_api_duration_ms_total") - provider_duration_ms: Final = ( - total_provider_duration_ms + provider_duration_ms: Final[float | None] = ( + _union_duration_ms( + metadata.get("llm_api_timing_windows"), + window_start.timestamp(), + end_time.timestamp(), + ) if receive_anchored - and isinstance(total_provider_duration_ms, float) - and isinstance(llm_api_duration_ms, (int, float)) - and total_provider_duration_ms >= llm_api_duration_ms - else llm_api_duration_ms - ) - overhead_ms = ( - round(total_response_time_ms - provider_duration_ms, 4) - if isinstance(provider_duration_ms, (int, float)) else None ) + effective: Final = provider_duration_ms if provider_duration_ms is not None else llm_api_duration_ms + overhead_ms = round(total_response_time_ms - effective, 4) if isinstance(effective, (int, float)) else None else: overhead_ms = None if overhead_ms is None: @@ -178,7 +204,8 @@ class ResponseMetadata: # pre-processing = time from request start to LLM API call start api_call_start: Final[datetime.datetime | None] = logging_obj.model_call_details.get("api_call_start_time") if api_call_start is not None and start_time is not None: - pre_ms: Final = (api_call_start - start_time).total_seconds() * 1000 + anchor: Final = _timing_window_start(start_time, logging_obj)[0] + pre_ms: Final = (api_call_start - anchor).total_seconds() * 1000 detailed["timing_pre_processing_ms"] = round(pre_ms, 4) # post-processing = total - pre - llm_api diff --git a/litellm/litellm_core_utils/logging_utils.py b/litellm/litellm_core_utils/logging_utils.py index 91cc13c8315..5be9dd7be2f 100644 --- a/litellm/litellm_core_utils/logging_utils.py +++ b/litellm/litellm_core_utils/logging_utils.py @@ -288,10 +288,19 @@ def _set_duration_in_model_call_details( if logging_obj and hasattr(logging_obj, "model_call_details"): logging_obj.model_call_details["llm_api_duration_ms"] = duration_ms metadata: Final[dict[str, object]] = get_litellm_metadata_from_kwargs(logging_obj.model_call_details) - existing_total: Final = metadata.get("llm_api_duration_ms_total") - metadata["llm_api_duration_ms_total"] = ( - existing_total if isinstance(existing_total, float) else 0.0 - ) + duration_ms + recorded: Final = metadata.get("llm_api_timing_windows") + earlier: Final[tuple[tuple[float, float], ...]] = tuple( + (float(window[0]), float(window[1])) + for window in (recorded if isinstance(recorded, (list, tuple)) else ()) + if isinstance(window, (list, tuple)) + and len(window) == 2 + and isinstance(window[0], (int, float)) + and isinstance(window[1], (int, float)) + ) + metadata["llm_api_timing_windows"] = ( + *earlier, + (start_time.timestamp(), end_time.timestamp()), + ) else: verbose_logger.debug("`logging_obj` not found - unable to track `llm_api_duration_ms") except Exception as e: diff --git a/tests/test_litellm/litellm_core_utils/llm_response_utils/test_response_metadata.py b/tests/test_litellm/litellm_core_utils/llm_response_utils/test_response_metadata.py index 832be1a12d9..97c154db783 100644 --- a/tests/test_litellm/litellm_core_utils/llm_response_utils/test_response_metadata.py +++ b/tests/test_litellm/litellm_core_utils/llm_response_utils/test_response_metadata.py @@ -16,6 +16,7 @@ import litellm.proxy.common_request_processing as common_request_processing_mod from litellm.litellm_core_utils.litellm_logging import Logging from litellm.litellm_core_utils.llm_response_utils.response_metadata import ( ResponseMetadata, + _union_duration_ms, response_timing_metrics, update_response_metadata, ) @@ -234,7 +235,7 @@ class TestResponseTimingMetrics: def _make_logging_obj( self, llm_api_duration_ms: float | None = None, - llm_api_duration_ms_total: float | None = None, + llm_api_timing_windows: object = None, caching_details: dict[str, object] | None = None, received_at: datetime.datetime | str | None = None, ) -> MagicMock: @@ -242,12 +243,12 @@ class TestResponseTimingMetrics: logging_obj.model_call_details = {} if llm_api_duration_ms is not None: logging_obj.model_call_details["llm_api_duration_ms"] = llm_api_duration_ms - if received_at is not None or llm_api_duration_ms_total is not None: + if received_at is not None or llm_api_timing_windows is not None: metadata = {} if received_at is not None: metadata["litellm_received_at"] = received_at - if llm_api_duration_ms_total is not None: - metadata["llm_api_duration_ms_total"] = llm_api_duration_ms_total + if llm_api_timing_windows is not None: + metadata["llm_api_timing_windows"] = llm_api_timing_windows logging_obj.model_call_details["litellm_params"] = {"metadata": metadata} logging_obj.caching_details = caching_details return logging_obj @@ -271,7 +272,10 @@ class TestResponseTimingMetrics: def test_receive_anchored_window_subtracts_all_provider_attempts(self): logging_obj = self._make_logging_obj( llm_api_duration_ms=300.0, - llm_api_duration_ms_total=700.0, + llm_api_timing_windows=( + (self.START.timestamp(), self.START.timestamp() + 0.3), + (self.START.timestamp() + 0.4, self.START.timestamp() + 0.8), + ), received_at=self.START, ) @@ -283,7 +287,7 @@ class TestResponseTimingMetrics: def test_sdk_window_subtracts_current_provider_attempt(self): logging_obj = self._make_logging_obj( llm_api_duration_ms=300.0, - llm_api_duration_ms_total=700.0, + llm_api_timing_windows=((self.START.timestamp(), self.START.timestamp() + 0.3),), ) result = response_timing_metrics(self.START, self.END, logging_obj) @@ -291,6 +295,38 @@ class TestResponseTimingMetrics: assert result["_response_ms"] == pytest.approx(1000.0) assert result["litellm_overhead_time_ms"] == pytest.approx(700.0) + def test_receive_anchored_window_unions_nested_and_retry_windows(self): + windows = ( + (self.START.timestamp(), self.START.timestamp() + 0.3), + (self.START.timestamp(), self.START.timestamp() + 0.3), + (self.START.timestamp() + 0.4, self.START.timestamp() + 0.7), + ) + logging_obj = self._make_logging_obj( + llm_api_duration_ms=300.0, + llm_api_timing_windows=windows, + received_at=self.START, + ) + + result = response_timing_metrics(self.START, self.END, logging_obj) + + assert result["litellm_overhead_time_ms"] == pytest.approx(400.0) + assert _union_duration_ms(windows, self.START.timestamp(), self.END.timestamp()) == pytest.approx(600.0) + + def test_receive_anchored_window_ignores_seeded_windows_outside_window(self): + windows = ( + (self.START.timestamp() - 10.0, self.START.timestamp() - 1.0), + (self.END.timestamp() + 1.0, self.END.timestamp() + 2.0), + ) + logging_obj = self._make_logging_obj( + llm_api_duration_ms=300.0, + llm_api_timing_windows=windows, + received_at=self.START, + ) + + result = response_timing_metrics(self.START, self.END, logging_obj) + + assert result["litellm_overhead_time_ms"] == pytest.approx(700.0) + def test_receive_anchored_window_falls_back_to_current_provider_attempt(self): logging_obj = self._make_logging_obj( llm_api_duration_ms=300.0, @@ -434,6 +470,27 @@ class TestDetailedTiming: assert hidden.get("timing_pre_processing_ms") == 20.0 assert hidden.get("timing_post_processing_ms") == 10.0 # 530 - 20 - 500 + def test_detailed_timing_pre_processing_uses_receive_anchor(self, monkeypatch): + monkeypatch.setattr(response_metadata_mod, "LITELLM_DETAILED_TIMING", True) + + result = ModelResponse() + start = datetime.datetime(2025, 1, 1, 0, 0, 0) + received_at = start - datetime.timedelta(milliseconds=200) + end = start + datetime.timedelta(milliseconds=530) + logging_obj = self._make_logging_obj( + llm_api_duration_ms=500.0, + api_call_start_time=start, + ) + logging_obj.model_call_details["litellm_params"] = {"metadata": {"litellm_received_at": received_at}} + + metadata = ResponseMetadata(result) + metadata.set_timing_metrics(start, end, logging_obj) + metadata.apply() + + hidden = result._hidden_params + assert hidden.get("timing_pre_processing_ms") == pytest.approx(200.0) + assert hidden.get("timing_post_processing_ms") == pytest.approx(30.0) + def test_detailed_timing_absent_when_disabled(self, monkeypatch): """When LITELLM_DETAILED_TIMING is false, no detailed timing keys.""" monkeypatch.setattr(response_metadata_mod, "LITELLM_DETAILED_TIMING", False) diff --git a/tests/test_litellm/litellm_core_utils/test_logging_utils.py b/tests/test_litellm/litellm_core_utils/test_logging_utils.py index f669ff86c13..672595b85d6 100644 --- a/tests/test_litellm/litellm_core_utils/test_logging_utils.py +++ b/tests/test_litellm/litellm_core_utils/test_logging_utils.py @@ -19,7 +19,7 @@ from litellm.litellm_core_utils.logging_utils import ( class TestSetDurationInModelCallDetails: - def test_accumulates_provider_attempts_in_shared_metadata(self): + def test_records_provider_attempt_windows_in_shared_metadata(self): metadata = {"request_id": "test"} logging_obj = MagicMock() logging_obj.model_call_details = {"litellm_params": {"metadata": metadata}} @@ -31,7 +31,10 @@ class TestSetDurationInModelCallDetails: _set_duration_in_model_call_details(logging_obj, first_start, first_end) _set_duration_in_model_call_details(logging_obj, second_start, second_end) - assert metadata["llm_api_duration_ms_total"] == pytest.approx(1000.0) + assert metadata["llm_api_timing_windows"] == ( + (first_start.timestamp(), first_end.timestamp()), + (second_start.timestamp(), second_end.timestamp()), + ) assert logging_obj.model_call_details["llm_api_duration_ms"] == pytest.approx(700.0) diff --git a/tests/test_litellm/test_router_retry_non_retryable_errors.py b/tests/test_litellm/test_router_retry_non_retryable_errors.py index c797f0f96a6..98a7db7a079 100644 --- a/tests/test_litellm/test_router_retry_non_retryable_errors.py +++ b/tests/test_litellm/test_router_retry_non_retryable_errors.py @@ -13,7 +13,7 @@ Regression tests for https://github.com/BerriAI/litellm/issues/21343 import asyncio import datetime from collections.abc import Awaitable, Callable -from typing import Final, cast +from typing import Final from unittest.mock import AsyncMock, patch import pytest @@ -21,6 +21,10 @@ import pytest import litellm from litellm import Router from litellm.litellm_core_utils.litellm_logging import Logging +from litellm.litellm_core_utils.llm_response_utils.response_metadata import ( + _union_duration_ms, + response_timing_metrics, +) from litellm.litellm_core_utils.logging_utils import track_llm_api_timing from litellm.litellm_core_utils.rules import Rules from litellm.utils import function_setup @@ -286,7 +290,11 @@ async def test_not_found_error_in_retry_loop_raises_immediately(): @pytest.mark.asyncio async def test_retry_attempts_accumulate_timing_in_shared_request_metadata(): - metadata: dict[str, object] = {"model_group": "test-model"} + received_at: Final = datetime.datetime.now() + metadata: dict[str, object] = { + "model_group": "test-model", + "litellm_received_at": received_at, + } logging_obj_raw, _ = function_setup( "acompletion", Rules(), @@ -297,7 +305,8 @@ async def test_retry_attempts_accumulate_timing_in_shared_request_metadata(): litellm_call_id="retry-timing-test", is_async_call=True, ) - logging_obj: Final[Logging] = cast(Logging, logging_obj_raw) + assert isinstance(logging_obj_raw, Logging) + logging_obj: Final[Logging] = logging_obj_raw attempt_numbers: list[int] = [] metadata_ids: list[int] = [] @@ -334,8 +343,17 @@ async def test_retry_attempts_accumulate_timing_in_shared_request_metadata(): ) request_metadata: Final = logging_obj.model_call_details["litellm_params"]["metadata"] + windows: Final = request_metadata["llm_api_timing_windows"] + end_time: Final = datetime.datetime.fromtimestamp(max(window[1] for window in windows)) + timing_metrics: Final = response_timing_metrics(received_at, end_time, logging_obj) assert result == "success" assert attempt_numbers == [1, 2] assert request_metadata is metadata assert metadata_ids == [id(metadata), id(metadata)] - assert request_metadata["llm_api_duration_ms_total"] > logging_obj.model_call_details["llm_api_duration_ms"] + assert len(windows) == 2 + union_duration_ms: Final = _union_duration_ms(windows, received_at.timestamp(), end_time.timestamp()) + assert union_duration_ms is not None + total_response_time_ms: Final = (end_time.timestamp() - received_at.timestamp()) * 1000 + assert timing_metrics["litellm_overhead_time_ms"] == pytest.approx( + round(total_response_time_ms - union_duration_ms, 4) + ) From 6b0ad3bed3f933311f59f9743b76182774328d45 Mon Sep 17 00:00:00 2001 From: kerry Date: Sat, 19 Sep 2026 00:49:34 +0000 Subject: [PATCH 107/144] feat(xai): add speech-to-text via /v1/audio/transcriptions Route xai audio transcription through a provider config hitting POST https://api.x.ai/v1/stt instead of the openai-compatible chat handler which targets /audio/transcriptions. Supports language, diarize, keyterm, filler_words and other provider fields as passthrough kwargs Resolves LIT-8153 Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/constants.py | 5 + .../llms/xai/audio_transcription/__init__.py | 3 + .../xai/audio_transcription/transformation.py | 187 ++++++++++++++++++ litellm/main.py | 6 +- ...odel_prices_and_context_window_backup.json | 28 +++ litellm/utils.py | 6 + model_prices_and_context_window.json | 28 +++ ..._xai_audio_transcription_transformation.py | 172 ++++++++++++++++ 8 files changed, 434 insertions(+), 1 deletion(-) create mode 100644 litellm/llms/xai/audio_transcription/__init__.py create mode 100644 litellm/llms/xai/audio_transcription/transformation.py create mode 100644 tests/test_litellm/llms/xai/test_xai_audio_transcription_transformation.py diff --git a/litellm/constants.py b/litellm/constants.py index a7d4eba0f15..624828d6eb0 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -998,6 +998,11 @@ openai_compatible_providers: Final[list] = [ "cognition", "scx-ai", ] + +# Providers that are openai-compatible for chat but have their own audio +# transcription endpoint, so litellm.transcription must route them through +# their provider config instead of the OpenAI SDK handler. +OPENAI_COMPATIBLE_PROVIDERS_WITH_NATIVE_AUDIO_TRANSCRIPTION: Final = frozenset({"xai"}) openai_text_completion_compatible_providers: Final[list] = [ # providers that support `/v1/completions` "together_ai", "fireworks_ai", diff --git a/litellm/llms/xai/audio_transcription/__init__.py b/litellm/llms/xai/audio_transcription/__init__.py new file mode 100644 index 00000000000..c7910cf1f6b --- /dev/null +++ b/litellm/llms/xai/audio_transcription/__init__.py @@ -0,0 +1,3 @@ +from .transformation import XAIAudioTranscriptionConfig + +__all__ = ["XAIAudioTranscriptionConfig"] diff --git a/litellm/llms/xai/audio_transcription/transformation.py b/litellm/llms/xai/audio_transcription/transformation.py new file mode 100644 index 00000000000..8f977dd99fa --- /dev/null +++ b/litellm/llms/xai/audio_transcription/transformation.py @@ -0,0 +1,187 @@ +""" +Translates from OpenAI's `/v1/audio/transcriptions` to xAI's `/v1/stt` +""" + +from collections.abc import Iterable, Mapping +from typing import Final, cast + +from httpx import Headers, Response +from pydantic import BaseModel, ConfigDict + +import litellm +from litellm.litellm_core_utils.audio_utils.utils import process_audio_file +from litellm.llms.base_llm.chat.transformation import BaseLLMException +from litellm.types.llms.openai import ( + AllMessageValues, + OpenAIAudioTranscriptionOptionalParams, +) +from litellm.types.utils import FileTypes, TranscriptionResponse + +from ...base_llm.audio_transcription.transformation import ( + AudioTranscriptionRequestData, + BaseAudioTranscriptionConfig, +) +from ..common_utils import XAIModelInfo + + +class XAIAudioTranscriptionError(BaseLLMException): + pass + + +class _XAISttWord(BaseModel): + model_config = ConfigDict(extra="allow") + text: str = "" + start: float = 0.0 + end: float = 0.0 + speaker: str | None = None + + +class _XAISttResponse(BaseModel): + model_config = ConfigDict(extra="allow") + text: str = "" + language: str = "unknown" + duration: float | None = None + words: list[_XAISttWord] | None = None + + +def _serialize_form_value(value: object) -> str | list[str]: + if isinstance(value, bool): + return "true" if value else "false" + if isinstance(value, (list, tuple)): + return [str(item) for item in cast(Iterable[object], value)] + return str(value) + + +class XAIAudioTranscriptionConfig(BaseAudioTranscriptionConfig): + @property + def custom_llm_provider(self) -> str: + return litellm.LlmProviders.XAI.value + + def get_supported_openai_params(self, model: str) -> list[OpenAIAudioTranscriptionOptionalParams]: + return ["language"] + + def map_openai_params( + self, + non_default_params: dict[str, object], + optional_params: dict[str, object], + model: str, + drop_params: bool, + ) -> dict[str, object]: + supported_params: Final = self.get_supported_openai_params(model) + for k, v in non_default_params.items(): + if k in supported_params: + optional_params[k] = v + return optional_params + + def get_error_class( + self, error_message: str, status_code: int, headers: dict[str, object] | Headers + ) -> BaseLLMException: + return XAIAudioTranscriptionError(message=error_message, status_code=status_code, headers=headers) + + def transform_audio_transcription_request( + self, + model: str, + audio_file: FileTypes, + optional_params: dict[str, object], + litellm_params: dict[str, object], + ) -> AudioTranscriptionRequestData: + processed_audio: Final = process_audio_file(audio_file) + + # Provider kwargs land in `extra_body` for openai_compatible_providers + extra_body: Final = optional_params.get("extra_body") + flat_params: Final[dict[str, object]] = { + **(dict(cast(Mapping[str, object], extra_body)) if isinstance(extra_body, Mapping) else {}), + **{k: v for k, v in optional_params.items() if k != "extra_body"}, + } + + openai_params: Final = self.get_supported_openai_params(model) + excluded_params: Final = frozenset({"model", "OPENAI_TRANSCRIPTION_PARAMS", *openai_params}) + provider_specific_params: Final[dict[str, object]] = { + k: v for k, v in flat_params.items() if v is not None and k not in excluded_params + } + + form_data: Final[dict[str, str | list[str]]] = {"model": model} + for key, value in provider_specific_params.items(): + form_data[key] = _serialize_form_value(value) + for key in openai_params: + value = flat_params.get(key) + if value is not None: + form_data[key] = _serialize_form_value(value) + + files: Final = { + "file": ( + processed_audio.filename, + processed_audio.file_content, + processed_audio.content_type, + ) + } + + return AudioTranscriptionRequestData(data=form_data, files=files) + + def transform_audio_transcription_response( + self, + raw_response: Response, + ) -> TranscriptionResponse: + try: + payload: Final = _XAISttResponse.model_validate_json(raw_response.content) + except Exception as e: + raise XAIAudioTranscriptionError( + message=f"Error parsing xAI response: {e}", + status_code=raw_response.status_code, + headers=dict(raw_response.headers), + ) + + response: Final = TranscriptionResponse(text=payload.text) + response["task"] = "transcribe" + response["language"] = payload.language + + if payload.duration is not None: + response["duration"] = payload.duration + + if payload.words is not None: + response["words"] = [ + { + "word": word.text, + "start": word.start, + "end": word.end, + **({"speaker": word.speaker} if word.speaker is not None else {}), + } + for word in payload.words + ] + + hidden_params: Final[dict[str, object]] = dict(payload.model_dump(mode="json")) + if payload.duration is not None: + hidden_params["audio_transcription_duration"] = payload.duration + response._hidden_params = hidden_params # pyright: ignore[reportPrivateUsage] # TranscriptionResponse exposes no public hidden-params setter + + return response + + def get_complete_url( + self, + api_base: str | None, + api_key: str | None, + model: str, + optional_params: dict[str, object], + litellm_params: dict[str, object], + stream: bool | None = None, + ) -> str: + base: Final = (XAIModelInfo.get_api_base(api_base) or "").rstrip("/") + normalized: Final = base.removesuffix("/v1") + return f"{normalized}/v1/stt" + + def validate_environment( + self, + headers: dict[str, object], + model: str, + messages: list[AllMessageValues], + optional_params: dict[str, object], + litellm_params: dict[str, object], + api_key: str | None = None, + api_base: str | None = None, + ) -> dict[str, object]: + resolved_key: Final = XAIModelInfo.get_api_key(api_key) + if resolved_key is None: + raise ValueError("xAI API key is required. Set XAI_API_KEY environment variable.") + + headers["Authorization"] = f"Bearer {resolved_key}" + return headers diff --git a/litellm/main.py b/litellm/main.py index ac8fa507728..1c3de8e766b 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -64,6 +64,7 @@ from litellm.constants import ( AZURE_OPENAI_AUDIO_PROVIDERS, DEFAULT_MOCK_RESPONSE_COMPLETION_TOKEN_COUNT, DEFAULT_MOCK_RESPONSE_PROMPT_TOKEN_COUNT, + OPENAI_COMPATIBLE_PROVIDERS_WITH_NATIVE_AUDIO_TRANSCRIPTION, ) from litellm.exceptions import LiteLLMUnknownProvider from litellm.integrations.custom_logger import CustomLogger @@ -7859,7 +7860,10 @@ def transcription( litellm_params=litellm_params_dict, custom_llm_provider=custom_llm_provider, ) - elif custom_llm_provider == "openai" or (custom_llm_provider in litellm.openai_compatible_providers): + elif custom_llm_provider == "openai" or ( + custom_llm_provider in litellm.openai_compatible_providers + and custom_llm_provider not in OPENAI_COMPATIBLE_PROVIDERS_WITH_NATIVE_AUDIO_TRANSCRIPTION + ): api_base = ( api_base or litellm.api_base diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 30b08e54410..f8dc79e6c04 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -63375,6 +63375,34 @@ "video" ] }, + "xai/grok-voice-transcribe-1.0": { + "input_cost_per_second": 2.778e-05, + "litellm_provider": "xai", + "metadata": { + "calculation": "$0.10/3600 seconds = $0.00002778 per second", + "original_pricing_per_hour": 0.1 + }, + "mode": "audio_transcription", + "output_cost_per_second": 0.0, + "source": "https://docs.x.ai/developers/pricing", + "supported_endpoints": [ + "/v1/audio/transcriptions" + ] + }, + "xai/grok-voice-transcribe-2.0": { + "input_cost_per_second": 2.778e-05, + "litellm_provider": "xai", + "metadata": { + "calculation": "$0.10/3600 seconds = $0.00002778 per second", + "original_pricing_per_hour": 0.1 + }, + "mode": "audio_transcription", + "output_cost_per_second": 0.0, + "source": "https://docs.x.ai/developers/pricing", + "supported_endpoints": [ + "/v1/audio/transcriptions" + ] + }, "low/1024-x-1024/grok-imagine-image-2.0": { "input_cost_per_image": 0.04, "litellm_provider": "xai", diff --git a/litellm/utils.py b/litellm/utils.py index f2315651a53..97c074ce112 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -8730,6 +8730,12 @@ class ProviderConfigManager: ) return ElevenLabsAudioTranscriptionConfig() + elif litellm.LlmProviders.XAI == provider: + from litellm.llms.xai.audio_transcription.transformation import ( + XAIAudioTranscriptionConfig, + ) + + return XAIAudioTranscriptionConfig() elif litellm.LlmProviders.OPENAI == provider: if "gpt-4o" in model: return litellm.OpenAIGPTAudioTranscriptionConfig() diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 30b08e54410..f8dc79e6c04 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -63375,6 +63375,34 @@ "video" ] }, + "xai/grok-voice-transcribe-1.0": { + "input_cost_per_second": 2.778e-05, + "litellm_provider": "xai", + "metadata": { + "calculation": "$0.10/3600 seconds = $0.00002778 per second", + "original_pricing_per_hour": 0.1 + }, + "mode": "audio_transcription", + "output_cost_per_second": 0.0, + "source": "https://docs.x.ai/developers/pricing", + "supported_endpoints": [ + "/v1/audio/transcriptions" + ] + }, + "xai/grok-voice-transcribe-2.0": { + "input_cost_per_second": 2.778e-05, + "litellm_provider": "xai", + "metadata": { + "calculation": "$0.10/3600 seconds = $0.00002778 per second", + "original_pricing_per_hour": 0.1 + }, + "mode": "audio_transcription", + "output_cost_per_second": 0.0, + "source": "https://docs.x.ai/developers/pricing", + "supported_endpoints": [ + "/v1/audio/transcriptions" + ] + }, "low/1024-x-1024/grok-imagine-image-2.0": { "input_cost_per_image": 0.04, "litellm_provider": "xai", diff --git a/tests/test_litellm/llms/xai/test_xai_audio_transcription_transformation.py b/tests/test_litellm/llms/xai/test_xai_audio_transcription_transformation.py new file mode 100644 index 00000000000..0fce47050b5 --- /dev/null +++ b/tests/test_litellm/llms/xai/test_xai_audio_transcription_transformation.py @@ -0,0 +1,172 @@ +import httpx +import pytest + +import litellm +from litellm.llms.base_llm.audio_transcription.transformation import ( + AudioTranscriptionRequestData, +) +from litellm.llms.custom_httpx.http_handler import HTTPHandler +from litellm.llms.xai.audio_transcription.transformation import ( + XAIAudioTranscriptionConfig, +) +from litellm.types.utils import LlmProviders +from litellm.utils import ProviderConfigManager + +CONFIG = XAIAudioTranscriptionConfig() + +WAV_BYTES = b"RIFF" + b"\x00" * 64 + + +def test_transform_request_serializes_provider_params(): + result = CONFIG.transform_audio_transcription_request( + model="grok-voice-transcribe-2.0", + audio_file=WAV_BYTES, + optional_params={ + "language": "en", + "diarize": True, + "keyterm": ["LiteLLM", "Grok"], + }, + litellm_params={}, + ) + + assert isinstance(result, AudioTranscriptionRequestData) + data = result.data + assert data["model"] == "grok-voice-transcribe-2.0" + assert data["language"] == "en" + assert data["diarize"] == "true" + assert data["keyterm"] == ["LiteLLM", "Grok"] + filename, content, content_type = result.files["file"] + assert content == WAV_BYTES + assert isinstance(filename, str) + assert isinstance(content_type, str) + + +def test_transform_request_flattens_extra_body(): + result = CONFIG.transform_audio_transcription_request( + model="grok-voice-transcribe-1.0", + audio_file=WAV_BYTES, + optional_params={ + "language": "en", + "extra_body": {"diarize": False, "channels": 2}, + }, + litellm_params={}, + ) + assert result.data["diarize"] == "false" + assert result.data["channels"] == "2" + assert "extra_body" not in result.data + + +@pytest.mark.parametrize( + "api_base,expected", + [ + (None, "https://api.x.ai/v1/stt"), + ("https://api.x.ai/v1", "https://api.x.ai/v1/stt"), + ("https://api.x.ai/v1/", "https://api.x.ai/v1/stt"), + ("https://proxy.example/", "https://proxy.example/v1/stt"), + ], +) +def test_get_complete_url(api_base, expected): + url = CONFIG.get_complete_url( + api_base=api_base, + api_key=None, + model="grok-voice-transcribe-2.0", + optional_params={}, + litellm_params={}, + ) + assert url == expected + + +def test_validate_environment_sets_bearer_header(): + headers = CONFIG.validate_environment( + headers={}, + model="grok-voice-transcribe-2.0", + messages=[], + optional_params={}, + litellm_params={}, + api_key="sk-test", + ) + assert headers["Authorization"] == "Bearer sk-test" + assert "Content-Type" not in headers + + +def test_validate_environment_requires_key(monkeypatch): + monkeypatch.delenv("XAI_API_KEY", raising=False) + monkeypatch.setattr(litellm, "xai_key", None) + with pytest.raises(ValueError): + CONFIG.validate_environment( + headers={}, + model="grok-voice-transcribe-2.0", + messages=[], + optional_params={}, + litellm_params={}, + api_key=None, + ) + + +def test_transform_response_maps_xai_shape(): + raw = httpx.Response( + 200, + json={ + "text": "hello world", + "language": "en", + "duration": 3.2, + "words": [ + {"text": "hello", "start": 0.0, "end": 0.5, "speaker": "1"}, + {"text": "world", "start": 0.5, "end": 1.0}, + ], + }, + request=httpx.Request("POST", "https://api.x.ai/v1/stt"), + ) + response = CONFIG.transform_audio_transcription_response(raw_response=raw) + + assert response.text == "hello world" + assert response["language"] == "en" + assert response["duration"] == 3.2 + assert response["task"] == "transcribe" + assert response["words"] == [ + {"word": "hello", "start": 0.0, "end": 0.5, "speaker": "1"}, + {"word": "world", "start": 0.5, "end": 1.0}, + ] + assert response._hidden_params["audio_transcription_duration"] == 3.2 + + +def test_transcription_routes_to_xai_stt(monkeypatch): + monkeypatch.delenv("XAI_API_KEY", raising=False) + monkeypatch.setattr(litellm, "xai_key", None) + captured: dict = {} + + def handler(request: httpx.Request) -> httpx.Response: + captured["request"] = request + return httpx.Response( + 200, + json={"text": "transcribed text", "language": "en", "duration": 1.5}, + request=request, + ) + + http_handler = HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(handler))) + response = litellm.transcription( + model="xai/grok-voice-transcribe-2.0", + file=("sample.wav", WAV_BYTES, "audio/wav"), + api_key="sk-test", + diarize=True, + keyterm=["LiteLLM"], + client=http_handler, + ) + + request = captured["request"] + assert str(request.url) == "https://api.x.ai/v1/stt" + assert request.headers["Authorization"] == "Bearer sk-test" + body = request.content.decode("utf-8", errors="replace") + assert 'name="model"' in body and "grok-voice-transcribe-2.0" in body + assert 'name="diarize"' in body and "true" in body + assert 'name="keyterm"' in body and "LiteLLM" in body + assert 'name="file"' in body + assert response.text == "transcribed text" + + +def test_provider_config_manager_returns_xai_config(): + config = ProviderConfigManager.get_provider_audio_transcription_config( + model="grok-voice-transcribe-2.0", + provider=LlmProviders.XAI, + ) + assert isinstance(config, XAIAudioTranscriptionConfig) From c65f11bf0ffc1237de081ceeaa7d9d94c7874667 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 18 Sep 2026 17:50:22 -0700 Subject: [PATCH 108/144] style(websearch): wrap three long lines in the interception handler --- .../integrations/websearch_interception/handler.py | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/litellm/integrations/websearch_interception/handler.py b/litellm/integrations/websearch_interception/handler.py index 29a586eaf20..4558d6c2c04 100644 --- a/litellm/integrations/websearch_interception/handler.py +++ b/litellm/integrations/websearch_interception/handler.py @@ -349,7 +349,9 @@ class WebSearchInterceptionLogger(CustomLogger): "input": {"query": query}, } ) - content.append(WebSearchTransformation.build_web_search_outcome_block(tool_use_id=tool_use_id, outcome=outcome)) + content.append( + WebSearchTransformation.build_web_search_outcome_block(tool_use_id=tool_use_id, outcome=outcome) + ) # Keep the text block so non-native short-circuit callers (Claude Code, # github_copilot, etc.) see the same payload they always have. content.append({"type": "text", "text": search_result_text}) @@ -953,7 +955,9 @@ class WebSearchInterceptionLogger(CustomLogger): isinstance(outcome, SearchFailed) for outcome in search_outcomes ) if every_search_failed: - return AgenticLoopPlan(run_agentic_loop=False, terminate=True, stop_reason="web_search_failed", metadata=metadata) + return AgenticLoopPlan( + run_agentic_loop=False, terminate=True, stop_reason="web_search_failed", metadata=metadata + ) return AgenticLoopPlan(run_agentic_loop=True, request_patch=request_patch, metadata=metadata) async def async_post_agentic_loop_response_hook( @@ -1424,7 +1428,9 @@ class WebSearchInterceptionLogger(CustomLogger): async def _short_circuit_search_outcome(self, query: str, kwargs: Mapping[str, object] | None) -> SearchOutcome: try: result: Final = ( - await self._execute_search(query) if kwargs is None else await self._execute_search(query, kwargs=kwargs) + await self._execute_search(query) + if kwargs is None + else await self._execute_search(query, kwargs=kwargs) ) except Exception as e: return WebSearchTransformation.search_outcome(e) From 6f54ad5166ab55358a91bef2ad96cce3a4efba9b Mon Sep 17 00:00:00 2001 From: kerry Date: Sat, 19 Sep 2026 00:53:02 +0000 Subject: [PATCH 109/144] fix(xai): parse integer speaker ids and simplify stt form build Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/constants.py | 4 +- .../xai/audio_transcription/transformation.py | 90 ++++++++++--------- ..._xai_audio_transcription_transformation.py | 4 +- 3 files changed, 49 insertions(+), 49 deletions(-) diff --git a/litellm/constants.py b/litellm/constants.py index 624828d6eb0..56d3f5450d7 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -999,10 +999,8 @@ openai_compatible_providers: Final[list] = [ "scx-ai", ] -# Providers that are openai-compatible for chat but have their own audio -# transcription endpoint, so litellm.transcription must route them through -# their provider config instead of the OpenAI SDK handler. OPENAI_COMPATIBLE_PROVIDERS_WITH_NATIVE_AUDIO_TRANSCRIPTION: Final = frozenset({"xai"}) + openai_text_completion_compatible_providers: Final[list] = [ # providers that support `/v1/completions` "together_ai", "fireworks_ai", diff --git a/litellm/llms/xai/audio_transcription/transformation.py b/litellm/llms/xai/audio_transcription/transformation.py index 8f977dd99fa..b5c5d9c522d 100644 --- a/litellm/llms/xai/audio_transcription/transformation.py +++ b/litellm/llms/xai/audio_transcription/transformation.py @@ -2,11 +2,11 @@ Translates from OpenAI's `/v1/audio/transcriptions` to xAI's `/v1/stt` """ -from collections.abc import Iterable, Mapping -from typing import Final, cast +from collections.abc import Mapping, Sequence +from typing import Final from httpx import Headers, Response -from pydantic import BaseModel, ConfigDict +from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError import litellm from litellm.litellm_core_utils.audio_utils.utils import process_audio_file @@ -33,7 +33,7 @@ class _XAISttWord(BaseModel): text: str = "" start: float = 0.0 end: float = 0.0 - speaker: str | None = None + speaker: int | None = None class _XAISttResponse(BaseModel): @@ -41,14 +41,18 @@ class _XAISttResponse(BaseModel): text: str = "" language: str = "unknown" duration: float | None = None - words: list[_XAISttWord] | None = None + words: tuple[_XAISttWord, ...] | None = None -def _serialize_form_value(value: object) -> str | list[str]: +_OBJECT_TUPLE: Final = TypeAdapter(tuple[object, ...]) +_STRING_OBJECT_DICT: Final = TypeAdapter(dict[str, object]) + + +def _serialize_form_value(value: object) -> str | list[str]: # mutable-ok: httpx multipart data takes list values for repeated form fields if isinstance(value, bool): return "true" if value else "false" if isinstance(value, (list, tuple)): - return [str(item) for item in cast(Iterable[object], value)] + return [str(item) for item in _OBJECT_TUPLE.validate_python(value)] return str(value) @@ -57,24 +61,24 @@ class XAIAudioTranscriptionConfig(BaseAudioTranscriptionConfig): def custom_llm_provider(self) -> str: return litellm.LlmProviders.XAI.value - def get_supported_openai_params(self, model: str) -> list[OpenAIAudioTranscriptionOptionalParams]: + def get_supported_openai_params(self, model: str) -> list[OpenAIAudioTranscriptionOptionalParams]: # mutable-ok: base class signature returns list return ["language"] def map_openai_params( self, - non_default_params: dict[str, object], - optional_params: dict[str, object], + non_default_params: Mapping[str, object], + optional_params: Mapping[str, object], model: str, drop_params: bool, - ) -> dict[str, object]: + ) -> dict[str, object]: # mutable-ok: base class signature returns dict supported_params: Final = self.get_supported_openai_params(model) - for k, v in non_default_params.items(): - if k in supported_params: - optional_params[k] = v - return optional_params + return { + **optional_params, + **{k: v for k, v in non_default_params.items() if k in supported_params}, + } def get_error_class( - self, error_message: str, status_code: int, headers: dict[str, object] | Headers + self, error_message: str, status_code: int, headers: dict[str, object] | Headers # mutable-ok: base class signature takes dict ) -> BaseLLMException: return XAIAudioTranscriptionError(message=error_message, status_code=status_code, headers=headers) @@ -82,32 +86,31 @@ class XAIAudioTranscriptionConfig(BaseAudioTranscriptionConfig): self, model: str, audio_file: FileTypes, - optional_params: dict[str, object], - litellm_params: dict[str, object], + optional_params: Mapping[str, object], + litellm_params: Mapping[str, object], ) -> AudioTranscriptionRequestData: processed_audio: Final = process_audio_file(audio_file) - # Provider kwargs land in `extra_body` for openai_compatible_providers extra_body: Final = optional_params.get("extra_body") - flat_params: Final[dict[str, object]] = { - **(dict(cast(Mapping[str, object], extra_body)) if isinstance(extra_body, Mapping) else {}), + flat_params: Final[Mapping[str, object]] = { + **( + _STRING_OBJECT_DICT.validate_python(extra_body) + if isinstance(extra_body, Mapping) + else {} + ), **{k: v for k, v in optional_params.items() if k != "extra_body"}, } - openai_params: Final = self.get_supported_openai_params(model) - excluded_params: Final = frozenset({"model", "OPENAI_TRANSCRIPTION_PARAMS", *openai_params}) - provider_specific_params: Final[dict[str, object]] = { - k: v for k, v in flat_params.items() if v is not None and k not in excluded_params + excluded_params: Final = frozenset({"model", "OPENAI_TRANSCRIPTION_PARAMS", "extra_body"}) + form_data: Final[dict[str, str | list[str]]] = { # mutable-ok: AudioTranscriptionRequestData.data requires dict and httpx needs list values + "model": model, + **{ + k: _serialize_form_value(v) + for k, v in flat_params.items() + if v is not None and k not in excluded_params + }, } - form_data: Final[dict[str, str | list[str]]] = {"model": model} - for key, value in provider_specific_params.items(): - form_data[key] = _serialize_form_value(value) - for key in openai_params: - value = flat_params.get(key) - if value is not None: - form_data[key] = _serialize_form_value(value) - files: Final = { "file": ( processed_audio.filename, @@ -124,7 +127,7 @@ class XAIAudioTranscriptionConfig(BaseAudioTranscriptionConfig): ) -> TranscriptionResponse: try: payload: Final = _XAISttResponse.model_validate_json(raw_response.content) - except Exception as e: + except ValidationError as e: raise XAIAudioTranscriptionError( message=f"Error parsing xAI response: {e}", status_code=raw_response.status_code, @@ -149,7 +152,7 @@ class XAIAudioTranscriptionConfig(BaseAudioTranscriptionConfig): for word in payload.words ] - hidden_params: Final[dict[str, object]] = dict(payload.model_dump(mode="json")) + hidden_params: Final[dict[str, object]] = dict(payload.model_dump(mode="json")) # mutable-ok: TranscriptionResponse._hidden_params is a dict if payload.duration is not None: hidden_params["audio_transcription_duration"] = payload.duration response._hidden_params = hidden_params # pyright: ignore[reportPrivateUsage] # TranscriptionResponse exposes no public hidden-params setter @@ -161,8 +164,8 @@ class XAIAudioTranscriptionConfig(BaseAudioTranscriptionConfig): api_base: str | None, api_key: str | None, model: str, - optional_params: dict[str, object], - litellm_params: dict[str, object], + optional_params: Mapping[str, object], + litellm_params: Mapping[str, object], stream: bool | None = None, ) -> str: base: Final = (XAIModelInfo.get_api_base(api_base) or "").rstrip("/") @@ -171,17 +174,16 @@ class XAIAudioTranscriptionConfig(BaseAudioTranscriptionConfig): def validate_environment( self, - headers: dict[str, object], + headers: dict[str, object], # mutable-ok: base class signature takes and returns dict model: str, - messages: list[AllMessageValues], - optional_params: dict[str, object], - litellm_params: dict[str, object], + messages: Sequence[AllMessageValues], + optional_params: Mapping[str, object], + litellm_params: Mapping[str, object], api_key: str | None = None, api_base: str | None = None, - ) -> dict[str, object]: + ) -> dict[str, object]: # mutable-ok: base class signature returns dict resolved_key: Final = XAIModelInfo.get_api_key(api_key) if resolved_key is None: raise ValueError("xAI API key is required. Set XAI_API_KEY environment variable.") - headers["Authorization"] = f"Bearer {resolved_key}" - return headers + return {**headers, "Authorization": f"Bearer {resolved_key}"} diff --git a/tests/test_litellm/llms/xai/test_xai_audio_transcription_transformation.py b/tests/test_litellm/llms/xai/test_xai_audio_transcription_transformation.py index 0fce47050b5..f2365f00ba3 100644 --- a/tests/test_litellm/llms/xai/test_xai_audio_transcription_transformation.py +++ b/tests/test_litellm/llms/xai/test_xai_audio_transcription_transformation.py @@ -111,7 +111,7 @@ def test_transform_response_maps_xai_shape(): "language": "en", "duration": 3.2, "words": [ - {"text": "hello", "start": 0.0, "end": 0.5, "speaker": "1"}, + {"text": "hello", "start": 0.0, "end": 0.5, "speaker": 1}, {"text": "world", "start": 0.5, "end": 1.0}, ], }, @@ -124,7 +124,7 @@ def test_transform_response_maps_xai_shape(): assert response["duration"] == 3.2 assert response["task"] == "transcribe" assert response["words"] == [ - {"word": "hello", "start": 0.0, "end": 0.5, "speaker": "1"}, + {"word": "hello", "start": 0.0, "end": 0.5, "speaker": 1}, {"word": "world", "start": 0.5, "end": 1.0}, ] assert response._hidden_params["audio_transcription_duration"] == 3.2 From 8d2476465fb4b6a84110ec5e9a264602c0e6e6d1 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Fri, 18 Sep 2026 17:58:07 -0700 Subject: [PATCH 110/144] fix(rust): honor environment proxies by default and name the cause in transport errors Python's aiohttp transport reads HTTP(S)_PROXY on every request unless disable_aiohttp_trust_env is set, so the Rust clients now do the same instead of requiring aiohttp_trust_env. Transport error messages include reqwest's source chain, so a rejected certificate or refused connection is no longer reported as just 'error sending request' --- litellm-rust/crates/http/src/config.rs | 25 +++++++++--- litellm-rust/crates/http/src/settings.rs | 11 ++++-- .../crates/llms/src/custom_httpx/transport.rs | 38 ++++++++++++++++++- .../crates/python-bridge/python_settings.json | 1 + litellm-rust/crates/python-bridge/src/http.rs | 13 ++++++- litellm/rust_bridge/settings.py | 2 + .../test_litellm/rust_bridge/test_settings.py | 2 + 7 files changed, 79 insertions(+), 13 deletions(-) diff --git a/litellm-rust/crates/http/src/config.rs b/litellm-rust/crates/http/src/config.rs index 7772e6cce5b..24a52315f7c 100644 --- a/litellm-rust/crates/http/src/config.rs +++ b/litellm-rust/crates/http/src/config.rs @@ -55,7 +55,10 @@ impl HttpClientConfig { force_ipv4: settings.force_ipv4, http2: settings.http2, user_agent: settings.user_agent.clone(), - trust_proxy_env: settings.trust_proxy_env || settings.http2 || settings.httpx_transport, + trust_proxy_env: !settings.ignore_proxy_env + || settings.trust_proxy_env + || settings.http2 + || settings.httpx_transport, connect_timeout: settings.connect_timeout, }) } @@ -237,11 +240,21 @@ mod tests { } #[rstest] - #[case::aiohttp_default(HttpSettings::default(), false)] - #[case::aiohttp_trust_env(HttpSettings { trust_proxy_env: true, ..HttpSettings::default() }, true)] - #[case::http2_uses_httpx(HttpSettings { http2: true, ..HttpSettings::default() }, true)] - #[case::aiohttp_disabled(HttpSettings { httpx_transport: true, ..HttpSettings::default() }, true)] - fn environment_proxies_apply_whenever_python_would_use_httpx( + #[case::aiohttp_default(HttpSettings::default(), true)] + #[case::aiohttp_opted_out(HttpSettings { ignore_proxy_env: true, ..HttpSettings::default() }, false)] + #[case::session_trust_env_beats_opt_out( + HttpSettings { ignore_proxy_env: true, trust_proxy_env: true, ..HttpSettings::default() }, + true + )] + #[case::http2_uses_httpx( + HttpSettings { ignore_proxy_env: true, http2: true, ..HttpSettings::default() }, + true + )] + #[case::aiohttp_disabled( + HttpSettings { ignore_proxy_env: true, httpx_transport: true, ..HttpSettings::default() }, + true + )] + fn environment_proxies_apply_unless_the_aiohttp_transport_opts_out( #[case] settings: HttpSettings, #[case] expected: bool, ) { diff --git a/litellm-rust/crates/http/src/settings.rs b/litellm-rust/crates/http/src/settings.rs index 55ac471bba9..c572c56ef3a 100644 --- a/litellm-rust/crates/http/src/settings.rs +++ b/litellm-rust/crates/http/src/settings.rs @@ -32,6 +32,7 @@ pub struct HttpSettings { pub httpx_transport: bool, pub user_agent: Option, pub trust_proxy_env: bool, + pub ignore_proxy_env: bool, pub connect_timeout: Duration, } @@ -48,6 +49,7 @@ impl Default for HttpSettings { httpx_transport: false, user_agent: None, trust_proxy_env: false, + ignore_proxy_env: false, connect_timeout: Duration::from_secs(5), } } @@ -78,6 +80,7 @@ impl HttpSettings { httpx_transport: self.httpx_transport || enabled("DISABLE_AIOHTTP_TRANSPORT"), user_agent: env("LITELLM_USER_AGENT").or(self.user_agent), trust_proxy_env: self.trust_proxy_env || enabled("AIOHTTP_TRUST_ENV"), + ignore_proxy_env: self.ignore_proxy_env || enabled("DISABLE_AIOHTTP_TRUST_ENV"), ..self } } @@ -214,14 +217,16 @@ mod tests { #[case("1", false)] fn boolean_switches_only_turn_on_for_true(#[case] value: &'static str, #[case] expected: bool) { let env = move |name: &str| match name { - "LITELLM_HTTP2" | "AIOHTTP_TRUST_ENV" | "DISABLE_AIOHTTP_TRANSPORT" => { - Some(value.to_string()) - } + "LITELLM_HTTP2" + | "AIOHTTP_TRUST_ENV" + | "DISABLE_AIOHTTP_TRANSPORT" + | "DISABLE_AIOHTTP_TRUST_ENV" => Some(value.to_string()), _ => None, }; let settings = HttpSettings::default().with_environment(&env); assert_eq!(settings.http2, expected); assert_eq!(settings.httpx_transport, expected); assert_eq!(settings.trust_proxy_env, expected); + assert_eq!(settings.ignore_proxy_env, expected); } } diff --git a/litellm-rust/crates/llms/src/custom_httpx/transport.rs b/litellm-rust/crates/llms/src/custom_httpx/transport.rs index 172dd96476a..c42cdf410f6 100644 --- a/litellm-rust/crates/llms/src/custom_httpx/transport.rs +++ b/litellm-rust/crates/llms/src/custom_httpx/transport.rs @@ -11,7 +11,7 @@ pub enum Error { impl Error { pub fn from_reqwest_before_dispatch(error: reqwest::Error) -> Self { let before_dispatch = !error.is_timeout() && (error.is_connect() || error.is_builder()); - let message = error.without_url().to_string(); + let message = describe(error); if before_dispatch { Self::Connect(message) } else { @@ -22,10 +22,18 @@ impl Error { impl From for Error { fn from(error: reqwest::Error) -> Self { - Self::Network(error.without_url().to_string()) + Self::Network(describe(error)) } } +fn describe(error: reqwest::Error) -> String { + let error = error.without_url(); + std::iter::successors(std::error::Error::source(&error), |cause| cause.source()) + .fold(error.to_string(), |message, cause| { + format!("{message}: {cause}") + }) +} + #[cfg(test)] mod tests { #[tokio::test] @@ -47,6 +55,32 @@ mod tests { assert!(!error.to_string().contains("private")); } + fn root_cause(error: &dyn std::error::Error) -> Option { + match error.source() { + Some(cause) => root_cause(cause).or_else(|| Some(cause.to_string())), + None => None, + } + } + + #[tokio::test] + async fn network_error_message_names_the_underlying_cause() { + let listener = std::net::TcpListener::bind("127.0.0.1:0").expect("bind"); + let address = listener.local_addr().expect("address"); + drop(listener); + let error = reqwest::Client::builder() + .no_proxy() + .build() + .expect("client") + .get(format!("http://{address}/private?api_key=secret")) + .send() + .await + .expect_err("nothing listens on the port"); + let root_cause = root_cause(&error).expect("reqwest reports a cause"); + let message = crate::custom_httpx::transport::Error::from(error).to_string(); + assert!(message.contains(&root_cause), "{message}"); + assert!(!message.contains("secret")); + } + #[tokio::test] async fn request_timeout_is_not_safe_to_retry_as_a_connect_failure() { use std::time::Duration; diff --git a/litellm-rust/crates/python-bridge/python_settings.json b/litellm-rust/crates/python-bridge/python_settings.json index 64dd01a0a84..40e36a900d3 100644 --- a/litellm-rust/crates/python-bridge/python_settings.json +++ b/litellm-rust/crates/python-bridge/python_settings.json @@ -7,6 +7,7 @@ "force_ipv4", "http2", "aiohttp_trust_env", + "disable_aiohttp_trust_env", "disable_aiohttp_transport", "user_agent" ] diff --git a/litellm-rust/crates/python-bridge/src/http.rs b/litellm-rust/crates/python-bridge/src/http.rs index cf7ca05515e..2ab6517b61d 100644 --- a/litellm-rust/crates/python-bridge/src/http.rs +++ b/litellm-rust/crates/python-bridge/src/http.rs @@ -72,6 +72,7 @@ struct PythonHttpSettings<'py> { force_ipv4: bool, http2: bool, aiohttp_trust_env: bool, + disable_aiohttp_trust_env: bool, disable_aiohttp_transport: bool, user_agent: String, } @@ -92,6 +93,7 @@ fn settings(value: &Bound<'_, PyAny>) -> PyResult { httpx_transport: python.disable_aiohttp_transport, user_agent: Some(python.user_agent), trust_proxy_env: python.aiohttp_trust_env, + ignore_proxy_env: python.disable_aiohttp_trust_env, ..HttpSettings::default() }) } @@ -133,6 +135,7 @@ defaults = dict( force_ipv4=False, http2=False, aiohttp_trust_env=False, + disable_aiohttp_trust_env=False, disable_aiohttp_transport=False, user_agent='litellm/test', ) @@ -177,6 +180,7 @@ ssl_ecdh_curve='X25519', force_ipv4=True, http2=True, aiohttp_trust_env=True, +disable_aiohttp_trust_env=True, disable_aiohttp_transport=True, user_agent='litellm/9.9.9', ", @@ -194,6 +198,7 @@ user_agent='litellm/9.9.9', httpx_transport: true, user_agent: Some("litellm/9.9.9".into()), trust_proxy_env: true, + ignore_proxy_env: true, ..HttpSettings::default() } ); @@ -295,11 +300,15 @@ user_agent='litellm/9.9.9', #[rstest] #[case::asynchronous(true, false)] #[case::synchronous(false, true)] - fn synchronous_calls_honor_environment_proxies_like_httpx( + fn synchronous_calls_honor_environment_proxies_even_when_aiohttp_opts_out( #[case] asynchronous: bool, #[case] expected: bool, ) { - let settings = for_call(HttpSettings::default(), None, asynchronous); + let opted_out = HttpSettings { + ignore_proxy_env: true, + ..HttpSettings::default() + }; + let settings = for_call(opted_out, None, asynchronous); let config = HttpClientConfig::resolve(&settings).unwrap(); assert_eq!(config.trust_proxy_env, expected); } diff --git a/litellm/rust_bridge/settings.py b/litellm/rust_bridge/settings.py index ad478fb28b5..491312c97b6 100644 --- a/litellm/rust_bridge/settings.py +++ b/litellm/rust_bridge/settings.py @@ -12,6 +12,7 @@ class HttpSettings: force_ipv4: bool http2: bool aiohttp_trust_env: bool + disable_aiohttp_trust_env: bool disable_aiohttp_transport: bool user_agent: str @@ -28,6 +29,7 @@ def http_settings() -> HttpSettings: force_ipv4=litellm.force_ipv4, http2=litellm.http2, aiohttp_trust_env=litellm.aiohttp_trust_env, + disable_aiohttp_trust_env=litellm.disable_aiohttp_trust_env, disable_aiohttp_transport=litellm.disable_aiohttp_transport, user_agent=default_user_agent(), ) diff --git a/tests/test_litellm/rust_bridge/test_settings.py b/tests/test_litellm/rust_bridge/test_settings.py index dce2324de08..f4f9cbc8eec 100644 --- a/tests/test_litellm/rust_bridge/test_settings.py +++ b/tests/test_litellm/rust_bridge/test_settings.py @@ -26,6 +26,7 @@ def test_http_settings_reads_the_litellm_globals(monkeypatch: pytest.MonkeyPatch monkeypatch.setattr(litellm, "force_ipv4", True) monkeypatch.setattr(litellm, "http2", True) monkeypatch.setattr(litellm, "aiohttp_trust_env", True) + monkeypatch.setattr(litellm, "disable_aiohttp_trust_env", True) monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) assert settings.http_settings() == settings.HttpSettings( @@ -36,6 +37,7 @@ def test_http_settings_reads_the_litellm_globals(monkeypatch: pytest.MonkeyPatch force_ipv4=True, http2=True, aiohttp_trust_env=True, + disable_aiohttp_trust_env=True, disable_aiohttp_transport=True, user_agent=default_user_agent(), ) From 4565bbee2f2837e2acde26f969fb4b52f739e62c Mon Sep 17 00:00:00 2001 From: kerry Date: Sat, 19 Sep 2026 00:58:44 +0000 Subject: [PATCH 111/144] style(xai): ruff format stt transformation Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../xai/audio_transcription/transformation.py | 27 ++++++++++++------- 1 file changed, 17 insertions(+), 10 deletions(-) diff --git a/litellm/llms/xai/audio_transcription/transformation.py b/litellm/llms/xai/audio_transcription/transformation.py index b5c5d9c522d..03c06f24a2d 100644 --- a/litellm/llms/xai/audio_transcription/transformation.py +++ b/litellm/llms/xai/audio_transcription/transformation.py @@ -48,7 +48,9 @@ _OBJECT_TUPLE: Final = TypeAdapter(tuple[object, ...]) _STRING_OBJECT_DICT: Final = TypeAdapter(dict[str, object]) -def _serialize_form_value(value: object) -> str | list[str]: # mutable-ok: httpx multipart data takes list values for repeated form fields +def _serialize_form_value( + value: object, +) -> str | list[str]: # mutable-ok: httpx multipart data takes list values for repeated form fields if isinstance(value, bool): return "true" if value else "false" if isinstance(value, (list, tuple)): @@ -61,7 +63,9 @@ class XAIAudioTranscriptionConfig(BaseAudioTranscriptionConfig): def custom_llm_provider(self) -> str: return litellm.LlmProviders.XAI.value - def get_supported_openai_params(self, model: str) -> list[OpenAIAudioTranscriptionOptionalParams]: # mutable-ok: base class signature returns list + def get_supported_openai_params( + self, model: str + ) -> list[OpenAIAudioTranscriptionOptionalParams]: # mutable-ok: base class signature returns list return ["language"] def map_openai_params( @@ -78,7 +82,10 @@ class XAIAudioTranscriptionConfig(BaseAudioTranscriptionConfig): } def get_error_class( - self, error_message: str, status_code: int, headers: dict[str, object] | Headers # mutable-ok: base class signature takes dict + self, + error_message: str, + status_code: int, + headers: dict[str, object] | Headers, # mutable-ok: base class signature takes dict ) -> BaseLLMException: return XAIAudioTranscriptionError(message=error_message, status_code=status_code, headers=headers) @@ -93,16 +100,14 @@ class XAIAudioTranscriptionConfig(BaseAudioTranscriptionConfig): extra_body: Final = optional_params.get("extra_body") flat_params: Final[Mapping[str, object]] = { - **( - _STRING_OBJECT_DICT.validate_python(extra_body) - if isinstance(extra_body, Mapping) - else {} - ), + **(_STRING_OBJECT_DICT.validate_python(extra_body) if isinstance(extra_body, Mapping) else {}), **{k: v for k, v in optional_params.items() if k != "extra_body"}, } excluded_params: Final = frozenset({"model", "OPENAI_TRANSCRIPTION_PARAMS", "extra_body"}) - form_data: Final[dict[str, str | list[str]]] = { # mutable-ok: AudioTranscriptionRequestData.data requires dict and httpx needs list values + form_data: Final[ + dict[str, str | list[str]] + ] = { # mutable-ok: AudioTranscriptionRequestData.data requires dict and httpx needs list values "model": model, **{ k: _serialize_form_value(v) @@ -152,7 +157,9 @@ class XAIAudioTranscriptionConfig(BaseAudioTranscriptionConfig): for word in payload.words ] - hidden_params: Final[dict[str, object]] = dict(payload.model_dump(mode="json")) # mutable-ok: TranscriptionResponse._hidden_params is a dict + hidden_params: Final[dict[str, object]] = dict( + payload.model_dump(mode="json") + ) # mutable-ok: TranscriptionResponse._hidden_params is a dict if payload.duration is not None: hidden_params["audio_transcription_duration"] = payload.duration response._hidden_params = hidden_params # pyright: ignore[reportPrivateUsage] # TranscriptionResponse exposes no public hidden-params setter From 80b0ea6a2f4601e0217786019d2fe37f4b1da83b Mon Sep 17 00:00:00 2001 From: kerry Date: Sat, 19 Sep 2026 01:00:53 +0000 Subject: [PATCH 112/144] test(xai): narrow raises match for missing api key Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../llms/xai/test_xai_audio_transcription_transformation.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_litellm/llms/xai/test_xai_audio_transcription_transformation.py b/tests/test_litellm/llms/xai/test_xai_audio_transcription_transformation.py index f2365f00ba3..0f3445eb400 100644 --- a/tests/test_litellm/llms/xai/test_xai_audio_transcription_transformation.py +++ b/tests/test_litellm/llms/xai/test_xai_audio_transcription_transformation.py @@ -92,7 +92,7 @@ def test_validate_environment_sets_bearer_header(): def test_validate_environment_requires_key(monkeypatch): monkeypatch.delenv("XAI_API_KEY", raising=False) monkeypatch.setattr(litellm, "xai_key", None) - with pytest.raises(ValueError): + with pytest.raises(ValueError, match="xAI API key is required"): CONFIG.validate_environment( headers={}, model="grok-voice-transcribe-2.0", From 5e0ad06725cb97b0c87b27da7072926938e244c4 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 18 Sep 2026 18:02:25 -0700 Subject: [PATCH 113/144] docs(websearch): drop the docstring paragraphs the fix reworded --- litellm/integrations/websearch_interception/handler.py | 5 ----- litellm/llms/anthropic/common_utils.py | 7 ------- 2 files changed, 12 deletions(-) diff --git a/litellm/integrations/websearch_interception/handler.py b/litellm/integrations/websearch_interception/handler.py index 4558d6c2c04..b7238629c85 100644 --- a/litellm/integrations/websearch_interception/handler.py +++ b/litellm/integrations/websearch_interception/handler.py @@ -1354,11 +1354,6 @@ class WebSearchInterceptionLogger(CustomLogger): ) -> tuple[AgenticLoopRequestPatch, tuple[SearchOutcome, ...]]: """ Execute litellm.search() and build follow-up request patch. - - Returns the patch alongside the parallel tuple of search outcomes (one - per tool_call). The caller uses these to optionally build - Anthropic-native ``web_search_tool_result`` content blocks for the - final response and to decide whether a follow-up call is worth making. """ # Extract search queries from tool_use blocks diff --git a/litellm/llms/anthropic/common_utils.py b/litellm/llms/anthropic/common_utils.py index 06561faae8c..05e22ecb55e 100644 --- a/litellm/llms/anthropic/common_utils.py +++ b/litellm/llms/anthropic/common_utils.py @@ -1446,13 +1446,6 @@ def _flattenable_web_search_tool_result(block: object) -> _ReplayedWebSearchTool """ The parsed block when it is a ``web_search_tool_result`` carrying no ``encrypted_content``, else None for anything Anthropic itself issued. - - An empty ``content`` list is flattenable too. It is what the interceptor emits - when a search legitimately returns nothing, and it carries neither evidence to - preserve nor an ``encrypted_content`` to respect, so leaving it in place only - buys the 400 this whole function exists to avoid. The same goes for the - ``web_search_tool_result_error`` object the interceptor emits when a search - raises: it never carries ``encrypted_content``, so it is flattened as well. """ try: parsed: Final = _WEB_SEARCH_TOOL_RESULT_ADAPTER.validate_python(block) From f3bbeed82ff9a9936e1227018dfe794ce778f0dd Mon Sep 17 00:00:00 2001 From: ryan Date: Sat, 19 Sep 2026 01:02:50 +0000 Subject: [PATCH 114/144] feat(proxy): let team admins manage projects via team_admin_editable_team_fields Adds a projects entry to the team_admin_editable_team_fields setting. When set, team admins (legacy admins list or members_with_roles role admin) can call /project/new and /project/update for the teams they administer. The two routes join self_managed_routes so the endpoint check runs instead of the route gate's blanket 401. /project/delete stays proxy admin only Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../management_endpoints/project_endpoints.py | 43 ++++--- litellm/proxy/_types.py | 3 + .../team_admin_field_permissions.py | 17 ++- .../proxy_setting_endpoints.py | 11 +- .../proxy/auth/test_route_checks.py | 34 ++++++ .../test_project_org_authz.py | 111 +++++++++++++++++- .../test_team_admin_field_permissions.py | 20 ++++ .../test_proxy_setting_endpoints.py | 23 +++- .../team/teamAdminEditAccess.test.ts | 1 + .../components/team/teamAdminEditAccess.ts | 1 + 10 files changed, 233 insertions(+), 31 deletions(-) diff --git a/enterprise/litellm_enterprise/proxy/management_endpoints/project_endpoints.py b/enterprise/litellm_enterprise/proxy/management_endpoints/project_endpoints.py index f40ced302ce..2114dfd9849 100644 --- a/enterprise/litellm_enterprise/proxy/management_endpoints/project_endpoints.py +++ b/enterprise/litellm_enterprise/proxy/management_endpoints/project_endpoints.py @@ -11,7 +11,7 @@ Endpoints for /project operations #### PROJECT MANAGEMENT #### import json -from collections.abc import Sequence +from collections.abc import Mapping, Sequence from typing import TYPE_CHECKING, Final from fastapi import APIRouter, Depends, HTTPException, Request @@ -22,7 +22,11 @@ from litellm._uuid import uuid from litellm.proxy._types import * from litellm.proxy.auth.auth_checks import delete_cached_project_object from litellm.proxy.auth.user_api_key_auth import user_api_key_auth -from litellm.proxy.management_endpoints.common_utils import _set_object_metadata_field +from litellm.proxy.management_endpoints.common_utils import ( + _is_user_team_admin, # pyright: ignore[reportPrivateUsage] # shared owner of team-admin membership + _set_object_metadata_field, +) +from litellm.proxy.management_endpoints.team_admin_field_permissions import team_admin_may_manage_projects from litellm.proxy.management_helpers.utils import ( management_endpoint_wrapper, ) @@ -82,37 +86,38 @@ async def _check_user_permission_for_project( user_api_key_dict: UserAPIKeyAuth, team_id: str | None, prisma_client: PrismaClient, + general_settings: Mapping[str, object], require_admin: bool = False, team_object: LiteLLM_TeamTable | None = None, ) -> bool: """ Check if user has permission to manage a project. - Returns True if user is proxy admin or team admin (when team_id provided). + Returns True if user is proxy admin, or a team admin of ``team_id`` when the + ``team_admin_editable_team_fields`` setting grants team admins the ``projects`` permission. If require_admin=True, only proxy admins are allowed. If team_object is provided, it will be used instead of fetching from DB (avoids duplicate DB queries when team was already fetched for validation). """ - is_proxy_admin = user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN + is_proxy_admin: Final = user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN - if require_admin: + if require_admin or is_proxy_admin: return is_proxy_admin - if is_proxy_admin: - return True - - if not team_id or not user_api_key_dict.user_id: + if not team_id or not user_api_key_dict.user_id or not team_admin_may_manage_projects(general_settings): return False - team = team_object - if team is None: - team = await _team_table(prisma_client).find_unique(where={"team_id": team_id}) + team_row: Final = ( + team_object + if team_object is not None + else await _team_table(prisma_client).find_unique(where={"team_id": team_id}) + ) + if team_row is None: + return False - if team and team.admins: - return user_api_key_dict.user_id in team.admins - - return False + team: Final = LiteLLM_TeamTable.model_validate(team_row.model_dump()) + return _is_user_team_admin(user_api_key_dict, team) or user_api_key_dict.user_id in (team.admins or []) async def _validate_team_exists( @@ -531,6 +536,7 @@ async def new_project( user_api_key_dict=user_api_key_dict, team_id=data.team_id, prisma_client=prisma_client, + general_settings=general_settings, team_object=LiteLLM_TeamTable.model_validate(team_object.model_dump()), ) @@ -735,6 +741,7 @@ async def update_project( user_api_key_dict=user_api_key_dict, team_id=existing_project.team_id, prisma_client=prisma_client, + general_settings=general_settings, ) if not has_permission: @@ -751,6 +758,7 @@ async def update_project( user_api_key_dict=user_api_key_dict, team_id=data.team_id, prisma_client=prisma_client, + general_settings=general_settings, team_object=( LiteLLM_TeamTable.model_validate(target_team_obj.model_dump()) if target_team_obj else None ), @@ -877,7 +885,7 @@ async def delete_project( }' ``` """ - from litellm.proxy.proxy_server import premium_user, prisma_client, user_api_key_cache + from litellm.proxy.proxy_server import general_settings, premium_user, prisma_client, user_api_key_cache try: if not premium_user: @@ -899,6 +907,7 @@ async def delete_project( user_api_key_dict=user_api_key_dict, team_id=None, prisma_client=prisma_client, + general_settings=general_settings, require_admin=True, ) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 0097d4b6e92..dd7d17d0f48 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -897,6 +897,9 @@ class LiteLLMRoutes(enum.Enum): # Project read routes - endpoint scopes results to caller's teams (non-admin) "/project/list", "/project/info", + # Project write routes - endpoint checks team admin + team_admin_editable_team_fields "projects" + "/project/new", + "/project/update", # Endpoint enforces proxy-admin vs team-admin model access itself. "/health/test_connection", # Invitation routes - org/team admins checked in endpoint via _user_has_admin_privileges diff --git a/litellm/proxy/management_endpoints/team_admin_field_permissions.py b/litellm/proxy/management_endpoints/team_admin_field_permissions.py index 56d455494c6..6038775d96b 100644 --- a/litellm/proxy/management_endpoints/team_admin_field_permissions.py +++ b/litellm/proxy/management_endpoints/team_admin_field_permissions.py @@ -1,4 +1,5 @@ -"""Proxy-wide allow-list of team-settings fields a team admin may change on /team/update.""" +"""Proxy-wide allow-list of what a team admin may do on the teams they administer: team-settings fields on +/team/update, plus the ``projects`` permission for /project/new and /project/update.""" from collections.abc import Mapping from dataclasses import dataclass @@ -21,6 +22,10 @@ TEAM_ADMIN_EDITABLE_TEAM_FIELDS_SETTING: Final = "team_admin_editable_team_field # TODO(LIT-5722): add the remaining team settings one per PR, each with its value-diff tests and dashboard field SUPPORTED_TEAM_ADMIN_EDITABLE_TEAM_FIELDS: Final[frozenset[str]] = frozenset({"tpm_limit", "rpm_limit", "max_budget"}) +TEAM_ADMIN_PROJECTS_PERMISSION: Final = "projects" +SUPPORTED_TEAM_ADMIN_PERMISSIONS: Final[frozenset[str]] = SUPPORTED_TEAM_ADMIN_EDITABLE_TEAM_FIELDS | { + TEAM_ADMIN_PROJECTS_PERMISSION +} _FIELD_LIST: Final = TypeAdapter(list[str]) _JSON_OBJECT: Final = TypeAdapter(dict[str, object]) @@ -67,17 +72,23 @@ def resolve_team_admin_editable_fields( "%s must be a list of field names; ignoring %r", TEAM_ADMIN_EDITABLE_TEAM_FIELDS_SETTING, raw ) return frozenset() - unsupported: Final = configured - supported + unsupported: Final = configured - supported - SUPPORTED_TEAM_ADMIN_PERMISSIONS if unsupported: verbose_proxy_logger.warning( "%s ignores unsupported field(s) %s; supported: %s", TEAM_ADMIN_EDITABLE_TEAM_FIELDS_SETTING, sorted(unsupported), - sorted(supported), + sorted(supported | SUPPORTED_TEAM_ADMIN_PERMISSIONS), ) return configured & supported +def team_admin_may_manage_projects(general_settings: Mapping[str, object]) -> bool: + return TEAM_ADMIN_PROJECTS_PERMISSION in resolve_team_admin_editable_fields( + general_settings, frozenset({TEAM_ADMIN_PROJECTS_PERMISSION}) + ) + + def _as_object(value: object) -> Mapping[str, object]: try: return _JSON_OBJECT.validate_json(value) if isinstance(value, str) else _JSON_OBJECT.validate_python(value) diff --git a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py index fd160636d46..75431383fbd 100644 --- a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py +++ b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py @@ -30,7 +30,7 @@ from litellm.proxy.config_resolvers.sso import ( resolve_sso_config, ) from litellm.proxy.management_endpoints.team_admin_field_permissions import ( - SUPPORTED_TEAM_ADMIN_EDITABLE_TEAM_FIELDS, + SUPPORTED_TEAM_ADMIN_PERMISSIONS, TEAM_ADMIN_EDITABLE_TEAM_FIELDS_SETTING, ) from litellm.proxy.spend_tracking.ptu_feature_flag import is_ptu_cost_attribution_enabled @@ -216,7 +216,7 @@ class UIThemeSettingsResponse(SettingsResponse): """Response model for UI theme settings""" -_TEAM_ADMIN_FIELD_ENUM: Final = tuple(sorted(SUPPORTED_TEAM_ADMIN_EDITABLE_TEAM_FIELDS)) +_TEAM_ADMIN_FIELD_ENUM: Final = tuple(sorted(SUPPORTED_TEAM_ADMIN_PERMISSIONS)) class UISettings(BaseModel): @@ -315,7 +315,8 @@ class UISettings(BaseModel): default=(), description=( "Team settings fields a team admin may change on the teams they administer. " - "Empty means team admins cannot edit team settings at all. " + "Include 'projects' to let team admins create and update projects for those teams. " + "Empty means team admins cannot edit team settings or manage projects at all. " "Proxy admins and org admins are not affected." ), json_schema_extra={ # mutable-ok: pydantic only merges json_schema_extra when it is a plain dict @@ -1626,7 +1627,7 @@ async def update_ui_settings( raise HTTPException(status_code=422, detail=e.errors()) unsupported_team_fields: Final = sorted( - frozenset(settings.team_admin_editable_team_fields) - SUPPORTED_TEAM_ADMIN_EDITABLE_TEAM_FIELDS + frozenset(settings.team_admin_editable_team_fields) - SUPPORTED_TEAM_ADMIN_PERMISSIONS ) if unsupported_team_fields: raise HTTPException( @@ -1634,7 +1635,7 @@ async def update_ui_settings( detail={ # mutable-ok: HTTPException detail must be a plain dict for FastAPI JSON serialization "error": ( f"{TEAM_ADMIN_EDITABLE_TEAM_FIELDS_SETTING} does not support {unsupported_team_fields}. " - f"Supported fields: {sorted(SUPPORTED_TEAM_ADMIN_EDITABLE_TEAM_FIELDS)}." + f"Supported fields: {sorted(SUPPORTED_TEAM_ADMIN_PERMISSIONS)}." ) }, ) diff --git a/tests/test_litellm/proxy/auth/test_route_checks.py b/tests/test_litellm/proxy/auth/test_route_checks.py index 72c59223549..2e0d2c710a4 100644 --- a/tests/test_litellm/proxy/auth/test_route_checks.py +++ b/tests/test_litellm/proxy/auth/test_route_checks.py @@ -4038,3 +4038,37 @@ def test_team_key_without_service_account_marker_still_rejected(): valid_token=valid_token, request_data={}, ) + + +@pytest.mark.parametrize("route", ["/project/new", "/project/update"]) +def test_project_write_routes_reach_endpoint_for_internal_user(route): + """The route gate lets a non-admin through so /project/new and /project/update can apply the + team_admin_editable_team_fields projects permission themselves, instead of a blanket 401.""" + valid_token = UserAPIKeyAuth(user_id="team_admin", user_role=LitellmUserRoles.INTERNAL_USER.value) + request = MagicMock(spec=Request) + request.query_params = {} + + RouteChecks.non_proxy_admin_allowed_routes_check( + user_obj=None, + _user_role=LitellmUserRoles.INTERNAL_USER.value, + route=route, + request=request, + valid_token=valid_token, + request_data={}, + ) + + +def test_project_delete_route_stays_proxy_admin_only(): + valid_token = UserAPIKeyAuth(user_id="team_admin", user_role=LitellmUserRoles.INTERNAL_USER.value) + request = MagicMock(spec=Request) + request.query_params = {} + + with pytest.raises(Exception, match="Only proxy admin can be used to generate, delete, update"): + RouteChecks.non_proxy_admin_allowed_routes_check( + user_obj=None, + _user_role=LitellmUserRoles.INTERNAL_USER.value, + route="/project/delete", + request=request, + valid_token=valid_token, + request_data={}, + ) diff --git a/tests/test_litellm/proxy/management_endpoints/test_project_org_authz.py b/tests/test_litellm/proxy/management_endpoints/test_project_org_authz.py index a06d79306ab..ce08030739d 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_project_org_authz.py +++ b/tests/test_litellm/proxy/management_endpoints/test_project_org_authz.py @@ -7,12 +7,18 @@ Unit tests for the VERIA-55 fixes: member of. """ +from types import MappingProxyType +from typing import Final from unittest.mock import AsyncMock, MagicMock import pytest from fastapi import HTTPException -from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth +from litellm.models.team import LiteLLM_TeamTable +from litellm.proxy._types import LitellmUserRoles, Member, UserAPIKeyAuth + +_PROJECTS_ENABLED: Final = MappingProxyType({"team_admin_editable_team_fields": ["projects"]}) +_PROJECTS_DISABLED: Final = MappingProxyType({"team_admin_editable_team_fields": ["max_budget"]}) # --------------------------------------------------------------------------- @@ -20,11 +26,9 @@ from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth # --------------------------------------------------------------------------- -def _make_prisma_with_team(team_id: str, admins: list): +def _make_prisma_with_team(team_id: str, admins: list, members_with_roles: tuple[Member, ...] = ()): prisma = MagicMock() - team_row = MagicMock() - team_row.team_id = team_id - team_row.admins = admins + team_row = LiteLLM_TeamTable(team_id=team_id, admins=admins, members_with_roles=list(members_with_roles)) prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=team_row) return prisma @@ -49,6 +53,7 @@ async def test_project_perm_check_uses_current_team_not_caller_supplied(): user_api_key_dict=caller, team_id="team-A", prisma_client=prisma, + general_settings=_PROJECTS_ENABLED, ) assert has_perm is False prisma.db.litellm_teamtable.find_unique.assert_awaited_once() @@ -70,10 +75,105 @@ async def test_project_perm_check_allows_team_admin_of_existing_team(): user_api_key_dict=alice, team_id="team-A", prisma_client=prisma, + general_settings=_PROJECTS_ENABLED, ) assert has_perm is True +@pytest.mark.asyncio +async def test_project_perm_check_allows_members_with_roles_admin(): + """Team admins added through /team/member_add live in members_with_roles, not the legacy admins list.""" + from litellm_enterprise.proxy.management_endpoints.project_endpoints import ( + _check_user_permission_for_project, + ) + + prisma = _make_prisma_with_team( + team_id="team-A", + admins=[], + members_with_roles=(Member(user_id="carol", role="admin"), Member(user_id="dave", role="user")), + ) + carol = UserAPIKeyAuth(user_id="carol", user_role=LitellmUserRoles.INTERNAL_USER.value) + dave = UserAPIKeyAuth(user_id="dave", user_role=LitellmUserRoles.INTERNAL_USER.value) + + assert ( + await _check_user_permission_for_project( + user_api_key_dict=carol, team_id="team-A", prisma_client=prisma, general_settings=_PROJECTS_ENABLED + ) + is True + ) + assert ( + await _check_user_permission_for_project( + user_api_key_dict=dave, team_id="team-A", prisma_client=prisma, general_settings=_PROJECTS_ENABLED + ) + is False + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("general_settings", [MappingProxyType({}), _PROJECTS_DISABLED]) +async def test_project_perm_check_denies_team_admin_unless_projects_permission_configured(general_settings): + from litellm_enterprise.proxy.management_endpoints.project_endpoints import ( + _check_user_permission_for_project, + ) + + prisma = _make_prisma_with_team( + team_id="team-A", admins=["alice"], members_with_roles=(Member(user_id="carol", role="admin"),) + ) + + for user_id in ("alice", "carol"): + caller = UserAPIKeyAuth(user_id=user_id, user_role=LitellmUserRoles.INTERNAL_USER.value) + has_perm = await _check_user_permission_for_project( + user_api_key_dict=caller, + team_id="team-A", + prisma_client=prisma, + general_settings=general_settings, + ) + assert has_perm is False + prisma.db.litellm_teamtable.find_unique.assert_not_called() + + +@pytest.mark.asyncio +async def test_project_perm_check_require_admin_denies_team_admin_even_when_configured(): + """/project/delete passes require_admin=True, so the projects permission must not open it up.""" + from litellm_enterprise.proxy.management_endpoints.project_endpoints import ( + _check_user_permission_for_project, + ) + + prisma = _make_prisma_with_team(team_id="team-A", admins=["alice"]) + alice = UserAPIKeyAuth(user_id="alice", user_role=LitellmUserRoles.INTERNAL_USER.value) + + has_perm = await _check_user_permission_for_project( + user_api_key_dict=alice, + team_id=None, + prisma_client=prisma, + general_settings=_PROJECTS_ENABLED, + require_admin=True, + ) + assert has_perm is False + prisma.db.litellm_teamtable.find_unique.assert_not_called() + + +@pytest.mark.asyncio +async def test_project_perm_check_uses_injected_team_object_for_reassignment_target(): + from litellm_enterprise.proxy.management_endpoints.project_endpoints import ( + _check_user_permission_for_project, + ) + + prisma = _make_prisma_with_team(team_id="team-A", admins=["alice"]) + alice = UserAPIKeyAuth(user_id="alice", user_role=LitellmUserRoles.INTERNAL_USER.value) + target_team = LiteLLM_TeamTable(team_id="team-B", members_with_roles=[Member(user_id="erin", role="admin")]) + + has_perm = await _check_user_permission_for_project( + user_api_key_dict=alice, + team_id="team-B", + prisma_client=prisma, + general_settings=_PROJECTS_ENABLED, + team_object=target_team, + ) + assert has_perm is False + prisma.db.litellm_teamtable.find_unique.assert_not_called() + + @pytest.mark.asyncio async def test_project_perm_check_proxy_admin_always_allowed(): from litellm_enterprise.proxy.management_endpoints.project_endpoints import ( @@ -90,6 +190,7 @@ async def test_project_perm_check_proxy_admin_always_allowed(): user_api_key_dict=admin, team_id="team-A", prisma_client=prisma, + general_settings=MappingProxyType({}), ) assert has_perm is True # Admin shortcut should not even hit the DB. diff --git a/tests/test_litellm/proxy/management_endpoints/test_team_admin_field_permissions.py b/tests/test_litellm/proxy/management_endpoints/test_team_admin_field_permissions.py index 5b31089f91e..1a72d1de393 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_admin_field_permissions.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_admin_field_permissions.py @@ -9,6 +9,7 @@ from litellm.proxy.management_endpoints.team_admin_field_permissions import ( changed_team_fields, resolve_team_admin_editable_fields, team_admin_edit_verdict, + team_admin_may_manage_projects, team_admin_request_or_raise, ) @@ -31,6 +32,25 @@ class TestResolveTeamAdminEditableFields: def test_malformed_setting_fails_closed(self, raw): assert resolve_team_admin_editable_fields({"team_admin_editable_team_fields": raw}, _SUPPORTED) == frozenset() + def test_projects_permission_is_not_a_team_field(self): + configured = {"team_admin_editable_team_fields": ["projects", "tpm_limit"]} + assert resolve_team_admin_editable_fields(configured, _SUPPORTED) == frozenset({"tpm_limit"}) + + +class TestTeamAdminMayManageProjects: + def test_missing_setting_denies(self): + assert team_admin_may_manage_projects({}) is False + + def test_team_fields_alone_do_not_grant_projects(self): + assert team_admin_may_manage_projects({"team_admin_editable_team_fields": ["tpm_limit", "max_budget"]}) is False + + def test_projects_entry_grants(self): + assert team_admin_may_manage_projects({"team_admin_editable_team_fields": ["max_budget", "projects"]}) is True + + @pytest.mark.parametrize("raw", ["projects", 7, [1, 2]]) + def test_malformed_setting_denies(self, raw): + assert team_admin_may_manage_projects({"team_admin_editable_team_fields": raw}) is False + class TestChangedTeamFields: def test_team_id_alone_changes_nothing(self): diff --git a/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py b/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py index fc5733b1a82..9860d1bf94a 100644 --- a/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py +++ b/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py @@ -3291,7 +3291,7 @@ class TestTeamAdminEditableTeamFieldsSetting: def test_patch_rejects_field_names_the_proxy_does_not_support(self, monkeypatch): mock_prisma = self._as_proxy_admin(monkeypatch) monkeypatch.setattr( - "litellm.proxy.ui_crud_endpoints.proxy_setting_endpoints.SUPPORTED_TEAM_ADMIN_EDITABLE_TEAM_FIELDS", + "litellm.proxy.ui_crud_endpoints.proxy_setting_endpoints.SUPPORTED_TEAM_ADMIN_PERMISSIONS", frozenset({"tpm_limit"}), ) @@ -3336,6 +3336,26 @@ class TestTeamAdminEditableTeamFieldsSetting: assert stored["team_admin_editable_team_fields"] == enabled assert general_settings["team_admin_editable_team_fields"] == enabled + def test_patch_accepts_the_projects_permission_and_project_endpoints_see_it(self, monkeypatch): + from litellm.proxy.management_endpoints.team_admin_field_permissions import ( + team_admin_may_manage_projects, + ) + + mock_prisma = self._as_proxy_admin(monkeypatch) + general_settings: dict = {} + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", general_settings) + assert team_admin_may_manage_projects(general_settings) is False + + try: + response = client.patch("/update/ui_settings", json={"team_admin_editable_team_fields": ["projects"]}) + finally: + app.dependency_overrides.clear() + + assert response.status_code == 200 + stored = json.loads(mock_prisma.db.litellm_uisettings.upsert.call_args.kwargs["data"]["create"]["ui_settings"]) + assert stored["team_admin_editable_team_fields"] == ["projects"] + assert team_admin_may_manage_projects(general_settings) is True + def test_patch_with_an_empty_list_turns_team_admin_editing_off_again(self, monkeypatch): mock_prisma = self._as_proxy_admin(monkeypatch) general_settings: dict = {"team_admin_editable_team_fields": ["tpm_limit"]} @@ -3372,6 +3392,7 @@ class TestTeamAdminEditableTeamFieldsSetting: assert field_schema["type"] == "array" assert field_schema["items"]["type"] == "string" assert "tpm_limit" in field_schema["items"]["enum"] + assert "projects" in field_schema["items"]["enum"] class TestSyncUiSettingsToGeneralSettings: diff --git a/ui/litellm-dashboard/src/components/team/teamAdminEditAccess.test.ts b/ui/litellm-dashboard/src/components/team/teamAdminEditAccess.test.ts index da3f9bf8289..1ef3816b5c4 100644 --- a/ui/litellm-dashboard/src/components/team/teamAdminEditAccess.test.ts +++ b/ui/litellm-dashboard/src/components/team/teamAdminEditAccess.test.ts @@ -13,6 +13,7 @@ describe("teamAdminFieldLabel", () => { ["tpm_limit", "Tokens per minute Limit (TPM)"], ["rpm_limit", "Requests per minute Limit (RPM)"], ["max_budget", "Max Budget (USD)"], + ["projects", "Create and update projects"], ])("names %s the way the team settings form does", (field, label) => { expect(teamAdminFieldLabel(field)).toBe(label); }); diff --git a/ui/litellm-dashboard/src/components/team/teamAdminEditAccess.ts b/ui/litellm-dashboard/src/components/team/teamAdminEditAccess.ts index b878af03df6..5706eeafe82 100644 --- a/ui/litellm-dashboard/src/components/team/teamAdminEditAccess.ts +++ b/ui/litellm-dashboard/src/components/team/teamAdminEditAccess.ts @@ -47,6 +47,7 @@ const TEAM_ADMIN_FIELD_LABELS: ReadonlyMap = new Map([ ["tpm_limit", "Tokens per minute Limit (TPM)"], ["rpm_limit", "Requests per minute Limit (RPM)"], ["max_budget", "Max Budget (USD)"], + ["projects", "Create and update projects"], ]); export const teamAdminFieldLabel = (field: string): string => TEAM_ADMIN_FIELD_LABELS.get(field) ?? field; From a124e079c369723fd236c8ae59cf5ed80624a6c3 Mon Sep 17 00:00:00 2001 From: kerry Date: Sat, 19 Sep 2026 01:06:09 +0000 Subject: [PATCH 115/144] refactor(xai): use derived provider set for transcription routing Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/constants.py | 4 ++++ litellm/main.py | 7 ++----- litellm/utils.py | 4 +--- 3 files changed, 7 insertions(+), 8 deletions(-) diff --git a/litellm/constants.py b/litellm/constants.py index 56d3f5450d7..55f92f29c96 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1001,6 +1001,10 @@ openai_compatible_providers: Final[list] = [ OPENAI_COMPATIBLE_PROVIDERS_WITH_NATIVE_AUDIO_TRANSCRIPTION: Final = frozenset({"xai"}) +OPENAI_AUDIO_TRANSCRIPTION_PROVIDERS: Final = frozenset( + {"openai"} | (frozenset(openai_compatible_providers) - OPENAI_COMPATIBLE_PROVIDERS_WITH_NATIVE_AUDIO_TRANSCRIPTION) +) + openai_text_completion_compatible_providers: Final[list] = [ # providers that support `/v1/completions` "together_ai", "fireworks_ai", diff --git a/litellm/main.py b/litellm/main.py index 1c3de8e766b..bd10c3924f7 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -64,7 +64,7 @@ from litellm.constants import ( AZURE_OPENAI_AUDIO_PROVIDERS, DEFAULT_MOCK_RESPONSE_COMPLETION_TOKEN_COUNT, DEFAULT_MOCK_RESPONSE_PROMPT_TOKEN_COUNT, - OPENAI_COMPATIBLE_PROVIDERS_WITH_NATIVE_AUDIO_TRANSCRIPTION, + OPENAI_AUDIO_TRANSCRIPTION_PROVIDERS, ) from litellm.exceptions import LiteLLMUnknownProvider from litellm.integrations.custom_logger import CustomLogger @@ -7860,10 +7860,7 @@ def transcription( litellm_params=litellm_params_dict, custom_llm_provider=custom_llm_provider, ) - elif custom_llm_provider == "openai" or ( - custom_llm_provider in litellm.openai_compatible_providers - and custom_llm_provider not in OPENAI_COMPATIBLE_PROVIDERS_WITH_NATIVE_AUDIO_TRANSCRIPTION - ): + elif custom_llm_provider in OPENAI_AUDIO_TRANSCRIPTION_PROVIDERS: api_base = ( api_base or litellm.api_base diff --git a/litellm/utils.py b/litellm/utils.py index 97c074ce112..e30e8cde86d 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -8731,9 +8731,7 @@ class ProviderConfigManager: return ElevenLabsAudioTranscriptionConfig() elif litellm.LlmProviders.XAI == provider: - from litellm.llms.xai.audio_transcription.transformation import ( - XAIAudioTranscriptionConfig, - ) + from litellm.llms.xai.audio_transcription.transformation import XAIAudioTranscriptionConfig return XAIAudioTranscriptionConfig() elif litellm.LlmProviders.OPENAI == provider: From a1560936f794a8bcd394ea02cfa0c8f9ab01ab10 Mon Sep 17 00:00:00 2001 From: kerry Date: Sat, 19 Sep 2026 01:06:53 +0000 Subject: [PATCH 116/144] fix(timing): use epoch math for detailed pre-processing and drop client-supplied timing windows Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../llm_response_utils/response_metadata.py | 2 +- litellm/proxy/litellm_pre_call_utils.py | 1 + .../test_response_metadata.py | 7 +++-- .../proxy/test_litellm_pre_call_utils.py | 30 +++++++++++++++++++ 4 files changed, 36 insertions(+), 4 deletions(-) diff --git a/litellm/litellm_core_utils/llm_response_utils/response_metadata.py b/litellm/litellm_core_utils/llm_response_utils/response_metadata.py index cc0d10ee7a6..93701b3c1e7 100644 --- a/litellm/litellm_core_utils/llm_response_utils/response_metadata.py +++ b/litellm/litellm_core_utils/llm_response_utils/response_metadata.py @@ -205,7 +205,7 @@ class ResponseMetadata: api_call_start: Final[datetime.datetime | None] = logging_obj.model_call_details.get("api_call_start_time") if api_call_start is not None and start_time is not None: anchor: Final = _timing_window_start(start_time, logging_obj)[0] - pre_ms: Final = (api_call_start - anchor).total_seconds() * 1000 + pre_ms: Final = (api_call_start.timestamp() - anchor.timestamp()) * 1000 detailed["timing_pre_processing_ms"] = round(pre_ms, 4) # post-processing = total - pre - llm_api diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index b03f1e4348c..9a973755894 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -2367,6 +2367,7 @@ async def add_litellm_data_to_request( # OTel layer can compute pre-request latency, including on the failure # path after the logging object is popped. data[_metadata_variable_name]["litellm_received_at"] = getattr(request.state, "litellm_received_at", None) + data[_metadata_variable_name]["llm_api_timing_windows"] = () # OTEL Controls / Tracing # Add the OTEL Parent Trace before sending it LiteLLM diff --git a/tests/test_litellm/litellm_core_utils/llm_response_utils/test_response_metadata.py b/tests/test_litellm/litellm_core_utils/llm_response_utils/test_response_metadata.py index 97c154db783..3379879a8a6 100644 --- a/tests/test_litellm/litellm_core_utils/llm_response_utils/test_response_metadata.py +++ b/tests/test_litellm/litellm_core_utils/llm_response_utils/test_response_metadata.py @@ -474,12 +474,13 @@ class TestDetailedTiming: monkeypatch.setattr(response_metadata_mod, "LITELLM_DETAILED_TIMING", True) result = ModelResponse() - start = datetime.datetime(2025, 1, 1, 0, 0, 0) - received_at = start - datetime.timedelta(milliseconds=200) + received_at = datetime.datetime.now(datetime.timezone.utc) + start = received_at + datetime.timedelta(milliseconds=200) + api_call_start = start.replace(tzinfo=None) end = start + datetime.timedelta(milliseconds=530) logging_obj = self._make_logging_obj( llm_api_duration_ms=500.0, - api_call_start_time=start, + api_call_start_time=api_call_start, ) logging_obj.model_call_details["litellm_params"] = {"metadata": {"litellm_received_at": received_at}} diff --git a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py index f4490519554..88d38d74f49 100644 --- a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py +++ b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py @@ -333,6 +333,36 @@ async def test_arrival_time_prefers_litellm_received_at_over_time_time(): assert updated_data["proxy_server_request"]["arrival_time"] == received_at.timestamp() +@pytest.mark.asyncio +async def test_proxy_clears_client_supplied_timing_windows(): + request_mock = MagicMock(spec=Request) + request_mock.url = MagicMock() + request_mock.url.path = "/v1/chat/completions" + request_mock.url.__str__.return_value = "http://localhost/v1/chat/completions" + request_mock.method = "POST" + request_mock.query_params = {} + request_mock.headers = {"Content-Type": "application/json"} + request_mock.client = MagicMock() + request_mock.client.host = "127.0.0.1" + request_mock.state = SimpleNamespace(litellm_received_at=datetime.now(timezone.utc)) + + user_api_key_dict = UserAPIKeyAuth(api_key="hashed-key", metadata={}, team_metadata={}) + + updated_data = await add_litellm_data_to_request( + data={ + "model": "gpt-3.5-turbo", + "metadata": {"llm_api_timing_windows": ((0.0, 1.0),)}, + }, + request=request_mock, + user_api_key_dict=user_api_key_dict, + proxy_config=MagicMock(), + general_settings={}, + version="test-version", + ) + + assert updated_data["metadata"]["llm_api_timing_windows"] == () + + @pytest.mark.asyncio async def test_arrival_time_falls_back_to_time_time_without_litellm_received_at(): """Callers that never went through user_api_key_auth (no stamp on request.state) From 1b305cd6b960ba0d71271b65bf86ab8bc11b62d2 Mon Sep 17 00:00:00 2001 From: kerry Date: Sat, 19 Sep 2026 01:07:07 +0000 Subject: [PATCH 117/144] fix(cost_calc): default fireworks cached input to the documented 50% discount when the map has no cache-read rate Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/constants.py | 3 + .../litellm_core_utils/llm_cost_calc/utils.py | 59 ++++++-- litellm/llms/fireworks_ai/cost_calculator.py | 29 +--- .../llm_cost_calc/test_llm_cost_calc_utils.py | 36 ++++- .../test_fireworks_ai_cost_calculator.py | 126 ++++++++++++++++-- 5 files changed, 204 insertions(+), 49 deletions(-) diff --git a/litellm/constants.py b/litellm/constants.py index e7cb21a3a7d..53e3d032356 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -576,6 +576,9 @@ FIREWORKS_AI_176_B_MOE: Final = int(os.getenv("FIREWORKS_AI_176_B_MOE", 176)) FIREWORKS_AI_4_B: Final = int(os.getenv("FIREWORKS_AI_4_B", 4)) FIREWORKS_AI_16_B: Final = int(os.getenv("FIREWORKS_AI_16_B", 16)) FIREWORKS_AI_80_B: Final = int(os.getenv("FIREWORKS_AI_80_B", 80)) +# https://docs.fireworks.ai/guides/prompt-caching (accessed 2026-09-19): serverless cached prompt tokens +# default to a 50% discount off the input rate +FIREWORKS_AI_DEFAULT_CACHE_READ_RATE_RATIO: Final = 0.5 #### Logging callback constants #### REDACTED_BY_LITELM_STRING: Final = "REDACTED_BY_LITELM" MAX_LANGFUSE_INITIALIZED_CLIENTS: Final = int(os.getenv("MAX_LANGFUSE_INITIALIZED_CLIENTS", 50)) diff --git a/litellm/litellm_core_utils/llm_cost_calc/utils.py b/litellm/litellm_core_utils/llm_cost_calc/utils.py index f8eb15dca88..7de4534ef2c 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/utils.py +++ b/litellm/litellm_core_utils/llm_cost_calc/utils.py @@ -14,6 +14,7 @@ from typing_extensions import ReadOnly import litellm from litellm._internal_context import current_billing_time from litellm._logging import verbose_logger +from litellm.constants import FIREWORKS_AI_DEFAULT_CACHE_READ_RATE_RATIO from litellm.litellm_core_utils.llm_cost_calc.tiered_pricing import ( select_tier_for_input, tier_rate, @@ -72,6 +73,34 @@ def _uses_inclusive_token_thresholds(custom_llm_provider: str | None) -> bool: return custom_llm_provider in _INCLUSIVE_THRESHOLD_PROVIDERS +def apply_provider_cache_read_default(model_info: ModelInfo, custom_llm_provider: str | None) -> ModelInfo: + """Apply provider-specific defaults for cache-read pricing.""" + if custom_llm_provider != "fireworks_ai": + return model_info + input_rate: Final = model_info.get("input_cost_per_token") + if model_info.get("cache_read_input_token_cost") is not None or input_rate is None: + return model_info + cache_read_rate: Final = input_rate * FIREWORKS_AI_DEFAULT_CACHE_READ_RATE_RATIO + off_peak: Final = model_info.get("off_peak_pricing") + if off_peak is None or "cache_read_input_token_cost" in off_peak: + return cast(ModelInfo, {**model_info, "cache_read_input_token_cost": cache_read_rate}) + return cast( + ModelInfo, + { + **model_info, + "cache_read_input_token_cost": cache_read_rate, + "off_peak_pricing": { + **off_peak, + "cache_read_input_token_cost": ( + off_peak["input_cost_per_token"] * FIREWORKS_AI_DEFAULT_CACHE_READ_RATE_RATIO + if "input_cost_per_token" in off_peak + else cache_read_rate + ), + }, + }, + ) + + def _get_token_detail_value(details: object, key: str) -> int | None: if isinstance(details, dict): value = details.get(key) @@ -1170,8 +1199,10 @@ def generic_cost_per_token( # rather than handing back a name for this to re-resolve. A name cannot express a # per-deployment override: those are registered under the deployment id and kept off # the shared model-name key, so resolving from the name here reads the public rate. - if model_info is None: - model_info = get_model_info(model=model, custom_llm_provider=custom_llm_provider) + resolved_model_info: Final = apply_provider_cache_read_default( + get_model_info(model=model, custom_llm_provider=custom_llm_provider) if model_info is None else model_info, + custom_llm_provider, + ) ## CALCULATE INPUT COST ### Cost of processing (non-cache hit + cache hit) + Cost of cache-writing (cache writing) @@ -1236,7 +1267,7 @@ def generic_cost_per_token( cache_creation_cost_above_1hr, cache_read_cost, ) = _get_token_base_cost( - model_info=model_info, + model_info=resolved_model_info, usage=usage, service_tier=service_tier, current_time=billing_time, @@ -1245,7 +1276,7 @@ def generic_cost_per_token( prompt_cost = _calculate_input_cost( prompt_tokens_details=prompt_tokens_details, - model_info=model_info, + model_info=resolved_model_info, prompt_base_cost=prompt_base_cost, cache_read_cost=cache_read_cost, cache_creation_cost=cache_creation_cost, @@ -1290,7 +1321,7 @@ def generic_cost_per_token( ## AUDIO COST if not is_text_tokens_total and audio_tokens is not None and audio_tokens > 0: - _output_cost_per_audio_token = _get_cost_per_unit(model_info, "output_cost_per_audio_token", None) + _output_cost_per_audio_token = _get_cost_per_unit(resolved_model_info, "output_cost_per_audio_token", None) _output_cost_per_audio_token = ( _output_cost_per_audio_token if _output_cost_per_audio_token is not None else completion_base_cost ) @@ -1299,7 +1330,7 @@ def generic_cost_per_token( ## REASONING COST if not is_text_tokens_total and reasoning_tokens and reasoning_tokens > 0: completion_cost += float(reasoning_tokens) * _resolve_billed_reasoning_rate( - model_info=model_info, + model_info=resolved_model_info, usage=usage, service_tier=service_tier, completion_base_cost=completion_base_cost, @@ -1308,7 +1339,7 @@ def generic_cost_per_token( ## IMAGE COST if not is_text_tokens_total and image_tokens and image_tokens > 0: - _output_cost_per_image_token = _get_cost_per_unit(model_info, "output_cost_per_image_token", None) + _output_cost_per_image_token = _get_cost_per_unit(resolved_model_info, "output_cost_per_image_token", None) _output_cost_per_image_token = ( _output_cost_per_image_token if _output_cost_per_image_token is not None else completion_base_cost ) @@ -1316,7 +1347,7 @@ def generic_cost_per_token( ## VIDEO COST if not is_text_tokens_total and video_tokens and video_tokens > 0: - _output_cost_per_video_token = _get_cost_per_unit(model_info, "output_cost_per_video_token", None) + _output_cost_per_video_token = _get_cost_per_unit(resolved_model_info, "output_cost_per_video_token", None) _output_cost_per_video_token = ( _output_cost_per_video_token if _output_cost_per_video_token is not None else completion_base_cost ) @@ -1325,12 +1356,12 @@ def generic_cost_per_token( ## REGIONAL DATA-RESIDENCY UPLIFT # Applied as a flat multiplier across all token costs for the request # when the upstream is a regionalized OpenAI host (eu./us.api.openai.com). - uplift: Final = _get_regional_uplift_multiplier(model_info, data_residency) + uplift: Final = _get_regional_uplift_multiplier(resolved_model_info, data_residency) if uplift != 1.0: prompt_cost *= uplift completion_cost *= uplift - vertex_uplift: Final = get_vertex_regional_endpoint_uplift(model_info, vertex_location) + vertex_uplift: Final = get_vertex_regional_endpoint_uplift(resolved_model_info, vertex_location) if vertex_uplift != 1.0: prompt_cost *= vertex_uplift completion_cost *= vertex_uplift @@ -1487,7 +1518,10 @@ def get_billed_token_rates( if custom_cost_per_token is not None: return _custom_pricing_rates(custom_cost_per_token) try: - model_info: Final = get_model_info(model=model, custom_llm_provider=custom_llm_provider) + model_info: Final = apply_provider_cache_read_default( + get_model_info(model=model, custom_llm_provider=custom_llm_provider), + custom_llm_provider, + ) except Exception: # noqa: BLE001 # get_model_info raises a bare Exception for an unmapped model: no rates return None return _cost_map_billed_rates( @@ -1578,8 +1612,9 @@ def calculate_prompt_caching_savings( ``billed_at`` is the request's completion time, so off-peak windows resolve as the biller saw them rather than at the later spend write. """ + model_info_with_cache_read_default: Final = apply_provider_cache_read_default(model_info, custom_llm_provider) prompt_base_cost, _, cache_creation_cost, cache_creation_cost_above_1hr, cache_read_cost = _get_token_base_cost( - model_info=model_info, + model_info=model_info_with_cache_read_default, usage=usage, service_tier=service_tier, current_time=billed_at, diff --git a/litellm/llms/fireworks_ai/cost_calculator.py b/litellm/llms/fireworks_ai/cost_calculator.py index 1795a700d25..4b6ca7c9896 100644 --- a/litellm/llms/fireworks_ai/cost_calculator.py +++ b/litellm/llms/fireworks_ai/cost_calculator.py @@ -3,10 +3,7 @@ For calculating cost of fireworks ai serverless inference models. """ from datetime import datetime -from typing import ( - Final, - cast, # noqa: TID251 # the fallback entry is a dict copy of a ReadOnly TypedDict; no cast-free way to retype it -) +from typing import Final from litellm.constants import ( FIREWORKS_AI_4_B, @@ -67,28 +64,6 @@ def _resolve_model_info(model: str) -> ModelInfo: return get_model_info(model=base_model, custom_llm_provider="fireworks_ai") -def _with_cache_read_fallback(model_info: ModelInfo) -> ModelInfo: - """Entries without a cache-read rate keep the previous calculator's input-rate fallback for cached - reads (LIT-7845 tracks the documented discount); the shared map is never mutated, so a copy carries it.""" - input_rate: Final = model_info.get("input_cost_per_token") - if model_info.get("cache_read_input_token_cost") is not None or input_rate is None: - return model_info - off_peak: Final = model_info.get("off_peak_pricing") - if off_peak is None or "cache_read_input_token_cost" in off_peak: - return cast(ModelInfo, {**model_info, "cache_read_input_token_cost": input_rate}) - return cast( - ModelInfo, - { - **model_info, - "cache_read_input_token_cost": input_rate, - "off_peak_pricing": { - **off_peak, - "cache_read_input_token_cost": off_peak.get("input_cost_per_token", input_rate), - }, - }, - ) - - def cost_per_token(model: str, usage: Usage, current_time: datetime | None = None) -> tuple[float, float]: """ Calculates the cost per token for a given model, prompt tokens, and completion tokens, @@ -102,7 +77,7 @@ def cost_per_token(model: str, usage: Usage, current_time: datetime | None = Non Returns: Tuple[float, float] - prompt_cost_in_usd, completion_cost_in_usd """ - model_info: Final = _with_cache_read_fallback(_resolve_model_info(model)) + model_info: Final = _resolve_model_info(model) return generic_cost_per_token( model=model, usage=usage, diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py index 686a792fa0f..32d7dd1d0c2 100644 --- a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py +++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py @@ -1,4 +1,5 @@ from collections.abc import Mapping +from copy import deepcopy from datetime import datetime, timezone import pytest @@ -15,6 +16,7 @@ from litellm.litellm_core_utils.llm_cost_calc.utils import ( _is_off_peak, _is_within_off_peak_window, apply_off_peak_pricing, + apply_provider_cache_read_default, calculate_cache_writing_cost, generic_cost_per_token, get_billed_token_rates, @@ -96,6 +98,38 @@ def test_generic_cost_per_token_bills_cache_reads_at_input_rate_when_no_cache_re assert completion_cost == pytest.approx(380 * 9.7e-7) +def test_apply_provider_cache_read_default_preserves_identity_and_input_data() -> None: + openai_info: ModelInfo = {"input_cost_per_token": 2e-6} + explicit_fireworks_info: ModelInfo = { + "input_cost_per_token": 2e-6, + "cache_read_input_token_cost": 1e-6, + } + fireworks_info: ModelInfo = { + "input_cost_per_token": 2e-6, + "off_peak_pricing": { + "hours_utc": "14:00-00:00", + "input_cost_per_token": 1e-6, + "output_cost_per_token": 3e-6, + }, + } + original_fireworks_info: ModelInfo = deepcopy(fireworks_info) + + assert apply_provider_cache_read_default(openai_info, "openai") is openai_info + assert apply_provider_cache_read_default(explicit_fireworks_info, "fireworks_ai") is explicit_fireworks_info + + processed_fireworks_info = apply_provider_cache_read_default(fireworks_info, "fireworks_ai") + + assert fireworks_info == original_fireworks_info + assert processed_fireworks_info is not fireworks_info + assert processed_fireworks_info["cache_read_input_token_cost"] == pytest.approx(2e-6 * 0.5) + assert processed_fireworks_info["off_peak_pricing"] == { + "hours_utc": "14:00-00:00", + "input_cost_per_token": 1e-6, + "output_cost_per_token": 3e-6, + "cache_read_input_token_cost": 1e-6 * 0.5, + } + + def test_generic_cost_per_token_prefers_audio_per_second_rate() -> None: model_info: ModelInfo = { "key": "gemini-embedding-2", @@ -239,9 +273,7 @@ def test_reasoning_tokens_no_price_set(_local_model_cost_map): model_cost_map["input_cost_per_token"] * usage.prompt_tokens, 10, ) - print(f"completion_cost: {completion_cost}") expected_completion_cost = model_cost_map["output_cost_per_token"] * usage.completion_tokens - print(f"expected_completion_cost: {expected_completion_cost}") assert round(completion_cost, 10) == round( expected_completion_cost, 10, diff --git a/tests/test_litellm/llms/fireworks_ai/test_fireworks_ai_cost_calculator.py b/tests/test_litellm/llms/fireworks_ai/test_fireworks_ai_cost_calculator.py index 1bee310d9d3..52222f22a51 100644 --- a/tests/test_litellm/llms/fireworks_ai/test_fireworks_ai_cost_calculator.py +++ b/tests/test_litellm/llms/fireworks_ai/test_fireworks_ai_cost_calculator.py @@ -5,6 +5,11 @@ from typing import Final import pytest import litellm +from litellm.litellm_core_utils.llm_cost_calc.utils import ( + calculate_prompt_caching_savings, + generic_cost_per_token, + get_token_type_cost_breakdown, +) from litellm.llms.fireworks_ai.cost_calculator import cost_per_token from litellm.types.utils import ( CompletionTokensDetailsWrapper, @@ -48,11 +53,13 @@ STANDARD_CACHE_READ_COST = 1.5e-08 def _register_off_peak_model( - off_peak_pricing: OffPeakPricing, cache_read_cost: float | None = STANDARD_CACHE_READ_COST + off_peak_pricing: OffPeakPricing, + cache_read_cost: float | None = STANDARD_CACHE_READ_COST, + model: str = OFF_PEAK_MODEL, ) -> None: - litellm.model_cost = { # test-quality-ok: the save/restore conftest returns litellm.model_cost to the original object after each test, so replacing the map for this entry leaks nothing + litellm.model_cost = { # test-quality-ok: conftest restores litellm.model_cost after each test **litellm.model_cost, - f"fireworks_ai/{OFF_PEAK_MODEL}": { + f"fireworks_ai/{model}": { "litellm_provider": "fireworks_ai", "mode": "chat", "input_cost_per_token": STANDARD_INPUT_COST, @@ -103,9 +110,8 @@ def test_off_peak_rates_left_unset_keep_the_standard_rates(): assert math.isclose(completion_cost, 200 * STANDARD_OUTPUT_COST, rel_tol=1e-10) -def test_off_peak_window_bills_cached_tokens_at_the_off_peak_input_rate_without_a_cache_read_rate(): - """Most fireworks_ai price-map entries carry no cache_read_input_token_cost, so cached tokens - fall back to the input rate, and inside the window that has to be the off-peak one.""" +def test_off_peak_window_bills_cached_tokens_at_the_discounted_off_peak_input_rate_without_a_cache_read_rate(): + """Entries without a cache-read rate use Fireworks' documented 50% cached-token discount.""" _register_off_peak_model( {"hours_utc": OFF_PEAK_WINDOW, "input_cost_per_token": 1e-08, "output_cost_per_token": 2e-08}, cache_read_cost=None, @@ -114,12 +120,116 @@ def test_off_peak_window_bills_cached_tokens_at_the_off_peak_input_rate_without_ prompt_cost, completion_cost = cost_per_token(model=OFF_PEAK_MODEL, usage=usage, current_time=INSIDE_WINDOW) - assert math.isclose(prompt_cost, 1000 * 1e-08, rel_tol=1e-10) + assert math.isclose(prompt_cost, (700 * 1e-08) + (300 * 1e-08 * 0.5), rel_tol=1e-10) assert math.isclose(completion_cost, 200 * 2e-08, rel_tol=1e-10) peak_prompt_cost, _ = cost_per_token(model=OFF_PEAK_MODEL, usage=usage, current_time=OUTSIDE_WINDOW) - assert math.isclose(peak_prompt_cost, 1000 * STANDARD_INPUT_COST, rel_tol=1e-10) + assert math.isclose( + peak_prompt_cost, + (700 * STANDARD_INPUT_COST) + (300 * STANDARD_INPUT_COST * 0.5), + rel_tol=1e-10, + ) + + no_input_rate_model = "accounts/fireworks/models/off-peak-no-input-rate-test" + _register_off_peak_model( + {"hours_utc": OFF_PEAK_WINDOW, "output_cost_per_token": 2e-08}, + cache_read_cost=None, + model=no_input_rate_model, + ) + + standard_cache_prompt_cost, _ = cost_per_token(model=no_input_rate_model, usage=usage, current_time=INSIDE_WINDOW) + + assert math.isclose( + standard_cache_prompt_cost, + (700 * STANDARD_INPUT_COST) + (300 * STANDARD_INPUT_COST * 0.5), + rel_tol=1e-10, + ) + + +def test_an_entry_without_a_cache_read_rate_bills_cached_tokens_at_the_documented_default_discount(): + """Fireworks documents a default 50% cached-token discount for serverless models: + https://docs.fireworks.ai/guides/prompt-caching, accessed 2026-09-19.""" + model = "accounts/fireworks/models/default-cache-read-test" + litellm.model_cost = { # test-quality-ok: conftest restores litellm.model_cost after each test + **litellm.model_cost, + f"fireworks_ai/{model}": { + "litellm_provider": "fireworks_ai", + "mode": "chat", + "input_cost_per_token": INPUT_COST, + "output_cost_per_token": OUTPUT_COST, + }, + } + usage = _usage(prompt_tokens=1000, cached_tokens=300, completion_tokens=200) + + prompt_cost, completion_cost = cost_per_token(model=model, usage=usage) + + assert math.isclose(prompt_cost, (700 * INPUT_COST) + (300 * INPUT_COST * 0.5), rel_tol=1e-10) + assert prompt_cost < 1000 * INPUT_COST + assert math.isclose(completion_cost, 200 * OUTPUT_COST, rel_tol=1e-10) + + +def test_fireworks_cache_read_rates_match_breakdown_and_caching_savings(): + model = "accounts/fireworks/models/breakdown-cache-read-test" + litellm.model_cost = { # test-quality-ok: the save/restore conftest returns litellm.model_cost to the original object after each test, so replacing the map for this entry leaks nothing + **litellm.model_cost, + f"fireworks_ai/{model}": { + "litellm_provider": "fireworks_ai", + "mode": "chat", + "input_cost_per_token": INPUT_COST, + "output_cost_per_token": OUTPUT_COST, + }, + } + usage = _usage(prompt_tokens=1000, cached_tokens=300, completion_tokens=200) + + breakdown = get_token_type_cost_breakdown( + model=model, + custom_llm_provider="fireworks_ai", + usage=usage, + ) + prompt_cost, _ = cost_per_token(model=model, usage=usage) + savings = calculate_prompt_caching_savings( + model_info=litellm.get_model_info(model=model, custom_llm_provider="fireworks_ai"), + usage=usage, + custom_llm_provider="fireworks_ai", + ) + + assert math.isclose(breakdown.cache_read_cost, 300 * INPUT_COST * 0.5, rel_tol=1e-10) + assert math.isclose(breakdown.rates.cache_read_input_token_cost, INPUT_COST * 0.5, rel_tol=1e-10) + assert math.isclose( + (700 * breakdown.rates.input_cost_per_token) + breakdown.cache_read_cost, prompt_cost, rel_tol=1e-10 + ) + assert math.isclose(savings, 300 * INPUT_COST * 0.5, rel_tol=1e-10) + + +def test_generic_cost_per_token_applies_fireworks_cache_read_default_with_or_without_model_info(): + model = "accounts/fireworks/models/generic-cache-read-test" + litellm.model_cost = { # test-quality-ok: conftest restores litellm.model_cost after each test + **litellm.model_cost, + f"fireworks_ai/{model}": { + "litellm_provider": "fireworks_ai", + "mode": "chat", + "input_cost_per_token": INPUT_COST, + "output_cost_per_token": OUTPUT_COST, + }, + } + usage = _usage(prompt_tokens=1000, cached_tokens=300, completion_tokens=200) + expected_prompt_cost = (700 * INPUT_COST) + (300 * INPUT_COST * 0.5) + + implicit_model_info_cost, _ = generic_cost_per_token( + model=model, + usage=usage, + custom_llm_provider="fireworks_ai", + ) + explicit_model_info_cost, _ = generic_cost_per_token( + model=model, + usage=usage, + custom_llm_provider="fireworks_ai", + model_info=litellm.get_model_info(model=model, custom_llm_provider="fireworks_ai"), + ) + + assert math.isclose(implicit_model_info_cost, expected_prompt_cost, rel_tol=1e-10) + assert math.isclose(explicit_model_info_cost, expected_prompt_cost, rel_tol=1e-10) def test_off_peak_defaults_to_the_current_time(): From cc2db3887107a716e931a31a2f81aec302559ab7 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 18 Sep 2026 18:07:20 -0700 Subject: [PATCH 118/144] chore(proxy): keep the OpenAPI snapshot as CI generates it --- litellm/proxy/_lazy_openapi_snapshot.json | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index 80527f50d10..d609016f442 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -19400,7 +19400,7 @@ } } }, - "description": "\nUnified rate-limit error.\n\nEvery rate-limit condition surfaced by litellm \u2014 whether it originated from\nan upstream LLM provider, a vendor batch endpoint, or one of litellm's own\nproxy-side limiters (parallel-requests, dynamic-rate, batch-rate, budget,\nmax-iterations, etc.) \u2014 is raised as an instance of this class.\n\nThe :attr:`category` attribute lets callers distinguish the source. See\n:class:`RateLimitErrorCategory` for the available values.\n" + "description": "\n Unified rate-limit error.\n\n Every rate-limit condition surfaced by litellm \u2014 whether it originated from\n an upstream LLM provider, a vendor batch endpoint, or one of litellm's own\n proxy-side limiters (parallel-requests, dynamic-rate, batch-rate, budget,\n max-iterations, etc.) \u2014 is raised as an instance of this class.\n\n The :attr:`category` attribute lets callers distinguish the source. See\n :class:`RateLimitErrorCategory` for the available values.\n " }, "500": { "content": { From ca8062e506135081e1d7805c23780be40a6b3f6b Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 18 Sep 2026 18:07:21 -0700 Subject: [PATCH 119/144] fix(claude_code_gateway): keep the device secret out of the browser URL and validate the login before claiming it --- .../anthropic_endpoints/gateway_endpoints.py | 117 +++++++++++------- .../test_gateway_endpoints.py | 97 ++++++++++++--- 2 files changed, 146 insertions(+), 68 deletions(-) diff --git a/litellm/proxy/anthropic_endpoints/gateway_endpoints.py b/litellm/proxy/anthropic_endpoints/gateway_endpoints.py index 08579186f5e..259d5202db6 100644 --- a/litellm/proxy/anthropic_endpoints/gateway_endpoints.py +++ b/litellm/proxy/anthropic_endpoints/gateway_endpoints.py @@ -18,6 +18,7 @@ import hashlib import json import secrets from collections.abc import Mapping +from dataclasses import dataclass from types import MappingProxyType from typing import Final @@ -32,13 +33,16 @@ from litellm.constants import ( CLI_SSO_SESSION_TTL_SECONDS, LITELLM_CLI_SOURCE_IDENTIFIER, ) +from litellm.proxy._types import LiteLLM_UserTable, LitellmUserRoles from litellm.proxy.anthropic_endpoints.endpoints import anthropic_response, count_tokens from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_utils.http_parsing_utils import _safe_set_request_parsed_body +from litellm.proxy.management_endpoints.ui_sso import CliSsoTeamDetail GATEWAY_PREFIX: Final = "/claude_code_gateway" _DEVICE_CODE_GRANT: Final = "urn:ietf:params:oauth:grant-type:device_code" _REFRESH_TOKEN_GRANT: Final = "refresh_token" +_DEVICE_CODE_SEPARATOR: Final = "." _DEVICE_POLL_INTERVAL_SECONDS: Final = 5 _SECONDS_PER_HOUR: Final = 3600 _MANAGED_SETTINGS_ADAPTER: Final = TypeAdapter(dict[str, object]) @@ -48,12 +52,19 @@ _POST_ONLY: Final = ["POST"] # mutable-ok: FastAPI's add_api_route only accepts class _GatewaySessionData(BaseModel): user_id: str - user_role: str | None + user_role: LitellmUserRoles models: list[str] = Field(default_factory=list) teams: tuple[str, ...] = () team_details: object | None = None +@dataclass(frozen=True, slots=True) +class _GatewayLogin: + user_info: LiteLLM_UserTable + team_id: str | None + team: CliSsoTeamDetail + + class _OAuthErrorBody(BaseModel): error: str error_description: str | None = None @@ -70,7 +81,7 @@ class _DeviceAuthorizationBody(BaseModel): device_code: str user_code: str verification_uri: str - verification_uri_complete: str + verification_uri_complete: str | None = None expires_in: int interval: int @@ -111,15 +122,11 @@ def _managed_settings() -> dict[str, object] | None: return _MANAGED_SETTINGS_ADAPTER.validate_python(settings) -def _oauth_error(*, status_code: int, error: str, description: str | None = None) -> "_OAuthError": - return _OAuthError(status_code=status_code, error=error, description=description) - - -class _OAuthError(Exception): - def __init__(self, *, status_code: int, error: str, description: str | None) -> None: - self.status_code = status_code - self.error = error - self.description = description +@dataclass(frozen=True, slots=True) +class _OAuthError: + status_code: int + error: str + description: str | None = None def _oauth_error_response(err: _OAuthError) -> JSONResponse: @@ -153,7 +160,7 @@ router.add_api_route( @router.get("/.well-known/oauth-authorization-server", include_in_schema=False) async def oauth_authorization_server(request: Request) -> JSONResponse: if not _is_gateway_enabled(): - return _oauth_error_response(_oauth_error(status_code=404, error="not_found")) + return _oauth_error_response(_OAuthError(status_code=404, error="not_found")) from litellm.proxy.utils import get_custom_url @@ -175,6 +182,7 @@ async def device_authorization(request: Request) -> JSONResponse: from litellm.proxy.management_endpoints.ui_sso import ( _check_cli_sso_start_rate_limit, # pyright: ignore[reportPrivateUsage] # shared device-flow helper + _cli_sso_verification_uri_complete_enabled, # pyright: ignore[reportPrivateUsage] # shared device-flow helper _generate_cli_sso_user_code, # pyright: ignore[reportPrivateUsage] # shared device-flow helper _hash_cli_sso_secret, # pyright: ignore[reportPrivateUsage] # shared device-flow helper _normalize_cli_sso_user_code, # pyright: ignore[reportPrivateUsage] # shared device-flow helper @@ -184,7 +192,7 @@ async def device_authorization(request: Request) -> JSONResponse: from litellm.proxy.utils import get_custom_url if not _is_gateway_enabled(): - return _oauth_error_response(_oauth_error(status_code=404, error="not_found")) + return _oauth_error_response(_OAuthError(status_code=404, error="not_found")) _check_cli_sso_start_rate_limit( request=request, @@ -192,50 +200,51 @@ async def device_authorization(request: Request) -> JSONResponse: use_x_forwarded_for=bool(_general_settings().get("use_x_forwarded_for", False)), ) - device_code: Final = f"cli-{secrets.token_urlsafe(24)}" + login_id: Final = f"cli-{secrets.token_urlsafe(24)}" + poll_secret: Final = secrets.token_urlsafe(32) user_code: Final = _generate_cli_sso_user_code() flow: Final = { # mutable-ok: the shared CLI SSO cache entry is a dict the browser leg mutates - "poll_secret_hash": _hash_cli_sso_secret(device_code), + "poll_secret_hash": _hash_cli_sso_secret(poll_secret), "user_code_hash": _hash_cli_sso_secret(_normalize_cli_sso_user_code(user_code)), "sso_complete": False, "user_code_verified": False, "session_data": None, } - _set_cli_sso_flow(login_id=device_code, cache=cli_sso_session_cache, flow=flow) + _set_cli_sso_flow(login_id=login_id, cache=cli_sso_session_cache, flow=flow) request_base_url: Final = str(request.base_url) verification_uri: Final = get_custom_url(request_base_url=request_base_url, route="sso/key/generate") - query: Final = MappingProxyType({"source": LITELLM_CLI_SOURCE_IDENTIFIER, "key": device_code}) + query: Final = MappingProxyType({"source": LITELLM_CLI_SOURCE_IDENTIFIER, "key": login_id}) body: Final = _DeviceAuthorizationBody( - device_code=device_code, + device_code=f"{login_id}{_DEVICE_CODE_SEPARATOR}{poll_secret}", user_code=user_code, verification_uri=f"{verification_uri}?{urlencode(query)}", verification_uri_complete=( f"{verification_uri}?{urlencode(MappingProxyType({**query, 'user_code': user_code}))}" + if _cli_sso_verification_uri_complete_enabled() + else None ), expires_in=CLI_SSO_SESSION_TTL_SECONDS, interval=_DEVICE_POLL_INTERVAL_SECONDS, ) - return JSONResponse(content=body.model_dump()) + return JSONResponse(content=body.model_dump(exclude_none=True)) -def _mint_access_token_from_flow(flow: Mapping[str, object]) -> str: - from litellm.proxy._types import LiteLLM_UserTable - from litellm.proxy.auth.auth_checks import ExperimentalUIJWTToken +def _validate_login(flow: Mapping[str, object]) -> _GatewayLogin | _OAuthError: from litellm.proxy.management_endpoints.ui_sso import selected_cli_sso_team_detail try: session_data: Final = _GatewaySessionData.model_validate(flow.get("session_data")) except ValidationError as err: verbose_proxy_logger.warning("Claude Code gateway login session is malformed: %s", err) - raise _oauth_error( + return _OAuthError( status_code=400, error="invalid_grant", description="The login session is malformed; sign in again" - ) from err + ) team_id: Final = session_data.teams[0] if session_data.teams else None selected_team: Final = selected_cli_sso_team_detail(team_details=session_data.team_details, team_id=team_id) if selected_team is None: - raise _oauth_error( + return _OAuthError( status_code=400, error="invalid_grant", description=f"Could not resolve the model grants for team {team_id}; sign in again", @@ -243,26 +252,32 @@ def _mint_access_token_from_flow(flow: Mapping[str, object]) -> str: user_info: Final = LiteLLM_UserTable( user_id=session_data.user_id, - user_role=session_data.user_role, + user_role=session_data.user_role.value, models=session_data.models, ) + return _GatewayLogin(user_info=user_info, team_id=team_id, team=selected_team) + + +def _mint_access_token(login: _GatewayLogin) -> str: + from litellm.proxy.auth.auth_checks import ExperimentalUIJWTToken + return ExperimentalUIJWTToken.get_cli_jwt_auth_token( - user_info=user_info, - team_id=team_id, - team_alias=selected_team.team_alias, - team_models=selected_team.team_models, - team_model_aliases=selected_team.team_model_aliases, + user_info=login.user_info, + team_id=login.team_id, + team_alias=login.team.team_alias, + team_models=login.team.team_models, + team_model_aliases=login.team.team_model_aliases, max_budget=None, ) -async def _claim_device_code(device_code: str, cache: DualCache) -> bool: +async def _claim_device_code(login_id: str, cache: DualCache) -> bool: from litellm.proxy.management_endpoints.ui_sso import ( _get_cli_sso_flow_cache_key, # pyright: ignore[reportPrivateUsage] # shared device-flow helper ) claims: Final = await cache.async_increment_cache( - key=f"{_get_cli_sso_flow_cache_key(device_code)}:claimed", + key=f"{_get_cli_sso_flow_cache_key(login_id)}:claimed", value=1, ttl=CLI_SSO_SESSION_TTL_SECONDS, ) @@ -275,39 +290,45 @@ async def _handle_device_code_grant(device_code: str | None) -> JSONResponse: from litellm.proxy.management_endpoints.ui_sso import ( _get_cli_sso_flow_cache_key, # pyright: ignore[reportPrivateUsage] # shared device-flow helper _get_cli_sso_flow_or_raise, # pyright: ignore[reportPrivateUsage] # shared device-flow helper + _verify_cli_sso_poll_secret, # pyright: ignore[reportPrivateUsage] # shared device-flow helper ) from litellm.proxy.proxy_server import cli_sso_session_cache if not device_code: return _oauth_error_response( - _oauth_error(status_code=400, error="invalid_request", description="device_code is required") + _OAuthError(status_code=400, error="invalid_request", description="device_code is required") ) + login_id, _, poll_secret = device_code.partition(_DEVICE_CODE_SEPARATOR) try: - flow: Final = _get_cli_sso_flow_or_raise(login_id=device_code, cache=cli_sso_session_cache) + flow: Final = _get_cli_sso_flow_or_raise(login_id=login_id, cache=cli_sso_session_cache) except HTTPException: - return _oauth_error_response(_oauth_error(status_code=400, error="expired_token")) + return _oauth_error_response(_OAuthError(status_code=400, error="expired_token")) + + if not _verify_cli_sso_poll_secret(flow, poll_secret): + return _oauth_error_response(_OAuthError(status_code=400, error="expired_token")) if not flow.get("sso_complete") or not flow.get("user_code_verified"): - return _oauth_error_response(_oauth_error(status_code=400, error="authorization_pending")) + return _oauth_error_response(_OAuthError(status_code=400, error="authorization_pending")) - if not await _claim_device_code(device_code, cli_sso_session_cache): - return _oauth_error_response(_oauth_error(status_code=400, error="expired_token")) + login: Final = _validate_login(flow) + if isinstance(login, _OAuthError): + return _oauth_error_response(login) - await cli_sso_session_cache.async_delete_cache(key=_get_cli_sso_flow_cache_key(device_code)) - try: - access_token: Final = _mint_access_token_from_flow(flow) - except _OAuthError as err: - return _oauth_error_response(err) + if not await _claim_device_code(login_id, cli_sso_session_cache): + return _oauth_error_response(_OAuthError(status_code=400, error="expired_token")) - body: Final = _AccessTokenBody(access_token=access_token, expires_in=CLI_JWT_EXPIRATION_HOURS * _SECONDS_PER_HOUR) + await cli_sso_session_cache.async_delete_cache(key=_get_cli_sso_flow_cache_key(login_id)) + body: Final = _AccessTokenBody( + access_token=_mint_access_token(login), expires_in=CLI_JWT_EXPIRATION_HOURS * _SECONDS_PER_HOUR + ) return JSONResponse(content=body.model_dump()) @router.post("/oauth/token", include_in_schema=False) async def oauth_token(request: Request) -> JSONResponse: if not _is_gateway_enabled(): - return _oauth_error_response(_oauth_error(status_code=404, error="not_found")) + return _oauth_error_response(_OAuthError(status_code=404, error="not_found")) form: Final = await request.form() grant_type: Final = form.get("grant_type") @@ -318,7 +339,7 @@ async def oauth_token(request: Request) -> JSONResponse: if grant_type == _REFRESH_TOKEN_GRANT: return _oauth_error_response( - _oauth_error( + _OAuthError( status_code=401, error="invalid_grant", description="This gateway does not issue refresh tokens; sign in again", @@ -326,7 +347,7 @@ async def oauth_token(request: Request) -> JSONResponse: ) return _oauth_error_response( - _oauth_error( + _OAuthError( status_code=400, error="unsupported_grant_type", description=f"Unsupported grant_type: {grant_type}" ) ) diff --git a/tests/test_litellm/proxy/anthropic_endpoints/test_gateway_endpoints.py b/tests/test_litellm/proxy/anthropic_endpoints/test_gateway_endpoints.py index e49047634bc..d442ac21307 100644 --- a/tests/test_litellm/proxy/anthropic_endpoints/test_gateway_endpoints.py +++ b/tests/test_litellm/proxy/anthropic_endpoints/test_gateway_endpoints.py @@ -20,11 +20,18 @@ from fastapi.testclient import TestClient from litellm.caching.dual_cache import DualCache from litellm.proxy._types import ProxyException from litellm.proxy.anthropic_endpoints import gateway_endpoints -from litellm.proxy.management_endpoints.ui_sso import _get_cli_sso_flow_cache_key, _set_cli_sso_flow +from litellm.proxy.management_endpoints.ui_sso import ( + _get_cli_sso_flow_cache_key, + _hash_cli_sso_secret, + _set_cli_sso_flow, +) from litellm.proxy.middleware.prometheus_auth_middleware import PrometheusAuthMiddleware _DEVICE_CODE_GRANT: Final = "urn:ietf:params:oauth:grant-type:device_code" _MASTER_KEY: Final = "sk-master-key" +_SHARED_LOGIN_ID: Final = "cli-shared-login-code" +_SHARED_POLL_SECRET: Final = "shared-poll-secret" +_SHARED_DEVICE_CODE: Final = f"{_SHARED_LOGIN_ID}.{_SHARED_POLL_SECRET}" _MINT: Final = "litellm.proxy.auth.auth_checks.ExperimentalUIJWTToken.get_cli_jwt_auth_token" _PROTOBUF_BODY: Final = b"\x0a\x05hello\x12\x03{{{" _COMPLETED_SESSION: Final = MappingProxyType( @@ -100,10 +107,12 @@ def _gateway_env( managed_settings: Mapping[str, object] | None = None, cache: DualCache | None = None, real_auth: bool = False, + extra_settings: Mapping[str, object] = MappingProxyType({}), ) -> Iterator[tuple[TestClient, DualCache]]: general_settings: Final = { "enable_claude_code_gateway": enabled, **({} if managed_settings is None else {"claude_code_gateway_managed_settings": dict(managed_settings)}), + **extra_settings, } session_cache: Final = cache or DualCache(default_in_memory_ttl=600) @@ -147,7 +156,7 @@ def _request_token(client: TestClient, device_code: str) -> httpx.Response: def _completed_flow(session_data: Mapping[str, object] = _COMPLETED_SESSION) -> dict[str, object]: return { - "poll_secret_hash": "unused", + "poll_secret_hash": _hash_cli_sso_secret(_SHARED_POLL_SECRET), "user_code_hash": "unused", "sso_complete": True, "user_code_verified": True, @@ -155,13 +164,18 @@ def _completed_flow(session_data: Mapping[str, object] = _COMPLETED_SESSION) -> } +def _login_id(device_code: str) -> str: + return device_code.partition(".")[0] + + def _complete_flow( cache: DualCache, device_code: str, session_data: Mapping[str, object] = _COMPLETED_SESSION ) -> None: - key: Final = _get_cli_sso_flow_cache_key(device_code) + key: Final = _get_cli_sso_flow_cache_key(_login_id(device_code)) flow: Final = cache.get_cache(key=key) assert isinstance(flow, dict) - cache.set_cache(key=key, value={**flow, **_completed_flow(session_data)}, ttl=600) + completed: Final = {**flow, **_completed_flow(session_data), "poll_secret_hash": flow["poll_secret_hash"]} + cache.set_cache(key=key, value=completed, ttl=600) def test_discovery_shape(): @@ -194,18 +208,35 @@ def test_device_authorization_returns_rfc8628_shape_and_persists_flow(): assert resp.status_code == 200 body = resp.json() device_code = body["device_code"] - assert device_code.startswith("cli-") + login_id, separator, poll_secret = device_code.partition(".") + assert login_id.startswith("cli-") + assert separator == "." + assert len(poll_secret) >= 32 assert body["user_code"] assert body["expires_in"] == 600 assert body["interval"] == 5 - # verification_uri_complete carries the user_code; the short uri does not. - assert f"user_code={body['user_code']}" in body["verification_uri_complete"] - assert "user_code=" not in body["verification_uri"] - assert f"key={device_code}" in body["verification_uri"] - # The device flow is stored under the device_code so the browser SSO leg can complete it. - stored = cache.get_cache(key=_get_cli_sso_flow_cache_key(device_code)) + assert "verification_uri_complete" not in body + assert body["verification_uri"].endswith(f"/sso/key/generate?source=litellm-cli&key={login_id}") + assert poll_secret not in body["verification_uri"] + stored = cache.get_cache(key=_get_cli_sso_flow_cache_key(login_id)) assert isinstance(stored, dict) assert stored["sso_complete"] is False + assert stored["poll_secret_hash"] == _hash_cli_sso_secret(poll_secret) + assert cache.get_cache(key=_get_cli_sso_flow_cache_key(device_code)) is None + + +@pytest.mark.parametrize("opted_in", [True, False]) +def test_verification_uri_complete_carries_the_user_code_only_when_the_operator_opts_in(opted_in: bool): + with _gateway_env(extra_settings={"allow_cli_sso_verification_uri_complete": opted_in}) as (client, _): + body = client.post("/claude_code_gateway/oauth/device_authorization").json() + login_id = _login_id(body["device_code"]) + if not opted_in: + assert "verification_uri_complete" not in body + return + assert body["verification_uri_complete"].endswith( + f"/sso/key/generate?source=litellm-cli&key={login_id}&user_code={body['user_code']}" + ) + assert "user_code=" not in body["verification_uri"] def test_token_authorization_pending_before_browser_completes(): @@ -215,6 +246,22 @@ def test_token_authorization_pending_before_browser_completes(): assert resp.json()["error"] == "authorization_pending" +@pytest.mark.parametrize("tamper", ["login_id_only", "wrong_secret"]) +def test_token_refuses_the_browser_login_id_without_the_client_secret(tamper: str): + with _gateway_env() as (client, cache): + device_code = _start_device_flow(client) + _complete_flow(cache, device_code) + login_id = _login_id(device_code) + presented = login_id if tamper == "login_id_only" else f"{login_id}.not-the-secret" + with patch(_MINT, return_value="sk-session") as mint: + resp = _request_token(client, presented) + assert resp.status_code == 400 + assert resp.json()["error"] == "expired_token" + mint.assert_not_called() + with_secret = _request_token(client, device_code) + assert with_secret.status_code == 200 + + def test_token_success_mints_bearer_and_is_single_use(): with _gateway_env() as (client, cache): device_code = _start_device_flow(client) @@ -251,14 +298,25 @@ def test_token_teamless_user_mints_without_a_team(): assert mint.call_args.kwargs["team_models"] == () -def test_token_malformed_session_is_invalid_grant(): +@pytest.mark.parametrize( + "session_data", + [ + {"user_role": "internal_user"}, + {**_COMPLETED_SESSION, "user_role": None}, + {**_COMPLETED_SESSION, "user_role": "not-a-role"}, + ], + ids=["missing_user_id", "no_role", "unknown_role"], +) +def test_token_malformed_session_is_invalid_grant_and_does_not_consume_the_login(session_data: Mapping[str, object]): with _gateway_env() as (client, cache): device_code = _start_device_flow(client) - _complete_flow(cache, device_code, session_data={"user_role": "internal_user"}) + _complete_flow(cache, device_code, session_data=session_data) with patch(_MINT) as mint: resp = _request_token(client, device_code) + again = _request_token(client, device_code) assert resp.status_code == 400 assert resp.json()["error"] == "invalid_grant" + assert again.json()["error"] == "invalid_grant" mint.assert_not_called() @@ -275,25 +333,24 @@ def test_token_unknown_team_grants_is_invalid_grant(): def test_token_mints_on_a_replica_that_did_not_start_the_login(): redis: Final = _SharedRedisFake() - device_code: Final = "cli-shared-login-code" - _set_cli_sso_flow(login_id=device_code, cache=_replica(redis), flow=_completed_flow()) + _set_cli_sso_flow(login_id=_SHARED_LOGIN_ID, cache=_replica(redis), flow=_completed_flow()) with _gateway_env(cache=_replica(redis)) as (client, _), patch(_MINT, return_value="sk-session") as mint: - resp = _request_token(client, device_code) + resp = _request_token(client, _SHARED_DEVICE_CODE) assert resp.status_code == 200 assert resp.json()["access_token"] == "sk-session" assert mint.call_args.kwargs["team_id"] == "team-a" + assert mint.call_args.kwargs["user_info"].user_role == "internal_user" def test_token_refuses_a_device_code_another_replica_already_claimed(): redis: Final = _SharedRedisFake() replica_a: Final = _replica(redis) - device_code: Final = "cli-shared-login-code" - _set_cli_sso_flow(login_id=device_code, cache=replica_a, flow=_completed_flow()) - assert asyncio.run(gateway_endpoints._claim_device_code(device_code, replica_a)) is True + _set_cli_sso_flow(login_id=_SHARED_LOGIN_ID, cache=replica_a, flow=_completed_flow()) + assert asyncio.run(gateway_endpoints._claim_device_code(_SHARED_LOGIN_ID, replica_a)) is True with _gateway_env(cache=_replica(redis)) as (client, _), patch(_MINT) as mint: - resp = _request_token(client, device_code) + resp = _request_token(client, _SHARED_DEVICE_CODE) assert resp.status_code == 400 assert resp.json()["error"] == "expired_token" mint.assert_not_called() From 157fa589478f37c0e5fbcf4c92b933bfaddfa4b8 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Fri, 18 Sep 2026 18:09:14 -0700 Subject: [PATCH 120/144] fix(rust): leave calls with a custom URL policy on the Python route litellm.user_url_validation and litellm.user_url_allowed_hosts are only implemented by the Python document fetcher, so an allowlisted internal document was rejected by the Rust route's network policy. The bridge now declines when either is changed from its default --- .../crates/python-bridge/python_settings.json | 4 ++ litellm-rust/crates/python-bridge/src/http.rs | 49 +++++++++++++++++++ .../python-bridge/src/python_settings.rs | 4 +- litellm/rust_bridge/settings.py | 16 ++++++ .../test_litellm/rust_bridge/test_settings.py | 15 +++++- 5 files changed, 86 insertions(+), 2 deletions(-) diff --git a/litellm-rust/crates/python-bridge/python_settings.json b/litellm-rust/crates/python-bridge/python_settings.json index 40e36a900d3..a6f5ee9c6f4 100644 --- a/litellm-rust/crates/python-bridge/python_settings.json +++ b/litellm-rust/crates/python-bridge/python_settings.json @@ -10,5 +10,9 @@ "disable_aiohttp_trust_env", "disable_aiohttp_transport", "user_agent" + ], + "url_policy": [ + "user_url_validation", + "user_url_allowed_hosts" ] } diff --git a/litellm-rust/crates/python-bridge/src/http.rs b/litellm-rust/crates/python-bridge/src/http.rs index 2ab6517b61d..77542855fec 100644 --- a/litellm-rust/crates/python-bridge/src/http.rs +++ b/litellm-rust/crates/python-bridge/src/http.rs @@ -24,6 +24,7 @@ pub(crate) fn call_config( asynchronous: bool, ) -> PyResult { decline_live_clients(kwargs)?; + decline_custom_url_policy(&PythonSettings::UrlPolicy.read(py)?)?; let configured = settings(&PythonSettings::Http.read(py)?)? .with_environment(&|name| std::env::var(name).ok()); let settings = for_call(configured, call_ssl_verify(kwargs)?, asynchronous) @@ -63,6 +64,23 @@ pub(crate) fn decline_live_clients(kwargs: &Bound<'_, PyDict>) -> PyResult<()> { Ok(()) } +#[derive(FromPyObject)] +struct PythonUrlPolicy { + user_url_validation: bool, + user_url_allowed_hosts: Vec, +} + +fn decline_custom_url_policy(value: &Bound<'_, PyAny>) -> PyResult<()> { + match value.extract::() { + Ok(policy) if policy.user_url_validation && policy.user_url_allowed_hosts.is_empty() => { + Ok(()) + } + Ok(_) | Err(_) => Err(RustBridgeDeclined::new_err( + "litellm.user_url_validation / user_url_allowed_hosts are applied by the Python route", + )), + } +} + #[derive(FromPyObject)] struct PythonHttpSettings<'py> { ssl_verify: Bound<'py, PyAny>, @@ -245,6 +263,37 @@ user_agent='litellm/9.9.9', }); } + fn url_policy<'py>(py: Python<'py>, fields: &str) -> Bound<'py, PyAny> { + let source = std::ffi::CString::new(format!( + "import types\npolicy = types.SimpleNamespace({fields})" + )) + .unwrap(); + let locals = PyDict::new(py); + py.run(&source, Some(&locals), Some(&locals)).unwrap(); + locals.get_item("policy").unwrap().unwrap() + } + + #[test] + fn default_url_policy_stays_on_the_rust_route() { + Python::initialize(); + Python::attach(|py| { + let policy = url_policy(py, "user_url_validation=True, user_url_allowed_hosts=[]"); + decline_custom_url_policy(&policy).unwrap(); + }); + } + + #[rstest] + #[case::validation_off("user_url_validation=False, user_url_allowed_hosts=[]")] + #[case::allowlist("user_url_validation=True, user_url_allowed_hosts=['docs.internal']")] + #[case::mistyped("user_url_validation=True, user_url_allowed_hosts=None")] + fn custom_url_policy_declines_so_python_applies_it(#[case] fields: &str) { + Python::initialize(); + Python::attach(|py| { + let error = decline_custom_url_policy(&url_policy(py, fields)).unwrap_err(); + assert!(error.is_instance_of::(py)); + }); + } + #[test] fn mistyped_python_settings_decline_instead_of_raising() { Python::initialize(); diff --git a/litellm-rust/crates/python-bridge/src/python_settings.rs b/litellm-rust/crates/python-bridge/src/python_settings.rs index dcb46e7d2b5..b7855566850 100644 --- a/litellm-rust/crates/python-bridge/src/python_settings.rs +++ b/litellm-rust/crates/python-bridge/src/python_settings.rs @@ -5,15 +5,17 @@ const MODULE: &str = "litellm.rust_bridge.settings"; #[derive(Clone, Copy, Debug, PartialEq, Eq)] pub(crate) enum PythonSettings { Http, + UrlPolicy, } impl PythonSettings { #[cfg(test)] - pub(crate) const ALL: [Self; 1] = [Self::Http]; + pub(crate) const ALL: [Self; 2] = [Self::Http, Self::UrlPolicy]; pub(crate) fn name(self) -> &'static str { match self { Self::Http => "http_settings", + Self::UrlPolicy => "url_policy", } } diff --git a/litellm/rust_bridge/settings.py b/litellm/rust_bridge/settings.py index 491312c97b6..bccfd01ec73 100644 --- a/litellm/rust_bridge/settings.py +++ b/litellm/rust_bridge/settings.py @@ -1,5 +1,6 @@ from __future__ import annotations +from collections.abc import Sequence from dataclasses import dataclass @@ -17,6 +18,21 @@ class HttpSettings: user_agent: str +@dataclass(frozen=True, slots=True) +class UrlPolicy: + user_url_validation: bool + user_url_allowed_hosts: Sequence[str] + + +def url_policy() -> UrlPolicy: + import litellm + + return UrlPolicy( + user_url_validation=litellm.user_url_validation, + user_url_allowed_hosts=litellm.user_url_allowed_hosts, + ) + + def http_settings() -> HttpSettings: import litellm from litellm.llms.custom_httpx.http_handler import default_user_agent diff --git a/tests/test_litellm/rust_bridge/test_settings.py b/tests/test_litellm/rust_bridge/test_settings.py index f4f9cbc8eec..7e7b1c6743b 100644 --- a/tests/test_litellm/rust_bridge/test_settings.py +++ b/tests/test_litellm/rust_bridge/test_settings.py @@ -15,7 +15,20 @@ CONTRACT_PATH: Final = Path(__file__).parents[3] / "litellm-rust/crates/python-b def test_the_rust_contract_matches_the_returned_fields() -> None: contract: Final = TypeAdapter(dict[str, list[str]]).validate_json(CONTRACT_PATH.read_text()) - assert contract == {"http_settings": [field.name for field in dataclasses.fields(settings.http_settings())]} + assert contract == { + "http_settings": [field.name for field in dataclasses.fields(settings.http_settings())], + "url_policy": [field.name for field in dataclasses.fields(settings.url_policy())], + } + + +def test_url_policy_reads_the_litellm_globals(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "user_url_validation", False) + monkeypatch.setattr(litellm, "user_url_allowed_hosts", ["docs.internal:8443"]) + + assert settings.url_policy() == settings.UrlPolicy( + user_url_validation=False, + user_url_allowed_hosts=["docs.internal:8443"], + ) def test_http_settings_reads_the_litellm_globals(monkeypatch: pytest.MonkeyPatch) -> None: From 17c519c40a3c98671eac5966eaed3ab9cb8b2261 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 18 Sep 2026 18:09:33 -0700 Subject: [PATCH 121/144] test(custom_httpx): pass the token resolver and two-argument client factory in the realtime bridge test --- .../llms/custom_httpx/test_llm_http_handler.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py index 1f71ffd43f6..82220d53375 100644 --- a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py +++ b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py @@ -3780,11 +3780,17 @@ async def test_async_realtime_bridges_a_transcription_session_through_the_provid yield script.pop(0) speech_client = FakeSpeechClient() + + async def resolve_access_token() -> str: + return "token" + provider_config = VertexChirpRealtimeConfig( - access_token="token", + resolve_access_token=resolve_access_token, project="proj-1", location="us", - backend_factory=lambda target: SpeechStreamingBackend(target, client_factory=lambda target: speech_client), + backend_factory=lambda target: SpeechStreamingBackend( + target, client_factory=lambda target, access_token: speech_client + ), ) audio = base64.b64encode(b"\x00\x01" * 800).decode() client_ws = _ScriptedClientWebSocket( From fcc7efa4db5701f271d812287d2d044d0ca1fb02 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 18 Sep 2026 18:10:26 -0700 Subject: [PATCH 122/144] fix(responses): forward the routed input and report routing rejections on the websocket --- .../proxy/response_api_endpoints/endpoints.py | 20 +++++- litellm/responses/main.py | 28 +++++++- .../response_api_endpoints/test_endpoints.py | 65 +++++++++++++++++++ .../test_responses_api_request_body.py | 64 ++++++++++++++++++ 4 files changed, 175 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/response_api_endpoints/endpoints.py b/litellm/proxy/response_api_endpoints/endpoints.py index 4b178c52de8..1b1fc466046 100644 --- a/litellm/proxy/response_api_endpoints/endpoints.py +++ b/litellm/proxy/response_api_endpoints/endpoints.py @@ -1395,6 +1395,15 @@ def _routing_hints_from_first_ws_frame(first_message: str) -> Mapping[str, objec return MappingProxyType({key: value for key, value in hints.items() if value is not None}) +def _responses_ws_failure_frame(failure: Exception) -> str: + raw_status: Final = getattr(failure, "status_code", None) + status: Final = raw_status if isinstance(raw_status, int) and not isinstance(raw_status, bool) else 500 + error_type: Final = ( + "rate_limit_exceeded" if status == 429 else "invalid_request_error" if 400 <= status < 500 else "server_error" + ) + return json.dumps({"type": "error", "status": status, "error": {"type": error_type, "message": str(failure)}}) + + async def _enforce_responses_ws_first_frame_model_auth( request: Request, model: str, @@ -1574,6 +1583,15 @@ async def responses_websocket_endpoint( original_exception=failure, request_data=data, ) - except Exception: + except Exception as e: verbose_proxy_logger.exception("Responses WebSocket error") + try: + await websocket.send_text(_responses_ws_failure_frame(e)) + except Exception: + pass + await proxy_logging_obj.post_call_failure_hook( + user_api_key_dict=user_api_key_dict, + original_exception=e, + request_data=data, + ) await websocket.close(code=1011, reason="Internal server error") diff --git a/litellm/responses/main.py b/litellm/responses/main.py index 85dec4f11e2..5a4a08b760c 100644 --- a/litellm/responses/main.py +++ b/litellm/responses/main.py @@ -1,5 +1,6 @@ import asyncio import contextvars +import json from collections.abc import Coroutine, Generator, Iterable, Mapping, Sequence from contextlib import contextmanager from dataclasses import dataclass @@ -8,7 +9,7 @@ from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, NoReturn, Optional, TypeAlias, cast import httpx -from pydantic import BaseModel, TypeAdapter +from pydantic import BaseModel, TypeAdapter, ValidationError from typing_extensions import assert_never import litellm @@ -2277,6 +2278,24 @@ def _deployment_reasoning_default(kwargs: Mapping[str, object]) -> Reasoning | d _RESPONSES_WS_ROUTING_HINT_KEYS: Final = frozenset({"input", "previous_response_id"}) +def _first_ws_frame_with_routed_input(first_message: str, routed_input: object) -> str: + try: + frame: Final = _JSON_OBJECT_ADAPTER.validate_json(first_message) + except ValidationError: + return first_message + if frame is None or routed_input is None: + return first_message + raw_nested: Final = frame.get("response") + nested: Final = _JSON_OBJECT_ADAPTER.validate_python(raw_nested) if isinstance(raw_nested, Mapping) else None + if nested is not None and nested.get("input") is not None: + if nested["input"] == routed_input: + return first_message + return json.dumps({**frame, "response": {**nested, "input": routed_input}}) + if frame.get("input") == routed_input: + return first_message + return json.dumps({**frame, "input": routed_input}) + + def _build_responses_websocket_request_defaults(kwargs: Mapping[str, object]) -> ResponsesWebSocketRequestDefaults: default_reasoning: Final = _deployment_reasoning_default(kwargs) candidate_params: Final[dict[str, object]] = { @@ -2367,10 +2386,12 @@ async def _aresponses_websocket( "api_base", "api_key", "timeout", + "first_message", *_RESPONSES_WS_ROUTING_HINT_KEYS, } remaining_kwargs: Final = {k: v for k, v in kwargs.items() if k not in _explicit_keys} deployment_kwargs: Final = {k: v for k, v in kwargs.items() if k not in _RESPONSES_WS_ROUTING_HINT_KEYS} + first_message: Final = kwargs.get("first_message") return await base_llm_http_handler.async_responses_websocket( model=resolved_model, @@ -2380,6 +2401,11 @@ async def _aresponses_websocket( api_base=resolved_api_base, api_key=resolved_api_key, timeout=timeout, + first_message=( + _first_ws_frame_with_routed_input(first_message, kwargs.get("input")) + if isinstance(first_message, str) + else None + ), user_api_key_dict=kwargs.get("user_api_key_dict"), litellm_metadata=_build_litellm_metadata_for_ws(kwargs), custom_llm_provider=_custom_llm_provider, diff --git a/tests/test_litellm/proxy/response_api_endpoints/test_endpoints.py b/tests/test_litellm/proxy/response_api_endpoints/test_endpoints.py index 45ec529ce7d..8c3bf27c88d 100644 --- a/tests/test_litellm/proxy/response_api_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/response_api_endpoints/test_endpoints.py @@ -638,6 +638,71 @@ class TestResponsesWSFirstFrameModelAuth: assert booked["user_api_key_dict"] is user_api_key_dict assert booked["request_data"]["model"] == "gpt-4o-mini" + @pytest.mark.asyncio + async def test_endpoint_sends_an_error_frame_when_routing_rejects_the_connection(self): + from litellm.proxy.response_api_endpoints.endpoints import ( + responses_websocket_endpoint, + ) + + ws = MagicMock() + ws.headers = {} + ws.query_params = {} + ws.scope = {"headers": []} + ws.url = "ws://testserver/v1/responses" + ws.accept = AsyncMock() + ws.receive_text = AsyncMock( + return_value=json.dumps({"type": "response.create", "model": "gpt-4o-mini", "input": []}) + ) + ws.send_text = AsyncMock() + ws.close = AsyncMock() + + processor = MagicMock() + processor.common_processing_pre_call_logic = AsyncMock( + return_value=({"model": "gpt-4o-mini", "litellm_metadata": {}}, MagicMock()) + ) + rejection = litellm.RateLimitError( + message="origin deployment is cooling down", model="gpt-4o-mini", llm_provider="openai" + ) + proxy_logging_obj = MagicMock() + proxy_logging_obj.post_call_failure_hook = AsyncMock(return_value=None) + user_api_key_dict = MagicMock() + + with ( + patch( # test-quality-ok: first-frame model auth needs a live router and key table and has its own tests above + "litellm.proxy.response_api_endpoints.endpoints._enforce_responses_ws_first_frame_model_auth", + new_callable=AsyncMock, + ), + patch( # test-quality-ok: the pre-call processor needs a live proxy; what the endpoint tells the client is under test + "litellm.proxy.response_api_endpoints.endpoints.ProxyBaseLLMRequestProcessing", + return_value=processor, + ), + patch( # test-quality-ok: routing is the seam that raises the affinity rejection + "litellm.proxy.route_llm_request.route_request", + new_callable=AsyncMock, + side_effect=rejection, + ), + patch( # test-quality-ok: the failure hook is the proxy's only path to a failed spend log row + "litellm.proxy.proxy_server.proxy_logging_obj", + proxy_logging_obj, + ), + ): + await responses_websocket_endpoint( + websocket=ws, + model=None, + user_api_key_dict=user_api_key_dict, + ) + + frame = json.loads(ws.send_text.await_args.args[0]) + assert frame["type"] == "error" + assert frame["status"] == 429 + assert frame["error"]["type"] == "rate_limit_exceeded" + assert "cooling down" in frame["error"]["message"] + ws.close.assert_awaited_once_with(code=1011, reason="Internal server error") + booked = proxy_logging_obj.post_call_failure_hook.await_args.kwargs + assert booked["original_exception"] is rejection + assert booked["user_api_key_dict"] is user_api_key_dict + assert booked["request_data"]["model"] == "gpt-4o-mini" + @pytest.mark.asyncio async def test_reruns_model_auth_for_first_frame_model(self): from starlette.requests import Request diff --git a/tests/test_litellm/responses/test_responses_api_request_body.py b/tests/test_litellm/responses/test_responses_api_request_body.py index 743ad237e45..6c1348f2350 100644 --- a/tests/test_litellm/responses/test_responses_api_request_body.py +++ b/tests/test_litellm/responses/test_responses_api_request_body.py @@ -448,6 +448,70 @@ async def test_aresponses_websocket_keeps_routing_hints_out_of_the_relay_kwargs( assert "previous_response_id" not in mock_ws.call_args.kwargs +_STRIPPED_WS_INPUT = [{"role": "user", "content": "hi"}] +_ORIGINAL_WS_INPUT = [ + {"type": "reasoning", "id": "rs_1", "encrypted_content": "blob-from-a-removed-deployment", "summary": []}, + *_STRIPPED_WS_INPUT, +] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("nested", [False, True]) +async def test_aresponses_websocket_forwards_the_routed_input_in_the_first_frame(nested): # test-quality-ok: the first frame handed to the relay is the only place the routed input is observable before the provider socket + from unittest.mock import MagicMock + + from litellm.responses.main import _aresponses_websocket + + body = {"model": "gpt-5.6", "input": _ORIGINAL_WS_INPUT, "store": False} + first_message = json.dumps( + {"type": "response.create", "response": body} if nested else {"type": "response.create", **body} + ) + + with patch.object( + import_module("litellm.responses.main").base_llm_http_handler, "async_responses_websocket", + new_callable=AsyncMock, + ) as mock_ws: + await _aresponses_websocket( + model="openai/gpt-5.6", + websocket=MagicMock(), + api_key="sk-test", + litellm_logging_obj=MagicMock(), + input=list(_STRIPPED_WS_INPUT), + first_message=first_message, + ) + + forwarded = json.loads(mock_ws.call_args.kwargs["first_message"]) + container = forwarded["response"] if nested else forwarded + assert container["input"] == _STRIPPED_WS_INPUT + assert container["store"] is False + assert container["model"] == "gpt-5.6" + assert forwarded["type"] == "response.create" + + +@pytest.mark.asyncio +async def test_aresponses_websocket_forwards_the_first_frame_verbatim_when_routing_left_the_input_alone(): # test-quality-ok: the relay kwargs are the boundary; byte-identical passthrough is only observable there + from unittest.mock import MagicMock + + from litellm.responses.main import _aresponses_websocket + + first_message = '{"type": "response.create", "model": "gpt-5.6", "input": [{"role": "user", "content": "hi"}]}' + + with patch.object( + import_module("litellm.responses.main").base_llm_http_handler, "async_responses_websocket", + new_callable=AsyncMock, + ) as mock_ws: + await _aresponses_websocket( + model="openai/gpt-5.6", + websocket=MagicMock(), + api_key="sk-test", + litellm_logging_obj=MagicMock(), + input=list(_STRIPPED_WS_INPUT), + first_message=first_message, + ) + + assert mock_ws.call_args.kwargs["first_message"] == first_message + + _INJECTION_POINT_INPUT = [{"role": "system", "content": "You are terse."}, {"role": "user", "content": "hi"}] _SYSTEM_POINT = {"location": "message", "role": "system"} _USER_POINT = {"location": "message", "role": "user"} From c9cd666b368c0194c6d9d3c04bbb77da7b2e630f Mon Sep 17 00:00:00 2001 From: kerry Date: Sat, 19 Sep 2026 01:18:09 +0000 Subject: [PATCH 123/144] refactor(cost_calc): move the fireworks cache-read default under litellm/llms Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../litellm_core_utils/llm_cost_calc/utils.py | 33 +++--------- litellm/llms/fireworks_ai/cache_pricing.py | 38 ++++++++++++++ .../llm_cost_calc/test_llm_cost_calc_utils.py | 25 ++-------- .../test_fireworks_ai_cache_pricing.py | 50 +++++++++++++++++++ 4 files changed, 98 insertions(+), 48 deletions(-) create mode 100644 litellm/llms/fireworks_ai/cache_pricing.py create mode 100644 tests/test_litellm/llms/fireworks_ai/test_fireworks_ai_cache_pricing.py diff --git a/litellm/litellm_core_utils/llm_cost_calc/utils.py b/litellm/litellm_core_utils/llm_cost_calc/utils.py index 7de4534ef2c..cf1b1962df8 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/utils.py +++ b/litellm/litellm_core_utils/llm_cost_calc/utils.py @@ -14,11 +14,11 @@ from typing_extensions import ReadOnly import litellm from litellm._internal_context import current_billing_time from litellm._logging import verbose_logger -from litellm.constants import FIREWORKS_AI_DEFAULT_CACHE_READ_RATE_RATIO from litellm.litellm_core_utils.llm_cost_calc.tiered_pricing import ( select_tier_for_input, tier_rate, ) +from litellm.llms.fireworks_ai.cache_pricing import with_default_cache_read_rate from litellm.types.utils import ( CacheCreationTokenDetails, CallTypes, @@ -74,31 +74,12 @@ def _uses_inclusive_token_thresholds(custom_llm_provider: str | None) -> bool: def apply_provider_cache_read_default(model_info: ModelInfo, custom_llm_provider: str | None) -> ModelInfo: - """Apply provider-specific defaults for cache-read pricing.""" - if custom_llm_provider != "fireworks_ai": - return model_info - input_rate: Final = model_info.get("input_cost_per_token") - if model_info.get("cache_read_input_token_cost") is not None or input_rate is None: - return model_info - cache_read_rate: Final = input_rate * FIREWORKS_AI_DEFAULT_CACHE_READ_RATE_RATIO - off_peak: Final = model_info.get("off_peak_pricing") - if off_peak is None or "cache_read_input_token_cost" in off_peak: - return cast(ModelInfo, {**model_info, "cache_read_input_token_cost": cache_read_rate}) - return cast( - ModelInfo, - { - **model_info, - "cache_read_input_token_cost": cache_read_rate, - "off_peak_pricing": { - **off_peak, - "cache_read_input_token_cost": ( - off_peak["input_cost_per_token"] * FIREWORKS_AI_DEFAULT_CACHE_READ_RATE_RATIO - if "input_cost_per_token" in off_peak - else cache_read_rate - ), - }, - }, - ) + """Dispatch to the provider's cache-read pricing default; providers without one keep their entry as is.""" + match custom_llm_provider: + case "fireworks_ai": + return with_default_cache_read_rate(model_info) + case _: + return model_info def _get_token_detail_value(details: object, key: str) -> int | None: diff --git a/litellm/llms/fireworks_ai/cache_pricing.py b/litellm/llms/fireworks_ai/cache_pricing.py new file mode 100644 index 00000000000..c28a8684c66 --- /dev/null +++ b/litellm/llms/fireworks_ai/cache_pricing.py @@ -0,0 +1,38 @@ +""" +Fireworks AI serverless cache-read pricing defaults. +""" + +from typing import ( + Final, + cast, # noqa: TID251 # the derived entry is a dict copy of a ReadOnly TypedDict; no cast-free way to retype it +) + +from litellm.constants import FIREWORKS_AI_DEFAULT_CACHE_READ_RATE_RATIO +from litellm.types.utils import ModelInfo + + +def with_default_cache_read_rate(model_info: ModelInfo) -> ModelInfo: + """Entries without a cache-read rate get the documented discount off the input rate; the shared map is + never mutated, so a copy carries it.""" + input_rate: Final = model_info.get("input_cost_per_token") + if model_info.get("cache_read_input_token_cost") is not None or input_rate is None: + return model_info + cache_read_rate: Final = input_rate * FIREWORKS_AI_DEFAULT_CACHE_READ_RATE_RATIO + off_peak: Final = model_info.get("off_peak_pricing") + if off_peak is None or "cache_read_input_token_cost" in off_peak: + return cast(ModelInfo, {**model_info, "cache_read_input_token_cost": cache_read_rate}) + return cast( + ModelInfo, + { + **model_info, + "cache_read_input_token_cost": cache_read_rate, + "off_peak_pricing": { + **off_peak, + "cache_read_input_token_cost": ( + off_peak["input_cost_per_token"] * FIREWORKS_AI_DEFAULT_CACHE_READ_RATE_RATIO + if "input_cost_per_token" in off_peak + else cache_read_rate + ), + }, + }, + ) diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py index 32d7dd1d0c2..e85cbe65b18 100644 --- a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py +++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py @@ -98,36 +98,17 @@ def test_generic_cost_per_token_bills_cache_reads_at_input_rate_when_no_cache_re assert completion_cost == pytest.approx(380 * 9.7e-7) -def test_apply_provider_cache_read_default_preserves_identity_and_input_data() -> None: +def test_apply_provider_cache_read_default_only_derives_a_rate_for_fireworks() -> None: openai_info: ModelInfo = {"input_cost_per_token": 2e-6} - explicit_fireworks_info: ModelInfo = { - "input_cost_per_token": 2e-6, - "cache_read_input_token_cost": 1e-6, - } - fireworks_info: ModelInfo = { - "input_cost_per_token": 2e-6, - "off_peak_pricing": { - "hours_utc": "14:00-00:00", - "input_cost_per_token": 1e-6, - "output_cost_per_token": 3e-6, - }, - } - original_fireworks_info: ModelInfo = deepcopy(fireworks_info) + fireworks_info: ModelInfo = {"input_cost_per_token": 2e-6} assert apply_provider_cache_read_default(openai_info, "openai") is openai_info - assert apply_provider_cache_read_default(explicit_fireworks_info, "fireworks_ai") is explicit_fireworks_info + assert apply_provider_cache_read_default(openai_info, None) is openai_info processed_fireworks_info = apply_provider_cache_read_default(fireworks_info, "fireworks_ai") - assert fireworks_info == original_fireworks_info assert processed_fireworks_info is not fireworks_info assert processed_fireworks_info["cache_read_input_token_cost"] == pytest.approx(2e-6 * 0.5) - assert processed_fireworks_info["off_peak_pricing"] == { - "hours_utc": "14:00-00:00", - "input_cost_per_token": 1e-6, - "output_cost_per_token": 3e-6, - "cache_read_input_token_cost": 1e-6 * 0.5, - } def test_generic_cost_per_token_prefers_audio_per_second_rate() -> None: diff --git a/tests/test_litellm/llms/fireworks_ai/test_fireworks_ai_cache_pricing.py b/tests/test_litellm/llms/fireworks_ai/test_fireworks_ai_cache_pricing.py new file mode 100644 index 00000000000..2ab273b27c6 --- /dev/null +++ b/tests/test_litellm/llms/fireworks_ai/test_fireworks_ai_cache_pricing.py @@ -0,0 +1,50 @@ +from copy import deepcopy + +import pytest + +from litellm.constants import FIREWORKS_AI_DEFAULT_CACHE_READ_RATE_RATIO +from litellm.llms.fireworks_ai.cache_pricing import with_default_cache_read_rate +from litellm.types.utils import ModelInfo + + +def test_explicit_cache_read_rate_and_missing_input_rate_keep_the_entry_untouched() -> None: + explicit_info: ModelInfo = {"input_cost_per_token": 2e-6, "cache_read_input_token_cost": 1e-6} + no_input_rate_info: ModelInfo = {"output_cost_per_token": 3e-6} + + assert with_default_cache_read_rate(explicit_info) is explicit_info + assert with_default_cache_read_rate(no_input_rate_info) is no_input_rate_info + + +def test_missing_cache_read_rate_is_derived_for_standard_and_off_peak_without_mutating_the_entry() -> None: + model_info: ModelInfo = { + "input_cost_per_token": 2e-6, + "off_peak_pricing": { + "hours_utc": "14:00-00:00", + "input_cost_per_token": 1e-6, + "output_cost_per_token": 3e-6, + }, + } + original: ModelInfo = deepcopy(model_info) + + derived = with_default_cache_read_rate(model_info) + + assert model_info == original + assert derived is not model_info + assert derived["cache_read_input_token_cost"] == pytest.approx(2e-6 * FIREWORKS_AI_DEFAULT_CACHE_READ_RATE_RATIO) + assert derived["off_peak_pricing"] == { + "hours_utc": "14:00-00:00", + "input_cost_per_token": 1e-6, + "output_cost_per_token": 3e-6, + "cache_read_input_token_cost": 1e-6 * FIREWORKS_AI_DEFAULT_CACHE_READ_RATE_RATIO, + } + + +def test_off_peak_window_without_its_own_input_rate_reuses_the_standard_derived_rate() -> None: + model_info: ModelInfo = { + "input_cost_per_token": 2e-6, + "off_peak_pricing": {"hours_utc": "14:00-00:00", "output_cost_per_token": 3e-6}, + } + + derived = with_default_cache_read_rate(model_info) + + assert derived["off_peak_pricing"]["cache_read_input_token_cost"] == derived["cache_read_input_token_cost"] From 8836410c4c00ba115e13a8912011b9bf56a1b24d Mon Sep 17 00:00:00 2001 From: kerry Date: Sat, 19 Sep 2026 01:18:42 +0000 Subject: [PATCH 124/144] test(cost_calc): drop the unused deepcopy import Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py index e85cbe65b18..4fe3d410ef8 100644 --- a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py +++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py @@ -1,5 +1,4 @@ from collections.abc import Mapping -from copy import deepcopy from datetime import datetime, timezone import pytest From b1b7af884abe764c03dece34c8b621b5c0b19a55 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 18 Sep 2026 18:19:48 -0700 Subject: [PATCH 125/144] fix(websearch): forward the deployment api_base to agentic follow-up calls on /v1/messages --- litellm/llms/custom_httpx/llm_http_handler.py | 25 +++--- .../custom_httpx/test_llm_http_handler.py | 90 +++++++++++++++++++ 2 files changed, 104 insertions(+), 11 deletions(-) diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 477d10a3cbd..cd76f0d0b54 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -2223,6 +2223,7 @@ class BaseLLMHTTPHandler: # Prepare headers kwargs = kwargs or {} + kwargs_for_agentic: Final = self._agentic_hook_kwargs(kwargs=kwargs, api_key=api_key, api_base=api_base) provider_specific_header: Final = cast( litellm.types.utils.ProviderSpecificHeader | Sequence[litellm.types.utils.ProviderSpecificHeader] | None, kwargs.get("provider_specific_header", None), @@ -2410,7 +2411,7 @@ class BaseLLMHTTPHandler: anthropic_messages_optional_request_params=anthropic_messages_optional_request_params, logging_obj=logging_obj, custom_llm_provider=custom_llm_provider, - kwargs={**kwargs, "api_key": api_key} if api_key else kwargs, + kwargs=kwargs_for_agentic, hold_back=bool(held_back_tool_names), server_fulfilled_tool_names=held_back_tool_names, ) @@ -2433,8 +2434,7 @@ class BaseLLMHTTPHandler: anthropic_messages_optional_request_params=anthropic_messages_optional_request_params, logging_obj=logging_obj, custom_llm_provider=custom_llm_provider, - api_key=api_key, - kwargs=kwargs, + kwargs=kwargs_for_agentic, ) async def _finalize_anthropic_messages_response( @@ -2447,14 +2447,8 @@ class BaseLLMHTTPHandler: anthropic_messages_optional_request_params: dict, logging_obj: LiteLLMLoggingObj, custom_llm_provider: str, - api_key: str | None, - kwargs: dict, + kwargs: dict[str, object], ) -> AnthropicMessagesResponse | AsyncIterator: - # Inject api_key into kwargs so follow-up calls in agentic hooks can - # authenticate. api_key is a named param here (not in kwargs), so - # _prepare_followup_kwargs would miss it otherwise. - kwargs_for_agentic: Final = {**kwargs, "api_key": api_key} if api_key else kwargs - # Call agentic completion hooks (non-streaming path only) final_response: Final = await self._call_agentic_completion_hooks( response=initial_response, model=model, @@ -2464,7 +2458,7 @@ class BaseLLMHTTPHandler: logging_obj=logging_obj, stream=False, custom_llm_provider=custom_llm_provider, - kwargs=kwargs_for_agentic, + kwargs=kwargs, ) return self._maybe_wrap_in_fake_stream( @@ -5312,6 +5306,15 @@ class BaseLLMHTTPHandler: fingerprints: Final = list(kwargs.get("_agentic_loop_fingerprints", []) or []) return depth, max_loops, fingerprints + @staticmethod + def _agentic_hook_kwargs( + kwargs: Mapping[str, object], api_key: str | None, api_base: str | None + ) -> dict[str, object]: + """``api_key`` and ``api_base`` are named parameters of ``anthropic_messages`` rather than kwargs, so the + follow-up call an agentic hook makes only reaches the same deployment if they are re-added here.""" + deployment_params: Final = {"api_key": api_key, "api_base": api_base} + return {**kwargs, **{key: value for key, value in deployment_params.items() if value}} + @staticmethod def _has_agentic_completion_hook(logging_obj: LiteLLMLoggingObj) -> bool: """ diff --git a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py index 95dceccb2f5..6dc457d26ec 100644 --- a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py +++ b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py @@ -1956,6 +1956,96 @@ async def test_async_anthropic_messages_handler_passes_api_key_to_agentic_hooks( ) +_FOUNDRY_API_BASE: Final = "https://lit5418.services.ai.azure.com/anthropic" +_FOUNDRY_SSE_BODY: Final = ( + b'event: message_start\ndata: {"type": "message_start", "message": {"id": "msg_1", "type": "message", ' + b'"role": "assistant", "model": "claude-fable-5-1", "content": [], "stop_reason": null, ' + b'"usage": {"input_tokens": 1, "output_tokens": 0}}}\n\n' + b'event: content_block_start\ndata: {"type": "content_block_start", "index": 0, ' + b'"content_block": {"type": "text", "text": ""}}\n\n' + b'event: content_block_delta\ndata: {"type": "content_block_delta", "index": 0, ' + b'"delta": {"type": "text_delta", "text": "ready"}}\n\n' + b'event: content_block_stop\ndata: {"type": "content_block_stop", "index": 0}\n\n' + b'event: message_delta\ndata: {"type": "message_delta", "delta": {"stop_reason": "end_turn"}, ' + b'"usage": {"output_tokens": 1}}\n\n' + b'event: message_stop\ndata: {"type": "message_stop"}\n\n' +) + + +@pytest.mark.parametrize("stream", [False, True]) +@pytest.mark.asyncio +async def test_async_anthropic_messages_handler_passes_deployment_api_base_to_agentic_hooks(stream, monkeypatch): + """ + Regression for LIT-5418: an azure_ai deployment carries its Foundry endpoint as + ``api_base``, a named parameter that never lands in kwargs. The agentic hooks + (websearch interception's follow-up call after the search) must receive it on + both the non-streaming and the streaming path, or the follow-up fails with + "Missing Azure API Base" and the client gets the dangling tool_use back. + """ + from litellm.integrations.custom_logger import CustomLogger + from litellm.llms.azure_ai.anthropic.messages_transformation import AzureAnthropicMessagesConfig + + monkeypatch.delenv("AZURE_API_BASE", raising=False) + + class CapturingAgenticCallback(CustomLogger): + def __init__(self): + super().__init__() + self.hook_kwargs: dict | None = None + + async def async_should_run_agentic_loop(self, response, model, messages, tools, stream, custom_llm_provider, kwargs): + self.hook_kwargs = dict(kwargs) + return False, {} + + callback = CapturingAgenticCallback() + handler = BaseLLMHTTPHandler() + upstream_request = httpx.Request("POST", f"{_FOUNDRY_API_BASE}/v1/messages") + upstream_response = ( + httpx.Response(200, content=_FOUNDRY_SSE_BODY, request=upstream_request) + if stream + else httpx.Response( + 200, + json={ + "id": "msg_1", + "type": "message", + "role": "assistant", + "model": "claude-fable-5-1", + "content": [{"type": "text", "text": "ready"}], + "stop_reason": "end_turn", + "usage": {"input_tokens": 1, "output_tokens": 1}, + }, + request=upstream_request, + ) + ) + mock_client = AsyncMock(spec=AsyncHTTPHandler) + mock_client.post = AsyncMock(return_value=upstream_response) + + mock_logging_obj = Mock() + mock_logging_obj.model_call_details = {} + mock_logging_obj.dynamic_success_callbacks = [callback] + + result = await handler.async_anthropic_messages_handler( + model="claude-fable-5-1", + messages=[{"role": "user", "content": "Say ready"}], + anthropic_messages_provider_config=AzureAnthropicMessagesConfig(), + anthropic_messages_optional_request_params={"max_tokens": 32}, + custom_llm_provider="azure_ai", + litellm_params=GenericLiteLLMParams(api_key="foundry-key", api_base=_FOUNDRY_API_BASE), + logging_obj=mock_logging_obj, + client=mock_client, + api_key="foundry-key", + api_base=_FOUNDRY_API_BASE, + stream=stream, + kwargs={}, + ) + if stream: + _ = [chunk async for chunk in result] + + assert mock_client.post.call_args.kwargs["url"] == f"{_FOUNDRY_API_BASE}/v1/messages" + assert callback.hook_kwargs is not None, "agentic hook never ran" + assert callback.hook_kwargs.get("api_base") == _FOUNDRY_API_BASE + assert callback.hook_kwargs.get("api_key") == "foundry-key" + + class _FakeWSExceptions: class WebSocketException(Exception): pass From 99d91d72056a43a95e6f377e9dda02df69d111d5 Mon Sep 17 00:00:00 2001 From: kerry Date: Sat, 19 Sep 2026 01:22:25 +0000 Subject: [PATCH 126/144] fix(xai): reject non-success stt responses before parsing Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../xai/audio_transcription/transformation.py | 7 +++++++ ...est_xai_audio_transcription_transformation.py | 16 ++++++++++++++++ 2 files changed, 23 insertions(+) diff --git a/litellm/llms/xai/audio_transcription/transformation.py b/litellm/llms/xai/audio_transcription/transformation.py index 03c06f24a2d..7b648fc8084 100644 --- a/litellm/llms/xai/audio_transcription/transformation.py +++ b/litellm/llms/xai/audio_transcription/transformation.py @@ -130,6 +130,13 @@ class XAIAudioTranscriptionConfig(BaseAudioTranscriptionConfig): self, raw_response: Response, ) -> TranscriptionResponse: + if raw_response.status_code >= 400: + raise self.get_error_class( + error_message=raw_response.text, + status_code=raw_response.status_code, + headers=raw_response.headers, + ) + try: payload: Final = _XAISttResponse.model_validate_json(raw_response.content) except ValidationError as e: diff --git a/tests/test_litellm/llms/xai/test_xai_audio_transcription_transformation.py b/tests/test_litellm/llms/xai/test_xai_audio_transcription_transformation.py index 0f3445eb400..e2e3fc3d4dc 100644 --- a/tests/test_litellm/llms/xai/test_xai_audio_transcription_transformation.py +++ b/tests/test_litellm/llms/xai/test_xai_audio_transcription_transformation.py @@ -8,6 +8,7 @@ from litellm.llms.base_llm.audio_transcription.transformation import ( from litellm.llms.custom_httpx.http_handler import HTTPHandler from litellm.llms.xai.audio_transcription.transformation import ( XAIAudioTranscriptionConfig, + XAIAudioTranscriptionError, ) from litellm.types.utils import LlmProviders from litellm.utils import ProviderConfigManager @@ -130,6 +131,21 @@ def test_transform_response_maps_xai_shape(): assert response._hidden_params["audio_transcription_duration"] == 3.2 +def test_transform_response_raises_on_error_status(): + raw = httpx.Response( + 400, + json={ + "code": "Client specified an invalid argument", + "error": "Incorrect API key provided", + }, + request=httpx.Request("POST", "https://api.x.ai/v1/stt"), + ) + with pytest.raises(XAIAudioTranscriptionError) as exc: + CONFIG.transform_audio_transcription_response(raw_response=raw) + assert exc.value.status_code == 400 + assert "Incorrect API key provided" in exc.value.message + + def test_transcription_routes_to_xai_stt(monkeypatch): monkeypatch.delenv("XAI_API_KEY", raising=False) monkeypatch.setattr(litellm, "xai_key", None) From 06a5594bb615471fd4c3fc125e72cdb619ff80ca Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 18 Sep 2026 18:25:55 -0700 Subject: [PATCH 127/144] fix(claude_code_gateway): mint the bearer before consuming the device code so a signing failure never spends the login --- .../anthropic_endpoints/gateway_endpoints.py | 5 ++--- .../test_gateway_endpoints.py | 17 ++++++++++++++--- 2 files changed, 16 insertions(+), 6 deletions(-) diff --git a/litellm/proxy/anthropic_endpoints/gateway_endpoints.py b/litellm/proxy/anthropic_endpoints/gateway_endpoints.py index 259d5202db6..0446992ae43 100644 --- a/litellm/proxy/anthropic_endpoints/gateway_endpoints.py +++ b/litellm/proxy/anthropic_endpoints/gateway_endpoints.py @@ -315,13 +315,12 @@ async def _handle_device_code_grant(device_code: str | None) -> JSONResponse: if isinstance(login, _OAuthError): return _oauth_error_response(login) + access_token: Final = _mint_access_token(login) if not await _claim_device_code(login_id, cli_sso_session_cache): return _oauth_error_response(_OAuthError(status_code=400, error="expired_token")) await cli_sso_session_cache.async_delete_cache(key=_get_cli_sso_flow_cache_key(login_id)) - body: Final = _AccessTokenBody( - access_token=_mint_access_token(login), expires_in=CLI_JWT_EXPIRATION_HOURS * _SECONDS_PER_HOUR - ) + body: Final = _AccessTokenBody(access_token=access_token, expires_in=CLI_JWT_EXPIRATION_HOURS * _SECONDS_PER_HOUR) return JSONResponse(content=body.model_dump()) diff --git a/tests/test_litellm/proxy/anthropic_endpoints/test_gateway_endpoints.py b/tests/test_litellm/proxy/anthropic_endpoints/test_gateway_endpoints.py index d442ac21307..7c3e8f56a21 100644 --- a/tests/test_litellm/proxy/anthropic_endpoints/test_gateway_endpoints.py +++ b/tests/test_litellm/proxy/anthropic_endpoints/test_gateway_endpoints.py @@ -320,6 +320,18 @@ def test_token_malformed_session_is_invalid_grant_and_does_not_consume_the_login mint.assert_not_called() +def test_token_mint_failure_leaves_the_login_unconsumed(): + with _gateway_env() as (client, cache): + device_code = _start_device_flow(client) + _complete_flow(cache, device_code) + with patch(_MINT, side_effect=RuntimeError("signing key unavailable")), pytest.raises(RuntimeError): + _request_token(client, device_code) + with patch(_MINT, return_value="sk-session"): + retry = _request_token(client, device_code) + assert retry.status_code == 200 + assert retry.json()["access_token"] == "sk-session" + + def test_token_unknown_team_grants_is_invalid_grant(): with _gateway_env() as (client, cache): device_code = _start_device_flow(client) @@ -349,11 +361,10 @@ def test_token_refuses_a_device_code_another_replica_already_claimed(): _set_cli_sso_flow(login_id=_SHARED_LOGIN_ID, cache=replica_a, flow=_completed_flow()) assert asyncio.run(gateway_endpoints._claim_device_code(_SHARED_LOGIN_ID, replica_a)) is True - with _gateway_env(cache=_replica(redis)) as (client, _), patch(_MINT) as mint: + with _gateway_env(cache=_replica(redis)) as (client, _), patch(_MINT, return_value="sk-session"): resp = _request_token(client, _SHARED_DEVICE_CODE) assert resp.status_code == 400 - assert resp.json()["error"] == "expired_token" - mint.assert_not_called() + assert resp.json() == {"error": "expired_token"} def test_token_unknown_device_code_is_expired_token(): From 05cefb1480bdac1943abea819060db222cceee8f Mon Sep 17 00:00:00 2001 From: kerry Date: Sat, 19 Sep 2026 01:33:41 +0000 Subject: [PATCH 128/144] fix(cost_calc): coerce string fireworks rates and drop the match fall-through Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../litellm_core_utils/llm_cost_calc/utils.py | 9 +++----- litellm/llms/fireworks_ai/cache_pricing.py | 22 +++++++++++-------- .../test_fireworks_ai_cache_pricing.py | 15 +++++++++++++ 3 files changed, 31 insertions(+), 15 deletions(-) diff --git a/litellm/litellm_core_utils/llm_cost_calc/utils.py b/litellm/litellm_core_utils/llm_cost_calc/utils.py index cf1b1962df8..e24fa004448 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/utils.py +++ b/litellm/litellm_core_utils/llm_cost_calc/utils.py @@ -74,12 +74,9 @@ def _uses_inclusive_token_thresholds(custom_llm_provider: str | None) -> bool: def apply_provider_cache_read_default(model_info: ModelInfo, custom_llm_provider: str | None) -> ModelInfo: - """Dispatch to the provider's cache-read pricing default; providers without one keep their entry as is.""" - match custom_llm_provider: - case "fireworks_ai": - return with_default_cache_read_rate(model_info) - case _: - return model_info + if custom_llm_provider == "fireworks_ai": + return with_default_cache_read_rate(model_info) + return model_info def _get_token_detail_value(details: object, key: str) -> int | None: diff --git a/litellm/llms/fireworks_ai/cache_pricing.py b/litellm/llms/fireworks_ai/cache_pricing.py index c28a8684c66..f5e49cad01a 100644 --- a/litellm/llms/fireworks_ai/cache_pricing.py +++ b/litellm/llms/fireworks_ai/cache_pricing.py @@ -1,7 +1,3 @@ -""" -Fireworks AI serverless cache-read pricing defaults. -""" - from typing import ( Final, cast, # noqa: TID251 # the derived entry is a dict copy of a ReadOnly TypedDict; no cast-free way to retype it @@ -11,16 +7,24 @@ from litellm.constants import FIREWORKS_AI_DEFAULT_CACHE_READ_RATE_RATIO from litellm.types.utils import ModelInfo +def _as_rate(value: object) -> float | None: + if isinstance(value, bool) or not isinstance(value, (int, float, str)): + return None + try: + return float(value) + except ValueError: + return None + + def with_default_cache_read_rate(model_info: ModelInfo) -> ModelInfo: - """Entries without a cache-read rate get the documented discount off the input rate; the shared map is - never mutated, so a copy carries it.""" - input_rate: Final = model_info.get("input_cost_per_token") + input_rate: Final = _as_rate(model_info.get("input_cost_per_token")) if model_info.get("cache_read_input_token_cost") is not None or input_rate is None: return model_info cache_read_rate: Final = input_rate * FIREWORKS_AI_DEFAULT_CACHE_READ_RATE_RATIO off_peak: Final = model_info.get("off_peak_pricing") if off_peak is None or "cache_read_input_token_cost" in off_peak: return cast(ModelInfo, {**model_info, "cache_read_input_token_cost": cache_read_rate}) + off_peak_input_rate: Final = _as_rate(off_peak.get("input_cost_per_token")) return cast( ModelInfo, { @@ -29,8 +33,8 @@ def with_default_cache_read_rate(model_info: ModelInfo) -> ModelInfo: "off_peak_pricing": { **off_peak, "cache_read_input_token_cost": ( - off_peak["input_cost_per_token"] * FIREWORKS_AI_DEFAULT_CACHE_READ_RATE_RATIO - if "input_cost_per_token" in off_peak + off_peak_input_rate * FIREWORKS_AI_DEFAULT_CACHE_READ_RATE_RATIO + if off_peak_input_rate is not None else cache_read_rate ), }, diff --git a/tests/test_litellm/llms/fireworks_ai/test_fireworks_ai_cache_pricing.py b/tests/test_litellm/llms/fireworks_ai/test_fireworks_ai_cache_pricing.py index 2ab273b27c6..c21943cfe75 100644 --- a/tests/test_litellm/llms/fireworks_ai/test_fireworks_ai_cache_pricing.py +++ b/tests/test_litellm/llms/fireworks_ai/test_fireworks_ai_cache_pricing.py @@ -48,3 +48,18 @@ def test_off_peak_window_without_its_own_input_rate_reuses_the_standard_derived_ derived = with_default_cache_read_rate(model_info) assert derived["off_peak_pricing"]["cache_read_input_token_cost"] == derived["cache_read_input_token_cost"] + + +def test_string_rates_from_config_are_coerced_before_the_discount_is_applied() -> None: + model_info: ModelInfo = { + "input_cost_per_token": "2e-6", + "off_peak_pricing": { + "hours_utc": "14:00-00:00", + "input_cost_per_token": "1e-6", + }, + } + + derived = with_default_cache_read_rate(model_info) + + assert derived["cache_read_input_token_cost"] == pytest.approx(1e-6) + assert derived["off_peak_pricing"]["cache_read_input_token_cost"] == pytest.approx(5e-7) From 9662b2a35c0ab65150bcab4bc45131bc06438371 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 18 Sep 2026 18:35:48 -0700 Subject: [PATCH 129/144] refactor(responses): type the websocket test parameters and suppress the error-frame send explicitly --- litellm/proxy/response_api_endpoints/endpoints.py | 5 ++--- .../litellm_core_utils/test_litellm_logging.py | 2 +- .../proxy/response_api_endpoints/test_endpoints.py | 6 ++++-- .../responses/test_responses_api_request_body.py | 2 +- .../test_responses_websocket_all_providers.py | 10 +++++++--- 5 files changed, 15 insertions(+), 10 deletions(-) diff --git a/litellm/proxy/response_api_endpoints/endpoints.py b/litellm/proxy/response_api_endpoints/endpoints.py index 1b1fc466046..b3d6a928a78 100644 --- a/litellm/proxy/response_api_endpoints/endpoints.py +++ b/litellm/proxy/response_api_endpoints/endpoints.py @@ -1,4 +1,5 @@ import asyncio +import contextlib import json import time from collections.abc import AsyncIterator, Awaitable, Mapping @@ -1585,10 +1586,8 @@ async def responses_websocket_endpoint( ) except Exception as e: verbose_proxy_logger.exception("Responses WebSocket error") - try: + with contextlib.suppress(Exception): await websocket.send_text(_responses_ws_failure_frame(e)) - except Exception: - pass await proxy_logging_obj.post_call_failure_hook( user_api_key_dict=user_api_key_dict, original_exception=e, diff --git a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py index 91c334692ee..836ac42e1f5 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -1069,7 +1069,7 @@ async def test_arealtime_marks_litellm_params_async(monkeypatch): @pytest.mark.asyncio -async def test_aresponses_websocket_hands_back_the_provider_failure_without_a_success_log(monkeypatch): +async def test_aresponses_websocket_hands_back_the_provider_failure_without_a_success_log(monkeypatch: pytest.MonkeyPatch): from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER from litellm.responses.main import base_llm_http_handler diff --git a/tests/test_litellm/proxy/response_api_endpoints/test_endpoints.py b/tests/test_litellm/proxy/response_api_endpoints/test_endpoints.py index 8c3bf27c88d..1560b7c32a6 100644 --- a/tests/test_litellm/proxy/response_api_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/response_api_endpoints/test_endpoints.py @@ -513,7 +513,9 @@ class TestResponsesWSFirstFrameModelAuth: @pytest.mark.asyncio @pytest.mark.parametrize("nested", [False, True]) @pytest.mark.parametrize("query_model", [None, "gpt-4o-mini"]) - async def test_endpoint_routes_on_first_frame_input_and_previous_response_id(self, nested, query_model): + async def test_endpoint_routes_on_first_frame_input_and_previous_response_id( + self, nested: bool, query_model: str | None + ): from litellm.proxy.response_api_endpoints.endpoints import ( responses_websocket_endpoint, ) @@ -572,7 +574,7 @@ class TestResponsesWSFirstFrameModelAuth: @pytest.mark.asyncio @pytest.mark.parametrize("provider_rejected", [True, False]) - async def test_endpoint_books_a_provider_rejected_connection_as_a_failed_request(self, provider_rejected): + async def test_endpoint_books_a_provider_rejected_connection_as_a_failed_request(self, provider_rejected: bool): from litellm.proxy.response_api_endpoints.endpoints import ( responses_websocket_endpoint, ) diff --git a/tests/test_litellm/responses/test_responses_api_request_body.py b/tests/test_litellm/responses/test_responses_api_request_body.py index 6c1348f2350..6b5aab932ec 100644 --- a/tests/test_litellm/responses/test_responses_api_request_body.py +++ b/tests/test_litellm/responses/test_responses_api_request_body.py @@ -457,7 +457,7 @@ _ORIGINAL_WS_INPUT = [ @pytest.mark.asyncio @pytest.mark.parametrize("nested", [False, True]) -async def test_aresponses_websocket_forwards_the_routed_input_in_the_first_frame(nested): # test-quality-ok: the first frame handed to the relay is the only place the routed input is observable before the provider socket +async def test_aresponses_websocket_forwards_the_routed_input_in_the_first_frame(nested: bool): # test-quality-ok: the first frame handed to the relay is the only place the routed input is observable before the provider socket from unittest.mock import MagicMock from litellm.responses.main import _aresponses_websocket diff --git a/tests/test_litellm/responses/test_responses_websocket_all_providers.py b/tests/test_litellm/responses/test_responses_websocket_all_providers.py index b6d4d9e93a6..2fe9f231f14 100644 --- a/tests/test_litellm/responses/test_responses_websocket_all_providers.py +++ b/tests/test_litellm/responses/test_responses_websocket_all_providers.py @@ -1503,7 +1503,9 @@ class TestNativeWebSocketDeploymentDefaults: assert dict(request_defaults.overrides) == {"provider_default": "configured"} @pytest.mark.asyncio - async def test_aresponses_websocket_keeps_first_frame_routing_hints_out_of_the_defaults(self, monkeypatch): + async def test_aresponses_websocket_keeps_first_frame_routing_hints_out_of_the_defaults( + self, monkeypatch: pytest.MonkeyPatch + ): import importlib from unittest.mock import AsyncMock @@ -2976,7 +2978,7 @@ class TestNativeWebSocketEncryptedContentAffinity: @pytest.mark.asyncio @pytest.mark.parametrize("nested", [False, True]) - async def test_client_to_backend_restores_wrapped_ids(self, nested): + async def test_client_to_backend_restores_wrapped_ids(self, nested: bool): from unittest.mock import AsyncMock from litellm.responses.utils import ResponsesAPIRequestUtils @@ -3138,7 +3140,9 @@ class TestNativeWebSocketEncryptedContentAffinity: ), ], ) - async def test_backend_to_client_books_failure_frames_as_failures(self, failure_frame, expected_status): + async def test_backend_to_client_books_failure_frames_as_failures( + self, failure_frame: dict[str, object], expected_status: int + ): import asyncio from unittest.mock import AsyncMock From 4e38f1845d8e972b330009fb1f70cbde35afb838 Mon Sep 17 00:00:00 2001 From: kerry Date: Sat, 19 Sep 2026 01:40:00 +0000 Subject: [PATCH 130/144] refactor(xai): move native stt routing opt-out behind the provider config Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/constants.py | 6 +----- .../llms/base_llm/audio_transcription/transformation.py | 9 +++++++++ litellm/llms/xai/audio_transcription/transformation.py | 4 ++++ litellm/main.py | 6 +++++- 4 files changed, 19 insertions(+), 6 deletions(-) diff --git a/litellm/constants.py b/litellm/constants.py index 55f92f29c96..83dd91c9b7d 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -999,11 +999,7 @@ openai_compatible_providers: Final[list] = [ "scx-ai", ] -OPENAI_COMPATIBLE_PROVIDERS_WITH_NATIVE_AUDIO_TRANSCRIPTION: Final = frozenset({"xai"}) - -OPENAI_AUDIO_TRANSCRIPTION_PROVIDERS: Final = frozenset( - {"openai"} | (frozenset(openai_compatible_providers) - OPENAI_COMPATIBLE_PROVIDERS_WITH_NATIVE_AUDIO_TRANSCRIPTION) -) +OPENAI_AUDIO_TRANSCRIPTION_PROVIDERS: Final = frozenset({"openai"} | frozenset(openai_compatible_providers)) openai_text_completion_compatible_providers: Final[list] = [ # providers that support `/v1/completions` "together_ai", diff --git a/litellm/llms/base_llm/audio_transcription/transformation.py b/litellm/llms/base_llm/audio_transcription/transformation.py index b323c4812b5..2296909cfe1 100644 --- a/litellm/llms/base_llm/audio_transcription/transformation.py +++ b/litellm/llms/base_llm/audio_transcription/transformation.py @@ -52,6 +52,15 @@ class BaseAudioTranscriptionConfig(BaseConfig, ABC): """ return False + @property + def has_native_transcription_endpoint(self) -> bool: + """ + Opt-in for OpenAI-compatible providers whose transcription lives on a + non-OpenAI route: when True the request skips the OpenAI SDK transport + and goes through this config via the shared http handler. + """ + return False + def get_complete_url( self, api_base: str | None, diff --git a/litellm/llms/xai/audio_transcription/transformation.py b/litellm/llms/xai/audio_transcription/transformation.py index 7b648fc8084..feeabed0d9c 100644 --- a/litellm/llms/xai/audio_transcription/transformation.py +++ b/litellm/llms/xai/audio_transcription/transformation.py @@ -63,6 +63,10 @@ class XAIAudioTranscriptionConfig(BaseAudioTranscriptionConfig): def custom_llm_provider(self) -> str: return litellm.LlmProviders.XAI.value + @property + def has_native_transcription_endpoint(self) -> bool: + return True + def get_supported_openai_params( self, model: str ) -> list[OpenAIAudioTranscriptionOptionalParams]: # mutable-ok: base class signature returns list diff --git a/litellm/main.py b/litellm/main.py index bd10c3924f7..38184db1d10 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -7831,6 +7831,10 @@ def transcription( provider=LlmProviders(custom_llm_provider), ) + uses_openai_transport: Final = custom_llm_provider in OPENAI_AUDIO_TRANSCRIPTION_PROVIDERS and not ( + provider_config is not None and provider_config.has_native_transcription_endpoint + ) + if custom_llm_provider in AZURE_OPENAI_AUDIO_PROVIDERS and provider_config is None: # azure configs api_base = api_base or litellm.api_base or get_secret_str("AZURE_API_BASE") @@ -7860,7 +7864,7 @@ def transcription( litellm_params=litellm_params_dict, custom_llm_provider=custom_llm_provider, ) - elif custom_llm_provider in OPENAI_AUDIO_TRANSCRIPTION_PROVIDERS: + elif uses_openai_transport: api_base = ( api_base or litellm.api_base From c635c35b3d2968dbf76ebec98eee833e2b5c8a0f Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Fri, 18 Sep 2026 18:47:36 -0700 Subject: [PATCH 131/144] fix(rust): keep native OCR on the proxy by declining only a supplied client The proxy attaches its shared aiohttp session to every request as shared_session, so declining on it sent every proxy OCR call to Python, which never uses that session for OCR. aclient_session is a litellm global and never a call argument, so that check could not match. The proxy-shaped lifecycle test now asserts the call was served by Rust --- litellm-rust/crates/python-bridge/src/http.rs | 46 +++++++++---------- tests/test_litellm_rust/ocr/test_lifecycle.py | 1 + 2 files changed, 22 insertions(+), 25 deletions(-) diff --git a/litellm-rust/crates/python-bridge/src/http.rs b/litellm-rust/crates/python-bridge/src/http.rs index 77542855fec..c5952c53132 100644 --- a/litellm-rust/crates/python-bridge/src/http.rs +++ b/litellm-rust/crates/python-bridge/src/http.rs @@ -12,8 +12,6 @@ use crate::{errors::RustBridgeDeclined, python_settings::PythonSettings}; static POOL: LazyLock = LazyLock::new(|| HttpClientPool::new(Arc::new(PublicDnsResolver))); -const LIVE_CLIENT_ARGUMENTS: [&str; 3] = ["client", "shared_session", "aclient_session"]; - pub(crate) fn pool() -> &'static HttpClientPool { &POOL } @@ -23,7 +21,7 @@ pub(crate) fn call_config( kwargs: &Bound<'_, PyDict>, asynchronous: bool, ) -> PyResult { - decline_live_clients(kwargs)?; + decline_live_client(kwargs)?; decline_custom_url_policy(&PythonSettings::UrlPolicy.read(py)?)?; let configured = settings(&PythonSettings::Http.read(py)?)? .with_environment(&|name| std::env::var(name).ok()); @@ -53,13 +51,14 @@ fn for_call( } } -pub(crate) fn decline_live_clients(kwargs: &Bound<'_, PyDict>) -> PyResult<()> { - for name in LIVE_CLIENT_ARGUMENTS { - if kwargs.get_item(name)?.is_some_and(|value| !value.is_none()) { - return Err(RustBridgeDeclined::new_err(format!( - "{name} is a live Python HTTP client and cannot be used by the Rust route" - ))); - } +fn decline_live_client(kwargs: &Bound<'_, PyDict>) -> PyResult<()> { + if kwargs + .get_item("client")? + .is_some_and(|value| !value.is_none()) + { + return Err(RustBridgeDeclined::new_err( + "client is a live Python HTTP client and cannot be used by the Rust route", + )); } Ok(()) } @@ -362,32 +361,29 @@ user_agent='litellm/9.9.9', assert_eq!(config.trust_proxy_env, expected); } - #[rstest] - #[case::client("client")] - #[case::shared_session("shared_session")] - #[case::aclient_session("aclient_session")] - fn live_python_clients_decline_before_dispatch(#[case] name: &str) { + #[test] + fn live_python_client_declines_before_dispatch() { Python::initialize(); Python::attach(|py| { let kwargs = PyDict::new(py); kwargs - .set_item(name, py.eval(c"object()", None, None).unwrap()) + .set_item("client", py.eval(c"object()", None, None).unwrap()) .unwrap(); - let error = decline_live_clients(&kwargs).unwrap_err(); + let error = decline_live_client(&kwargs).unwrap_err(); assert!(error.is_instance_of::(py)); - assert!(error.value(py).to_string().contains(name)); }); } - #[test] - fn none_valued_client_arguments_are_not_live_clients() { + #[rstest] + #[case::absent_client("{}")] + #[case::none_client("{'client': None}")] + #[case::proxy_shared_session("{'shared_session': object()}")] + fn calls_without_a_python_client_stay_on_the_rust_route(#[case] kwargs: &str) { Python::initialize(); Python::attach(|py| { - let kwargs = PyDict::new(py); - for name in LIVE_CLIENT_ARGUMENTS { - kwargs.set_item(name, py.None()).unwrap(); - } - decline_live_clients(&kwargs).unwrap(); + let source = std::ffi::CString::new(kwargs).unwrap(); + let kwargs = py.eval(&source, None, None).unwrap(); + decline_live_client(kwargs.cast::().unwrap()).unwrap(); }); } } diff --git a/tests/test_litellm_rust/ocr/test_lifecycle.py b/tests/test_litellm_rust/ocr/test_lifecycle.py index 5fca927bea3..264a666c685 100644 --- a/tests/test_litellm_rust/ocr/test_lifecycle.py +++ b/tests/test_litellm_rust/ocr/test_lifecycle.py @@ -39,6 +39,7 @@ async def test_proxy_metadata_remains_python_owned(ocr_server: RecordingServer) ) events: Final = await recorder.wait_for_async("async_log_success_event") assert response.pages[0].markdown == "native OCR response" + assert response._hidden_params["additional_headers"]["x-litellm-rust"] == "true" assert events[0].kwargs["litellm_params"]["metadata"]["user_api_key_auth"].user_id == "ocr-user" assert "metadata" not in ocr_server.requests[0].body From 0119f5001581eea9e202ddc6e8b6543da9424422 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Fri, 18 Sep 2026 18:47:36 -0700 Subject: [PATCH 132/144] fix(rust): restore the 10s connect timeout and share media clients across proxy settings Python OCR passes the call timeout per request, so its connect timeout is the call timeout and never the 5s handler default. 10s is what every Rust route uses on main. The media client never uses a proxy, so trust_proxy_env no longer splits its pool key --- litellm-rust/crates/http/src/pool.rs | 17 ++++++++++++++++- litellm-rust/crates/http/src/settings.rs | 2 +- 2 files changed, 17 insertions(+), 2 deletions(-) diff --git a/litellm-rust/crates/http/src/pool.rs b/litellm-rust/crates/http/src/pool.rs index 0d9b1abf504..b51e6711419 100644 --- a/litellm-rust/crates/http/src/pool.rs +++ b/litellm-rust/crates/http/src/pool.rs @@ -51,6 +51,7 @@ impl HttpClientPool { let effective = match variant { ClientVariant::Media => HttpClientConfig { client_certificate: None, + trust_proxy_env: false, ..config.clone() }, ClientVariant::Provider | ClientVariant::NoRedirect => config.clone(), @@ -86,7 +87,6 @@ impl HttpClientPool { ClientVariant::NoRedirect => builder.redirect(reqwest::redirect::Policy::none()), ClientVariant::Media => builder .redirect(reqwest::redirect::Policy::none()) - .no_proxy() .dns_resolver2(Arc::clone(&self.media_resolver)), } } @@ -206,6 +206,21 @@ mod tests { assert_eq!(connections.load(Ordering::SeqCst), 2); } + #[tokio::test] + async fn media_clients_are_shared_across_proxy_settings_they_never_use() { + let (address, connections, _) = serve("HTTP/1.1 204 No Content").await; + let pool = HttpClientPool::new(Arc::new(FixedResolver(address))); + let url = format!("http://media.invalid:{}/doc", address.port()); + for trust_proxy_env in [true, false] { + let config = HttpClientConfig { + trust_proxy_env, + ..config("a") + }; + get(&pool, &config, ClientVariant::Media, &url).await; + } + assert_eq!(connections.load(Ordering::SeqCst), 1); + } + #[test] fn media_variant_never_loads_the_client_certificate() { let pool = pool(); diff --git a/litellm-rust/crates/http/src/settings.rs b/litellm-rust/crates/http/src/settings.rs index c572c56ef3a..8aaf7f21f2c 100644 --- a/litellm-rust/crates/http/src/settings.rs +++ b/litellm-rust/crates/http/src/settings.rs @@ -50,7 +50,7 @@ impl Default for HttpSettings { user_agent: None, trust_proxy_env: false, ignore_proxy_env: false, - connect_timeout: Duration::from_secs(5), + connect_timeout: Duration::from_secs(10), } } } From bf7d1c07330a7d5676fb4a461f7b7f72bc4098d3 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Sat, 19 Sep 2026 02:30:35 +0000 Subject: [PATCH 133/144] chore: consolidate CLAUDE.md into AGENTS.md Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- AGENTS.md | 130 +++++++++++++++++- CLAUDE.md | 129 ----------------- CONTRIBUTING.md | 2 +- GEMINI.md | 2 +- litellm-rust/crates/python-bridge/AGENTS.md | 44 +++++- litellm-rust/crates/python-bridge/CLAUDE.md | 43 ------ .../proxy/_experimental/mcp_server/AGENTS.md | 9 +- .../proxy/_experimental/mcp_server/CLAUDE.md | 1 - .../check_e2e_no_raw_requests.py | 2 +- .../test_e2e_changed_gate.py | 2 +- tests/e2e/{CLAUDE.md => AGENTS.md} | 2 +- tests/e2e/CONTRIBUTING.md | 4 +- tests/e2e/batches/COVERAGE.md | 2 +- tests/e2e/claude_code/cron_vm/run_daily.sh | 2 +- tests/e2e/coverage_registry/README.md | 2 +- tests/e2e/coverage_registry/__init__.py | 2 +- tests/e2e/coverage_registry/mcp.yaml | 2 +- .../realtime/REALTIME_COVERAGE_MATRIX.md | 2 +- .../e2e/llm_translation/realtime/conftest.py | 2 +- .../realtime/realtime_client.py | 4 +- ui/litellm-dashboard/{CLAUDE.md => AGENTS.md} | 0 21 files changed, 193 insertions(+), 195 deletions(-) delete mode 100644 CLAUDE.md delete mode 100644 litellm-rust/crates/python-bridge/CLAUDE.md delete mode 100644 litellm/proxy/_experimental/mcp_server/CLAUDE.md rename tests/e2e/{CLAUDE.md => AGENTS.md} (99%) rename ui/litellm-dashboard/{CLAUDE.md => AGENTS.md} (100%) diff --git a/AGENTS.md b/AGENTS.md index a1e8f6f618d..cade08bdd02 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -1,3 +1,131 @@ -Read @CLAUDE.md for coding guidelines +Do not write comments unless they are any of: +- absolutely necessary to explain some very complex business logic (in which case, keep it concise and clear) +- used as an input for tools to read and act on. For example: + - entries in `.git-blame-ignore-revs` saying which commit is excluded from git blame + - a lint or type checker suppression like `# mutable-ok` or `# pyright: ignore[reportArgumentType] # ` when introducing a truly unavoidable violation +- a TODO or FIXME + - Not great to have those, but if it's unavoidable, make sure to include a strong, concise reason for why it's there or, better yet, link to a GitHub issue for the follow-up work + +Explanation: The point of this rule is to keep out AI slop comments. AI writes way too many and way too verbose comments. Code comments are, in a way, a violation of DRY code. You must update logic in two locations to change the code, and "hard to change" is literally the definition of tech debt. We should instead aim to write code that is intuitive and clear, even at a glance, to the reader, being both easy to maintain and high performance + +Don't assume that the existing code is correct or the right way of doing things / good coding patterns. In fact, there are a lot of bad coding practices, overly complex code, code smells, etc. If something doesn't look right, speak up. Feel free to break existing patterns or question weird existing code to make new code high quality, as in: + +- correct +- secure +- performant +- readable +- easy to maintain/change +- modern + +In descending order of importance + +When adding new features, add meaningful tests. Don't add tests that don't check anything substantial and is there just to make the code coverage pass. Yes, code coverage is important, but I'd rather have no signal whether the code is working than tests that don't fail when code is broken. The goal is to have tests that would fail before the feature was added/if the code was mutated in a way that breaks the feature and succeed only when the feature is fully working. I should run mutation testing and see > 90% kill rate + +Same thing for bug fixes. The tests should make it so that this specific bug can never happen again without failing tests (i.e., regression) + +Never test structure of code only function of it + +A test must only fail when litellm code changes. Never pin facts we don't own (a vendor's price, a third party's field, an upstream default, today's date) as literals or as "X must be absent"; assert the invariant our code guarantees instead, e.g. two rows agree, a value is within range, a field is derived from another. If an outside fact is truly load-bearing, cite its source and date next to the assertion so a reader can tell stale from broken + +`tests/test_litellm/` mirrors `litellm/` in a parallel path (see `tests/test_litellm/readme.md`). Name tests `test_.py`, but always match the existing test file in the directory you touch — many provider dirs use longer descriptive names (e.g. `test_anthropic_chat_transformation.py`) to avoid ambiguity across sibling folders. For bug fixes, extend the existing mapped test file rather than creating a new one. Only create a new test file for a new feature (provider, endpoint, or transformation module) that has no mapped test yet, following that directory's naming convention (or `test_.py` if you're the first test there). One focused regression test beats many shallow ones + +End-to-end tests belong in `tests/e2e/` and must follow the harness conventions documented in that directory's `AGENTS.md` + +When creating PRs, target the repository's current default branch for both internal and external / OSS contributions. Check it with `python3 scripts/default_branch.py --branch` instead of assuming a branch name or relying on cached `origin/HEAD` + +When writing a PR body, treat the comments and imperative instructions inside .github/pull_request_template.md as rules to follow, not just layout. Agent harnesses may strip HTML comments from copies of that file injected into context, so read .github/pull_request_template.md from disk before writing a PR body to make sure you see every comment rule + +Same applies for filing bug reports and feature requests, with .github/ISSUE_TEMPLATE/bug_report.yml and .github/ISSUE_TEMPLATE/feature_request.yml, respectively + +If you're resolving a linear ticket, in the "## Linear ticket" section of the PR, say "Resolves LIT-1234", replacing "LIT-1234" with the actual ticket id that you're resolving. If you don't have the ticket id, don't make one up or search for it. Just leave the section blank + +Never use `pytest` commands or the like as "Screenshots / Proof of Fix". We prefer curl'ing a live proxy instance running on localhost:4000 (I like to run it with `python litellm/proxy/proxy_cli.py --config litellm/proxy/dev_config.yaml --detailed_debug --reload --use_v2_migration_resolver 2>&1 | tee litellm.log`; the Admin UI dev server is `npm run dev` in `ui/litellm-dashboard`, served on port 3000) and showing both the command run and the output. Also, it should hit real LLM provider APIs, not mocks, and cost real $$$ because that is the most realistic test. The proof of fix should be exactly what the end user / customer would see / do. The run logs in PR #27703 is a prime example of how to do it (not a huge fan of using a python test script that future me and the team will have no visibility into; I prefer just curl commands or a short list of bash commands (e.g., using `for`)). If it's a UI thing, or the main use case runs through a headful agentic coding tool like Claude Code or Codex, drive that surface yourself and embed your own before and after screenshots of it in the PR (the Admin UI page, or what the coding tool shows), next to an ordered list of the URLs to go to (e.g., http://localhost:4000/ui/?page=logs), where to click, and what fields to fill out so a reviewer can reproduce it + +If you ever write any human-facing text (pull requests, issues, commit messages, discussion posts, github comments, release notes, docs, etc.), always follow these guidelines to sound less AI-y: +- don't use emojis +- don't use "—". Instead, reach for ",", ".", conjunction words, ":", ";", etc. in descending order of preference: vary among them, weighted toward the front of the list, and skip "," where it would cause a comma splice or the sentence is getting long. Overusing any one of them, ";" especially, also feels AI-y. A word cap does not penalize you for adding more sentences: when writing under tight word budgets, prefer a period split or a conjunction over ";", and keep to at most one ";" per message +- don't use the pattern "It's not X, it's Y", "You're not X, you're Y", etc. +- unless explicitly asked, don't use bulleted or numbered lists unless it would be nonsensical not to. Instead, prefer prose +- don't add a trailing "." at the end of paragraphs (just like this file). That means every paragraph, not just the last one (of the markdown file, PR description, GitHub comment, etc.). Rule of thumb: if you're adding new line(s) before the next sentence, don't add a "." +- don't use →. Instead, prefer not to use arrows, and if need be, use -> instead +- use plain, simple, everyday engineering language: the common phrase engineers actually say over rare compact phrasing, in grammatically complete sentences. When explicitly asked to use bullets or ordered lists and structure legitimately helps the reader, prefer nested bullets (any depth is fine) over dense lines in a flat structure + +Don't hesitate to use values in .env to get needed API keys and other secrets, as long as you never add them to conversation history, commit them, or include them in GitHub issues / PRs + +Python max line length is 120, not 88 + +Never edit or commit `ruff-strict-budget.json`, `type-discipline-budget.json`, `basedpyright-code-budget.json`, or `test-quality-budget.json` on a PR branch, and don't run `make lint-budget-update` there. A scheduled Devin automation lowers the limits on the default branch in its own PR by exactly what landed since the last ratchet, so concurrent PRs don't fight over the same `"limit"` lines. Keep the hosted automation's target in sync when the repository default changes. If your branch already carries a budget edit, drop it before opening the PR + +`make check` (f.k.a. `make pre-commit`, which still works identically as an alias) saves its complete output to a log file in .git (overwriting previous logs) and prints that path as its first and last output lines. To inspect a run, read or grep that log instead of re-running the multi-minute checks just to see a different slice + +`make check`, `make lint`, `scripts/pre_commit_lint.sh`, and the standalone budget gates (`scripts/ruff_strict_gate.py`, `scripts/type_discipline_gate.py`, `scripts/type_check_gate.py`) each hold one of 2 machine-wide slots, so when other sessions or worktrees on the same box are already running heavy work, yours prints "all N machine-wide slots are busy; queueing" and then stays quiet until a slot frees. Give the command a long timeout and let it wait rather than killing it, retrying it, or assuming it hung. Don't change the # of machine-wide slots or make it unlimited by setting `LITELLM_GATE_SLOTS=0` + +If you're trying to create a new function that relies on untyped stuff, instead of adding more Any's and pushing `reportAny` / `reportExplicitAny` closer to their basedpyright ceilings, just validate it in the caller with Pydantic (a model or `TypeAdapter` that returns the typed thing or raises will do) and then pass the now typed variable in + +If you get an LIT001 or LIT002 fail, refactor the code to follow functional programming best practices rather than introducing mutable data structures. For example, build values in one shot with comprehensions or generators wrapped in `tuple()` / `MappingProxyType()` / `frozenset()` instead of seeding an empty `list`/`dict`/`set` and mutating it over time. Ideally, `# mutable-ok` is never used; reach for it only as a genuine last resort when an immutable rewrite is truly impossible, and always pair it with a real reason + +Every lint or type suppression must name the exact rule inside brackets and carry a reason comment, e.g. `# pyright: ignore[reportArgumentType] # stubs lack async overload` or `# noqa: TID251 # `. `# type: ignore` is banned (LIT009): pyrightconfig.json sets `enableTypeIgnoreComments` to false, so it silently does nothing + +Commit and push your work when you're done without asking + +When referencing or running models (coding, QA'ing, writing docs, writing tests, etc.), use the latest model in that model family unless otherwise specified; treat your training knowledge, memories, configs, and tests as stale, and determine the family's latest with model_prices_and_context_window.json or the web + +Always pull before starting any work. The checkout or worktree may be sitting on a stale branch + +If you're an internal contributor, when creating a new PR, the typical flow is to branch off the repository's current default branch and create a branch prefixed with litellm_. Do not create a branch prefixed with claude/ and generally do not have / in your branch names + +Do not add `Co-Authored-By: Claude` or any Claude attribution to commit messages. Never use a `claude/` prefix or put a `/` in a branch name. Do not add "Generated with Claude Code" (or any similar attribution) to PR descriptions or comments. Do not create a new PR/branch off the existing PR to fix/add something that is related and could've just been committed directly to the existing PR's branch + +When working on a PR, keep the PR description in sync with new commits being made + +All GitHub comments must be human-readable and 15-25 words max + +Monkeypatching attributes of a class to do testing is an anti-pattern. Prefer dependency-injecting things into classes. That way, at unit test time, you can pass a mocked dependency in + +Do not put names of customers or customer company names in code, PR descriptions, issue bodies, etc. This means never mention literally any company name. Especially if you're about to say a sentence mentioning that the reason the PR exists was a feature/model/bug fix/etc. requested by a company. That's the indication that you should replace that company name with "the customer". e.g. not "Model request from Acme (Pylon #1234)" but "Model request from a customer (Pylon #1234)". This is because the codebase is public. The only exception is for publicly known providers or vendors such as OpenAI, Anthropic, AWS Bedrock, etc. only IF we're adding support for that provider/vendor in general and NOT if that PR or whatnot was a request by one of them, and they're actually one of our customers. + +CI supply-chain safety: Never pipe a remote script into a shell (`curl ... | bash`, `wget ... | sh`); download the artifact to a file, verify its SHA-256 checksum, then install. Pin every external tool to a specific version with a full URL (not `latest` or `stable`). Verify checksums for all downloaded binaries, using the provider's official `.sha256` / `.sha256sum` sidecar when available. These rules apply to every download in CI + +Prisma migrations apply synchronously at proxy boot, before it serves traffic, so a migration must only change schema, never rewrite rows. No `UPDATE`, `DELETE` or `MERGE`, and no `INSERT ... SELECT`: on a spend-log-sized table any of those is minutes of downtime plus a doubled heap that plain autovacuum won't give back. `tests/code_coverage_tests/check_migrations_no_data_rewrites.py` enforces this. When a rewrite is genuinely bounded and has to ship inside the migration, mark the statement `-- data-migration-ok: ` + +Follow these coding conventions for new/updated code (a three-line fix in a legacy file shouldn't trigger huge drive-by refactors): + +- Composition over inheritance +- Never-nester: early returns over deep nesting +- Don't throw; model failures as values (One function (e.g., raise_public) maps error union to existing public exception contracts via exhaustive match + assert_never) +- No mutation; don't reassign variables, global or local. Instead of mutable lists and dicts, prefer tuples, frozen dataclasses (with slots=True), `MappingProxyType`, etc. + - Annotate every variable with `: Final` (LIT010). Unpacking and walrus targets cannot carry the annotation, so they are implicitly final. Don't rebind them. Never rebind or mutate function parameters (LIT011); `self`/`cls` attribute stores are the exception. If rebinding or in-place mutation is truly unavoidable, suppress with `# rebind-ok: ` + - Qualify every TypedDict field with `ReadOnly[...]` (LIT012), which nests freely with `Required` / `NotRequired` / `Annotated` in any order. If making the key writable is truly unavoidable, suppress with `# writable-ok: ` +- Use dependency injection +- Fully typed; no `Any` or coarse types like `dict[str, Any]` or just `dict`. Every function parameter must be strongly typed +- Use tagged unions + match +- No monster files or god objects +- No file sprawl: deliberate file and folder structure +- Standard over hand-rolled: use the official SDK or a library where one exists; where none does, follow industry standards instead of inventing local conventions +- API-fragmentation-aware: when logic must branch on which API surface produced or consumes data (e.g. chat completions vs Anthropic Messages vs Responses API shapes), proactively look for an existing shared helper (e.g. `litellm_core_utils/prompt_templates/factory.py`) before writing per-surface parsing in the new module; if none exists, add one there instead of duplicating the same format-detection logic in every new guardrail/integration + +Follow conventional commits for commit names and PR titles + +## Think Before Coding + +**Don't assume. Don't hide confusion. Surface tradeoffs** + +Before implementing: +- State your assumptions explicitly. If uncertain, ask +- If multiple interpretations exist, present them. Don't pick silently +- If a simpler approach exists, say so. Push back when warranted +- If something is unclear, stop. Name what's confusing. Ask + +## Simplicity First + +**Minimum code that solves the problem. Nothing speculative** + +- No features beyond what was asked +- No abstractions for single-use code +- No "flexibility" or "configurability" that wasn't requested +- No error handling for impossible scenarios +- If you write 200 lines and it could be 50, rewrite it + +Ask yourself: "Would a senior engineer say this is overcomplicated?" If yes, simplify Before requesting maintainer review, verify the current PR tip passes required CI and code coverage, meets Greptile confidence of at least 4/5, and has acceptable Veria and Bugbot reviews. Inspect warnings and findings, fix actionable issues, and rerun the affected checks and reviewers after changes. Record evidence for any false positive or unavailable review; never treat a pending or missing bot result as a pass. Do not lower coverage thresholds or lint budgets to satisfy a check diff --git a/CLAUDE.md b/CLAUDE.md deleted file mode 100644 index b9753ab864b..00000000000 --- a/CLAUDE.md +++ /dev/null @@ -1,129 +0,0 @@ -Do not write comments unless they are any of: -- absolutely necessary to explain some very complex business logic (in which case, keep it concise and clear) -- used as an input for tools to read and act on. For example: - - entries in `.git-blame-ignore-revs` saying which commit is excluded from git blame - - a lint or type checker suppression like `# mutable-ok` or `# pyright: ignore[reportArgumentType] # ` when introducing a truly unavoidable violation -- a TODO or FIXME - - Not great to have those, but if it's unavoidable, make sure to include a strong, concise reason for why it's there or, better yet, link to a GitHub issue for the follow-up work - -Explanation: The point of this rule is to keep out AI slop comments. AI writes way too many and way too verbose comments. Code comments are, in a way, a violation of DRY code. You must update logic in two locations to change the code, and "hard to change" is literally the definition of tech debt. We should instead aim to write code that is intuitive and clear, even at a glance, to the reader, being both easy to maintain and high performance - -Don't assume that the existing code is correct or the right way of doing things / good coding patterns. In fact, there are a lot of bad coding practices, overly complex code, code smells, etc. If something doesn't look right, speak up. Feel free to break existing patterns or question weird existing code to make new code high quality, as in: - -- correct -- secure -- performant -- readable -- easy to maintain/change -- modern - -In descending order of importance - -When adding new features, add meaningful tests. Don't add tests that don't check anything substantial and is there just to make the code coverage pass. Yes, code coverage is important, but I'd rather have no signal whether the code is working than tests that don't fail when code is broken. The goal is to have tests that would fail before the feature was added/if the code was mutated in a way that breaks the feature and succeed only when the feature is fully working. I should run mutation testing and see > 90% kill rate - -Same thing for bug fixes. The tests should make it so that this specific bug can never happen again without failing tests (i.e., regression) - -Never test structure of code only function of it - -A test must only fail when litellm code changes. Never pin facts we don't own (a vendor's price, a third party's field, an upstream default, today's date) as literals or as "X must be absent"; assert the invariant our code guarantees instead, e.g. two rows agree, a value is within range, a field is derived from another. If an outside fact is truly load-bearing, cite its source and date next to the assertion so a reader can tell stale from broken - -`tests/test_litellm/` mirrors `litellm/` in a parallel path (see `tests/test_litellm/readme.md`). Name tests `test_.py`, but always match the existing test file in the directory you touch — many provider dirs use longer descriptive names (e.g. `test_anthropic_chat_transformation.py`) to avoid ambiguity across sibling folders. For bug fixes, extend the existing mapped test file rather than creating a new one. Only create a new test file for a new feature (provider, endpoint, or transformation module) that has no mapped test yet, following that directory's naming convention (or `test_.py` if you're the first test there). One focused regression test beats many shallow ones - -End-to-end tests belong in `tests/e2e/` and must follow the harness conventions documented in that directory's `CLAUDE.md` - -When creating PRs, target the repository's current default branch for both internal and external / OSS contributions. Check it with `python3 scripts/default_branch.py --branch` instead of assuming a branch name or relying on cached `origin/HEAD` - -When writing a PR body, treat the comments and imperative instructions inside .github/pull_request_template.md as rules to follow, not just layout. Agent harnesses may strip HTML comments from copies of that file injected into context, so read .github/pull_request_template.md from disk before writing a PR body to make sure you see every comment rule - -Same applies for filing bug reports and feature requests, with .github/ISSUE_TEMPLATE/bug_report.yml and .github/ISSUE_TEMPLATE/feature_request.yml, respectively - -If you're resolving a linear ticket, in the "## Linear ticket" section of the PR, say "Resolves LIT-1234", replacing "LIT-1234" with the actual ticket id that you're resolving. If you don't have the ticket id, don't make one up or search for it. Just leave the section blank - -Never use `pytest` commands or the like as "Screenshots / Proof of Fix". We prefer curl'ing a live proxy instance running on localhost:4000 (I like to run it with `python litellm/proxy/proxy_cli.py --config litellm/proxy/dev_config.yaml --detailed_debug --reload --use_v2_migration_resolver 2>&1 | tee litellm.log`; the Admin UI dev server is `npm run dev` in `ui/litellm-dashboard`, served on port 3000) and showing both the command run and the output. Also, it should hit real LLM provider APIs, not mocks, and cost real $$$ because that is the most realistic test. The proof of fix should be exactly what the end user / customer would see / do. The run logs in PR #27703 is a prime example of how to do it (not a huge fan of using a python test script that future me and the team will have no visibility into; I prefer just curl commands or a short list of bash commands (e.g., using `for`)). If it's a UI thing, or the main use case runs through a headful agentic coding tool like Claude Code or Codex, drive that surface yourself and embed your own before and after screenshots of it in the PR (the Admin UI page, or what the coding tool shows), next to an ordered list of the URLs to go to (e.g., http://localhost:4000/ui/?page=logs), where to click, and what fields to fill out so a reviewer can reproduce it - -If you ever write any human-facing text (pull requests, issues, commit messages, discussion posts, github comments, release notes, docs, etc.), always follow these guidelines to sound less AI-y: -- don't use emojis -- don't use "—". Instead, reach for ",", ".", conjunction words, ":", ";", etc. in descending order of preference: vary among them, weighted toward the front of the list, and skip "," where it would cause a comma splice or the sentence is getting long. Overusing any one of them, ";" especially, also feels AI-y. A word cap does not penalize you for adding more sentences: when writing under tight word budgets, prefer a period split or a conjunction over ";", and keep to at most one ";" per message -- don't use the pattern "It's not X, it's Y", "You're not X, you're Y", etc. -- unless explicitly asked, don't use bulleted or numbered lists unless it would be nonsensical not to. Instead, prefer prose -- don't add a trailing "." at the end of paragraphs (just like this file). That means every paragraph, not just the last one (of the markdown file, PR description, GitHub comment, etc.). Rule of thumb: if you're adding new line(s) before the next sentence, don't add a "." -- don't use →. Instead, prefer not to use arrows, and if need be, use -> instead -- use plain, simple, everyday engineering language: the common phrase engineers actually say over rare compact phrasing, in grammatically complete sentences. When explicitly asked to use bullets or ordered lists and structure legitimately helps the reader, prefer nested bullets (any depth is fine) over dense lines in a flat structure - -Don't hesitate to use values in .env to get needed API keys and other secrets, as long as you never add them to conversation history, commit them, or include them in GitHub issues / PRs - -Python max line length is 120, not 88 - -Never edit or commit `ruff-strict-budget.json`, `type-discipline-budget.json`, `basedpyright-code-budget.json`, or `test-quality-budget.json` on a PR branch, and don't run `make lint-budget-update` there. A scheduled Devin automation lowers the limits on the default branch in its own PR by exactly what landed since the last ratchet, so concurrent PRs don't fight over the same `"limit"` lines. Keep the hosted automation's target in sync when the repository default changes. If your branch already carries a budget edit, drop it before opening the PR - -`make check` (f.k.a. `make pre-commit`, which still works identically as an alias) saves its complete output to a log file in .git (overwriting previous logs) and prints that path as its first and last output lines. To inspect a run, read or grep that log instead of re-running the multi-minute checks just to see a different slice - -`make check`, `make lint`, `scripts/pre_commit_lint.sh`, and the standalone budget gates (`scripts/ruff_strict_gate.py`, `scripts/type_discipline_gate.py`, `scripts/type_check_gate.py`) each hold one of 2 machine-wide slots, so when other sessions or worktrees on the same box are already running heavy work, yours prints "all N machine-wide slots are busy; queueing" and then stays quiet until a slot frees. Give the command a long timeout and let it wait rather than killing it, retrying it, or assuming it hung. Don't change the # of machine-wide slots or make it unlimited by setting `LITELLM_GATE_SLOTS=0` - -If you're trying to create a new function that relies on untyped stuff, instead of adding more Any's and pushing `reportAny` / `reportExplicitAny` closer to their basedpyright ceilings, just validate it in the caller with Pydantic (a model or `TypeAdapter` that returns the typed thing or raises will do) and then pass the now typed variable in - -If you get an LIT001 or LIT002 fail, refactor the code to follow functional programming best practices rather than introducing mutable data structures. For example, build values in one shot with comprehensions or generators wrapped in `tuple()` / `MappingProxyType()` / `frozenset()` instead of seeding an empty `list`/`dict`/`set` and mutating it over time. Ideally, `# mutable-ok` is never used; reach for it only as a genuine last resort when an immutable rewrite is truly impossible, and always pair it with a real reason - -Every lint or type suppression must name the exact rule inside brackets and carry a reason comment, e.g. `# pyright: ignore[reportArgumentType] # stubs lack async overload` or `# noqa: TID251 # `. `# type: ignore` is banned (LIT009): pyrightconfig.json sets `enableTypeIgnoreComments` to false, so it silently does nothing - -Commit and push your work when you're done without asking - -When referencing or running models (coding, QA'ing, writing docs, writing tests, etc.), use the latest model in that model family unless otherwise specified; treat your training knowledge, memories, configs, and tests as stale, and determine the family's latest with model_prices_and_context_window.json or the web - -Always pull before starting any work. The checkout or worktree may be sitting on a stale branch - -If you're an internal contributor, when creating a new PR, the typical flow is to branch off the repository's current default branch and create a branch prefixed with litellm_. Do not create a branch prefixed with claude/ and generally do not have / in your branch names - -Do not add `Co-Authored-By: Claude` or any Claude attribution to commit messages. Never use a `claude/` prefix or put a `/` in a branch name. Do not add "Generated with Claude Code" (or any similar attribution) to PR descriptions or comments. Do not create a new PR/branch off the existing PR to fix/add something that is related and could've just been committed directly to the existing PR's branch - -When working on a PR, keep the PR description in sync with new commits being made - -All GitHub comments must be human-readable and 15-25 words max - -Monkeypatching attributes of a class to do testing is an anti-pattern. Prefer dependency-injecting things into classes. That way, at unit test time, you can pass a mocked dependency in - -Do not put names of customers or customer company names in code, PR descriptions, issue bodies, etc. This means never mention literally any company name. Especially if you're about to say a sentence mentioning that the reason the PR exists was a feature/model/bug fix/etc. requested by a company. That's the indication that you should replace that company name with "the customer". e.g. not "Model request from Acme (Pylon #1234)" but "Model request from a customer (Pylon #1234)". This is because the codebase is public. The only exception is for publicly known providers or vendors such as OpenAI, Anthropic, AWS Bedrock, etc. only IF we're adding support for that provider/vendor in general and NOT if that PR or whatnot was a request by one of them, and they're actually one of our customers. - -CI supply-chain safety: Never pipe a remote script into a shell (`curl ... | bash`, `wget ... | sh`); download the artifact to a file, verify its SHA-256 checksum, then install. Pin every external tool to a specific version with a full URL (not `latest` or `stable`). Verify checksums for all downloaded binaries, using the provider's official `.sha256` / `.sha256sum` sidecar when available. These rules apply to every download in CI - -Prisma migrations apply synchronously at proxy boot, before it serves traffic, so a migration must only change schema, never rewrite rows. No `UPDATE`, `DELETE` or `MERGE`, and no `INSERT ... SELECT`: on a spend-log-sized table any of those is minutes of downtime plus a doubled heap that plain autovacuum won't give back. `tests/code_coverage_tests/check_migrations_no_data_rewrites.py` enforces this. When a rewrite is genuinely bounded and has to ship inside the migration, mark the statement `-- data-migration-ok: ` - -Follow these coding conventions for new/updated code (a three-line fix in a legacy file shouldn't trigger huge drive-by refactors): - -- Composition over inheritance -- Never-nester: early returns over deep nesting -- Don't throw; model failures as values (One function (e.g., raise_public) maps error union to existing public exception contracts via exhaustive match + assert_never) -- No mutation; don't reassign variables, global or local. Instead of mutable lists and dicts, prefer tuples, frozen dataclasses (with slots=True), `MappingProxyType`, etc. - - Annotate every variable with `: Final` (LIT010). Unpacking and walrus targets cannot carry the annotation, so they are implicitly final. Don't rebind them. Never rebind or mutate function parameters (LIT011); `self`/`cls` attribute stores are the exception. If rebinding or in-place mutation is truly unavoidable, suppress with `# rebind-ok: ` - - Qualify every TypedDict field with `ReadOnly[...]` (LIT012), which nests freely with `Required` / `NotRequired` / `Annotated` in any order. If making the key writable is truly unavoidable, suppress with `# writable-ok: ` -- Use dependency injection -- Fully typed; no `Any` or coarse types like `dict[str, Any]` or just `dict`. Every function parameter must be strongly typed -- Use tagged unions + match -- No monster files or god objects -- No file sprawl: deliberate file and folder structure -- Standard over hand-rolled: use the official SDK or a library where one exists; where none does, follow industry standards instead of inventing local conventions -- API-fragmentation-aware: when logic must branch on which API surface produced or consumes data (e.g. chat completions vs Anthropic Messages vs Responses API shapes), proactively look for an existing shared helper (e.g. `litellm_core_utils/prompt_templates/factory.py`) before writing per-surface parsing in the new module; if none exists, add one there instead of duplicating the same format-detection logic in every new guardrail/integration - -Follow conventional commits for commit names and PR titles - -## Think Before Coding - -**Don't assume. Don't hide confusion. Surface tradeoffs** - -Before implementing: -- State your assumptions explicitly. If uncertain, ask -- If multiple interpretations exist, present them. Don't pick silently -- If a simpler approach exists, say so. Push back when warranted -- If something is unclear, stop. Name what's confusing. Ask - -## Simplicity First - -**Minimum code that solves the problem. Nothing speculative** - -- No features beyond what was asked -- No abstractions for single-use code -- No "flexibility" or "configurability" that wasn't requested -- No error handling for impossible scenarios -- If you write 200 lines and it could be 50, rewrite it - -Ask yourself: "Would a senior engineer say this is overcomplicated?" If yes, simplify diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 0443f1bed75..153ca040e27 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -162,7 +162,7 @@ make format > **Black formatting is enforced in CI.** All PRs must pass the Black formatting check. > -> - **AI coding agents** (Claude Code, Copilot, Cursor, etc.): `AGENTS.md` and `CLAUDE.md` instruct agents to run `poetry run black .` before committing. +> - **AI coding agents** (Claude Code, Copilot, Cursor, etc.): `AGENTS.md` instructs agents to run `poetry run black .` before committing. > - **VS Code users**: Install the [Black Formatter extension](https://marketplace.visualstudio.com/items?itemName=ms-python.black-formatter) and enable format-on-save: > ```json > { diff --git a/GEMINI.md b/GEMINI.md index 41921fdff4d..5fc00e0b5ae 100644 --- a/GEMINI.md +++ b/GEMINI.md @@ -1 +1 @@ -Read @CLAUDE.md for coding guidelines +Read @AGENTS.md for coding guidelines diff --git a/litellm-rust/crates/python-bridge/AGENTS.md b/litellm-rust/crates/python-bridge/AGENTS.md index 5dccfb4aca8..a19a709e60c 100644 --- a/litellm-rust/crates/python-bridge/AGENTS.md +++ b/litellm-rust/crates/python-bridge/AGENTS.md @@ -1,4 +1,4 @@ -- Target invariants, not completion claims; these supersede older conflicting bridge guidance +- Target invariants, not completion claims; these supersede the crate guidance below where they conflict - Keep this crate the product-specific PyO3 consumer of `litellm-host-python` - Own registration, input projection, the route host and the caller callables it answers operations with (file readers, token providers), public response/error construction and the per-call composition of machine, route host and callback contract - Legacy callback sharing (the caller's args, kwargs and request object, body/header roots, re-aliasing unchanged body keys) lives in `litellm-callbacks-legacy` behind `PublicCall` and `run_legacy_call`; the bridge hands the public call over and keeps no copy @@ -34,3 +34,45 @@ - References: [ownership](https://pyo3.rs/v0.29.2/types.html), [GC](https://pyo3.rs/v0.29.2/class/protocols.html#garbage-collector-integration), [exception transfer](https://docs.rs/pyo3/0.29.2/pyo3/struct.PyErr.html#method.into_value), [re-entry](https://pyo3.rs/v0.29.2/class/call.html) - [GIL policy](https://pyo3.rs/v0.29.2/free-threading.html), [experimental async limits](https://pyo3.rs/v0.29.2/async-await.html), [task conversion](https://docs.rs/pyo3-async-runtimes/0.29.0/pyo3_async_runtimes/fn.into_future_with_locals.html), [native cancellation/delivery](https://docs.rs/pyo3-async-runtimes/0.29.0/pyo3_async_runtimes/tokio/fn.future_into_py.html) - [performance](https://pyo3.rs/v0.29.2/performance.html), [PyBackedBytes](https://docs.rs/pyo3/0.29.2/pyo3/pybacked/struct.PyBackedBytes.html), [typing](https://pyo3.rs/v0.29.2/python-typing-hints.html) + +Rules for `litellm-rust/crates/python-bridge`. + +## Responsibility + +`python-bridge` is the PyO3 boundary between Python LiteLLM and Rust transforms. +Keep this crate thin. It exposes LiteLLM Rust APIs, assembles domain requests, +maps domain errors to Python exceptions, and delegates generic conversion and +GIL handling to `litellm-host-python`. + +## Bridge Shape + +- Prefer one stable method per top-level LiteLLM route, for example + `messages(...)`, calling the matching `litellm-core` entrypoint. +- Do not add one exported PyO3 function per provider helper unless there is a + measured reason. +- Provider dispatch belongs in the `litellm-core` route module (e.g. + `litellm_core::messages`), not in this PyO3 crate. +- Python owns rollout state and fallback. Rust should return errors; Python + decides whether to raise or fall back. For a rust-only provider/route (no + Python reference), the Python side is a thin dispatch that calls Rust and + raises when the bridge is unavailable, with no fallback. +- Keep the Python interface minimal (well under 100 lines per route): it only + marshals inputs and calls Rust. Do not add per-route feature flags, and do + not put provider dispatch in `litellm/main.py`; it lives in a thin dispatch + class under `litellm/llms///`. + +## Data Handling + +- OCR payloads can contain personal data and large base64 images. Do not log + payloads or provider responses. +- Avoid copying large payloads more than needed. The current JSON round-trip is + acceptable for the first scaffold, but future performance work should evaluate + direct PyO3 conversion before expanding Rust coverage to image-heavy paths. +- Do not expose raw Rust errors that include document contents or upstream + bodies. + +## Tests + +- `cargo test --workspace` must compile this crate. +- Python tests must cover bridge disabled, bridge enabled, and module-missing + fallback behavior for every exposed route. diff --git a/litellm-rust/crates/python-bridge/CLAUDE.md b/litellm-rust/crates/python-bridge/CLAUDE.md deleted file mode 100644 index e55bb192cdd..00000000000 --- a/litellm-rust/crates/python-bridge/CLAUDE.md +++ /dev/null @@ -1,43 +0,0 @@ -# CLAUDE.md - -Rules for `litellm-rust/crates/python-bridge`. - -## Responsibility - -`python-bridge` is the PyO3 boundary between Python LiteLLM and Rust transforms. -Keep this crate thin. It exposes LiteLLM Rust APIs, assembles domain requests, -maps domain errors to Python exceptions, and delegates generic conversion and -GIL handling to `litellm-host-python`. - -## Bridge Shape - -- Prefer one stable method per top-level LiteLLM route, for example - `messages(...)`, calling the matching `litellm-core` entrypoint. -- Do not add one exported PyO3 function per provider helper unless there is a - measured reason. -- Provider dispatch belongs in the `litellm-core` route module (e.g. - `litellm_core::messages`), not in this PyO3 crate. -- Python owns rollout state and fallback. Rust should return errors; Python - decides whether to raise or fall back. For a rust-only provider/route (no - Python reference), the Python side is a thin dispatch that calls Rust and - raises when the bridge is unavailable, with no fallback. -- Keep the Python interface minimal (well under 100 lines per route): it only - marshals inputs and calls Rust. Do not add per-route feature flags, and do - not put provider dispatch in `litellm/main.py`; it lives in a thin dispatch - class under `litellm/llms///`. - -## Data Handling - -- OCR payloads can contain personal data and large base64 images. Do not log - payloads or provider responses. -- Avoid copying large payloads more than needed. The current JSON round-trip is - acceptable for the first scaffold, but future performance work should evaluate - direct PyO3 conversion before expanding Rust coverage to image-heavy paths. -- Do not expose raw Rust errors that include document contents or upstream - bodies. - -## Tests - -- `cargo test --workspace` must compile this crate. -- Python tests must cover bridge disabled, bridge enabled, and module-missing - fallback behavior for every exposed route. diff --git a/litellm/proxy/_experimental/mcp_server/AGENTS.md b/litellm/proxy/_experimental/mcp_server/AGENTS.md index fa83f86a675..d9e0bfa3589 100644 --- a/litellm/proxy/_experimental/mcp_server/AGENTS.md +++ b/litellm/proxy/_experimental/mcp_server/AGENTS.md @@ -1,6 +1,6 @@ # Experimental MCP Server Change Guidelines -Read @../../../../CLAUDE.md and @CLAUDE.md before changing this package. +Read @../../../../AGENTS.md before changing this package. This directory owns the proxy-hosted MCP server implementation. Keep changes inside the module that owns the behavior, and only reach outside this package @@ -14,7 +14,6 @@ Respect the current package boundaries: ```text litellm/proxy/_experimental/mcp_server/ AGENTS.md - CLAUDE.md server.py # ASGI/MCP route handling, sessions, tool calls [PR7: 7-arm only — move BYOK/OAuth pre-fetch into resolver] mcp_server_manager.py # upstream server registry, clients, tool routing [PR7: _create_mcp_client swaps resolve_mcp_auth -> resolve_credentials] auth/ @@ -68,8 +67,10 @@ module materially harder to understand. auth, SSE, streamable HTTP, and stdio as separate flows. Do not collapse them behind a single generic branch unless tests prove every mode still behaves correctly. -- Be especially careful with legacy `delegate_auth_to_upstream: true`. The local - `CLAUDE.md` explains its admitted replacement and public discovery contract. +- Be especially careful with legacy `delegate_auth_to_upstream: true`. `auth_type: oauth2` + with `delegate_auth_to_upstream: true` is deprecated: LiteLLM admission is required + for matching MCP routes. Use `auth_type: oauth_delegate` for client-forwarded OAuth. + OAuth discovery endpoints stay public so clients can start the RFC 9728 flow. - Keep database-backed fields in sync across migrations, typed models under `litellm/types/mcp.py` or `litellm/types/mcp_server/`, config loading, this package, and dashboard state when the field is user-visible. diff --git a/litellm/proxy/_experimental/mcp_server/CLAUDE.md b/litellm/proxy/_experimental/mcp_server/CLAUDE.md deleted file mode 100644 index 7f8d06b4570..00000000000 --- a/litellm/proxy/_experimental/mcp_server/CLAUDE.md +++ /dev/null @@ -1 +0,0 @@ -MCP note: **`auth_type: oauth2` with `delegate_auth_to_upstream: true` is deprecated** - LiteLLM admission is required for matching MCP routes. Use `auth_type: oauth_delegate` for client-forwarded OAuth. OAuth discovery endpoints stay public so clients can start the RFC 9728 flow diff --git a/tests/code_coverage_tests/check_e2e_no_raw_requests.py b/tests/code_coverage_tests/check_e2e_no_raw_requests.py index fe6a77fc26c..3f40cc3ee1e 100644 --- a/tests/code_coverage_tests/check_e2e_no_raw_requests.py +++ b/tests/code_coverage_tests/check_e2e_no_raw_requests.py @@ -5,7 +5,7 @@ anywhere; a small allowlist grandfathers the files that legitimately make raw ca (the transport itself, the root conftest liveness probe, the claude_code version resolver's constant registry URL fetch, and the mcp OAuth client, whose httpx client is the object the official mcp SDK's streamable_http_client requires and so -cannot go through the sync requests transport). Referenced by tests/e2e/CLAUDE.md.""" +cannot go through the sync requests transport). Referenced by tests/e2e/AGENTS.md.""" from __future__ import annotations diff --git a/tests/code_coverage_tests/test_e2e_changed_gate.py b/tests/code_coverage_tests/test_e2e_changed_gate.py index 101816c7f11..5ae0863baf0 100644 --- a/tests/code_coverage_tests/test_e2e_changed_gate.py +++ b/tests/code_coverage_tests/test_e2e_changed_gate.py @@ -142,7 +142,7 @@ def select_tests(changed: tuple[str, ...]) -> tuple[str, ...]: ("tests/e2e/guardrails/test_bedrock_guardrail_e2e.py",), ("tests/e2e/guardrails/test_bedrock_guardrail_e2e.py",), ), - (("tests/e2e/logging/helpers.py", "docs/my-website/docs/index.md", "tests/e2e/CLAUDE.md"), ()), + (("tests/e2e/logging/helpers.py", "docs/my-website/docs/index.md", "tests/e2e/AGENTS.md"), ()), ( ("tests/e2e/logging/test_datadog_e2e.py", "tests/e2e/logging/test_datadog_e2e.py"), ("tests/e2e/logging/test_datadog_e2e.py",), diff --git a/tests/e2e/CLAUDE.md b/tests/e2e/AGENTS.md similarity index 99% rename from tests/e2e/CLAUDE.md rename to tests/e2e/AGENTS.md index 0541ce25d4b..8a56e8673c4 100644 --- a/tests/e2e/CLAUDE.md +++ b/tests/e2e/AGENTS.md @@ -1,6 +1,6 @@ # e2e harness conventions -Code-style rules for writing tests under `tests/e2e/`. The harness already encodes the plumbing; your job is the feature-specific behavior, not reinventing it. For what a complete test must do (the lifecycle contract, asserting both recorded state and enforced behavior) and how to run a suite, see `CONTRIBUTING.md` in this directory. Repo-wide conventions live in the root `CLAUDE.md` +Code-style rules for writing tests under `tests/e2e/`. The harness already encodes the plumbing; your job is the feature-specific behavior, not reinventing it. For what a complete test must do (the lifecycle contract, asserting both recorded state and enforced behavior) and how to run a suite, see `CONTRIBUTING.md` in this directory. Repo-wide conventions live in the root `AGENTS.md` ## Suite folders diff --git a/tests/e2e/CONTRIBUTING.md b/tests/e2e/CONTRIBUTING.md index 20073e5d68f..2afcc563824 100644 --- a/tests/e2e/CONTRIBUTING.md +++ b/tests/e2e/CONTRIBUTING.md @@ -2,7 +2,7 @@ This directory holds the live end-to-end suites that prove product correctness against a real running proxy and real provider APIs. The goal of this guide is simple: when you ship a feature, you add e2e coverage that walks that feature the way production does, across every route and edge case it touches, so a later change that breaks it fails here first -Read this before adding a test and i recommend reading through CLAUDE.md +Read this before adding a test and i recommend reading through AGENTS.md When contributing to this directory, please first discuss the change you wish to make via issue or pull request. We require screenshots and proof of your tests working on a live proxy. @@ -134,7 +134,7 @@ One sharp edge: a replayed response reuses the recorded provider response id, an Another sharp edge, same root: record and replay derive every per-test token deterministically (the model name included, so a replay regenerates the exact requests the record run sent), which means an edge-wired deployment left in the database by an interrupted earlier run carries the same model name as the fresh one the current run registers. The proxy then holds two deployments under one model group and load-balances across both, and because the leftover's `api_base` points at the earlier run's edge process, which is gone, the calls that land on it fail with a connection error that reads like a transport bug rather than the stale row it is. Give each record or replay run a fresh database, or let a run finish so its own teardown deletes what it registered, and never reuse one long-lived proxy across back-to-back record/replay sessions. CI hands every job its own empty database and its own proxy, so it never sees this -Replay answers any provider call that drifted from the recording with an HTTP 599 whose body names the computed and closest recorded keys, so the test fails loudly instead of silently going live, and a bundle older than seven days fails at collection time naming its age; either way the fix is to re-record. Only tests that register edge-wired deployments participate: everything else hits its provider live in every mode, so record exactly the suite you replay. If the proxy runs in a container, set `E2E_PROVIDER_EDGE_ADVERTISE_HOST` (e.g. `host.docker.internal`) so the api_base the proxy stores can reach the edge on the pytest host, and `E2E_PROVIDER_EDGE_BIND_HOST=0.0.0.0` so the edge accepts it. The suites wired to the edge today are `quota_management/spend_tracking/test_provider_edge_spend_e2e.py`, `llm_translation/test_chat_completions_contract_e2e.py`, the OpenAI registrations in `llm_translation/test_embeddings_endpoint_e2e.py`, the Anthropic tests in `llm_translation/test_messages_e2e.py`, streamed and not, and the OpenAI batch deployment behind `batches/`. A streamed response replays as the chunk sequence the provider sent rather than one buffered body. See `CLAUDE.md` in this directory for the bundle format, the edge design, and the current limits (Bedrock). The scheduled CI record/replay lane is described above +Replay answers any provider call that drifted from the recording with an HTTP 599 whose body names the computed and closest recorded keys, so the test fails loudly instead of silently going live, and a bundle older than seven days fails at collection time naming its age; either way the fix is to re-record. Only tests that register edge-wired deployments participate: everything else hits its provider live in every mode, so record exactly the suite you replay. If the proxy runs in a container, set `E2E_PROVIDER_EDGE_ADVERTISE_HOST` (e.g. `host.docker.internal`) so the api_base the proxy stores can reach the edge on the pytest host, and `E2E_PROVIDER_EDGE_BIND_HOST=0.0.0.0` so the edge accepts it. The suites wired to the edge today are `quota_management/spend_tracking/test_provider_edge_spend_e2e.py`, `llm_translation/test_chat_completions_contract_e2e.py`, the OpenAI registrations in `llm_translation/test_embeddings_endpoint_e2e.py`, the Anthropic tests in `llm_translation/test_messages_e2e.py`, streamed and not, and the OpenAI batch deployment behind `batches/`. A streamed response replays as the chunk sequence the provider sent rather than one buffered body. See `AGENTS.md` in this directory for the bundle format, the edge design, and the current limits (Bedrock). The scheduled CI record/replay lane is described above Tests marked `@pytest.mark.e2e` hard-fail when no proxy answers `/health/liveliness`, so a run that goes red with `No live proxy` at setup means the proxy isn't up; they never skip for a missing proxy, so an absent proxy can't be mistaken for a pass diff --git a/tests/e2e/batches/COVERAGE.md b/tests/e2e/batches/COVERAGE.md index b36d8937ad0..d18bed6c088 100644 --- a/tests/e2e/batches/COVERAGE.md +++ b/tests/e2e/batches/COVERAGE.md @@ -12,7 +12,7 @@ cost write-back via a cross-run marker baton (design below). Only supported cells are tested. The capability table in `capabilities.py` holds one row per supported (provider, scenario) pair, so there are no skipped cells in the parametrized run. The batches suite never skips: missing provider creds or upstream -failures are hard test failures (see `tests/e2e/CLAUDE.md`). +failures are hard test failures (see `tests/e2e/AGENTS.md`). | Provider | create | retrieve | cancel | list | content download | file backing | |-----------|--------|----------|--------|------|------------------|--------------| diff --git a/tests/e2e/claude_code/cron_vm/run_daily.sh b/tests/e2e/claude_code/cron_vm/run_daily.sh index 00d3e66e5bc..e878007d8a3 100755 --- a/tests/e2e/claude_code/cron_vm/run_daily.sh +++ b/tests/e2e/claude_code/cron_vm/run_daily.sh @@ -288,7 +288,7 @@ else # Download the tarball and Astral's official .sha256 sidecar to disk # and verify the digest before extracting/executing anything. This # closes the supply-chain trust gap of piping a remote binary - # straight into `tar -xzO ... > file ; chmod +x` (see CLAUDE.md + # straight into `tar -xzO ... > file ; chmod +x` (see AGENTS.md # "CI Supply-Chain Safety"). curl -fsSL --output "${UV_TMPDIR}/${UV_TARBALL_NAME}" "${UV_DOWNLOAD_URL}" curl -fsSL --output "${UV_TMPDIR}/${UV_TARBALL_NAME}.sha256" "${UV_DOWNLOAD_URL}.sha256" diff --git a/tests/e2e/coverage_registry/README.md b/tests/e2e/coverage_registry/README.md index da6aee84cc4..4f9845bab87 100644 --- a/tests/e2e/coverage_registry/README.md +++ b/tests/e2e/coverage_registry/README.md @@ -3,7 +3,7 @@ This directory is the **denominator** for e2e test coverage: the set of behaviors we want covered, one row per behavior, checked into the repo so coverage is a number we can track instead of a guess. It implements the plan in the "E2E Coverage Tracking" -note; the naming grammar lives in `tests/e2e/CLAUDE.md`. +note; the naming grammar lives in `tests/e2e/AGENTS.md`. ## The model diff --git a/tests/e2e/coverage_registry/__init__.py b/tests/e2e/coverage_registry/__init__.py index 959b3327194..0eb153011a6 100644 --- a/tests/e2e/coverage_registry/__init__.py +++ b/tests/e2e/coverage_registry/__init__.py @@ -3,6 +3,6 @@ `schema.py` defines one validated row per customer-noticeable behavior (a "cell"). The `*.yaml` files hold the rows, one file per id-prefix. `registry.py` loads and validates them; `collector.py` diffs the registry against the `@pytest.mark.covers` -markers on the live tests and reports coverage per module. See tests/e2e/CLAUDE.md +markers on the live tests and reports coverage per module. See tests/e2e/AGENTS.md for the naming grammar. """ diff --git a/tests/e2e/coverage_registry/mcp.yaml b/tests/e2e/coverage_registry/mcp.yaml index a7d4135d550..85ace835144 100644 --- a/tests/e2e/coverage_registry/mcp.yaml +++ b/tests/e2e/coverage_registry/mcp.yaml @@ -1,4 +1,4 @@ -# MCP module. Grounded in litellm/proxy/_experimental/mcp_server/. See tests/e2e/CLAUDE.md for the grammar. +# MCP module. Grounded in litellm/proxy/_experimental/mcp_server/. See tests/e2e/AGENTS.md for the grammar. - id: mcp.list_tools.api_key.succeeds module: mcp tier: P0 diff --git a/tests/e2e/llm_translation/realtime/REALTIME_COVERAGE_MATRIX.md b/tests/e2e/llm_translation/realtime/REALTIME_COVERAGE_MATRIX.md index a6e32b88479..c85471da90d 100644 --- a/tests/e2e/llm_translation/realtime/REALTIME_COVERAGE_MATRIX.md +++ b/tests/e2e/llm_translation/realtime/REALTIME_COVERAGE_MATRIX.md @@ -49,7 +49,7 @@ kept commented out in `PROVIDERS` until they pass end-to-end here; re-enable the uncommenting their entry. Every provider is provisioned and asserted; the suite never skips a provider. Per -`tests/e2e/CLAUDE.md` there is no sanctioned skip: the whole-suite proxy-liveness +`tests/e2e/AGENTS.md` there is no sanctioned skip: the whole-suite proxy-liveness probe hard-fails when no proxy answers, and a provider whose credentials or upstream realtime model are missing on the gateway is likewise a hard failure, not a skip. Give the gateway each provider's credentials to turn its tests green. diff --git a/tests/e2e/llm_translation/realtime/conftest.py b/tests/e2e/llm_translation/realtime/conftest.py index 752737e830e..804a9b9c649 100644 --- a/tests/e2e/llm_translation/realtime/conftest.py +++ b/tests/e2e/llm_translation/realtime/conftest.py @@ -29,7 +29,7 @@ def realtime_models(client: RealtimeClient) -> Iterator[dict[str, str]]: provider-id -> model-name map the tests connect with; delete them on teardown. Every provider is provisioned (never skipped): a provider whose credentials or upstream model are missing on the gateway hard-fails its test, per the suite's - fail-on-behavior contract in tests/e2e/CLAUDE.md.""" + fail-on-behavior contract in tests/e2e/AGENTS.md.""" records = tuple((provider.id, *client.provision(provider)) for provider in PROVIDERS) try: yield {provider_id: model_name for provider_id, model_name, _ in records} diff --git a/tests/e2e/llm_translation/realtime/realtime_client.py b/tests/e2e/llm_translation/realtime/realtime_client.py index 3ffca7e8b88..7c4a9cc4af9 100644 --- a/tests/e2e/llm_translation/realtime/realtime_client.py +++ b/tests/e2e/llm_translation/realtime/realtime_client.py @@ -38,7 +38,7 @@ class RealtimeProvider: the suite registers through /model/new (the gateway resolves the os.environ/* credential refs), so the suite is self-contained and never depends on a static gateway model_list. Every provider here is provisioned and asserted: per - tests/e2e/CLAUDE.md the suite never skips a provider, so a provider whose + tests/e2e/AGENTS.md the suite never skips a provider, so a provider whose credentials or upstream realtime model are missing on the gateway is a hard failure, not a skip.""" @@ -98,7 +98,7 @@ PROVIDERS = ( def realtime_model(provider: RealtimeProvider, provisioned: Mapping[str, str]) -> str: """Return the provisioned deployment name for this provider. Every provider in PROVIDERS is provisioned at session start, so a missing entry is a harness bug, - never an environment skip - the suite hard-fails instead (see tests/e2e/CLAUDE.md).""" + never an environment skip - the suite hard-fails instead (see tests/e2e/AGENTS.md).""" model = provisioned.get(provider.id) assert model is not None, ( f"{provider.id} was not provisioned; the realtime_models fixture is broken" diff --git a/ui/litellm-dashboard/CLAUDE.md b/ui/litellm-dashboard/AGENTS.md similarity index 100% rename from ui/litellm-dashboard/CLAUDE.md rename to ui/litellm-dashboard/AGENTS.md From fb41bc3ed658b9023937c5f477c6311256c2dd68 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Fri, 18 Sep 2026 19:39:07 -0700 Subject: [PATCH 134/144] revert(ocr): stop forwarding client= on the Python path Python becomes a thin SDK interface over Rust, so a live Python HTTP client has no effect on either route. This puts the Python OCR path back to what main does --- litellm/ocr/main.py | 8 -------- tests/test_litellm/ocr/test_main.py | 29 ----------------------------- 2 files changed, 37 deletions(-) diff --git a/litellm/ocr/main.py b/litellm/ocr/main.py index 851d9162964..06830ed4b53 100644 --- a/litellm/ocr/main.py +++ b/litellm/ocr/main.py @@ -25,7 +25,6 @@ from litellm.llms.base_llm.ocr.transformation import ( OCRResponse, parse_ocr_request_format, ) -from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import CustomPricingLiteLLMParams @@ -53,11 +52,6 @@ class _PreparedOCRRequest: litellm_logging_obj: LiteLLMLoggingObj -def _supplied_client(kwargs: Mapping[str, object]) -> HTTPHandler | AsyncHTTPHandler | None: - candidate: Final = kwargs.get("client") - return candidate if isinstance(candidate, (HTTPHandler, AsyncHTTPHandler)) else None - - def _prepare_ocr_request( model: str, document: Mapping[str, object], @@ -244,7 +238,6 @@ async def aocr( api_key=prepared.api_key, api_base=prepared.api_base, custom_llm_provider=prepared.custom_llm_provider, - client=_supplied_client(kwargs), aocr=True, headers=prepared.extra_headers, provider_config=prepared.provider_config, @@ -411,7 +404,6 @@ def ocr( api_key=prepared.api_key, api_base=prepared.api_base, custom_llm_provider=prepared.custom_llm_provider, - client=_supplied_client(kwargs), aocr=_is_async, headers=prepared.extra_headers, provider_config=prepared.provider_config, diff --git a/tests/test_litellm/ocr/test_main.py b/tests/test_litellm/ocr/test_main.py index 32e5637ee09..5531a2639c0 100644 --- a/tests/test_litellm/ocr/test_main.py +++ b/tests/test_litellm/ocr/test_main.py @@ -113,35 +113,6 @@ async def test_python_request_response_and_callbacks( assert logger.log_pre_api_call.call_count == 1 -@pytest.mark.asyncio -@pytest.mark.parametrize("asynchronous", [False, True]) -async def test_python_uses_the_supplied_client(provider: Mock, asynchronous: bool) -> None: - supplied: Final = Mock(return_value=provider.return_value) - transport: Final = httpx.MockTransport(supplied) - arguments: Final = { - "model": "mistral/mistral-ocr-latest", - "document": dict(PRICING_DOCUMENT), - "api_key": "test-key", - "api_base": "https://ocr.test/v1", - } - - async def call() -> OCRResponse: - if not asynchronous: - with httpx.Client(transport=transport) as sync_client: - return litellm.ocr(**arguments, client=HTTPHandler(client=sync_client)) - async with httpx.AsyncClient(transport=transport) as async_client: - handler: Final = AsyncHTTPHandler() - await handler.client.aclose() - handler.client = async_client - return await litellm.aocr(**arguments, client=handler) - - response: Final = await call() - assert response.pages[0].markdown == "parsed document" - assert supplied.call_count == 1 - assert str(supplied.call_args.args[0].url) == "https://ocr.test/v1/ocr" - assert provider.call_count == 0 - - @pytest.mark.asyncio @pytest.mark.parametrize("asynchronous", [False, True]) async def test_python_provider_errors_keep_public_exception(provider: Mock, asynchronous: bool) -> None: From 51010ea486666b736c9d289e9df6b17eff5b7d7d Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Fri, 18 Sep 2026 19:39:07 -0700 Subject: [PATCH 135/144] feat(rust): serve every gateway HTTP setting natively instead of declining to Python litellm-http now builds the rustls config itself, so one route-neutral place covers roots, the client certificate, ALPN, ssl_ecdh_curve and ssl_security_level. A curve picks the single key exchange group. A cipher string restricts the TLS 1.2 suites it names, and entries rustls cannot express, such as @SECLEVEL=1, are logged once and skipped. user_url_validation and user_url_allowed_hosts are applied by the media fetcher. Document downloads honor the environment proxy whenever provider calls do, keeping the per-hop address check, and stay on the pinned resolver when no proxy applies. AIOHTTP_SO_KEEPALIVE, AIOHTTP_TCP_KEEPIDLE, AIOHTTP_TCP_KEEPINTVL, AIOHTTP_TCP_KEEPCNT and AIOHTTP_KEEPALIVE_TIMEOUT map onto the client. A client= argument and a live SSLContext are ignored --- litellm-rust/Cargo.lock | 4 + litellm-rust/Cargo.toml | 3 + litellm-rust/crates/core/tests/ocr.rs | 8 +- litellm-rust/crates/http/Cargo.toml | 4 + litellm-rust/crates/http/src/config.rs | 235 +++++----- litellm-rust/crates/http/src/error.rs | 5 - litellm-rust/crates/http/src/lib.rs | 8 +- litellm-rust/crates/http/src/pool.rs | 29 +- litellm-rust/crates/http/src/proxy.rs | 15 + litellm-rust/crates/http/src/settings.rs | 52 +++ litellm-rust/crates/http/src/tls.rs | 402 ++++++++++++++++++ .../llms/src/custom_httpx/llm_http_handler.rs | 5 +- .../crates/llms/src/custom_httpx/media.rs | 207 ++++++++- litellm-rust/crates/python-bridge/src/http.rs | 174 +++----- .../python-bridge/src/python_settings.rs | 5 + .../python-bridge/src/routes/ocr/mod.rs | 9 +- litellm/rust_bridge/settings.py | 6 + .../test_litellm/rust_bridge/test_settings.py | 8 + 18 files changed, 944 insertions(+), 235 deletions(-) create mode 100644 litellm-rust/crates/http/src/proxy.rs create mode 100644 litellm-rust/crates/http/src/tls.rs diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 7f2d2e6b28b..d4b32659ba1 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -2135,10 +2135,14 @@ dependencies = [ name = "litellm-http" version = "0.1.0" dependencies = [ + "http 1.4.2", + "hyper-util", "reqwest 0.12.28", "rstest", + "rustls 0.23.42", "thiserror 2.0.19", "tokio", + "webpki-roots", ] [[package]] diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index 02f4cc6b3ab..8634dce92d0 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -27,6 +27,8 @@ litellm-token-counter = { path = "crates/token-counter" } litellm-host-python = { path = "crates/host-python" } bytes = "1" +http = "1" +hyper-util = { version = "0.1.20", default-features = false, features = ["client-proxy"] } proptest = "1.7.0" pyo3 = "0.29.2" pyo3-async-runtimes = { version = "0.29.0", features = ["tokio-runtime"] } @@ -50,6 +52,7 @@ base64 = "0.22" moka = { version = "0.12.16", features = ["future"] } strum = { version = "0.28.0", features = ["derive"] } url = "2.5.8" +webpki-roots = "1" time = { version = "0.3.53", features = ["parsing"] } criterion = "0.8.2" fancy-regex = "0.19.2" diff --git a/litellm-rust/crates/core/tests/ocr.rs b/litellm-rust/crates/core/tests/ocr.rs index 2ae162d964f..a6b26bd8a27 100644 --- a/litellm-rust/crates/core/tests/ocr.rs +++ b/litellm-rust/crates/core/tests/ocr.rs @@ -12,7 +12,10 @@ use litellm_llms::{ error::Error as OcrError, transformation::{LiteLLMOcrResponse, OCR_RESPONSE_MAX_BYTES, OcrTransportConfig}, }, - custom_httpx::{llm_http_handler::OcrClient, media::PublicDnsResolver}, + custom_httpx::{ + llm_http_handler::OcrClient, + media::{PublicDnsResolver, UrlPolicy}, + }, }; use rstest::rstest; use serde_json::{Value, json}; @@ -181,7 +184,8 @@ async fn ocr_client_uses_the_injected_http_pool_configuration() { }; let client = OcrClient::new( &HttpClientPool::new(Arc::new(PublicDnsResolver)), - &HttpClientConfig::resolve(&settings).unwrap(), + &HttpClientConfig::resolve(&settings).config, + UrlPolicy::default(), VertexAuth::default(), ) .unwrap(); diff --git a/litellm-rust/crates/http/Cargo.toml b/litellm-rust/crates/http/Cargo.toml index 48ea4e66cef..0ac09a9d155 100644 --- a/litellm-rust/crates/http/Cargo.toml +++ b/litellm-rust/crates/http/Cargo.toml @@ -6,8 +6,12 @@ license.workspace = true repository.workspace = true [dependencies] +http.workspace = true +hyper-util.workspace = true reqwest.workspace = true +rustls.workspace = true thiserror.workspace = true +webpki-roots.workspace = true [dev-dependencies] rstest.workspace = true diff --git a/litellm-rust/crates/http/src/config.rs b/litellm-rust/crates/http/src/config.rs index 24a52315f7c..ebe557788ca 100644 --- a/litellm-rust/crates/http/src/config.rs +++ b/litellm-rust/crates/http/src/config.rs @@ -1,12 +1,13 @@ use std::{ net::{IpAddr, Ipv4Addr}, - path::{Path, PathBuf}, + path::PathBuf, time::Duration, }; use crate::{ error::Error, - settings::{HttpSettings, SslVerify}, + settings::{HttpSettings, SslVerify, TcpKeepalive}, + tls::{self, KeyExchangeGroup, Tls12CipherSuite, Unsupported}, }; #[derive(Clone, Debug, PartialEq, Eq, Hash)] @@ -20,27 +21,38 @@ pub enum Verify { pub struct HttpClientConfig { pub verify: Verify, pub client_certificate: Option, + pub key_exchange_group: Option, + pub tls12_cipher_suites: Option>, pub force_ipv4: bool, pub http2: bool, pub user_agent: Option, pub trust_proxy_env: bool, pub connect_timeout: Duration, + pub tcp_keepalive: Option, + pub pool_idle_timeout: Duration, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct Resolution { + pub config: HttpClientConfig, + pub unsupported: Vec, } impl HttpClientConfig { - pub fn resolve(settings: &HttpSettings) -> Result { - if let Some(level) = &settings.ssl_security_level { - return Err(Error::Unsupported { - setting: "ssl_security_level", - reason: format!("OpenSSL cipher string {level:?} has no rustls equivalent"), - }); - } - if let Some(curve) = &settings.ssl_ecdh_curve { - return Err(Error::Unsupported { - setting: "ssl_ecdh_curve", - reason: format!("key exchange group {curve:?} is fixed by the rustls provider"), - }); - } + pub fn resolve(settings: &HttpSettings) -> Resolution { + let (key_exchange_group, unsupported_curve) = match settings + .ssl_ecdh_curve + .as_deref() + .map(KeyExchangeGroup::from_openssl_name) + { + None => (None, None), + Some(Ok(group)) => (Some(group), None), + Some(Err(unsupported)) => (None, Some(unsupported)), + }; + let ciphers = settings + .ssl_security_level + .as_deref() + .map(tls::parse_cipher_string); let verify = match &settings.ssl_verify { Some(SslVerify::Disabled) => Verify::Disabled, Some(SslVerify::CaBundle(path)) => Verify::CaBundle(path.clone()), @@ -49,62 +61,50 @@ impl HttpClientConfig { .clone() .map_or(Verify::BuiltInRoots, Verify::CaBundle), }; - Ok(Self { - verify, - client_certificate: settings.ssl_certificate.clone(), - force_ipv4: settings.force_ipv4, - http2: settings.http2, - user_agent: settings.user_agent.clone(), - trust_proxy_env: !settings.ignore_proxy_env - || settings.trust_proxy_env - || settings.http2 - || settings.httpx_transport, - connect_timeout: settings.connect_timeout, - }) + let (tls12_cipher_suites, unsupported_ciphers) = ciphers + .map_or((None, Vec::new()), |ciphers| { + (ciphers.tls12_cipher_suites, ciphers.unsupported) + }); + Resolution { + config: Self { + verify, + client_certificate: settings.ssl_certificate.clone(), + key_exchange_group, + tls12_cipher_suites, + force_ipv4: settings.force_ipv4, + http2: settings.http2, + user_agent: settings.user_agent.clone(), + trust_proxy_env: !settings.ignore_proxy_env + || settings.trust_proxy_env + || settings.http2 + || settings.httpx_transport, + connect_timeout: settings.connect_timeout, + tcp_keepalive: settings.tcp_keepalive, + pool_idle_timeout: settings.pool_idle_timeout, + }, + unsupported: unsupported_curve + .into_iter() + .chain(unsupported_ciphers) + .collect(), + } } pub fn client_builder(&self) -> Result { - let base = reqwest::Client::builder().connect_timeout(self.connect_timeout); - let with_roots = match &self.verify { - Verify::Disabled => base.danger_accept_invalid_certs(true), - Verify::BuiltInRoots => base, - Verify::CaBundle(path) => { - let pem = read(path)?; - let certificates = - reqwest::Certificate::from_pem_bundle(&pem).map_err(|error| { - Error::InvalidPem { - path: path.clone(), - message: error.without_url().to_string(), - } - })?; - if certificates.is_empty() { - return Err(Error::InvalidPem { - path: path.clone(), - message: "no certificates found".into(), - }); - } - certificates.into_iter().fold( - base.tls_built_in_root_certs(false), - |builder, certificate| builder.add_root_certificate(certificate), - ) - } - }; - let with_identity = match &self.client_certificate { - None => with_roots, - Some(path) => { - let identity = reqwest::Identity::from_pem(&read(path)?).map_err(|error| { - Error::InvalidPem { - path: path.clone(), - message: error.without_url().to_string(), - } - })?; - with_roots.identity(identity) - } + let base = reqwest::Client::builder() + .use_preconfigured_tls(tls::client_config(self)?) + .connect_timeout(self.connect_timeout) + .pool_idle_timeout(self.pool_idle_timeout); + let with_keepalive = match self.tcp_keepalive { + None => base, + Some(keepalive) => base + .tcp_keepalive(keepalive.idle) + .tcp_keepalive_interval(keepalive.interval) + .tcp_keepalive_retries(keepalive.retries), }; let with_address = if self.force_ipv4 { - with_identity.local_address(IpAddr::V4(Ipv4Addr::UNSPECIFIED)) + with_keepalive.local_address(IpAddr::V4(Ipv4Addr::UNSPECIFIED)) } else { - with_identity + with_keepalive }; let with_protocol = if self.http2 { with_address @@ -123,13 +123,6 @@ impl HttpClientConfig { } } -fn read(path: &Path) -> Result, Error> { - std::fs::read(path).map_err(|error| Error::Read { - path: path.to_path_buf(), - message: error.to_string(), - }) -} - #[cfg(test)] mod tests { use rstest::rstest; @@ -168,7 +161,7 @@ mod tests { #[case] settings: HttpSettings, #[case] expected: Verify, ) { - let config = HttpClientConfig::resolve(&settings).unwrap(); + let config = HttpClientConfig::resolve(&settings).config; assert_eq!(config.verify, expected); } @@ -179,42 +172,88 @@ mod tests { ..HttpSettings::default() } .with_environment(&|name: &str| (name == "SSL_VERIFY").then(|| "true".to_string())); - let config = HttpClientConfig::resolve(&settings).unwrap(); + let config = HttpClientConfig::resolve(&settings).config; assert_eq!(config.verify, Verify::BuiltInRoots); } + #[rstest] + #[case::x25519("X25519", Some(KeyExchangeGroup::X25519))] + #[case::openssl_p256("prime256v1", Some(KeyExchangeGroup::Secp256r1))] + #[case::p384("secp384r1", Some(KeyExchangeGroup::Secp384r1))] + fn ecdh_curve_selects_the_single_key_exchange_group( + #[case] curve: &str, + #[case] expected: Option, + ) { + let settings = HttpSettings { + ssl_ecdh_curve: Some(curve.into()), + ..HttpSettings::default() + }; + let resolution = HttpClientConfig::resolve(&settings); + assert_eq!(resolution.config.key_exchange_group, expected); + assert_eq!(resolution.unsupported, []); + } + #[test] - fn cipher_strings_are_rejected_rather_than_ignored() { + fn unsupported_ecdh_curve_keeps_the_defaults_and_is_reported() { + let settings = HttpSettings { + ssl_ecdh_curve: Some("secp521r1".into()), + ..HttpSettings::default() + }; + let resolution = HttpClientConfig::resolve(&settings); + assert_eq!(resolution.config.key_exchange_group, None); + assert_eq!( + resolution.unsupported, + [Unsupported::EcdhCurve("secp521r1".into())] + ); + } + + #[test] + fn legacy_security_level_keeps_every_suite_and_is_reported_unsupported() { let settings = HttpSettings { ssl_security_level: Some("DEFAULT@SECLEVEL=1".into()), ..HttpSettings::default() }; - assert!(matches!( - HttpClientConfig::resolve(&settings), - Err(Error::Unsupported { - setting: "ssl_security_level", - .. - }) - )); + let resolution = HttpClientConfig::resolve(&settings); + assert_eq!(resolution.config.tls12_cipher_suites, None); + assert_eq!( + resolution.unsupported, + [Unsupported::SecurityLevel("@SECLEVEL=1".into())] + ); } #[test] - fn ecdh_curves_are_rejected_rather_than_ignored() { + fn named_suites_restrict_tls12_and_unsupported_entries_are_reported() { let settings = HttpSettings { - ssl_ecdh_curve: Some("X25519".into()), + ssl_security_level: Some( + "ECDHE-RSA-AES256-GCM-SHA384:ECDHE-ECDSA-AES128-GCM-SHA256:!aNULL:AES256-SHA@SECLEVEL=2" + .into(), + ), ..HttpSettings::default() }; - assert!(matches!( - HttpClientConfig::resolve(&settings), - Err(Error::Unsupported { - setting: "ssl_ecdh_curve", - .. - }) - )); + let resolution = HttpClientConfig::resolve(&settings); + assert_eq!( + resolution.config.tls12_cipher_suites, + Some(vec![ + Tls12CipherSuite::EcdheEcdsaAes128Gcm, + Tls12CipherSuite::EcdheRsaAes256Gcm + ]) + ); + assert_eq!( + resolution.unsupported, + [ + Unsupported::CipherToken("!aNULL".into()), + Unsupported::CipherToken("AES256-SHA".into()) + ] + ); } #[test] fn connection_settings_carry_over_unchanged() { + let keepalive = TcpKeepalive { + idle: Duration::from_secs(60), + interval: Duration::from_secs(30), + retries: 5, + }; let settings = HttpSettings { ssl_certificate: Some("/client.pem".into()), force_ipv4: true, @@ -222,19 +261,25 @@ mod tests { user_agent: Some("litellm/1.0".into()), trust_proxy_env: true, connect_timeout: Duration::from_secs(7), + tcp_keepalive: Some(keepalive), + pool_idle_timeout: Duration::from_secs(45), ..HttpSettings::default() }; - let config = HttpClientConfig::resolve(&settings).unwrap(); + let config = HttpClientConfig::resolve(&settings).config; assert_eq!( config, HttpClientConfig { verify: Verify::BuiltInRoots, client_certificate: Some("/client.pem".into()), + key_exchange_group: None, + tls12_cipher_suites: None, force_ipv4: true, http2: true, user_agent: Some("litellm/1.0".into()), trust_proxy_env: true, connect_timeout: Duration::from_secs(7), + tcp_keepalive: Some(keepalive), + pool_idle_timeout: Duration::from_secs(45), } ); } @@ -258,7 +303,7 @@ mod tests { #[case] settings: HttpSettings, #[case] expected: bool, ) { - let config = HttpClientConfig::resolve(&settings).unwrap(); + let config = HttpClientConfig::resolve(&settings).config; assert_eq!(config.trust_proxy_env, expected); } @@ -267,7 +312,7 @@ mod tests { let path = std::env::temp_dir().join("litellm-http-missing-bundle.pem"); let config = HttpClientConfig { verify: Verify::CaBundle(path.clone()), - ..HttpClientConfig::resolve(&HttpSettings::default()).unwrap() + ..HttpClientConfig::resolve(&HttpSettings::default()).config }; assert!(matches!( config.client_builder(), @@ -282,7 +327,7 @@ mod tests { std::fs::write(&path, b"not a certificate").unwrap(); let config = HttpClientConfig { verify: Verify::CaBundle(path.clone()), - ..HttpClientConfig::resolve(&HttpSettings::default()).unwrap() + ..HttpClientConfig::resolve(&HttpSettings::default()).config }; let result = config.client_builder().map(drop); std::fs::remove_file(&path).unwrap(); diff --git a/litellm-rust/crates/http/src/error.rs b/litellm-rust/crates/http/src/error.rs index 27899f06cf1..697d0cf59c8 100644 --- a/litellm-rust/crates/http/src/error.rs +++ b/litellm-rust/crates/http/src/error.rs @@ -2,11 +2,6 @@ use std::path::PathBuf; #[derive(Clone, Debug, thiserror::Error, PartialEq, Eq)] pub enum Error { - #[error("{setting} cannot be expressed with rustls: {reason}")] - Unsupported { - setting: &'static str, - reason: String, - }, #[error("could not read {}: {message}", path.display())] Read { path: PathBuf, message: String }, #[error("{} is not a PEM file: {message}", path.display())] diff --git a/litellm-rust/crates/http/src/lib.rs b/litellm-rust/crates/http/src/lib.rs index c02a82539ff..45f370d3a9c 100644 --- a/litellm-rust/crates/http/src/lib.rs +++ b/litellm-rust/crates/http/src/lib.rs @@ -1,9 +1,13 @@ mod config; mod error; mod pool; +mod proxy; mod settings; +mod tls; -pub use config::{HttpClientConfig, Verify}; +pub use config::{HttpClientConfig, Resolution, Verify}; pub use error::Error; pub use pool::{ClientVariant, HttpClientPool}; -pub use settings::{HttpSettings, SslVerify}; +pub use proxy::EnvironmentProxies; +pub use settings::{HttpSettings, SslVerify, TcpKeepalive}; +pub use tls::{KeyExchangeGroup, Tls12CipherSuite, Unsupported, client_config}; diff --git a/litellm-rust/crates/http/src/pool.rs b/litellm-rust/crates/http/src/pool.rs index b51e6711419..e6e0de9bc5f 100644 --- a/litellm-rust/crates/http/src/pool.rs +++ b/litellm-rust/crates/http/src/pool.rs @@ -13,6 +13,7 @@ pub enum ClientVariant { Provider, NoRedirect, Media, + UnpinnedMedia, } const CLIENT_TTL: Duration = Duration::from_secs(3600); @@ -54,6 +55,10 @@ impl HttpClientPool { trust_proxy_env: false, ..config.clone() }, + ClientVariant::UnpinnedMedia => HttpClientConfig { + client_certificate: None, + ..config.clone() + }, ClientVariant::Provider | ClientVariant::NoRedirect => config.clone(), }; let key = (effective, variant); @@ -84,7 +89,9 @@ impl HttpClientPool { ) -> reqwest::ClientBuilder { match variant { ClientVariant::Provider => builder, - ClientVariant::NoRedirect => builder.redirect(reqwest::redirect::Policy::none()), + ClientVariant::NoRedirect | ClientVariant::UnpinnedMedia => { + builder.redirect(reqwest::redirect::Policy::none()) + } ClientVariant::Media => builder .redirect(reqwest::redirect::Policy::none()) .dns_resolver2(Arc::clone(&self.media_resolver)), @@ -125,7 +132,7 @@ mod tests { fn config(user_agent: &str) -> HttpClientConfig { HttpClientConfig { user_agent: Some(user_agent.into()), - ..HttpClientConfig::resolve(&HttpSettings::default()).unwrap() + ..HttpClientConfig::resolve(&HttpSettings::default()).config } } @@ -233,6 +240,10 @@ mod tests { .is_err() ); assert!(pool.client(&with_identity, ClientVariant::Media).is_ok()); + assert!( + pool.client(&with_identity, ClientVariant::UnpinnedMedia) + .is_ok() + ); } #[test] @@ -277,6 +288,20 @@ mod tests { assert_eq!(response.headers()["location"], "/elsewhere"); } + #[tokio::test] + async fn unpinned_media_variant_uses_the_system_resolver_and_returns_redirects() { + let (address, _, _) = serve("HTTP/1.1 302 Found").await; + let pool = HttpClientPool::new(Arc::new(FixedResolver(([192, 0, 2, 1], 80).into()))); + let response = get( + &pool, + &config("a"), + ClientVariant::UnpinnedMedia, + &format!("http://localhost:{}/doc", address.port()), + ) + .await; + assert_eq!(response.status(), 302); + } + #[tokio::test] async fn media_variant_resolves_through_the_injected_resolver() { let (address, _, requests) = serve("HTTP/1.1 204 No Content").await; diff --git a/litellm-rust/crates/http/src/proxy.rs b/litellm-rust/crates/http/src/proxy.rs new file mode 100644 index 00000000000..4dc4bf778b8 --- /dev/null +++ b/litellm-rust/crates/http/src/proxy.rs @@ -0,0 +1,15 @@ +use hyper_util::client::proxy::matcher::Matcher; + +pub struct EnvironmentProxies(Matcher); + +impl EnvironmentProxies { + pub fn from_environment() -> Self { + Self(Matcher::from_system()) + } + + pub fn apply_to(&self, url: &reqwest::Url) -> bool { + url.as_str() + .parse::() + .is_ok_and(|uri| self.0.intercept(&uri).is_some()) + } +} diff --git a/litellm-rust/crates/http/src/settings.rs b/litellm-rust/crates/http/src/settings.rs index 8aaf7f21f2c..be2f4f42fb4 100644 --- a/litellm-rust/crates/http/src/settings.rs +++ b/litellm-rust/crates/http/src/settings.rs @@ -20,6 +20,13 @@ impl SslVerify { } } +#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)] +pub struct TcpKeepalive { + pub idle: Duration, + pub interval: Duration, + pub retries: u32, +} + #[derive(Clone, Debug, PartialEq, Eq)] pub struct HttpSettings { pub ssl_verify: Option, @@ -34,6 +41,8 @@ pub struct HttpSettings { pub trust_proxy_env: bool, pub ignore_proxy_env: bool, pub connect_timeout: Duration, + pub tcp_keepalive: Option, + pub pool_idle_timeout: Duration, } impl Default for HttpSettings { @@ -51,6 +60,8 @@ impl Default for HttpSettings { trust_proxy_env: false, ignore_proxy_env: false, connect_timeout: Duration::from_secs(10), + tcp_keepalive: None, + pool_idle_timeout: Duration::from_secs(120), } } } @@ -59,6 +70,10 @@ impl HttpSettings { pub fn with_environment(self, env: &(dyn Fn(&str) -> Option + Sync)) -> Self { let enabled = |name: &str| env(name).is_some_and(|value| value.trim().eq_ignore_ascii_case("true")); + let number = |name: &str| env(name).and_then(|value| value.trim().parse::().ok()); + let seconds = |name: &str, default: u32| { + Duration::from_secs(u64::from(number(name).unwrap_or(default))) + }; Self { ssl_verify: env("SSL_VERIFY") .map(|value| SslVerify::parse(&value)) @@ -81,6 +96,17 @@ impl HttpSettings { user_agent: env("LITELLM_USER_AGENT").or(self.user_agent), trust_proxy_env: self.trust_proxy_env || enabled("AIOHTTP_TRUST_ENV"), ignore_proxy_env: self.ignore_proxy_env || enabled("DISABLE_AIOHTTP_TRUST_ENV"), + tcp_keepalive: enabled("AIOHTTP_SO_KEEPALIVE") + .then(|| TcpKeepalive { + idle: seconds("AIOHTTP_TCP_KEEPIDLE", 60), + interval: seconds("AIOHTTP_TCP_KEEPINTVL", 30), + retries: number("AIOHTTP_TCP_KEEPCNT").unwrap_or(5), + }) + .or(self.tcp_keepalive), + pool_idle_timeout: number("AIOHTTP_KEEPALIVE_TIMEOUT") + .map_or(self.pool_idle_timeout, |timeout| { + Duration::from_secs(u64::from(timeout)) + }), ..self } } @@ -188,6 +214,32 @@ mod tests { assert_eq!(settings.ssl_ecdh_curve, None); } + #[test] + fn socket_keepalive_follows_the_aiohttp_variables_with_python_defaults() { + let tuned = HttpSettings::default().with_environment(&env_of(&[ + ("AIOHTTP_SO_KEEPALIVE", "True"), + ("AIOHTTP_TCP_KEEPIDLE", "45"), + ("AIOHTTP_KEEPALIVE_TIMEOUT", "30"), + ])); + assert_eq!( + tuned.tcp_keepalive, + Some(TcpKeepalive { + idle: Duration::from_secs(45), + interval: Duration::from_secs(30), + retries: 5, + }) + ); + assert_eq!(tuned.pool_idle_timeout, Duration::from_secs(30)); + } + + #[test] + fn socket_keepalive_stays_off_unless_enabled() { + let settings = + HttpSettings::default().with_environment(&env_of(&[("AIOHTTP_TCP_KEEPIDLE", "45")])); + assert_eq!(settings.tcp_keepalive, None); + assert_eq!(settings.pool_idle_timeout, Duration::from_secs(120)); + } + #[test] fn missing_files_fall_back_to_default_verification() { let settings = HttpSettings { diff --git a/litellm-rust/crates/http/src/tls.rs b/litellm-rust/crates/http/src/tls.rs new file mode 100644 index 00000000000..605200c0be6 --- /dev/null +++ b/litellm-rust/crates/http/src/tls.rs @@ -0,0 +1,402 @@ +use std::{fmt, path::Path, sync::Arc}; + +use rustls::{ + CipherSuite, ClientConfig, DigitallySignedStruct, RootCertStore, SignatureScheme, + client::danger::{HandshakeSignatureValid, ServerCertVerified, ServerCertVerifier}, + crypto::{CryptoProvider, SupportedKxGroup, ring}, + pki_types::{CertificateDer, PrivateKeyDer, ServerName, UnixTime, pem::PemObject}, +}; + +use crate::{ + config::{HttpClientConfig, Verify}, + error::Error, +}; + +#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)] +pub enum KeyExchangeGroup { + X25519, + Secp256r1, + Secp384r1, +} + +impl KeyExchangeGroup { + pub(crate) fn from_openssl_name(name: &str) -> Result { + match name.trim().to_ascii_lowercase().as_str() { + "x25519" => Ok(Self::X25519), + "prime256v1" | "secp256r1" | "p-256" => Ok(Self::Secp256r1), + "secp384r1" | "p-384" => Ok(Self::Secp384r1), + _ => Err(Unsupported::EcdhCurve(name.to_owned())), + } + } + + fn supported(self) -> &'static dyn SupportedKxGroup { + match self { + Self::X25519 => ring::kx_group::X25519, + Self::Secp256r1 => ring::kx_group::SECP256R1, + Self::Secp384r1 => ring::kx_group::SECP384R1, + } + } +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord)] +pub enum Tls12CipherSuite { + EcdheEcdsaAes128Gcm, + EcdheEcdsaAes256Gcm, + EcdheEcdsaChacha20, + EcdheRsaAes128Gcm, + EcdheRsaAes256Gcm, + EcdheRsaChacha20, +} + +impl Tls12CipherSuite { + fn from_openssl_name(name: &str) -> Option { + match name { + "ECDHE-ECDSA-AES128-GCM-SHA256" => Some(Self::EcdheEcdsaAes128Gcm), + "ECDHE-ECDSA-AES256-GCM-SHA384" => Some(Self::EcdheEcdsaAes256Gcm), + "ECDHE-ECDSA-CHACHA20-POLY1305" => Some(Self::EcdheEcdsaChacha20), + "ECDHE-RSA-AES128-GCM-SHA256" => Some(Self::EcdheRsaAes128Gcm), + "ECDHE-RSA-AES256-GCM-SHA384" => Some(Self::EcdheRsaAes256Gcm), + "ECDHE-RSA-CHACHA20-POLY1305" => Some(Self::EcdheRsaChacha20), + _ => None, + } + } + + fn suite(self) -> CipherSuite { + match self { + Self::EcdheEcdsaAes128Gcm => CipherSuite::TLS_ECDHE_ECDSA_WITH_AES_128_GCM_SHA256, + Self::EcdheEcdsaAes256Gcm => CipherSuite::TLS_ECDHE_ECDSA_WITH_AES_256_GCM_SHA384, + Self::EcdheEcdsaChacha20 => CipherSuite::TLS_ECDHE_ECDSA_WITH_CHACHA20_POLY1305_SHA256, + Self::EcdheRsaAes128Gcm => CipherSuite::TLS_ECDHE_RSA_WITH_AES_128_GCM_SHA256, + Self::EcdheRsaAes256Gcm => CipherSuite::TLS_ECDHE_RSA_WITH_AES_256_GCM_SHA384, + Self::EcdheRsaChacha20 => CipherSuite::TLS_ECDHE_RSA_WITH_CHACHA20_POLY1305_SHA256, + } + } +} + +#[derive(Clone, Debug, PartialEq, Eq, Hash, thiserror::Error)] +pub enum Unsupported { + #[error( + "ssl_ecdh_curve {0:?} is not supported: rustls with ring only offers X25519, prime256v1 and secp384r1, so the default key exchange groups are used" + )] + EcdhCurve(String), + #[error( + "ssl_security_level {0:?} is not supported: rustls has one fixed security level, comparable to OpenSSL level 2, so legacy servers that need a lower level cannot be reached" + )] + SecurityLevel(String), + #[error( + "ssl_security_level entry {0:?} is not supported: rustls only offers ECDHE AEAD cipher suites, so the entry is ignored" + )] + CipherToken(String), +} + +pub(crate) struct CipherSelection { + pub(crate) tls12_cipher_suites: Option>, + pub(crate) unsupported: Vec, +} + +enum CipherToken { + Suite(Tls12CipherSuite), + EverySuite, + Ordering, + Unsupported(Unsupported), +} + +fn cipher_token(token: &str) -> CipherToken { + if let Some(suite) = Tls12CipherSuite::from_openssl_name(token) { + return CipherToken::Suite(suite); + } + match token { + "DEFAULT" | "ALL" | "HIGH" => CipherToken::EverySuite, + "@STRENGTH" | "@SECLEVEL=2" => CipherToken::Ordering, + level if level.starts_with("@SECLEVEL=") => { + CipherToken::Unsupported(Unsupported::SecurityLevel(level.to_owned())) + } + other => CipherToken::Unsupported(Unsupported::CipherToken(other.to_owned())), + } +} + +pub(crate) fn parse_cipher_string(value: &str) -> CipherSelection { + let tokens: Vec = tokenize(value) + .iter() + .map(|token| cipher_token(token)) + .collect(); + let every_suite = tokens + .iter() + .any(|token| matches!(token, CipherToken::EverySuite)); + let mut suites: Vec = tokens + .iter() + .filter_map(|token| match token { + CipherToken::Suite(suite) => Some(*suite), + _ => None, + }) + .collect(); + suites.sort_unstable(); + suites.dedup(); + CipherSelection { + tls12_cipher_suites: (!every_suite && !suites.is_empty()).then_some(suites), + unsupported: tokens + .into_iter() + .filter_map(|token| match token { + CipherToken::Unsupported(unsupported) => Some(unsupported), + _ => None, + }) + .collect(), + } +} + +fn tokenize(value: &str) -> Vec { + value + .split([':', ',', ' ']) + .flat_map(|entry| match entry.split_once('@') { + Some((name, command)) => vec![name.to_owned(), format!("@{command}")], + None => vec![entry.to_owned()], + }) + .filter(|token| !token.is_empty()) + .collect() +} + +pub fn client_config(config: &HttpClientConfig) -> Result { + let base = ring::default_provider(); + let provider = Arc::new(CryptoProvider { + kx_groups: config + .key_exchange_group + .map_or_else(|| base.kx_groups.clone(), |group| vec![group.supported()]), + cipher_suites: base + .cipher_suites + .iter() + .copied() + .filter(|suite| { + suite.tls13().is_some() + || config + .tls12_cipher_suites + .as_ref() + .is_none_or(|allowed| allowed.iter().any(|a| a.suite() == suite.suite())) + }) + .collect(), + ..base + }); + let builder = ClientConfig::builder_with_provider(Arc::clone(&provider)) + .with_safe_default_protocol_versions() + .map_err(|error| Error::Client(error.to_string()))?; + let verified = match &config.verify { + Verify::Disabled => builder + .dangerous() + .with_custom_certificate_verifier(Arc::new(NoVerification(provider))), + Verify::BuiltInRoots => builder.with_root_certificates(built_in_roots()), + Verify::CaBundle(path) => builder.with_root_certificates(bundle_roots(path)?), + }; + let mut tls = match &config.client_certificate { + None => verified.with_no_client_auth(), + Some(path) => { + let (chain, key) = identity(path)?; + verified + .with_client_auth_cert(chain, key) + .map_err(|error| invalid_pem(path, error))? + } + }; + tls.alpn_protocols = if config.http2 { + vec![b"h2".to_vec(), b"http/1.1".to_vec()] + } else { + vec![b"http/1.1".to_vec()] + }; + Ok(tls) +} + +fn built_in_roots() -> RootCertStore { + let mut store = RootCertStore::empty(); + store.extend(webpki_roots::TLS_SERVER_ROOTS.iter().cloned()); + store +} + +fn bundle_roots(path: &Path) -> Result { + let certificates = certificates(path)?; + if certificates.is_empty() { + return Err(invalid_pem(path, "no certificates found")); + } + let mut store = RootCertStore::empty(); + for certificate in certificates { + store + .add(certificate) + .map_err(|error| invalid_pem(path, error))?; + } + Ok(store) +} + +fn identity(path: &Path) -> Result<(Vec>, PrivateKeyDer<'static>), Error> { + let chain = certificates(path)?; + if chain.is_empty() { + return Err(invalid_pem(path, "no certificates found")); + } + let key = + PrivateKeyDer::from_pem_slice(&read(path)?).map_err(|error| invalid_pem(path, error))?; + Ok((chain, key)) +} + +fn certificates(path: &Path) -> Result>, Error> { + CertificateDer::pem_slice_iter(&read(path)?) + .collect::>() + .map_err(|error| invalid_pem(path, error)) +} + +fn read(path: &Path) -> Result, Error> { + std::fs::read(path).map_err(|error| Error::Read { + path: path.to_path_buf(), + message: error.to_string(), + }) +} + +fn invalid_pem(path: &Path, message: impl fmt::Display) -> Error { + Error::InvalidPem { + path: path.to_path_buf(), + message: message.to_string(), + } +} + +#[derive(Debug)] +struct NoVerification(Arc); + +impl ServerCertVerifier for NoVerification { + fn verify_server_cert( + &self, + _end_entity: &CertificateDer<'_>, + _intermediates: &[CertificateDer<'_>], + _server_name: &ServerName<'_>, + _ocsp_response: &[u8], + _now: UnixTime, + ) -> Result { + Ok(ServerCertVerified::assertion()) + } + + fn verify_tls12_signature( + &self, + _message: &[u8], + _cert: &CertificateDer<'_>, + _dss: &DigitallySignedStruct, + ) -> Result { + Ok(HandshakeSignatureValid::assertion()) + } + + fn verify_tls13_signature( + &self, + _message: &[u8], + _cert: &CertificateDer<'_>, + _dss: &DigitallySignedStruct, + ) -> Result { + Ok(HandshakeSignatureValid::assertion()) + } + + fn supported_verify_schemes(&self) -> Vec { + self.0.signature_verification_algorithms.supported_schemes() + } +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + use rustls::NamedGroup; + + use super::*; + use crate::HttpSettings; + + fn config(settings: HttpSettings) -> HttpClientConfig { + HttpClientConfig::resolve(&settings).config + } + + fn offered_groups(tls: &ClientConfig) -> Vec { + tls.crypto_provider() + .kx_groups + .iter() + .map(|group| group.name()) + .collect() + } + + fn offered_tls12_suites(tls: &ClientConfig) -> Vec { + tls.crypto_provider() + .cipher_suites + .iter() + .filter(|suite| suite.tls13().is_none()) + .map(|suite| suite.suite()) + .collect() + } + + #[rstest] + #[case("X25519", NamedGroup::X25519)] + #[case("prime256v1", NamedGroup::secp256r1)] + #[case("secp384r1", NamedGroup::secp384r1)] + fn ecdh_curve_is_the_only_key_exchange_group_offered( + #[case] curve: &str, + #[case] expected: NamedGroup, + ) { + let tls = client_config(&config(HttpSettings { + ssl_ecdh_curve: Some(curve.into()), + ..HttpSettings::default() + })) + .unwrap(); + assert_eq!(offered_groups(&tls), [expected]); + } + + #[test] + fn default_settings_offer_every_group_and_suite_of_the_provider() { + let tls = client_config(&config(HttpSettings::default())).unwrap(); + let provider = ring::default_provider(); + assert_eq!(offered_groups(&tls).len(), provider.kx_groups.len()); + assert_eq!( + tls.crypto_provider().cipher_suites.len(), + provider.cipher_suites.len() + ); + } + + #[test] + fn named_suites_are_the_only_tls12_suites_offered_and_tls13_stays() { + let tls = client_config(&config(HttpSettings { + ssl_security_level: Some("ECDHE-RSA-AES256-GCM-SHA384".into()), + ..HttpSettings::default() + })) + .unwrap(); + assert_eq!( + offered_tls12_suites(&tls), + [CipherSuite::TLS_ECDHE_RSA_WITH_AES_256_GCM_SHA384] + ); + assert!( + tls.crypto_provider() + .cipher_suites + .iter() + .any(|suite| suite.tls13().is_some()) + ); + } + + #[rstest] + #[case(true, &[b"h2".as_slice(), b"http/1.1".as_slice()])] + #[case(false, &[b"http/1.1".as_slice()])] + fn alpn_offers_h2_only_when_http2_is_on(#[case] http2: bool, #[case] expected: &[&[u8]]) { + let tls = client_config(&config(HttpSettings { + http2, + ..HttpSettings::default() + })) + .unwrap(); + assert_eq!(tls.alpn_protocols, expected); + } + + #[test] + fn client_certificate_without_a_private_key_is_an_invalid_pem_error() { + let path = std::env::temp_dir().join(format!( + "litellm-http-cert-without-key-{}.pem", + std::process::id() + )); + std::fs::write( + &path, + b"-----BEGIN CERTIFICATE-----\nAA==\n-----END CERTIFICATE-----\n", + ) + .unwrap(); + let result = client_config(&HttpClientConfig { + client_certificate: Some(path.clone()), + ..config(HttpSettings::default()) + }) + .map(drop); + std::fs::remove_file(&path).unwrap(); + assert!(matches!( + result, + Err(Error::InvalidPem { path: reported, .. }) if reported == path + )); + } +} diff --git a/litellm-rust/crates/llms/src/custom_httpx/llm_http_handler.rs b/litellm-rust/crates/llms/src/custom_httpx/llm_http_handler.rs index 876fa0aae87..58dc03eea2d 100644 --- a/litellm-rust/crates/llms/src/custom_httpx/llm_http_handler.rs +++ b/litellm-rust/crates/llms/src/custom_httpx/llm_http_handler.rs @@ -16,7 +16,7 @@ use crate::{ }, custom_httpx::{ http_handler::{HeaderPolicy, execute_http_request, with_headers}, - media::MediaFetcher, + media::{MediaFetcher, UrlPolicy}, transport, }, }; @@ -41,12 +41,13 @@ impl OcrClient { pub fn new( pool: &HttpClientPool, config: &HttpClientConfig, + url_policy: UrlPolicy, vertex_auth: VertexAuth, ) -> Result { Ok(Self { provider_http: pool.client(config, ClientVariant::Provider)?, polling_http: pool.client(config, ClientVariant::NoRedirect)?, - document_fetcher: MediaFetcher::new(pool, config)?, + document_fetcher: MediaFetcher::new(pool, config, url_policy)?, vertex_auth, }) } diff --git a/litellm-rust/crates/llms/src/custom_httpx/media.rs b/litellm-rust/crates/llms/src/custom_httpx/media.rs index a1c4fe68734..02d152d3ef1 100644 --- a/litellm-rust/crates/llms/src/custom_httpx/media.rs +++ b/litellm-rust/crates/llms/src/custom_httpx/media.rs @@ -7,7 +7,7 @@ use std::{ time::Duration, }; -use litellm_http::{ClientVariant, HttpClientConfig, HttpClientPool}; +use litellm_http::{ClientVariant, EnvironmentProxies, HttpClientConfig, HttpClientPool}; use reqwest::{ Url, dns::{Addrs, Name, Resolve, Resolving}, @@ -35,10 +35,45 @@ pub enum Error { Transport(#[from] crate::custom_httpx::transport::Error), } +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct UrlPolicy { + pub validate: bool, + pub allowed_hosts: Vec, +} + +impl Default for UrlPolicy { + fn default() -> Self { + Self { + validate: true, + allowed_hosts: Vec::new(), + } + } +} + +impl UrlPolicy { + fn allows(&self, host: &str, port: u16) -> bool { + let host = normalize_host(host); + let with_port = format!("{host}:{port}"); + self.allowed_hosts + .iter() + .map(|entry| normalize_host(entry)) + .any(|entry| entry == host || entry == with_port) + } +} + +fn normalize_host(host: &str) -> String { + host.to_ascii_lowercase().trim_end_matches('.').to_owned() +} + +type ProxyMatch = Arc bool + Send + Sync>; + #[derive(Clone)] pub struct MediaFetcher { - client: reqwest::Client, + pinned: reqwest::Client, + unpinned: reqwest::Client, + uses_proxy: ProxyMatch, address_resolver: Arc, + url_policy: UrlPolicy, allow_private_network: bool, } @@ -65,19 +100,36 @@ impl MediaFetcher { pub fn new( pool: &HttpClientPool, config: &HttpClientConfig, + url_policy: UrlPolicy, ) -> Result { - Self::with_address_resolver(pool, config, Arc::new(SystemAddressResolver)) + let uses_proxy: ProxyMatch = if config.trust_proxy_env { + let proxies = EnvironmentProxies::from_environment(); + Arc::new(move |url| proxies.apply_to(url)) + } else { + Arc::new(|_| false) + }; + Self::with_resolution( + pool, + config, + url_policy, + Arc::new(SystemAddressResolver), + uses_proxy, + ) } - fn with_address_resolver( + fn with_resolution( pool: &HttpClientPool, config: &HttpClientConfig, + url_policy: UrlPolicy, address_resolver: Arc, + uses_proxy: ProxyMatch, ) -> Result { - let client = pool.client(config, ClientVariant::Media)?; Ok(Self { - client, + pinned: pool.client(config, ClientVariant::Media)?, + unpinned: pool.client(config, ClientVariant::UnpinnedMedia)?, + uses_proxy, address_resolver, + url_policy, allow_private_network: false, }) } @@ -85,8 +137,11 @@ impl MediaFetcher { #[cfg(any(test, feature = "test-support"))] pub fn for_test(client: reqwest::Client) -> Self { Self { - client, + pinned: client.clone(), + unpinned: client, + uses_proxy: Arc::new(|_| false), address_resolver: Arc::new(AllowPrivateResolver), + url_policy: UrlPolicy::default(), allow_private_network: true, } } @@ -107,9 +162,9 @@ impl MediaFetcher { ) -> Result { let mut redirects_followed = 0; loop { - self.validate_url(&url).await?; let mut response = self - .client + .client_for(&url) + .await? .get(url.clone()) .send() .await @@ -156,7 +211,10 @@ impl MediaFetcher { } } - async fn validate_url(&self, url: &Url) -> Result<(), Error> { + async fn client_for(&self, url: &Url) -> Result<&reqwest::Client, Error> { + if !self.url_policy.validate { + return Ok(&self.unpinned); + } if !matches!(url.scheme(), "http" | "https") || !url.username().is_empty() || url.password().is_some() @@ -165,12 +223,28 @@ impl MediaFetcher { } let host = url.host_str().ok_or(Error::BlockedUrl)?; if self.allow_private_network { - return Ok(()); - } - if let Ok(ip) = host.parse::() { - return (!is_blocked_ip(ip)).then_some(()).ok_or(Error::BlockedUrl); + return Ok(&self.pinned); } let port = url.port_or_known_default().ok_or(Error::BlockedUrl)?; + if self.url_policy.allows(host, port) { + return Ok(&self.unpinned); + } + self.validate_host(host, port).await?; + Ok(if (self.uses_proxy)(url) { + &self.unpinned + } else { + &self.pinned + }) + } + + async fn validate_host(&self, host: &str, port: u16) -> Result<(), Error> { + if let Ok(ip) = host + .trim_start_matches('[') + .trim_end_matches(']') + .parse::() + { + return (!is_blocked_ip(ip)).then_some(()).ok_or(Error::BlockedUrl); + } let addresses = self .address_resolver .resolve(host, port) @@ -360,14 +434,35 @@ mod tests { address: SocketAddr, blocked_hosts: HashSet<&'static str>, ) -> MediaFetcher { - MediaFetcher::with_address_resolver( - &HttpClientPool::new(Arc::new(LoopbackDnsResolver(address))), - &HttpClientConfig::resolve(&HttpSettings::default()).unwrap(), + fetcher(address, blocked_hosts, UrlPolicy::default(), false) + } + + fn fetcher( + pinned_address: SocketAddr, + blocked_hosts: HashSet<&'static str>, + url_policy: UrlPolicy, + uses_proxy: bool, + ) -> MediaFetcher { + let direct = HttpClientConfig { + trust_proxy_env: false, + ..HttpClientConfig::resolve(&HttpSettings::default()).config + }; + MediaFetcher::with_resolution( + &HttpClientPool::new(Arc::new(LoopbackDnsResolver(pinned_address))), + &direct, + url_policy, Arc::new(TestAddressResolver { blocked_hosts }), + Arc::new(move |_| uses_proxy), ) .expect("test fetcher builds") } + const UNROUTABLE: SocketAddr = + SocketAddr::new(IpAddr::V4(std::net::Ipv4Addr::new(192, 0, 2, 1)), 9); + + const OK_RESPONSE: &[u8] = + b"HTTP/1.1 200 OK\r\nContent-Length: 2\r\nConnection: close\r\n\r\nok"; + fn policy(max_bytes: u64, max_redirects: usize) -> DownloadPolicy { DownloadPolicy { timeout: Duration::from_secs(1), @@ -541,14 +636,88 @@ mod tests { async fn rejects_url_credentials_before_network_access() { let fetcher = MediaFetcher::new( &HttpClientPool::new(Arc::new(PublicDnsResolver)), - &HttpClientConfig::resolve(&HttpSettings::default()).expect("default settings resolve"), + &HttpClientConfig::resolve(&HttpSettings::default()).config, + UrlPolicy::default(), ) .expect("media fetcher builds"); let url = Url::parse("https://user:password@8.8.8.8/document").expect("credentialed URL parses"); assert!(matches!( - fetcher.validate_url(&url).await, + fetcher.fetch(url, policy(1, 0)).await, Err(Error::BlockedUrl) )); } + + #[tokio::test] + async fn allowlisted_private_host_is_fetched_without_the_pinned_resolver() { + let (url, server, _) = serve_named("localhost", vec![OK_RESPONSE]).await; + let port = url.port().expect("test URL has a port"); + let allowed = UrlPolicy { + validate: true, + allowed_hosts: vec![format!("LOCALHOST:{port}")], + }; + let media = fetcher(UNROUTABLE, HashSet::from(["localhost"]), allowed, false) + .fetch(url, policy(2, 0)) + .await + .expect("allowlisted host downloads"); + server.await.expect("server completes"); + assert_eq!(media.bytes, b"ok"); + } + + #[tokio::test] + async fn allowlist_entry_for_another_port_does_not_open_the_host() { + let (url, _server, _) = serve_named("localhost", vec![OK_RESPONSE]).await; + let other_port = UrlPolicy { + validate: true, + allowed_hosts: vec!["localhost:1".into()], + }; + let result = fetcher(UNROUTABLE, HashSet::from(["localhost"]), other_port, false) + .fetch(url, policy(2, 0)) + .await; + assert!(matches!(result, Err(Error::BlockedUrl))); + } + + #[tokio::test] + async fn validation_off_fetches_private_hosts_and_follows_redirects() { + let (url, server, _) = serve_named( + "localhost", + vec![ + b"HTTP/1.1 302 Found\r\nLocation: /moved\r\nContent-Length: 0\r\nConnection: close\r\n\r\n", + OK_RESPONSE, + ], + ) + .await; + let off = UrlPolicy { + validate: false, + allowed_hosts: Vec::new(), + }; + let media = fetcher(UNROUTABLE, HashSet::from(["localhost"]), off, false) + .fetch(url, policy(2, 1)) + .await + .expect("unvalidated download succeeds"); + let requests = server.await.expect("server completes"); + assert_eq!(media.bytes, b"ok"); + assert!(requests[1].starts_with("GET /moved ")); + } + + #[tokio::test] + async fn proxied_urls_skip_the_pinned_resolver_but_keep_the_address_check() { + let (url, server, _) = serve_named("localhost", vec![OK_RESPONSE]).await; + let media = fetcher(UNROUTABLE, HashSet::new(), UrlPolicy::default(), true) + .fetch(url.clone(), policy(2, 0)) + .await + .expect("public host behind a proxy downloads"); + server.await.expect("server completes"); + assert_eq!(media.bytes, b"ok"); + + let blocked = fetcher( + UNROUTABLE, + HashSet::from(["localhost"]), + UrlPolicy::default(), + true, + ) + .fetch(url, policy(2, 0)) + .await; + assert!(matches!(blocked, Err(Error::BlockedUrl))); + } } diff --git a/litellm-rust/crates/python-bridge/src/http.rs b/litellm-rust/crates/python-bridge/src/http.rs index c5952c53132..118f4669b62 100644 --- a/litellm-rust/crates/python-bridge/src/http.rs +++ b/litellm-rust/crates/python-bridge/src/http.rs @@ -1,10 +1,11 @@ use std::{ + collections::HashSet, path::{Path, PathBuf}, - sync::{Arc, LazyLock}, + sync::{Arc, LazyLock, Mutex, PoisonError}, }; -use litellm_http::{HttpClientConfig, HttpClientPool, HttpSettings, SslVerify}; -use litellm_llms::custom_httpx::media::PublicDnsResolver; +use litellm_http::{HttpClientConfig, HttpClientPool, HttpSettings, SslVerify, Unsupported}; +use litellm_llms::custom_httpx::media::{PublicDnsResolver, UrlPolicy}; use pyo3::{prelude::*, types::PyDict}; use crate::{errors::RustBridgeDeclined, python_settings::PythonSettings}; @@ -12,6 +13,8 @@ use crate::{errors::RustBridgeDeclined, python_settings::PythonSettings}; static POOL: LazyLock = LazyLock::new(|| HttpClientPool::new(Arc::new(PublicDnsResolver))); +static REPORTED_UNSUPPORTED: LazyLock>> = LazyLock::new(Mutex::default); + pub(crate) fn pool() -> &'static HttpClientPool { &POOL } @@ -21,22 +24,48 @@ pub(crate) fn call_config( kwargs: &Bound<'_, PyDict>, asynchronous: bool, ) -> PyResult { - decline_live_client(kwargs)?; - decline_custom_url_policy(&PythonSettings::UrlPolicy.read(py)?)?; let configured = settings(&PythonSettings::Http.read(py)?)? .with_environment(&|name| std::env::var(name).ok()); let settings = for_call(configured, call_ssl_verify(kwargs)?, asynchronous) .without_missing_files(&|path: &Path| path.exists()); - HttpClientConfig::resolve(&settings) - .map_err(|error| RustBridgeDeclined::new_err(error.to_string())) + let resolution = HttpClientConfig::resolve(&settings); + for unsupported in unreported(&REPORTED_UNSUPPORTED, resolution.unsupported) { + PythonSettings::warn(py, &unsupported.to_string())?; + } + Ok(resolution.config) +} + +fn unreported( + reported: &Mutex>, + unsupported: Vec, +) -> Vec { + let mut reported = reported.lock().unwrap_or_else(PoisonError::into_inner); + unsupported + .into_iter() + .filter(|unsupported| reported.insert(unsupported.clone())) + .collect() +} + +pub(crate) fn url_policy(py: Python<'_>) -> PyResult { + let policy: PythonUrlPolicy = + PythonSettings::UrlPolicy + .read(py)? + .extract() + .map_err(|error: PyErr| { + RustBridgeDeclined::new_err(format!( + "litellm URL policy cannot be used by the Rust route: {error}" + )) + })?; + Ok(UrlPolicy { + validate: policy.user_url_validation, + allowed_hosts: policy.user_url_allowed_hosts, + }) } fn call_ssl_verify(kwargs: &Bound<'_, PyDict>) -> PyResult> { - kwargs + Ok(kwargs .get_item("ssl_verify")? - .filter(|value| !value.is_none()) - .map(|value| ssl_verify(&value, "the ssl_verify argument")) - .transpose() + .and_then(|value| ssl_verify(&value))) } fn for_call( @@ -51,35 +80,12 @@ fn for_call( } } -fn decline_live_client(kwargs: &Bound<'_, PyDict>) -> PyResult<()> { - if kwargs - .get_item("client")? - .is_some_and(|value| !value.is_none()) - { - return Err(RustBridgeDeclined::new_err( - "client is a live Python HTTP client and cannot be used by the Rust route", - )); - } - Ok(()) -} - #[derive(FromPyObject)] struct PythonUrlPolicy { user_url_validation: bool, user_url_allowed_hosts: Vec, } -fn decline_custom_url_policy(value: &Bound<'_, PyAny>) -> PyResult<()> { - match value.extract::() { - Ok(policy) if policy.user_url_validation && policy.user_url_allowed_hosts.is_empty() => { - Ok(()) - } - Ok(_) | Err(_) => Err(RustBridgeDeclined::new_err( - "litellm.user_url_validation / user_url_allowed_hosts are applied by the Python route", - )), - } -} - #[derive(FromPyObject)] struct PythonHttpSettings<'py> { ssl_verify: Bound<'py, PyAny>, @@ -101,7 +107,7 @@ fn settings(value: &Bound<'_, PyAny>) -> PyResult { )) })?; Ok(HttpSettings { - ssl_verify: Some(ssl_verify(&python.ssl_verify, "litellm.ssl_verify")?), + ssl_verify: ssl_verify(&python.ssl_verify), ssl_certificate: python.ssl_certificate.map(PathBuf::from), ssl_security_level: python.ssl_security_level, ssl_ecdh_curve: python.ssl_ecdh_curve, @@ -115,20 +121,18 @@ fn settings(value: &Bound<'_, PyAny>) -> PyResult { }) } -fn ssl_verify(value: &Bound<'_, PyAny>, source: &str) -> PyResult { +fn ssl_verify(value: &Bound<'_, PyAny>) -> Option { if let Ok(enabled) = value.extract::() { - return Ok(if enabled { + return Some(if enabled { SslVerify::Enabled } else { SslVerify::Disabled }); } - if let Ok(path) = value.extract::() { - return Ok(SslVerify::parse(&path)); - } - Err(RustBridgeDeclined::new_err(format!( - "{source} is a live Python object and cannot be used by the Rust route" - ))) + value + .extract::() + .ok() + .map(|path| SslVerify::parse(&path)) } #[cfg(test)] @@ -247,50 +251,30 @@ user_agent='litellm/9.9.9', Python::initialize(); Python::attach(|py| { let settings = settings(&python_settings(py, overrides)).unwrap(); - let config = HttpClientConfig::resolve(&settings).unwrap(); + let config = HttpClientConfig::resolve(&settings).config; assert_eq!(config.verify, expected); }); } #[test] - fn ssl_context_global_declines_instead_of_being_dropped() { + fn ssl_context_global_is_ignored_so_environment_and_defaults_apply() { Python::initialize(); Python::attach(|py| { - let error = settings(&python_settings(py, "ssl_verify=object()")).unwrap_err(); - assert!(error.is_instance_of::(py)); - assert!(error.value(py).to_string().contains("litellm.ssl_verify")); + let settings = settings(&python_settings(py, "ssl_verify=object()")).unwrap(); + assert_eq!(settings.ssl_verify, None); }); } - fn url_policy<'py>(py: Python<'py>, fields: &str) -> Bound<'py, PyAny> { - let source = std::ffi::CString::new(format!( - "import types\npolicy = types.SimpleNamespace({fields})" - )) - .unwrap(); - let locals = PyDict::new(py); - py.run(&source, Some(&locals), Some(&locals)).unwrap(); - locals.get_item("policy").unwrap().unwrap() - } - #[test] - fn default_url_policy_stays_on_the_rust_route() { - Python::initialize(); - Python::attach(|py| { - let policy = url_policy(py, "user_url_validation=True, user_url_allowed_hosts=[]"); - decline_custom_url_policy(&policy).unwrap(); - }); - } - - #[rstest] - #[case::validation_off("user_url_validation=False, user_url_allowed_hosts=[]")] - #[case::allowlist("user_url_validation=True, user_url_allowed_hosts=['docs.internal']")] - #[case::mistyped("user_url_validation=True, user_url_allowed_hosts=None")] - fn custom_url_policy_declines_so_python_applies_it(#[case] fields: &str) { - Python::initialize(); - Python::attach(|py| { - let error = decline_custom_url_policy(&url_policy(py, fields)).unwrap_err(); - assert!(error.is_instance_of::(py)); - }); + fn unsupported_settings_are_reported_once_per_process() { + let reported = Mutex::default(); + let curve = Unsupported::EcdhCurve("secp521r1".into()); + let level = Unsupported::SecurityLevel("@SECLEVEL=1".into()); + assert_eq!( + unreported(&reported, vec![curve.clone(), level.clone()]), + [curve.clone(), level] + ); + assert_eq!(unreported(&reported, vec![curve]), []); } #[test] @@ -333,15 +317,19 @@ user_agent='litellm/9.9.9', } #[test] - fn live_ssl_context_argument_declines() { + fn live_ssl_context_argument_is_ignored_so_the_configured_value_applies() { Python::initialize(); Python::attach(|py| { let kwargs = PyDict::new(py); kwargs .set_item("ssl_verify", py.eval(c"object()", None, None).unwrap()) .unwrap(); - let error = call_ssl_verify(&kwargs).unwrap_err(); - assert!(error.is_instance_of::(py)); + let configured = HttpSettings { + ssl_verify: Some(SslVerify::Disabled), + ..HttpSettings::default() + }; + let settings = for_call(configured.clone(), call_ssl_verify(&kwargs).unwrap(), true); + assert_eq!(settings, configured); }); } @@ -357,33 +345,7 @@ user_agent='litellm/9.9.9', ..HttpSettings::default() }; let settings = for_call(opted_out, None, asynchronous); - let config = HttpClientConfig::resolve(&settings).unwrap(); + let config = HttpClientConfig::resolve(&settings).config; assert_eq!(config.trust_proxy_env, expected); } - - #[test] - fn live_python_client_declines_before_dispatch() { - Python::initialize(); - Python::attach(|py| { - let kwargs = PyDict::new(py); - kwargs - .set_item("client", py.eval(c"object()", None, None).unwrap()) - .unwrap(); - let error = decline_live_client(&kwargs).unwrap_err(); - assert!(error.is_instance_of::(py)); - }); - } - - #[rstest] - #[case::absent_client("{}")] - #[case::none_client("{'client': None}")] - #[case::proxy_shared_session("{'shared_session': object()}")] - fn calls_without_a_python_client_stay_on_the_rust_route(#[case] kwargs: &str) { - Python::initialize(); - Python::attach(|py| { - let source = std::ffi::CString::new(kwargs).unwrap(); - let kwargs = py.eval(&source, None, None).unwrap(); - decline_live_client(kwargs.cast::().unwrap()).unwrap(); - }); - } } diff --git a/litellm-rust/crates/python-bridge/src/python_settings.rs b/litellm-rust/crates/python-bridge/src/python_settings.rs index b7855566850..79921d67452 100644 --- a/litellm-rust/crates/python-bridge/src/python_settings.rs +++ b/litellm-rust/crates/python-bridge/src/python_settings.rs @@ -22,6 +22,11 @@ impl PythonSettings { pub(crate) fn read(self, py: Python<'_>) -> PyResult> { py.import(MODULE)?.getattr(self.name())?.call0() } + + pub(crate) fn warn(py: Python<'_>, message: &str) -> PyResult<()> { + py.import(MODULE)?.getattr("warn")?.call1((message,))?; + Ok(()) + } } #[cfg(test)] diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs b/litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs index 174d0ff18c8..f9d7024c824 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs @@ -38,8 +38,13 @@ fn run_ocr( asynchronous: bool, ) -> PyResult> { let config = http::call_config(py, &kwargs, asynchronous)?; - let client = OcrClient::new(http::pool(), &config, VERTEX_AUTH.clone()) - .map_err(|error| RustBridgeDeclined::new_err(error.to_string()))?; + let client = OcrClient::new( + http::pool(), + &config, + http::url_policy(py)?, + VERTEX_AUTH.clone(), + ) + .map_err(|error| RustBridgeDeclined::new_err(error.to_string()))?; run_legacy_call( py, if asynchronous { ASYNC_SURFACE } else { SURFACE }, diff --git a/litellm/rust_bridge/settings.py b/litellm/rust_bridge/settings.py index bccfd01ec73..e170f93b198 100644 --- a/litellm/rust_bridge/settings.py +++ b/litellm/rust_bridge/settings.py @@ -24,6 +24,12 @@ class UrlPolicy: user_url_allowed_hosts: Sequence[str] +def warn(message: str) -> None: + from litellm._logging import verbose_logger + + verbose_logger.warning("%s", message) + + def url_policy() -> UrlPolicy: import litellm diff --git a/tests/test_litellm/rust_bridge/test_settings.py b/tests/test_litellm/rust_bridge/test_settings.py index 7e7b1c6743b..f75145c2b2c 100644 --- a/tests/test_litellm/rust_bridge/test_settings.py +++ b/tests/test_litellm/rust_bridge/test_settings.py @@ -1,4 +1,5 @@ import dataclasses +import logging from pathlib import Path from typing import Final @@ -65,3 +66,10 @@ def test_http_settings_ignores_environment_overrides(monkeypatch: pytest.MonkeyP assert result.user_agent == default_user_agent() assert result.ssl_verify is True + + +def test_warn_reaches_the_litellm_logger(caplog: pytest.LogCaptureFixture) -> None: + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + settings.warn("ssl_ecdh_curve 'secp521r1' is not supported") + + assert [record.getMessage() for record in caplog.records] == ["ssl_ecdh_curve 'secp521r1' is not supported"] From ffbfe7205fa10c1f56b2205583728bf614a95419 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Fri, 18 Sep 2026 19:40:16 -0700 Subject: [PATCH 136/144] refactor(rust): parse TLS settings through FromStr, From and TryFrom KeyExchangeGroup and Tls12CipherSuite parse with FromStr and fail with Unsupported, so a setting rustls cannot honor is a typed error instead of a missing value. The cipher string conversions cannot fail and use From. The rustls ClientConfig is built with TryFrom<&HttpClientConfig>, and the built-in root store is constructed in one expression --- litellm-rust/crates/http/src/config.rs | 8 +- litellm-rust/crates/http/src/lib.rs | 2 +- litellm-rust/crates/http/src/tls.rs | 208 +++++++++++++------------ 3 files changed, 113 insertions(+), 105 deletions(-) diff --git a/litellm-rust/crates/http/src/config.rs b/litellm-rust/crates/http/src/config.rs index ebe557788ca..30216405fc8 100644 --- a/litellm-rust/crates/http/src/config.rs +++ b/litellm-rust/crates/http/src/config.rs @@ -7,7 +7,7 @@ use std::{ use crate::{ error::Error, settings::{HttpSettings, SslVerify, TcpKeepalive}, - tls::{self, KeyExchangeGroup, Tls12CipherSuite, Unsupported}, + tls::{CipherSelection, KeyExchangeGroup, Tls12CipherSuite, Unsupported}, }; #[derive(Clone, Debug, PartialEq, Eq, Hash)] @@ -43,7 +43,7 @@ impl HttpClientConfig { let (key_exchange_group, unsupported_curve) = match settings .ssl_ecdh_curve .as_deref() - .map(KeyExchangeGroup::from_openssl_name) + .map(str::parse::) { None => (None, None), Some(Ok(group)) => (Some(group), None), @@ -52,7 +52,7 @@ impl HttpClientConfig { let ciphers = settings .ssl_security_level .as_deref() - .map(tls::parse_cipher_string); + .map(CipherSelection::from); let verify = match &settings.ssl_verify { Some(SslVerify::Disabled) => Verify::Disabled, Some(SslVerify::CaBundle(path)) => Verify::CaBundle(path.clone()), @@ -91,7 +91,7 @@ impl HttpClientConfig { pub fn client_builder(&self) -> Result { let base = reqwest::Client::builder() - .use_preconfigured_tls(tls::client_config(self)?) + .use_preconfigured_tls(rustls::ClientConfig::try_from(self)?) .connect_timeout(self.connect_timeout) .pool_idle_timeout(self.pool_idle_timeout); let with_keepalive = match self.tcp_keepalive { diff --git a/litellm-rust/crates/http/src/lib.rs b/litellm-rust/crates/http/src/lib.rs index 45f370d3a9c..e222d0e3f50 100644 --- a/litellm-rust/crates/http/src/lib.rs +++ b/litellm-rust/crates/http/src/lib.rs @@ -10,4 +10,4 @@ pub use error::Error; pub use pool::{ClientVariant, HttpClientPool}; pub use proxy::EnvironmentProxies; pub use settings::{HttpSettings, SslVerify, TcpKeepalive}; -pub use tls::{KeyExchangeGroup, Tls12CipherSuite, Unsupported, client_config}; +pub use tls::{KeyExchangeGroup, Tls12CipherSuite, Unsupported}; diff --git a/litellm-rust/crates/http/src/tls.rs b/litellm-rust/crates/http/src/tls.rs index 605200c0be6..49405b97366 100644 --- a/litellm-rust/crates/http/src/tls.rs +++ b/litellm-rust/crates/http/src/tls.rs @@ -1,4 +1,4 @@ -use std::{fmt, path::Path, sync::Arc}; +use std::{fmt, path::Path, str::FromStr, sync::Arc}; use rustls::{ CipherSuite, ClientConfig, DigitallySignedStruct, RootCertStore, SignatureScheme, @@ -19,8 +19,10 @@ pub enum KeyExchangeGroup { Secp384r1, } -impl KeyExchangeGroup { - pub(crate) fn from_openssl_name(name: &str) -> Result { +impl FromStr for KeyExchangeGroup { + type Err = Unsupported; + + fn from_str(name: &str) -> Result { match name.trim().to_ascii_lowercase().as_str() { "x25519" => Ok(Self::X25519), "prime256v1" | "secp256r1" | "p-256" => Ok(Self::Secp256r1), @@ -28,7 +30,9 @@ impl KeyExchangeGroup { _ => Err(Unsupported::EcdhCurve(name.to_owned())), } } +} +impl KeyExchangeGroup { fn supported(self) -> &'static dyn SupportedKxGroup { match self { Self::X25519 => ring::kx_group::X25519, @@ -48,19 +52,23 @@ pub enum Tls12CipherSuite { EcdheRsaChacha20, } -impl Tls12CipherSuite { - fn from_openssl_name(name: &str) -> Option { +impl FromStr for Tls12CipherSuite { + type Err = Unsupported; + + fn from_str(name: &str) -> Result { match name { - "ECDHE-ECDSA-AES128-GCM-SHA256" => Some(Self::EcdheEcdsaAes128Gcm), - "ECDHE-ECDSA-AES256-GCM-SHA384" => Some(Self::EcdheEcdsaAes256Gcm), - "ECDHE-ECDSA-CHACHA20-POLY1305" => Some(Self::EcdheEcdsaChacha20), - "ECDHE-RSA-AES128-GCM-SHA256" => Some(Self::EcdheRsaAes128Gcm), - "ECDHE-RSA-AES256-GCM-SHA384" => Some(Self::EcdheRsaAes256Gcm), - "ECDHE-RSA-CHACHA20-POLY1305" => Some(Self::EcdheRsaChacha20), - _ => None, + "ECDHE-ECDSA-AES128-GCM-SHA256" => Ok(Self::EcdheEcdsaAes128Gcm), + "ECDHE-ECDSA-AES256-GCM-SHA384" => Ok(Self::EcdheEcdsaAes256Gcm), + "ECDHE-ECDSA-CHACHA20-POLY1305" => Ok(Self::EcdheEcdsaChacha20), + "ECDHE-RSA-AES128-GCM-SHA256" => Ok(Self::EcdheRsaAes128Gcm), + "ECDHE-RSA-AES256-GCM-SHA384" => Ok(Self::EcdheRsaAes256Gcm), + "ECDHE-RSA-CHACHA20-POLY1305" => Ok(Self::EcdheRsaChacha20), + _ => Err(Unsupported::CipherToken(name.to_owned())), } } +} +impl Tls12CipherSuite { fn suite(self) -> CipherSuite { match self { Self::EcdheEcdsaAes128Gcm => CipherSuite::TLS_ECDHE_ECDSA_WITH_AES_128_GCM_SHA256, @@ -101,46 +109,47 @@ enum CipherToken { Unsupported(Unsupported), } -fn cipher_token(token: &str) -> CipherToken { - if let Some(suite) = Tls12CipherSuite::from_openssl_name(token) { - return CipherToken::Suite(suite); - } - match token { - "DEFAULT" | "ALL" | "HIGH" => CipherToken::EverySuite, - "@STRENGTH" | "@SECLEVEL=2" => CipherToken::Ordering, - level if level.starts_with("@SECLEVEL=") => { - CipherToken::Unsupported(Unsupported::SecurityLevel(level.to_owned())) +impl From<&str> for CipherToken { + fn from(token: &str) -> Self { + match token { + "DEFAULT" | "ALL" | "HIGH" => Self::EverySuite, + "@STRENGTH" | "@SECLEVEL=2" => Self::Ordering, + level if level.starts_with("@SECLEVEL=") => { + Self::Unsupported(Unsupported::SecurityLevel(level.to_owned())) + } + name => name.parse().map_or_else(Self::Unsupported, Self::Suite), } - other => CipherToken::Unsupported(Unsupported::CipherToken(other.to_owned())), } } -pub(crate) fn parse_cipher_string(value: &str) -> CipherSelection { - let tokens: Vec = tokenize(value) - .iter() - .map(|token| cipher_token(token)) - .collect(); - let every_suite = tokens - .iter() - .any(|token| matches!(token, CipherToken::EverySuite)); - let mut suites: Vec = tokens - .iter() - .filter_map(|token| match token { - CipherToken::Suite(suite) => Some(*suite), - _ => None, - }) - .collect(); - suites.sort_unstable(); - suites.dedup(); - CipherSelection { - tls12_cipher_suites: (!every_suite && !suites.is_empty()).then_some(suites), - unsupported: tokens - .into_iter() +impl From<&str> for CipherSelection { + fn from(value: &str) -> Self { + let tokens: Vec = tokenize(value) + .iter() + .map(|token| CipherToken::from(token.as_str())) + .collect(); + let every_suite = tokens + .iter() + .any(|token| matches!(token, CipherToken::EverySuite)); + let mut suites: Vec = tokens + .iter() .filter_map(|token| match token { - CipherToken::Unsupported(unsupported) => Some(unsupported), + CipherToken::Suite(suite) => Some(*suite), _ => None, }) - .collect(), + .collect(); + suites.sort_unstable(); + suites.dedup(); + CipherSelection { + tls12_cipher_suites: (!every_suite && !suites.is_empty()).then_some(suites), + unsupported: tokens + .into_iter() + .filter_map(|token| match token { + CipherToken::Unsupported(unsupported) => Some(unsupported), + _ => None, + }) + .collect(), + } } } @@ -155,57 +164,56 @@ fn tokenize(value: &str) -> Vec { .collect() } -pub fn client_config(config: &HttpClientConfig) -> Result { - let base = ring::default_provider(); - let provider = Arc::new(CryptoProvider { - kx_groups: config - .key_exchange_group - .map_or_else(|| base.kx_groups.clone(), |group| vec![group.supported()]), - cipher_suites: base - .cipher_suites - .iter() - .copied() - .filter(|suite| { - suite.tls13().is_some() - || config - .tls12_cipher_suites - .as_ref() - .is_none_or(|allowed| allowed.iter().any(|a| a.suite() == suite.suite())) - }) - .collect(), - ..base - }); - let builder = ClientConfig::builder_with_provider(Arc::clone(&provider)) - .with_safe_default_protocol_versions() - .map_err(|error| Error::Client(error.to_string()))?; - let verified = match &config.verify { - Verify::Disabled => builder - .dangerous() - .with_custom_certificate_verifier(Arc::new(NoVerification(provider))), - Verify::BuiltInRoots => builder.with_root_certificates(built_in_roots()), - Verify::CaBundle(path) => builder.with_root_certificates(bundle_roots(path)?), - }; - let mut tls = match &config.client_certificate { - None => verified.with_no_client_auth(), - Some(path) => { - let (chain, key) = identity(path)?; - verified - .with_client_auth_cert(chain, key) - .map_err(|error| invalid_pem(path, error))? - } - }; - tls.alpn_protocols = if config.http2 { - vec![b"h2".to_vec(), b"http/1.1".to_vec()] - } else { - vec![b"http/1.1".to_vec()] - }; - Ok(tls) -} +impl TryFrom<&HttpClientConfig> for ClientConfig { + type Error = Error; -fn built_in_roots() -> RootCertStore { - let mut store = RootCertStore::empty(); - store.extend(webpki_roots::TLS_SERVER_ROOTS.iter().cloned()); - store + fn try_from(config: &HttpClientConfig) -> Result { + let base = ring::default_provider(); + let provider = Arc::new(CryptoProvider { + kx_groups: config + .key_exchange_group + .map_or_else(|| base.kx_groups.clone(), |group| vec![group.supported()]), + cipher_suites: base + .cipher_suites + .iter() + .copied() + .filter(|suite| { + suite.tls13().is_some() + || config.tls12_cipher_suites.as_ref().is_none_or(|allowed| { + allowed.iter().any(|a| a.suite() == suite.suite()) + }) + }) + .collect(), + ..base + }); + let builder = ClientConfig::builder_with_provider(Arc::clone(&provider)) + .with_safe_default_protocol_versions() + .map_err(|error| Error::Client(error.to_string()))?; + let verified = match &config.verify { + Verify::Disabled => builder + .dangerous() + .with_custom_certificate_verifier(Arc::new(NoVerification(provider))), + Verify::BuiltInRoots => builder.with_root_certificates(RootCertStore { + roots: webpki_roots::TLS_SERVER_ROOTS.to_vec(), + }), + Verify::CaBundle(path) => builder.with_root_certificates(bundle_roots(path)?), + }; + let mut tls = match &config.client_certificate { + None => verified.with_no_client_auth(), + Some(path) => { + let (chain, key) = identity(path)?; + verified + .with_client_auth_cert(chain, key) + .map_err(|error| invalid_pem(path, error))? + } + }; + tls.alpn_protocols = if config.http2 { + vec![b"h2".to_vec(), b"http/1.1".to_vec()] + } else { + vec![b"http/1.1".to_vec()] + }; + Ok(tls) + } } fn bundle_roots(path: &Path) -> Result { @@ -327,7 +335,7 @@ mod tests { #[case] curve: &str, #[case] expected: NamedGroup, ) { - let tls = client_config(&config(HttpSettings { + let tls = ClientConfig::try_from(&config(HttpSettings { ssl_ecdh_curve: Some(curve.into()), ..HttpSettings::default() })) @@ -337,7 +345,7 @@ mod tests { #[test] fn default_settings_offer_every_group_and_suite_of_the_provider() { - let tls = client_config(&config(HttpSettings::default())).unwrap(); + let tls = ClientConfig::try_from(&config(HttpSettings::default())).unwrap(); let provider = ring::default_provider(); assert_eq!(offered_groups(&tls).len(), provider.kx_groups.len()); assert_eq!( @@ -348,7 +356,7 @@ mod tests { #[test] fn named_suites_are_the_only_tls12_suites_offered_and_tls13_stays() { - let tls = client_config(&config(HttpSettings { + let tls = ClientConfig::try_from(&config(HttpSettings { ssl_security_level: Some("ECDHE-RSA-AES256-GCM-SHA384".into()), ..HttpSettings::default() })) @@ -369,7 +377,7 @@ mod tests { #[case(true, &[b"h2".as_slice(), b"http/1.1".as_slice()])] #[case(false, &[b"http/1.1".as_slice()])] fn alpn_offers_h2_only_when_http2_is_on(#[case] http2: bool, #[case] expected: &[&[u8]]) { - let tls = client_config(&config(HttpSettings { + let tls = ClientConfig::try_from(&config(HttpSettings { http2, ..HttpSettings::default() })) @@ -388,7 +396,7 @@ mod tests { b"-----BEGIN CERTIFICATE-----\nAA==\n-----END CERTIFICATE-----\n", ) .unwrap(); - let result = client_config(&HttpClientConfig { + let result = ClientConfig::try_from(&HttpClientConfig { client_certificate: Some(path.clone()), ..config(HttpSettings::default()) }) From 80dbb2a28a57660bd4c57a55ee4abb90d10a6d2f Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Fri, 18 Sep 2026 19:47:10 -0700 Subject: [PATCH 137/144] refactor(rust): resolve the http client config through From and TryFrom HttpClientConfig::resolve becomes From<&HttpSettings> for Resolution and client_builder becomes TryFrom<&HttpClientConfig> for reqwest::ClientBuilder, matching the rustls conversion. The verify decision moves into From<&HttpSettings> for Verify, and the proxy environment rule moves next to its flags as HttpSettings::trusts_proxy_env. The curve and cipher results are read with transpose and a default selection, which removes the tuple destructuring --- litellm-rust/crates/core/tests/ocr.rs | 4 +- litellm-rust/crates/http/AGENTS.md | 1 + litellm-rust/crates/http/src/config.rs | 112 +++++++++--------- litellm-rust/crates/http/src/pool.rs | 8 +- litellm-rust/crates/http/src/settings.rs | 4 + litellm-rust/crates/http/src/tls.rs | 5 +- .../crates/llms/src/custom_httpx/media.rs | 6 +- litellm-rust/crates/python-bridge/src/http.rs | 10 +- 8 files changed, 78 insertions(+), 72 deletions(-) create mode 100644 litellm-rust/crates/http/AGENTS.md diff --git a/litellm-rust/crates/core/tests/ocr.rs b/litellm-rust/crates/core/tests/ocr.rs index a6b26bd8a27..e7a8fc0abc1 100644 --- a/litellm-rust/crates/core/tests/ocr.rs +++ b/litellm-rust/crates/core/tests/ocr.rs @@ -6,7 +6,7 @@ use litellm_host::{ host::{Host, HostOp, HostResult}, machine::{HostFailure, Machine, MachineStep}, }; -use litellm_http::{HttpClientConfig, HttpClientPool, HttpSettings}; +use litellm_http::{HttpClientPool, HttpSettings, Resolution}; use litellm_llms::{ base_llm::ocr::{ error::Error as OcrError, @@ -184,7 +184,7 @@ async fn ocr_client_uses_the_injected_http_pool_configuration() { }; let client = OcrClient::new( &HttpClientPool::new(Arc::new(PublicDnsResolver)), - &HttpClientConfig::resolve(&settings).config, + &Resolution::from(&settings).config, UrlPolicy::default(), VertexAuth::default(), ) diff --git a/litellm-rust/crates/http/AGENTS.md b/litellm-rust/crates/http/AGENTS.md new file mode 100644 index 00000000000..08fa34bd799 --- /dev/null +++ b/litellm-rust/crates/http/AGENTS.md @@ -0,0 +1 @@ +- https://github.com/BerriAI/litellm-docs/blob/main/docs/guides/security_settings.md diff --git a/litellm-rust/crates/http/src/config.rs b/litellm-rust/crates/http/src/config.rs index 30216405fc8..a6cc08c210d 100644 --- a/litellm-rust/crates/http/src/config.rs +++ b/litellm-rust/crates/http/src/config.rs @@ -38,84 +38,80 @@ pub struct Resolution { pub unsupported: Vec, } -impl HttpClientConfig { - pub fn resolve(settings: &HttpSettings) -> Resolution { - let (key_exchange_group, unsupported_curve) = match settings - .ssl_ecdh_curve - .as_deref() - .map(str::parse::) - { - None => (None, None), - Some(Ok(group)) => (Some(group), None), - Some(Err(unsupported)) => (None, Some(unsupported)), - }; - let ciphers = settings - .ssl_security_level - .as_deref() - .map(CipherSelection::from); - let verify = match &settings.ssl_verify { - Some(SslVerify::Disabled) => Verify::Disabled, - Some(SslVerify::CaBundle(path)) => Verify::CaBundle(path.clone()), +impl From<&HttpSettings> for Verify { + fn from(settings: &HttpSettings) -> Self { + match &settings.ssl_verify { + Some(SslVerify::Disabled) => Self::Disabled, + Some(SslVerify::CaBundle(path)) => Self::CaBundle(path.clone()), Some(SslVerify::Enabled) | None => settings .ssl_cert_file .clone() - .map_or(Verify::BuiltInRoots, Verify::CaBundle), - }; - let (tls12_cipher_suites, unsupported_ciphers) = ciphers - .map_or((None, Vec::new()), |ciphers| { - (ciphers.tls12_cipher_suites, ciphers.unsupported) - }); - Resolution { - config: Self { - verify, + .map_or(Self::BuiltInRoots, Self::CaBundle), + } + } +} + +impl From<&HttpSettings> for Resolution { + fn from(settings: &HttpSettings) -> Self { + let curve = settings + .ssl_ecdh_curve + .as_deref() + .map(str::parse::) + .transpose(); + let ciphers = settings + .ssl_security_level + .as_deref() + .map(CipherSelection::from) + .unwrap_or_default(); + Self { + config: HttpClientConfig { + verify: Verify::from(settings), client_certificate: settings.ssl_certificate.clone(), - key_exchange_group, - tls12_cipher_suites, + key_exchange_group: curve.clone().ok().flatten(), + tls12_cipher_suites: ciphers.tls12_cipher_suites, force_ipv4: settings.force_ipv4, http2: settings.http2, user_agent: settings.user_agent.clone(), - trust_proxy_env: !settings.ignore_proxy_env - || settings.trust_proxy_env - || settings.http2 - || settings.httpx_transport, + trust_proxy_env: settings.trusts_proxy_env(), connect_timeout: settings.connect_timeout, tcp_keepalive: settings.tcp_keepalive, pool_idle_timeout: settings.pool_idle_timeout, }, - unsupported: unsupported_curve - .into_iter() - .chain(unsupported_ciphers) - .collect(), + unsupported: curve.err().into_iter().chain(ciphers.unsupported).collect(), } } +} - pub fn client_builder(&self) -> Result { +impl TryFrom<&HttpClientConfig> for reqwest::ClientBuilder { + type Error = Error; + + fn try_from(config: &HttpClientConfig) -> Result { let base = reqwest::Client::builder() - .use_preconfigured_tls(rustls::ClientConfig::try_from(self)?) - .connect_timeout(self.connect_timeout) - .pool_idle_timeout(self.pool_idle_timeout); - let with_keepalive = match self.tcp_keepalive { + .use_preconfigured_tls(rustls::ClientConfig::try_from(config)?) + .connect_timeout(config.connect_timeout) + .pool_idle_timeout(config.pool_idle_timeout); + let with_keepalive = match config.tcp_keepalive { None => base, Some(keepalive) => base .tcp_keepalive(keepalive.idle) .tcp_keepalive_interval(keepalive.interval) .tcp_keepalive_retries(keepalive.retries), }; - let with_address = if self.force_ipv4 { + let with_address = if config.force_ipv4 { with_keepalive.local_address(IpAddr::V4(Ipv4Addr::UNSPECIFIED)) } else { with_keepalive }; - let with_protocol = if self.http2 { + let with_protocol = if config.http2 { with_address } else { with_address.http1_only() }; - let with_agent = match &self.user_agent { + let with_agent = match &config.user_agent { Some(agent) => with_protocol.user_agent(agent), None => with_protocol, }; - Ok(if self.trust_proxy_env { + Ok(if config.trust_proxy_env { with_agent } else { with_agent.no_proxy() @@ -161,7 +157,7 @@ mod tests { #[case] settings: HttpSettings, #[case] expected: Verify, ) { - let config = HttpClientConfig::resolve(&settings).config; + let config = Resolution::from(&settings).config; assert_eq!(config.verify, expected); } @@ -172,7 +168,7 @@ mod tests { ..HttpSettings::default() } .with_environment(&|name: &str| (name == "SSL_VERIFY").then(|| "true".to_string())); - let config = HttpClientConfig::resolve(&settings).config; + let config = Resolution::from(&settings).config; assert_eq!(config.verify, Verify::BuiltInRoots); } @@ -188,7 +184,7 @@ mod tests { ssl_ecdh_curve: Some(curve.into()), ..HttpSettings::default() }; - let resolution = HttpClientConfig::resolve(&settings); + let resolution = Resolution::from(&settings); assert_eq!(resolution.config.key_exchange_group, expected); assert_eq!(resolution.unsupported, []); } @@ -199,7 +195,7 @@ mod tests { ssl_ecdh_curve: Some("secp521r1".into()), ..HttpSettings::default() }; - let resolution = HttpClientConfig::resolve(&settings); + let resolution = Resolution::from(&settings); assert_eq!(resolution.config.key_exchange_group, None); assert_eq!( resolution.unsupported, @@ -213,7 +209,7 @@ mod tests { ssl_security_level: Some("DEFAULT@SECLEVEL=1".into()), ..HttpSettings::default() }; - let resolution = HttpClientConfig::resolve(&settings); + let resolution = Resolution::from(&settings); assert_eq!(resolution.config.tls12_cipher_suites, None); assert_eq!( resolution.unsupported, @@ -230,7 +226,7 @@ mod tests { ), ..HttpSettings::default() }; - let resolution = HttpClientConfig::resolve(&settings); + let resolution = Resolution::from(&settings); assert_eq!( resolution.config.tls12_cipher_suites, Some(vec![ @@ -265,7 +261,7 @@ mod tests { pool_idle_timeout: Duration::from_secs(45), ..HttpSettings::default() }; - let config = HttpClientConfig::resolve(&settings).config; + let config = Resolution::from(&settings).config; assert_eq!( config, HttpClientConfig { @@ -303,7 +299,7 @@ mod tests { #[case] settings: HttpSettings, #[case] expected: bool, ) { - let config = HttpClientConfig::resolve(&settings).config; + let config = Resolution::from(&settings).config; assert_eq!(config.trust_proxy_env, expected); } @@ -312,10 +308,10 @@ mod tests { let path = std::env::temp_dir().join("litellm-http-missing-bundle.pem"); let config = HttpClientConfig { verify: Verify::CaBundle(path.clone()), - ..HttpClientConfig::resolve(&HttpSettings::default()).config + ..Resolution::from(&HttpSettings::default()).config }; assert!(matches!( - config.client_builder(), + reqwest::ClientBuilder::try_from(&config), Err(Error::Read { path: reported, .. }) if reported == path )); } @@ -327,9 +323,9 @@ mod tests { std::fs::write(&path, b"not a certificate").unwrap(); let config = HttpClientConfig { verify: Verify::CaBundle(path.clone()), - ..HttpClientConfig::resolve(&HttpSettings::default()).config + ..Resolution::from(&HttpSettings::default()).config }; - let result = config.client_builder().map(drop); + let result = reqwest::ClientBuilder::try_from(&config).map(drop); std::fs::remove_file(&path).unwrap(); assert!(matches!( result, diff --git a/litellm-rust/crates/http/src/pool.rs b/litellm-rust/crates/http/src/pool.rs index e6e0de9bc5f..330d6de29e8 100644 --- a/litellm-rust/crates/http/src/pool.rs +++ b/litellm-rust/crates/http/src/pool.rs @@ -67,7 +67,9 @@ impl HttpClientPool { { return Ok(pooled.client.clone()); } - let client = self.apply(variant, key.0.client_builder()?).build()?; + let client = self + .apply(variant, reqwest::ClientBuilder::try_from(&key.0)?) + .build()?; self.lock().insert( key, PooledClient { @@ -114,7 +116,7 @@ mod tests { }; use super::*; - use crate::{HttpSettings, Verify}; + use crate::{HttpSettings, Resolution, Verify}; struct FixedResolver(SocketAddr); @@ -132,7 +134,7 @@ mod tests { fn config(user_agent: &str) -> HttpClientConfig { HttpClientConfig { user_agent: Some(user_agent.into()), - ..HttpClientConfig::resolve(&HttpSettings::default()).config + ..Resolution::from(&HttpSettings::default()).config } } diff --git a/litellm-rust/crates/http/src/settings.rs b/litellm-rust/crates/http/src/settings.rs index be2f4f42fb4..6d7bf934e4b 100644 --- a/litellm-rust/crates/http/src/settings.rs +++ b/litellm-rust/crates/http/src/settings.rs @@ -111,6 +111,10 @@ impl HttpSettings { } } + pub fn trusts_proxy_env(&self) -> bool { + !self.ignore_proxy_env || self.trust_proxy_env || self.http2 || self.httpx_transport + } + pub fn without_missing_files(self, exists: &dyn Fn(&Path) -> bool) -> Self { Self { ssl_verify: match self.ssl_verify { diff --git a/litellm-rust/crates/http/src/tls.rs b/litellm-rust/crates/http/src/tls.rs index 49405b97366..aaae2b659e3 100644 --- a/litellm-rust/crates/http/src/tls.rs +++ b/litellm-rust/crates/http/src/tls.rs @@ -97,6 +97,7 @@ pub enum Unsupported { CipherToken(String), } +#[derive(Default)] pub(crate) struct CipherSelection { pub(crate) tls12_cipher_suites: Option>, pub(crate) unsupported: Vec, @@ -304,10 +305,10 @@ mod tests { use rustls::NamedGroup; use super::*; - use crate::HttpSettings; + use crate::{HttpSettings, Resolution}; fn config(settings: HttpSettings) -> HttpClientConfig { - HttpClientConfig::resolve(&settings).config + Resolution::from(&settings).config } fn offered_groups(tls: &ClientConfig) -> Vec { diff --git a/litellm-rust/crates/llms/src/custom_httpx/media.rs b/litellm-rust/crates/llms/src/custom_httpx/media.rs index 02d152d3ef1..572e7f12e54 100644 --- a/litellm-rust/crates/llms/src/custom_httpx/media.rs +++ b/litellm-rust/crates/llms/src/custom_httpx/media.rs @@ -350,7 +350,7 @@ impl Resolve for PublicDnsResolver { mod tests { use std::collections::HashSet; - use litellm_http::HttpSettings; + use litellm_http::{HttpSettings, Resolution}; use tokio::{ io::{AsyncReadExt, AsyncWriteExt}, net::TcpListener, @@ -445,7 +445,7 @@ mod tests { ) -> MediaFetcher { let direct = HttpClientConfig { trust_proxy_env: false, - ..HttpClientConfig::resolve(&HttpSettings::default()).config + ..Resolution::from(&HttpSettings::default()).config }; MediaFetcher::with_resolution( &HttpClientPool::new(Arc::new(LoopbackDnsResolver(pinned_address))), @@ -636,7 +636,7 @@ mod tests { async fn rejects_url_credentials_before_network_access() { let fetcher = MediaFetcher::new( &HttpClientPool::new(Arc::new(PublicDnsResolver)), - &HttpClientConfig::resolve(&HttpSettings::default()).config, + &Resolution::from(&HttpSettings::default()).config, UrlPolicy::default(), ) .expect("media fetcher builds"); diff --git a/litellm-rust/crates/python-bridge/src/http.rs b/litellm-rust/crates/python-bridge/src/http.rs index 118f4669b62..385f78be9fa 100644 --- a/litellm-rust/crates/python-bridge/src/http.rs +++ b/litellm-rust/crates/python-bridge/src/http.rs @@ -4,7 +4,9 @@ use std::{ sync::{Arc, LazyLock, Mutex, PoisonError}, }; -use litellm_http::{HttpClientConfig, HttpClientPool, HttpSettings, SslVerify, Unsupported}; +use litellm_http::{ + HttpClientConfig, HttpClientPool, HttpSettings, Resolution, SslVerify, Unsupported, +}; use litellm_llms::custom_httpx::media::{PublicDnsResolver, UrlPolicy}; use pyo3::{prelude::*, types::PyDict}; @@ -28,7 +30,7 @@ pub(crate) fn call_config( .with_environment(&|name| std::env::var(name).ok()); let settings = for_call(configured, call_ssl_verify(kwargs)?, asynchronous) .without_missing_files(&|path: &Path| path.exists()); - let resolution = HttpClientConfig::resolve(&settings); + let resolution = Resolution::from(&settings); for unsupported in unreported(&REPORTED_UNSUPPORTED, resolution.unsupported) { PythonSettings::warn(py, &unsupported.to_string())?; } @@ -251,7 +253,7 @@ user_agent='litellm/9.9.9', Python::initialize(); Python::attach(|py| { let settings = settings(&python_settings(py, overrides)).unwrap(); - let config = HttpClientConfig::resolve(&settings).config; + let config = Resolution::from(&settings).config; assert_eq!(config.verify, expected); }); } @@ -345,7 +347,7 @@ user_agent='litellm/9.9.9', ..HttpSettings::default() }; let settings = for_call(opted_out, None, asynchronous); - let config = HttpClientConfig::resolve(&settings).config; + let config = Resolution::from(&settings).config; assert_eq!(config.trust_proxy_env, expected); } } From bae4f22d3a201bf6c3d84c6017b44851dc5d07f0 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Fri, 18 Sep 2026 19:53:14 -0700 Subject: [PATCH 138/144] refactor(rust): merge http settings from per-source layers Each source (per-call kwargs, environment variables, the Python module) now builds an HttpSettingsLayer, and HttpSettings::from_layers merges them with explicit precedence. The aiohttp and httpx proxy-env rule is resolved once in the merge, so HttpSettings carries a single trust_proxy_env flag --- litellm-rust/crates/http/src/config.rs | 41 +-- litellm-rust/crates/http/src/lib.rs | 2 +- litellm-rust/crates/http/src/settings.rs | 286 +++++++++++++----- litellm-rust/crates/python-bridge/src/http.rs | 137 ++++----- 4 files changed, 278 insertions(+), 188 deletions(-) diff --git a/litellm-rust/crates/http/src/config.rs b/litellm-rust/crates/http/src/config.rs index a6cc08c210d..10f28b44eec 100644 --- a/litellm-rust/crates/http/src/config.rs +++ b/litellm-rust/crates/http/src/config.rs @@ -72,7 +72,7 @@ impl From<&HttpSettings> for Resolution { force_ipv4: settings.force_ipv4, http2: settings.http2, user_agent: settings.user_agent.clone(), - trust_proxy_env: settings.trusts_proxy_env(), + trust_proxy_env: settings.trust_proxy_env, connect_timeout: settings.connect_timeout, tcp_keepalive: settings.tcp_keepalive, pool_idle_timeout: settings.pool_idle_timeout, @@ -125,17 +125,12 @@ mod tests { use super::*; - fn no_env(_: &str) -> Option { - None - } - fn settings(ssl_verify: Option, ssl_cert_file: Option<&str>) -> HttpSettings { HttpSettings { ssl_verify, ssl_cert_file: ssl_cert_file.map(PathBuf::from), ..HttpSettings::default() } - .with_environment(&no_env) } #[rstest] @@ -161,17 +156,6 @@ mod tests { assert_eq!(config.verify, expected); } - #[test] - fn ssl_verify_environment_variable_beats_the_configured_setting() { - let settings = HttpSettings { - ssl_verify: Some(SslVerify::Disabled), - ..HttpSettings::default() - } - .with_environment(&|name: &str| (name == "SSL_VERIFY").then(|| "true".to_string())); - let config = Resolution::from(&settings).config; - assert_eq!(config.verify, Verify::BuiltInRoots); - } - #[rstest] #[case::x25519("X25519", Some(KeyExchangeGroup::X25519))] #[case::openssl_p256("prime256v1", Some(KeyExchangeGroup::Secp256r1))] @@ -280,29 +264,6 @@ mod tests { ); } - #[rstest] - #[case::aiohttp_default(HttpSettings::default(), true)] - #[case::aiohttp_opted_out(HttpSettings { ignore_proxy_env: true, ..HttpSettings::default() }, false)] - #[case::session_trust_env_beats_opt_out( - HttpSettings { ignore_proxy_env: true, trust_proxy_env: true, ..HttpSettings::default() }, - true - )] - #[case::http2_uses_httpx( - HttpSettings { ignore_proxy_env: true, http2: true, ..HttpSettings::default() }, - true - )] - #[case::aiohttp_disabled( - HttpSettings { ignore_proxy_env: true, httpx_transport: true, ..HttpSettings::default() }, - true - )] - fn environment_proxies_apply_unless_the_aiohttp_transport_opts_out( - #[case] settings: HttpSettings, - #[case] expected: bool, - ) { - let config = Resolution::from(&settings).config; - assert_eq!(config.trust_proxy_env, expected); - } - #[test] fn missing_ca_bundle_is_a_read_error() { let path = std::env::temp_dir().join("litellm-http-missing-bundle.pem"); diff --git a/litellm-rust/crates/http/src/lib.rs b/litellm-rust/crates/http/src/lib.rs index e222d0e3f50..ddbc3b63b08 100644 --- a/litellm-rust/crates/http/src/lib.rs +++ b/litellm-rust/crates/http/src/lib.rs @@ -9,5 +9,5 @@ pub use config::{HttpClientConfig, Resolution, Verify}; pub use error::Error; pub use pool::{ClientVariant, HttpClientPool}; pub use proxy::EnvironmentProxies; -pub use settings::{HttpSettings, SslVerify, TcpKeepalive}; +pub use settings::{HttpSettings, HttpSettingsLayer, SslVerify, TcpKeepalive}; pub use tls::{KeyExchangeGroup, Tls12CipherSuite, Unsupported}; diff --git a/litellm-rust/crates/http/src/settings.rs b/litellm-rust/crates/http/src/settings.rs index 6d7bf934e4b..8ac7ef92568 100644 --- a/litellm-rust/crates/http/src/settings.rs +++ b/litellm-rust/crates/http/src/settings.rs @@ -27,6 +27,79 @@ pub struct TcpKeepalive { pub retries: u32, } +#[derive(Clone, Debug, Default, PartialEq, Eq)] +pub struct HttpSettingsLayer { + pub ssl_verify: Option, + pub ssl_cert_file: Option, + pub ssl_certificate: Option, + pub ssl_security_level: Option, + pub ssl_ecdh_curve: Option, + pub force_ipv4: Option, + pub http2: Option, + pub aiohttp_trust_env: Option, + pub disable_aiohttp_trust_env: Option, + pub disable_aiohttp_transport: Option, + pub user_agent: Option, + pub tcp_keepalive: Option, + pub pool_idle_timeout: Option, +} + +impl HttpSettingsLayer { + pub fn from_environment(env: &(dyn Fn(&str) -> Option + Sync)) -> Self { + let enabled = |name: &str| { + env(name) + .is_some_and(|value| value.trim().eq_ignore_ascii_case("true")) + .then_some(true) + }; + let number = |name: &str| env(name).and_then(|value| value.trim().parse::().ok()); + let seconds = |name: &str, default: u32| { + Duration::from_secs(u64::from(number(name).unwrap_or(default))) + }; + Self { + ssl_verify: env("SSL_VERIFY").map(|value| SslVerify::parse(&value)), + ssl_cert_file: env("SSL_CERT_FILE").map(PathBuf::from), + ssl_certificate: env("SSL_CERTIFICATE").map(PathBuf::from), + ssl_security_level: env("SSL_SECURITY_LEVEL"), + ssl_ecdh_curve: env("SSL_ECDH_CURVE"), + force_ipv4: None, + http2: enabled("LITELLM_HTTP2"), + aiohttp_trust_env: enabled("AIOHTTP_TRUST_ENV"), + disable_aiohttp_trust_env: enabled("DISABLE_AIOHTTP_TRUST_ENV"), + disable_aiohttp_transport: enabled("DISABLE_AIOHTTP_TRANSPORT"), + user_agent: env("LITELLM_USER_AGENT"), + tcp_keepalive: enabled("AIOHTTP_SO_KEEPALIVE").map(|_| TcpKeepalive { + idle: seconds("AIOHTTP_TCP_KEEPIDLE", 60), + interval: seconds("AIOHTTP_TCP_KEEPINTVL", 30), + retries: number("AIOHTTP_TCP_KEEPCNT").unwrap_or(5), + }), + pool_idle_timeout: number("AIOHTTP_KEEPALIVE_TIMEOUT") + .map(|timeout| Duration::from_secs(u64::from(timeout))), + } + } + + fn or(self, lower: Self) -> Self { + Self { + ssl_verify: self.ssl_verify.or(lower.ssl_verify), + ssl_cert_file: self.ssl_cert_file.or(lower.ssl_cert_file), + ssl_certificate: self.ssl_certificate.or(lower.ssl_certificate), + ssl_security_level: self.ssl_security_level.or(lower.ssl_security_level), + ssl_ecdh_curve: self.ssl_ecdh_curve.or(lower.ssl_ecdh_curve), + force_ipv4: self.force_ipv4.or(lower.force_ipv4), + http2: self.http2.or(lower.http2), + aiohttp_trust_env: self.aiohttp_trust_env.or(lower.aiohttp_trust_env), + disable_aiohttp_trust_env: self + .disable_aiohttp_trust_env + .or(lower.disable_aiohttp_trust_env), + disable_aiohttp_transport: self + .disable_aiohttp_transport + .or(lower.disable_aiohttp_transport), + user_agent: self.user_agent.or(lower.user_agent), + tcp_keepalive: self.tcp_keepalive.or(lower.tcp_keepalive), + pool_idle_timeout: self.pool_idle_timeout.or(lower.pool_idle_timeout), + } + } +} + #[derive(Clone, Debug, PartialEq, Eq)] pub struct HttpSettings { pub ssl_verify: Option, @@ -36,10 +109,8 @@ pub struct HttpSettings { pub ssl_ecdh_curve: Option, pub force_ipv4: bool, pub http2: bool, - pub httpx_transport: bool, pub user_agent: Option, pub trust_proxy_env: bool, - pub ignore_proxy_env: bool, pub connect_timeout: Duration, pub tcp_keepalive: Option, pub pool_idle_timeout: Duration, @@ -55,10 +126,8 @@ impl Default for HttpSettings { ssl_ecdh_curve: None, force_ipv4: false, http2: false, - httpx_transport: false, user_agent: None, - trust_proxy_env: false, - ignore_proxy_env: false, + trust_proxy_env: true, connect_timeout: Duration::from_secs(10), tcp_keepalive: None, pool_idle_timeout: Duration::from_secs(120), @@ -67,54 +136,38 @@ impl Default for HttpSettings { } impl HttpSettings { - pub fn with_environment(self, env: &(dyn Fn(&str) -> Option + Sync)) -> Self { - let enabled = - |name: &str| env(name).is_some_and(|value| value.trim().eq_ignore_ascii_case("true")); - let number = |name: &str| env(name).and_then(|value| value.trim().parse::().ok()); - let seconds = |name: &str, default: u32| { - Duration::from_secs(u64::from(number(name).unwrap_or(default))) - }; + pub fn from_layers( + highest_precedence_first: impl IntoIterator, + ) -> Self { + let merged = highest_precedence_first + .into_iter() + .reduce(HttpSettingsLayer::or) + .unwrap_or_default(); + let defaults = Self::default(); + let http2 = merged.http2.unwrap_or(defaults.http2); Self { - ssl_verify: env("SSL_VERIFY") - .map(|value| SslVerify::parse(&value)) - .or(self.ssl_verify), - ssl_cert_file: env("SSL_CERT_FILE") - .map(PathBuf::from) - .or(self.ssl_cert_file), - ssl_certificate: env("SSL_CERTIFICATE") - .map(PathBuf::from) - .or(self.ssl_certificate) + ssl_verify: merged.ssl_verify, + ssl_cert_file: merged.ssl_cert_file, + ssl_certificate: merged + .ssl_certificate .filter(|path| !path.as_os_str().is_empty()), - ssl_security_level: env("SSL_SECURITY_LEVEL") - .or(self.ssl_security_level) - .filter(|level| !level.is_empty()), - ssl_ecdh_curve: env("SSL_ECDH_CURVE") - .or(self.ssl_ecdh_curve) - .filter(|curve| !curve.is_empty()), - http2: self.http2 || enabled("LITELLM_HTTP2"), - httpx_transport: self.httpx_transport || enabled("DISABLE_AIOHTTP_TRANSPORT"), - user_agent: env("LITELLM_USER_AGENT").or(self.user_agent), - trust_proxy_env: self.trust_proxy_env || enabled("AIOHTTP_TRUST_ENV"), - ignore_proxy_env: self.ignore_proxy_env || enabled("DISABLE_AIOHTTP_TRUST_ENV"), - tcp_keepalive: enabled("AIOHTTP_SO_KEEPALIVE") - .then(|| TcpKeepalive { - idle: seconds("AIOHTTP_TCP_KEEPIDLE", 60), - interval: seconds("AIOHTTP_TCP_KEEPINTVL", 30), - retries: number("AIOHTTP_TCP_KEEPCNT").unwrap_or(5), - }) - .or(self.tcp_keepalive), - pool_idle_timeout: number("AIOHTTP_KEEPALIVE_TIMEOUT") - .map_or(self.pool_idle_timeout, |timeout| { - Duration::from_secs(u64::from(timeout)) - }), - ..self + ssl_security_level: merged.ssl_security_level.filter(|level| !level.is_empty()), + ssl_ecdh_curve: merged.ssl_ecdh_curve.filter(|curve| !curve.is_empty()), + force_ipv4: merged.force_ipv4.unwrap_or(defaults.force_ipv4), + http2, + user_agent: merged.user_agent, + trust_proxy_env: !merged.disable_aiohttp_trust_env.unwrap_or(false) + || merged.aiohttp_trust_env.unwrap_or(false) + || merged.disable_aiohttp_transport.unwrap_or(false) + || http2, + tcp_keepalive: merged.tcp_keepalive, + pool_idle_timeout: merged + .pool_idle_timeout + .unwrap_or(defaults.pool_idle_timeout), + ..defaults } } - pub fn trusts_proxy_env(&self) -> bool { - !self.ignore_proxy_env || self.trust_proxy_env || self.http2 || self.httpx_transport - } - pub fn without_missing_files(self, exists: &dyn Fn(&Path) -> bool) -> Self { Self { ssl_verify: match self.ssl_verify { @@ -161,15 +214,15 @@ mod tests { } #[test] - fn environment_overrides_configured_ssl_values() { - let settings = HttpSettings { + fn higher_layers_override_lower_ones() { + let configured = HttpSettingsLayer { ssl_verify: Some(SslVerify::Enabled), ssl_certificate: Some("/configured/client.pem".into()), ssl_security_level: Some("configured".into()), user_agent: Some("configured/1".into()), - ..HttpSettings::default() - } - .with_environment(&env_of(&[ + ..HttpSettingsLayer::default() + }; + let environment = HttpSettingsLayer::from_environment(&env_of(&[ ("SSL_VERIFY", "false"), ("SSL_CERT_FILE", "/env/roots.pem"), ("SSL_CERTIFICATE", "/env/client.pem"), @@ -177,6 +230,7 @@ mod tests { ("SSL_ECDH_CURVE", "X25519"), ("LITELLM_USER_AGENT", "env/2"), ])); + let settings = HttpSettings::from_layers([environment, configured]); assert_eq!(settings.ssl_verify, Some(SslVerify::Disabled)); assert_eq!(settings.ssl_cert_file, Some("/env/roots.pem".into())); assert_eq!(settings.ssl_certificate, Some("/env/client.pem".into())); @@ -189,30 +243,60 @@ mod tests { } #[test] - fn missing_environment_keeps_configured_values() { - let configured = HttpSettings { - ssl_verify: Some(SslVerify::CaBundle("/configured/roots.pem".into())), - http2: true, - trust_proxy_env: true, - user_agent: Some("configured/1".into()), - ..HttpSettings::default() + fn an_explicit_false_in_a_higher_layer_beats_a_lower_true() { + let higher = HttpSettingsLayer { + http2: Some(false), + force_ipv4: Some(false), + ..HttpSettingsLayer::default() }; - assert_eq!(configured.clone().with_environment(&no_env), configured); + let lower = HttpSettingsLayer { + http2: Some(true), + force_ipv4: Some(true), + ..HttpSettingsLayer::default() + }; + let settings = HttpSettings::from_layers([higher, lower]); + assert!(!settings.http2); + assert!(!settings.force_ipv4); + } + + #[test] + fn an_empty_environment_is_an_empty_layer_so_lower_layers_and_defaults_apply() { + assert_eq!( + HttpSettingsLayer::from_environment(&no_env), + HttpSettingsLayer::default() + ); + let configured = HttpSettingsLayer { + ssl_verify: Some(SslVerify::CaBundle("/configured/roots.pem".into())), + http2: Some(true), + user_agent: Some("configured/1".into()), + ..HttpSettingsLayer::default() + }; + assert_eq!( + HttpSettings::from_layers([HttpSettingsLayer::default(), configured]), + HttpSettings { + ssl_verify: Some(SslVerify::CaBundle("/configured/roots.pem".into())), + http2: true, + user_agent: Some("configured/1".into()), + ..HttpSettings::default() + } + ); + assert_eq!(HttpSettings::from_layers([]), HttpSettings::default()); } #[test] fn empty_environment_values_clear_the_setting_like_python_truthiness() { - let settings = HttpSettings { + let configured = HttpSettingsLayer { ssl_certificate: Some("/configured/client.pem".into()), ssl_security_level: Some("configured".into()), ssl_ecdh_curve: Some("X25519".into()), - ..HttpSettings::default() - } - .with_environment(&env_of(&[ + ..HttpSettingsLayer::default() + }; + let environment = HttpSettingsLayer::from_environment(&env_of(&[ ("SSL_CERTIFICATE", ""), ("SSL_SECURITY_LEVEL", ""), ("SSL_ECDH_CURVE", ""), ])); + let settings = HttpSettings::from_layers([environment, configured]); assert_eq!(settings.ssl_certificate, None); assert_eq!(settings.ssl_security_level, None); assert_eq!(settings.ssl_ecdh_curve, None); @@ -220,11 +304,11 @@ mod tests { #[test] fn socket_keepalive_follows_the_aiohttp_variables_with_python_defaults() { - let tuned = HttpSettings::default().with_environment(&env_of(&[ + let tuned = HttpSettings::from_layers([HttpSettingsLayer::from_environment(&env_of(&[ ("AIOHTTP_SO_KEEPALIVE", "True"), ("AIOHTTP_TCP_KEEPIDLE", "45"), ("AIOHTTP_KEEPALIVE_TIMEOUT", "30"), - ])); + ]))]); assert_eq!( tuned.tcp_keepalive, Some(TcpKeepalive { @@ -238,12 +322,53 @@ mod tests { #[test] fn socket_keepalive_stays_off_unless_enabled() { - let settings = - HttpSettings::default().with_environment(&env_of(&[("AIOHTTP_TCP_KEEPIDLE", "45")])); + let settings = HttpSettings::from_layers([HttpSettingsLayer::from_environment(&env_of( + &[("AIOHTTP_TCP_KEEPIDLE", "45")], + ))]); assert_eq!(settings.tcp_keepalive, None); assert_eq!(settings.pool_idle_timeout, Duration::from_secs(120)); } + fn proxy_flags( + aiohttp_trust_env: bool, + disable_aiohttp_trust_env: bool, + disable_aiohttp_transport: bool, + http2: bool, + ) -> HttpSettingsLayer { + HttpSettingsLayer { + aiohttp_trust_env: Some(aiohttp_trust_env), + disable_aiohttp_trust_env: Some(disable_aiohttp_trust_env), + disable_aiohttp_transport: Some(disable_aiohttp_transport), + http2: Some(http2), + ..HttpSettingsLayer::default() + } + } + + #[rstest] + #[case::aiohttp_default(proxy_flags(false, false, false, false), true)] + #[case::aiohttp_opted_out(proxy_flags(false, true, false, false), false)] + #[case::session_trust_env_beats_opt_out(proxy_flags(true, true, false, false), true)] + #[case::http2_uses_httpx(proxy_flags(false, true, false, true), true)] + #[case::aiohttp_disabled(proxy_flags(false, true, true, false), true)] + fn environment_proxies_apply_unless_the_aiohttp_transport_opts_out( + #[case] layer: HttpSettingsLayer, + #[case] expected: bool, + ) { + assert_eq!(HttpSettings::from_layers([layer]).trust_proxy_env, expected); + } + + #[test] + fn a_proxy_opt_out_in_one_source_still_yields_to_trust_env_from_another() { + let environment = + HttpSettingsLayer::from_environment(&env_of(&[("DISABLE_AIOHTTP_TRUST_ENV", "true")])); + let configured = HttpSettingsLayer { + aiohttp_trust_env: Some(true), + ..HttpSettingsLayer::default() + }; + assert!(!HttpSettings::from_layers([environment.clone()]).trust_proxy_env); + assert!(HttpSettings::from_layers([environment, configured]).trust_proxy_env); + } + #[test] fn missing_files_fall_back_to_default_verification() { let settings = HttpSettings { @@ -267,11 +392,14 @@ mod tests { } #[rstest] - #[case("true", true)] - #[case("True", true)] - #[case("false", false)] - #[case("1", false)] - fn boolean_switches_only_turn_on_for_true(#[case] value: &'static str, #[case] expected: bool) { + #[case("true", Some(true))] + #[case("True", Some(true))] + #[case("false", None)] + #[case("1", None)] + fn boolean_switches_only_turn_on_for_true( + #[case] value: &'static str, + #[case] expected: Option, + ) { let env = move |name: &str| match name { "LITELLM_HTTP2" | "AIOHTTP_TRUST_ENV" @@ -279,10 +407,10 @@ mod tests { | "DISABLE_AIOHTTP_TRUST_ENV" => Some(value.to_string()), _ => None, }; - let settings = HttpSettings::default().with_environment(&env); - assert_eq!(settings.http2, expected); - assert_eq!(settings.httpx_transport, expected); - assert_eq!(settings.trust_proxy_env, expected); - assert_eq!(settings.ignore_proxy_env, expected); + let layer = HttpSettingsLayer::from_environment(&env); + assert_eq!(layer.http2, expected); + assert_eq!(layer.aiohttp_trust_env, expected); + assert_eq!(layer.disable_aiohttp_transport, expected); + assert_eq!(layer.disable_aiohttp_trust_env, expected); } } diff --git a/litellm-rust/crates/python-bridge/src/http.rs b/litellm-rust/crates/python-bridge/src/http.rs index 385f78be9fa..d174dccaa56 100644 --- a/litellm-rust/crates/python-bridge/src/http.rs +++ b/litellm-rust/crates/python-bridge/src/http.rs @@ -5,7 +5,8 @@ use std::{ }; use litellm_http::{ - HttpClientConfig, HttpClientPool, HttpSettings, Resolution, SslVerify, Unsupported, + HttpClientConfig, HttpClientPool, HttpSettings, HttpSettingsLayer, Resolution, SslVerify, + Unsupported, }; use litellm_llms::custom_httpx::media::{PublicDnsResolver, UrlPolicy}; use pyo3::{prelude::*, types::PyDict}; @@ -26,10 +27,12 @@ pub(crate) fn call_config( kwargs: &Bound<'_, PyDict>, asynchronous: bool, ) -> PyResult { - let configured = settings(&PythonSettings::Http.read(py)?)? - .with_environment(&|name| std::env::var(name).ok()); - let settings = for_call(configured, call_ssl_verify(kwargs)?, asynchronous) - .without_missing_files(&|path: &Path| path.exists()); + let settings = HttpSettings::from_layers([ + for_call(call_ssl_verify(kwargs)?, asynchronous), + HttpSettingsLayer::from_environment(&|name| std::env::var(name).ok()), + configured(&PythonSettings::Http.read(py)?)?, + ]) + .without_missing_files(&|path: &Path| path.exists()); let resolution = Resolution::from(&settings); for unsupported in unreported(&REPORTED_UNSUPPORTED, resolution.unsupported) { PythonSettings::warn(py, &unsupported.to_string())?; @@ -70,15 +73,11 @@ fn call_ssl_verify(kwargs: &Bound<'_, PyDict>) -> PyResult> { .and_then(|value| ssl_verify(&value))) } -fn for_call( - configured: HttpSettings, - call_ssl_verify: Option, - asynchronous: bool, -) -> HttpSettings { - HttpSettings { - ssl_verify: call_ssl_verify.or(configured.ssl_verify), - httpx_transport: configured.httpx_transport || !asynchronous, - ..configured +fn for_call(call_ssl_verify: Option, asynchronous: bool) -> HttpSettingsLayer { + HttpSettingsLayer { + ssl_verify: call_ssl_verify, + disable_aiohttp_transport: (!asynchronous).then_some(true), + ..HttpSettingsLayer::default() } } @@ -102,24 +101,24 @@ struct PythonHttpSettings<'py> { user_agent: String, } -fn settings(value: &Bound<'_, PyAny>) -> PyResult { +fn configured(value: &Bound<'_, PyAny>) -> PyResult { let python: PythonHttpSettings = value.extract().map_err(|error: PyErr| { RustBridgeDeclined::new_err(format!( "litellm HTTP settings cannot be used by the Rust route: {error}" )) })?; - Ok(HttpSettings { + Ok(HttpSettingsLayer { ssl_verify: ssl_verify(&python.ssl_verify), ssl_certificate: python.ssl_certificate.map(PathBuf::from), ssl_security_level: python.ssl_security_level, ssl_ecdh_curve: python.ssl_ecdh_curve, - force_ipv4: python.force_ipv4, - http2: python.http2, - httpx_transport: python.disable_aiohttp_transport, + force_ipv4: Some(python.force_ipv4), + http2: Some(python.http2), + aiohttp_trust_env: Some(python.aiohttp_trust_env), + disable_aiohttp_trust_env: Some(python.disable_aiohttp_trust_env), + disable_aiohttp_transport: Some(python.disable_aiohttp_transport), user_agent: Some(python.user_agent), - trust_proxy_env: python.aiohttp_trust_env, - ignore_proxy_env: python.disable_aiohttp_trust_env, - ..HttpSettings::default() + ..HttpSettingsLayer::default() }) } @@ -174,12 +173,12 @@ settings = types.SimpleNamespace(**{{name: defaults[name] for name in json.loads } #[test] - fn default_python_settings_produce_default_settings_with_verification_on() { + fn default_python_settings_resolve_to_default_settings_with_verification_on() { Python::initialize(); Python::attach(|py| { - let settings = settings(&python_settings(py, "")).unwrap(); + let layer = configured(&python_settings(py, "")).unwrap(); assert_eq!( - settings, + HttpSettings::from_layers([layer]), HttpSettings { ssl_verify: Some(SslVerify::Enabled), user_agent: Some("litellm/test".into()), @@ -190,10 +189,10 @@ settings = types.SimpleNamespace(**{{name: defaults[name] for name in json.loads } #[test] - fn python_settings_flow_into_settings() { + fn python_settings_flow_into_the_configured_layer() { Python::initialize(); Python::attach(|py| { - let settings = settings(&python_settings( + let layer = configured(&python_settings( py, " ssl_verify='/etc/ssl/corp.pem', @@ -210,19 +209,19 @@ user_agent='litellm/9.9.9', )) .unwrap(); assert_eq!( - settings, - HttpSettings { + layer, + HttpSettingsLayer { ssl_verify: Some(SslVerify::CaBundle("/etc/ssl/corp.pem".into())), ssl_certificate: Some("/etc/ssl/client.pem".into()), ssl_security_level: Some("2".into()), ssl_ecdh_curve: Some("X25519".into()), - force_ipv4: true, - http2: true, - httpx_transport: true, + force_ipv4: Some(true), + http2: Some(true), + aiohttp_trust_env: Some(true), + disable_aiohttp_trust_env: Some(true), + disable_aiohttp_transport: Some(true), user_agent: Some("litellm/9.9.9".into()), - trust_proxy_env: true, - ignore_proxy_env: true, - ..HttpSettings::default() + ..HttpSettingsLayer::default() } ); }); @@ -232,11 +231,12 @@ user_agent='litellm/9.9.9', fn user_agent_environment_variable_beats_the_python_default() { Python::initialize(); Python::attach(|py| { - let settings = settings(&python_settings(py, "")) - .unwrap() - .with_environment(&|name| { + let settings = HttpSettings::from_layers([ + HttpSettingsLayer::from_environment(&|name| { (name == "LITELLM_USER_AGENT").then(|| "operator/1".to_string()) - }); + }), + configured(&python_settings(py, "")).unwrap(), + ]); assert_eq!(settings.user_agent.as_deref(), Some("operator/1")); }); } @@ -252,8 +252,8 @@ user_agent='litellm/9.9.9', ) { Python::initialize(); Python::attach(|py| { - let settings = settings(&python_settings(py, overrides)).unwrap(); - let config = Resolution::from(&settings).config; + let layer = configured(&python_settings(py, overrides)).unwrap(); + let config = Resolution::from(&HttpSettings::from_layers([layer])).config; assert_eq!(config.verify, expected); }); } @@ -262,8 +262,8 @@ user_agent='litellm/9.9.9', fn ssl_context_global_is_ignored_so_environment_and_defaults_apply() { Python::initialize(); Python::attach(|py| { - let settings = settings(&python_settings(py, "ssl_verify=object()")).unwrap(); - assert_eq!(settings.ssl_verify, None); + let layer = configured(&python_settings(py, "ssl_verify=object()")).unwrap(); + assert_eq!(layer.ssl_verify, None); }); } @@ -283,22 +283,27 @@ user_agent='litellm/9.9.9', fn mistyped_python_settings_decline_instead_of_raising() { Python::initialize(); Python::attach(|py| { - let error = settings(&python_settings(py, "force_ipv4='yes'")).unwrap_err(); + let error = configured(&python_settings(py, "force_ipv4='yes'")).unwrap_err(); assert!(error.is_instance_of::(py)); }); } + fn configured_ssl_verify(ssl_verify: SslVerify) -> HttpSettingsLayer { + HttpSettingsLayer { + ssl_verify: Some(ssl_verify), + ..HttpSettingsLayer::default() + } + } + #[test] - fn call_ssl_verify_beats_the_configured_and_environment_value() { + fn call_ssl_verify_beats_the_configured_value() { Python::initialize(); Python::attach(|py| { let kwargs = PyDict::new(py); kwargs.set_item("ssl_verify", false).unwrap(); - let configured = HttpSettings { - ssl_verify: Some(SslVerify::Enabled), - ..HttpSettings::default() - }; - let settings = for_call(configured, call_ssl_verify(&kwargs).unwrap(), true); + let call = for_call(call_ssl_verify(&kwargs).unwrap(), true); + let settings = + HttpSettings::from_layers([call, configured_ssl_verify(SslVerify::Enabled)]); assert_eq!(settings.ssl_verify, Some(SslVerify::Disabled)); }); } @@ -309,12 +314,10 @@ user_agent='litellm/9.9.9', Python::attach(|py| { let kwargs = PyDict::new(py); kwargs.set_item("ssl_verify", py.None()).unwrap(); - let configured = HttpSettings { - ssl_verify: Some(SslVerify::Disabled), - ..HttpSettings::default() - }; - let settings = for_call(configured.clone(), call_ssl_verify(&kwargs).unwrap(), true); - assert_eq!(settings, configured); + let call = for_call(call_ssl_verify(&kwargs).unwrap(), true); + let settings = + HttpSettings::from_layers([call, configured_ssl_verify(SslVerify::Disabled)]); + assert_eq!(settings.ssl_verify, Some(SslVerify::Disabled)); }); } @@ -326,12 +329,10 @@ user_agent='litellm/9.9.9', kwargs .set_item("ssl_verify", py.eval(c"object()", None, None).unwrap()) .unwrap(); - let configured = HttpSettings { - ssl_verify: Some(SslVerify::Disabled), - ..HttpSettings::default() - }; - let settings = for_call(configured.clone(), call_ssl_verify(&kwargs).unwrap(), true); - assert_eq!(settings, configured); + let call = for_call(call_ssl_verify(&kwargs).unwrap(), true); + let settings = + HttpSettings::from_layers([call, configured_ssl_verify(SslVerify::Disabled)]); + assert_eq!(settings.ssl_verify, Some(SslVerify::Disabled)); }); } @@ -342,12 +343,12 @@ user_agent='litellm/9.9.9', #[case] asynchronous: bool, #[case] expected: bool, ) { - let opted_out = HttpSettings { - ignore_proxy_env: true, - ..HttpSettings::default() + let opted_out = HttpSettingsLayer { + disable_aiohttp_trust_env: Some(true), + disable_aiohttp_transport: Some(false), + ..HttpSettingsLayer::default() }; - let settings = for_call(opted_out, None, asynchronous); - let config = Resolution::from(&settings).config; - assert_eq!(config.trust_proxy_env, expected); + let settings = HttpSettings::from_layers([for_call(None, asynchronous), opted_out]); + assert_eq!(settings.trust_proxy_env, expected); } } From 92ea8adb3bff060911a65a4b4e811c7419863257 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Sat, 19 Sep 2026 03:18:23 +0000 Subject: [PATCH 139/144] docs: replace stale Black formatting instructions with ruff format Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- CONTRIBUTING.md | 17 ++++++++--------- 1 file changed, 8 insertions(+), 9 deletions(-) diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 153ca040e27..82cad680a70 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -148,7 +148,7 @@ make lint Individual linting commands: ```bash -make format-check # Check Black formatting +make format-check # Check ruff format formatting make lint-ruff # Run Ruff linting make lint-basedpyright # Run basedpyright type checking make check-circular-imports # Check for circular imports @@ -160,14 +160,14 @@ Apply formatting (auto-fixes issues): make format ``` -> **Black formatting is enforced in CI.** All PRs must pass the Black formatting check. +> **Formatting is enforced in CI.** All PRs must pass the `ruff format --check` step. > -> - **AI coding agents** (Claude Code, Copilot, Cursor, etc.): `AGENTS.md` instructs agents to run `poetry run black .` before committing. -> - **VS Code users**: Install the [Black Formatter extension](https://marketplace.visualstudio.com/items?itemName=ms-python.black-formatter) and enable format-on-save: +> - **AI coding agents** (Claude Code, Copilot, Cursor, etc.): follow `AGENTS.md` and run `make format` before committing. +> - **VS Code users**: Install the [Ruff extension](https://marketplace.visualstudio.com/items?itemName=charliermarsh.ruff) and enable format-on-save: > ```json > { > "[python]": { -> "editor.defaultFormatter": "ms-python.black-formatter", +> "editor.defaultFormatter": "charliermarsh.ruff", > "editor.formatOnSave": true > } > } @@ -197,8 +197,8 @@ make help # Show all available commands make install-dev # Install development dependencies make install-proxy-dev # Install proxy development dependencies make install-test-deps # Install the full local test environment -make format # Apply Black code formatting -make format-check # Check Black formatting (matches CI) +make format # Apply ruff format code formatting +make format-check # Check ruff format formatting (matches CI) make lint # Run all linting checks make test-unit # Run unit tests make test-integration # Run integration tests @@ -210,8 +210,7 @@ make test-unit-helm # Run Helm unit tests LiteLLM follows the [Google Python Style Guide](https://google.github.io/styleguide/pyguide.html). Our automated quality checks include: -- **Black** for consistent code formatting -- **Ruff** for linting and code quality +- **Ruff** for formatting, linting, and code quality - **basedpyright** for static type checking - **Circular import detection** - **Import safety validation** From 3157a8a3ca142ff18cdf40d65a302522dce1aa9e Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Sat, 19 Sep 2026 03:26:54 +0000 Subject: [PATCH 140/144] docs: replace poetry run with uv run in script instructions Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- scripts/test_tool_allowlist_script.py | 6 +++--- tests/test_litellm/test_utils.py | 2 +- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/scripts/test_tool_allowlist_script.py b/scripts/test_tool_allowlist_script.py index 9503a21219c..f94aac60f80 100644 --- a/scripts/test_tool_allowlist_script.py +++ b/scripts/test_tool_allowlist_script.py @@ -3,10 +3,10 @@ Standalone script to test tool allowlist enforcement and tool name extraction. Run from repo root: - poetry run python scripts/test_tool_allowlist_script.py + uv run python scripts/test_tool_allowlist_script.py Or run the unit tests: - poetry run pytest tests/test_litellm/proxy/test_tools_allowlist_enforcement.py -v + uv run pytest tests/test_litellm/proxy/test_tools_allowlist_enforcement.py -v """ import asyncio @@ -148,7 +148,7 @@ def main(): asyncio.run(test_check_tools_allowlist()) print("Done. For full unit tests run:") print( - " poetry run pytest tests/test_litellm/proxy/test_tools_allowlist_enforcement.py -v" + " uv run pytest tests/test_litellm/proxy/test_tools_allowlist_enforcement.py -v" ) diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index 5adfb2aea4c..2fda5dfc490 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -1060,7 +1060,7 @@ def test_max_tokens_consistency(): if len(inconsistencies) > 10: error_msg += f"\n ... and {len(inconsistencies) - 10} more\n" - error_msg += "\nTo fix these inconsistencies, run: poetry run python fix_max_tokens_inconsistencies.py" + error_msg += "\nTo fix these inconsistencies, run: uv run python fix_max_tokens_inconsistencies.py" raise AssertionError(error_msg) From d8d0e343e1111a5ba523d1ae5343967b52756665 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Sat, 19 Sep 2026 03:28:18 +0000 Subject: [PATCH 141/144] docs: drop stale Black, MyPy, and isort mentions from README and pyproject Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- README.md | 5 ++--- pyproject.toml | 3 --- 2 files changed, 2 insertions(+), 6 deletions(-) diff --git a/README.md b/README.md index 901cc5b0cea..3f3ea0bd60b 100644 --- a/README.md +++ b/README.md @@ -633,9 +633,8 @@ For detailed contributing guidelines, see [CONTRIBUTING.md](CONTRIBUTING.md). LiteLLM follows the [Google Python Style Guide](https://google.github.io/styleguide/pyguide.html). Our automated checks include: -- **Black** for code formatting -- **Ruff** for linting and code quality -- **MyPy** for type checking +- **Ruff** for formatting, linting, and code quality +- **basedpyright** for type checking - **Circular import detection** - **Import safety checks** diff --git a/pyproject.toml b/pyproject.toml index dfe84a28d52..11eae05c213 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -329,9 +329,6 @@ litellm-enterprise = { workspace = true } [tool.uv.workspace] members = ["enterprise", "litellm-proxy-extras"] -[tool.isort] -profile = "black" - [tool.commitizen] version = "1.103.0" version_files = [ From 33223920caf9b45e1546fac4f876c023647007a0 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 18 Sep 2026 20:57:56 -0700 Subject: [PATCH 142/144] refactor(responses): build the routed websocket request and relay frames without in-place mutation --- .../proxy/response_api_endpoints/endpoints.py | 15 ++-- litellm/responses/streaming_iterator.py | 82 ++++++++++--------- 2 files changed, 50 insertions(+), 47 deletions(-) diff --git a/litellm/proxy/response_api_endpoints/endpoints.py b/litellm/proxy/response_api_endpoints/endpoints.py index b3d6a928a78..b07458bb5ed 100644 --- a/litellm/proxy/response_api_endpoints/endpoints.py +++ b/litellm/proxy/response_api_endpoints/endpoints.py @@ -2,7 +2,7 @@ import asyncio import contextlib import json import time -from collections.abc import AsyncIterator, Awaitable, Mapping +from collections.abc import AsyncIterator, Awaitable, Mapping, Sequence from enum import Enum from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, NamedTuple, Protocol, cast, get_args @@ -1376,7 +1376,7 @@ def _extract_model_from_first_ws_event(first_event: Any) -> str | None: class _ResponseCreateRoutingHints(BaseModel): model_config = ConfigDict(extra="ignore", frozen=True) - input: str | list[object] | None = None + input: str | Sequence[object] | None = None previous_response_id: str | None = None response: "_ResponseCreateRoutingHints | None" = None @@ -1567,12 +1567,13 @@ async def responses_websocket_endpoint( await websocket.close(code=1008, reason="Pre-call error") return + routed_data: Final = dict( + data, user_api_key_dict=user_api_key_dict, **_routing_hints_from_first_ws_frame(first_message) + ) # Phase 2: route to upstream provider try: - data["user_api_key_dict"] = user_api_key_dict - data.update(_routing_hints_from_first_ws_frame(first_message)) llm_call: Final = await route_request( - data=data, + data=routed_data, route_type="_aresponses_websocket", llm_router=llm_router, user_model=user_model, @@ -1582,7 +1583,7 @@ async def responses_websocket_endpoint( await proxy_logging_obj.post_call_failure_hook( user_api_key_dict=user_api_key_dict, original_exception=failure, - request_data=data, + request_data=routed_data, ) except Exception as e: verbose_proxy_logger.exception("Responses WebSocket error") @@ -1591,6 +1592,6 @@ async def responses_websocket_endpoint( await proxy_logging_obj.post_call_failure_hook( user_api_key_dict=user_api_key_dict, original_exception=e, - request_data=data, + request_data=routed_data, ) await websocket.close(code=1011, reason="Internal server error") diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 32a36ffe4e8..122476a1be3 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -1,6 +1,7 @@ from __future__ import annotations import asyncio +import copy import json import time import traceback @@ -154,7 +155,7 @@ def _load_json_value(payload: str | bytes) -> object: return json.loads(payload) -def _model_id_from_metadata(litellm_metadata: dict[str, object] | None) -> str | None: +def _model_id_from_metadata(litellm_metadata: Mapping[str, object] | None) -> str | None: model_info: Final = litellm_metadata.get("model_info") if litellm_metadata else None model_id: Final = model_info.get("id") if _is_json_object(model_info) else None return model_id if isinstance(model_id, str) else None @@ -1701,59 +1702,59 @@ _RESPONSES_WS_FAILURE_EVENT_TYPES: Final = frozenset({"error", "response.failed" _RESPONSES_WS_OUTPUT_ITEM_EVENT_TYPES: Final = frozenset({"response.output_item.added", "response.output_item.done"}) -def _ws_event_error(event: _MutableJsonObject) -> object: +def _ws_event_error(event: Mapping[str, object]) -> object: if event.get("type") == "error": return event.get("error") response: Final = event.get("response") return response.get("error") if _is_json_object(response) else None -def _item_id_fields(item: object) -> tuple[object, object]: - return (item.get("id"), item.get("encrypted_content")) if _is_json_object(item) else (None, None) +def _restore_input_item_ids(items: Sequence[object]) -> Sequence[object]: + return ResponsesAPIRequestUtils._restore_encrypted_content_item_ids_in_input(copy.deepcopy(list(items))) # pyright: ignore[reportPrivateUsage] # same restore the HTTP responses path runs -def _restore_input_item_ids(items: list[object]) -> bool: - before: Final = tuple(_item_id_fields(item) for item in items) - ResponsesAPIRequestUtils._restore_encrypted_content_item_ids_in_input(items) # pyright: ignore[reportPrivateUsage] # same restore the HTTP responses path runs - return before != tuple(_item_id_fields(item) for item in items) - - -def _restore_wrapped_ids_in_container(container: _MutableJsonObject) -> bool: +def _restored_container_fields(container: Mapping[str, object]) -> Mapping[str, object]: input_items: Final = container.get("input") - input_restored: Final = _is_json_array(input_items) and _restore_input_item_ids(input_items) previous_response_id: Final = container.get("previous_response_id") - if not isinstance(previous_response_id, str): - return input_restored - original_previous_response_id: Final = ( - ResponsesAPIRequestUtils.decode_previous_response_id_to_original_previous_response_id(previous_response_id) - ) - if original_previous_response_id == previous_response_id: - return input_restored - container["previous_response_id"] = original_previous_response_id - return True + restored: Final = { + "input": _restore_input_item_ids(input_items) if _is_json_array(input_items) else input_items, + "previous_response_id": ( + ResponsesAPIRequestUtils.decode_previous_response_id_to_original_previous_response_id(previous_response_id) + if isinstance(previous_response_id, str) + else previous_response_id + ), + } + return MappingProxyType({key: value for key, value in restored.items() if value != container.get(key)}) -def _restore_wrapped_ids_in_response_create(msg_obj: _MutableJsonObject) -> bool: +def _restore_wrapped_ids_in_response_create(msg_obj: Mapping[str, object]) -> dict[str, object] | None: nested: Final = msg_obj.get("response") - containers: Final = (msg_obj, nested) if _is_json_object(nested) else (msg_obj,) - restored: Final = tuple(_restore_wrapped_ids_in_container(container) for container in containers) - return any(restored) + nested_fields: Final = _restored_container_fields(nested) if _is_json_object(nested) else EMPTY_MAPPING + top_fields: Final = _restored_container_fields(msg_obj) + if not nested_fields and not top_fields: + return None + restored_nested: Final = ( + {"response": {**nested, **nested_fields}} if _is_json_object(nested) and nested_fields else EMPTY_MAPPING + ) + return {**msg_obj, **top_fields, **restored_nested} -def _wrap_output_item_encrypted_content(event_obj: _MutableJsonObject, litellm_metadata: dict[str, object]) -> bool: +def _wrap_output_item_encrypted_content( + event_obj: Mapping[str, object], litellm_metadata: Mapping[str, object] +) -> dict[str, object] | None: if not litellm_metadata.get("encrypted_content_affinity_enabled"): - return False + return None model_id: Final = _model_id_from_metadata(litellm_metadata) item: Final = event_obj.get("item") if model_id is None or not _is_json_object(item): - return False + return None encrypted_content: Final = item.get("encrypted_content") if not isinstance(encrypted_content, str) or not encrypted_content: - return False - item["encrypted_content"] = ResponsesAPIRequestUtils._wrap_encrypted_content_with_model_id( # pyright: ignore[reportPrivateUsage] # same wrap the HTTP streaming path applies + return None + wrapped_content: Final = ResponsesAPIRequestUtils._wrap_encrypted_content_with_model_id( # pyright: ignore[reportPrivateUsage] # same wrap the HTTP streaming path applies encrypted_content=encrypted_content, model_id=model_id ) - return True + return {**event_obj, "item": {**item, "encrypted_content": wrapped_content}} class ResponsesWebSocketStreaming: @@ -1909,16 +1910,16 @@ class ResponsesWebSocketStreaming: return response_str response: Final = event_obj.get("response") if _is_json_object(response): - event_obj["response"] = ResponsesAPIRequestUtils._update_responses_api_response_id_with_model_id( # pyright: ignore[reportPrivateUsage] # same wrap the HTTP streaming path applies + wrapped_response: Final = ResponsesAPIRequestUtils._update_responses_api_response_id_with_model_id( # pyright: ignore[reportPrivateUsage] # same wrap the HTTP streaming path applies responses_api_response=response, custom_llm_provider=self.custom_llm_provider, litellm_metadata=self.litellm_metadata, ) - return json.dumps(event_obj) + return json.dumps({**event_obj, "response": wrapped_response}) if event_obj.get("type") not in _RESPONSES_WS_OUTPUT_ITEM_EVENT_TYPES: return response_str - item_wrapped: Final = _wrap_output_item_encrypted_content(event_obj, self.litellm_metadata) - return json.dumps(event_obj) if item_wrapped else response_str + wrapped_event: Final = _wrap_output_item_encrypted_content(event_obj, self.litellm_metadata) + return response_str if wrapped_event is None else json.dumps(wrapped_event) async def backend_to_client(self) -> None: """Forward events from backend WebSocket to the client.""" @@ -2030,13 +2031,14 @@ class ResponsesWebSocketStreaming: if parsed.get("type") != "response.create": return message - msg_obj: Final = self._with_request_defaults(parsed) - defaults_applied: Final = msg_obj != parsed + authorized_obj: Final = self._with_request_defaults(parsed) + defaults_applied: Final = authorized_obj != parsed # Always enforce the authorized model, even when PII masking is off. - model_modified: Final = self._enforce_authorized_model(msg_obj) - ids_restored: Final = _restore_wrapped_ids_in_response_create(msg_obj) - frame_modified: Final = model_modified or ids_restored or defaults_applied + model_modified: Final = self._enforce_authorized_model(authorized_obj) + restored_obj: Final = _restore_wrapped_ids_in_response_create(authorized_obj) + msg_obj: Final = authorized_obj if restored_obj is None else restored_obj + frame_modified: Final = model_modified or restored_obj is not None or defaults_applied if not self.guardrail_callbacks: return json.dumps(msg_obj) if frame_modified else message From ef34e44d8ba4edda04904903630099b67a01bbfb Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Fri, 18 Sep 2026 21:00:18 -0700 Subject: [PATCH 143/144] fix(proxy): parse role_permissions where it is read load_config used to return a local general_settings dict that it had normalized in place, turning the configured role_permissions entries into RoleBasedPermissions objects. It now returns the SettingsStore, which never saw that write, so JWT auth received raw dicts and every request failed with "'dict' object has no attribute 'role'" whenever role_permissions was set. Convert the entries in the consumer instead, with a TypeAdapter, so the value is parsed wherever it comes from. load_config keeps validating at boot, so a malformed entry still fails startup rather than the first request. --- litellm/proxy/auth/auth_checks.py | 22 +++--- litellm/proxy/proxy_server.py | 6 +- tests/test_litellm/proxy/test_proxy_server.py | 74 +++++++++++++++++++ 3 files changed, 87 insertions(+), 15 deletions(-) diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 4de00f19db3..5db8b62e84d 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -15,10 +15,10 @@ import re import time from collections.abc import Awaitable, Callable, Iterator, Mapping, Sequence from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Protocol, cast +from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Protocol from fastapi import HTTPException, Request, status -from pydantic import BaseModel +from pydantic import BaseModel, TypeAdapter from typing_extensions import ReadOnly, TypedDict import litellm @@ -2414,22 +2414,22 @@ def _update_last_db_access_time(key: str, value: object | None, last_db_access_t last_db_access_time[key] = (value, time.time()) +ROLE_BASED_PERMISSIONS_ADAPTER: Final[TypeAdapter[list[RoleBasedPermissions]]] = TypeAdapter(list[RoleBasedPermissions]) + + def _get_role_based_permissions( rbac_role: RBAC_ROLES, - general_settings: dict, + general_settings: Mapping[str, object], key: Literal["models", "routes"], ) -> list[str] | None: """ Get the role based permissions from the general settings. """ - role_based_permissions: Final = cast( - list[RoleBasedPermissions] | None, - general_settings.get("role_permissions", []), - ) - if role_based_permissions is None: + configured: Final = general_settings.get("role_permissions") + if configured is None: return None - for role_based_permission in role_based_permissions: + for role_based_permission in ROLE_BASED_PERMISSIONS_ADAPTER.validate_python(configured): if role_based_permission.role == rbac_role: return role_based_permission.models if key == "models" else role_based_permission.routes @@ -2438,7 +2438,7 @@ def _get_role_based_permissions( def get_role_based_models( rbac_role: RBAC_ROLES, - general_settings: dict, + general_settings: Mapping[str, object], ) -> list[str] | None: """ Get the models allowed for a user role. @@ -2455,7 +2455,7 @@ def get_role_based_models( def get_role_based_routes( rbac_role: RBAC_ROLES, - general_settings: dict, + general_settings: Mapping[str, object], ) -> list[str] | None: """ Get the routes allowed for a user role. diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 7791f034fba..daf94b79849 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -111,7 +111,6 @@ from litellm.proxy._types import ( PassThroughGenericEndpoint, ProxyErrorTypes, ProxyException, - RoleBasedPermissions, SpecialModelNames, SupportedDBObjectType, TeamDefaultSettings, @@ -317,6 +316,7 @@ from litellm.proxy.analytics_endpoints.analytics_endpoints import ( router as analytics_router, ) from litellm.proxy.auth.auth_checks import ( + ROLE_BASED_PERMISSIONS_ADAPTER, ExperimentalUIJWTToken, can_key_call_resolved_model, get_team_object, @@ -6307,9 +6307,7 @@ class ProxyConfig: ### RBAC ### rbac_role_permissions: Final = general_settings.get("role_permissions", None) if rbac_role_permissions is not None: - general_settings["role_permissions"] = [ # validate role permissions - RoleBasedPermissions(**role_permission) for role_permission in rbac_role_permissions - ] + ROLE_BASED_PERMISSIONS_ADAPTER.validate_python(rbac_role_permissions) ### SSRF URL VALIDATION SETTINGS ### _apply_ssrf_general_settings(general_settings) diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 8ec24e25326..263300d12b1 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -3237,6 +3237,80 @@ async def test_load_config_max_budget_env_var_coerced_to_float(tmp_path, monkeyp litellm.max_budget = original_max_budget +@pytest.mark.asyncio +async def test_load_config_role_permissions_usable_by_jwt_auth(tmp_path): + from litellm.proxy.auth.auth_checks import get_role_based_models, get_role_based_routes + from litellm.proxy.proxy_server import ProxyConfig + + config_file: Final = tmp_path / "config.yaml" + config_file.write_text( + yaml.dump( + { + "model_list": [], + "general_settings": { + "role_permissions": [ + { + "role": "proxy_admin", + "models": ["admin-only-model"], + "routes": ["/v1/embeddings"], + }, + { + "role": "internal_user", + "models": ["shared-model"], + "routes": ["/v1/chat/completions"], + }, + ] + }, + } + ) + ) + + _, _, settings = await ProxyConfig().load_config(router=MagicMock(), config_file_path=str(config_file)) + + assert get_role_based_models(rbac_role="internal_user", general_settings=settings) == ["shared-model"] + assert get_role_based_routes(rbac_role="internal_user", general_settings=settings) == ["/v1/chat/completions"] + assert get_role_based_models(rbac_role="proxy_admin", general_settings=settings) == ["admin-only-model"] + assert get_role_based_routes(rbac_role="proxy_admin", general_settings=settings) == ["/v1/embeddings"] + assert get_role_based_models(rbac_role="team", general_settings=settings) is None + + +@pytest.mark.asyncio +async def test_load_config_without_role_permissions_leaves_every_role_unrestricted(tmp_path): + from litellm.proxy.auth.auth_checks import get_role_based_models, get_role_based_routes + from litellm.proxy.proxy_server import ProxyConfig + + config_file: Final = tmp_path / "config.yaml" + config_file.write_text( + yaml.dump({"model_list": [], "general_settings": {"max_parallel_requests": 7}}) + ) + + _, _, settings = await ProxyConfig().load_config(router=MagicMock(), config_file_path=str(config_file)) + + assert settings["max_parallel_requests"] == 7 + assert get_role_based_models(rbac_role="internal_user", general_settings=settings) is None + assert get_role_based_routes(rbac_role="internal_user", general_settings=settings) is None + + +@pytest.mark.asyncio +async def test_load_config_rejects_malformed_role_permissions(tmp_path): + from pydantic import ValidationError + + from litellm.proxy.proxy_server import ProxyConfig + + config_file: Final = tmp_path / "config.yaml" + config_file.write_text( + yaml.dump( + { + "model_list": [], + "general_settings": {"role_permissions": [{"role": "not_a_real_role", "models": ["gpt-4o"]}]}, + } + ) + ) + + with pytest.raises(ValidationError): + await ProxyConfig().load_config(router=MagicMock(), config_file_path=str(config_file)) + + def test_max_ui_session_budget_default_is_one_dollar(): """LIT-4662: the dashboard session budget default is a product decision; the old 0.25 default locked admins out of auto router Test Connection and the From 4468c9fdcb717cb54a6dcd685a14cdec88f6afe9 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 18 Sep 2026 21:37:01 -0700 Subject: [PATCH 144/144] fix(responses): close the reasoning item before announcing the message item --- .../streaming_iterator.py | 6 ----- .../test_streaming_iterator_transformation.py | 26 +++++++++++++++++++ 2 files changed, 26 insertions(+), 6 deletions(-) diff --git a/litellm/responses/litellm_completion_transformation/streaming_iterator.py b/litellm/responses/litellm_completion_transformation/streaming_iterator.py index 0cda83d979d..5173cd04a89 100644 --- a/litellm/responses/litellm_completion_transformation/streaming_iterator.py +++ b/litellm/responses/litellm_completion_transformation/streaming_iterator.py @@ -927,12 +927,6 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): def _ensure_output_item_for_chunk(self, chunk: ModelResponseStream) -> None: # Change: Never return a value, just enqueue output item events if self.sent_output_item_added_event: - if ( - not self.sent_message_item_added_event - and chunk.choices - and self._get_delta_string_from_streaming_choices(chunk.choices) - ): - self._queue_message_item_added_events() return if not chunk.choices: return diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py b/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py index 3581771bc63..8fbba0dbf87 100644 --- a/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py +++ b/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py @@ -1058,6 +1058,32 @@ async def test_reasoning_then_text_announces_message_item_before_text_events(syn assert announced_indexes_by_item_type["message"] != announced_indexes_by_item_type["reasoning"] +@pytest.mark.asyncio +async def test_reasoning_item_closes_before_message_item_opens(): + iterator: Final = _build_iterator( + [ + _reasoning_chunk("let me think"), + _chunk("Hello"), + _chunk("!", finish_reason="stop"), + ] + ) + + events: Final = await _collect_events(iterator, sync_mode=False) + + item_lifecycle: Final = [ + (event.type, event.item.type) + for event in events + if getattr(event, "type", None) + in (ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED, ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE) + ] + assert item_lifecycle == [ + (ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED, "reasoning"), + (ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE, "reasoning"), + (ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED, "message"), + (ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE, "message"), + ] + + @pytest.mark.parametrize("sync_mode", [True, False]) @pytest.mark.asyncio async def test_tool_then_reasoning_then_text_gives_message_its_own_output_index(sync_mode: bool):