mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-09 08:18:12 +00:00
Merge remote-tracking branch 'origin/main' into HEAD
# Conflicts: # litellm/proxy/proxy_server.py # litellm/text_completions/streaming_iterator.py # tests/unit/proxy/test_proxy_server_endpoints_and_startup.py
This commit is contained in:
commit
d7d3054ceb
42 changed files with 2203 additions and 298 deletions
2
.github/workflows/test-integration.yml
vendored
2
.github/workflows/test-integration.yml
vendored
|
|
@ -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
|
||||
|
|
|
|||
35
litellm/litellm_core_utils/stream_usage_chunks.py
Normal file
35
litellm/litellm_core_utils/stream_usage_chunks.py
Normal file
|
|
@ -0,0 +1,35 @@
|
|||
from collections.abc import Iterable
|
||||
from typing import Final
|
||||
|
||||
from litellm.constants import VLLM_CHOICE_FIELDS
|
||||
from litellm.types.utils import ModelResponseStream, StreamingChoices
|
||||
|
||||
|
||||
def is_empty_streaming_choice(choice: StreamingChoices) -> bool:
|
||||
if choice.finish_reason is not None:
|
||||
return False
|
||||
if getattr(choice, "logprobs", None) is not None:
|
||||
return False
|
||||
if not _all_values_none(_vllm_choice_values(choice)):
|
||||
return False
|
||||
delta: Final = choice.delta if hasattr(choice, "delta") else None
|
||||
if delta is None:
|
||||
return True
|
||||
return _all_values_none(delta.model_dump().values())
|
||||
|
||||
|
||||
def _vllm_choice_values(choice: StreamingChoices) -> Iterable[object]:
|
||||
choice_extras: Final = choice.model_extra or {}
|
||||
return (choice_extras.get(field) for field in VLLM_CHOICE_FIELDS)
|
||||
|
||||
|
||||
def _all_values_none(values: Iterable[object]) -> bool:
|
||||
return all(value is None for value in values)
|
||||
|
||||
|
||||
def is_injected_stream_usage_artifact(chunk: object) -> bool:
|
||||
if not isinstance(chunk, ModelResponseStream):
|
||||
return False
|
||||
if chunk.provider_specific_fields is not None:
|
||||
return False
|
||||
return all(is_empty_streaming_choice(choice) for choice in chunk.choices or [])
|
||||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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` "
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
}
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -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}")
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
]
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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``).
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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 | {
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
@ -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
|
||||
|
|
@ -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"
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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 (
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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 == []
|
||||
|
|
|
|||
|
|
@ -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}}
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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),
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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"),
|
||||
]
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue