diff --git a/.github/workflows/test-integration.yml b/.github/workflows/test-integration.yml index 0c5b70a67fe..73159bd64d2 100644 --- a/.github/workflows/test-integration.yml +++ b/.github/workflows/test-integration.yml @@ -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 diff --git a/litellm/litellm_core_utils/stream_usage_chunks.py b/litellm/litellm_core_utils/stream_usage_chunks.py new file mode 100644 index 00000000000..d073ddc2724 --- /dev/null +++ b/litellm/litellm_core_utils/stream_usage_chunks.py @@ -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 []) diff --git a/litellm/llms/anthropic/common_utils.py b/litellm/llms/anthropic/common_utils.py index 3b7fcdea0e4..0bc519187f4 100644 --- a/litellm/llms/anthropic/common_utils.py +++ b/litellm/llms/anthropic/common_utils.py @@ -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 diff --git a/litellm/llms/base_llm/anthropic_messages/transformation.py b/litellm/llms/base_llm/anthropic_messages/transformation.py index bda7015ab2b..d234ac40c3d 100644 --- a/litellm/llms/base_llm/anthropic_messages/transformation.py +++ b/litellm/llms/base_llm/anthropic_messages/transformation.py @@ -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( diff --git a/litellm/llms/bedrock/chat/native_output_fallback.py b/litellm/llms/bedrock/chat/native_output_fallback.py index 6892cd37beb..8069c3cf7e7 100644 --- a/litellm/llms/bedrock/chat/native_output_fallback.py +++ b/litellm/llms/bedrock/chat/native_output_fallback.py @@ -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: diff --git a/litellm/llms/bedrock/common_utils.py b/litellm/llms/bedrock/common_utils.py index 4175479f068..d88d7e6280c 100644 --- a/litellm/llms/bedrock/common_utils.py +++ b/litellm/llms/bedrock/common_utils.py @@ -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` " diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 4178095c9ef..1f25235a651 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -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, ) diff --git a/litellm/main.py b/litellm/main.py index 50ab274a060..0162742de18 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -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 diff --git a/litellm/messages/dispatch.py b/litellm/messages/dispatch.py index 8e6becc60b4..35136c41dc2 100644 --- a/litellm/messages/dispatch.py +++ b/litellm/messages/dispatch.py @@ -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, } ) diff --git a/litellm/proxy/auth/oauth2_check.py b/litellm/proxy/auth/oauth2_check.py index 3633f7b6127..970f046c3d6 100644 --- a/litellm/proxy/auth/oauth2_check.py +++ b/litellm/proxy/auth/oauth2_check.py @@ -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}") diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index e5b0beadebf..2b6471952fd 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -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: diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index e810a5ab329..666269be2eb 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -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 diff --git a/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py b/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py index ecd23d86540..7e052796545 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py +++ b/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py @@ -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, ] diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 8f60311e491..a87fc815402 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -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 diff --git a/litellm/proxy/mcp_server/AGENTS.md b/litellm/proxy/mcp_server/AGENTS.md index 25150c11f1d..b66128454b8 100644 --- a/litellm/proxy/mcp_server/AGENTS.md +++ b/litellm/proxy/mcp_server/AGENTS.md @@ -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 diff --git a/litellm/proxy/mcp_server/auth/user_api_key_auth_mcp.py b/litellm/proxy/mcp_server/auth/user_api_key_auth_mcp.py index cc7080738d0..db640b6d689 100644 --- a/litellm/proxy/mcp_server/auth/user_api_key_auth_mcp.py +++ b/litellm/proxy/mcp_server/auth/user_api_key_auth_mcp.py @@ -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``). diff --git a/litellm/proxy/mcp_server/bridge_token_flow.py b/litellm/proxy/mcp_server/bridge_token_flow.py index 63428dfb28a..dd9be4cb3e1 100644 --- a/litellm/proxy/mcp_server/bridge_token_flow.py +++ b/litellm/proxy/mcp_server/bridge_token_flow.py @@ -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 diff --git a/litellm/proxy/mcp_server/mcp_server_manager.py b/litellm/proxy/mcp_server/mcp_server_manager.py index ceac2b2a8e1..c772e2e9d09 100644 --- a/litellm/proxy/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/mcp_server/mcp_server_manager.py @@ -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) diff --git a/litellm/proxy/mcp_server/server.py b/litellm/proxy/mcp_server/server.py index 32f790a8e81..2105d875b58 100644 --- a/litellm/proxy/mcp_server/server.py +++ b/litellm/proxy/mcp_server/server.py @@ -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 diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 557df3e5450..57ff1a6fe66 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -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, diff --git a/litellm/rust_bridge/catalog.py b/litellm/rust_bridge/catalog.py index fe52f8416d0..629d9098970 100644 --- a/litellm/rust_bridge/catalog.py +++ b/litellm/rust_bridge/catalog.py @@ -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 | { diff --git a/litellm/rust_bridge/messages/detectors.py b/litellm/rust_bridge/messages/detectors.py index 52917126b3b..8edd2707405 100644 --- a/litellm/rust_bridge/messages/detectors.py +++ b/litellm/rust_bridge/messages/detectors.py @@ -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) diff --git a/litellm/text_completions/streaming_iterator.py b/litellm/text_completions/streaming_iterator.py index 7b65bfbd249..6bba56a5afe 100644 --- a/litellm/text_completions/streaming_iterator.py +++ b/litellm/text_completions/streaming_iterator.py @@ -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 diff --git a/litellm/types/mcp_server/mcp_server_manager.py b/litellm/types/mcp_server/mcp_server_manager.py index d29e4db55a4..0d39b0a0616 100644 --- a/litellm/types/mcp_server/mcp_server_manager.py +++ b/litellm/types/mcp_server/mcp_server_manager.py @@ -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 diff --git a/litellm/types/proxy/auth/user_api_key_auth.py b/litellm/types/proxy/auth/user_api_key_auth.py index 2c1c75fbe04..1183c36eda7 100644 --- a/litellm/types/proxy/auth/user_api_key_auth.py +++ b/litellm/types/proxy/auth/user_api_key_auth.py @@ -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) diff --git a/scripts/openapi_docs_gate.py b/scripts/openapi_docs_gate.py index 5f5d1021c22..de340bf84e7 100755 --- a/scripts/openapi_docs_gate.py +++ b/scripts/openapi_docs_gate.py @@ -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. diff --git a/tests/integration/configuration/test_openapi_docs_gate.py b/tests/integration/configuration/test_openapi_docs_gate.py new file mode 100644 index 00000000000..eefa663cb10 --- /dev/null +++ b/tests/integration/configuration/test_openapi_docs_gate.py @@ -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 diff --git a/tests/integration/spend/test_text_completion_stream_usage_cost.py b/tests/integration/spend/test_text_completion_stream_usage_cost.py new file mode 100644 index 00000000000..4d66e60f3e7 --- /dev/null +++ b/tests/integration/spend/test_text_completion_stream_usage_cost.py @@ -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 diff --git a/tests/test_litellm_rust/messages/test_dispatch.py b/tests/test_litellm_rust/messages/test_dispatch.py index 15f58afb4e5..5b11e628b8d 100644 --- a/tests/test_litellm_rust/messages/test_dispatch.py +++ b/tests/test_litellm_rust/messages/test_dispatch.py @@ -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" + ), ) diff --git a/tests/unit/llms/anthropic/messages/native/test_anthropic_messages_handler.py b/tests/unit/llms/anthropic/messages/native/test_anthropic_messages_handler.py index c11f58faec4..5ce6d03f477 100644 --- a/tests/unit/llms/anthropic/messages/native/test_anthropic_messages_handler.py +++ b/tests/unit/llms/anthropic/messages/native/test_anthropic_messages_handler.py @@ -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" diff --git a/tests/unit/llms/anthropic/test_anthropic_common_utils.py b/tests/unit/llms/anthropic/test_anthropic_common_utils.py index b37b2534b73..3ec72f29d34 100644 --- a/tests/unit/llms/anthropic/test_anthropic_common_utils.py +++ b/tests/unit/llms/anthropic/test_anthropic_common_utils.py @@ -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 + ) diff --git a/tests/unit/llms/bedrock/messages/invoke_transformations/test_anthropic_claude_transformation.py b/tests/unit/llms/bedrock/messages/invoke_transformations/test_anthropic_claude_transformation.py index 74f6f3030cb..0bd480dcdbc 100644 --- a/tests/unit/llms/bedrock/messages/invoke_transformations/test_anthropic_claude_transformation.py +++ b/tests/unit/llms/bedrock/messages/invoke_transformations/test_anthropic_claude_transformation.py @@ -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 ( diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py index 8308a7ba386..b5f794cb136 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py @@ -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" diff --git a/tests/unit/proxy/mcp_server/auth/test_user_api_key_auth_mcp.py b/tests/unit/proxy/mcp_server/auth/test_user_api_key_auth_mcp.py index b7795b4d2f0..ee3ef3f5323 100644 --- a/tests/unit/proxy/mcp_server/auth/test_user_api_key_auth_mcp.py +++ b/tests/unit/proxy/mcp_server/auth/test_user_api_key_auth_mcp.py @@ -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, diff --git a/tests/unit/proxy/mcp_server/test_mcp_server_tool_calls_and_headers.py b/tests/unit/proxy/mcp_server/test_mcp_server_tool_calls_and_headers.py index 78c37184d6c..621a7b65005 100644 --- a/tests/unit/proxy/mcp_server/test_mcp_server_tool_calls_and_headers.py +++ b/tests/unit/proxy/mcp_server/test_mcp_server_tool_calls_and_headers.py @@ -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 == [] diff --git a/tests/unit/proxy/test_common_request_processing.py b/tests/unit/proxy/test_common_request_processing.py index 55b913aa7e4..cc1fc5cd6c3 100644 --- a/tests/unit/proxy/test_common_request_processing.py +++ b/tests/unit/proxy/test_common_request_processing.py @@ -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}} diff --git a/tests/unit/proxy/test_litellm_pre_call_utils.py b/tests/unit/proxy/test_litellm_pre_call_utils.py index 00051462b3c..1ce61a8e9b3 100644 --- a/tests/unit/proxy/test_litellm_pre_call_utils.py +++ b/tests/unit/proxy/test_litellm_pre_call_utils.py @@ -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 diff --git a/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py b/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py index 5750d73b926..43288094a17 100644 --- a/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py +++ b/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py @@ -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: diff --git a/tests/unit/rust_bridge/test_detectors.py b/tests/unit/rust_bridge/test_detectors.py index c02d2c12dbf..634e69910b9 100644 --- a/tests/unit/rust_bridge/test_detectors.py +++ b/tests/unit/rust_bridge/test_detectors.py @@ -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), diff --git a/tests/unit/test_main.py b/tests/unit/test_main.py index e10352e0b3d..11104cf7bca 100644 --- a/tests/unit/test_main.py +++ b/tests/unit/test_main.py @@ -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( diff --git a/tests/unit/test_openapi_docs_gate.py b/tests/unit/test_openapi_docs_gate.py index fa8a5b1c7e3..b376178c8a9 100644 --- a/tests/unit/test_openapi_docs_gate.py +++ b/tests/unit/test_openapi_docs_gate.py @@ -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: diff --git a/tests/unit/text_completions/test_streaming_iterator.py b/tests/unit/text_completions/test_streaming_iterator.py index fbcb7f6ae3f..38d7f99e449 100644 --- a/tests/unit/text_completions/test_streaming_iterator.py +++ b/tests/unit/text_completions/test_streaming_iterator.py @@ -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"), + ]