Merge remote-tracking branch 'origin/main' into HEAD

# Conflicts:
#	litellm/proxy/proxy_server.py
#	litellm/text_completions/streaming_iterator.py
#	tests/unit/proxy/test_proxy_server_endpoints_and_startup.py
This commit is contained in:
mateo-berri 2026-10-29 17:26:41 -07:00
commit d7d3054ceb
42 changed files with 2203 additions and 298 deletions

View file

@ -51,7 +51,7 @@ jobs:
suite: ${{ fromJSON((inputs.suite || 'all') == 'all' && '["providers","extensions","accounting","management","management-replica","database","database-replica","mcp","security","cost","sdk","browser"]' || format('["{0}"]', inputs.suite)) }}
env:
UV_PYTHON: "3.12"
INTEGRATION_WORKERS: ${{ inputs.workers || fromJSON('{"providers":"32","extensions":"32","accounting":"16","management":"16","management-replica":"16","cost":"16","mcp":"8","security":"8","database":"4","database-replica":"4"}')[matrix.suite] || '' }}
INTEGRATION_WORKERS: ${{ inputs.workers || fromJSON('{"providers":"32","extensions":"48","accounting":"16","management":"16","management-replica":"16","cost":"16","mcp":"8","security":"8","database":"4","database-replica":"4"}')[matrix.suite] || '' }}
services:
postgres:
image: postgres:16@sha256:e17e86066e5ef83e0952a9347f5c792b7ece00972e2aa787a6986f471b3dd3d5

View file

@ -0,0 +1,35 @@
from collections.abc import Iterable
from typing import Final
from litellm.constants import VLLM_CHOICE_FIELDS
from litellm.types.utils import ModelResponseStream, StreamingChoices
def is_empty_streaming_choice(choice: StreamingChoices) -> bool:
if choice.finish_reason is not None:
return False
if getattr(choice, "logprobs", None) is not None:
return False
if not _all_values_none(_vllm_choice_values(choice)):
return False
delta: Final = choice.delta if hasattr(choice, "delta") else None
if delta is None:
return True
return _all_values_none(delta.model_dump().values())
def _vllm_choice_values(choice: StreamingChoices) -> Iterable[object]:
choice_extras: Final = choice.model_extra or {}
return (choice_extras.get(field) for field in VLLM_CHOICE_FIELDS)
def _all_values_none(values: Iterable[object]) -> bool:
return all(value is None for value in values)
def is_injected_stream_usage_artifact(chunk: object) -> bool:
if not isinstance(chunk, ModelResponseStream):
return False
if chunk.provider_specific_fields is not None:
return False
return all(is_empty_streaming_choice(choice) for choice in chunk.choices or [])

View file

@ -6,6 +6,7 @@ import copy
import json
import re
from collections.abc import Container, Mapping, MutableMapping, Sequence
from dataclasses import dataclass
from datetime import datetime, timezone
from types import MappingProxyType
from typing import Annotated, Any, ClassVar, Final, Literal, TypeVar
@ -1774,6 +1775,48 @@ def anthropic_request_errors_on_thinking_binding_mismatch(request_body: object)
return binding is not None and binding.get("prefix_mismatch_behavior") == "error"
_CLAUDE_STRICT_TOOLS_REJECTION: Final = re.compile(
r"Too many strict tools \(\d+\)"
r"|Schemas contains too many (?:optional parameters|parameters with union types)"
r"|[Cc]ompiled grammar (?:is too large|size)"
r"|tools\.\d+\.custom(?:\.strict)?: "
r"|Grammar compilation timed out"
)
_CLAUDE_TOOLS_ADAPTER: Final = TypeAdapter(tuple[Mapping[str, object], ...])
def is_claude_strict_tools_rejection(status_code: int, error_text: str) -> bool:
return status_code == 400 and _CLAUDE_STRICT_TOOLS_REJECTION.search(error_text) is not None
@dataclass(frozen=True, slots=True)
class ToolsWithoutStrict:
tools: tuple[Mapping[str, object], ...]
dropped_count: int
def tools_without_strict(tools: object) -> ToolsWithoutStrict | None:
try:
validated: Final = _CLAUDE_TOOLS_ADAPTER.validate_python(tools)
except ValidationError:
return None
dropped_count: Final = sum(tool.get("strict") is True for tool in validated)
if dropped_count == 0:
return None
return ToolsWithoutStrict(
tools=tuple({key: value for key, value in tool.items() if key != "strict"} for tool in validated),
dropped_count=dropped_count,
)
def warn_claude_strict_tools_dropped(dropped_count: int) -> None:
litellm.verbose_logger.warning(
"Claude rejected the request's strict schemas, so litellm removed `strict` from every tool (%d had "
"strict true) and sent the request again because drop_params is set.",
dropped_count,
)
def strip_thinking_blocks_from_anthropic_messages(messages: Sequence[object]) -> list[object]:
"""
Return a new message list with thinking / redacted_thinking content blocks removed

View file

@ -4,6 +4,7 @@ from typing import TYPE_CHECKING, Any, Final
import httpx
from litellm.litellm_core_utils.core_helpers import drops_unsupported_params
from litellm.types.llms.anthropic_messages.anthropic_response import (
AnthropicMessagesResponse,
)
@ -201,11 +202,23 @@ class BaseAnthropicMessagesConfig(ABC):
anthropic_request_errors_on_thinking_binding_mismatch,
is_anthropic_invalid_thinking_block_error,
is_anthropic_thinking_binding_mismatch_error,
is_claude_strict_tools_rejection,
tools_without_strict,
validated_object_mapping,
)
if e.response.status_code != 400:
return False
error_text: Final[str] = e.response.text
body: Final = validated_object_mapping(request_body)
drops_rejected_strict_tools: Final = (
drops_unsupported_params(litellm_params.get("drop_params"))
and is_claude_strict_tools_rejection(e.response.status_code, error_text)
and body is not None
and tools_without_strict(body.get("tools")) is not None
)
if drops_rejected_strict_tools:
return True
caller_asked_for_this_error: Final = anthropic_request_errors_on_thinking_binding_mismatch(
request_body
) and is_anthropic_thinking_binding_mismatch_error(error_text)
@ -219,11 +232,22 @@ class BaseAnthropicMessagesConfig(ABC):
"""
from litellm.llms.anthropic.common_utils import (
is_anthropic_invalid_thinking_block_error,
is_claude_strict_tools_rejection,
strip_thinking_blocks_from_anthropic_messages_request_dict,
tools_without_strict,
warn_claude_strict_tools_dropped,
)
if e.response.status_code == 400 and is_anthropic_invalid_thinking_block_error(e.response.text):
strip_thinking_blocks_from_anthropic_messages_request_dict(request_data)
without_strict: Final = (
tools_without_strict(request_data.get("tools"))
if is_claude_strict_tools_rejection(e.response.status_code, e.response.text)
else None
)
if without_strict is not None:
warn_claude_strict_tools_dropped(without_strict.dropped_count)
request_data["tools"] = list(without_strict.tools) # rebind-ok: the retry loop re-signs this same dict
return request_data
def native_output_fallback_litellm_params(

View file

@ -2,11 +2,11 @@ from collections.abc import Callable, Coroutine
from functools import partial
from typing import Final, TypeAlias
from litellm.llms.anthropic.common_utils import is_claude_strict_tools_rejection
from litellm.llms.base_llm.chat.transformation import BaseLLMException
from litellm.llms.bedrock.common_utils import (
NativeOutputFallback,
is_bedrock_native_output_rejection,
is_bedrock_strict_tools_rejection,
warn_bedrock_native_output_rejected,
warn_bedrock_strict_tools_dropped,
)
@ -34,7 +34,7 @@ def _native_output_retry(
def _strict_tools_retry(
error: BaseLLMException, fallback: NativeOutputFallbackFactory | None, model: str
) -> NativeOutputFallback | None:
if fallback is None or not is_bedrock_strict_tools_rejection(error.status_code, error.message):
if fallback is None or not is_claude_strict_tools_rejection(error.status_code, error.message):
return None
retry: Final = fallback()
if retry is not None:

View file

@ -366,21 +366,10 @@ class NativeOutputFallback:
litellm_params: Mapping[str, object]
_BEDROCK_STRICT_TOOLS_REJECTION: Final = re.compile(
r"Too many strict tools \(\d+\)"
r"|Schemas contains too many (?:optional parameters|parameters with union types)"
r"|[Cc]ompiled grammar (?:is too large|size)"
r"|tools\.\d+\.custom(?:\.strict)?: "
r"|Grammar compilation timed out"
)
_TOOLS_ADAPTER: Final = TypeAdapter(tuple[Mapping[str, object], ...])
_TOOL_SPEC_ADAPTER: Final = TypeAdapter(Mapping[str, object])
def is_bedrock_strict_tools_rejection(status_code: int, error_text: str) -> bool:
return status_code == 400 and _BEDROCK_STRICT_TOOLS_REJECTION.search(error_text) is not None
def warn_bedrock_strict_tools_dropped(model: str) -> None:
verbose_logger.warning(
"Bedrock rejected the strict tool schemas for model=%s, so litellm sent the request again without `strict` "

View file

@ -2467,8 +2467,7 @@ class BaseLLMHTTPHandler:
)
else:
verbose_logger.debug(
"Anthropic /v1/messages: invalid thinking signature; "
"stripping thinking blocks and retrying (attempt %s/%s).",
"Anthropic /v1/messages: retrying after a recoverable HTTP error (attempt %s/%s).",
len(recovered_errors) + 2,
max_attempts,
)

View file

@ -7975,6 +7975,8 @@ async def atext_completion(*args, **kwargs) -> TextCompletionResponse | TextComp
model=model,
custom_llm_provider=custom_llm_provider,
stream_options=kwargs.get("stream_options"),
usage_injected_by_proxy="_litellm_strip_stream_usage" in kwargs
and kwargs["_litellm_strip_stream_usage"] is True,
)
else:
## OpenAI / Azure Text Completion Returns here

View file

@ -121,6 +121,7 @@ _ASYNC_FEATURES: Final[detectors.Detectors] = MappingProxyType(
Feature.PRE_REQUEST_HOOK: route_detectors.pre_request_hook,
Feature.WEBSEARCH_INTERCEPTION: route_detectors.websearch_interception,
Feature.MESSAGES_INTERCEPTOR: route_detectors.intercepted,
Feature.STRICT_TOOLS_RETRY: route_detectors.strict_tools_retry,
}
)

View file

@ -12,6 +12,11 @@ from litellm.llms.custom_httpx.http_handler import (
from litellm.proxy._types import CommonProxyErrors, LitellmUserRoles, UserAPIKeyAuth
class OAuth2TokenValidationError(ValueError):
"""The token-info or introspection endpoint did not vouch for the token: it rejected it, reported it
inactive, or could not be reached."""
class Oauth2Handler:
"""
Handles OAuth2 token validation.
@ -207,7 +212,7 @@ class Oauth2Handler:
)
except httpx.HTTPStatusError as e:
# This will catch any 4xx or 5xx errors
raise ValueError(f"Oauth 2.0 Token validation failed: {e}")
raise OAuth2TokenValidationError(f"Oauth 2.0 Token validation failed: {e}")
except Exception as e:
# This will catch any other errors (like network issues)
raise ValueError(f"An error occurred during token validation: {e}")
raise OAuth2TokenValidationError(f"An error occurred during token validation: {e}")

View file

@ -1160,6 +1160,24 @@ def _matches_routing_override(token_claims: dict, override: "JWTRoutingOverride"
)
def oauth2_introspection_applies(
api_key: str, route: str, *, enable_oauth2_auth: bool, enable_jwt_auth: bool, jwt_handler: JWTHandler
) -> bool:
"""Whether LiteLLM validates ``api_key`` on ``route`` by OAuth2 token introspection, on LLM API, info and
lens routes only: a non-JWT token when OAuth2 auth is on globally, or a JWT matched by a routing override."""
is_jwt: Final = enable_jwt_auth and jwt_handler.is_jwt(token=api_key)
# Routing uses unverified JWT claims only to choose auth path.
# Final authentication is enforced by the selected validator.
route_jwt_to_oauth2: Final = is_jwt and _should_route_jwt_to_oauth2_override(token=api_key, jwt_handler=jwt_handler)
covered_route: Final = (
RouteChecks.is_llm_api_route(route=route)
or RouteChecks.is_info_route(route=route)
or route == "/lens"
or route.startswith("/lens/")
)
return covered_route and ((enable_oauth2_auth and not is_jwt) or route_jwt_to_oauth2)
def _should_route_jwt_to_oauth2_override(token: str, jwt_handler: JWTHandler) -> bool:
routing_overrides: Final = jwt_handler.litellm_jwtauth.routing_overrides
if not routing_overrides:
@ -2278,23 +2296,13 @@ async def user_api_key_auth_builder(
enable_oauth2_auth: Final = general_settings.get("enable_oauth2_auth", False) is True
enable_jwt_auth: Final = general_settings.get("enable_jwt_auth", False) is True
is_jwt = jwt_handler.is_jwt(token=api_key) if enable_jwt_auth else False
# Routing uses unverified JWT claims only to choose auth path.
# Final authentication is enforced by the selected validator.
route_jwt_to_oauth2 = is_jwt and _should_route_jwt_to_oauth2_override(token=api_key, jwt_handler=jwt_handler)
# OAuth2 applies for:
# 1) when global OAuth2 auth is enabled on LLM + info routes
# 2) JWT tokens that explicitly match routing_overrides on LLM + info routes
route_is_lens: Final = route == "/lens" or route.startswith("/lens/")
should_apply_override_oauth2: Final = route_jwt_to_oauth2 and (
RouteChecks.is_llm_api_route(route=route) or RouteChecks.is_info_route(route=route) or route_is_lens
)
should_apply_global_oauth2: Final = enable_oauth2_auth and (
RouteChecks.is_llm_api_route(route=route) or RouteChecks.is_info_route(route=route) or route_is_lens
)
if (should_apply_global_oauth2 and not is_jwt) or should_apply_override_oauth2:
if oauth2_introspection_applies(
api_key,
route,
enable_oauth2_auth=enable_oauth2_auth,
enable_jwt_auth=enable_jwt_auth,
jwt_handler=jwt_handler,
):
from litellm.proxy.proxy_server import premium_user
if premium_user is not True:

View file

@ -588,6 +588,7 @@ def model_deployments_support_stream_options(
_STREAM_USAGE_FLAGS: Final = ("_litellm_strip_stream_usage", STREAM_USAGE_INJECTED_FIELD)
_STREAM_USAGE_TRACKED_ROUTES: Final = frozenset({"acompletion", "atext_completion"})
def stream_usage_tracking_updates(
@ -613,7 +614,7 @@ def stream_usage_tracking_updates(
return scrub
if always_include is False and route_type == "agenerate_content_stream" and "stream_options" not in data:
return {**scrub, "stream_options": {"include_usage": False}}
if always_include is False or route_type != "acompletion":
if always_include is False or route_type not in _STREAM_USAGE_TRACKED_ROUTES:
return scrub
if isinstance(stream_options, dict) and stream_options.get("include_usage") is True:
return scrub

View file

@ -664,4 +664,7 @@ class GenericGuardrailAPI(CustomGuardrail):
GuardrailEventHooks.pre_call,
GuardrailEventHooks.post_call,
GuardrailEventHooks.during_call,
GuardrailEventHooks.pre_mcp_call,
GuardrailEventHooks.during_mcp_call,
GuardrailEventHooks.post_mcp_call,
]

View file

@ -2505,6 +2505,7 @@ def refresh_proxy_server_request_body_snapshot(
"litellm_logging_obj",
VECTOR_STORE_CALLER_METADATA_KWARG,
STREAM_USAGE_INJECTED_FIELD,
"_litellm_strip_stream_usage",
}
)
| _TRANSPORT_ONLY_CREDENTIAL_KEYS

View file

@ -68,8 +68,11 @@ module materially harder to understand.
behind a single generic branch unless tests prove every mode still behaves
correctly.
- 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.
with `delegate_auth_to_upstream: true` is deprecated but still supported. A bare
`Authorization` that is a LiteLLM credential is validated and never forwarded. Any other
bearer is the upstream token: it is forwarded, and the caller is admitted as the JWT or
custom-auth user it also proves when that user is granted the server, else anonymously.
New servers should use `auth_type: true_passthrough` or `auth_type: oauth_delegate`.
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

View file

@ -2,14 +2,15 @@ import asyncio
import functools
import re
import secrets
from collections.abc import Awaitable, Callable, Mapping, Sequence
from collections.abc import Awaitable, Callable, Iterator, Mapping, Sequence
from collections.abc import Set as AbstractSet
from dataclasses import dataclass
from datetime import datetime, timezone
from itertools import chain
from types import MappingProxyType
from typing import TYPE_CHECKING, Final, Literal
from typing import TYPE_CHECKING, Final, Literal, TypeAlias
import jwt
from fastapi import HTTPException
from pydantic import SecretStr, TypeAdapter
from starlette.datastructures import Headers
@ -45,9 +46,11 @@ from litellm.proxy.agent_endpoints.auth.agent_access_groups import (
from litellm.proxy.agent_endpoints.auth.agent_caller import agent_caller_auth
from litellm.proxy.agent_endpoints.auth.managed_authorization import managed_agent_policy
from litellm.proxy.auth.ip_address_utils import IPAddressUtils
from litellm.proxy.auth.oauth2_check import OAuth2TokenValidationError
from litellm.proxy.auth.user_api_key_auth import ( # noqa: F401 # legacy module exports
_get_bearer_token_or_received_api_key, # pyright: ignore[reportPrivateUsage] # shared x-litellm-api-key parser lives with user_api_key_auth
_run_centralized_common_checks, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export
oauth2_introspection_applies,
run_centralized_common_checks,
user_api_key_auth,
)
@ -100,6 +103,9 @@ if TYPE_CHECKING:
_EMPTY_TOOLSET_GRANTS: Final[Mapping[str, Sequence[str]]] = MappingProxyType({})
OPTIONAL_STRING_ADAPTER: Final[TypeAdapter[str | None]] = TypeAdapter(str | None)
_PROXY_SETTINGS_ADAPTER: Final = TypeAdapter(Mapping[str, object])
_CAUSE_CHAIN_MAX_DEPTH: Final = 16
_IdentityCheck: TypeAlias = Literal["oauth2", "oauth2_proxy", "jwt"]
_CACHED_VALUE: Final = TypeAdapter(object)
@ -233,6 +239,136 @@ def _is_legacy_delegate_cold_start(mcp_servers: list[str] | None, client_ip: str
return True
def _cause_chain(exc: BaseException) -> Iterator[BaseException]:
link: BaseException | None = exc # rebind-ok: advances one link of the cause chain per iteration
for _ in range(_CAUSE_CHAIN_MAX_DEPTH):
if link is None:
return
yield link
link = link.__cause__ or link.__context__
def _caused_by(exc: BaseException, error_type: type[BaseException]) -> bool:
return any(isinstance(link, error_type) for link in _cause_chain(exc))
def _signed_with(token: str, secret: str) -> bool:
"""True when ``token`` is a JWT whose HS256 signature verifies with ``secret``, whatever its claims say."""
try:
jwt.decode(
token,
secret,
algorithms=["HS256"],
options={"verify_exp": False, "verify_nbf": False, "verify_iat": False, "verify_aud": False},
)
except jwt.PyJWTError:
return False
return True
async def _bearer_is_litellm_credential(token: str) -> bool:
"""True when ``token`` is, or may be, a credential LiteLLM issued: the master key or a JWT it signed
(dashboard sessions, BYOK sessions), a sealed gateway bearer, or a virtual key the key table holds (a
lookup fault counts as held). Only ``sk-`` tokens are looked up, because no other token can be a
virtual key."""
from litellm.proxy.mcp_server.bridge_token_flow import ( # noqa: PLC0415 # bridge flow imports this module
is_sealed_gateway_bearer,
virtual_key_may_exist,
)
from litellm.proxy.mcp_server.mcp_server_manager import ( # noqa: PLC0415 # manager imports this module
LITELLM_VIRTUAL_KEY_PREFIX,
)
from litellm.proxy.proxy_server import master_key # noqa: PLC0415 # startup owns the master key
if master_key and (secrets.compare_digest(token.encode(), master_key.encode()) or _signed_with(token, master_key)):
return True
if is_sealed_gateway_bearer(token):
return True
return token.startswith(LITELLM_VIRTUAL_KEY_PREFIX) and await virtual_key_may_exist(token)
def _identity_check(token: str, route: str, settings: Mapping[str, object]) -> _IdentityCheck | None:
"""The identity check LiteLLM's own auth would run on ``token`` for ``route``, in its order: OAuth2
introspection, then trusted reverse-proxy identity headers, then a JWT whose issuer may be LiteLLM's."""
from litellm.proxy.auth.handle_jwt import JWTHandler # noqa: PLC0415 # proxy import cycle
from litellm.proxy.mcp_server.bridge_token_flow import ( # noqa: PLC0415 # bridge flow imports this module
jwt_issuer_may_be_litellm,
)
from litellm.proxy.proxy_server import ( # noqa: PLC0415 # startup owns the active auth configuration
jwt_handler,
premium_user,
)
enable_jwt_auth: Final = settings.get("enable_jwt_auth") is True
if oauth2_introspection_applies(
token,
route,
enable_oauth2_auth=settings.get("enable_oauth2_auth") is True,
enable_jwt_auth=enable_jwt_auth,
jwt_handler=jwt_handler,
):
return "oauth2" if premium_user is True else None
if settings.get("enable_oauth2_proxy_auth") is True:
return "oauth2_proxy"
if enable_jwt_auth and premium_user is True and JWTHandler.is_jwt(token=token) and jwt_issuer_may_be_litellm(token):
return "jwt"
return None
async def _verify_identity(check: Literal["oauth2_proxy", "jwt"], token: str, request: Request) -> None:
from litellm.proxy.auth.oauth2_proxy_hook import ( # noqa: PLC0415 # proxy import cycle
handle_oauth2_proxy_request,
)
from litellm.proxy.proxy_server import jwt_handler # noqa: PLC0415 # startup owns the active auth configuration
match check:
case "oauth2_proxy":
await handle_oauth2_proxy_request(request=request)
case "jwt":
await jwt_handler.auth_jwt(token=token) # pyright: ignore[reportUnknownMemberType] # returns a bare dict; only whether it raises matters
case _: # pragma: no cover - exhaustiveness guard, unreachable while the Literal holds
assert_never(check) # pragma: no cover
async def _upstream_bearer_identity(token: str, litellm_api_key: str, request: Request) -> UserAPIKeyAuth | None:
"""The LiteLLM user an upstream bearer also proves, or ``None`` when it proves none.
Custom auth decides alone, so any rejection is ``None``. Otherwise the identity check LiteLLM would run
must vouch for the caller. Reverse-proxy headers and a JWT are verified first, and OAuth2 introspection
runs once inside the full check, where a provider that does not vouch is ``None``. Past that, only a
user LiteLLM does not know is ``None``; every other refusal (budget, blocked team, disallowed route)
propagates."""
from litellm.proxy.auth.auth_utils import get_request_route # noqa: PLC0415 # proxy import cycle
from litellm.proxy.mcp_server.bridge_token_flow import ( # noqa: PLC0415 # bridge flow imports this module
custom_auth_configured,
)
from litellm.proxy.proxy_server import ( # noqa: PLC0415 # startup owns the active auth configuration
general_settings, # pyright: ignore[reportUnknownVariableType] # legacy global is declared as a bare dict
)
if custom_auth_configured():
try:
return await user_api_key_auth(api_key=litellm_api_key, request=request)
except (HTTPException, ProxyException):
return None
check: Final = _identity_check(
token, get_request_route(request), _PROXY_SETTINGS_ADAPTER.validate_python(general_settings)
)
if check is None:
return None
if check != "oauth2":
try:
await _verify_identity(check, token, request)
except Exception: # noqa: BLE001 # any verification failure means LiteLLM does not vouch for this bearer
return None
try:
return await user_api_key_auth(api_key=litellm_api_key, request=request)
except (HTTPException, ProxyException, OAuth2TokenValidationError) as exc:
if _caused_by(exc, UserNotFoundError) or _caused_by(exc, OAuth2TokenValidationError):
return None
raise
def _is_litellm_auth_admission_error(exc: Exception) -> bool:
if isinstance(exc, HTTPException):
return exc.status_code == 401
@ -704,6 +840,19 @@ class MCPRequestHandler:
route=request_route,
mcp_servers=mcp_servers,
)
elif oauth2_headers and (
legacy_delegate_targets := MCPRequestHandler._legacy_delegate_targets(
path=request_route,
mcp_servers=mcp_servers,
client_ip=IPAddressUtils.get_mcp_client_ip(request),
)
):
validated_user_api_key_auth = await MCPRequestHandler._admit_legacy_delegate_bearer( # rebind-ok: one assignment per admission arm, like its siblings
authorization_value=oauth2_headers["Authorization"],
litellm_api_key=litellm_api_key,
targets=legacy_delegate_targets,
request=request,
)
elif oauth2_headers:
# Authorization on a non-delegated server: the bearer must be a real
# LiteLLM credential, so a failed validation is a genuine 401/403 and
@ -741,10 +890,14 @@ class MCPRequestHandler:
custom_key_header_name: Final = OPTIONAL_STRING_ADAPTER.validate_python(
general_settings.get("litellm_key_header_name")
)
admitted_credential: Final = MCPRequestHandler.caller_admission_credential(
headers,
validated_user_api_key_auth,
custom_key_header_name=custom_key_header_name,
admitted_credential: Final = (
None
if validated_user_api_key_auth.mcp_forwards_caller_authorization
else MCPRequestHandler.caller_admission_credential(
headers,
validated_user_api_key_auth,
custom_key_header_name=custom_key_header_name,
)
)
(
oauth2_headers,
@ -934,6 +1087,53 @@ class MCPRequestHandler:
return False
return True
@staticmethod
def _legacy_delegate_targets(
path: str, mcp_servers: Sequence[str] | None, client_ip: str | None
) -> tuple[MCPServer, ...]:
"""Every server the request targets when ALL of them are legacy ``delegate_auth_to_upstream``
servers, else ``()``. Names resolve the way downstream routing resolves them, so a permissive
``x-mcp-servers`` header cannot open the gate for a stricter path target."""
from litellm.proxy.mcp_server.mcp_server_manager import ( # noqa: PLC0415 # manager imports this module
global_mcp_server_manager,
is_legacy_upstream_delegate,
)
servers: Final = tuple(
global_mcp_server_manager.get_mcp_server_by_name(name, client_ip=client_ip)
for name in MCPRequestHandler._resolve_target_server_names(path=path, mcp_servers_header=mcp_servers)
)
delegates: Final = tuple(
server for server in servers if server is not None and is_legacy_upstream_delegate(server)
)
return delegates if servers and len(delegates) == len(servers) else ()
@staticmethod
async def _admit_legacy_delegate_bearer(
authorization_value: str,
litellm_api_key: str,
targets: tuple[MCPServer, ...],
request: Request,
) -> UserAPIKeyAuth:
"""Admit a bare ``Authorization`` bearer sent to legacy ``delegate_auth_to_upstream`` servers.
A LiteLLM credential runs normal admission, and its errors propagate. Any other bearer is the
upstream token and is marked for forwarding: a user it also proves who is granted every target is
admitted as that user, and every other caller is admitted anonymously."""
from litellm.proxy.mcp_server.mcp_server_manager import ( # noqa: PLC0415 # manager imports this module
global_mcp_server_manager,
)
token: Final = _normalize_caller_admission_credential(authorization_value)
if await _bearer_is_litellm_credential(token):
return await user_api_key_auth(api_key=litellm_api_key, request=request)
identity: Final = await _upstream_bearer_identity(token, litellm_api_key, request)
if identity is not None:
allowed: Final = frozenset(await global_mcp_server_manager.get_allowed_mcp_servers(identity))
if all(server.server_id in allowed for server in targets):
return identity.model_copy(update={"mcp_forwards_caller_authorization": True})
return UserAPIKeyAuth().model_copy(update={"mcp_forwards_caller_authorization": True})
@staticmethod
def _single_dcr_bridge_delegate_target(
path: str, mcp_servers: list[str] | None, client_ip: str | None
@ -1553,7 +1753,7 @@ class MCPRequestHandler:
return expiry >= datetime.now(timezone.utc)
@staticmethod
def _resolve_target_server_names(path: str, mcp_servers_header: list[str] | None) -> list[str]:
def _resolve_target_server_names(path: str, mcp_servers_header: Sequence[str] | None) -> Sequence[str]:
"""
Resolve the target MCP server names exactly as downstream routing
does (``server.py::extract_mcp_auth_context``).

View file

@ -61,7 +61,6 @@ async def oauth_authorization_uses_gateway_credential(request: Request) -> bool:
from litellm.proxy.proxy_server import ( # noqa: PLC0415 # startup owns the active auth configuration
jwt_handler,
master_key,
user_custom_auth,
)
if "x-litellm-api-key" in request.headers:
@ -71,10 +70,26 @@ async def oauth_authorization_uses_gateway_credential(request: Request) -> bool:
return "authorization" in request.headers
if token.startswith("sk-") or (master_key and secrets.compare_digest(token.encode(), master_key.encode())):
return True
if user_custom_auth is not None or jwt_handler.litellm_jwtauth.oidc_userinfo_enabled:
if custom_auth_configured() or jwt_handler.litellm_jwtauth.oidc_userinfo_enabled:
return True
if not JWTHandler.is_jwt(token):
return await _opaque_bearer_is_gateway_credential(token)
return is_sealed_gateway_bearer(token) or await virtual_key_may_exist(token)
return jwt_issuer_may_be_litellm(token)
def custom_auth_configured() -> bool:
from litellm.proxy.proxy_server import ( # noqa: PLC0415 # startup owns the active auth configuration
user_custom_auth,
)
return user_custom_auth is not None
def jwt_issuer_may_be_litellm(token: str) -> bool:
"""False only when the JWT names an issuer LiteLLM's JWT auth is not configured to accept."""
from litellm.proxy.auth.handle_jwt import JWTHandler # noqa: PLC0415 # proxy import cycle
from litellm.proxy.proxy_server import jwt_handler # noqa: PLC0415 # startup owns the active auth configuration
claims: Final = JWTHandler.get_unverified_claims(token)
issuer: Final = claims.get("iss") if claims is not None else None
global_issuer: Final = os.getenv("JWT_ISSUER")
@ -86,25 +101,35 @@ async def oauth_authorization_uses_gateway_credential(request: Request) -> bool:
)
async def _opaque_bearer_is_gateway_credential(token: str) -> bool:
from litellm.proxy._types import hash_token # noqa: PLC0415 # proxy import cycle
def is_sealed_gateway_bearer(token: str) -> bool:
"""True for a bearer only the gateway issues: a bridge or refresh envelope, a v2 GCM blob, or a
dashboard UI token. A UI token that cannot be read still counts, so a lookup fault never reads as foreign."""
from litellm.proxy.auth.auth_checks import ExperimentalUIJWTToken # noqa: PLC0415 # proxy import cycle
from litellm.proxy.auth.resolvers.exceptions import KeyNotFoundError # noqa: PLC0415 # proxy import cycle
from litellm.proxy.auth.resolvers.store import IdentityStore # noqa: PLC0415 # proxy import cycle
from litellm.proxy.mcp_server.outbound_credentials.envelope import ( # noqa: PLC0415 # envelope imports bridge types
is_envelope,
is_refresh_envelope,
)
from litellm.proxy.proxy_server import ( # noqa: PLC0415 # startup owns the identity store dependencies
prisma_client,
user_api_key_cache,
)
if is_envelope(token) or is_refresh_envelope(token) or token.startswith(V2_GCM_PREFIX):
return True
try:
if ExperimentalUIJWTToken.get_key_object_from_ui_hash_key(token) is not None:
return True
return ExperimentalUIJWTToken.get_key_object_from_ui_hash_key(token) is not None
except Exception as exc: # noqa: BLE001 # an identity lookup fault must not permit cookie fallback
verbose_logger.debug("OAuth bearer ownership could not be checked (%s)", type(exc).__name__)
return True
async def virtual_key_may_exist(token: str) -> bool:
"""False only when the key table answers that ``token`` is not a virtual key; a lookup fault counts as held."""
from litellm.proxy._types import hash_token # noqa: PLC0415 # proxy import cycle
from litellm.proxy.auth.resolvers.exceptions import KeyNotFoundError # noqa: PLC0415 # proxy import cycle
from litellm.proxy.auth.resolvers.store import IdentityStore # noqa: PLC0415 # proxy import cycle
from litellm.proxy.proxy_server import ( # noqa: PLC0415 # startup owns the identity store dependencies
prisma_client,
user_api_key_cache,
)
try:
await IdentityStore(prisma_client, user_api_key_cache).resolve(hashed_token=hash_token(token))
except KeyNotFoundError:
return False

View file

@ -1196,8 +1196,11 @@ def should_strip_caller_authorization(
pass-through cold-start case (RFC 9728) the bearer in
``Authorization`` is the upstream OAuth token and must be
forwarded, so we keep it.
- **Delegated OAuth servers**: admission always runs and there is no
anonymous path, so the caller's separate ``Authorization`` is
- **Legacy delegate servers**: forward when MCP admission marked the
caller's ``Authorization`` as the upstream token
(``mcp_forwards_caller_authorization``); otherwise follow the
delegated OAuth rule below.
- **Delegated OAuth servers**: the caller's separate ``Authorization`` is
forwarded only when a distinct ``x-litellm-api-key`` carried
admission. Without that header the ``Authorization`` *was* the
admission credential — a virtual key, an IdP JWT, or an SSO / OIDC /
@ -1206,6 +1209,12 @@ def should_strip_caller_authorization(
- **Every other auth-enabled mode**: strip when ``Authorization`` is the
credential admission validated.
"""
if (
user_api_key_auth is not None
and user_api_key_auth.mcp_forwards_caller_authorization
and is_legacy_upstream_delegate(mcp_server)
):
return False
if _authorization_is_admission_jwt(raw_headers, user_api_key_auth):
return True
if mcp_server.auth_type == MCPAuth.oauth2_token_exchange:
@ -1557,6 +1566,18 @@ async def _materialize_auth_headers(auth: httpx2.Auth | None) -> dict[str, str]
return {header_name: header_value} if header_value else None
def is_legacy_upstream_delegate(server: MCPServer) -> bool:
"""True for an interactive ``auth_type=oauth2`` server with ``delegate_auth_to_upstream=true``, whose
caller ``Authorization`` is the upstream token. The flow is resolved, not read from the column, so an
unstamped M2M-shaped row stays excluded: its stored client credentials drive egress, and admitting it
anonymously would let any caller act as LiteLLM's service account."""
return (
server.auth_type == MCPAuth.oauth2
and server.delegate_auth_to_upstream is True
and MCPServerManager.effective_oauth2_flow(server) != "client_credentials"
)
def _consumes_caller_authorization(server: MCPServer) -> bool:
"""True when this server's egress forwards the caller's request-wide ``Authorization`` upstream:
the client-forwarded token modes, legacy OAuth pass-through, and legacy upstream-delegated
@ -1726,18 +1747,15 @@ def _warn_overlapping_tool_policies(server: MCPServer) -> None:
def _warn_legacy_delegate_auth_if_applicable(server: MCPServer, *, source: str) -> None:
"""Direct legacy delegated OAuth configurations to the admitted replacement."""
if server.auth_type != MCPAuth.oauth2:
return
if getattr(server, "delegate_auth_to_upstream", False) is not True:
return
if server.has_client_credentials:
"""Point legacy delegated OAuth configurations at the client-forwarded replacements."""
if not is_legacy_upstream_delegate(server):
return
label: Final = get_server_prefix(server)
verbose_logger.warning(
"MCP server %r (id=%s, source=%s) uses deprecated auth_type=oauth2 with "
"delegate_auth_to_upstream=true. LiteLLM admission is now required; migrate to "
"auth_type=oauth_delegate for client-forwarded OAuth.",
"delegate_auth_to_upstream=true. It still works: a caller without a LiteLLM credential is "
"admitted anonymously and its upstream token is forwarded. Migrate to auth_type=true_passthrough "
"or auth_type=oauth_delegate.",
label,
server.server_id,
source,
@ -3766,12 +3784,12 @@ class MCPServerManager:
)
)
if is_anonymous:
passthrough_server_ids: Final = [
anonymous_server_ids: Final = [
server.server_id
for server in self.get_registry().values()
if server.auth_type == MCPAuth.true_passthrough
if server.auth_type == MCPAuth.true_passthrough or is_legacy_upstream_delegate(server)
]
combined_servers.update(passthrough_server_ids)
combined_servers.update(anonymous_server_ids)
restrict_allow_all: Final = (
resolved_general_settings.get("mcp_allow_all_keys_respects_mcp_scope", False)

View file

@ -588,6 +588,7 @@ if MCP_AVAILABLE:
from litellm.proxy.mcp_server.mcp_server_manager import (
MCPServerManager,
global_mcp_server_manager,
is_legacy_upstream_delegate,
)
######################################################
@ -2162,8 +2163,9 @@ if MCP_AVAILABLE:
oauth2_headers: Mapping[str, str] | None,
mcp_server_auth_headers: dict[str, dict[str, str]] | None,
raw_headers: Mapping[str, str] | None,
requested_servers: Sequence[str] | None = None,
) -> None:
"""Probe pass-through upstream servers in parallel before the MCP session starts.
"""Probe pass-through and legacy delegate upstream servers in parallel before the MCP session starts.
Only servers the caller's key is already authorized to reach are probed —
the list is derived from _get_allowed_mcp_servers so that a user cannot
@ -2186,6 +2188,18 @@ if MCP_AVAILABLE:
mcp_servers=mcp_servers,
client_ip=client_ip,
)
# A legacy delegate is probed only when the single requested name resolves to it the way admission
# resolves it, so neither a multi-server connect nor an access group can be failed by one rejected
# token; its challenge echoes the requested name, matching the tokenless challenge.
requested_scope: Final = requested_servers if requested_servers is not None else mcp_servers
single_requested_name: Final = (
requested_scope[0] if requested_scope is not None and len(requested_scope) == 1 else None
)
requested_server: Final = (
operations.global_mcp_server_manager.get_mcp_server_by_name(single_requested_name, client_ip=client_ip)
if single_requested_name is not None
else None
)
prepared_headers: Final = tuple(
(
srv,
@ -2201,6 +2215,11 @@ if MCP_AVAILABLE:
)
for srv in allowed_servers
if srv.is_oauth_passthrough
or (
requested_server is not None
and srv.server_id == requested_server.server_id
and is_legacy_upstream_delegate(srv)
)
)
probe_headers: Final = tuple(
(
@ -2213,7 +2232,7 @@ if MCP_AVAILABLE:
for srv, (auth, extra) in prepared_headers
)
probe_targets: Final[tuple[tuple[MCPServer, str, str], ...]] = tuple(
(srv, auth_header, srv.name)
(srv, auth_header, srv.name if srv.is_oauth_passthrough else single_requested_name or srv.name)
for srv, headers in probe_headers
if (
auth_header := next(
@ -2311,6 +2330,7 @@ if MCP_AVAILABLE:
if context.mcp_server_auth_headers is not None
else None,
raw_headers=context.raw_headers,
requested_servers=context.mcp_servers,
)
return None

View file

@ -202,7 +202,6 @@ from litellm.types.utils import (
CountContextManagement,
ModelResponse,
ModelResponseStream,
StreamingChoices,
TextCompletionResponse,
TokenCountMessagesFormat,
TokenCountResponse,
@ -385,6 +384,7 @@ from litellm.litellm_core_utils.sensitive_data_masker import (
SensitiveDataMasker,
mask_sensitive_keys,
)
from litellm.litellm_core_utils.stream_usage_chunks import is_injected_stream_usage_artifact
from litellm.llms.anthropic.messages.convertible_blocks import (
deployment_drops_params,
drop_non_string_tool_result_texts,
@ -10250,31 +10250,6 @@ def _dump_set_fields(model: BaseModel, exclude: AbstractSet[str] = frozenset())
return model.model_dump(mode="json", exclude_none=True, exclude_unset=True, exclude=set(exclude))
def _is_injected_stream_usage_artifact(chunk: object) -> bool:
if not isinstance(chunk, ModelResponseStream):
return False
if chunk.provider_specific_fields is not None:
return False
return all(_is_empty_streaming_choice(choice) for choice in chunk.choices or [])
def _is_empty_streaming_choice(choice: StreamingChoices) -> bool:
if choice.finish_reason is not None:
return False
if getattr(choice, "logprobs", None) is not None:
return False
if not _all_values_none(_vllm_choice_fields(choice).values()):
return False
delta: Final = choice.delta if hasattr(choice, "delta") else None
if delta is None:
return True
return _all_values_none(delta.model_dump().values())
def _all_values_none(values: Iterable[object]) -> bool:
return all(value is None for value in values)
async def _apply_streaming_chunk_hooks(
*,
chunk: Any,
@ -10654,7 +10629,7 @@ async def async_data_generator(
fallback_model_from_metadata=fallback_model_from_metadata,
)
if strip_stream_usage and _is_injected_stream_usage_artifact(chunk):
if strip_stream_usage and is_injected_stream_usage_artifact(chunk):
if pending_fallback_event:
yield _format_fallback_metadata_sse_event(
fallback_model=fallback_model_from_metadata,

View file

@ -98,6 +98,8 @@ class Feature(str, Enum):
"""Replayed thinking blocks, which Python strips and retries once if Anthropic rejects their signature."""
TOOL_CUSTOM_FIELD = "tool_custom_field"
"""Tools carrying a non-spec custom object, which Python drops or hoists for the provider."""
STRICT_TOOLS_RETRY = "strict_tools_retry"
"""Strict tools under drop_params, which Python sends again without strict if Claude rejects their limits."""
CLIENT_HEADERS = "client_headers"
"""Proxy-forwarded headers that Python merges into extra_headers."""
DEPLOYMENT_HOOK = "deployment_hook"
@ -179,6 +181,7 @@ MESSAGES_PYTHON_FEATURES: Final = COMMON_PYTHON_FEATURES | {
Feature.STREAM_COST_INJECTION,
Feature.THINKING_SIGNATURE_RETRY,
Feature.TOOL_CUSTOM_FIELD,
Feature.STRICT_TOOLS_RETRY,
}
RESPONSES_PYTHON_FEATURES: Final = COMMON_PYTHON_FEATURES | {

View file

@ -5,6 +5,7 @@ from __future__ import annotations
from typing import Final
import litellm
from litellm.litellm_core_utils.core_helpers import drops_unsupported_params
from litellm.llms.opencode.common_utils import is_opencode_api_base
from litellm.rust_bridge.detectors import api_base, overridden, present, streaming
from litellm.rust_bridge.public_call import NativeCall, optional_mapping, optional_sequence, optional_str
@ -48,6 +49,13 @@ def tool_custom_field(request: NativeCall) -> bool:
return any("custom" in (optional_mapping(tool) or {}) for tool in tools)
def strict_tools_retry(request: NativeCall) -> bool:
tools: Final = optional_sequence(request.resolved.get("tools")) or ()
return drops_unsupported_params(request.resolved.get("drop_params")) and any(
(optional_mapping(tool) or {}).get("strict") is True for tool in tools
)
def stream_cost_injection(request: NativeCall) -> bool:
"""Python adds response cost to streaming usage when the global setting is truthy."""
return streaming(request) and bool(litellm.include_cost_in_streaming_usage)

View file

@ -1,7 +1,8 @@
from typing import Final
from litellm.constants import VLLM_CHOICE_FIELDS
from litellm.types.utils import Choices, ModelResponse, TextChoices, TextCompletionResponse
from litellm.litellm_core_utils.stream_usage_chunks import is_injected_stream_usage_artifact
from litellm.types.utils import Choices, ModelResponse, ModelResponseStream, TextChoices, TextCompletionResponse
class TextCompletionStreamWrapper:
@ -11,11 +12,18 @@ class TextCompletionStreamWrapper:
model,
stream_options: dict | None = None,
custom_llm_provider: str | None = None,
usage_injected_by_proxy: bool = False,
):
self.completion_stream = completion_stream
self.model = model
self.stream_options = stream_options
self.custom_llm_provider = custom_llm_provider
self.usage_injected_by_proxy: Final = usage_injected_by_proxy
self.include_usage: Final = (
not usage_injected_by_proxy
and stream_options is not None
and stream_options.get("include_usage", False) is True
)
def __iter__(self):
return self
@ -23,7 +31,7 @@ class TextCompletionStreamWrapper:
def __aiter__(self):
return self
def convert_to_text_completion_object(self, chunk: ModelResponse):
def convert_to_text_completion_object(self, chunk: ModelResponse | ModelResponseStream):
try:
response: Final = TextCompletionResponse()
response["id"] = chunk.get("id", None)
@ -47,8 +55,7 @@ class TextCompletionStreamWrapper:
else:
response["choices"] = []
# only pass usage when stream_options["include_usage"] is True
if self.stream_options and self.stream_options.get("include_usage", False) is True:
if self.include_usage:
response.usage = chunk.get("usage", None)
else:
response.usage = None
@ -57,6 +64,9 @@ class TextCompletionStreamWrapper:
except Exception as e:
raise Exception(f"Error occurred converting to text completion object - chunk: {chunk}; Error: {e}")
def _is_injected_usage_artifact(self, chunk: ModelResponseStream) -> bool:
return self.usage_injected_by_proxy and is_injected_stream_usage_artifact(chunk)
def __next__(self):
# model_response = ModelResponse(stream=True, model=self.model)
TextCompletionResponse()
@ -64,8 +74,9 @@ class TextCompletionStreamWrapper:
for chunk in self.completion_stream:
if chunk == "None" or chunk is None:
raise Exception
processed_chunk = self.convert_to_text_completion_object(chunk=chunk)
return processed_chunk
if isinstance(chunk, ModelResponseStream) and self._is_injected_usage_artifact(chunk):
continue
return self.convert_to_text_completion_object(chunk=chunk)
raise StopIteration
except StopIteration:
raise StopIteration
@ -85,8 +96,9 @@ class TextCompletionStreamWrapper:
async for chunk in self.completion_stream:
if chunk == "None" or chunk is None:
raise Exception
processed_chunk = self.convert_to_text_completion_object(chunk=chunk)
return processed_chunk
if isinstance(chunk, ModelResponseStream) and self._is_injected_usage_artifact(chunk):
continue
return self.convert_to_text_completion_object(chunk=chunk)
raise StopIteration
except StopIteration:
raise StopAsyncIteration

View file

@ -204,9 +204,11 @@ class MCPServer(LiteLLMBaseModel):
access_groups: list[str] | None = None
allow_all_keys: bool = False
available_on_public_internet: bool = True
# Legacy opt-in to upstream-delegated authentication for ``oauth2``
# servers. LiteLLM admission still applies; use ``oauth_delegate`` for the
# supported client-forwarded OAuth flow.
# Deprecated opt-in to upstream-delegated authentication for ``oauth2``
# servers. A bare ``Authorization`` that is not a LiteLLM credential is the
# upstream token and is forwarded; see
# ``MCPRequestHandler._admit_legacy_delegate_bearer``. New servers should use
# ``true_passthrough`` or ``oauth_delegate``.
#
# Honored only for ``auth_type == oauth2``; ignored for any other
# ``auth_type``. OAuth pass-through for non-oauth2 servers

View file

@ -169,6 +169,10 @@ class UserAPIKeyAuth(LiteLLM_VerificationTokenView): # the expected response ob
# above; a forged value could at most narrow, but the stripping keeps the field's provenance
# single-owner so its meaning stays trustworthy.
mcp_session_resource_server_ids: tuple[str, ...] | None = Field(default=None, exclude=True)
# Set only by MCP admission when a legacy delegate_auth_to_upstream caller's Authorization is the
# upstream token rather than a LiteLLM credential, so egress forwards it. Stripped from validated
# input like the markers above: a forged value would forward a LiteLLM credential upstream.
mcp_forwards_caller_authorization: bool = Field(default=False, exclude=True)
mcp_toolset_id: str | None = Field(default=None, exclude=True)
authenticated_by_custom_auth: bool = Field(default=False, exclude=True)
via_virtual_key: bool = Field(
@ -230,6 +234,7 @@ class UserAPIKeyAuth(LiteLLM_VerificationTokenView): # the expected response ob
values.pop("mcp_explicit_grants_only", None)
values.pop("mcp_source_team_rpm_limits", None)
values.pop("mcp_session_resource_server_ids", None)
values.pop("mcp_forwards_caller_authorization", None)
values.pop("mcp_toolset_id", None)
values.pop("via_virtual_key", None)
values.pop("authenticated_by_custom_auth", None)

View file

@ -101,6 +101,23 @@ def _install_termination_handlers() -> None:
signal.signal(termination, _exit_on_termination)
def _install_base_native(worktree: Path, repo_root: Path) -> None:
env: Final = {
**os.environ,
"CARGO_TARGET_DIR": str(
Path(os.environ.get("CARGO_TARGET_DIR", repo_root / "litellm-rust" / "target")).resolve()
),
}
_run(
[
"uv", "pip", "install", "--python", sys.executable, "--no-deps",
"--target", str(worktree.parent / "native-install"), "--editable", str(worktree),
],
cwd=worktree,
env=env,
)
def base_counts(ref: str, repo_root: Path = REPO_ROOT, checker: Path = CHECKER) -> Mapping[str, int]:
"""Rule counts at `ref`, measured with the *current* checker over the base tree."""
_install_termination_handlers()
@ -111,6 +128,7 @@ def base_counts(ref: str, repo_root: Path = REPO_ROOT, checker: Path = CHECKER)
(worktree / "scripts").mkdir(parents=True, exist_ok=True)
base_checker: Final = worktree / "scripts" / "check_openapi_docs.py"
shutil.copy(checker, base_checker)
_install_base_native(worktree, repo_root)
return count_by_rule(_check(worktree, base_checker))
finally:
# Teardown must never raise, or it masks the real error when the body failed.

View file

@ -0,0 +1,9 @@
from typing import Final
import openapi_docs_gate as gate
def test_base_counts_imports_native_extension_from_detached_worktree() -> None:
expected_counts: Final = dict(gate.count_by_rule(gate.head_violations()))
assert expected_counts
assert dict(gate.base_counts("HEAD")) == expected_counts

View file

@ -0,0 +1,185 @@
import json
from pathlib import Path
from typing import Final
from uuid import uuid4
import pytest
from integration._support.client import Gateway, eventually
from integration._support.database import read_rows
from integration._support.process import owned_proxy
from integration._support.wire import Reply, Request, wire_server
from pydantic import JsonValue
PROMPT_TOKENS: Final = 19
COMPLETION_TOKENS: Final = 5
INPUT_RATE: Final = 0.001
OUTPUT_RATE: Final = 0.002
EXPECTED_SPEND: Final = PROMPT_TOKENS * INPUT_RATE + COMPLETION_TOKENS * OUTPUT_RATE
USAGE: Final[dict[str, JsonValue]] = {
"prompt_tokens": PROMPT_TOKENS,
"completion_tokens": COMPLETION_TOKENS,
"total_tokens": PROMPT_TOKENS + COMPLETION_TOKENS,
}
PROMPT: Final = "Say hi"
def _sse_frame(payload: dict[str, JsonValue]) -> bytes:
return f"data: {json.dumps(payload, separators=(',', ':'))}\n\n".encode()
def _chunk(request_id: str, model: str, delta: dict[str, JsonValue], finish_reason: str | None) -> dict[str, JsonValue]:
return {
"id": request_id,
"object": "chat.completion.chunk",
"created": 1,
"model": model,
"choices": [{"index": 0, "delta": delta, "finish_reason": finish_reason}],
}
def _respond(request: Request) -> Reply:
if request.method == "GET" and request.target == "/v1/models":
return Reply(body=json.dumps({"object": "list", "data": [{"id": "gpt-5.4-mini", "object": "model"}]}).encode())
assert request.method == "POST", request.method
assert request.target == "/v1/chat/completions", request.target
body: Final = json.loads(request.body)
assert body["messages"] == [{"role": "user", "content": PROMPT}], body
request_id: Final = f"chatcmpl-{uuid4().hex}"
model: Final = str(body["model"])
if body.get("stream") is not True:
return Reply(
body=json.dumps(
{
"id": request_id,
"object": "chat.completion",
"created": 1,
"model": model,
"choices": [
{"index": 0, "message": {"role": "assistant", "content": "Hi there"}, "finish_reason": "length"}
],
"usage": USAGE,
}
).encode()
)
stream_options: Final = body.get("stream_options") or {}
usage_frames: Final = (
(
_sse_frame(
{
"id": request_id,
"object": "chat.completion.chunk",
"created": 1,
"model": model,
"choices": [],
"usage": USAGE,
}
),
)
if stream_options.get("include_usage") is True
else ()
)
return Reply(
content_type="text/event-stream",
chunks=(
_sse_frame(_chunk(request_id, model, {"role": "assistant", "content": "Hi"}, None)),
_sse_frame(_chunk(request_id, model, {"content": " there"}, None)),
_sse_frame(_chunk(request_id, model, {}, "length")),
*usage_frames,
b"data: [DONE]\n\n",
),
)
def _proxy_config(directory: Path, model: str, upstream_url: str) -> Path:
config: Final = directory / "text_completion_stream_usage_cost_config.yaml"
config.write_text(
json.dumps(
{
"model_list": [
{
"model_name": model,
"litellm_params": {
"model": "openai/gpt-5.4-mini",
"api_base": f"{upstream_url}/v1",
"api_key": "scripted-upstream",
"input_cost_per_token": INPUT_RATE,
"output_cost_per_token": OUTPUT_RATE,
},
}
],
"general_settings": {
"master_key": "os.environ/LITELLM_MASTER_KEY",
"database_url": "os.environ/DATABASE_URL",
"store_model_in_db": True,
"disable_spend_logs": False,
"proxy_batch_write_at": 1,
"proxy_batch_polling_interval": 1,
},
"router_settings": {"disable_cooldowns": True},
}
)
)
return config
def _data_events(body: str) -> tuple[dict[str, JsonValue], ...]:
return tuple(
json.loads(line.removeprefix("data:"))
for line in body.splitlines()
if line.startswith("data:") and not line.endswith("[DONE]")
)
def _carries_text_or_finish(frame: dict[str, JsonValue]) -> bool:
choices: Final = frame["choices"]
assert isinstance(choices, list) and len(choices) == 1, frame
choice: Final = choices[0]
assert isinstance(choice, dict), frame
return choice.get("text") is not None or choice.get("finish_reason") is not None
def _billed(row: dict[str, object]) -> tuple[object, object, float]:
return (row["prompt_tokens"], row["completion_tokens"], float(str(row["spend"])))
def test_text_completion_stream_without_stream_options_bills_like_its_twins(gateway: Gateway, tmp_path: Path) -> None:
model: Final = f"gpt-5.4-mini-{uuid4().hex}"
body: Final[dict[str, JsonValue]] = {"model": model, "prompt": PROMPT, "max_tokens": COMPLETION_TOKENS}
with wire_server(_respond) as wire:
config: Final = _proxy_config(tmp_path, model, wire.url)
with owned_proxy(gateway, tmp_path, {}, config=config) as candidate:
silent: Final = candidate.request("POST", "/v1/completions", {**body, "stream": True})
asked: Final = candidate.request(
"POST", "/v1/completions", {**body, "stream": True, "stream_options": {"include_usage": True}}
)
whole: Final = candidate.request("POST", "/v1/completions", body)
assert (silent.status_code, asked.status_code, whole.status_code) == (200, 200, 200), (
silent.text,
asked.text,
whole.text,
)
silent_frames: Final = _data_events(silent.text)
asked_frames: Final = _data_events(asked.text)
request_ids: Final = (silent_frames[0]["id"], asked_frames[0]["id"], whole.json()["id"])
rows: Final = eventually(
lambda: read_rows(
'SELECT request_id, prompt_tokens, completion_tokens, spend FROM "LiteLLM_SpendLogs" '
"WHERE request_id = ANY(%s)",
(list(request_ids),),
database_url=candidate.database_url,
),
lambda values: len(values) == 3,
seconds=70,
)
upstream_stream_options: Final = [
json.loads(sent.body).get("stream_options") for sent in wire.drain() if sent.method == "POST"
]
assert upstream_stream_options == [{"include_usage": True}, {"include_usage": True}, None]
assert all(_carries_text_or_finish(frame) for frame in silent_frames), silent.text
assert all("usage" not in frame for frame in silent_frames), silent.text
assert asked_frames[-1]["usage"]["total_tokens"] == PROMPT_TOKENS + COMPLETION_TOKENS, asked.text
billed: Final = {str(row["request_id"]): _billed(row) for row in rows}
assert billed[str(request_ids[0])] == billed[str(request_ids[1])] == billed[str(request_ids[2])], billed
assert billed[str(request_ids[0])] == (PROMPT_TOKENS, COMPLETION_TOKENS, pytest.approx(EXPECTED_SPEND)), billed

View file

@ -442,6 +442,7 @@ class _Agentic(CustomLogger):
ADVISOR_TOOL: Final = {"type": ANTHROPIC_ADVISOR_TOOL_TYPE, "name": "advisor", "model": "claude-opus-5-5"}
STRICT_TOOL: Final = {"name": "a", "input_schema": {"type": "object"}, "strict": True}
@pytest.mark.parametrize(
@ -527,6 +528,45 @@ def test_each_detected_feature_reaches_its_rule(
assert served(kwargs) is python_result
@pytest.mark.asyncio
@pytest.mark.parametrize(
"kwargs",
(
pytest.param({"tools": [STRICT_TOOL]}, id="drop-params-unset"),
pytest.param({"drop_params": False, "tools": [STRICT_TOOL]}, id="drop-params-off"),
pytest.param(
{"drop_params": True, "tools": [{"name": "a", "input_schema": {"type": "object"}}]}, id="no-strict"
),
),
)
async def test_strict_tools_stay_native_unless_python_could_drop_strict(
monkeypatch: pytest.MonkeyPatch, kwargs: Mapping[str, object]
) -> None:
rules: Final[Rules] = (
RouteRule(Route.MESSAGES, Rollout.PYTHON_ONLY, features=frozenset({Feature.STRICT_TOOLS_RETRY})),
RouteRule(Route.MESSAGES, Rollout.RUST_REQUIRED),
)
native_result: Final = response()
monkeypatch.setattr(litellm, "drop_params", False)
async def python(*call_args: object, **python_kwargs: object) -> AnthropicMessagesResponse: # kwargs-ok: stub
pytest.fail("Python cannot drop strict here")
async def native(request: NativeCall) -> AnthropicMessagesResponse:
return native_result
served: Final = await _ADISPATCH.arun(
(16, MESSAGES, "anthropic/claude-sonnet-4-5"),
kwargs,
python=python,
binding=amessages_binding(native),
native=lambda hook, request, call_args, native_kwargs: hook(request),
rules=rules,
)
assert served is native_result
ASYNC_ONLY_FEATURES: Final = (
pytest.param(Feature.PRE_REQUEST_HOOK, {}, {"callbacks": [_Rewrite()]}, id="pre-request-hook"),
pytest.param(
@ -538,6 +578,12 @@ ASYNC_ONLY_FEATURES: Final = (
{},
id="messages-interceptor",
),
pytest.param(
Feature.STRICT_TOOLS_RETRY, {"drop_params": True, "tools": [STRICT_TOOL]}, {}, id="strict-tools-retry"
),
pytest.param(
Feature.STRICT_TOOLS_RETRY, {"tools": [STRICT_TOOL]}, {"drop_params": True}, id="strict-tools-retry-global"
),
)

View file

@ -3,6 +3,7 @@ import copy
import json
import os
import uuid
from dataclasses import dataclass
from typing import Any, Dict, Final, List
import httpx
@ -10,7 +11,7 @@ import responses
import respx
import pytest
from fastapi.testclient import TestClient
from pydantic import ValidationError
from pydantic import TypeAdapter, ValidationError
from unittest.mock import AsyncMock, MagicMock, patch
@ -2396,3 +2397,61 @@ def test_native_messages_tool_config_point_marks_last_tool_and_drops_the_key(mon
assert "cache_control" not in captured["params"]["tools"][0]
assert captured["messages"][0]["content"][-1]["cache_control"] == {"type": "ephemeral"}
assert tools == tools_before
# Anthropic's 400 for a request past Claude's strict tool limits, seen live on claude-sonnet-4-6 on 2026-10-29
_TOO_MANY_STRICT_TOOLS: Final = {
"type": "error",
"error": {
"type": "invalid_request_error",
"message": "Too many strict tools (21). The maximum number of strict tools supported is 20. Try reducing the "
"number of tools marked as strict.",
},
"request_id": "req_011CgWXzv8cmAr9vFs786WV4",
}
_JSON_BODY: Final = TypeAdapter(dict[str, object])
@dataclass(slots=True)
class _RejectsStrictToolsOnce:
sent: tuple[dict[str, object], ...] = ()
def __call__(self, request: httpx.Request) -> httpx.Response:
self.sent = (*self.sent, _JSON_BODY.validate_json(request.content))
if len(self.sent) == 1:
return httpx.Response(400, json=_TOO_MANY_STRICT_TOOLS, request=request)
return httpx.Response(
200,
json={
"id": "msg_strict_retry",
"type": "message",
"role": "assistant",
"model": "claude-sonnet-4-6",
"content": [{"type": "text", "text": "done"}],
"stop_reason": "end_turn",
"stop_sequence": None,
"usage": {"input_tokens": 1, "output_tokens": 1},
},
request=request,
)
@pytest.mark.asyncio
async def test_anthropic_messages_strict_tools_rejection_is_sent_again_without_strict_under_drop_params() -> None:
anthropic: Final = _RejectsStrictToolsOnce()
tool: Final = {"name": "lookup_record", "input_schema": {"type": "object"}, "strict": True}
response: Final = await litellm.anthropic_messages(
max_tokens=16,
messages=[{"role": "user", "content": "Use lookup_record to find invoice 42"}],
model="anthropic/claude-sonnet-4-6",
api_key="sk-test",
tools=[tool],
drop_params=True,
client=AsyncHTTPHandler(transport=httpx.MockTransport(anthropic)),
)
first, second = anthropic.sent
assert first["tools"] == [tool]
assert second == {**first, "tools": [{"name": "lookup_record", "input_schema": {"type": "object"}}]}
assert _JSON_BODY.validate_python(response)["id"] == "msg_strict_retry"

View file

@ -25,7 +25,7 @@ import pytest
import respx
import litellm
from litellm.llms.anthropic.common_utils import is_claude_tool_target
from litellm.llms.anthropic.common_utils import ToolsWithoutStrict, is_claude_tool_target, tools_without_strict
from litellm.llms.custom_httpx.http_handler import HTTPHandler
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../../..")))
@ -4559,3 +4559,27 @@ def test_shared_legacy_thinking_translation_drops_a_non_numeric_budget_under_dro
)
def test_is_claude_tool_target(custom_llm_provider: str, model: str, expected: bool):
assert is_claude_tool_target(custom_llm_provider=custom_llm_provider, model=model) is expected
@pytest.mark.parametrize(
"tools",
[
None,
"lookup_record",
[],
[{"name": "a"}],
[{"name": "a", "strict": False}],
[{"name": "a", "strict": True}, "b"],
],
ids=["absent", "not_a_list", "empty", "no_strict_key", "strict_false_only", "malformed_tool"],
)
def test_tools_without_strict_is_none_unless_every_tool_is_an_object_and_one_is_strict(tools: object) -> None:
assert tools_without_strict(tools) is None
def test_tools_without_strict_drops_only_the_strict_key_and_counts_the_strict_tools() -> None:
tools: Final = [{"name": "a", "strict": True}, {"name": "b"}, {"name": "c", "strict": False}]
assert tools_without_strict(tools) == ToolsWithoutStrict(
tools=({"name": "a"}, {"name": "b"}, {"name": "c"}), dropped_count=1
)

View file

@ -8,7 +8,7 @@ import zlib
from collections.abc import AsyncIterator, Mapping, Sequence
from dataclasses import dataclass
from datetime import datetime
from types import SimpleNamespace
from types import MappingProxyType, SimpleNamespace
from typing import Final
from unittest.mock import Mock
@ -4051,12 +4051,14 @@ _THINKING_HISTORY: Final = (
@dataclass(slots=True)
class ScriptedBedrockInvoke:
replies: tuple[tuple[int, Mapping[str, object]], ...]
replies: tuple[tuple[int, Mapping[str, object] | bytes], ...]
bodies: tuple[dict[str, object], ...] = ()
def __call__(self, request: httpx.Request) -> httpx.Response:
self.bodies = (*self.bodies, REASONING_BODY_ADAPTER.validate_json(request.content))
status_code, body = self.replies[len(self.bodies) - 1]
if isinstance(body, bytes):
return httpx.Response(status_code, content=body)
return httpx.Response(status_code, json=body)
@ -4065,6 +4067,7 @@ async def _send_bedrock_messages(
model_info: Mapping[str, object] | None,
replies: Sequence[tuple[int, Mapping[str, object]]],
messages: Sequence[Mapping[str, object]] = ({"role": "user", "content": "Summarise the weather as JSON."},),
extra: Mapping[str, object] = MappingProxyType({}),
) -> tuple[ScriptedBedrockInvoke, Exception | None]:
bedrock: Final = ScriptedBedrockInvoke(replies=tuple(replies))
try:
@ -4080,6 +4083,7 @@ async def _send_bedrock_messages(
aws_region_name="us-east-1",
client=AsyncHTTPHandler(transport=httpx.MockTransport(bedrock)),
**({} if model_info is None else {"model_info": model_info}),
**extra,
)
except Exception as error:
return bedrock, error
@ -4171,6 +4175,158 @@ async def test_bedrock_messages_400_is_not_retried_unless_it_rejects_opt_in_nati
assert len(bedrock.bodies) == 1
# Bedrock's 400 bodies for requests past Claude's strict tool limits, seen live on us.anthropic.claude-sonnet-4-6 on 2026-10-29
_STRICT_TOOLS_REJECTIONS: Final = (
{
"message": "Too many strict tools (21). The maximum number of strict tools supported is 20. Try reducing the "
"number of tools marked as strict."
},
{
"message": "Schemas contains too many optional parameters (96), which would make grammar compilation "
"inefficient. Reduce the number of optional parameters in your tool schemas (limit: 24)."
},
{
"message": "Schemas contains too many parameters with union types (24 parameters with type arrays or anyOf). "
"This causes exponential compilation cost. Reduce the number of nullable or union-typed parameters (limit: 16 "
"parameters with unions)."
},
{
"message": "The compiled grammar is too large, which would cause performance issues. Simplify your tool schemas "
"or reduce the number of strict tools."
},
)
_LOOKUP_TOOL: Final = {
"name": "lookup_record",
"description": "Look up records",
"input_schema": {
"type": "object",
"additionalProperties": False,
"required": ["query"],
"properties": {"query": {"type": "string"}},
},
}
_STRICT_TOOLS: Final = ({**_LOOKUP_TOOL, "strict": True}, {**_LOOKUP_TOOL, "name": "lookup_invoice", "strict": True})
_STREAMED_MESSAGE: Final = _bedrock_event_frame(
{"type": "message_start", "message": {**_ANTHROPIC_MESSAGE, "content": [], "stop_reason": None}}
) + _bedrock_event_frame({"type": "message_stop"})
async def _send_strict_tools_to_bedrock(
replies: Sequence[tuple[int, Mapping[str, object] | bytes]],
drop_params: bool | None,
tools: Sequence[Mapping[str, object]] = _STRICT_TOOLS,
stream: bool = False,
) -> tuple[ScriptedBedrockInvoke, object]:
bedrock: Final = ScriptedBedrockInvoke(replies=tuple(replies))
try:
response: Final = await litellm.anthropic_messages(
model="bedrock/invoke/us.anthropic.claude-sonnet-4-6",
messages=[{"role": "user", "content": "Use lookup_record to find invoice 42"}],
max_tokens=200,
tools=[dict(tool) for tool in tools],
stream=stream,
drop_params=drop_params,
aws_access_key_id="test-access-key",
aws_secret_access_key="test-secret-key",
aws_region_name="us-east-1",
client=AsyncHTTPHandler(transport=httpx.MockTransport(bedrock)),
)
except litellm.BadRequestError as error:
return bedrock, error
if isinstance(response, AsyncIterator):
return bedrock, b"".join([chunk async for chunk in response])
return bedrock, response
def _assert_resent_without_strict(bedrock: ScriptedBedrockInvoke) -> None:
first, second = bedrock.bodies
first_tools: Final = TypeAdapter(list[dict[str, object]]).validate_python(first["tools"])
assert [tool["strict"] for tool in first_tools] == [True, True]
assert second == {
**first,
"tools": [{key: value for key, value in tool.items() if key != "strict"} for tool in first_tools],
}
@pytest.mark.usefixtures("local_model_cost_map")
@pytest.mark.parametrize("rejection", _STRICT_TOOLS_REJECTIONS, ids=["count", "optional", "union", "grammar"])
async def test_bedrock_messages_strict_tools_rejection_is_sent_again_without_strict_under_drop_params(
rejection: Mapping[str, object],
) -> None:
bedrock, response = await _send_strict_tools_to_bedrock([(400, rejection), (200, _ANTHROPIC_MESSAGE)], True)
_assert_resent_without_strict(bedrock)
assert REASONING_BODY_ADAPTER.validate_python(response)["id"] == _ANTHROPIC_MESSAGE["id"]
@pytest.mark.usefixtures("local_model_cost_map")
async def test_bedrock_messages_strict_tools_rejection_is_sent_again_under_global_drop_params(
monkeypatch: pytest.MonkeyPatch,
) -> None:
monkeypatch.setattr(litellm, "drop_params", True)
bedrock, response = await _send_strict_tools_to_bedrock(
[(400, _STRICT_TOOLS_REJECTIONS[0]), (200, _ANTHROPIC_MESSAGE)], None
)
_assert_resent_without_strict(bedrock)
assert REASONING_BODY_ADAPTER.validate_python(response)["id"] == _ANTHROPIC_MESSAGE["id"]
@pytest.mark.usefixtures("local_model_cost_map")
async def test_bedrock_messages_streamed_strict_tools_rejection_is_sent_again_without_strict_under_drop_params() -> (
None
):
bedrock, streamed = await _send_strict_tools_to_bedrock(
[(400, _STRICT_TOOLS_REJECTIONS[0]), (200, _STREAMED_MESSAGE)], True, stream=True
)
_assert_resent_without_strict(bedrock)
assert isinstance(streamed, bytes) and b"event: message_stop" in streamed
@pytest.mark.usefixtures("local_model_cost_map")
async def test_bedrock_messages_body_rebuilt_without_native_output_still_drops_the_rejected_strict() -> None:
bedrock, error = await _send_bedrock_messages(
"us.anthropic.claude-opus-4-8",
None,
[(400, _STRICT_TOOLS_REJECTIONS[0]), (400, _NATIVE_OUTPUT_REJECTION), (200, _ANTHROPIC_MESSAGE)],
extra={"tools": [dict(tool) for tool in _STRICT_TOOLS], "drop_params": True},
)
assert error is None
first, _, rebuilt = bedrock.bodies
assert "format" in REASONING_BODY_ADAPTER.validate_python(first["output_config"])
assert "format" not in REASONING_BODY_ADAPTER.validate_python(rebuilt.get("output_config", {}))
assert rebuilt["tools"] == [
{key: value for key, value in tool.items() if key != "strict"} for tool in _STRICT_TOOLS
]
@pytest.mark.usefixtures("local_model_cost_map")
@pytest.mark.parametrize(
("drop_params", "tools", "rejection"),
[
(None, _STRICT_TOOLS, _STRICT_TOOLS_REJECTIONS[0]),
(False, _STRICT_TOOLS, _STRICT_TOOLS_REJECTIONS[0]),
(True, _STRICT_TOOLS, {"message": "messages.0.content: Field required"}),
(True, (_LOOKUP_TOOL,), _STRICT_TOOLS_REJECTIONS[0]),
(True, ({**_LOOKUP_TOOL, "strict": False},), _STRICT_TOOLS_REJECTIONS[1]),
],
ids=["drop_params_unset", "drop_params_off", "other_400", "no_strict_tool", "strict_false_only"],
)
async def test_bedrock_messages_400_reaches_the_caller_unless_drop_params_can_drop_the_rejected_strict(
drop_params: bool | None, tools: Sequence[Mapping[str, object]], rejection: Mapping[str, object]
) -> None:
bedrock, raised = await _send_strict_tools_to_bedrock(
[(400, rejection), (200, _ANTHROPIC_MESSAGE)], drop_params, tools
)
assert isinstance(raised, litellm.BadRequestError)
assert str(rejection["message"]) in str(raised)
assert len(bedrock.bodies) == 1
@pytest.mark.parametrize("trace", ["enabled", "disabled"])
def test_bedrock_messages_applies_the_deployment_guardrail_config_as_invoke_headers(trace: str) -> None:
from litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude_transformation import (

View file

@ -16,12 +16,16 @@ from unittest.mock import AsyncMock, MagicMock, patch
import httpx
import pytest
from mcp.types import CallToolResult, TextContent
import litellm
from litellm import ModelResponse
from litellm._version import version as litellm_version
from litellm.exceptions import GuardrailRaisedException, Timeout
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, MaskedHTTPStatusError
from litellm.proxy._experimental.mcp_server.guardrail_translation.handler import (
MCPGuardrailTranslationHandler,
)
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.guardrails.guardrail_hooks.generic_guardrail_api import (
GenericGuardrailAPI,
@ -29,6 +33,8 @@ from litellm.proxy.guardrails.guardrail_hooks.generic_guardrail_api import (
from litellm.proxy.guardrails.guardrail_hooks.generic_guardrail_api.generic_guardrail_api import (
_HEADER_PRESENT_PLACEHOLDER,
)
from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import UnifiedLLMGuardrails
from litellm.types.guardrails import GuardrailEventHooks
from litellm.types.utils import Choices, GenericGuardrailAPIInputs, Message
@ -2621,3 +2627,68 @@ class TestNumRetries:
assert sent.qsize() == num_retries + 1
assert result == {"texts": ["hi"]}
def _mcp_guardrail(mode: str) -> GenericGuardrailAPI:
return GenericGuardrailAPI(api_base="http://x", guardrail_name="g", event_hook=mode, default_on=True)
def _mock_api_response(payload: dict[str, object]) -> MagicMock:
response: Final = MagicMock()
response.json.return_value = payload
response.raise_for_status = MagicMock()
return response
@pytest.mark.parametrize("mode", ["pre_mcp_call", "during_mcp_call", "post_mcp_call"])
def test_mcp_call_modes_pass_event_hook_validation(mode: str) -> None:
guardrail: Final = _mcp_guardrail(mode)
assert guardrail.should_run_guardrail(data={}, event_type=GuardrailEventHooks(mode)) is True
@pytest.mark.asyncio
async def test_pre_mcp_call_blocked_raises() -> None:
guardrail: Final = _mcp_guardrail("pre_mcp_call")
response: Final = _mock_api_response({"action": "BLOCKED", "blocked_reason": "credential detected in tool input"})
with patch.object(guardrail.async_handler, "post", return_value=response):
with pytest.raises(GuardrailRaisedException) as exc_info:
await MCPGuardrailTranslationHandler().process_input_messages(
data={"mcp_tool_name": "t", "mcp_arguments": {"a": "secret"}},
guardrail_to_apply=guardrail,
)
assert exc_info.value.message == "credential detected in tool input"
@pytest.mark.asyncio
async def test_post_mcp_call_redacts_result() -> None:
guardrail: Final = _mcp_guardrail("post_mcp_call")
response: Final = _mock_api_response({"action": "GUARDRAIL_INTERVENED", "texts": ["x [R] y"]})
result: Final = CallToolResult(content=[TextContent(type="text", text="x secret y")], isError=False)
with patch.object(guardrail.async_handler, "post", return_value=response):
returned: Final = await MCPGuardrailTranslationHandler().process_output_response(
response=result,
guardrail_to_apply=guardrail,
request_data={"mcp_tool_name": "t"},
)
assert returned.content[0].text == "x [R] y"
@pytest.mark.asyncio
async def test_during_mcp_call_blocked_raises() -> None:
guardrail: Final = _mcp_guardrail("during_mcp_call")
response: Final = _mock_api_response({"action": "BLOCKED", "blocked_reason": "credential detected in tool input"})
with patch.object(guardrail.async_handler, "post", return_value=response):
with pytest.raises(GuardrailRaisedException) as exc_info:
await UnifiedLLMGuardrails().async_moderation_hook(
data={"mcp_tool_name": "t", "mcp_arguments": {"a": "secret"}, "guardrail_to_apply": guardrail},
user_api_key_dict=UserAPIKeyAuth(),
call_type="call_mcp_tool",
)
assert exc_info.value.message == "credential detected in tool input"

View file

@ -1,7 +1,7 @@
import contextlib
import json
import os
from collections.abc import Sequence
from collections.abc import Callable, Iterator, Sequence
from collections.abc import Set as AbstractSet
from datetime import datetime, timedelta, timezone
from types import SimpleNamespace
@ -1599,9 +1599,7 @@ class TestMCPRequestHandler:
team_id="test-team-id",
)
with patch(
"litellm.proxy.mcp_server.auth.user_api_key_auth_mcp.user_api_key_auth"
) as mock_user_api_key_auth:
with patch("litellm.proxy.mcp_server.auth.user_api_key_auth_mcp.user_api_key_auth") as mock_user_api_key_auth:
mock_user_api_key_auth.return_value = mock_auth_result
# Call the method
@ -1663,10 +1661,10 @@ class TestMCPOAuth2AuthFlow:
as LiteLLM API keys, causing auth failures and empty tool listings.
"""
async def test_oauth2_token_in_authorization_header_requires_litellm_admission(self):
async def test_oauth2_upstream_token_on_legacy_delegate_is_admitted_anonymously(self) -> None:
"""
A bare Authorization token on the legacy delegated mode must establish
a LiteLLM principal rather than entering anonymously.
A bare upstream OAuth token on the legacy delegated mode is admitted
anonymously without LiteLLM key validation and kept for egress.
"""
from litellm.types.mcp import MCPAuth
@ -1702,10 +1700,11 @@ class TestMCPOAuth2AuthFlow:
raw_headers,
) = await MCPRequestHandler.process_mcp_request(scope)
assert isinstance(auth_result, UserAPIKeyAuth)
assert auth_result.user_id == "admitted-user"
mock_auth.assert_awaited_once()
assert oauth2_headers.get("Authorization") == "Bearer atlassian-oauth2-access-token-xyz"
assert auth_result.user_id is None
assert auth_result.api_key is None
assert auth_result.mcp_forwards_caller_authorization is True
mock_auth.assert_not_awaited()
assert oauth2_headers == {"Authorization": "Bearer atlassian-oauth2-access-token-xyz"}
async def test_explicit_litellm_key_with_oauth2_authorization(self):
"""
@ -2545,9 +2544,10 @@ class TestMCPDelegateAuthToUpstream:
"""
Tests for the ``delegate_auth_to_upstream`` per-server flag.
The legacy flag no longer bypasses LiteLLM admission. OAuth discovery may
still use the anonymous cold-start challenge, but a presented bearer must
authenticate to LiteLLM unless a separate admission credential is supplied.
A bare bearer that is a LiteLLM credential takes normal admission and is
never forwarded. Any other bearer is the upstream token: it is forwarded,
and the caller is admitted as the JWT or custom-auth user it also proves
when that user is granted the server, else anonymously.
"""
@staticmethod
@ -2652,37 +2652,497 @@ class TestMCPDelegateAuthToUpstream:
assert isinstance(auth_result, UserAPIKeyAuth)
mock_auth.assert_awaited_once()
async def test_delegate_with_only_upstream_token_requires_litellm_auth(self):
"""
An upstream token cannot establish a LiteLLM principal and must not
reopen anonymous admission.
"""
from litellm.types.mcp import MCPAuth
scope = {
@staticmethod
def _bearer_scope(token: str) -> dict[str, object]:
return {
"type": "http",
"method": "POST",
"path": "/mcp/delegated_oauth_server",
"headers": [(b"authorization", b"Bearer upstream-pkce-token")],
"headers": [(b"authorization", f"Bearer {token}".encode())],
}
@staticmethod
def _jwt(issuer: str = "https://idp.example") -> str:
import jwt
return jwt.encode({"iss": issuer, "sub": "jwt-user"}, "test-signing-secret", algorithm="HS256")
@staticmethod
@contextlib.contextmanager
def _delegate_admission(
admission: AsyncMock, allowed_server_ids: Sequence[str] = ("test-server-id",)
) -> Iterator[None]:
from litellm.types.mcp import MCPAuth
with (
patch(
"litellm.proxy.mcp_server.auth.user_api_key_auth_mcp.user_api_key_auth",
new_callable=AsyncMock,
side_effect=HTTPException(status_code=401, detail="Invalid API key"),
) as mock_auth,
patch("litellm.proxy.mcp_server.auth.user_api_key_auth_mcp.user_api_key_auth", admission),
patch("litellm.proxy.mcp_server.mcp_server_manager.global_mcp_server_manager") as mock_mgr,
):
mock_mgr.get_mcp_server_by_name.return_value = TestMCPDelegateAuthToUpstream._make_server(
auth_type=MCPAuth.oauth2,
delegate_auth_to_upstream=True,
)
with pytest.raises(HTTPException) as exc_info:
await MCPRequestHandler.process_mcp_request(scope)
mock_mgr.get_allowed_mcp_servers = AsyncMock(return_value=list(allowed_server_ids))
yield
assert exc_info.value.status_code == 401
mock_auth.assert_awaited_once()
@staticmethod
@contextlib.contextmanager
def _jwt_auth_enabled(auth_jwt: AsyncMock) -> Iterator[None]:
from litellm.proxy import proxy_server
from litellm.proxy._types import LiteLLM_JWTAuth
with (
patch.dict(proxy_server.general_settings, {"enable_jwt_auth": True}),
patch.object(proxy_server, "premium_user", True),
patch.object(proxy_server, "user_custom_auth", None),
patch.object(proxy_server.jwt_handler, "auth_jwt", auth_jwt),
patch.object(proxy_server.jwt_handler, "litellm_jwtauth", LiteLLM_JWTAuth(), create=True),
patch.dict(os.environ),
):
os.environ.pop("JWT_ISSUER", None)
yield
async def test_delegate_forwards_unknown_sk_token_after_key_table_miss(self) -> None:
"""An sk- token the key table does not hold, such as a provider API key, belongs to the upstream:
it is forwarded anonymously rather than rejected for its prefix."""
from litellm.proxy.auth.resolvers.exceptions import KeyNotFoundError
admission: Final = AsyncMock()
with (
TestMCPDelegateAuthToUpstream._delegate_admission(admission),
patch(
"litellm.proxy.auth.resolvers.store.IdentityStore.resolve",
new_callable=AsyncMock,
side_effect=KeyNotFoundError("hashed-token"),
),
):
auth_result, _, _, _, oauth2_headers, _ = await MCPRequestHandler.process_mcp_request(
TestMCPDelegateAuthToUpstream._bearer_scope("sk-proj-upstream-provider-key")
)
assert auth_result.api_key is None
assert auth_result.mcp_forwards_caller_authorization is True
assert oauth2_headers == {"Authorization": "Bearer sk-proj-upstream-provider-key"}
admission.assert_not_awaited()
@pytest.mark.parametrize(
"lookup_effect",
[None, RuntimeError("key table unreachable")],
ids=["key-held", "lookup-fault"],
)
async def test_delegate_litellm_key_refusal_propagates(self, lookup_effect: Exception | None) -> None:
"""A key the table holds, or one the lookup could not check, takes normal admission, so its
refusal reaches the caller instead of becoming an anonymous request that forwards the key."""
refusal: Final = ProxyException(
message="Budget has been exceeded", type="budget_exceeded", param=None, code=400
)
with (
TestMCPDelegateAuthToUpstream._delegate_admission(AsyncMock(side_effect=refusal)),
patch(
"litellm.proxy.auth.resolvers.store.IdentityStore.resolve",
new_callable=AsyncMock,
side_effect=lookup_effect,
),
):
with pytest.raises(ProxyException) as exc_info:
await MCPRequestHandler.process_mcp_request(TestMCPDelegateAuthToUpstream._bearer_scope("sk-held-key"))
assert exc_info.value is refusal
async def test_delegate_master_key_is_admitted_and_never_forwarded(self) -> None:
"""The master key is not in the key table, so only recognizing it directly stops it from being
treated as the upstream's token and sent to the upstream server."""
from litellm.proxy import proxy_server
from litellm.proxy.auth.resolvers.exceptions import KeyNotFoundError
admission: Final = AsyncMock(return_value=UserAPIKeyAuth(api_key="sk-master-secret", user_role="proxy_admin"))
with (
patch.object(proxy_server, "master_key", "sk-master-secret"),
TestMCPDelegateAuthToUpstream._delegate_admission(admission),
patch(
"litellm.proxy.auth.resolvers.store.IdentityStore.resolve",
new_callable=AsyncMock,
side_effect=KeyNotFoundError("hashed-master-key"),
),
):
auth_result, _, _, _, oauth2_headers, _ = await MCPRequestHandler.process_mcp_request(
TestMCPDelegateAuthToUpstream._bearer_scope("sk-master-secret")
)
admission.assert_awaited_once()
assert auth_result.mcp_forwards_caller_authorization is False
assert oauth2_headers is None
@pytest.mark.parametrize(
"signing_secret, expires_in, forwarded",
[
("sk-master-secret", timedelta(hours=1), False),
("sk-master-secret", timedelta(hours=-1), False),
("upstream-idp-signing-secret", timedelta(hours=1), True),
],
ids=["master-key-signed", "master-key-signed-expired", "signed-elsewhere"],
)
async def test_delegate_master_key_signed_jwt_is_never_forwarded(
self, signing_secret: str, expires_in: timedelta, forwarded: bool
) -> None:
"""LiteLLM signs dashboard and OAuth tokens with the master key, so such a JWT is LiteLLM's own even
once expired: it takes normal admission and never reaches the upstream. A JWT signed by anyone
else is the upstream's token."""
import jwt
from litellm.proxy import proxy_server
token: Final = jwt.encode(
{"sub": "dashboard-user", "exp": datetime.now(timezone.utc) + expires_in}, signing_secret, algorithm="HS256"
)
admission: Final = AsyncMock(return_value=UserAPIKeyAuth(user_id="dashboard-user"))
with (
patch.object(proxy_server, "master_key", "sk-master-secret"),
TestMCPDelegateAuthToUpstream._delegate_admission(admission),
):
auth_result, *_ = await MCPRequestHandler.process_mcp_request(
TestMCPDelegateAuthToUpstream._bearer_scope(token)
)
assert auth_result.mcp_forwards_caller_authorization is forwarded
assert admission.await_count == (0 if forwarded else 1)
@pytest.mark.parametrize(
"token, ui_token_lookup",
[
("llm_env_sealed-bridge-envelope", None),
("llm_refresh_sealed-refresh-envelope", None),
("v2:gcm:sealed-credential-blob", None),
("dashboard-login-token", lambda: MagicMock(return_value=UserAPIKeyAuth(user_id="dashboard-user"))),
("unreadable-dashboard-token", lambda: MagicMock(side_effect=ValueError("Invalid hash key"))),
],
ids=["bridge-envelope", "refresh-envelope", "gcm-blob", "dashboard-token", "unreadable-dashboard-token"],
)
async def test_delegate_sealed_litellm_token_takes_normal_admission(
self, token: str, ui_token_lookup: Callable[[], MagicMock] | None
) -> None:
"""A token only LiteLLM issues is LiteLLM's own, even one it fails to read, so it goes through
normal admission and is never marked for forwarding upstream."""
from litellm.proxy.auth.auth_checks import ExperimentalUIJWTToken
admission: Final = AsyncMock(return_value=UserAPIKeyAuth(user_id="litellm-user"))
ui_lookup: Final = (
patch.object(ExperimentalUIJWTToken, "get_key_object_from_ui_hash_key", ui_token_lookup())
if ui_token_lookup is not None
else contextlib.nullcontext()
)
with TestMCPDelegateAuthToUpstream._delegate_admission(admission), ui_lookup:
auth_result, *_ = await MCPRequestHandler.process_mcp_request(
TestMCPDelegateAuthToUpstream._bearer_scope(token)
)
admission.assert_awaited_once()
assert auth_result.mcp_forwards_caller_authorization is False
@staticmethod
def _unknown_user_refusal() -> ProxyException:
from litellm.types.proxy.auth.auth_checks import UserNotFoundError
refusal: Final = ProxyException(message="User doesn't exist in db", type="auth_error", param=None, code=401)
refusal.__cause__ = UserNotFoundError(user_id="jwt-user")
return refusal
@pytest.mark.parametrize(
"auth_jwt_effect, admission, allowed_server_ids, expected_user_id",
[
(
None,
lambda: AsyncMock(return_value=UserAPIKeyAuth(user_id="jwt-user", jwt_claims={"sub": "jwt-user"})),
["test-server-id"],
"jwt-user",
),
(
None,
lambda: AsyncMock(return_value=UserAPIKeyAuth(user_id="jwt-user", jwt_claims={"sub": "jwt-user"})),
[],
None,
),
(
None,
lambda: AsyncMock(side_effect=TestMCPDelegateAuthToUpstream._unknown_user_refusal()),
["test-server-id"],
None,
),
(
ValueError("Signature verification failed"),
lambda: AsyncMock(return_value=UserAPIKeyAuth(user_id="jwt-user")),
["test-server-id"],
None,
),
],
ids=["granted-user", "user-not-granted", "user-unknown-to-litellm", "signature-rejected"],
)
async def test_delegate_jwt_is_forwarded_as_granted_user_or_anonymously(
self,
auth_jwt_effect: Exception | None,
admission: Callable[[], AsyncMock],
allowed_server_ids: Sequence[str],
expected_user_id: str | None,
) -> None:
"""A JWT LiteLLM also accepts admits that user when granted the server; a user who is not
granted, unknown to LiteLLM, or whose signature fails is admitted anonymously. The JWT is always
forwarded, because on this server it is the upstream's token."""
token: Final = TestMCPDelegateAuthToUpstream._jwt()
with (
TestMCPDelegateAuthToUpstream._jwt_auth_enabled(AsyncMock(side_effect=auth_jwt_effect)),
TestMCPDelegateAuthToUpstream._delegate_admission(admission(), allowed_server_ids),
):
auth_result, _, _, _, oauth2_headers, _ = await MCPRequestHandler.process_mcp_request(
TestMCPDelegateAuthToUpstream._bearer_scope(token)
)
assert auth_result.user_id == expected_user_id
assert auth_result.mcp_forwards_caller_authorization is True
assert oauth2_headers == {"Authorization": f"Bearer {token}"}
@pytest.mark.parametrize(
"refusal",
[
ProxyException(message="Budget has been exceeded", type="budget_exceeded", param=None, code=400),
ProxyException(message="Team=t1 is blocked", type="auth_error", param=None, code=401),
],
ids=["over-budget", "blocked-team"],
)
async def test_delegate_jwt_refusal_after_valid_signature_propagates(self, refusal: ProxyException) -> None:
"""Once the signature verifies, LiteLLM knows who the caller is, so a budget or an explicit
disallowment is enforced and the request fails instead of continuing anonymously."""
with (
TestMCPDelegateAuthToUpstream._jwt_auth_enabled(AsyncMock(return_value={"sub": "jwt-user"})),
TestMCPDelegateAuthToUpstream._delegate_admission(AsyncMock(side_effect=refusal)),
):
with pytest.raises(ProxyException) as exc_info:
await MCPRequestHandler.process_mcp_request(
TestMCPDelegateAuthToUpstream._bearer_scope(TestMCPDelegateAuthToUpstream._jwt())
)
assert exc_info.value is refusal
@staticmethod
@contextlib.contextmanager
def _idp_auth_enabled(
mechanism: Literal["oauth2", "oauth2_proxy"], verify: AsyncMock, premium: bool = True
) -> Iterator[None]:
from litellm.proxy import proxy_server
from litellm.proxy.auth.oauth2_check import Oauth2Handler
settings: Final = {"enable_oauth2_auth": True} if mechanism == "oauth2" else {"enable_oauth2_proxy_auth": True}
verifier: Final = (
patch.object(Oauth2Handler, "check_oauth2_token", verify)
if mechanism == "oauth2"
else patch("litellm.proxy.auth.oauth2_proxy_hook.handle_oauth2_proxy_request", verify)
)
with (
patch.dict(proxy_server.general_settings, settings),
patch.object(proxy_server, "premium_user", premium),
patch.object(proxy_server, "user_custom_auth", None),
verifier,
):
yield
@pytest.mark.parametrize(
"verify_effect, allowed_server_ids, expected_user_id",
[
(None, ["test-server-id"], "idp-user"),
(None, [], None),
(ValueError("Missing trusted proxy header"), ["test-server-id"], None),
],
ids=["granted-user", "user-not-granted", "identity-rejected"],
)
async def test_delegate_proxy_header_identity_follows_the_jwt_rules(
self,
verify_effect: Exception | None,
allowed_server_ids: Sequence[str],
expected_user_id: str | None,
) -> None:
"""Trusted reverse-proxy headers admit a granted user by name, and a user who is not granted or not
vouched for anonymously. The bearer is forwarded either way."""
admission: Final = AsyncMock(return_value=UserAPIKeyAuth(user_id="idp-user"))
with (
TestMCPDelegateAuthToUpstream._idp_auth_enabled("oauth2_proxy", AsyncMock(side_effect=verify_effect)),
TestMCPDelegateAuthToUpstream._delegate_admission(admission, allowed_server_ids),
):
auth_result, _, _, _, oauth2_headers, _ = await MCPRequestHandler.process_mcp_request(
TestMCPDelegateAuthToUpstream._bearer_scope("opaque-idp-token")
)
assert auth_result.user_id == expected_user_id
assert auth_result.mcp_forwards_caller_authorization is True
assert oauth2_headers == {"Authorization": "Bearer opaque-idp-token"}
@pytest.mark.parametrize(
"provider_status, active, allowed_server_ids, expected_user_id",
[
(200, True, ["test-server-id"], "idp-user"),
(200, True, [], None),
(200, False, ["test-server-id"], None),
(401, True, ["test-server-id"], None),
],
ids=["granted-user", "user-not-granted", "token-inactive", "provider-rejects-token"],
)
async def test_delegate_oauth2_introspects_once_and_follows_the_jwt_rules(
self, provider_status: int, active: bool, allowed_server_ids: Sequence[str], expected_user_id: str | None
) -> None:
"""The provider is asked about the token once per request. A token it vouches for admits a granted
user by name, and one it reports inactive or rejects outright is forwarded anonymously instead of
failing the request."""
import httpx
from litellm.proxy import proxy_server
from litellm.types.mcp import MCPAuth
async def introspect(outbound: httpx.Request, **kwargs: object) -> httpx.Response:
return httpx.Response(provider_status, request=outbound, json={"active": active, "sub": "idp-user"})
with (
patch.dict(proxy_server.general_settings, {"enable_oauth2_auth": True}),
patch.object(proxy_server, "premium_user", True),
patch.object(proxy_server, "user_custom_auth", None),
patch.object(proxy_server, "master_key", "sk-fixture-master"),
patch.object(proxy_server, "prisma_client", None),
patch.object(proxy_server, "user_api_key_cache", DualCache()),
patch.dict(
os.environ,
{
"OAUTH_TOKEN_INFO_ENDPOINT": "https://idp.fixture.test/introspect",
"OAUTH_CLIENT_ID": "fixture-client",
"OAUTH_CLIENT_SECRET": "fixture-client-secret",
},
),
patch("httpx.AsyncClient.send", side_effect=introspect) as network,
patch("litellm.proxy.mcp_server.mcp_server_manager.global_mcp_server_manager") as mock_mgr,
):
mock_mgr.get_mcp_server_by_name.return_value = TestMCPDelegateAuthToUpstream._make_server(
auth_type=MCPAuth.oauth2, delegate_auth_to_upstream=True
)
mock_mgr.get_allowed_mcp_servers = AsyncMock(return_value=list(allowed_server_ids))
auth_result, _, _, _, oauth2_headers, _ = await MCPRequestHandler.process_mcp_request(
TestMCPDelegateAuthToUpstream._bearer_scope("opaque-idp-token")
)
assert network.call_count == 1
assert auth_result.user_id == expected_user_id
assert auth_result.mcp_forwards_caller_authorization is True
assert oauth2_headers == {"Authorization": "Bearer opaque-idp-token"}
async def test_delegate_oauth2_rejection_admits_anonymously_whether_or_not_auth_wraps_it(self) -> None:
"""A provider rejection is the provider declining to vouch, so the bearer is forwarded anonymously
even when it reaches admission as the bare validation error rather than a wrapped auth error."""
from litellm.proxy.auth.oauth2_check import OAuth2TokenValidationError
with (
TestMCPDelegateAuthToUpstream._idp_auth_enabled("oauth2", AsyncMock()),
TestMCPDelegateAuthToUpstream._delegate_admission(
AsyncMock(side_effect=OAuth2TokenValidationError("Token is not active"))
),
):
auth_result, _, _, _, oauth2_headers, _ = await MCPRequestHandler.process_mcp_request(
TestMCPDelegateAuthToUpstream._bearer_scope("opaque-idp-token")
)
assert auth_result.user_id is None
assert auth_result.mcp_forwards_caller_authorization is True
assert oauth2_headers == {"Authorization": "Bearer opaque-idp-token"}
@pytest.mark.parametrize("mechanism", ["oauth2", "oauth2_proxy"])
async def test_delegate_idp_user_refusal_propagates(self, mechanism: Literal["oauth2", "oauth2_proxy"]) -> None:
"""Once the provider vouches for the caller, LiteLLM's budget and permission refusals apply."""
refusal: Final = ProxyException(
message="Budget has been exceeded", type="budget_exceeded", param=None, code=400
)
with (
TestMCPDelegateAuthToUpstream._idp_auth_enabled(mechanism, AsyncMock()),
TestMCPDelegateAuthToUpstream._delegate_admission(AsyncMock(side_effect=refusal)),
):
with pytest.raises(ProxyException) as exc_info:
await MCPRequestHandler.process_mcp_request(
TestMCPDelegateAuthToUpstream._bearer_scope("opaque-idp-token")
)
assert exc_info.value is refusal
async def test_delegate_oauth2_without_premium_never_calls_the_provider(self) -> None:
"""OAuth2 introspection is enterprise-only, so without a license the bearer is the upstream's alone."""
introspect: Final = AsyncMock()
with (
TestMCPDelegateAuthToUpstream._idp_auth_enabled("oauth2", introspect, premium=False),
TestMCPDelegateAuthToUpstream._delegate_admission(
AsyncMock(return_value=UserAPIKeyAuth(user_id="idp-user"))
),
):
auth_result, *_ = await MCPRequestHandler.process_mcp_request(
TestMCPDelegateAuthToUpstream._bearer_scope("opaque-idp-token")
)
assert auth_result.user_id is None
assert auth_result.mcp_forwards_caller_authorization is True
introspect.assert_not_awaited()
async def test_delegate_jwt_from_unconfigured_issuer_is_forwarded_without_jwt_auth(self) -> None:
"""A JWT from an issuer LiteLLM's JWT auth is not configured for never reaches JWT auth."""
auth_jwt: Final = AsyncMock(return_value={"sub": "jwt-user"})
admission: Final = AsyncMock(return_value=UserAPIKeyAuth(user_id="jwt-user"))
with (
TestMCPDelegateAuthToUpstream._jwt_auth_enabled(auth_jwt),
patch.dict(os.environ, {"JWT_ISSUER": "https://litellm-idp.example"}),
TestMCPDelegateAuthToUpstream._delegate_admission(admission),
):
auth_result, *_ = await MCPRequestHandler.process_mcp_request(
TestMCPDelegateAuthToUpstream._bearer_scope(
TestMCPDelegateAuthToUpstream._jwt(issuer="https://upstream-idp.example")
)
)
assert auth_result.user_id is None
assert auth_result.mcp_forwards_caller_authorization is True
auth_jwt.assert_not_awaited()
@pytest.mark.parametrize(
"admission, expected_user_id",
[
(
lambda: AsyncMock(return_value=UserAPIKeyAuth(user_id="custom-user", api_key="upstream-custom-token")),
"custom-user",
),
(
lambda: AsyncMock(
side_effect=ProxyException(message="custom auth rejected", type="auth_error", param=None, code=401)
),
None,
),
],
ids=["accepted", "rejected"],
)
async def test_delegate_custom_auth_decides_alone(
self, admission: Callable[[], AsyncMock], expected_user_id: str | None
) -> None:
"""With custom auth configured, its answer alone decides: accepted admits that user, any failure
admits anonymously. The bearer is forwarded either way."""
from litellm.proxy import proxy_server
with (
patch.object(proxy_server, "user_custom_auth", AsyncMock()),
TestMCPDelegateAuthToUpstream._delegate_admission(admission()),
):
auth_result, _, _, _, oauth2_headers, _ = await MCPRequestHandler.process_mcp_request(
TestMCPDelegateAuthToUpstream._bearer_scope("upstream-custom-token")
)
assert auth_result.user_id == expected_user_id
assert auth_result.mcp_forwards_caller_authorization is True
assert oauth2_headers == {"Authorization": "Bearer upstream-custom-token"}
def test_forwarding_marker_cannot_be_forged_from_input(self) -> None:
"""Only admission may mark a caller's Authorization for forwarding; caller-shaped input never can."""
assert UserAPIKeyAuth(mcp_forwards_caller_authorization=True).mcp_forwards_caller_authorization is False
assert (
UserAPIKeyAuth.model_validate({"mcp_forwards_caller_authorization": True}).mcp_forwards_caller_authorization
is False
)
async def test_delegate_off_still_requires_litellm_auth(self):
"""
@ -2753,7 +3213,12 @@ class TestMCPDelegateAuthToUpstream:
await MCPRequestHandler.process_mcp_request(scope)
assert exc_info.value.status_code == 401
async def test_delegate_mixed_targets_fail_closed(self):
@pytest.mark.parametrize(
"bearer_headers",
[[], [(b"authorization", b"Bearer upstream-pkce-token")]],
ids=["no-bearer", "upstream-bearer"],
)
async def test_delegate_mixed_targets_fail_closed(self, bearer_headers: list[tuple[bytes, bytes]]) -> None:
"""
x-mcp-servers can list multiple targets. If ANY of them does not opt in
to delegate_auth_to_upstream, the bypass must NOT fire — otherwise an
@ -2769,6 +3234,7 @@ class TestMCPDelegateAuthToUpstream:
"path": "/mcp",
"headers": [
(b"x-mcp-servers", b"delegated_oauth,plain_oauth"),
*bearer_headers,
],
}
@ -2863,9 +3329,8 @@ class TestMCPDelegateAuthToUpstream:
async def test_authorization_bearer_on_delegate_server_establishes_litellm_principal(self):
"""
A bare Authorization bearer now follows normal LiteLLM admission. A
separate x-litellm-api-key is required when Authorization is intended
for the upstream server.
A bare sk- bearer the key table may hold takes normal LiteLLM admission
and is never marked for forwarding upstream.
"""
from litellm.types.mcp import MCPAuth
@ -2898,6 +3363,7 @@ class TestMCPDelegateAuthToUpstream:
) = await MCPRequestHandler.process_mcp_request(scope)
assert isinstance(auth_result, UserAPIKeyAuth)
assert auth_result.user_id == "real-user"
assert auth_result.mcp_forwards_caller_authorization is False
assert oauth2_headers.get("Authorization") == "Bearer sk-9876"
mock_auth.assert_awaited_once()
@ -3079,8 +3545,9 @@ class TestMCPDelegateAuthToUpstream:
mock_auth.assert_awaited_once()
assert auth.api_key is None
async def test_get_allowed_servers_excludes_legacy_delegates(self):
"""Legacy delegated servers are never added to anonymous access."""
async def test_get_allowed_servers_includes_interactive_legacy_delegates_for_anonymous(self) -> None:
"""Anonymous callers reach interactive legacy delegates, never an M2M server whose stored
client credentials would let them act as LiteLLM's service account."""
from litellm.proxy.mcp_server.mcp_server_manager import (
MCPServerManager,
)
@ -3118,7 +3585,7 @@ class TestMCPDelegateAuthToUpstream:
):
result = await manager.get_allowed_mcp_servers(None)
assert result == []
assert result == ["pkce-server"]
async def test_true_passthrough_skips_litellm_auth_anonymously(self):
"""auth_type=true_passthrough performs no admission auth: the caller's Authorization is an
@ -7644,9 +8111,7 @@ class TestMCPDcrBridgeDualCredential:
new_callable=AsyncMock,
return_value=UserAPIKeyAuth(user_id=presented_user, api_key="any-hash"),
),
patch(
"litellm.proxy.mcp_server.mcp_server_manager.global_mcp_server_manager"
) as mock_mgr,
patch("litellm.proxy.mcp_server.mcp_server_manager.global_mcp_server_manager") as mock_mgr,
patch("litellm.proxy.proxy_server.master_key", self._MASTER_KEY),
):
mock_mgr.get_mcp_server_by_name.return_value = self._DELEGATE._bridge_delegate_server()
@ -7697,7 +8162,9 @@ class TestMCPDcrBridgeDualCredential:
@pytest.mark.parametrize("root_path", ["", "/litellm"])
@pytest.mark.parametrize("invalid_token", [False, True])
def test_gateway_challenge_metadata_contains_the_public_resource(
public_path: str, root_path: str, invalid_token: bool,
public_path: str,
root_path: str,
invalid_token: bool,
) -> None:
from mcp.shared.auth_utils import check_resource_allowed
from starlette.requests import Request
@ -7710,8 +8177,12 @@ def test_gateway_challenge_metadata_contains_the_public_resource(
)
scope: Final[Scope] = {
"type": "http", "method": "POST", "scheme": "https", "path": "/mcp",
"root_path": root_path, "_original_path": f"{root_path}{public_path}",
"type": "http",
"method": "POST",
"scheme": "https",
"path": "/mcp",
"root_path": root_path,
"_original_path": f"{root_path}{public_path}",
"headers": [(b"host", b"gateway.example.com")],
}
request: Final = Request(scope)
@ -7720,8 +8191,7 @@ def test_gateway_challenge_metadata_contains_the_public_resource(
metadata_url: Final = header.split('resource_metadata="', 1)[1].split('"', 1)[0]
toolset: Final = public_path in ("/toolset/analysis/mcp", "/analysis/mcp")
assert metadata_url == (
f"https://gateway.example.com{root_path}/.well-known/oauth-protected-resource"
+ ("" if toolset else "/mcp")
f"https://gateway.example.com{root_path}/.well-known/oauth-protected-resource" + ("" if toolset else "/mcp")
)
metadata: Final = (
oauth_protected_resource_root(request) if toolset else _build_aggregate_protected_resource_response(request)
@ -7924,9 +8394,7 @@ class TestAggregateGatewayDcrChallenge:
):
with (
patch(self._AUTH_PATCH_TARGET, side_effect=self._auth_401()),
patch(
"litellm.proxy.mcp_server.mcp_server_manager.global_mcp_server_manager"
) as mock_mgr,
patch("litellm.proxy.mcp_server.mcp_server_manager.global_mcp_server_manager") as mock_mgr,
):
mock_mgr.get_mcp_server_by_name.return_value = server
with pytest.raises(HTTPException) as exc_info:
@ -7972,9 +8440,7 @@ class TestAggregateGatewayDcrChallenge:
with (
patch.dict(os.environ, {"SERVER_ROOT_PATH": "/litellm"}),
patch(self._AUTH_PATCH_TARGET, side_effect=self._auth_401()),
patch(
"litellm.proxy.mcp_server.mcp_server_manager.global_mcp_server_manager"
) as mock_mgr,
patch("litellm.proxy.mcp_server.mcp_server_manager.global_mcp_server_manager") as mock_mgr,
):
mock_mgr.get_mcp_server_by_name.return_value = server
with pytest.raises(HTTPException) as exc_info:
@ -8005,9 +8471,7 @@ class TestAggregateGatewayDcrChallenge:
):
with (
patch(self._AUTH_PATCH_TARGET, side_effect=self._auth_401()),
patch(
"litellm.proxy.mcp_server.mcp_server_manager.global_mcp_server_manager"
) as mock_mgr,
patch("litellm.proxy.mcp_server.mcp_server_manager.global_mcp_server_manager") as mock_mgr,
):
mock_mgr.get_mcp_server_by_name.return_value = resolved
mock_mgr.get_mcp_server_by_identifier.return_value = resolved
@ -8052,9 +8516,7 @@ class TestAggregateGatewayDcrChallenge:
(None, None),
]
for resolved, expected in cases:
with patch(
"litellm.proxy.mcp_server.mcp_server_manager.global_mcp_server_manager"
) as mock_mgr:
with patch("litellm.proxy.mcp_server.mcp_server_manager.global_mcp_server_manager") as mock_mgr:
mock_mgr.get_mcp_server_by_identifier.return_value = resolved
assert _gateway_dcr_challenge_target("/mcp/srv", None, None) == expected, resolved
assert _gateway_dcr_challenge_target("/mcp/a,b", None, None) is None
@ -10950,14 +11412,15 @@ async def test_catalog_refresh_uses_current_virtual_key_policy_and_keeps_session
from litellm.proxy import proxy_server
from litellm.types.mcp_server.mcp_toolset import MCPToolset
permission = LiteLLM_ObjectPermissionTable(
object_permission_id="current-policy", mcp_toolsets=["session-toolset"]
)
permission = LiteLLM_ObjectPermissionTable(object_permission_id="current-policy", mcp_toolsets=["session-toolset"])
toolset: Final = MCPToolset(
toolset_id="session-toolset", toolset_name="session-toolset",
toolset_id="session-toolset",
toolset_name="session-toolset",
tools=[{"server_id": "current-server", "tool_name": "read"}],
)
tables: Final = SimpleNamespace(litellm_mcptoolsettable=SimpleNamespace(find_many=AsyncMock(return_value=[toolset])))
tables: Final = SimpleNamespace(
litellm_mcptoolsettable=SimpleNamespace(find_many=AsyncMock(return_value=[toolset]))
)
monkeypatch.setattr(proxy_server, "prisma_client", SimpleNamespace(db=tables, writer_db=tables))
monkeypatch.setattr(proxy_server, "user_api_key_cache", DualCache())
current = UserAPIKeyAuth(
@ -10984,10 +11447,13 @@ async def test_catalog_refresh_uses_current_virtual_key_policy_and_keeps_session
caller.mcp_toolset_id = "session-toolset"
refreshed = await MCPRequestHandler.refresh_catalog_authority(caller)
reload_key.assert_awaited_once_with("owned-key-hash", check_db_only=True)
assert refreshed.object_permission == permission.model_copy(update={
"mcp_servers": ["current-server"], "mcp_toolsets": [],
"mcp_tool_permissions": {"current-server": ["read"]},
})
assert refreshed.object_permission == permission.model_copy(
update={
"mcp_servers": ["current-server"],
"mcp_toolsets": [],
"mcp_tool_permissions": {"current-server": ["read"]},
}
)
assert refreshed.object_permission_id == "current-policy"
assert (refreshed.team_id, refreshed.org_id, refreshed.project_id, refreshed.user_id) == (
"new-team",
@ -12091,9 +12557,7 @@ async def test_personal_key_guardrail_auth_preserves_native_tool_name_matching_s
from litellm.proxy.mcp_server.mcp_server_manager import global_mcp_server_manager
async def _get_team_object(team_id: str, **_: object) -> LiteLLM_TeamTable:
return _personal_key_team_granting_tools(
team_id, ("math-multiply",) if team_id == "team-b" else ("other",)
)
return _personal_key_team_granting_tools(team_id, ("math-multiply",) if team_id == "team-b" else ("other",))
monkeypatch.setattr(auth_checks, "get_team_object", _get_team_object)
monkeypatch.setattr(auth_checks, "get_org_object", AsyncMock(return_value=None))
@ -12132,9 +12596,7 @@ async def test_personal_key_mcp_billing_inherits_adopted_organization_limits(
)
async def _get_team_object(team_id: str, **_: object) -> LiteLLM_TeamTable:
return _personal_key_team_granting_tools(team_id, ("divide",)).model_copy(
update={"organization_id": "org-1"}
)
return _personal_key_team_granting_tools(team_id, ("divide",)).model_copy(update={"organization_id": "org-1"})
org_loader: Final = AsyncMock(return_value=organization)
monkeypatch.setattr(auth_checks, "get_team_object", _get_team_object)
@ -12180,16 +12642,12 @@ async def test_personal_key_mcp_billing_keeps_key_organization_limits(
)
async def _get_team_object(team_id: str, **_: object) -> LiteLLM_TeamTable:
return _personal_key_team_granting_tools(team_id, ("divide",)).model_copy(
update={"organization_id": "org-1"}
)
return _personal_key_team_granting_tools(team_id, ("divide",)).model_copy(update={"organization_id": "org-1"})
org_loader: Final = AsyncMock(return_value=organization)
monkeypatch.setattr(auth_checks, "get_team_object", _get_team_object)
monkeypatch.setattr(auth_checks, "get_org_object_for_request", org_loader)
auth: Final = personal_key_owner_teams.model_copy(
update={"org_id": "org-k", "organization_rpm_limit": 9}
)
auth: Final = personal_key_owner_teams.model_copy(update={"org_id": "org-k", "organization_rpm_limit": 9})
server: Final = global_mcp_server_manager.registry["srv-b"]
with patch.object(global_mcp_server_manager, "get_mcp_server_from_tool_name", return_value=server):
@ -12321,9 +12779,7 @@ async def test_personal_key_mcp_billing_uses_fresh_team_and_membership_for_stric
}
)
async def _get_team_membership(
user_id: str, team_id: str, **kwargs: object
) -> LiteLLM_TeamMembership:
async def _get_team_membership(user_id: str, team_id: str, **kwargs: object) -> LiteLLM_TeamMembership:
database_value: Final = bool(kwargs.get("check_db_only"))
return LiteLLM_TeamMembership(
user_id=user_id,

View file

@ -3,6 +3,7 @@ import contextlib
import contextvars
import json
import os
from collections.abc import Coroutine, Iterator, Sequence
from datetime import datetime, timedelta
from types import SimpleNamespace
from typing import Callable, Final, Literal
@ -669,6 +670,75 @@ def test_prepare_mcp_server_headers_legacy_delegate_strips_repeated_admission_ke
assert extra_headers is None
@pytest.mark.parametrize(
"caller_fields",
[{}, {"user_id": "jwt-user", "api_key": "upstream-jwt", "jwt_claims": {"sub": "jwt-user"}}],
ids=["anonymous", "admitted-jwt-user"],
)
def test_prepare_mcp_server_headers_legacy_delegate_forwards_marked_upstream_authorization(
caller_fields: dict[str, object],
) -> None:
"""When admission marked the caller's Authorization as the upstream token, a legacy delegate forwards
it, including a JWT that also admitted the caller and would otherwise be stripped as an admission JWT."""
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.mcp_server.server import (
_prepare_mcp_server_headers,
)
server = MCPServer(
server_id="legacy-delegate-marked",
name="legacy-delegate",
transport=MCPTransport.http,
auth_type=MCPAuth.oauth2,
delegate_auth_to_upstream=True,
)
caller: Final = UserAPIKeyAuth.model_validate(caller_fields).model_copy(
update={"mcp_forwards_caller_authorization": True}
)
server_auth_header, extra_headers = _prepare_mcp_server_headers(
server=server,
mcp_server_auth_headers=None,
mcp_auth_header=None,
oauth2_headers={"Authorization": "Bearer upstream-jwt"},
raw_headers={"authorization": "Bearer upstream-jwt"},
user_api_key_auth=caller,
)
assert server_auth_header is None
assert extra_headers == {"Authorization": "Bearer upstream-jwt"}
def test_prepare_mcp_server_headers_forwarding_marker_is_ignored_off_legacy_delegates() -> None:
"""The forwarding marker only applies to legacy delegates, so it can never carry a LiteLLM key to an
oauth_delegate upstream."""
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.mcp_server.server import (
_prepare_mcp_server_headers,
)
server = MCPServer(
server_id="oauth-delegate-marked",
name="oauth-delegate",
transport=MCPTransport.http,
auth_type=MCPAuth.oauth_delegate,
)
caller: Final = UserAPIKeyAuth(api_key="sk-litellm-key").model_copy(
update={"mcp_forwards_caller_authorization": True}
)
_, extra_headers = _prepare_mcp_server_headers(
server=server,
mcp_server_auth_headers=None,
mcp_auth_header=None,
oauth2_headers={"Authorization": "Bearer sk-litellm-key"},
raw_headers={"authorization": "Bearer sk-litellm-key"},
user_api_key_auth=caller,
)
assert extra_headers is None
def test_prepare_mcp_server_headers_m2m_skips_authorization_from_raw_extra_headers():
"""M2M must not merge caller Authorization from raw_headers when extra_headers lists it."""
try:
@ -7265,15 +7335,19 @@ async def test_legacy_delegate_bare_token_is_not_probed_upstream(): # test-qual
@pytest.mark.asyncio
async def test_legacy_delegate_dual_credentials_are_not_probed_upstream(): # test-quality-ok: this removed security-sensitive egress has no return value; non-invocation is the contract
"""A separate upstream bearer never triggers the removed legacy probe."""
@pytest.mark.usefixtures("httpx_transport")
async def test_legacy_delegate_dual_credentials_probe_only_the_upstream_token() -> None:
"""With a separate x-litellm-api-key, the preflight probes the upstream with its own token and never
with the LiteLLM key."""
import respx
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.mcp_server.server import (
_check_passthrough_upstream_auth,
)
server = _delegate_auth_mcp_server()
scope = _delegate_scope(
server: Final = _delegate_auth_mcp_server()
scope: Final = _delegate_scope(
[
(b"x-litellm-api-key", b"sk-litellm-proxy-key"),
(b"authorization", b"Bearer upstream-token"),
@ -7281,15 +7355,14 @@ async def test_legacy_delegate_dual_credentials_are_not_probed_upstream(): # te
)
with (
_patch_delegate_resolver(server, "delegate_test"),
patch( # test-quality-ok: isolate authorized-server resolution so this test targets the preflight boundary
"litellm.proxy.mcp_server.operations._get_allowed_mcp_servers",
new=AsyncMock(return_value=[server]),
),
patch( # test-quality-ok: the removed probe call is the security regression under test
"litellm.proxy.mcp_server.server._probe_upstream_auth",
new=AsyncMock(),
) as probe,
respx.mock as upstream,
):
probe: Final = upstream.post(server.url).respond(200)
await _check_passthrough_upstream_auth(
scope=scope,
user_api_key_auth=UserAPIKeyAuth(user_id="admitted-user"),
@ -7300,7 +7373,160 @@ async def test_legacy_delegate_dual_credentials_are_not_probed_upstream(): # te
raw_headers={"x-litellm-api-key": "sk-litellm-proxy-key", "Authorization": "Bearer upstream-token"},
)
probe.assert_not_awaited()
assert probe.call_count == 1
sent_headers: Final = probe.calls[0].request.headers
assert sent_headers["Authorization"] == "Bearer upstream-token"
assert all("sk-litellm-proxy-key" not in value for value in sent_headers.values())
def _expired_delegate_token_preflight(mcp_servers: Sequence[str]) -> Coroutine[object, object, None]:
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.mcp_server.server import (
_check_passthrough_upstream_auth,
)
return _check_passthrough_upstream_auth(
scope=_delegate_scope([(b"authorization", b"Bearer expired-upstream-token")]),
user_api_key_auth=UserAPIKeyAuth().model_copy(update={"mcp_forwards_caller_authorization": True}),
mcp_servers=list(mcp_servers),
client_ip=None,
oauth2_headers={"Authorization": "Bearer expired-upstream-token"},
mcp_server_auth_headers=None,
raw_headers={"authorization": "Bearer expired-upstream-token"},
)
@contextlib.contextmanager
def _rejecting_delegate_upstream() -> Iterator[None]:
server: Final = _delegate_auth_mcp_server()
with (
_patch_delegate_resolver(server, "delegate_test"),
patch( # test-quality-ok: isolate authorized-server resolution so this test targets the preflight boundary
"litellm.proxy.mcp_server.operations._get_allowed_mcp_servers",
new=AsyncMock(return_value=[server]),
),
patch( # test-quality-ok: stands in for the upstream rejecting an expired token
"litellm.proxy.mcp_server.server._probe_upstream_auth",
new=AsyncMock(return_value=(401, 'Bearer error="invalid_token"')),
),
):
yield
@pytest.mark.asyncio
async def test_legacy_delegate_expired_upstream_token_gets_reauthorization_challenge() -> None:
"""An expired upstream token on a single-server route gets the RFC 9728 re-authorization challenge
for the requested name, so the client signs in again instead of seeing an empty tool list."""
with _rejecting_delegate_upstream(), pytest.raises(HTTPException) as exc_info:
await _expired_delegate_token_preflight(["delegate_test"])
assert exc_info.value.status_code == 401
challenge: Final = exc_info.value.headers["www-authenticate"]
assert 'error="invalid_token"' in challenge
assert (
'resource_metadata="http://localhost:4000/.well-known/oauth-protected-resource/mcp/delegate_test"' in challenge
)
@pytest.mark.asyncio
async def test_legacy_delegate_multi_server_scope_is_not_failed_by_one_upstream_token() -> None:
"""A multi-server scope is not probed, so one rejected upstream token cannot fail the whole connect."""
with _rejecting_delegate_upstream():
assert await _expired_delegate_token_preflight(["delegate_test", "other_server"]) is None
@pytest.mark.asyncio
@pytest.mark.usefixtures("httpx_transport")
async def test_legacy_delegate_probe_checks_the_callers_token_not_a_static_authorization() -> None:
"""The probe is what checks a delegate caller's own token, so a static Authorization the server sends on
tool calls cannot stand in for it: a junk token still gets the re-authorization challenge."""
import respx
server: Final = _delegate_auth_mcp_server().model_copy(
update={"static_headers": {"Authorization": "Bearer operator-token"}}
)
with (
_patch_delegate_resolver(server, "delegate_test"),
patch( # test-quality-ok: isolate authorized-server resolution so this test targets the preflight boundary
"litellm.proxy.mcp_server.operations._get_allowed_mcp_servers",
new=AsyncMock(return_value=[server]),
),
respx.mock as upstream,
):
upstream.post(server.url).mock(
side_effect=lambda request: httpx.Response(
200 if request.headers["Authorization"] == "Bearer operator-token" else 401
)
)
outcome: Final = await asyncio.gather(
_expired_delegate_token_preflight(["delegate_test"]), return_exceptions=True
)
assert getattr(outcome[0], "status_code", None) == 401
@pytest.mark.asyncio
@pytest.mark.usefixtures("httpx_transport")
@pytest.mark.parametrize(
"requested, expected_status",
[(("delegate_test",), 401), (("delegate_test", "other_server"), None)],
ids=["single-server", "two-requested-one-reachable"],
)
async def test_modern_interaction_probes_legacy_delegate_only_for_a_single_requested_server(
requested: tuple[str, ...], expected_status: int | None
) -> None:
"""The single-server rule follows what the caller asked for, not what it can reach, so a request
naming two servers is not failed by one expired delegate token."""
import respx
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.mcp_server import server
from litellm.proxy.mcp_server.contracts import OperationContext
target: Final = _delegate_auth_mcp_server()
scope: Final = _delegate_scope(
[
(b"authorization", b"Bearer expired-upstream-token"),
(b"mcp-protocol-version", b"2026-07-28"),
(b"mcp-method", b"tools/call"),
(b"mcp-name", b"delegate-add"),
]
)
context: Final = OperationContext(
_caller=UserAPIKeyAuth().model_copy(update={"mcp_forwards_caller_authorization": True}),
mcp_servers=requested,
oauth2_headers={"Authorization": "Bearer expired-upstream-token"},
raw_headers={"authorization": "Bearer expired-upstream-token"},
)
body: Final = json.dumps(
{
"jsonrpc": "2.0",
"id": 1,
"method": "tools/call",
"params": {
"name": "delegate-add",
"arguments": {"a": 2, "b": 3},
"_meta": {
"io.modelcontextprotocol/protocolVersion": "2026-07-28",
"io.modelcontextprotocol/clientCapabilities": {},
},
},
}
).encode()
with (
_patch_delegate_resolver(target, "delegate_test"),
patch( # test-quality-ok: isolate authorized-server resolution so this test targets the preflight boundary
"litellm.proxy.mcp_server.operations._get_allowed_mcp_servers",
new=AsyncMock(return_value=[target]),
),
respx.mock as upstream,
):
upstream.post(target.url).respond(401, headers={"www-authenticate": 'Bearer error="invalid_token"'})
outcome: Final = await asyncio.gather(
server._preflight_modern_interaction(scope, body, context), return_exceptions=True
)
assert getattr(outcome[0], "status_code", None) == expected_status
@pytest.mark.parametrize(
@ -7537,39 +7763,47 @@ async def test_delegate_not_probed_when_named_only_via_server_id():
@pytest.mark.asyncio
async def test_delegate_probe_not_fanned_out_to_access_group_members():
"""A single access-group name passes the one-target route gate but must not fan
the delegate probe out to group-expanded member servers; the group name resolves
to no server under admission's resolver, so no probe fires."""
@pytest.mark.usefixtures("httpx_transport")
async def test_delegate_probe_not_fanned_out_to_access_group_members() -> None:
"""A single access-group name resolves to no server under admission's resolver, so the delegate probe is
not fanned out to the group's members and one member rejecting the token cannot fail the whole group."""
import respx
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.mcp_server.server import (
_check_passthrough_upstream_auth,
)
group_member = _delegate_auth_mcp_server()
rejecting: Final = _delegate_auth_mcp_server()
accepting: Final = _delegate_auth_mcp_server("delegate-2").model_copy(
update={"name": "delegate_two", "url": "http://upstream-two:9402/mcp"}
)
with (
_patch_delegate_resolver(group_member, "delegate_test"),
patch(
_patch_delegate_resolver(rejecting, "delegate_test"),
patch( # test-quality-ok: isolate authorized-server resolution so this test targets the preflight boundary
"litellm.proxy.mcp_server.operations._get_allowed_mcp_servers",
new=AsyncMock(return_value=[group_member]),
new=AsyncMock(return_value=[rejecting, accepting]),
),
patch(
"litellm.proxy.mcp_server.server._probe_upstream_auth",
new=AsyncMock(return_value=(401, None)),
) as probe,
respx.mock as upstream,
):
await _check_passthrough_upstream_auth(
scope=_delegate_scope([(b"authorization", b"Bearer bogus-token")]),
user_api_key_auth=UserAPIKeyAuth(),
mcp_servers=["prod_tools_group"],
client_ip=None,
oauth2_headers=None,
mcp_server_auth_headers=None,
raw_headers=None,
upstream.post(rejecting.url).respond(401, headers={"www-authenticate": 'Bearer error="invalid_token"'})
upstream.post(accepting.url).respond(200)
assert (
await _check_passthrough_upstream_auth(
scope=_delegate_scope(
[(b"x-litellm-api-key", b"sk-litellm-proxy-key"), (b"authorization", b"Bearer upstream-token")]
),
user_api_key_auth=UserAPIKeyAuth(user_id="admitted-user"),
mcp_servers=["prod_tools_group"],
client_ip=None,
oauth2_headers={"Authorization": "Bearer upstream-token"},
mcp_server_auth_headers=None,
raw_headers={"x-litellm-api-key": "sk-litellm-proxy-key", "Authorization": "Bearer upstream-token"},
)
is None
)
probe.assert_not_awaited()
assert upstream.calls.call_count == 0
@pytest.mark.asyncio
@ -8846,8 +9080,7 @@ async def test_get_allowed_mcp_servers_from_mcp_server_names_unknown_name_fails_
]
with patch(
"litellm.proxy.mcp_server.auth.user_api_key_auth_mcp."
"MCPRequestHandler.get_mcp_servers_from_access_groups",
"litellm.proxy.mcp_server.auth.user_api_key_auth_mcp.MCPRequestHandler.get_mcp_servers_from_access_groups",
new_callable=AsyncMock,
return_value=[],
):
@ -8905,8 +9138,7 @@ async def test_get_allowed_mcp_servers_from_mcp_server_names_known_alias_returns
]
with patch(
"litellm.proxy.mcp_server.auth.user_api_key_auth_mcp."
"MCPRequestHandler.get_mcp_servers_from_access_groups",
"litellm.proxy.mcp_server.auth.user_api_key_auth_mcp.MCPRequestHandler.get_mcp_servers_from_access_groups",
new_callable=AsyncMock,
return_value=[],
):
@ -8938,8 +9170,7 @@ async def test_get_allowed_mcp_servers_from_mcp_server_names_mixed_known_and_unk
]
with patch(
"litellm.proxy.mcp_server.auth.user_api_key_auth_mcp."
"MCPRequestHandler.get_mcp_servers_from_access_groups",
"litellm.proxy.mcp_server.auth.user_api_key_auth_mcp.MCPRequestHandler.get_mcp_servers_from_access_groups",
new_callable=AsyncMock,
return_value=[],
):
@ -8971,8 +9202,7 @@ async def test_get_allowed_mcp_servers_from_mcp_server_names_access_group_resolv
]
with patch(
"litellm.proxy.mcp_server.auth.user_api_key_auth_mcp."
"MCPRequestHandler.get_mcp_servers_from_access_groups",
"litellm.proxy.mcp_server.auth.user_api_key_auth_mcp.MCPRequestHandler.get_mcp_servers_from_access_groups",
new_callable=AsyncMock,
return_value=["id-b"],
):
@ -9240,15 +9470,11 @@ async def test_stateful_mcp_tracing_follows_each_message_exclusion(
is_otel_v2_enabled.cache_clear()
notified: Final = asyncio.Event()
async def call_tool(
ctx: ServerRequestContext[object, object], params: CallToolRequestParams
) -> CallToolResult:
async def call_tool(ctx: ServerRequestContext[object, object], params: CallToolRequestParams) -> CallToolResult:
with tracer.start_as_current_span("tool-handler"):
return CallToolResult(content=[TextContent(type="text", text="done")])
async def notification(
ctx: ServerRequestContext[object, object], params: NotificationParams
) -> HandlerResult:
async def notification(ctx: ServerRequestContext[object, object], params: NotificationParams) -> HandlerResult:
with tracer.start_as_current_span("notification-handler"):
notified.set()
return None
@ -9266,51 +9492,85 @@ async def test_stateful_mcp_tracing_follows_each_message_exclusion(
try:
async with manager.run():
async with httpx2.AsyncClient(transport=httpx2.ASGITransport(app=app)) as client:
initialized: Final = await client.post(initial_url, headers=headers, json={
"jsonrpc": "2.0", "id": 0, "method": "initialize", "params": {
"protocolVersion": "2025-11-25", "capabilities": {},
"clientInfo": {"name": "trace-exclusion", "version": "1"},
},
})
assert initialized.status_code == 200, initialized.text
assert any(
span.name == "initialize" for span in exporter.get_finished_spans()
) is not initialized_excluded
exporter.clear()
session_headers: Final = {**headers, "MCP-Session-Id": initialized.headers["mcp-session-id"]}
ready: Final = await client.post(initial_url, headers=session_headers,
json={"jsonrpc": "2.0", "method": "notifications/initialized"})
assert ready.status_code == 202, ready.text
response: Final = await client.post(current_url, headers=session_headers, json={
"jsonrpc": "2.0", "id": 1, "method": "tools/call", "params": {
"name": "probe", "arguments": {}, "_meta": {
"traceparent": "00-11111111111111111111111111111111-2222222222222222-01",
initialized: Final = await client.post(
initial_url,
headers=headers,
json={
"jsonrpc": "2.0",
"id": 0,
"method": "initialize",
"params": {
"protocolVersion": "2025-11-25",
"capabilities": {},
"clientInfo": {"name": "trace-exclusion", "version": "1"},
},
},
})
)
assert initialized.status_code == 200, initialized.text
assert (
any(span.name == "initialize" for span in exporter.get_finished_spans()) is not initialized_excluded
)
exporter.clear()
session_headers: Final = {**headers, "MCP-Session-Id": initialized.headers["mcp-session-id"]}
ready: Final = await client.post(
initial_url, headers=session_headers, json={"jsonrpc": "2.0", "method": "notifications/initialized"}
)
assert ready.status_code == 202, ready.text
response: Final = await client.post(
current_url,
headers=session_headers,
json={
"jsonrpc": "2.0",
"id": 1,
"method": "tools/call",
"params": {
"name": "probe",
"arguments": {},
"_meta": {
"traceparent": "00-11111111111111111111111111111111-2222222222222222-01",
},
},
},
)
assert response.status_code == 200 and '"text":"done"' in response.text, response.text
names: Final = {span.name for span in exporter.get_finished_spans()}
assert ("tools/call probe" in names) is initialized_excluded
assert ("tool-handler" in names) is initialized_excluded
exporter.clear()
notification_response: Final = await client.post(current_url, headers=session_headers,
json={"jsonrpc": "2.0", "method": "notifications/exclusion-probe"})
notification_response: Final = await client.post(
current_url,
headers=session_headers,
json={"jsonrpc": "2.0", "method": "notifications/exclusion-probe"},
)
assert notification_response.status_code == 202, notification_response.text
await notified.wait()
assert any(
span.name == "notification-handler" for span in exporter.get_finished_spans()
) is initialized_excluded
assert (
any(span.name == "notification-handler" for span in exporter.get_finished_spans())
is initialized_excluded
)
exporter.clear()
modern_response: Final = await client.post(current_url, headers={
"Accept": "application/json, text/event-stream", "MCP-Protocol-Version": "2026-07-28",
"MCP-Method": "tools/call", "MCP-Name": "probe",
}, json={
"jsonrpc": "2.0", "id": 2, "method": "tools/call", "params": {
"name": "probe", "arguments": {}, "_meta": {
PROTOCOL_VERSION_META_KEY: "2026-07-28", CLIENT_CAPABILITIES_META_KEY: {},
modern_response: Final = await client.post(
current_url,
headers={
"Accept": "application/json, text/event-stream",
"MCP-Protocol-Version": "2026-07-28",
"MCP-Method": "tools/call",
"MCP-Name": "probe",
},
json={
"jsonrpc": "2.0",
"id": 2,
"method": "tools/call",
"params": {
"name": "probe",
"arguments": {},
"_meta": {
PROTOCOL_VERSION_META_KEY: "2026-07-28",
CLIENT_CAPABILITIES_META_KEY: {},
},
},
},
})
)
assert modern_response.status_code == 200, modern_response.text
assert '"text":"done"' in modern_response.text, modern_response.text
modern_names: Final = {span.name for span in exporter.get_finished_spans()}
@ -9413,8 +9673,7 @@ async def test_get_allowed_mcp_servers_includes_active_servers_submitted_by_user
return_value={"submitted-1": submitted_server},
),
patch(
"litellm.proxy.mcp_server.auth.user_api_key_auth_mcp."
"MCPRequestHandler.get_allowed_mcp_servers",
"litellm.proxy.mcp_server.auth.user_api_key_auth_mcp.MCPRequestHandler.get_allowed_mcp_servers",
AsyncMock(return_value=[]),
),
patch("litellm.proxy.proxy_server.prisma_client", MagicMock()),
@ -10374,9 +10633,7 @@ async def test_handle_list_tools_all_faults_on_continuation_stays_degraded(
outcomes={"a": ServerListFault(tag="upstream_error", status_code=429)},
)
result: Final = await _list_tools_with_listing(
_mcp_request_ctx, listing, PaginatedRequestParams(cursor="cursor-1")
)
result: Final = await _list_tools_with_listing(_mcp_request_ctx, listing, PaginatedRequestParams(cursor="cursor-1"))
assert isinstance(result, ListToolsResult)
assert result.tools == []

View file

@ -9400,6 +9400,14 @@ class TestApplyStreamUsageTracking:
assert data["_litellm_strip_stream_usage"] is True
assert data["_litellm_stream_usage_injected"] is True
def test_default_injects_usage_and_marks_strip_for_text_completions(self):
data: Final = {"stream": True, "model": "gpt-5.4-nano"}
_apply_stream_usage_tracking(data=data, general_settings={}, route_type="atext_completion")
assert data["stream_options"] == {"include_usage": True}
assert data["_litellm_strip_stream_usage"] is True
def test_default_preserves_other_client_stream_options_keys(self):
data = {"stream": True, "stream_options": {"include_obfuscation": True}}

View file

@ -1203,24 +1203,35 @@ def test_body_snapshot_excludes_team_callback_credentials() -> None:
}, proxy_request
def test_body_snapshot_keeps_the_stream_usage_injection_marker_out_of_the_stored_request() -> None:
@pytest.mark.parametrize(
("route_type", "request_input"),
[
("acompletion", {"messages": [{"role": "user", "content": "hi"}]}),
("atext_completion", {"prompt": "hi"}),
],
)
def test_body_snapshot_keeps_the_stream_usage_injection_marker_out_of_the_stored_request(
route_type: str, request_input: dict[str, object]
) -> None:
from litellm.proxy.common_request_processing import stream_usage_tracking_updates
from litellm.proxy.litellm_pre_call_utils import refresh_proxy_server_request_body_snapshot
from litellm.types.litellm_params import STREAM_USAGE_INJECTED_FIELD
request: Final = {"model": "gpt-5-mini", "messages": [{"role": "user", "content": "hi"}], "stream": True}
request: Final = {"model": "gpt-5-mini", **request_input, "stream": True}
proxy_request: Final = {"body": dict(request)}
data: Final = {
**request,
**stream_usage_tracking_updates(request, {}, "acompletion", lambda: True),
**stream_usage_tracking_updates(request, {}, route_type, lambda: True),
"proxy_server_request": proxy_request,
}
refresh_proxy_server_request_body_snapshot(data)
assert data[STREAM_USAGE_INJECTED_FIELD] is True
assert data["_litellm_strip_stream_usage"] is True
assert proxy_request["body"]["stream_options"] == {"include_usage": True}
assert STREAM_USAGE_INJECTED_FIELD not in proxy_request["body"]
assert "_litellm_strip_stream_usage" not in proxy_request["body"]
@pytest.mark.asyncio

View file

@ -13583,44 +13583,44 @@ def _stream_usage_generator_chunks():
def test_is_injected_stream_usage_artifact():
from litellm.proxy.proxy_server import _is_injected_stream_usage_artifact
from litellm.litellm_core_utils.stream_usage_chunks import is_injected_stream_usage_artifact
from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices, Usage
content_chunk, finish_chunk, empty_choices_usage_chunk = _stream_usage_test_chunks()
assert _is_injected_stream_usage_artifact(empty_choices_usage_chunk) is True
assert is_injected_stream_usage_artifact(empty_choices_usage_chunk) is True
synthetic_final_chunk = ModelResponseStream(model="gpt-5.4-nano")
synthetic_final_chunk.usage = Usage(prompt_tokens=50, completion_tokens=188, total_tokens=238)
assert _is_injected_stream_usage_artifact(synthetic_final_chunk) is True
assert is_injected_stream_usage_artifact(synthetic_final_chunk) is True
azure_prompt_filter_chunk = ModelResponseStream(model="gpt-5.4-nano", choices=[])
assert _is_injected_stream_usage_artifact(azure_prompt_filter_chunk) is True
assert is_injected_stream_usage_artifact(azure_prompt_filter_chunk) is True
assert _is_injected_stream_usage_artifact(content_chunk) is False
assert _is_injected_stream_usage_artifact(finish_chunk) is False
assert is_injected_stream_usage_artifact(content_chunk) is False
assert is_injected_stream_usage_artifact(finish_chunk) is False
content_chunk_with_usage, finish_chunk_with_usage, _ = _stream_usage_test_chunks()
content_chunk_with_usage.usage = Usage(prompt_tokens=50, completion_tokens=188, total_tokens=238)
finish_chunk_with_usage.usage = Usage(prompt_tokens=50, completion_tokens=188, total_tokens=238)
assert _is_injected_stream_usage_artifact(content_chunk_with_usage) is False
assert _is_injected_stream_usage_artifact(finish_chunk_with_usage) is False
assert is_injected_stream_usage_artifact(content_chunk_with_usage) is False
assert is_injected_stream_usage_artifact(finish_chunk_with_usage) is False
assert _is_injected_stream_usage_artifact({"usage": {"prompt_tokens": 1}}) is False
assert is_injected_stream_usage_artifact({"usage": {"prompt_tokens": 1}}) is False
held_back_token_ids_chunk: Final = ModelResponseStream(
model="qwen3", choices=[StreamingChoices(delta=Delta(content=None), token_ids=[16])]
)
assert _is_injected_stream_usage_artifact(held_back_token_ids_chunk) is False
assert is_injected_stream_usage_artifact(held_back_token_ids_chunk) is False
null_token_ids_chunk: Final = ModelResponseStream(
model="qwen3", choices=[StreamingChoices(delta=Delta(content=None), token_ids=None)]
)
assert _is_injected_stream_usage_artifact(null_token_ids_chunk) is True
assert is_injected_stream_usage_artifact(null_token_ids_chunk) is True
unrelated_extra_chunk: Final = ModelResponseStream(
model="qwen3", choices=[StreamingChoices(delta=Delta(content=None), text="")]
)
assert _is_injected_stream_usage_artifact(unrelated_extra_chunk) is True
assert is_injected_stream_usage_artifact(unrelated_extra_chunk) is True
async def _collect_async_data_generator_frames(request_data: dict) -> list:

View file

@ -33,6 +33,7 @@ MAPPED_FILE_INPUT: Final = [{"role": "user", "content": [{"type": "input_file",
CACHED_TEXT: Final = {"type": "input_text", "text": "hi", "cache_control": {"type": "ephemeral"}}
THINKING_TURN: Final = {"role": "assistant", "content": [{"type": "thinking", "thinking": "x", "signature": "s"}]}
REDACTED: Final = {"type": "redacted_thinking", "data": "x"}
STRICT_TOOL: Final = {"name": "b", "input_schema": {"type": "object"}, "strict": True}
TOOL_SEARCH_TYPE: Final = sorted(ANTHROPIC_TOOL_SEARCH_TOOL_TYPES)[0]
@ -192,6 +193,21 @@ def _wrap_nested_level(value: object, depth: int) -> object:
(messages_detectors.tool_custom_field, {"tools": [{"name": "a"}, {"name": "b", "custom": {}}]}, True),
(messages_detectors.tool_custom_field, {"tools": [{"name": "a", "defer_loading": True}]}, False),
(messages_detectors.tool_custom_field, {}, False),
(messages_detectors.strict_tools_retry, {"drop_params": True, "tools": [{"name": "a"}, STRICT_TOOL]}, True),
(
messages_detectors.strict_tools_retry,
{"drop_params": True, "tools": [{**STRICT_TOOL, "strict": False}]},
False,
),
(
messages_detectors.strict_tools_retry,
{"additional_drop_params": ["metadata.user_id"], "tools": [STRICT_TOOL]},
False,
),
(messages_detectors.strict_tools_retry, {"tools": [STRICT_TOOL]}, False),
(messages_detectors.strict_tools_retry, {"drop_params": False, "tools": [STRICT_TOOL]}, False),
(messages_detectors.strict_tools_retry, {"drop_params": True, "tools": [{"name": "a"}]}, False),
(messages_detectors.strict_tools_retry, {"drop_params": True}, False),
(detectors.drop_params, {"drop_params": True}, True),
(detectors.drop_params, {"drop_params": False}, False),
(detectors.drop_params, {"drop_params": " True "}, True),

View file

@ -48,7 +48,14 @@ from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
from litellm.types.litellm_params import ControlOptions
from litellm.types.llms.openai import AllMessageValues, HttpxBinaryResponseContent
from litellm.types.prompts.init_prompts import PromptSpec
from litellm.types.utils import Delta, ModelResponseStream, StandardCallbackDynamicParams, StreamingChoices, Usage
from litellm.types.utils import (
Delta,
ModelResponseStream,
StandardCallbackDynamicParams,
StreamingChoices,
TextCompletionResponse,
Usage,
)
from tests.unit.llms.bedrock.slow_upstream import (
CONVERSE_ANSWER,
STREAM_TIMEOUT_SECONDS,
@ -4136,6 +4143,95 @@ def test_mock_text_completion_stream_and_non_stream_report_the_same_zero_admissi
assert non_stream.usage.prompt_tokens == 0
def _carries_text_or_finish(frame: TextCompletionResponse) -> bool:
return frame.choices[0].text is not None or frame.choices[0].finish_reason is not None
@pytest.mark.asyncio
async def test_mock_atext_completion_stream_hides_the_usage_chunk_the_proxy_asked_for_on_the_clients_behalf():
response: Final = await litellm.atext_completion(
model="openai/gpt-5.4-mini",
prompt="hi",
mock_response="ok then",
api_key="mock",
stream=True,
stream_options={"include_usage": True},
_litellm_strip_stream_usage=True,
)
frames: Final[list[TextCompletionResponse]] = [frame async for frame in response]
assert len(frames) == 4
assert all(_carries_text_or_finish(frame) for frame in frames)
assert all(frame.usage is None for frame in frames)
def _native_completions_reply(outbound: list[dict[str, object]]) -> Callable[[httpx.Request], httpx.Response]:
def respond(request: httpx.Request) -> httpx.Response:
outbound.append(json.loads(request.content))
base: Final = {"id": "cmpl-native", "object": "text_completion", "created": 1, "model": "code-model"}
frames: Final = (
{**base, "choices": [{"index": 0, "text": "def ", "finish_reason": None, "logprobs": None}]},
{**base, "choices": [{"index": 0, "text": "main", "finish_reason": "length", "logprobs": None}]},
{**base, "choices": [], "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}},
)
content: Final = "".join(f"data: {json.dumps(frame)}\n\n" for frame in frames) + "data: [DONE]\n\n"
return httpx.Response(200, headers={"content-type": "text/event-stream"}, content=content.encode())
return respond
def _native_completions_upstream(outbound: list[dict[str, object]]) -> OpenAI:
transport: Final = httpx.MockTransport(_native_completions_reply(outbound))
return OpenAI(api_key="transport-only", http_client=httpx.Client(transport=transport))
def _async_native_completions_upstream(outbound: list[dict[str, object]]) -> AsyncOpenAI:
transport: Final = httpx.MockTransport(_native_completions_reply(outbound))
return AsyncOpenAI(api_key="transport-only", http_client=httpx.AsyncClient(transport=transport))
@pytest.mark.asyncio
async def test_native_atext_completion_stream_hides_the_usage_chunk_the_proxy_asked_for_on_the_clients_behalf():
outbound: Final[list[dict[str, object]]] = []
response: Final = await litellm.atext_completion(
model="text-completion-openai/code-model",
prompt="hi",
client=_async_native_completions_upstream(outbound),
api_key="transport-only",
stream=True,
stream_options={"include_usage": True},
_litellm_strip_stream_usage=True,
)
frames: Final[list[TextCompletionResponse]] = [frame async for frame in response]
assert [sent.get("stream_options") for sent in outbound] == [{"include_usage": True}]
assert [(frame.choices[0].text, frame.choices[0].finish_reason) for frame in frames] == [
("def ", None),
("main", None),
(None, "length"),
]
assert all(frame.usage is None for frame in frames)
def test_native_text_completion_stream_keeps_the_usage_chunk_the_client_asked_for():
frames: Final = list(
litellm.text_completion(
model="text-completion-openai/code-model",
prompt="hi",
client=_native_completions_upstream([]),
api_key="transport-only",
stream=True,
stream_options={"include_usage": True},
)
)
assert len(frames) == 4
assert all(_carries_text_or_finish(frame) for frame in frames[:3])
assert (frames[-1].usage.prompt_tokens, frames[-1].usage.completion_tokens) == (10, 5)
def test_mock_completion_stream_with_model_response():
"""Test that mock_completion correctly handles stream=True with a ModelResponse as mock_response."""
from litellm import completion
@ -4293,7 +4389,6 @@ class TestCallTypesOCR:
def test_stream_chunk_builder_text_completion_combines_text_and_usage():
from litellm.main import stream_chunk_builder_text_completion
from litellm.types.utils import TextCompletionResponse
chunks = [
TextCompletionResponse(

View file

@ -5,10 +5,35 @@ here is the checker output parse and `introduced`, the diff scan that names the
endpoint behind a breach.
"""
import subprocess
from pathlib import Path
from typing import Final
from unittest.mock import patch
import openapi_docs_gate as gate
import pytest
def test_base_counts_stops_when_the_native_build_fails(tmp_path: Path) -> None:
parent: Final = tmp_path / "oa_base"
parent.mkdir()
checker: Final = tmp_path / "check_openapi_docs.py"
checker.write_text("print('POST /v1/test OA001')\n")
with (
patch("openapi_docs_gate.tempfile.mkdtemp", return_value=str(parent)),
patch(
"subprocess.run",
side_effect=(
subprocess.CompletedProcess([], 0, stdout="", stderr=""),
subprocess.CompletedProcess([], 1, stdout="", stderr="build failed"),
subprocess.CompletedProcess([], 0, stdout="", stderr=""),
),
) as run,
patch("signal.signal"),
):
with pytest.raises(SystemExit, match="uv exited 1"):
gate.base_counts("HEAD", repo_root=tmp_path, checker=checker)
assert run.call_count == 3
def test_editing_the_checker_rekeys_the_base_counts(tmp_path: Path) -> None:

View file

@ -222,3 +222,90 @@ def test_sync_text_completion_stream_maps_upstream_failure_after_content() -> No
next(wrapper)
assert (error.value.model, error.value.llm_provider) == ("test-model", "openai")
def _stream_with_trailing_usage_chunk() -> tuple[ModelResponseStream, ...]:
return (
ModelResponseStream(choices=[StreamingChoices(index=0, delta=Delta(content="Hi"))]),
ModelResponseStream(choices=[StreamingChoices(index=0, delta=Delta(), finish_reason="length")]),
ModelResponseStream(
choices=[StreamingChoices(index=0, delta=Delta())],
usage=Usage(prompt_tokens=19, completion_tokens=5, total_tokens=24),
),
)
@pytest.mark.asyncio
@pytest.mark.parametrize("async_stream", (False, True))
@pytest.mark.parametrize(
("usage_injected_by_proxy", "expected_frames", "expected_usage"),
(
(True, [("Hi", None), (None, "length")], [None, None]),
(False, [("Hi", None), (None, "length"), (None, None)], [None, None, (19, 5)]),
),
)
async def test_text_completion_stream_hides_only_the_usage_chunk_the_proxy_asked_for(
async_stream: bool,
usage_injected_by_proxy: bool,
expected_frames: list[tuple[str | None, str | None]],
expected_usage: list[tuple[int, int] | None],
) -> None:
async def chunks() -> AsyncIterator[ModelResponseStream]:
for chunk in _stream_with_trailing_usage_chunk():
yield chunk
wrapper: Final = TextCompletionStreamWrapper(
completion_stream=chunks() if async_stream else iter(_stream_with_trailing_usage_chunk()),
model="gpt-5.4-mini",
stream_options={"include_usage": True},
usage_injected_by_proxy=usage_injected_by_proxy,
)
frames: Final = [frame async for frame in wrapper] if async_stream else list(wrapper)
assert [(frame.choices[0].text, frame.choices[0].finish_reason) for frame in frames] == expected_frames
assert [
None if frame.usage is None else (frame.usage.prompt_tokens, frame.usage.completion_tokens) for frame in frames
] == expected_usage
def test_text_completion_stream_keeps_a_finish_frame_but_hides_the_usage_the_proxy_asked_for() -> None:
finish_with_usage: Final = ModelResponseStream(
choices=[StreamingChoices(index=0, delta=Delta(), finish_reason="stop")],
usage=Usage(prompt_tokens=27, completion_tokens=5, total_tokens=32),
)
wrapper: Final = TextCompletionStreamWrapper(
completion_stream=iter((finish_with_usage,)),
model="claude-haiku-4-5",
stream_options={"include_usage": True},
usage_injected_by_proxy=True,
)
frames: Final = list(wrapper)
assert [(frame.choices[0].finish_reason, frame.usage) for frame in frames] == [("stop", None)]
@pytest.mark.asyncio
@pytest.mark.parametrize("async_stream", (False, True))
async def test_text_completion_stream_hides_the_empty_prompt_filter_frame_when_the_proxy_asked_for_usage(
async_stream: bool,
) -> None:
prompt_filter_frame: Final = ModelResponseStream(id="", choices=[])
upstream: Final = (prompt_filter_frame, *_stream_with_trailing_usage_chunk())
async def chunks() -> AsyncIterator[ModelResponseStream]:
for chunk in upstream:
yield chunk
wrapper: Final = TextCompletionStreamWrapper(
completion_stream=chunks() if async_stream else iter(upstream),
model="gpt-5.4-mini",
stream_options={"include_usage": True},
usage_injected_by_proxy=True,
)
frames: Final = [frame async for frame in wrapper] if async_stream else list(wrapper)
assert [(frame.choices[0].text, frame.choices[0].finish_reason) for frame in frames] == [
("Hi", None),
(None, "length"),
]